{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "29294f52",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Flexible Slider-Crank\n",
    "This example takes the [rigid slider-crank](../example_slider_crank/notebook.ipynb) and makes one of its links flexible. Hence, this notebooks showcases a simulation that contains flexible bodies and a loop constraint. \n",
    "\n",
    "In this example, an H5 file with the modal frequencies and mode shapes for the flexible body has already been created. This H5 file was created with `FEMBridge`. `FEMBridge` is tool available with enterprise installations of `kdflex` that allows one to do modal analyses on NASTRAN run decks to create H5 files like the one used here. If you do not currently have an enterprise `kdflex` installation, this example will not run, as `FEMBridge` will not be included. More information on `FEMBridge` will be included in an upcoming example. \n",
    "\n",
    "Requirements:\n",
    "- [slider-crank](../example_slider_crank/notebook.ipynb)\n",
    "\n",
    "In this tutorial we will:\n",
    "- Create the Flexible multibody using DataStruct.DataStruct\n",
    "- Initialize the state and constraints\n",
    "- Setup the **kdFlex** scene\n",
    "- Setup the state propagator and models\n",
    "- Run the simulation\n",
    "- Clean up the simulation\n",
    "\n",
    "![](../resources/nb_images/flex_slider_crank.png)\n",
    "\n",
    "For a more in-depth descriptions of **kdflex** concepts see [usage](usage_page).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4cc2204",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "First, let's verify that FEMBridge is installed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a072a5e5",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    from Karana import FEMBridge\n",
    "except:\n",
    "    raise ModuleNotFoundError(\n",
    "        \"This notebook only works with an enterprise installation of kdFlex. One that includes \"\n",
    "        \"FEMBridge. Please ensure your installation is an enterprise installation before \"\n",
    "        \"continuing.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e42db173",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import atexit\n",
    "from numpy.typing import NDArray\n",
    "import numpy as np\n",
    "from copy import deepcopy\n",
    "from typing import cast\n",
    "from pathlib import Path\n",
    "\n",
    "import Karana.Core as kc\n",
    "import Karana.Math as km\n",
    "import Karana.Integrators as ki\n",
    "import Karana.Frame as kf\n",
    "import Karana.Dynamics as kd\n",
    "import Karana.Dynamics.SOADyn_types as kdt\n",
    "import Karana.Scene as ks\n",
    "import Karana.Scene.Scene_types as kst\n",
    "import Karana.KUtils as ku\n",
    "import Karana.Models as kmdl\n",
    "import Karana.FEMBridge.ReadH5PhysicalModalBody as kfem\n",
    "\n",
    "resources = Path(\"../resources\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dfaba58",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Create the Flexible Multibody\n",
    "We will create the multibody using a SOADyn_types.SubGraphDS. We begin by creating the first link."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c880128f",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "fc = kf.FrameContainer(\"root\")\n",
    "\n",
    "# Create a BodyDS for the first link. This can be used to specify all the parameters for the body, and the way\n",
    "# the body attaches to it's parent, i.e., the hinge.\n",
    "link_1 = kdt.BodyDS(\n",
    "    name=\"link_1\",\n",
    "    mass=1.0,\n",
    "    body_to_cm=np.zeros(3),\n",
    "    inertia=np.eye(3),\n",
    "    hinge=kdt.HingeDS(\n",
    "        hinge_type=kd.HingeType.REVOLUTE,\n",
    "        subhinges=[kdt.PinSubhingeDS(prescribed=False, unit_axis=np.array([0.0, 1.0, 0.0]))],  # pyright: ignore - Types are okay\n",
    "    ),\n",
    "    body_to_joint=np.array([-0.25, 0.0, 0.0]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5ebc9b7",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Flexible body creation\n",
    "Now, we create the second link. This is our flexible body. We use the Karana.FEMBridge.ReadH5PhysicalModalBody.ReadH5PhysicalModalBody to do so. This class reads in a Karana.FEMBridge.ReadH5PhysicalModalBody.H5PhysicalModalBodyDS instance, and creates a SOADyn_types.BodyDS instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "20161697",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "hinge = deepcopy(link_1.hinge)\n",
    "\n",
    "ds = kfem.H5PhysicalModalBodyDS(\n",
    "    filename=resources / \"slider_crank\" / \"slider_crank_beam.h5\",\n",
    "    p_node_grid=1,\n",
    "    sensor_nodes={k: f\"sn_{k}\" for k in range(1, 51 + 1, 3)},\n",
    "    force_nodes={},\n",
    "    damping=0.2,\n",
    ")\n",
    "link_2_reader = kfem.ReadH5PhysicalModalBody(ds)\n",
    "link_2 = link_2_reader.getModalBodyDS(\"link_2\", hinge)\n",
    "link_2.inb_to_joint = np.array([0.25, 0.0, 0.0])\n",
    "\n",
    "link_2_context = kdt.BodyWithContextDS(body=link_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "287791dd",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Now, we create the `block` body. This body will be attached to `link_2`. When we attach a body to `link_2`, we need to make sure that the `onode` of it's joint gets the correct nodal matrix. The `onode` of the joint is the node attached to the parent body. Here, we are going to use the point corresponding with node 51 in the original NASTRAN run deck. The `link_2_reader` we created earlier has a convenience method Karana.FEMBridge.ReadH5PhysicalModalBody.ReadH5PhysicalModalBody.addInbNodeToDS for this purpose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7b6431cd",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# We create a block body. The block is the body that will have the slider constraint,\n",
    "# so we also add a constraint node via a NodeDS.\n",
    "block = deepcopy(link_1)\n",
    "block.name = \"block\"\n",
    "block.mass = 1.0\n",
    "block.body_to_joint = np.zeros(3)\n",
    "block.constraint_nodes = [\n",
    "    kdt.NodeDS(\n",
    "        name=\"cn\",\n",
    "        constraint_node=True,\n",
    "        force_node=True,\n",
    "    )\n",
    "]\n",
    "\n",
    "# The block body is attached to link_2, so we need to use the link_2_reader to add the appropriate\n",
    "# nodal matrix and inboard body to joint transform.\n",
    "link_2_reader.addInbNodeToDS(block, 51)\n",
    "link_2_context.children = [kdt.BodyWithContextDS(body=block, children=[])]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea7d2b8a",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Now, create the rest of the SubGraphDS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "eb7977cd",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Make link_1 prescribed\n",
    "link_1.hinge.subhinges[0].prescribed = True\n",
    "\n",
    "# Add visuals.\n",
    "importer = ks.AssimpImporter()\n",
    "block_part = importer.importFrom(str(resources / \"slider_crank\" / \"block_w_mat.glb\"))\n",
    "block.scene_parts = [kst.ScenePartSpecDS.fromScenePartSpec(block_part.parts[0])]\n",
    "del block_part\n",
    "\n",
    "# Now, we combine all the bodies together to form a SubGraphDS. In addition, we add our LoopConstraint\n",
    "# and a constraint frame for the slider constraint.\n",
    "mbody_ds = kdt.SubGraphDS(\n",
    "    name=\"mb\",\n",
    "    base_bodies=[\n",
    "        kdt.BodyWithContextDS(\n",
    "            body=link_1,\n",
    "            children=[link_2_context],\n",
    "        )\n",
    "    ],\n",
    "    loop_constraints=[\n",
    "        kdt.LoopConstraintCutJointDS(\n",
    "            name=\"block_constraint\",\n",
    "            hinge=kdt.HingeDS(\n",
    "                hinge_type=kd.HingeType.SLIDER,\n",
    "                subhinges=[\n",
    "                    kdt.LinearSubhingeDS(prescribed=False, unit_axis=np.array([1.0, 0.0, 0.0]))\n",
    "                ],\n",
    "            ),\n",
    "            get_oframe={\"frame\": \"constraint_frame\"},\n",
    "            get_pframe={\"body\": \"block\", \"node\": \"cn\"},\n",
    "        )\n",
    "    ],\n",
    "    constraint_frames=[\n",
    "        kdt.ConstraintFrameDS(\n",
    "            name=\"constraint_frame\",\n",
    "            translation=np.zeros(3),\n",
    "            unit_quaternion=km.UnitQuaternion(0, 0, 0, 1),\n",
    "            parent=\"newtonian\",\n",
    "        )\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1e2f11c",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Finally, create the Multibody and cleanup any variables we don't need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c5e73088",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "mb = mbody_ds.toMultibody(fc)\n",
    "del mbody_ds, block, link_1, link_2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec6d6508",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Initialize the State and Constraints\n",
    "We are going to initialize the state of this system. This system contains both constraints and flexible bodies. For the constraints, we will need to use a {py:class}`Karana.Dynamics.ConstraintKinematicsSolver`, just as we did in the [rigid slider-crank](../example_slider_crank/notebook.ipynb) example. \n",
    "\n",
    "We could finish the state initialization here. However, when we first start moving, there may be some transient behavior, as the flexible body itself is not in the correct deformed state. Here, we create a simple equilibrium solver to deform the flexible body such that it's initial modal coordinate acceleration is zero. In this case, the solution is trivial, as we are starting from rest, but this showcases how it can be done in general."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "750f0823",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Satisfying the constraints\n",
    "We being by setting an initial guess for the {py:class}`Karana.Dynamics.ConstraintKinematicsSolver`. This is done by setting all joints, body coordinates, and constraint coordinates to 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "61fdb952",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Once all bodies are created, we make the Multibody healthy. This sets up internal\n",
    "# variables for use in dynamics computations.\n",
    "mb.ensureHealthy()\n",
    "\n",
    "# Set all state data to zero\n",
    "mb.resetData()\n",
    "cd = mb.subhingeCoordData()\n",
    "\n",
    "# In addition to the subhinge data, we also need to set the constraint data and body coordinate data.\n",
    "cdc = mb.cutjointCoordData()\n",
    "cdc.setQ(0.0)\n",
    "cdc.setU(0.0)\n",
    "cdc.setUdot(0.0)\n",
    "cdc.setT(0.0)\n",
    "del cdc\n",
    "\n",
    "bcd = mb.bodyCoordData()\n",
    "bcd.setQ(0)\n",
    "bcd.setU(0)\n",
    "bcd.setUdot(0)\n",
    "bcd.setT(0)\n",
    "del bcd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd24a085",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Now, we use the Multibody's ConstraintKinematicsSolver to solve for coordinates and velocities that satisfy the constraints. We want `link_1`'s subhinge to remain unchanged, so we freeze this coordinate. This effectively removes it from the coordinates that the ConstraintKinematicsSolver is allowed to change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a338bc0c",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Freeze link_1's subhinge so that the ConstraintKinematicsSolver doesn't change it.\n",
    "cks = mb.cks()\n",
    "cks.freezeCoord(mb.getBody(\"link_1\").parentHinge().subhinge(0), 0)\n",
    "\n",
    "# Now, we create a ConstraintKinematicsSolver and solve for the initial state.\n",
    "errQ = cks.solveQ()\n",
    "errU = cks.solveU()\n",
    "del cks\n",
    "\n",
    "# Check the errors to ensure we got a good solution.\n",
    "assert abs(errQ) < 1e-6\n",
    "assert abs(errU) < 1e-6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf3e6dfc",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Now, let's ensure everything is ready."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "43158f3d",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Before we finalize everything (see below), we need to set the gravity. In this example, there is no gravity,\n",
    "# so we just set this to zero.\n",
    "mb.setUniformGravAccel(np.zeros(3))\n",
    "\n",
    "assert kc.allReady()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63b2049a",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Equilibrium\n",
    "As mentioned earlier, we solve for the flexible body's equilibrium. Here, the solution is trivial since we are starting from rest (so the solution should be a vector of all zeros). Regardless, we include it here to showcase how solving for equilibrium can be done.\n",
    "\n",
    "To solve for equilibrium we use a {py:class}`Karana.Dynamics.NonlinearSolver`. These solvers will be covered in detail in an upcoming example. For now, it is sufficient to note that they take in a function that takes in a guess for the current solution, `x`, and outputs values, `v` for that guess. The values must be modified in-place, which is why the code below uses `v[:] = ...` rather than `v = ...`. Here, our equilibrium function is simple: it takes in a guess for the flexible body coordinates, `q`, of `link_2`, runs the forward dynamics, and outputs the modal accelerations of `link_2`. The nonlinear solver tries to drive these outputs to zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ca3c3136",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "link_2_bd = cast(kd.PhysicalBody, mb.getBody(\"link_2\"))\n",
    "\n",
    "\n",
    "def equilibrium(x: NDArray, v: NDArray) -> None:\n",
    "    \"\"\"Solve for equilibrium.\n",
    "\n",
    "    The incoming x values will be the new Q states of the beam.\n",
    "    The output v values should be the accelerations of the beam as a\n",
    "    result of the incoming Q values.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : NDArray\n",
    "        New guess for beam Q values.\n",
    "    v : NDArray\n",
    "        Udot values of the beam after running the dynamics.\n",
    "    \"\"\"\n",
    "    link_2_bd.setQ(x)\n",
    "    kd.Algorithms.evalForwardDynamics(mb)\n",
    "    v[:] = link_2_bd.getUdot()\n",
    "\n",
    "\n",
    "# Create a Nonlinear solver that uses the equilibrium function\n",
    "solver = kd.NonlinearSolver(\"solver\", link_2_bd.nQ(), link_2_bd.nU(), equilibrium)\n",
    "solver.createSolver(kd.SolverType.HYBRID_NON_LINEAR_SOLVER)\n",
    "q = np.array(link_2_bd.getQ())\n",
    "\n",
    "# Solve for the modal coordinates that keep the modal accelerations zero\n",
    "solver.solve(q)\n",
    "link_2_bd.setQ(q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5644812",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Similar to the constraint solution, we can again check that the solution here is sufficient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "983fba7c",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "udot = np.empty(q.size)\n",
    "equilibrium(q, udot)\n",
    "assert np.linalg.norm(udot) < 1e-6\n",
    "\n",
    "del link_2_bd\n",
    "del solver"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4745435b",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Setup the kdFlex Scene\n",
    "Here we add graphics for the scene. We will use a combination of stick parts and meshes. For the flexible body, we will use a set of sensor nodes so we can also visualize deformation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "771bdce5",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <div style=\"height:300px; resize:vertical; overflow: auto;\">\n",
       "          <iframe src=\"http://localhost:39965\" style=\"width:100%; height:100%; border:0; display:block;\"></iframe>\n",
       "        </div>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cleanup_graphics, web_scene = mb.setupGraphics(port=0)\n",
    "web_scene.defaultCamera().pointCameraAt(\n",
    "    np.array([1.0, -2.0, 1.5]), np.array([1.0, 0.0, 0.0]), np.array([0.0, 0.0, 1.0])\n",
    ")\n",
    "del web_scene\n",
    "\n",
    "# Add stick parts\n",
    "mb.createStickParts()\n",
    "scene = mb.getScene()\n",
    "\n",
    "# Remove stick parts from the block. We added a nicer mesh when we defined the SubGraphDS\n",
    "# earlier on. In addition, we remove some of the stick parts for the constraint. The linear subhinge\n",
    "# for the constraint has an update function on the scene, so this must also be removed.\n",
    "nd = scene.getNodesAttachedToFrame(mb.enabledConstraints()[0].hinge().pframe(), False)[0]\n",
    "scene.update_callbacks.pop(\"stick_block_constraint_pframe_part_update\")\n",
    "kc.discard(nd)\n",
    "nd = scene.getNodesAttachedToFrame(cast(kd.PhysicalBody, mb.getBody(\"block\")).pnode(), False)[0]\n",
    "kc.discard(nd)\n",
    "\n",
    "# Rather than use the stick parts for the link_2 body, we will use a set of sensor nodes.\n",
    "# These sensor nodes will deform as the flexible body itself deforms.\n",
    "ndl = scene.getNodesAttachedToFrame(cast(kd.PhysicalBody, mb.getBody(\"link_2\")), False)\n",
    "del scene, ndl\n",
    "\n",
    "# Let's also add a nice mesh for the canal that the block slides in\n",
    "p = importer.importFrom(str(resources / \"slider_crank\" / \"canal_w_mat.glb\")).parts[0]  # pyright: ignore - importer will exist if graphics is on\n",
    "part = ks.ProxyScenePart(\"canal\", mb.getScene(), p.geometry, p.material)\n",
    "part.setTranslation(np.array([0.5 + 0.25, 0.0, 0.0]))\n",
    "part.setUnitQuaternion(p.transform.getUnitQuaternion())\n",
    "part.attachTo(fc.root())\n",
    "del p\n",
    "del part\n",
    "del importer  # pyright: ignore"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9505d2d0",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Setup the State Propagator and Models\n",
    "As in other examples, we use a {py:class}`Karana.Dynamics.StatePropagator` to drive the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5190c94",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Karana.Models.UpdateProxyScene at 0x7b281a5014b0>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set up state propagator.\n",
    "# Use cvode_stiff as the integrator.\n",
    "sp = kd.StatePropagator(\n",
    "    mb, ki.IntegratorType.CVODE_STIFF, None, None, kd.MMSolverType.TREE_AUGMENTED_DYNAMICS\n",
    ")\n",
    "\n",
    "# Initialize state. Use the state from the Multibody that we set up earlier\n",
    "t_init = np.timedelta64(0, \"ns\")\n",
    "sp.setTime(t_init)\n",
    "sp.setState(sp.assembleState())\n",
    "\n",
    "# Makes sure the visualization scene is updated after each state change.\n",
    "kmdl.UpdateProxyScene(\"update_proxy_scene\", sp, mb.getScene())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68b0fea5",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Constraint error projection\n",
    "Even though the individual forward dynamics calls themselves should not introduce much constraint error, as we integrate, constraint error will build up. Hence, periodically, we need to use a method to eliminate this error. Here, we use the {py:class}`Karana.Models.ProjectConstraintError` model, which uses a ConstraintKinematicsSolver to project out the error once it gets too large."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b69a002d",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "pce = kmdl.ProjectConstraintError(\"project_error\", sp, mb, mb.cks())\n",
    "pce.params.tol = 1e-8\n",
    "h = np.timedelta64(int(1e8), \"ns\")\n",
    "pce.setPeriod(h)\n",
    "del pce"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcaa918e",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Prescribing link_1 acceleration\n",
    "Currently, our simulation won't move, as nothing is changing the acceleration of the prescribed hinge for `link_1`. Here, we add a function that will modify the acceleration so that we increase the velocity for the first two seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "71bba626",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def ramp(t, _):\n",
    "    \"\"\"Apply prescribed Udot to move the system.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    t : float\n",
    "        The current time.\n",
    "    _ : NDArray[np.float64]\n",
    "       The current state (unused).\n",
    "    \"\"\"\n",
    "    ft = float(t) / 1e9\n",
    "    if ft < 1.0:\n",
    "        u = cd.getUdot()\n",
    "        u[0] = ft\n",
    "        cd.setUdot(u)\n",
    "    elif ft < 2.0:\n",
    "        u = cd.getUdot()\n",
    "        u[0] = 2.0 - ft\n",
    "        cd.setUdot(u)\n",
    "\n",
    "\n",
    "sp.fns.pre_deriv_fns[\"ramp\"] = ramp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42210fc0",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Run the Simulation\n",
    "Run the simulation for 2 revolutions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "229b0f15",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "sp.advanceTo(4 * np.pi + 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6673c41",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Cleanup the Simulation\n",
    "Cleanup the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d104455",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Cleanup\n",
    "def cleanup():\n",
    "    \"\"\"Cleanup the simulation.\"\"\"\n",
    "    global cd, link_2_context\n",
    "    del cd, link_2_context\n",
    "    kc.discard(sp)\n",
    "    cleanup_graphics()\n",
    "    mb.ensureNotHealthy()\n",
    "    kc.discard(mb.enabledConstraints())\n",
    "    kc.discard(fc.lookupFrame(\"constraint_frame\"))\n",
    "    kc.discard(mb)\n",
    "    kc.discard(fc)\n",
    "\n",
    "\n",
    "atexit.register(cleanup)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0790bab9",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Summary\n",
    "From this tutorial you learned how to create your own flexible multibody. Additionally, you know how to set the state and constraints, solve them, and register the shock, constraint error, and acceleration models to make the simulation work.\n",
    "\n",
    "## Further Readings\n",
    "[flexible double-wishbone](../example_flexible_double_wishbone/notebook.ipynb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  },
  "name": "notebook.ipynb"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
