3d-n-link-spring.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sm\n",
    "import sympy.physics.mechanics as me\n",
    "import time\n",
    "import numpy as np\n",
    "from scipy.integrate import solve_ivp\n",
    "from scipy.optimize import fsolve, minimize\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mp\n",
    "from matplotlib import patches\n",
    "%matplotlib inline\n",
    "from IPython.display import HTML\n",
    "mp.rcParams['animation.embed_limit'] = 2**128\n",
    "from matplotlib import animation\n",
    "\n",
    "import pythreejs as p3js\n",
    "from copy import deepcopy, copy\n",
    "\n",
    "from itertools import permutations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This models a n link pendulum, same like the one in JM's lecture, only in 3D. $n$ balls of radius $r$ and mass $m$ are located at the center of each massless rod. An observer of mass $m_o$ may be attached to each ball, its distance rom the center of each ball is $\\beta \\cdot r$ perpendcular to the massless rod the ball is attached to.\\\n",
    "Collisions between balls are handled with the **Hunt-Crossley** method, details see below.\n",
    "\n",
    "In order to fix the last point, $P_n$ at $l_n\\cdot N.x$ I use the method described below.\n",
    "\n",
    "*feder method*\\\n",
    "In order to fix $P_n$ at $l_n \\cdot N.x$, I introduce a point $Pn$ at the position $P_n$ should be at, and attach spring and, optionally, a dampening device from $P_n$ to $Pn$. The advantage of this approach is that it seems stable over longer times. Of course, it dissipates energy if the dampening is unequal to zero, which the original system does not.\n",
    "\n",
    "\n",
    "**Symbols used**\n",
    "\n",
    "- $q_{x1}..q_{zn}:\\hspace{78 mm} \\text{generalized coordinates}$\n",
    "- $u_{x1}..u_{zn}:\\hspace{78 mm} \\text{generalized speeds}$\n",
    "- $N:\\hspace{90 mm} \\text{inertial frame}$\n",
    "- $A_1,..., A_n:\\hspace{75 mm} \\text{body fixed frames of each each ball}$\n",
    "- $P_0.....P_n:\\hspace{75 mm} \\text{points at the end of each rod}$\n",
    "- $Dmc_i$:\\hspace{85 mm} \\text{Centers of the balls}\n",
    "- $Pn:\\hspace{89.5 mm} \\text{auxiliary point}$\n",
    "- $l_1, ....., l_n:\\hspace{74.5 mm} \\text{length of the rods and length from $P_{i-1}$ to $P_i$}$\n",
    "- $ln:\\hspace{90 mm} \\text{distance from $P_0$ to $Pn$ along the N.x axis}$\n",
    "- $m, m_o:\\hspace{84 mm} \\text{mass of ball, mass of observer}$\n",
    "- $r, k, d:\\hspace{83.5 mm} \\text{radius of ball, spring constant, dampening constant}$\n",
    "- $ k_b: \\hspace{90 mm} \\text{spring constant of elastic ball}$\n",
    "- $c_\\tau: \\hspace{90 mm} \\text{experimental constant needed for Hunt-Crossley}$\n",
    "- $l_0: \\hspace{91 mm} \\text{natural length of the spring}$\n",
    "- $aux_{0x}, aux_{0y}, aux_{0z}, aux_{nx}, aux_{ny}, aux_{nz}:\\hspace{23 mm} \\text{virtual speeds for reaction force}$\n",
    "- $f_{0x}, f_{0y}, f_{0z}, f_{nx}, f_{ny}, f_{nz}: \\hspace{50 mm} \\text{reaction forces}$\n",
    "- $rhodt_{max} \\hspace{85 mm} \\text{holds the speeds just before impact}$\n",
    "\n",
    "\n",
    "With *term_info* you can select whether information (dynamic symbols, free symbols, number of operations) of mass matrix, force vector and reaction forces vector (*eingepraegt*) should be printed or not. Not printing speeds up the process of setting up Kane's equations by around 15%...20%\\\n",
    "If *daempfung* is set to TRUE, a speed dependent dampening will be applied in parallel to the spring.\\\n",
    "If *frictionless = True* the balls are ideally slick.\n",
    "\n",
    "\n",
    "*observation*\\\n",
    "The number of operations increase a lot with increasing number of balls.\\\n",
    "Values below are *without* dampening, dampening increases the size of the force vector by a factor of almost 10, they are also **without friction** during the collisions:\n",
    "- $n = 2$:\n",
    "    - mass matrix 4,000 operations\n",
    "    - force vector 42,000 operations\n",
    "    - reaction force 25,000 operations\n",
    "- $n = 3$:\n",
    "     - mass matrix      70,000 operations\n",
    "     - force vector    790,000 operations\n",
    "     - reaction  force 225,000 operation\n",
    "- $n = 4$:\n",
    "    - mass matrix 560,000 operations\n",
    "    - force vector 8,500,000 operations\n",
    "    - reaction force 1,500,000 operations\n",
    "    \n",
    "**With friction** I get these number of operations:\n",
    "- $n = 3$:\n",
    "    - mass matrix 70,000\n",
    "    - force vector 23,000,000\n",
    "    - reaction force 7,000,000\n",
    " \n",
    "I obviously did not try for larger $n$.\n",
    "    \n",
    "For the speed constraints below I need the relationship between $d/dt(q_{ij})$ and $u_{ij}$, with $ i \\in \\{x, y, z \\}$ and $j \\in \\{1,.. ,n\\}$  where $u_{ij}$ are the speeds defined below. It will be stored in *antwort*\n",
    "For this I solve the linear equation at the very bottom. I use *sm.solve* as the linear system is small, und *sm.solve* returns the solution nicely as a dictionary. For large linear systems, Jason Moore's *LUsolve* method is much faster.\\\n",
    "This is also basically used to get the kinematic equations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "start5 = None # just to avoid an error in the pythreejs animation, if the matplotlib animation has not run.\n",
    "\n",
    "# here are n balls, and (n+1) frames, including the inertial frame N\n",
    "#===================\n",
    "n = 3\n",
    "#===================\n",
    "daempfung    = False\n",
    "#===================\n",
    "term_info    = True\n",
    "#===================\n",
    "frictionless = True\n",
    "#===================\n",
    "\n",
    "if n < 2 or isinstance(n, int) == False:\n",
    "    raise Exception(\"n must be an integer >= 2\")\n",
    "\n",
    "\n",
    "aux0x, aux0y, aux0z, auxnx, auxny, auxnz = me.dynamicsymbols('aux0x, aux0y, aux0z, auxnx, auxny, auxnz')\n",
    "f0x, f0y, f0z, fnx, fny, fnz = me.dynamicsymbols('f0x, f0y, f0z, fnx, fny, fnz')\n",
    "aux = [aux0x, aux0y, aux0z, auxnx, auxny, auxnz]\n",
    "F = [f0x, f0y, f0z, fnx, fny, fnz]\n",
    "rhs = [sm.symbols('rhs' + str(i)) for i in range(3*n, 6*n)]\n",
    "\n",
    "\n",
    "N = me.ReferenceFrame('N')\n",
    "P0 = me.Point('P0')\n",
    "P0.set_vel(N, aux0x*N.x + aux0y*N.y + aux0z*N.z)\n",
    "\n",
    "A = [N]\n",
    "Dmc = []\n",
    "Po = []\n",
    "P = [P0]\n",
    "q1 = []\n",
    "u1 = []\n",
    "l = []\n",
    "\n",
    "for i in range(1, n+1):\n",
    "    P.append(me.Point('P' + str(i)))\n",
    "    Dmc.append(me.Point('Dmc' + str(i)))\n",
    "    Po.append(me.Point('Po' + str(i)))\n",
    "    A.append(me.ReferenceFrame('A' + str(i)))\n",
    "\n",
    "    q1.append([me.dynamicsymbols('q' + k + str(i)) for k in ('x', 'y', 'z')])\n",
    "    u1.append([me.dynamicsymbols('u' + k + str(i)) for k in ('x', 'y', 'z')])\n",
    "    \n",
    "    l.append(sm.symbols('l' + str(i)))\n",
    "    \n",
    "    \n",
    "m, mo, g, r, k, kb, d, l0, ln, beta, ctau, mu = sm.symbols('m, mo, g, r, k, kb, d, l0, ln, beta, ctau, mu')\n",
    "iXX, iYY, iZZ, t = sm.symbols('iXX, iYY, iZZ, t')\n",
    "rhodtmax = [sm.symbols('rhodt_' + str(i) + str(j)) for i, j in permutations(range(n), r=2)]\n",
    "\n",
    "rot  = []    # used for the kinematic equations below.\n",
    "rot1 = []   # dto.\n",
    "\n",
    "for i in range(1, n+1):\n",
    "    A[i].orient_body_fixed(A[i-1], (q1[i-1][1], q1[i-1][0], q1[i-1][2]), '213')\n",
    "    rot.append(A[i].ang_vel_in(A[0]))\n",
    "    A[i].set_ang_vel(A[i-1], u1[i-1][0]*A[i].x + u1[i-1][1]*A[i].y + u1[i-1][2]*A[i].z)\n",
    "    rot1.append(A[i].ang_vel_in(A[0]))\n",
    "    \n",
    "    P[i].set_pos(P[i-1], l[i-1]*A[i].y)\n",
    "    P[i].v2pt_theory(P[i-1], A[0], A[i])\n",
    "    \n",
    "    Dmc[i-1].set_pos(P[i-1], l[i-1]/2.*A[i].y)\n",
    "    Dmc[i-1].v2pt_theory(P[i-1], A[0], A[i])\n",
    "    \n",
    "    Po[i-1].set_pos(Dmc[i-1], beta*r*A[i].x)\n",
    "    Po[i-1].v2pt_theory(Dmc[i-1], A[0], A[i])\n",
    "    \n",
    "Pn = me.Point('Pn')\n",
    "Pn.set_pos(P[0], ln * A[0].x)\n",
    "Pn.set_vel(A[0], auxnx*A[0].x + auxny*A[0].y + auxnz*A[0].z)\n",
    "\n",
    "\n",
    "gleichung = []\n",
    "for i in range(n):\n",
    "    for uv in A[i+1]:\n",
    "        gleichung.append(me.dot(rot[i]-rot1[i], uv))\n",
    "variablen = [sm.Derivative(j, t) for i in range(n) for j in q1[i]]\n",
    "\n",
    "antwort = sm.solve(gleichung, variablen)\n",
    "zaehler = 0.\n",
    "for zahl in variablen:\n",
    "    zaehler += antwort[zahl].count_ops(visual=False)\n",
    "print(f'antwort contains {zaehler:.0f} operations')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is needed to get initial generalized coordinates and speeds compatible with the configuration constraint and the resulting speed constraints.\\\n",
    "Here the dictionary *antwort* is needed.\\\n",
    "lnitial generalized coordinates and generalized speeds, which are compatible with the configuration constraint and the resulting speed constraints are found numerically before the integration starts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "distanz = [P[i+1].pos_from(P[i]) for i in range(n)]\n",
    "hol  = sum(distanz) - ln*A[0].x\n",
    "holm = me.dot(hol, hol)\n",
    "\n",
    "holx = me.dot(hol, A[0].x)\n",
    "holy = me.dot(hol, A[0].y)\n",
    "holz = me.dot(hol, A[0].z)\n",
    "\n",
    "hol_mat = sm.Matrix([holx, holy, holz])\n",
    "        \n",
    "holxdt = holx.diff(t).subs(antwort)\n",
    "holydt = holy.diff(t).subs(antwort)\n",
    "holzdt = holz.diff(t).subs(antwort)\n",
    "\n",
    "speed_mat = sm.Matrix([holxdt, holydt, holzdt])\n",
    "\n",
    "matrix_A = speed_mat.jacobian((u1[-2][2], u1[-1][0], u1[-1][2]))\n",
    "vector_b = speed_mat.subs({u1[-2][2]: 0., u1[-1][0]: 0., u1[-1][2]: 0.})\n",
    "loesung  = matrix_A.LUsolve(-vector_b)\n",
    "\n",
    "if term_info == True:\n",
    "    print('loesung DS', set().union(*[me.find_dynamicsymbols(loesung[i]) for i in range(3)]))\n",
    "    print('loesung has {} operations'.format(sum([loesung[i].count_ops(visual=False) for i in range(3)])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This function calculates the **forces** and the **torques** when two balls collide.\n",
    "\n",
    "1.\\\n",
    "The impact force is on the line $\\overline{P_1 P_2}$. I use Hunt Crossley's method to calculate it\\\n",
    "It's direction is $P_1 \\to P_2$.\\\n",
    "2.\\\n",
    "The friction force acting on the contact point $CP_2$ is proportional to the component of relative speeds of the contact points lying in the plane perpendicular to $\\overline{P_1 P_2}$, the impact force and a friction factor $m_{u}$.\\\n",
    "It is equivalent to a force through $P_2$ and to a torque acting on $A_2$, the ball's fixed frame.\n",
    "\n",
    "\n",
    "**Note about the force during the collisions**\n",
    "\n",
    " **Hunt Crossley's method**\n",
    " \n",
    "My reference is this article, given to me by JM\\\n",
    "https://www.sciencedirect.com/science/article/pii/S0094114X23000782 \\\n",
    "\n",
    " \n",
    "This is with dissipation during the collision, the general force is given in (63) as\\\n",
    "$f_n = k_0 \\cdot \\rho + \\chi \\cdot \\dot \\rho$, with $k_0$ as above, $\\rho$ the penetration, and $\\dot\\rho$ the speed of the penetration.\\\n",
    "In the article it is stated, that $n = \\frac{3}{2}$ is a good choice, it is derived in Hertz' approach. Of course, $\\rho, \\dot\\rho$ must be the signed magnitudes of the respective vectors.\n",
    "\n",
    "A more realistic force is given in (64) as:\\\n",
    "$f_n = k_0 \\cdot \\rho^n + \\chi \\cdot \\rho^n\\cdot \\dot \\rho$, as this avoids discontinuity at the moment of impact.\n",
    "\n",
    "**Hunt and Crossley** give this value for $\\chi$, see table 1:\n",
    "\n",
    "$\\chi = \\dfrac{3}{2} \\cdot(1 - c_\\tau) \\cdot \\dfrac{k_0}{\\dot \\rho^{(-)}}$, \n",
    "where $c_\\tau = \\dfrac{v_1^{(+)} - v_2^{(+)}}{v_1^{(-)} - v_2^{(-)}}$, where $v_i^{(-)}, v_i^{(+)}$ are the speeds of $body_i$, before and after the collosion, see (45), $\\dot\\rho^{(-)}$ is the speed right at the time the impact starts. $c_\\tau$ is an experimental factor, apparently around 0.8 for steel.\n",
    "\n",
    "Using (64), this results in their expression for the force:\n",
    "\n",
    "$f_n = k_0 \\cdot \\rho^n \\left[1 + \\dfrac{3}{2} \\cdot(1 - c_\\tau) \\cdot \\dfrac{\\dot\\rho}{\\dot\\rho^{(-)}}\\right]$\n",
    "\n",
    "with $k_0 = \\frac{4}{3\\cdot(\\sigma_1 + \\sigma_2)} \\cdot \\sqrt{\\frac{R_1 \\cdot R_2}{R_1 + R_2}}$, where $\\sigma_i = \\frac{1 - \\nu_i^2}{E_i}$, with $\\nu_i$ = Poisson's ratio, $E_i$ = Young\"s modulus, $R_1, R_2$ the radii of the colliding bodies, $\\rho$ the penetration depth. All is near equations (54) and (61) of this article.\n",
    "\n",
    "As per the article, $n = \\frac{3}{2}$ is always to be used.\n",
    "\n",
    "*spring energy* =   $ k_0 \\cdot \\int_{0}^{\\rho} k^{3/2}\\,dk$ = $k_0 \\cdot\\frac{2}{5} \\cdot \\rho^{5/2}$\\\n",
    "I assume, the dissipated energy cannot be given in closed form, at least the article does not give one.\n",
    "\n",
    "*Note*  \n",
    "$c_\\tau = 1.$ gives **Hertz's** solution to the impact problem, also described in the article.\n",
    "\n",
    "\n",
    "**Friction when a ball hits a ball**\n",
    "\n",
    "This website:\n",
    "https://math.stackexchange.com/questions/2195047/solve-the-vector-cross-product-equation\n",
    "\n",
    "gives: $b = a \\times x \\rightarrow x = \\dfrac{b \\times a}{|a|^2}$\n",
    "\n",
    "This way, I can easily get the force acting on $P_2$ without any further geometric considerations. Of course the friction force has opposite direction to the relative speeds\\\n",
    "The friction force on $CP_2$  is equivalent to a force on $P_2$ (center of the ball) and a torque on $A_2$ (ball fixed frame)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def HC_disc(N, A1, A2, P1, P2, r, ctau, rhodtmax, k0):\n",
    "    '''\n",
    "This function returns the forces, torques on P2, when colliding with P1.\n",
    "Very similar to HC_wall above\n",
    "    '''    \n",
    "    vektor = P2.pos_from(P1)\n",
    "    richtung = vektor.normalize()\n",
    "    abstand  = vektor.magnitude()\n",
    "    rho = 2. * r - abstand       # positive in the ball has penetrated the wall\n",
    "    CP01 = me.Point('CP01')\n",
    "    CP01.set_pos(P1, 0.5 * abstand * richtung)\n",
    "    vCP01 = CP01.v2pt_theory(P1, N, A1)\n",
    "    rhodt = me.dot(vCP01, richtung)             #me.dot(P1.pos_from(P0).diff(t, N), richtung)\n",
    "    rho = sm.Max(rho, sm.S(0))\n",
    "    \n",
    "    forcec = k0 * rho**(3/2) * ( 1. + 3./2. * (1 - ctau) * rhodt / rhodtmax) * (richtung *\n",
    "            sm.Heaviside(2. * r - abstand, 0.))\n",
    "    \n",
    "    CP02 = me.Point('CP02')\n",
    "    CP02.set_pos(P2, -0.5 * abstand * richtung)\n",
    "    vCP02 = CP02.v2pt_theory(P2, N, A2)\n",
    "    \n",
    "    \n",
    "    friction_force = forcec.magnitude() * mu * -(CP02.vel(N) - CP01.vel(N))        \n",
    "    hilfs = CP02.pos_from(P2)\n",
    "    torque = hilfs.cross(friction_force) * sm.Heaviside(2. * r - abstand, 0.)\n",
    "    forcef = 1./me.dot(hilfs, hilfs) * torque.cross(hilfs) * sm.Heaviside(2. * r - abstand, 0.)\n",
    "        \n",
    "    return [forcec, forcef, torque]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Body list and list of forces for Kane's equations**\n",
    "\n",
    "If I assume friction between the balls, that is *frictionless = False*, the force vector created by Kane's equations below becomes unmanagably large. My iPad ran out of memory when trying to count the operations in it.\\\n",
    "The dampening (of the attachment of $P_n$ to $Pn$) is proportional to the speed of $P_n$ relative to $Pn$, which is fixed at $l_n \\cdot N$.x. However, only the component of the speed in $P_n Pn$ direction is relevant. I get it using the dot product.\\\n",
    "As it does not seem to help, but drastically blows up the force term of Kane's equations, I do not use it.\n",
    "\n",
    "As there may be numerical problems if $P_n$ is very close to $Pn$, I use the intermediate variable $test$, instead of using .*magnitude()*\n",
    "\n",
    "Maybe I do not do this the most efficient way:\\\n",
    "I ignore, that the impact force of $ball_i$ on $ball_j$ is equal in magnitude and opposite in direction of the impact force of $ball_j$ on $ball_i$, and same with friction and with torque. But this way it is easier to keep the indices 'straight'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "I = [me.inertia(A[i], iXX, iYY, iZZ) for i in range(1, n+1)]\n",
    "\n",
    "body = [me.RigidBody('body' + str(i+1), Dmc[i], A[i+1], m, (I[i], Dmc[i])) for i in range(0, n)]\n",
    "obs = [me.Particle('obs' + str(i+1), Po[i], mo) for i in range(0, n)]\n",
    "BODY = body + obs\n",
    "\n",
    "# Gravitational forces\n",
    "punkte = Dmc + Po\n",
    "massen = [m] * n + [mo] * n\n",
    "FL1 = [(punkt, -masse*g*A[0].y) for punkt, masse in zip(punkte, massen)] \n",
    "\n",
    "# Reaction forces\n",
    "FL_react = [(P[0], f0x*A[0].x + f0y*A[0].y + f0z*A[0].z), (Pn, fnx*A[0].x + fny*A[0].y + fnz*A[0].z)]\n",
    "\n",
    "# forces due to the balls hitting each other. My assumption is that never will 3 balls hit each other\n",
    "# at the same time.\n",
    "FB = []\n",
    "distanz_Dmc = [Dmc[i].pos_from(Dmc[j]).magnitude() for i, j in permutations(range(n), r=2)]\n",
    "rhodtmax1 = [me.dot(Dmc[i].pos_from(Dmc[j]).diff(t, N).subs(antwort), Dmc[i].pos_from(Dmc[j]).normalize()) \n",
    "            for i, j in permutations(range(n), r=2)]\n",
    "print('rhodtmax1 DS', set.union(*[me.find_dynamicsymbols(rhodtmax1[j]) \n",
    "                    for j in range(len(list(permutations(range(n), r=2))))]))\n",
    "zaehler = -1\n",
    "\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    zaehler += 1\n",
    "    FB.append((Dmc[i], HC_disc(N, A[j+1], A[i+1], Dmc[j], Dmc[i], r, ctau, rhodtmax[zaehler], kb)[0]))\n",
    "    if frictionless == False:\n",
    "        FB.append((Dmc[i], HC_disc(N, A[j+1], A[i+1], Dmc[j], Dmc[i], r, ctau, rhodtmax[zaehler], kb)[1]))\n",
    "        FB.append((A[i+1], HC_disc(N, A[j+1], A[i+1], Dmc[j], Dmc[i], r, ctau, rhodtmax[zaehler], kb)[2]))\n",
    "\n",
    "# Forces due to the spring from P_n to Pn\n",
    "PPF  = Pn.pos_from(P[n])\n",
    "PPFx = sm.Abs(me.dot(PPF, A[0].x))\n",
    "PPFy = sm.Abs(me.dot(PPF, A[0].y))\n",
    "PPFz = sm.Abs(me.dot(PPF, A[0].z))\n",
    "PPF_len = sm.Abs(sm.sqrt(PPFx**2 + PPFy**2 + PPFz**2))\n",
    "\n",
    "test = 1. / sm.Max(PPF_len, 1.e-20)   # to avoid divisions by zero\n",
    "FL2 = [(P[n], -k * (l0 - PPF_len) * PPF * test), (Pn, k * (l0 - PPF_len) * PPF * test)]\n",
    "\n",
    "#friction forces\n",
    "# Adding this force term increased the number of operations in force by a factor of around 10.\n",
    "if daempfung == True:\n",
    "    PPFdt = PPF.diff(t, N)\n",
    "    PPFdt_wirk = me.dot(PPFdt, PPF * test ) # only the speed in PPF direction is relevant\n",
    "    FL3 = [(P[n], d * PPFdt_wirk * PPF * test), (Pn, -d * PPFdt_wirk * PPF * test)]\n",
    "else:\n",
    "    FL3 = []\n",
    "FL = FL1 + FL2 + FL3 + FL_react + FB"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Kane's equations**\n",
    "\n",
    "Note, that I do **not** use velocity constraints, but use the spring to enforce the configuration constraint.\\\n",
    "I tried velocity constraints before I came up with the spring, but the numerical integration would only work over short times. I do not know, why this is so.\n",
    "\n",
    "Note the frames used for setting up the kinematic equations. Using other frames may drastically increase the number of operations in the *force vektor*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# q1, u1 need to be 'flattened' for use in Kane's set up.\n",
    "q1h = [q1[i][j] for i in range(0, len(q1)) for j in range(3)]\n",
    "u1h = [u1[i][j] for i in range(0, len(u1)) for j in range(3)]\n",
    "\n",
    "q_ind = q1h\n",
    "u_ind = u1h\n",
    "kd = [me.dot(rot[i] - rot1[i], uv) for i in range(n) for uv in A[i+1]]\n",
    "#raise Exception()\n",
    "KM = me.KanesMethod(N, q_ind=q_ind, u_ind=u_ind, kd_eqs=kd, u_auxiliary=aux)\n",
    "(fr, frstar) = KM.kanes_equations(BODY, FL)\n",
    "\n",
    "MM = KM.mass_matrix_full\n",
    "if term_info == True:\n",
    "    print('MM DS', me.find_dynamicsymbols(MM))\n",
    "    print('MM free symbols', MM.free_symbols)\n",
    "    print('MM contains {} operations'.format(sum([MM[i, j].count_ops(visual=False) for i in range(MM.shape[0]) for j in range(MM.shape[1])])), '\\n')\n",
    "\n",
    "force = KM.forcing_full.subs({i: sm.S(0.) for i in aux})\n",
    "if term_info == True:\n",
    "    print('force DS', me.find_dynamicsymbols(force))\n",
    "    print('force free symbols', force.free_symbols)\n",
    "    print('force contains {} operations'.format(sum([force[i].count_ops(visual=False) for i in range(force.shape[0])])), '\\n')\n",
    "\n",
    "print('it took {:.3f} sec to get Kanes equations'.format(time.time() - start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the **reaction forces**\\\n",
    "For some reason the virtual speeds appear in *eingepraegt*. As per Jason Moore's advise, I set them to zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start1 = time.time()\n",
    "eingepraegt_dict = {sm.Derivative(i, t): j for i, j in zip(u1h, rhs)}\n",
    "eingepraegt = (KM.auxiliary_eqs.subs(eingepraegt_dict)).subs({i: sm.S(0) for i in aux})\n",
    "if term_info == True:\n",
    "    print('shape of eingepraegt', eingepraegt.shape)\n",
    "    print('eingepraegt DS', me.find_dynamicsymbols(eingepraegt))\n",
    "    print('eingepraegt free symbols', eingepraegt.free_symbols)\n",
    "    print('eingepraegt contains {} operations'.format(sum([eingepraegt[i].count_ops(visual=False) for i in range(eingepraegt.shape[0])])), '\\n')\n",
    "\n",
    "print('it took {:.3f} sec to get the reaction forces'.format(time.time() - start1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Always good to look at the energies of a system. Absent any friction, the total energy must be constant (a necessary, but of course not a sufficient condition). This has helped me often to correct the set up of Kane's equations.\\\n",
    "The virtual speeds in aux must be set to zero to for the kinetic energy.\n",
    "\n",
    "- *spring_energie:*$\\hspace{50 mm} \\text{energy stored in the spring connecting $P_n$ to $Pn$}$\n",
    "- *spring2_energie:*$\\hspace{48 mm} \\text{energy stored in the elastic balls during a collission}$\n",
    "\n",
    "P4pos used to check the accuracy of the location of $P_n$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kin_dict = {i: 0. for i in aux}\n",
    "kin_energie = sum([koerper.kinetic_energy(N).subs(kin_dict) for koerper in BODY])\n",
    "pot_energie = sum([masse*g*me.dot(teil.pos_from(P[0]), N.y) for teil, masse in zip(punkte, massen)])\n",
    "spring_energie = 0.5 * k * (l0 -  PPF_len)**2\n",
    "\n",
    "spring2_energie = 0.\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    abstand = Dmc[i].pos_from(Dmc[j]).magnitude()\n",
    "    rho = sm.Max(2.*r - abstand, sm.S(0))\n",
    "    spring2_energie += 2./5. * kb * rho**(5/2) * sm.Heaviside(rho, sm.S(0)) \n",
    "spring2_energie = spring2_energie / 2.\n",
    "              \n",
    "P4pos = [me.dot(Pn.pos_from(P[n]), uv) for uv in N]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert the sympy functions to numpy functions to do numeric calculations. *cse=True* speeds up the numerical integration a lot (in this case, without *cse=True* it would not be practical on my machine), it may slow down the lambdification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start1 = time.time()\n",
    "qL  = q1h + u1h\n",
    "pL  = [m, mo, g, r, k, kb, d, l0, ln, beta, mu, ctau] + l + [iXX, iYY, iZZ] + rhodtmax\n",
    "pL1 = [m, mo, g, r, k, kb, d, l0, ln, beta, mu, ctau] + l + [iXX, iYY, iZZ]\n",
    "\n",
    "# q_ind1, q_dep1, u_ind1, up_dep1 are needed further down in the lambdification \n",
    "u_ind1 = deepcopy(u1h) # deepcopy needed to avoid 'destroying' u1h\n",
    "hilfs1 = u_ind1.pop(-1)\n",
    "hilfs2 = u_ind1.pop(-2)\n",
    "hilfs3 = u_ind1.pop(-2) \n",
    "u_dep1 = [hilfs3, hilfs2, hilfs1]\n",
    "\n",
    "q_ind1 = deepcopy(q1h)\n",
    "hilfs1 = q_ind1.pop(-1)\n",
    "hilfs2 = q_ind1.pop(-2)\n",
    "hilfs3 = q_ind1.pop(-2)\n",
    "q_dep1 = [hilfs3, hilfs2, hilfs1]\n",
    "\n",
    "MM_lam = sm.lambdify(qL + pL, MM, cse=True)\n",
    "force_lam = sm.lambdify(qL + pL, force, cse=True)\n",
    "\n",
    "print(' it took {:.3f} sec for the lambdification of force and of MM'.format(time.time() - start1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eingepraegt_lam = sm.lambdify(F + qL + pL + rhs, eingepraegt, cse=True)\n",
    "\n",
    "# find initial values of the generalized coordinates which satisfy the configuration constraint\n",
    "holm_lam = sm.lambdify(q1h + pL, holm, cse=True)\n",
    "\n",
    "# to get the initial values of the dependent speeds\n",
    "speed_mat1_lam = sm.lambdify(u_dep1 + q1h + u_ind1 + pL, speed_mat, cse=True)\n",
    "speed_mat11_lam = sm.lambdify(qL + pL, speed_mat, cse=True)\n",
    "\n",
    "loesung_lam = sm.lambdify(q1h + u_ind1 + pL, loesung, cse=True)\n",
    "\n",
    "kin_lam = sm.lambdify(qL + pL, kin_energie, cse=True)\n",
    "pot_lam = sm.lambdify(qL + pL, pot_energie, cse=True)\n",
    "spring_lam = sm.lambdify(qL + pL, spring_energie, cse=True)\n",
    "spring2_lam = sm.lambdify(qL + pL, spring2_energie, cse=True)\n",
    "\n",
    "loc_lam = sm.lambdify(q_dep1 + q_ind1 + pL, hol_mat, cse=True)\n",
    "\n",
    "crash_lam = sm.lambdify(q1h + pL, distanz_Dmc, cse=True)\n",
    "\n",
    "matrix_A_lam = sm.lambdify(q1h + u_ind1 + pL, matrix_A, cse=True)\n",
    "P4pos_lam = sm.lambdify(qL + pL, P4pos, cse=True)\n",
    "\n",
    "rhodtmax1_lam   = sm.lambdify(qL + pL1, rhodtmax1, cse=True)\n",
    "distanz_Dmc_lam = sm.lambdify(qL + pL1, distanz_Dmc, cse=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Numerical Integration**\n",
    "\n",
    "1.\\\n",
    "In the case of a 3D link, it is not so trivial to guess initial conditions which meet the configuration constraint. I try to find good approximations by minimizing the length of the configuration constraint.\n",
    "Then this initial result is improved by numerically by solving the configuration constraint for three generalized coordinates. (This n link pendulum has $ 3 \\cdot n -3$ independent generalized coordinates and 3 dependent coordinates)\\\n",
    "There are \"many\" feasible initial conditions to meet the configuration constraint. Different initial guesses generally lead to dfferent solutions minimizing the length of the configuration constraint. Not all of them lead to 'reasonable' results. I assume these are numerical problems with such large equations - but of course I do not know. \n",
    "\n",
    "2.\\\n",
    "The initial generalized coordinates must be such that the balls are separated.\\\n",
    "if they are not, an *Exception* will be raised.\n",
    "\n",
    "3.\\\n",
    "I use the speed constraints to get initial generalized speeds which are consistent.\n",
    "\n",
    "\n",
    "4.\\\n",
    "The speed constraints are fulfilled better if $n = 3$ or $n = 4$ than with $n = 2$.\n",
    "\n",
    "5.\\\n",
    "Initial speeds much different from zero may cause problems.\n",
    "\n",
    "All else is standard."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#============================\n",
    "# Input variables\n",
    "# the names are the same as in setting up Kane's equations, except that a \"1\" is appended.\n",
    "#============================\n",
    "m1  = 1.e0\n",
    "mo1 = 1.e-1\n",
    "g1  = 9.8\n",
    "r1  = 1.\n",
    "k1  = 1.e7\n",
    "kb1 = 1.e4    \n",
    "d1  = 1.e3\n",
    "l01 = 1.e-7\n",
    "\n",
    "mu1  = 0.\n",
    "ctau1 = 0.8\n",
    "\n",
    "\n",
    "l1 = [4. + 1. * (i+1) for i in range(len(l))]\n",
    "ln1 = np.sum(l1) / 2.5\n",
    "print(\"length of the rods:\", l1)\n",
    "print(\"distance from P0 to Pn, that is ln:\", ln1, \"\\n\")\n",
    "beta1 = .95\n",
    "\n",
    "u11 = [(-1.)**i * 0.  for i in range(len(u1h))]\n",
    "rhodtmax1 = [1. for i, j in permutations(range(n), r=2)]\n",
    "intervall = 25.\n",
    "#=====================================================================================\n",
    "schritte = int(intervall * 3000)\n",
    "times = np.linspace(0., intervall, schritte)\n",
    "\n",
    "iXX1 = 2./5. * m1 * r1**2    #from the internet.\n",
    "iYY1, iZZ1 = iXX1, iXX1\n",
    "\n",
    "pL_vals  = [m1, mo1, g1, r1, k1, kb1, d1, l01, ln1, beta1, mu1, ctau1\n",
    "    ] + l1 +  [iXX1, iYY1, iZZ1] + rhodtmax1\n",
    "pL1_vals = [m1, mo1, g1, r1, k1, kb1, d1, l01, ln1, beta1, mu1, ctau1] + l1 +  [iXX1, iYY1, iZZ1]\n",
    "\n",
    "def func(x0, args):\n",
    "    return holm_lam(*x0, *args)\n",
    "\n",
    "# find good consistent initial generalized coordinates\n",
    "x0 = [1.] * len(q1h) # initial guess. Different guesses may lead to different solutions.\n",
    "for i in range(10):\n",
    "    anfangs_q = minimize(func, x0, pL_vals)\n",
    "    anfangs_q = list(anfangs_q.x)\n",
    "    x0 = anfangs_q\n",
    "print(f\"config constraint violated by initial guess: {f'{holm_lam(*anfangs_q, *pL_vals):0.3e}':>20}\")\n",
    "\n",
    "# improve the initial generalized coordi8nates.\n",
    "def func2(x0, args):\n",
    "    return loc_lam(*x0, *args).reshape(3)\n",
    "\n",
    "anfangs_q1 = deepcopy(anfangs_q) # anfangs_q1 may be manipulated without affecting anfangs_q\n",
    "hilfs1 = anfangs_q1.pop(-1)\n",
    "hilfs2 = anfangs_q1.pop(-2)\n",
    "hilfs3 = anfangs_q1.pop(-2)\n",
    "\n",
    "args = anfangs_q1 + pL_vals\n",
    "x0 = [hilfs3, hilfs2, hilfs1]\n",
    "AAA = fsolve(func2, x0, args)\n",
    "anfangs_q[-4] = AAA[0]\n",
    "anfangs_q[-3] = AAA[1]\n",
    "anfangs_q[-1] = AAA[2]\n",
    "print(f\"after improvement, config constraint violated by: {f'{holm_lam(*anfangs_q, *pL_vals):0.3e}':>15}\", '\\n')\n",
    "\n",
    "# Setting the natural length of the spring a bit smaller that the initial error of the location of\n",
    "# P_n from Pn may help.\n",
    "pL_vals[7] = holm_lam(*anfangs_q, *pL_vals) / 2.\n",
    "\n",
    "# avoid initial conditions, where the balls are not separated from eacb other\n",
    "zaehler = -1\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    if i < j:\n",
    "        zaehler += 1\n",
    "        print(f\"initial distance Dmc_{i} from Dmc_{j} is {f'{crash_lam(*anfangs_q, *pL_vals)[zaehler]:0.3f}':>5}\")\n",
    "if np.min(crash_lam(*anfangs_q, *pL_vals)) <= 2. * r1:\n",
    "    raise Exception('Some balls are not separated')\n",
    "\n",
    "def func1(x0, args):\n",
    "    return speed_mat1_lam(*x0, *args).reshape(3)\n",
    "\n",
    "anfangs_u_ind = deepcopy(u11)\n",
    "anfangs_u_ind.pop(-1)\n",
    "anfangs_u_ind.pop(-2)\n",
    "anfangs_u_ind.pop(-2)\n",
    "\n",
    "hilfs1, hilfs2, hilfs3 = loesung_lam(*anfangs_q, *anfangs_u_ind, *pL_vals)\n",
    "hilfs1 = hilfs1[0]\n",
    "hilfs2 = hilfs2[0]\n",
    "hilfs3 = hilfs3[0]\n",
    "anfangs_u_dep = [hilfs1, hilfs2, hilfs3]\n",
    "\n",
    "args1 = anfangs_q + anfangs_u_ind + pL_vals\n",
    "print('\\n', 'Violation of speed constraints:', [f'{speed_mat1_lam(*anfangs_u_dep, *args1)[i][0]:0.3e}'\n",
    "                    for i in range(3)], \"\\n\")\n",
    "\n",
    "matrix_A1_det = np.linalg.det(matrix_A_lam(*anfangs_q, *anfangs_u_ind, *pL_vals))\n",
    "print(f'Determinante of matrix_A, used to determine u_dep = {matrix_A1_det:.3f} \\n')\n",
    "\n",
    "u11[-4] = hilfs1\n",
    "u11[-3] = hilfs2\n",
    "u11[-1] = hilfs3\n",
    "y0 = anfangs_q + u11 \n",
    "print('initial conditions are:', '\\n', 'y0 =', [f'{y0[i]:.3f}' for i in range(len(y0))], '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The actual **numerical integration** starts here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start2 = time.time()\n",
    "t_span = (0., intervall)\n",
    "def gradient(t, y, args):\n",
    "\n",
    "    '''\n",
    "Here I try to find the speed just before a collission takes place.\n",
    "    '''\n",
    "    zaehler  = -1\n",
    "    laenge = -len(list(permutations(range(n), r=2)))\n",
    "    for i, j in permutations(range(n), r=2):\n",
    "        zaehler += 1 \n",
    "        if 0. < 2.*r1 - distanz_Dmc_lam(*y, *pL1_vals)[zaehler] < 0.001:\n",
    "            args[laenge] = rhodtmax1_lam(*y, *pL1_vals)[zaehler]\n",
    "        laenge += 1\n",
    "\n",
    "    sol = np.linalg.solve(MM_lam(*y, *args), force_lam(*y, *args))\n",
    "    return np.array(sol).T[0]\n",
    "        \n",
    "resultat1 = solve_ivp(gradient, t_span, y0, t_eval = times, args=(pL_vals,), method='Radau')\n",
    "#   , max_step=max_step)        \n",
    "#    ,atol = 1.e-5, rtol = 1.e-5)\n",
    "\n",
    "resultat = resultat1.y.T\n",
    "print('resultat shape', resultat.shape, '\\n')\n",
    "event_dict = {-1: 'Integration failed', 0: 'Integration finished successfully', 1: 'some termination event'}\n",
    "print(event_dict[resultat1.status], ', message is:', resultat1.message, '\\n')\n",
    "\n",
    "print(\"To numerically integrate an intervall of {} sec the routine cycled {} times and it took {:.3f} sec \"\n",
    "      .format(intervall, resultat1.nfev, time.time() - start2), '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot whichever generalized coordinates you want to see."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schritte1 = resultat.shape[0]\n",
    "times1 = times[0: schritte1]\n",
    "bezeichnung = [str(i) for i in qL]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10,5))\n",
    "for i in (10, 13, 16):\n",
    "    ax.plot(times1, [resultat[l, i] for l in range(resultat.shape[0])], label=bezeichnung[i])\n",
    "ax.set_xlabel('time (sec)')\n",
    "if daempfung == True:\n",
    "    ax.set_title('Selected coordinates, with friction = {}'.format(d1))\n",
    "else:\n",
    "    ax.set_title('Selected coordinates')\n",
    "ax.legend();   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the errors in the location of $P_n$ from the desired location, that is from $l_n \\cdot N$.x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schritte1 = resultat.shape[0]\n",
    "times1 = times[0: schritte1]\n",
    "\n",
    "P4x_np = np.empty(schritte1)\n",
    "P4y_np = np.empty(schritte1)\n",
    "P4z_np = np.empty(schritte1)\n",
    "\n",
    "for i in range(schritte1):\n",
    "    P4x_np[i] = P4pos_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)[0]\n",
    "    P4y_np[i] = P4pos_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)[1]\n",
    "    P4z_np[i] = P4pos_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)[2]\n",
    "    \n",
    "    \n",
    "fig, ax = plt.subplots(figsize=(10,5))\n",
    "for i, j in zip((P4x_np, P4y_np, P4z_np), ('P_n X', 'P_n Y', 'P_n Z')):\n",
    "    ax.plot(times1, i, label=j)\n",
    "ax.set_xlabel('time (sec)')\n",
    "if daempfung == True:\n",
    "    print('Error in P[n] coordinates, friction = {}'.format(d1))\n",
    "else:\n",
    "    ax.set_title('Error in P[n] coordinates')\n",
    "ax.legend();\n",
    "mean_average = np.sqrt(np.mean(P4x_np)**2 + np.mean(P4y_np)**2 + np.mean(P4z_np)**2)\n",
    "print('average distance of Pn from required position:', mean_average)\n",
    "print('average standard deviation                    ', (np.std(P4x_np) + np.std(P4y_np)\n",
    "                    + np.std(P4z_np)) / 3.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the **energies** of the system. Absent any friction, that is $c_{tau} = 1$ and $m_u = 0.$ it should be zero.\\\n",
    "Of course, the system will continue to dissipate energy if dampening is unequal to zero, but as this blows up the equations dramatically, I do not use it.\\"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schritte1 = resultat.shape[0]\n",
    "times1 = times[0: schritte1]\n",
    "\n",
    "kin_np = np.empty(schritte1)\n",
    "pot_np = np.empty(schritte1)\n",
    "spring_np = np.empty(schritte1)\n",
    "spring2_np = np.empty(schritte1)\n",
    "total_np = np.empty(schritte1)\n",
    "\n",
    "for i in range(schritte1):\n",
    "    kin_np[i]     = kin_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)\n",
    "    pot_np[i]     = pot_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)\n",
    "    spring_np[i]  = spring_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)\n",
    "    spring2_np[i] = spring2_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)\n",
    "    total_np[i]   = kin_np[i] + pot_np[i] + spring_np[i] + spring2_np[i]\n",
    "    \n",
    "fig, ax = plt.subplots(figsize=(10,5))\n",
    "for i, j in zip((kin_np, pot_np, spring_np, spring2_np, total_np), \n",
    "    ('kin. energy', 'pot. energy', 'spring energy', 'spring2 energy', 'total energy')):\n",
    "    ax.plot(times1, i, label=j)\n",
    "ax.set_xlabel('time (sec)')\n",
    "if daempfung == True:\n",
    "    ax.set_title('Energy of system, spring constant = {}, dampening constant = {}'.format(k1, d1))\n",
    "else:\n",
    "    ax.set_title('Energy of the system')\n",
    "ax.legend();\n",
    "    \n",
    "if d1 == 0.:\n",
    "    maxtotal = np.max(total_np)\n",
    "    mintotal = np.min(total_np)\n",
    "    print('max deviation of total energy from being constant is {:.3f} % of max. total energy'.\n",
    "              format(np.abs((maxt - mint)/maxt) * 100))\n",
    "else:\n",
    "    maxtotal = np.max(total_np)\n",
    "    mintotal = np.min(total_np)\n",
    "    print('Total energy dissipated is {:.3f} Nm'.format(maxtotal - mintotal))   \n",
    "            \n",
    "\n",
    "# find when the balls collided, which balls collided, and whatthe the energy of the collission \n",
    "zusammenstoss = np.empty((schritte1, 3))\n",
    "zusammenstoss1 = []\n",
    "for i, zeit in enumerate(times1):\n",
    "    boxen = crash_lam(*[resultat[i, k] for k in range(int(resultat.shape[1]/2.))], *pL_vals)\n",
    "    zusammenstoss[i] = (0., 0., 0.)\n",
    "    for k_1 in range(int(len(list(permutations(range(n), r=2))))):\n",
    "        if boxen[k_1] <= 2. * r1:\n",
    "            staerke = spring2_lam(*[resultat[i, l] for l in range(resultat.shape[1])], *pL_vals)\n",
    "            zusammenstoss[i] = (k_1, zeit, staerke)\n",
    "\n",
    "for i in range(schritte1 - 1):\n",
    "    if zusammenstoss[i, 1] == 0. and zusammenstoss[i+1, 1] != 0.:\n",
    "        zusammenstoss1.append(zusammenstoss[i+1])\n",
    "        staerke_max = -1.\n",
    "        k_1 = i + 1\n",
    "        while k_1 < schritte1 and zusammenstoss[k_1, 1] != 0.:\n",
    "            staerke_max = max(staerke_max, zusammenstoss[k_1, 2])\n",
    "            k_1 += 1\n",
    "        zusammenstoss1[-1][2] = staerke_max\n",
    "\n",
    "print('\\n')\n",
    "print('When did which balls collide, and what was the energy of the collision?', '\\n')\n",
    "bezeichnung = [\"ball\" + str(j) + \" collided with ball\" + str(i) for i, j in permutations(range(n), r=2)]\n",
    "for i in range(len(zusammenstoss1)):\n",
    "    k_1, t11, energie = zusammenstoss1[i][0], zusammenstoss1[i][1], zusammenstoss1[i][2]\n",
    "    k_1 = int(k_1)\n",
    "    print('at {:.2f} sec {} with energy around {:.0f} Nm'.format(t11, bezeichnung[k_1], energie))\n",
    "print('in total {} collisions took place'.format(len(zusammenstoss1)))   \n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10,5))\n",
    "ax.plot(times1, spring_np)\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_title('Energy in the spring used to pull P[n] towards Pn')\n",
    "print('\\n')\n",
    "print(f'Average energy stored in the spring used to pull P[n] towards Pn is {np.mean(spring_np):.5f} Nm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the **hysteresis curves** of $ball_i$ colliding with $ball_j$\\\n",
    "The red numbers on the curves indicate the time of the impact."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zaehler1 = -1\n",
    "for i1, j1 in permutations(range(n), r=2):\n",
    "    HC_kraft = []\n",
    "    HC_displ = []\n",
    "    HC_times = []\n",
    "    zaehler  = 0\n",
    "    i0 = 0\n",
    "    zaehler1 += 1\n",
    "    for i in range(resultat.shape[0]):\n",
    "        abstand = 2.*r1  - distanz_Dmc_lam(*[resultat[i, j] for j in range(resultat.shape[1])], \n",
    "            *pL1_vals)[zaehler1]\n",
    "        if abstand < 0.:\n",
    "            i0 = i+1\n",
    "\n",
    "        if abstand >= 0. and i0 == i:\n",
    "            walldt = rhodtmax1_lam(*[resultat[i, j] for j in range(resultat.shape[1])], \n",
    "                                        *pL1_vals)[zaehler1]\n",
    "        if abstand >= 0.:\n",
    "            rhodt = rhodtmax1_lam(*[resultat[i, j] for j in range(resultat.shape[1])], \n",
    "                                      *pL1_vals)[zaehler1]\n",
    "            kraft0 = kb1 * abstand**(3/2) * (1. + 3./2. * (1 - ctau1) * rhodt/walldt)\n",
    "            HC_displ.append(abstand)\n",
    "            HC_kraft.append(kraft0)\n",
    "            HC_times.append((zaehler, times[i]))\n",
    "            zaehler +=1\n",
    "\n",
    "    HC_displ = np.array(HC_displ)\n",
    "    HC_kraft = np.array(HC_kraft)\n",
    "# print only, if there were collisions.    \n",
    "    if len(HC_displ) != 0 and i1 < j1:\n",
    "        fig, ax = plt.subplots(figsize=(10,5))\n",
    "        ax.plot(HC_displ, HC_kraft, color='green')\n",
    "        ax.set_xlabel('penetration depth (m)')\n",
    "        ax.set_ylabel('contact force (N)')\n",
    "        ax.set_title(f'hysteresis curves of successive impacts of ball_{i1} with ball_{j1}, ctau = {ctau1}, mu = {mu1}')\n",
    "    \n",
    "        zeitpunkte = 10\n",
    "        reduction = max(1, int(len(HC_times)/zeitpunkte))\n",
    "        for k in range(len(HC_times)):\n",
    "            if k % reduction == 0:\n",
    "                coord  = HC_times[k][0]\n",
    "                ax.text(HC_displ[coord], HC_kraft[coord], f'{HC_times[k][1]:.2f}', color=\"red\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Calculate the reaction forces on $P_0$ and $Pn$.*\n",
    "\n",
    "- numerically solve $rhs = MM^{-1} \\cdot force$\n",
    "- numerically solve $eingepraegt = 0.$ for the reaction forces in $F$. I iterate here 3 times, sometimes this improves the accuracy of the solution.\n",
    "- plot the reaction forces\n",
    "\n",
    "I need many time points to get 'nice' hysteresis curves. This makes this part very slow. I limit the points considered to around *zeitpunkte*. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times1 = []\n",
    "resultat1 = []\n",
    "\n",
    "#=======================\n",
    "zeitpunkte = 25000\n",
    "#=======================\n",
    "\n",
    "reduction = max(1, int(len(times)/zeitpunkte))\n",
    "\n",
    "for i in range(len(times)):\n",
    "    if i % reduction == 0:\n",
    "        times1.append(times[i])\n",
    "        resultat1.append(resultat[i])\n",
    "\n",
    "schritte1 = len(times1)\n",
    "resultat1 = np.array(resultat1)\n",
    "times1 = np.array(times1)\n",
    "print('number of points considered:',len(times1))\n",
    "\n",
    "# RHS calculated numerically, too large to do it symbolically. Needed for reaction forces.\n",
    "RHS1 = np.zeros((schritte1, resultat.shape[1]))\n",
    "for i in range(schritte1):\n",
    "    RHS1[i, :] = np.linalg.solve(MM_lam(*[resultat[i, j]for j in range(resultat.shape[1])], *pL_vals), \n",
    "        force_lam(*[resultat[i, j] for j in range(resultat.shape[1])], \n",
    "        *pL_vals)).reshape(resultat.shape[1])\n",
    "\n",
    "#calculate implied forces numerically\n",
    "def func (x, *args):\n",
    "# just serves to 'modify' the arguments for fsolve.\n",
    "    return eingepraegt_lam(*x, *args).reshape(6)\n",
    "\n",
    "kraft1x = np.empty(schritte1)\n",
    "kraft1y = np.empty(schritte1)\n",
    "kraft1z = np.empty(schritte1)\n",
    "kraft5x = np.empty(schritte1)\n",
    "kraft5y = np.empty(schritte1)\n",
    "kraft5z = np.empty(schritte1)\n",
    "\n",
    "x0 = tuple((1., 1., 1., 1., 1., 1.))   # initial guess\n",
    "\n",
    "for i in range(schritte1):\n",
    "    for _ in range(3):\n",
    "        y0 = [resultat[i, j] for j in range(resultat.shape[1])]\n",
    "        rhs1 = [RHS1[i, j] for j in range(3*n, 6*n)]\n",
    "        args = tuple(y0 + pL_vals + rhs1)\n",
    "        AF = fsolve(func, x0, args=args).reshape(6)\n",
    "        x0 = tuple(AF)      # improved guess. Should speed up convergence of fsolve\n",
    "    kraft1x[i] = AF[0]\n",
    "    kraft1y[i] = AF[1]\n",
    "    kraft1z[i] = AF[2]\n",
    "    kraft5x[i] = AF[3]\n",
    "    kraft5y[i] = AF[4]\n",
    "    kraft5z[i] = AF[5]\n",
    "    \n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "plt.plot(times1, kraft1x, label='force in X direction')\n",
    "plt.plot(times1, kraft1y, label='force in Y direction')\n",
    "plt.plot(times1, kraft1z, label='force in Z direction')\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('force (N)')\n",
    "ax.set_title('Reaction Forces on P0')\n",
    "plt.legend();\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "plt.plot(times1, kraft5x, label='force in X direction')\n",
    "plt.plot(times1, kraft5y, label='force in Y direction')\n",
    "plt.plot(times1, kraft5z, label='force in Z direction')\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('force (N)')\n",
    "ax.set_title('Reaction Forces on Pn')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "Plot the motion of the system as a projection onto the X / Z plane. The large circles are the balls, the smaller circles are the points $P_1...P_n$. Their color indicates their height above the lowest point.\\\n",
    "The balls are shown 'correctly' in that the 'darker' ball is shown to be 'under' the lighter ball. I did not do this with the points $P_1....P_n$, just to many cases to distinguish.\n",
    "\n",
    "I need many time points to get nice hysteresis curves. This makes the animation very slow. I limit the time points considered to around *zeitpunkte*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start4 = time.time()\n",
    "\n",
    "Dmc_loc = [[me.dot(punkt.pos_from(P[0]), uv) for uv in N] for punkt in Dmc]\n",
    "Po_loc  = [[me.dot(punkt.pos_from(P[0]), uv) for uv in N] for punkt in Po]\n",
    "Pdc = copy(P)\n",
    "Pdc.pop(0)\n",
    "P_loc   = [[me.dot(punkt.pos_from(P[0]), uv) for uv in N] for punkt in Pdc]\n",
    "\n",
    "Dmc_lam = sm.lambdify(qL + pL, Dmc_loc, cse = True)\n",
    "Po_lam = sm.lambdify(qL + pL, Po_loc, cse = True)\n",
    "P_lam = sm.lambdify(qL + pL, P_loc, cse = True)\n",
    "\n",
    "times1 = []\n",
    "resultat1 = []\n",
    "\n",
    "#=======================\n",
    "zeitpunkte = 500\n",
    "#=======================\n",
    "\n",
    "reduction = max(1, int(len(times)/zeitpunkte))\n",
    "\n",
    "for i in range(len(times)):\n",
    "    if i % reduction == 0:\n",
    "        times1.append(times[i])\n",
    "        resultat1.append(resultat[i])\n",
    "\n",
    "schritte1 = len(times1)\n",
    "resultat1 = np.array(resultat1)\n",
    "times1 = np.array(times1)\n",
    "print('number of points considered:',len(times1))\n",
    "\n",
    "\n",
    "Dmc_np = np.empty((schritte1, n, 3))\n",
    "Po_np  = np.empty((schritte1, n, 3))\n",
    "P_np   = np.empty((schritte1, n, 3))\n",
    "\n",
    "# generate the data for the animation\n",
    "for i in range(schritte1):\n",
    "    for j in range(n):\n",
    "        for k_1 in range(3):\n",
    "            Dmc_np[i, j, k_1] = Dmc_lam(*[resultat1[i, l] for l in range(resultat1.shape[1])], *pL_vals)[j][k_1]\n",
    "            Po_np[i, j, k_1]  = Po_lam(*[resultat1[i, l] for l in range(resultat1.shape[1])], *pL_vals)[j][k_1]\n",
    "            P_np[i, j, k_1]   = P_lam(*[resultat1[i, l] for l in range(resultat1.shape[1])], *pL_vals)[j][k_1]\n",
    "\n",
    "# determine the size of the animation\n",
    "minx = min(np.min(Dmc_np[:, :, 0].flatten()), np.min(P_np[:, :, 0].flatten()))\n",
    "minz = min(np.min(Dmc_np[:, :, 2].flatten()), np.min(P_np[:, :, 2].flatten()))\n",
    "maxx = max(np.max(Dmc_np[:, :, 0].flatten()), np.max(P_np[:, :, 0].flatten()))\n",
    "maxz = max(np.max(Dmc_np[:, :, 2].flatten()), np.max(P_np[:, :, 2].flatten()))\n",
    "skala = max(np.abs(minx), np.abs(maxx), np.abs(minz), np.abs(maxz))\n",
    "\n",
    "# get the color scale. For some reason, integers seem to be needed. In order to have a finer scale,\n",
    "# I multiply the heights of the points / balls by 100. before converting them to integers\n",
    "CDmc = np.empty((schritte1, n), int)\n",
    "CP   = np.empty((schritte, n), int)\n",
    "for i in range(schritte1):\n",
    "    for k_1 in range(n):\n",
    "        CDmc[i, k_1] = int(Dmc_np[i, k_1, 1] * 100.)\n",
    "        CP[i, k_1]   = int(P_np[i, k_1, 1] * 100.)\n",
    "        \n",
    "ymin = min(np.min(CDmc.flatten()), np.min(CP.flatten()))\n",
    "ymax = max(np.max(CDmc.flatten()), np.max(CP.flatten()))\n",
    "\n",
    "# This is to asign colors of 'plasma' to the points.\n",
    "Test = mp.colors.Normalize(ymin, ymax)\n",
    "Farbe = mp.cm.ScalarMappable(Test, cmap='plasma')\n",
    "\n",
    "farbeDmc = [Farbe.to_rgba(CDmc[0, k]) for k in range(n)] # colprs of starting positions\n",
    "farbeP = [Farbe.to_rgba(CP[0, k]) for k in range(n)]\n",
    "\n",
    "\n",
    "def animate_pendulum(times, Dmc_np, Po_np, P_np):\n",
    "    \n",
    "    fig, ax = plt.subplots(figsize=(8, 8), subplot_kw={'aspect': 'equal'})\n",
    "    fig.colorbar(Farbe, label='Height of the bodies/pointd above lowest level, factor = 100 \\n The actual heights are the numbers divided by factor', shrink=0.9, ax=ax)\n",
    "\n",
    "    ax.axis('on')\n",
    "    ax.set_xlim(-skala - 2., skala + 2.)\n",
    "    ax.set_ylim(-skala - 2., skala + 2.)\n",
    "    ax.set_xlabel('X axis', fontsize=13)\n",
    "    ax.set_ylabel('Z axis', fontsize=13)\n",
    "    \n",
    "    LINE1 = [] # connection of P_i to P_i+1\n",
    "    for i in range(n):  \n",
    "        LINE1.append(ax.plot([], [], 'bo', linestyle = '-', linewidth=0.5, markersize=1)) \n",
    "    \n",
    "    LINE2 = [] # observer on ball_i\n",
    "    for i in range(n):\n",
    "        LINE2.append(ax.plot([], [], color='red', marker='o', markersize=5)) \n",
    "    \n",
    "    ax.plot([0., 0.], [0., 0.], color='green', marker='o', markersize=6) # P0, fixed point\n",
    "    \n",
    "    kreisDL = []       # the balls   \n",
    "    kreisPL = []       # the points\n",
    "    for i in range(n):\n",
    "        kreisD = patches.Circle((Dmc_np[0, i, 0], Dmc_np[0, i, 2]), radius = r1, fill=True, \n",
    "                                color=farbeDmc[i], ec='black')  # balls\n",
    "        ax.add_patch(kreisD)\n",
    "    \n",
    "        kreisP = patches.Circle((P_np[0, i, 0], P_np[0, i, 2]), radius = r1/4., fill=True, \n",
    "                                color=farbeP[i], ec='black')  # observers\n",
    "        ax.add_patch(kreisP)\n",
    "        \n",
    "        kreisDL.append(kreisD)\n",
    "        kreisPL.append(kreisP)\n",
    "\n",
    "    def animate(i):        \n",
    "        ax.set_title('Projection of the motion onto the X / Z plane. \\n Red dots are the particles. \\n The green dot is P0 \\n Running time {:.2f} sec'.format(times[i]), fontsize=13)\n",
    "        \n",
    "# this is to ensure that the lower balls are shown to be below the higher balls.\n",
    "        ordnung  = [Dmc_np[i, k] for k in range(n)]\n",
    "        ordnung1 = sorted(ordnung, key = lambda test: test[1])\n",
    "        farbe_ordnung = [int(ordnung1[k][1] * 100) for k in range(n)]\n",
    "        farbeD = [Farbe.to_rgba(farbe_ordnung[k]) for k in range(n)]\n",
    "        \n",
    "        for j, kreis in enumerate(kreisDL):\n",
    "            kreis.set_center((ordnung1[j][0], ordnung1[j][2]))\n",
    "            kreis.set_color(farbeD[j])\n",
    "            kreis.set_ec(\"black\")\n",
    "            \n",
    "        for j, kreis in enumerate(kreisPL):\n",
    "            kreis.set_center((P_np[i, j, 0], P_np[i, j, 2]))\n",
    "            kreis.set_color(Farbe.to_rgba(CP[i, j]))\n",
    "            kreis.set_ec(\"black\")\n",
    "            \n",
    "# for whatever reason, the ax.plot(....) ist stored in LINE1 / LINE2 as a one-element-list, no directly/\n",
    "# hence linie[0] is needed. Same applies for the return.\n",
    "        for j, linie in enumerate(LINE1):\n",
    "            if j == 0:\n",
    "                x_values = [0., P_np[i, 0, 0]]\n",
    "                z_values = [0., P_np[i, 0, 2]]\n",
    "            else:\n",
    "                x_values = [P_np[i, j-1, 0], P_np[i, j, 0]]\n",
    "                z_values = [P_np[i, j-1, 2], P_np[i, j, 2]]             \n",
    "            linie[0].set_data(x_values, z_values)                  \n",
    "        \n",
    "        for j, linie in enumerate(LINE2):\n",
    "            linie[0].set_data(Po_np[i, j, 0], Po_np[i, j, 2])\n",
    "                                \n",
    "        return *LINE1[0], *LINE2[0]\n",
    "\n",
    "    anim = animation.FuncAnimation(fig, animate, frames=schritte1,\n",
    "                                   interval=1000*max(times1) / schritte1,\n",
    "                                   blit=True)\n",
    "    plt.close(fig)\n",
    "    return anim\n",
    "\n",
    "anim = animate_pendulum(times1, Dmc_np, Po_np, P_np)\n",
    "start5 = time.time()\n",
    "print('it took {:.3f} sec to run the animation, before HTML(..)'.format(start5 - start4))\n",
    "HTML(anim.to_jshtml())    # needed, when run on an iPad, I know no other way to do it. It is SLOW!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Animation using pythreejs.\n",
    "This is basically copied from a program by Jason Moore, just adapted to my needs here.\n",
    "\n",
    "NOTE: the 'reference frame' for pythreejs seems to be:\n",
    "- X - axis downwards, color red\n",
    "- Y - axis to the right, color green (hence:)\n",
    "- Z - axis pointing to the observer, color blue\n",
    "\n",
    "Rotation is used to transform my coordinate system used above to set up the equations of motion \n",
    "to the one prescribed by pythreejs.\n",
    "\n",
    "If you know from the beginning, that you want to use pythreejs it is probably better to use its\n",
    "orientation of coordinates, when setting up Kane's equations. Saves the trouble of guessing, which\n",
    "rotation is correct. I am not sure my rotation is fully correct, just played around until it 'looked' reasonable.\n",
    "\n",
    "If you want to draw a line between two point, only way I know is CylinderGeometry. In this case, the two points must be on the Y axis of the coordinate system, because the length of the CylinderGeometry extends in this direction. (there may be other ways to do it, I do not know them.)\n",
    "\n",
    "The points in time are large to get nice hysteresis curves. This will cause the animation to take a long time. I reduce the number of points in time to around *zeitpunkte*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start6 = time.time()\n",
    "\n",
    "times1 = []\n",
    "resultat1 = []\n",
    "\n",
    "#=======================\n",
    "zeitpunkte = 500\n",
    "#=======================\n",
    "\n",
    "reduction = max(1, int(len(times)/zeitpunkte))\n",
    "\n",
    "for i in range(len(times)):\n",
    "    if i % reduction == 0:\n",
    "        times1.append(times[i])\n",
    "        resultat1.append(resultat[i])\n",
    "\n",
    "schritte1 = len(times1)\n",
    "resultat1 = np.array(resultat1)\n",
    "times1 = np.array(times1)\n",
    "print('number of points considered:',len(times1))\n",
    "\n",
    "Dmc_loc = [[me.dot(punkt.pos_from(P[0]), uv) for uv in N] for punkt in Dmc]\n",
    "punkt_loc = [[me.dot(punkt.pos_from(P[0]), uv) for uv in N] for punkt in Po]\n",
    "Dmc_loc_lam = sm.lambdify(qL + pL, Dmc_loc, cse=True)\n",
    "punkt_loc_lam = sm.lambdify(qL + pL, punkt_loc, cse=True)\n",
    "\n",
    "\n",
    "winkel = sm.symbols('winkel')\n",
    "Rotation1 = sm.Matrix([[sm.cos(winkel), -sm.sin(winkel), 0], [sm.sin(winkel), sm.cos(winkel), 0], [0., 0., 1]])\n",
    "Rot_lam = sm.lambdify(winkel, Rotation1.T)\n",
    "Rotation = Rot_lam(np.pi/2.)\n",
    "\n",
    "TC_store = []\n",
    "TR_store = []\n",
    "TP_store = []\n",
    "body_mesh_store = []\n",
    "track_store = []\n",
    "farben = ['orange', 'blue', 'green', 'yellow', 'red']\n",
    "\n",
    "laengen = l1\n",
    "\n",
    "# the list of frames A contains the inertial frame A[0] = N. This must not be here\n",
    "AB = copy(A) # for whatever reason, deepcopy does not work here\n",
    "AB.pop(0)\n",
    "\n",
    "for i in range(n):\n",
    "#for its mass center\n",
    "    TC = sm.eye(4)\n",
    "    TC[:3, :3] = (AB[i].dcm(N)) * Rotation\n",
    "    TC = TC.reshape(16, 1)\n",
    "    TC_lam = sm.lambdify(qL + pL, TC)   \n",
    "\n",
    "    TR = sm.eye(4)\n",
    "    TR[:3, :3] = (AB[i].dcm(N)) * Rotation\n",
    "    TR = TR.reshape(16, 1)\n",
    "    TR_lam = sm.lambdify(qL + pL, TR) \n",
    "    \n",
    "    TP = sm.eye(4)\n",
    "    TP[:3, :3] = (AB[i].dcm(N)) * Rotation\n",
    "    TP = TP.reshape(16, 1)\n",
    "    TP_lam = sm.lambdify(qL + pL, TP) \n",
    "\n",
    "\n",
    "# store the information about the body, expressed in TAc for every time step.\n",
    "    TCs = []   # for the ball\n",
    "    TRs = []   # for the rod\n",
    "    TPs = []   # for the red dot\n",
    "    \n",
    "# Create the TAs, containing 'one TA' for each time step\n",
    "# resultat contains the results of the numeric integration. \n",
    "# where the numeric integration was evaluated\n",
    "    for k_1 in range(resultat1.shape[0]):\n",
    "        zeit = times[i]\n",
    "        TCi = TC_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)  # the balls\n",
    "        TRi = TR_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)  # the rod\n",
    "        TPi = TP_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)  # the dot\n",
    "\n",
    "# TAi[12], TAi[13], TAi[14] hold the location of A[i] w.r.t. N.\n",
    "# As the axis chosen for solving the equations of motion, and the axis given by pythreejs do not \n",
    "# coincide, the values for TAi[..] must be permutated accordingly.\n",
    "# of course here different locations for center of ball and center of mass.\n",
    "        TRi[12] = -Dmc_loc_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)[i][1]\n",
    "        TRi[13] =  Dmc_loc_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)[i][0]\n",
    "        TRi[14] =  Dmc_loc_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)[i][2]\n",
    "    \n",
    "        TCi[12] = -Dmc_loc_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)[i][1]\n",
    "        TCi[13] =  Dmc_loc_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)[i][0]\n",
    "        TCi[14] =  Dmc_loc_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)[i][2] \n",
    "\n",
    "        TPi[12] = -punkt_loc_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)[i][1]\n",
    "        TPi[13] =  punkt_loc_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)[i][0]\n",
    "        TPi[14] =  punkt_loc_lam(*[resultat1[k_1, l] for l in range(resultat1.shape[1])], *pL_vals)[i][2]\n",
    "\n",
    "        TRs.append(TRi.squeeze().tolist())\n",
    "        TCs.append(TCi.squeeze().tolist())\n",
    "        TPs.append(TPi.squeeze().tolist())\n",
    "       \n",
    "    TC_store.append(TCs)\n",
    "    TR_store.append(TRs)\n",
    "    TP_store.append(TPs)\n",
    "    \n",
    "# Create the objects, which will move\n",
    "# 1. The ball\n",
    "    body_geom_C = p3js.SphereGeometry(r1, 12, 12)\n",
    "    body_material_C = p3js.MeshStandardMaterial(color=farben[i], wireframe=False)\n",
    "    body_mesh_C = p3js.Mesh(geometry=body_geom_C, material=body_material_C, name='ball_' + str(i))\n",
    "\n",
    "# 2. Rod\n",
    "    body_geom_R = p3js.CylinderGeometry(radiusTop=0.05, radiusBottom=0.05, height=laengen[i], \n",
    "                    radialSegments=6, heightSegments=10, openEnded=False)\n",
    "    body_material_R = p3js.MeshStandardMaterial(color='black', wireframe=False)\n",
    "    body_mesh_R = p3js.Mesh(geometry=body_geom_R, material=body_material_R, name='rod_' + str(i))\n",
    "    \n",
    "# 3. the dot    \n",
    "    body_geom_P = p3js.SphereGeometry(0.25, 12, 12)\n",
    "    body_material_P = p3js.MeshStandardMaterial(color='green', wireframe=False)\n",
    "    body_mesh_P = p3js.Mesh(geometry=body_geom_P, material=body_material_P, name='punkt_' + str(i))\n",
    "\n",
    "# locate the body in 3D space and add the coordinate system of the body\n",
    "    body_mesh_R.matrixAutoUpdate = False\n",
    "    body_mesh_R.add(p3js.AxesHelper(0.1))  # length of the axis of the ball system A2\n",
    "    body_mesh_R.matrix = TR_store[i][0]             # starting point of the animation\n",
    "\n",
    "    body_mesh_C.matrixAutoUpdate = False\n",
    "    body_mesh_C.add(p3js.AxesHelper(0.01))    # length of the axis of the center of mass system A2\n",
    "    body_mesh_C.matrix = TC_store[i][0]          # starting point of the animation\n",
    "    \n",
    "    body_mesh_P.matrixAutoUpdate = False\n",
    "    body_mesh_P.add(p3js.AxesHelper(0.01))    # length of the axis of the center of mass system A2\n",
    "    body_mesh_P.matrix = TP_store[i][0]          # starting point of the animation\n",
    "    \n",
    "    \n",
    "    body_mesh_store.append(body_mesh_C)\n",
    "    body_mesh_store.append(body_mesh_R)\n",
    "    body_mesh_store.append(body_mesh_P)\n",
    "    \n",
    "\n",
    "# Create the 'picture'.\n",
    "# all the 'paramters' are taken by trial and error.\n",
    "view_width = 1200\n",
    "view_height = 400\n",
    "\n",
    "# Values just found by trial an error.\n",
    "p1, p2 = 5, 5\n",
    "if n == 2:\n",
    "    p3 = 18\n",
    "elif n == 3:\n",
    "    p3 = 22\n",
    "else:\n",
    "    p3 = 30\n",
    "\n",
    "camera = p3js.PerspectiveCamera(position=[p1, p2, p3],\n",
    "                                up=[-1.0, 0.0, 0.0],\n",
    "                                aspect=view_width/view_height)\n",
    "\n",
    "key_light = p3js.DirectionalLight(position=[0, 0, 10])\n",
    "ambient_light = p3js.AmbientLight()\n",
    "\n",
    "axes = p3js.AxesHelper(20)\n",
    "\n",
    "children = []\n",
    "for i in range(3*n):\n",
    "    children = children + [body_mesh_store[i], axes, camera, key_light, ambient_light]\n",
    "\n",
    "scene = p3js.Scene(children=children)\n",
    "controller = p3js.OrbitControls(controlling=camera)\n",
    "renderer = p3js.Renderer(camera=camera, scene=scene, controls=[controller],\n",
    "                         width=view_width, height=view_height)\n",
    "\n",
    "# Create the action, simply copied from JM's lecture.\n",
    "\n",
    "for i in range(n):\n",
    "    eigenname = 'ball_'+str(i)\n",
    "    track_C = p3js.VectorKeyframeTrack(\n",
    "        name=\"scene/\" + eigenname + \".matrix\",\n",
    "        times=times1,\n",
    "        values=TC_store[i])\n",
    "    \n",
    "    eigenname = 'rod_' + str(i)\n",
    "    track_R = p3js.VectorKeyframeTrack(\n",
    "        name=\"scene/\" + eigenname + \".matrix\",\n",
    "        times=times1,\n",
    "        values=TR_store[i])\n",
    "    \n",
    "    eigenname = 'punkt_' + str(i)\n",
    "    track_P = p3js.VectorKeyframeTrack(\n",
    "        name=\"scene/\" + eigenname + \".matrix\",\n",
    "        times=times1,\n",
    "        values=TP_store[i])\n",
    "    \n",
    "    track_store += [track_C] + [track_R] + [track_P]\n",
    "\n",
    "duration = times1[-1] - times1[0]\n",
    "clip = p3js.AnimationClip(tracks=track_store, duration=duration)\n",
    "action = p3js.AnimationAction(p3js.AnimationMixer(scene), clip, scene)\n",
    "\n",
    "# avoid error if the animation has not run before this runs.\n",
    "if start5 != None:\n",
    "    print(' it tooks {:.3f} sec to run the complete program'.format(time.time() - start6 + start5 - start ))\n",
    "\n",
    "renderer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will start the animation.\\\n",
    "As I know nothing about pythreejs, I cannot explain how it works.\\\n",
    "For some reason, with python 3.11 some error is raised, which did not happen with python 3.9. This error does not seem to affect the animation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "action"
   ]
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}