n-3D-pendulum-sympy-linearized-controlled.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy.physics.mechanics as me\n",
    "import sympy as sm\n",
    "from scipy.integrate import solve_ivp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import matplotlib\n",
    "matplotlib.rcParams['animation.embed_limit'] = 2**128\n",
    "\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import HTML\n",
    "import time\n",
    "\n",
    "from scipy.linalg import solve_continuous_are, det, expm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This creates a decorator to test functions for usage of CPU time, line by line.\\\n",
    "To see the results, this line: *profiler.print_stats()* must be added. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from line_profiler import LineProfiler\n",
    "\n",
    "profiler = LineProfiler()\n",
    "\n",
    "def profile(func):\n",
    "    def inner(*args, **kwargs):\n",
    "        profiler.add_function(func)\n",
    "        profiler.enable_by_count()\n",
    "        return func(*args, **kwargs)\n",
    "    return inner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a **3D n body pendulum**\\\n",
    "I put it into a vertical position, linearize the equations of motion and try to stabilize in by\n",
    "- moving the point $P_{01}$, it is attached to the horizontal plane\n",
    "- adding a torque to every body. (This is needed, because I model the pendulum bodies as RigidBody, and their rotation cannot be stopped by moving around the point of attachment)\n",
    "\n",
    "The control variables are\n",
    "- $F_x, F_z$ forces applied to the point $P_0$ where the pendulum is attached to the plane.\n",
    "- Torques applied to each body fixed frame\n",
    "\n",
    "The 'suspension point' is modeled as a Particle, rest are RigidBody\\\n",
    "(I basically copied this program from some other program quite some time ago, and only enlarged it to 3D and rigid bodies)\\\n",
    "For some reason, which I did not investigate, it only works well for $n \\in \\{{ 1, 3}\\} $. Other values give problems.\n",
    "\n",
    "I do these calculations:\n",
    "- I calculate the full system without controls\n",
    "- I linearize to get the optimum feed back, and use this feed back to numerically solve the full EOMs\n",
    "- I numerically solve the linearized equations, using the optimum feed back\n",
    "- A linear system $\\dfrac{d}{dt} \\hat y = RHS \\cdot \\hat y$, with RHS = const., has the closed form solution $\\hat y = e^{\\displaystyle{RHS \\cdot t}} \\cdot \\hat y_0$, and I use scypi's *expm(..)* function to solve it.\n",
    "\n",
    "**NOTE**:\\\n",
    "At first, I rotated the body fixed frames as 'xyz'. This gave problems with n = 2 bodies. The I changed the rotation to 'yxz'. This seemed to solve the problem with n = 2 and n = 5.\\\n",
    "I think, it is related the the *gimbal lock*, describerd here: https://moorepants.github.io/learn-multibody-dynamics/angular.html?highlight=lock (you have to scroll down a bit to find it) but not really clear, why."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#===============================================\n",
    "n =  3                  #number of pendulums (bodies)\n",
    "plot_NL = True         # if True, the running of the note book is stopped after the full EOMs with controlls have been calculated, so this result may be animated.\n",
    "#===============================================\n",
    "\n",
    "start = time.time()    \n",
    "\n",
    "N = me.ReferenceFrame('N')\n",
    "P0 = me.Point('P0')\n",
    "P0.set_vel(N, 0)                                      \n",
    "\n",
    "# Generalized coordinates and velocities\n",
    "q = []\n",
    "u = []\n",
    "\n",
    "A = [N]\n",
    "\n",
    "Torque = []\n",
    "for i in range(1, n+1):                                             # Data for pendulun i at at position [i-1]\n",
    "    for j in ('x', 'y', 'z'):\n",
    "        q.append(me.dynamicsymbols('q'+str(i) + str(j)))            #Rotation of frame i relative to N\n",
    "        u.append(me.dynamicsymbols('u'+str(i) + str(j)))           \n",
    "    A.append(me.ReferenceFrame('A'+ str(i)))\n",
    "    Torque.append(me.dynamicsymbols('T' + str(i-1)))\n",
    "        \n",
    "q0x, q0z = me.dynamicsymbols('q0x q0z')                             # coordinates of attachment poi nt\n",
    "u0x, u0z = me.dynamicsymbols('u0x, u0z')\n",
    "\n",
    "g, l, m, c1, t = sm.symbols('g l m c1, t')                          # all masses, distances, time, friction are the same for all pendulums\n",
    "r1 = sm.symbols('r1')                                               # distance from Pi to Dmc\n",
    "iXX, iYY, iZZ = sm.symbols('iXX, iYY, iZZ')                         # moments of inertia around the center of mass, principal axes\n",
    "\n",
    "Fx, Fz = me.dynamicsymbols('Fx, Fz')                                       #Control parameters. forces at attachment point acting on P01\n",
    "\n",
    "P01 = P0.locatenew('P01', q0x*N.x + q0z*N.z)\n",
    "P01.set_vel(N, u0x*N.x + u0z*N.z)\n",
    "\n",
    "#Frame points, particles for each Pendulum, numbered 1, 2, ..., n\n",
    "rot = [0]                                                           # rot, rot1 needed for the kinematical differential equations\n",
    "rot1 = [0]\n",
    "P = [P01]\n",
    "DMC = [P01]                                                         # list of center of gravity of each body\n",
    "BODY = [me.Particle('P01a', P01, m)]                                # listof bodies\n",
    "\n",
    "# NOTE, see the comment above about gimbal lock\n",
    "for i in range(1, n+1):\n",
    "    A[i].orient_body_fixed(N, (q[3*i-2], q[3*i-3], q[3*i-1]), '213')  \n",
    "    rot.append(A[i].ang_vel_in(N))\n",
    "    A[i].set_ang_vel(N, u[3*i-3] * A[i].x + u[3*i-2] * A[i].y + u[3*i-1] * A[i].z )\n",
    "    rot1.append(A[i].ang_vel_in(N))\n",
    "\n",
    "    Pi = P[i-1].locatenew('P' + str(i), l * A[i].y)                 # old way to determine points\n",
    "    Pi.v2pt_theory(P[i-1], N, A[i])\n",
    "    P.append(Pi)\n",
    "    Dmci = P[i-1].locatenew('Dmc' + str(i), r1 * A[i].y)            # location of center of gravity\n",
    "    Dmci.v2pt_theory(P[i-1], N, A[i])\n",
    "    DMC.append(Dmci)\n",
    "\n",
    "    I = me.inertia(A[i], iXX, iYY, iZZ)\n",
    "    Bodyi = me.RigidBody('Body' + str(i), Dmci, A[i], m, (I, Dmci))\n",
    "    BODY.append(Bodyi)   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up **Kane's equations**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Kanes forrmalism\n",
    "start1 = time.time()\n",
    "\n",
    "#Kinematic equations\n",
    "kd = [(q0x.diff(t) - u0x)] + [(q0z.diff(t) - u0z)] + [me.dot(rot[i] - rot1[i], uv) for uv in N for i in range(1, n+1)]\n",
    "\n",
    "# forces, torques\n",
    "FL = [(P01, -m*g*N.y + Fx*N.x + Fz*N.z)]\n",
    "\n",
    "for i in range(1, n+1):\n",
    "    FL.append((DMC[i], -m*g*N.y))\n",
    "    FL.append((A[i], Torque[i-1]*A[i].y))\n",
    "\n",
    "q_ind = [q0x, q0z] + q\n",
    "u_ind = [u0x, u0z] + u\n",
    "\n",
    "KM = me.KanesMethod(N, q_ind=q_ind, u_ind=u_ind, kd_eqs=kd)\n",
    "(fr, frstar) = KM.kanes_equations(BODY, FL)\n",
    "\n",
    "MM = KM.mass_matrix_full\n",
    "print('MM DS', me.find_dynamicsymbols(MM))\n",
    "print('MM FS', MM.free_symbols)\n",
    "print(f'MM contains {sm.count_ops(MM):,} operations, {sm.count_ops(sm.cse(MM)):,} after cse, \\n')\n",
    "\n",
    "force = KM.forcing_full\n",
    "print('force DS', me.find_dynamicsymbols(force))\n",
    "print('force FS', force.free_symbols)\n",
    "print(f'force contains {sm.count_ops(force):,} operations, {sm.count_ops(sm.cse(force)):,} after cse \\n')\n",
    "print(\"it took {:.3f} sec to establish Kane's equations\".format(time.time() - start1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Lambdification**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Carthesian coordinates of the center points, needed only for plotting them\n",
    "orte_Dmc = ([ [me.dot(ort.pos_from(P0), uv) for uv in N] for ort in DMC] )\n",
    "\n",
    "qL = q_ind + u_ind\n",
    "pL = [g, m, l, r1, iXX, iYY, iZZ, c1, Fx, Fz] + Torque                                     \n",
    "\n",
    "MM_lam = sm.lambdify(qL + pL, MM, cse=True)\n",
    "force_lam = sm.lambdify(qL + pL, force, cse=True)\n",
    "\n",
    "orte_Dmc_lam = []\n",
    "for ort in orte_Dmc:\n",
    "    orte_Dmc_lam.append(sm.lambdify(qL + pL, ort))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Numerical Integration**\\\n",
    "Normally in my simulations I explicitly set all parameters and initial conditions.\\\n",
    "As here I am after something else, linearization and control, I simply pre_set them.\\\n",
    "Of course they may be changed, just a bit more complicated.\\\n",
    "This is the uncontrolled motion, below I linearize around the vertical position and control it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "start2 = time.time()\n",
    "intervall = 20\n",
    "schritte = 250\n",
    "times = np.linspace(0, intervall, schritte)\n",
    "\n",
    "#Anfangsbedingungen\n",
    "#qL = [q0x, q0z] + q + [u0x, u0z] + u\n",
    "#pL = [g, m, l, r1, iXX, iYY, iZZ, c1, Fx, Fz] + Torque    \n",
    "                                  \n",
    "pL_vals = [9.8, 1, 10, 5., 1/2, 1/3, 1/3, 0., 0., 0.] + [ 0. for _ in range(n)]\n",
    "y0 = [0., 0.] + [0, 0, 0] * int(len(q)/3) + [0., 0.] + [0.1, np.random.choice((-1, 1)) * 0.1, 0.1] * int(len(u)/3)\n",
    "                                            \n",
    "def gradient(t, y, args):                        \n",
    "    vals = np.concatenate((y, args))\n",
    "    sol = np.linalg.solve(MM_lam(*vals), force_lam(*vals))\n",
    "    return np.array(sol).T[0]\n",
    "\n",
    "times = np.linspace(0., intervall, int(schritte*intervall))\n",
    "t_span = (0., intervall)\n",
    "\n",
    "resultat1 = solve_ivp(gradient, t_span, y0, t_eval = times, args=(pL_vals,), atol=1.e-6, rtol=1.e-6) #, method='Radau') \n",
    "resultat = resultat1.y.T\n",
    "print('Shape of result: ', resultat.shape)\n",
    "event_dict = {-1: 'Integration failed', 0: 'Integration finished successfully', 1: 'some termination event'}\n",
    "print(event_dict[resultat1.status], ' the message is: ', resultat1.message)\n",
    "print(f'to simulate {intervall:.3f} sec, the routine made {resultat1.nfev:,} function calls and it took {(time.time() - start2):.3f} CPU - sec  ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "bez = ['X', 'Y', 'Z']\n",
    "fig,(ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10), sharex=True)\n",
    "for i in range(2, int(resultat.shape[1]/2)):\n",
    "    wahl  = (i+1) % 3\n",
    "    wahl1 = (i+1) // 3\n",
    "    ax1.plot(times, np.rad2deg(resultat[:, i]), label=f'{bez[wahl]} - gen. angle of body {wahl1}')\n",
    "ax1.set_title('Generalized angles')\n",
    "ax1.set_ylabel('Angle in degree [°]')\n",
    "ax1.legend();\n",
    "\n",
    "x_coords = np.zeros((len(times), n+1))\n",
    "y_coords = np.zeros((len(times), n+1))\n",
    "z_coords = np.zeros((len(times), n+1))\n",
    "\n",
    "for i in range(len(times)):\n",
    "    for k, ort in enumerate(orte_Dmc_lam):\n",
    "        x_coords[i, k] = ort(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)[0]\n",
    "        y_coords[i, k] = ort(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)[1]\n",
    "        z_coords[i, k] = ort(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)[2]\n",
    "    \n",
    "for j in range(0, n+1):\n",
    "    ax2.plot(times, x_coords[:,j], label='X-Koordinate von Punkt {}'.format(j))\n",
    "    ax2.plot(times, y_coords[:,j], label='Y-Koordinate von Punkt {}'.format(j))\n",
    "    ax2.plot(times, z_coords[:,j], label='Z-Koordinate von Punkt {}'.format(j))\n",
    "ax2.set_title('Coordinates of center of gravity')\n",
    "ax2.set_xlabel('Time [s]')\n",
    "ax2.set_ylabel('Coordinates [m]')\n",
    "ax2.legend(loc='upper right');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Linearisation**\n",
    "\n",
    "I use sympy's method as expalined here: https://docs.sympy.org/latest/modules/physics/mechanics/linearize.html\n",
    "\n",
    "Notes:\n",
    "- the op_point dictionary *must* also include the values of the *accelerations* of the coordinates. This is not very obvious from the documentatrion, JM explained this to me.\n",
    "- the controls variable must be *dynamic symbols*. This is clearly explained in the documentation. Here, they are $F_x, F_z, T_i$ and they are declared to be dynamic symbols above.\n",
    "\n",
    "This is much faster than trying to do it 'from first principles' using jacobian operators - as I did it in my related example.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start3 = time.time()\n",
    "dict_eq = {i: 0. for i in qL} | {i.diff(t): 0. for i in u_ind}\n",
    "\n",
    "F_A, F_B, input_vector = KM.linearize(A_and_B=True, op_point=dict_eq, new_method=True)\n",
    "print('input_vector', input_vector)                                                                     # input vector contains the control parameters\n",
    "print('F_A DS', me.find_dynamicsymbols(F_A))\n",
    "print('F_A FS', F_A.free_symbols)\n",
    "print(f'F_A contains {sm.count_ops(F_A):,} operations, {sm.count_ops(sm.cse(F_A)):,} after cse \\n')\n",
    "\n",
    "print('F_B', me.find_dynamicsymbols(F_B))\n",
    "print('F_B FS', F_B.free_symbols)   \n",
    "print(f'F_B contains {sm.count_ops(F_B):,} operations, {sm.count_ops(sm.cse(F_B)):,} after cse \\n')\n",
    "\n",
    "print(f'It took {(time.time() - start3):.3f} sec to linearize the EOMs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This funktion integrates the 'dynamic Riccati equation', equation (9.18) of this book:\\\n",
    "https://www.amazon.de/Mathematical-Control-Theory-Introduction-Applications/dp/3030447766/ref=sr_1_1?__mk_de_DE=%C3%85M%C3%85%C5%BD%C3%95%C3%91&crid=31QL7QCEXT61I&keywords=mathematical+control+theory&qid=1705340852&sprefix=mathematical+control+theory%2Caps%2C106&sr=8-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@profile\n",
    "def integrate_riccati(A, B, Q, R, end_zeit=10, schritte=500, method='RK45', **kwargs):\n",
    "    '''\n",
    "    This solves the Riccati equation, eq (9.18) of this book\n",
    "    https://www.amazon.de/Mathematical-Control-Theory-Introduction-Applications/dp/3030447766/ref=sr_1_1?__mk_de_DE=%C3%85M%C3%85%C5%BD%C3%95%C3%91&crid=31QL7QCEXT61I&keywords=mathematical+control+theory&qid=1705340852&sprefix=mathematical+control+theory%2Caps%2C106&sr=8-1\n",
    "    \n",
    "    d/dtP = Q + P A + A.T P - P B R^-1 B.T P, where I (arbitrarily) take P(0) = I\n",
    "    I integrate  from 0. to end_zeit,\n",
    "    schritte * end_zeit gives the number of integration steps to be returned \n",
    "    '''\n",
    "    n = Q.shape[0]\n",
    "    m = Q.shape[1]\n",
    "\n",
    "    P0 = np.identity(n).flatten()\n",
    "\n",
    "    def gradient(t, P0):\n",
    "        P0 = P0.reshape(n, m)           \n",
    "        RHS = (Q + P0 @ A + A.T @ P0 - P0 @ B @ np.linalg.inv(R) @ B.T @ P0).flatten()\n",
    "        P0 = P0.flatten()\n",
    "        return RHS\n",
    "\n",
    "    times = np.linspace(0., end_zeit, int(schritte*end_zeit))\n",
    "    t_span = (0., end_zeit)\n",
    "\n",
    "    resultat1 = solve_ivp(gradient, t_span, P0, t_eval = times, method=method, options=kwargs)\n",
    "    resultat = resultat1.y.T\n",
    "    print(resultat1.message)\n",
    "    resultat = resultat.reshape(len(resultat), n, m)\n",
    "    return resultat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The linearized sytem is: $\\dfrac{d}{dt} \\bar y = F_A \\cdot \\bar y + F_B \\cdot \\bar u$, where $\\bar u$ is the control, and $F_A, F_B$ are the matrices obtained above.\n",
    "\n",
    "Calculate the **optimal feed back** to stabilize the linearized system.\\\n",
    "Let P be the solution to the Riccati equation $Q + P \\cdot F_A + F_A^T \\cdot P - P \\cdot B \\cdot R^{-1} \\cdot F_B^T \\cdot P = 0$, where $Q, R$ are the weight matrices in the cost functional related to the coordinates and the control, respectively.\\\n",
    "Then $\\hat{\\overline{control}} = \\left[ R^{-1} \\cdot F_B^T \\cdot P \\right] \\cdot \\bar y$ ist the optimal control.\\\n",
    "For lack of better reason, I pick $R$ and $Q$ to be the unit matrices of appropriate dimensions.\\\n",
    "The scipy function **solve_continuous_are** numerically solves the *Ricatti equation*.\\\n",
    "Controllability means, I think, that any starting configuration of the system may be moved to another predetermined position via the controls. Intuitively clear, that with this pendulum this is only the case for n = 1\\\n",
    "All of this is in books about (linear) control theory. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start4 = time.time()\n",
    "\n",
    "pL = [g, m, l, r1, iXX, iYY, iZZ, c1, Fx, Fz] + Torque\n",
    "pL_vals = [9.8, 1., 10., 5., 1, 2, 3, 1.,  0, 0] + [0. for i in range(len(Torque))]\n",
    "\n",
    "A_lam = sm.lambdify(pL, F_A, cse=True)\n",
    "B_lam = sm.lambdify(pL, F_B, cse=True)\n",
    "\n",
    "Alin = A_lam(*pL_vals)\n",
    "Blin = B_lam(*pL_vals)\n",
    "\n",
    "print('determinant of Alin', det(Alin))\n",
    "#Is the system controllable? This seems to be only a sufficient condition (?)\n",
    "def controllable(a, b):\n",
    "    a = np.matrix(a)\n",
    "    b = np.matrix(b)\n",
    "    n = a.shape[0]\n",
    "    controllable_matrix = []\n",
    "    for i in range(n):\n",
    "        controllable_matrix.append(a ** i * b)\n",
    "    controllable_matrix = np.hstack(controllable_matrix)\n",
    "    return np.linalg.matrix_rank(controllable_matrix) == n\n",
    "print('System controllable?  ', controllable(Alin, Blin))\n",
    "\n",
    "#compute optimal gains K\n",
    "Q = np.eye(Alin.shape[1])\n",
    "R = np.eye(Blin.shape[1])\n",
    "S = solve_continuous_are(Alin, Blin, Q, R) #solves Riccati equation\n",
    "K = np.dot(np.dot(np.linalg.inv(R), Blin.T), S)\n",
    "\n",
    "print('it took {:.3f} sec to calculate the feed back'.format(time.time() - start4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Her I just want to see how close the solution of the integratred Riccati equation is to the solution of the algebraic Riccati equation, solved above.\\\n",
    "As the norm of the matrix, I use the largest entry of it: $ || M || := max_{i, j} |M_{ij}|$.\\\n",
    "At least here, it seems to converge quite rapidly.\\\n",
    "Just a side line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "end_zeit = 100.\n",
    "TEST = integrate_riccati(Alin, Blin, Q, R, end_zeit=end_zeit, schritte=500, kwargs={'rtol': 1.e-8, 'rtol': 1.3-8})\n",
    "fehler = np.max(np.abs(np.subtract(TEST[-1], S))) / np.max(np.abs(S)) * 100\n",
    "print(f'The maximal error of the integrated Riccati - ARE is {fehler:.3f} % of max entry of S, when it was integrated from 0 to {end_zeit} sec, and the last solution used \\n')\n",
    "#profiler.print_stats()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Numerical integration** of the **full system** using the optimum controls\\\n",
    "The optimal feed back is calculated each step of the integration.\n",
    "\n",
    "**Important note**\\\n",
    "Only integrating this cell will tell you whether the system, given the initial disturbance, will actually be stabilized.\\\n",
    "Calculating the linearized versions, by definition!, will always look like you can control it.\\\n",
    "(If this integration runs more than a few seconds, it is likely not stabilizeable)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pL1_vals = [9.8, 1., 10., 5.,  1., 2., 3., 1]\n",
    "equilibrium_point = np.zeros(len(qL))\n",
    "\n",
    "aa = []\n",
    "for i in range(int(len(u)/3)):\n",
    "    aa.append(0.0 * np.random.choice((-1., 1.)))\n",
    "    aa.append(0.15 * np.random.choice((-1., 1.)))\n",
    "    aa.append(0.)\n",
    "y0 = [0., 0.] + [0.145, 0., 0.195] * int(len(q)/3) + [0., 0.] + aa \n",
    "y_store = list(y0)                                                      # store the initial conditions, so I do the solution using matrix exponentials below with the same initial conditions\n",
    "\n",
    "\n",
    "start5 = time.time()\n",
    "\n",
    "# The controls, which originally were in pL_vals will now, of course be the optimal feed back, and added during each step of the integration\n",
    "# now: kraft = K @ (x(equilibrium) - x(t)). I store the forces and the times in lists, so that I can plot them later.\n",
    "kraefte = []\n",
    "zeiten  = []\n",
    "\n",
    "#@profile\n",
    "def gradient(t, x, args):\n",
    "# calculate the control force\n",
    "    kraft = np.dot(K, equilibrium_point - x)\n",
    "    kraefte.append(kraft)\n",
    "    zeiten.append(t)\n",
    "    arguments = np.hstack((x, args, kraft))\n",
    "    return  np.array(np.linalg.solve(MM_lam(*arguments), force_lam(*arguments))).T[0]\n",
    "\n",
    "times = np.linspace(0., intervall, int(schritte*intervall)) \n",
    "t_span = (0., intervall)\n",
    "\n",
    "resultat1 = solve_ivp(gradient, t_span, y0, t_eval = times, args=(pL1_vals,), atol=1.e-8, rtol=1.e-8) #, method='Radau') \n",
    "resultat = resultat1.y.T\n",
    "print('Shape of result: ', resultat.shape)\n",
    "event_dict = {-1: 'Integration failed', 0: 'Integration finished successfully', 1: 'some termination event'}\n",
    "print(event_dict[resultat1.status], ' the message is: ', resultat1.message)\n",
    "print(f'to simulate {intervall:.3f} sec, the routine made {resultat1.nfev:,} function calls and it took {(time.time() - start5):.3f} CPU - sec  ')\n",
    "#profiler.print_stats()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the results\\\n",
    "It can easily be seen, that the system is driven to its equilibrium point, if the initial disturbances are not too great.\\\n",
    "Note, that the results are a bit different from those obtained by solving the linearized EOMs below. I do not know why.\\\n",
    "**Note**: if you set **plot_NL = True** above, the program will be interrupted here, so the results obtained here wil not be overriden by the results of the subsequent calculations. So, you should start ther animation at the bottom of this program to start the animation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "bez = ['X', 'Y', 'Z']\n",
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 12), sharex=True)\n",
    "for i in range(2, int(resultat.shape[1]/2)):\n",
    "    wahl  = (i+1) % 3\n",
    "    wahl1 = (i+1) // 3\n",
    "    ax1.plot(times, np.rad2deg(resultat[:, i]), label=f'{bez[wahl]} - gen. angle of body {wahl1}')\n",
    "ax1.set_title('Generalized angles, integrating full EOMs, with controls')\n",
    "ax1.set_ylabel('Angle in degree [°]')\n",
    "ax1.legend();\n",
    "\n",
    "kraefte = np.array(kraefte)\n",
    "for i, j in enumerate(('Fx', 'Fy', *['Torque on body ' + str(kk) for kk in range(1, n+1)])):\n",
    "    ax2.plot(zeiten, kraefte[:, i], label=j)\n",
    "ax2.set_title('Control forces / torques')\n",
    "ax2.set_xlabel('Time [s]')\n",
    "ax2.set_ylabel('Force [N] / Torque [Nm]')   \n",
    "ax2.legend();\n",
    "\n",
    "if plot_NL:\n",
    "    raise Exception('go down to the last cell, the animation, and start it')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Numerically integrate** the **linearized system**.\n",
    "\n",
    "The sytem here is: $\\dfrac{d}{dt} \\bar y = F_A \\cdot \\bar y + F_B \\cdot \\bar u$, where $\\bar u$ is the control.\n",
    "\n",
    "Form the matrix $RHS$ to be calculated\n",
    " as follows:\n",
    "\n",
    "First, form new state vektor $\\hat y = \\left[ \\bar y, \\bar u \\right]$\n",
    "\n",
    "The new rhs of the system is $RHS = \\begin{bmatrix} F_A & F_B \\\\ 0_1 & 0_2 \\end{bmatrix}$, where $0_i = \\begin{bmatrix} 0. &....& 0. \\\\. &...&.\\\\ 0. &...& 0. \\end{bmatrix}$ of appropriate dimensions.\n",
    "\n",
    "Now I integrate the system numerically.\\\n",
    "( Of course, I could have done it the same way I did it with the full EOMs with the controls as parameters, I simply wanted to try another way.) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start6 = time.time()\n",
    "\n",
    "RHS = sm.Matrix.hstack(F_A, F_B)\n",
    "RHS = sm.Matrix.vstack(RHS, sm.zeros(F_B.shape[1], RHS.shape[1]))\n",
    "print('RHS shape', RHS.shape)\n",
    "print('RHS DS', me.find_dynamicsymbols(RHS))\n",
    "print('RHS FS', RHS.free_symbols)\n",
    "print(f'RHS contains {sm.count_ops(RHS):,} operations, {sm.count_ops(sm.cse(RHS)):,} after cse \\n')\n",
    "\n",
    "pL1 = [g, m, l, r1, iXX, iYY, iZZ, c1]\n",
    "p1L_vals = [9.8, 1., 10., 5., 1, 2, 3, 1.]\n",
    "\n",
    "# convert RHS from a sympy matrix to a numpy array\n",
    "RHS_lam = sm.lambdify(pL1, RHS, cse=True)\n",
    "RHS = RHS_lam(*pL1_vals)\n",
    "\n",
    "# extend the equilibrium point to include the control forces\n",
    "equilibrium_point1 = np.hstack((equilibrium_point, np.zeros(Blin.shape[1])))\n",
    "\n",
    "# form y_hat, the initial conditions for the solution of the linearized system\n",
    "y0 = np.array(list(y_store) + [0. for _ in range(Blin.shape[1])])      \n",
    "\n",
    "# solve the linearized system, using matrix exponentials\n",
    "resultat = []\n",
    "kraefte  = []\n",
    "zeiten   = [] \n",
    "\n",
    "#@profile\n",
    "def gradient(t, y):\n",
    "# get the optium control\n",
    "    kraft = K @ (equilibrium_point - y[: len(qL)])\n",
    "    kraefte.append(kraft)\n",
    "    zeiten.append(t)\n",
    "    y[len(qL):] = kraft\n",
    "    return RHS @ (y - equilibrium_point1)\n",
    "\n",
    "times = np.linspace(0., intervall, int(schritte*intervall))\n",
    "t_span = (0., intervall)\n",
    "\n",
    "resultat1 = solve_ivp(gradient, t_span, y0, t_eval = times, atol=1.e-10, rtol=1.e-10) #, method='Radau') \n",
    "resultat = resultat1.y.T\n",
    "print('Shape of result: ', resultat.shape)\n",
    "event_dict = {-1: 'Integration failed', 0: 'Integration finished successfully', 1: 'some termination event'}\n",
    "print(event_dict[resultat1.status], ' the message is: ', resultat1.message)\n",
    "print(f'to simulate {intervall:.3f} sec, the routine made {resultat1.nfev:,} function calls and it took {(time.time() - start6):.3f} CPU - sec  ')\n",
    "#profiler.print_stats()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bez = ['X', 'Y', 'Z']\n",
    "resultat_hilfs = []\n",
    "for i in range(len(times)):\n",
    "    resultat_hilfs.append(resultat[i, : len(qL)])\n",
    "resultat = np.array(resultat_hilfs)\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 12), sharex=True)\n",
    "for i in range(2, int(resultat.shape[1]/2)):\n",
    "    wahl  = (i+1) % 3\n",
    "    wahl1 = (i+1) // 3\n",
    "    ax1.plot(times, np.rad2deg(resultat[:, i]), label=f'{bez[wahl]} - gen. angle of body {wahl1}')\n",
    "ax1.set_title('Generalized angles, integrating the linearized system')\n",
    "ax1.set_ylabel('Angle in degree [°]')\n",
    "ax1.legend();\n",
    "\n",
    "kraefte = np.array(kraefte)\n",
    "for i, j in enumerate(('Fx', 'Fy', *['Torque on body ' + str(kk) for kk in range(1, n+1)])):\n",
    "    ax2.plot(zeiten, kraefte[:, i], label=j)\n",
    "ax2.set_title('Control forces / torques')\n",
    "ax2.set_xlabel('Time [s]')\n",
    "ax2.set_ylabel('Force [N] / Torque [Nm]')   \n",
    "ax2.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here I use scipy's **matrix exponential** solver *expm(..)* to solve the differential equation without integrating.\\\n",
    "I use the fact, that this idfferential equation $\\dfrac{d}{dt} \\hat y = RHS \\cdot \\hat y$, with RHS = const., has the solution $\\hat y = e^{\\displaystyle{RHS \\cdot t}} \\cdot \\hat y_0$\\\n",
    "I have to to it in steps, using the result of the previous step as input for the next step, otherwise of course the controls will have no effect.\\\n",
    "At least with this example, this is the *fastest* way of doing it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start7 = time.time()\n",
    "y = np.array(list(y_store) + [0. for _ in range(Blin.shape[1])])      \n",
    "# solve the linearized system, using matrix exponentials\n",
    "resultat = []\n",
    "kraefte  = []\n",
    "zeiten   = [] \n",
    "\n",
    "for i in range(len(times)-1):\n",
    "# get the optium control\n",
    "    kraft = K @ (equilibrium_point - y[: len(qL)])\n",
    "    kraefte.append(kraft)\n",
    "    zeiten.append(times[i])\n",
    "    y[len(qL):] = kraft\n",
    "# Here I do the stepwise solution.\n",
    "    y = expm(RHS * (times[i+1] - times[i])) @ y\n",
    "    resultat.append((y))\n",
    "    \n",
    "print(f'to calculate {intervall:.3f} sec took {(time.time() - start7):.3f} CPU - sec  ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resultat_hilfs = []\n",
    "resultat = np.array(resultat)\n",
    "for i in range(len(times)-1):\n",
    "    resultat_hilfs.append(resultat[i, : len(qL)])\n",
    "resultat = np.array(resultat_hilfs)\n",
    "times = times[:-1]\n",
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 12), sharex=True)\n",
    "for i in range(2, int(resultat.shape[1]/2)):\n",
    "    wahl  = (i+1) % 3\n",
    "    wahl1 = (i+1) // 3\n",
    "    ax1.plot(times, np.rad2deg(resultat[:, i]), label=f'{bez[wahl]} - gen. angle of body {wahl1}')\n",
    "ax1.set_title('Generalized angles, used matrix exponential')\n",
    "ax1.set_ylabel('Angle in degree [°]')\n",
    "ax1.legend();\n",
    "\n",
    "kraefte = np.array(kraefte)\n",
    "for i, j in enumerate(('Fx', 'Fy', *['Torque on body ' + str(kk) for kk in range(1, n+1)])):\n",
    "    ax2.plot(zeiten, kraefte[:, i], label=j)\n",
    "ax2.set_title('Control forces / torques')\n",
    "ax2.set_xlabel('Time [s]')\n",
    "ax2.set_ylabel('Force [N] / Torque [Nm]')   \n",
    "ax2.legend();\n",
    "print(f'running the complete program takes {(time.time() - start):.3f} CPU - sec')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Animation**\n",
    "\n",
    "In my model, gravitation points in the negative Y direction, in matplotlib 3D, gravitation points in the negative Z direction.\\\n",
    "Hence, I must make this transformation:\n",
    "- $N.x_{old} \\rightarrow N.y_{new}$\n",
    "- $N.y_{old} \\rightarrow N.z_{new}$\n",
    "- $N.z_{old} \\rightarrow N.x_{new}$\n",
    "\n",
    "The 'skeleton' of this program is from chat GPT\n",
    "\n",
    "matplotlib seems to have no intrinsic object *sphere*. Trying to do it from 'first principles' takes VERY long.\\\n",
    "Hence, I marked the bodies simply with large markers.\\\n",
    "I left the spheres in the code, if anyone wants to test it.\n",
    "\n",
    "Please note, that the *last* calculation is plotted.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reduce the number of points of time to around zeitpunkte\n",
    "#=======================\n",
    "zeitpunkte = 500\n",
    "#=======================\n",
    "\n",
    "# plot the horizontal plane H at Z = 0.\n",
    "def plot_3d_plane(x_min, x_max, y_min, y_max):\n",
    "    # Create a meshgrid for x and y values\n",
    "    x = np.linspace(x_min, x_max, 100)\n",
    "    y = np.linspace(y_min, y_max, 100)\n",
    "    x, y = np.meshgrid(x, y)\n",
    "    # Z values are set to 0 for the plane\n",
    "    z = np.zeros_like(x)\n",
    "    # Plot the 3D plane\n",
    "    ax.plot_surface(x, y, z, alpha=0.1, rstride=100, cstride=100, color='c')\n",
    "\n",
    "'''\n",
    "# draw a shere in matplotlib 3D\n",
    "def kugel(ax, radius, mittelpunkt, farbe):\n",
    "    u = np.linspace(0, 2 * np.pi, 50)\n",
    "    v = np.linspace(0, np.pi, 50)\n",
    "    x = mittelpunkt[0] + radius * np.outer(np.cos(u), np.sin(v))\n",
    "    y = mittelpunkt[1] + radius * np.outer(np.sin(u), np.sin(v))\n",
    "    z = mittelpunkt[2] + radius * np.outer(np.ones(np.size(u)), np.cos(v))\n",
    "    ax.plot_surface(x, y, z, color=farbe)\n",
    "'''\n",
    "times2 = []\n",
    "resultat2 = []\n",
    "\n",
    "reduction = max(1, int(len(times)/zeitpunkte))\n",
    "for i in range(len(times)):\n",
    "    if i % reduction == 0:\n",
    "        times2.append(times[i])\n",
    "        resultat2. append(resultat[i, 0: len(qL)])\n",
    "  \n",
    "schritte2 = len(times2)\n",
    "print(f'animation used {schritte2} points in time')\n",
    "resultat2 = np.array(resultat2)\n",
    "times2 = np.array(times2)  \n",
    "\n",
    "'''\n",
    "# I assume, the pendulum bodies are sheres. I put a point on each one, so the rotation may be seen more easily.\n",
    "rb  = sm.symbols('rb')     \n",
    "rb1 = 2.5                                          # radius of the spheres\n",
    "OBS = []\n",
    "for i in range(1, n+1):\n",
    "    OBS.append(me.Point('obs' + str(i)))\n",
    "    OBS[i-1].set_pos(DMC[i], rb * A[i].x)\n",
    "    OBS[i-1].v2pt_theory(DMC[i], N, A[i])\n",
    "'''\n",
    "trans = [N.z, N.x, N.y]                                         # aligns the coordinate system of my model with that of matplotlib\n",
    "\n",
    "P01_ort = [me.dot(P01.pos_from(P0), uv) for uv in trans]        # coordinates of attachment point\n",
    "P_ort   = []\n",
    "DMC_ort = []\n",
    "#OBS_ort = []\n",
    "for i in range(1, n+1):\n",
    "    P_ort.append([me.dot(P[i].pos_from(P0), uv) for uv in trans])\n",
    "    DMC_ort.append([me.dot(DMC[i].pos_from(P0), uv) for uv in trans])\n",
    "#    OBS_ort.append([me.dot(OBS[i-1].pos_from(P0), uv) for uv in trans])\n",
    "    \n",
    "pL_p = [g, m, l, r1, iXX, iYY, iZZ, c1] #, rb]\n",
    "pL_p_vals = [9.8, 1., 10., 5., 1, 2, 3, 1.] #, rb1]\n",
    "\n",
    "P01_ort_lam = sm.lambdify(qL + pL_p, P01_ort, cse=True)\n",
    "P_ort_lam   = sm.lambdify(qL + pL_p, P_ort, cse=True)\n",
    "DMC_ort_lam = sm.lambdify(qL + pL_p, DMC_ort, cse=True)\n",
    "#OBS_ort_lam = sm.lambdify(qL + pL_p, OBS_ort, cse=True)\n",
    "\n",
    "P_x = np.empty((schritte2, n))\n",
    "P_y = np.empty((schritte2, n))\n",
    "P_z = np.empty((schritte2, n))\n",
    "\n",
    "DMC_x = np.empty((schritte2, n))\n",
    "DMC_y = np.empty((schritte2, n))\n",
    "DMC_z = np.empty((schritte2, n))\n",
    "\n",
    "#OBS_x = np.empty((schritte2, n))\n",
    "#OBS_y = np.empty((schritte2, n))\n",
    "#OBS_z = np.empty((schritte2, n))\n",
    "\n",
    "P01_x = np.empty(schritte2)\n",
    "P01_y = np.empty(schritte2)\n",
    "P01_z = np.empty(schritte2)\n",
    "\n",
    "for i in range(schritte2):\n",
    "    for j in range(n):\n",
    "        P_x[i, j] = P_ort_lam(*resultat2[i, :], *pL_p_vals)[j][0]\n",
    "        P_y[i, j] = P_ort_lam(*resultat2[i, :], *pL_p_vals)[j][1]\n",
    "        P_z[i, j] = P_ort_lam(*resultat2[i, :], *pL_p_vals)[j][2]\n",
    "        \n",
    "        DMC_x[i, j] = DMC_ort_lam(*resultat2[i, :], *pL_p_vals)[j][0]\n",
    "        DMC_y[i, j] = DMC_ort_lam(*resultat2[i, :], *pL_p_vals)[j][1]\n",
    "        DMC_z[i, j] = DMC_ort_lam(*resultat2[i, :], *pL_p_vals)[j][2]\n",
    "        \n",
    "#        OBS_x[i, j] = OBS_ort_lam(*resultat2[i, :], *pL_p_vals)[j][0]\n",
    "#        OBS_y[i, j] = OBS_ort_lam(*resultat2[i, :], *pL_p_vals)[j][1]\n",
    "#        OBS_z[i, j] = OBS_ort_lam(*resultat2[i, :], *pL_p_vals)[j][2] \n",
    "\n",
    "    P01_x[i] = P01_ort_lam(*resultat2[i, :], *pL_p_vals)[0]\n",
    "    P01_y[i] = P01_ort_lam(*resultat2[i, :], *pL_p_vals)[1]\n",
    "    P01_z[i] = P01_ort_lam(*resultat2[i, :], *pL_p_vals)[2]\n",
    "\n",
    "\n",
    "x_min = min(np.min(P_x), np.min(DMC_x), np.min(P01_x))\n",
    "y_min = min(np.min(P_y), np.min(DMC_y), np.min(P01_y))\n",
    "z_min = min(np.min(P_z), np.min(DMC_z), np.min(P01_z))\n",
    "\n",
    "x_max = max(np.max(P_x), np.max(DMC_x), np.max(P01_x))\n",
    "y_max = max(np.max(P_y), np.max(DMC_y), np.max(P01_y))\n",
    "z_max = max(np.max(P_z), np.max(DMC_z), np.max(P01_z))\n",
    "\n",
    "farben = ['b', 'g', 'r', 'c', 'm', 'y', 'k', 'w']   # I need a color for each pendulum body\n",
    "\n",
    "fig = plt.figure(figsize=(10, 10))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.set_xlim([x_min - 1, x_max + 1])  # Adjust the limits based on your data\n",
    "ax.set_ylim([y_min - 1, y_max + 1])\n",
    "ax.set_zlim([z_min - 1, z_max + 1])\n",
    "ax.set_xlabel('X', fontsize = 15)\n",
    "ax.set_ylabel('Y', fontsize = 15)\n",
    "ax.set_zlabel('Z', fontsize = 15)\n",
    "ax.set_aspect('equal')\n",
    "\n",
    "ebene = plot_3d_plane(x_min-1, x_max+1, y_min-1, y_max+1)           # plot base plane, at Z = 0.\n",
    "\n",
    "# Initialize points and rods\n",
    "LINE_P = ['P' + str(i) for i in range(1, n+1)]\n",
    "LINE_DMC = ['DMC' + str(i) for i in range(1, n+1)]\n",
    "#LINE_OBS = ['OBS' + str(i) for i in range(1, n+1)]\n",
    "LINE_rod = ['rod' + str(i) for i in range(1, n+1)]\n",
    "\n",
    "LINE_P0 = ax.plot([], [], [], 'o', color=farben[2], markersize=8)\n",
    "for i in range(n):\n",
    "    LINE_P[i], = ax.plot([], [], [], 'o', color=farben[0], markersize=3)\n",
    "    LINE_DMC[i], = ax.plot([], [], [], 'o', color=farben[i], markersize=25)\n",
    "#    LINE_OBS[i], = ax.plot([], [], [], 'o', color=farben[i], markersize=5)\n",
    "    LINE_rod[i], = ax.plot([], [], [], '-', color=farben[0], linewidth=1)\n",
    "LINE_trace, = ax.plot([], [], [], '-', color=farben[2], linewidth=0.5)          # traces the movement of the suspension point P01 of the pendulum\n",
    "\n",
    "# Update function for animation\n",
    "def update(frame):\n",
    "\n",
    "    message = (f'Vertical Pendulum \\n Running time {times2[frame]:.2f} sec')\n",
    "    ax.set_title(message, fontsize=15)\n",
    "\n",
    "    for i in range(n):\n",
    "        LINE_P[i].set_data([P_x[frame, i]], [P_y[frame, i]])\n",
    "        LINE_P[i].set_3d_properties([P_z[frame, i]])\n",
    "\n",
    "        LINE_DMC[i].set_data([DMC_x[frame, i]], [DMC_y[frame, i]])\n",
    "        LINE_DMC[i].set_3d_properties([DMC_z[frame, i]])\n",
    "\n",
    "#        LINE_OBS[i].set_data([OBS_x[frame, i]], [OBS_y[frame, i]])\n",
    "#        LINE_OBS[i].set_3d_properties([OBS_z[frame, i]])\n",
    "    \n",
    "    LINE_rod[0].set_data([P01_x[frame], P_x[frame, 0]],\n",
    "                             [P01_y[frame], P_y[frame, 0]])\n",
    "    LINE_rod[0].set_3d_properties([P01_z[frame], P_z[frame, 0]])\n",
    "\n",
    "    for i in range(1, n):\n",
    "        LINE_rod[i].set_data([P_x[frame, i-1], P_x[frame, i]],\n",
    "                             [P_y[frame, i-1], P_y[frame, i]])\n",
    "        LINE_rod[i].set_3d_properties([P_z[frame, i-1], P_z[frame, i]])\n",
    "    \n",
    "    LINE_P0[0].set_data([P01_x[frame]], [P01_y[frame]])\n",
    "    LINE_P0[0].set_3d_properties([P01_z[frame]])\n",
    "\n",
    "    LINE_trace.set_data(P01_x[:frame], P01_y[:frame])\n",
    "    LINE_trace.set_3d_properties(P01_z[:frame])\n",
    "\n",
    "#    for i in range(n):\n",
    "#        kugel(ax, rb1, [DMC_x[frame, i], DMC_y[frame, i], DMC_z[frame, i]], farben[i])\n",
    "\n",
    "    return LINE_P + LINE_rod + LINE_P0 + LINE_DMC + [LINE_trace]\n",
    "\n",
    "# Create the animation\n",
    "num_frames = schritte2\n",
    "animation = FuncAnimation(fig, update, frames=num_frames, interval=50, blit=True)\n",
    "plt.close(fig)                                                                                      # avoids the 'second' picture.\n",
    "# Display the animation\n",
    "HTML(animation.to_jshtml())"
   ]
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}