rolling-ball-uneven-street-collision-2.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 mpl_toolkits.mplot3d import axes3d\n",
    "%matplotlib inline\n",
    "from IPython.display import HTML\n",
    "mp.rcParams['animation.embed_limit'] = 2**128\n",
    "from matplotlib import animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Rausspringen(Exception):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Two uniform, solid balls with radius r are running on an uneven street without slipping. The balls are not allowed to jump, but are in contact with the street at all times. The reaction forces to hold one of them on the street without slipping are calculated.\\\n",
    "A square of vertical walls ( the positive Y direction ) surrounds the street. the four corners of the square have (X / Z) coordinate ($-l_W, -l_W$), ($l_W, -l_W$), ($l_W, l_W$), ($-l_W, l_W$)\\\n",
    "$wall_0$ is the top horizontal wall, then I count the rest clockwise.\n",
    "\n",
    "The basic shape of the street is modeled by strassen_form, the superimposed unevenness is \n",
    "modeled by strasse. (Strasse = Street in German)\n",
    "\n",
    "An observer, a particle of mass $m_o$, may be attached anywhere inside each ball. \n",
    "\n",
    "**Variables / Parameters**\n",
    "- $q_{1i}, q_{2i}, q_{3i}$: generalized coordinates of  $ball_i$\n",
    "- $u_{1i}, u_{2i}, u_{3i}$: its angular velocities\n",
    "- $x_i, z_i$: coordinates of the contact point of the ball with the street\n",
    "- $u_{xi}, u_{zi}$: their speeds\n",
    "- $aux_x, aux_y, aux_z, f_x, f_y, f_z$: virtual speeds and rection forces on the center of the ball\n",
    "\n",
    "\n",
    "- N: inertial frame\n",
    "- A1: frame fixed to $ball_1$\n",
    "- A2: frame fixed to $ball_2$\n",
    "\n",
    "\n",
    "- $P_0$: point fixed in N\n",
    "- $CP_i$: contact point of $ball_i$ with the street\n",
    "- $Dmc_i$: center of the $ball_i$\n",
    "- $m_{Dmc_i}$: position of observer of $ball_i$\n",
    "\n",
    "\n",
    "- rhodtwall: max speed just before collission of $ball_i$ with $wall_j$\n",
    "- rhodtmax: max speed just before collision between the walls\n",
    "\n",
    "\n",
    "- m: mass of the ball\n",
    "- $m_o$ mass of the particle on the ball\n",
    "- r: radius of the ball\n",
    "- $i_{XX}, i_{YY}, i_{ZZ}$: moments of inertia of the ball\n",
    "- amplitude, frequenz: parameters for the stree\n",
    "- $m_u$: friction between balls\n",
    "- $m_{uW}$: friction between balls and walls\n",
    "\n",
    "- $\\alpha, \\beta, \\gamma$: give the positions of the particle relative to the center of the ball"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "start1 = time.time()\n",
    "q11, q12, q13, q21, q22, q23 = me.dynamicsymbols('q11, q12, q13, q21, q22, q23') #q is the angle of the ball\n",
    "u11, u12, u13, u21, u22, u23 = me.dynamicsymbols('u11, u12, u13, u21, u22, u23')\n",
    "\n",
    "x1, z1, x2, z2 = me.dynamicsymbols('x1, z1, x2, z2')             \n",
    "ux1, uz1, ux2, uz2 = me.dynamicsymbols('ux1, uz1, ux2, uz2')\n",
    "\n",
    "rhodtwall = [[sm.symbols('rhodtwall' + str(j) + str(i)) for i in range(4)] for j in (1, 2)]\n",
    "print(rhodtwall)\n",
    "rhodtmax  = [sm.symbols('rhodtmax' + str(i)) for i in (1, 2)]\n",
    "# for the reaction forces. \n",
    "auxx, auxy, auxz, fx, fy, fz = me.dynamicsymbols('auxx, auxy, auxz, fx, fy, fz')\n",
    "\n",
    "\n",
    "m, mo, g, r, iXX, iYY, iZZ, k0, k0W, lW, ctau, mu, muW = sm.symbols(\n",
    "    'm, mo, g, r, iXX, iYY, iZZ, k0, k0W, lW, ctau, mu, muW')\n",
    "amplitude, frequenz, reibung, alpha, beta, gamma, t = sm.symbols('amplitude frequenz reibung alpha beta gamma t')\n",
    "\n",
    "N, A1, A2, A3 = sm.symbols('N, A1, A2, A3', cls=me.ReferenceFrame)\n",
    "P0, CP1, CP2, Dmc1, Dmc2, m_Dmc1, m_Dmc2 = sm.symbols('P0, CP1, CP2, Dmc1, Dmc2, m_Dmc1, m_Dmc2'\n",
    "            , cls=me.Point) \n",
    "P0.set_vel(N, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Define the street and find its minimum osculating circle**\n",
    "\n",
    "The larger the integer *rumpel* is, the more uneven the street.\n",
    "\n",
    "The radius of the wheel must be smaller than the smallest osculating circle (Schmiegekreis) of the road, as it must only touch the street an exactly *one* point only.\n",
    "In this 3D case, I will just take the smaller one of the two osculating circles in X  and Z direction.\n",
    "(I did not bother to try and find the correct formula for higher dimensional functions)\n",
    "\n",
    "\n",
    "Define the rotations of the frame A1, A2, A3, and the contact points CP1, CP2\n",
    "\n",
    "$rot, rot_1$ are used for the kinematic equations, that is $\\frac{d}{dt}q_i = f_i(u_1, u_2, u_3)$\n",
    "\n",
    "As the mass matrix and the force vector are large, getting the term info, such as dynamic symbols and number of operations takes time. With *term_info = False*, this may be suppressed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#============================================\n",
    "rumpel    = 2     # the higher the number the more 'uneven the street'.\n",
    "term_info = False # if True, term info is displayed\n",
    "#============================================\n",
    "x, z = sm.symbols('x, z')\n",
    "def gesamt(x, z):\n",
    "    strasse = sum([amplitude/j * (sm.sin(j*frequenz * x) + sm.sin(j*frequenz * z)) for j in range(1, rumpel)])\n",
    "    strassen_form = (frequenz/2. * x)**2  + (frequenz/2. * z)**2\n",
    "    return strassen_form + strasse\n",
    "\n",
    "# osculating circles in X / Z directions.\n",
    "r_max_x = (sm.S(1.) + (gesamt(x, z).diff(x))**2 )**sm.S(3/2)/gesamt(x, z).diff(x, 2)\n",
    "r_max_z = (sm.S(1.) + (gesamt(x, z).diff(z))**2 )**sm.S(3/2)/gesamt(x, z).diff(z, 2)\n",
    "\n",
    "A1.orient_body_fixed(N, (q11, q12, q13), '123')\n",
    "rot11 = (A1.ang_vel_in(N))\n",
    "A1.set_ang_vel(N, u11*A1.x + u12*A1.y + u13*A1.z)\n",
    "rot12 = (A1.ang_vel_in(N))\n",
    "\n",
    "A2.orient_body_fixed(N, (q21, q22, q23), '123')\n",
    "rot21 = (A2.ang_vel_in(N))\n",
    "A2.set_ang_vel(N, u21*A2.x + u22*A2.y + u23*A2.z)\n",
    "rot22 = (A2.ang_vel_in(N))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Determine the location of the center of the ball and the speed of (subsequent!) contact points**\n",
    "\n",
    "vektor is the normal to the streel at point (x, z), the contact points $C_i$. \n",
    "I found this formula in the internet.\n",
    "It should point 'inwards', hence the leading minus sign.\n",
    "$Dmc_i$ is in the direction of this normal at distance r from $CP_i$\n",
    "\n",
    "As the contact point $CP_i$ are momentarily fixed in $A_i$, one may take *v2pt_theory* to determine the speed of Dmc.\n",
    "\n",
    "\n",
    "**Relationship of x(t) to q(t)**:\n",
    "\n",
    "\n",
    "Obviously, $ x(t) = function(q(t), gesamt(x(t), z(t)), r) $.\n",
    "When the ball is rotated through an angle $q$, the arc length is $r \\cdot q(t)$.\n",
    "\n",
    "The arc length of a function f(k(t)) from 0 to $x(t)$ is: $ \\int_{0}^{x(t)}\\sqrt{1 + \\left(\\frac{d}{dk}f(k(t))\\right) ^2} \\,dk \\ $ (I found this in the internet)\n",
    "\n",
    "This gives the sought after relationship between $q(t)$ and $x(t)$:\n",
    "$ r \\cdot (-q(t))  =  \\int_{0}^{x(t)} \\sqrt{1 + \\left(\\frac{d}{dk}(gesamt(k(t), z(t))\\right) ^2} \\,dk \\ $, differentiated:\n",
    "- $ r \\cdot (-u)  = \\sqrt{1 + \\left(\\frac{d}{dx}(gesamt(x(t), z(t))\\right) ^2} \\cdot \\frac{d}{dt}(x(t)) $, that is solved for $\\frac{d}{dt}(x(t))$:\n",
    "- $ \\frac{d}{dt}(x(t)) = \\dfrac{-(r \\cdot u)} {\\sqrt{1 + \\left(\\frac{d}{dx}(gesamt(x(t), z(t))\\right) ^2}} $\n",
    "\n",
    "A similar formula holds for the speed in Z direction, hence I get\n",
    "\n",
    "$ \\frac{d}{dt}(x(t)) = \\dfrac{- (u3 \\cdot r)} {\\sqrt{1 + \\left(\\frac{d}{dx}(gesamt(x(t), z(t))\\right) ^2}} $\n",
    "\n",
    "$ \\frac{d}{dt}(z(t)) =  \\dfrac{(u1 \\cdot r)} {\\sqrt{1 + \\left(\\frac{d}{dz}(gesamt(x(t), z(t)))\\right) ^2}} $\n",
    "\n",
    "As speeds are vectors, they can be added to give the resultant speed.\n",
    "\n",
    "The + / - signs are a consequence of the 'right hand rule' for frames. These are the sought after first order differential equations for $(x(t), z(t))$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# location og the contact points, where the balls touch the street\n",
    "CP1.set_pos(P0, x1*N.x + gesamt(x1, z1)*N.y + z1*N.z)\n",
    "CP2.set_pos(P0, x2*N.x + gesamt(x2, z2)*N.y + z2*N.z)\n",
    "\n",
    "# first ball\n",
    "vektor = -(gesamt(x1, z1).diff(x1)*N.x - N.y + gesamt(x1, z1).diff(z1)*N.z)\n",
    "Dmc1.set_pos(CP1, r * vektor.normalize()) \n",
    "m_Dmc1.set_pos(Dmc1, r*(alpha*A1.x + beta*A1.y + gamma*A1.z))\n",
    "    \n",
    "OMEGA = A1.ang_vel_in(N)  \n",
    "u3_wirk = me.dot(OMEGA, N.z)\n",
    "u1_wirk = me.dot(OMEGA, N.x)\n",
    "rhsx1 = (-u3_wirk * r / sm.sqrt(1. + gesamt(x1, z1).diff(x1)**2))\n",
    "rhsz1 =  (u1_wirk * r / sm.sqrt(1. + gesamt(x1, z1).diff(z1)**2))\n",
    "\n",
    "# second ball\n",
    "vektor = -(gesamt(x2, z2).diff(x2)*N.x - N.y + gesamt(x2, z2).diff(z2)*N.z)\n",
    "Dmc2.set_pos(CP2, r * vektor.normalize()) \n",
    "m_Dmc2.set_pos(Dmc2, r*(alpha*A2.x + beta*A2.y + gamma*A2.z))\n",
    "    \n",
    "OMEGA = A2.ang_vel_in(N)  \n",
    "u3_wirk = me.dot(OMEGA, N.z)\n",
    "u1_wirk = me.dot(OMEGA, N.x)\n",
    "rhsx2 = (-u3_wirk * r / sm.sqrt(1. + gesamt(x2, z2).diff(x2)**2))\n",
    "rhsz2 =  (u1_wirk * r / sm.sqrt(1. + gesamt(x2, z2).diff(z2)**2))\n",
    "\n",
    "# this is needed at various places below\n",
    "rhs_dict = {sm.Derivative(i, t): j for i, j in zip((x1, z1, x2, z2), (rhsx1, rhsz1, rhsx2, rhsz2))}\n",
    "\n",
    "CP1.set_vel(N, CP1.pos_from(P0).diff(t, N).subs(rhs_dict)) \n",
    "CP2.set_vel(N, CP2.pos_from(P0).diff(t, N).subs(rhs_dict))\n",
    "\n",
    "Dmc1.set_vel(N, Dmc1.pos_from(P0).diff(t, N).subs(rhs_dict) + auxx*N.x + auxy*N.y + auxz*N.z)\n",
    "m_Dmc1.v2pt_theory(Dmc1, N, A1)\n",
    "\n",
    "Dmc2.set_vel(N, Dmc2.pos_from(P0).diff(t, N).subs(rhs_dict))\n",
    "m_Dmc2.v2pt_theory(Dmc2, N, A2)\n",
    "\n",
    "if term_info == True:\n",
    "    print('CP1 DS', me.find_dynamicsymbols(CP1.vel(N), reference_frame=N))    \n",
    "    print('vel(Dmc2) DS:',me.find_dynamicsymbols(Dmc2.vel(N), reference_frame=N))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Function which calculates the forces of an impact between a ball and a wall**\\\n",
    "1.\\\n",
    "The impact force is on the line normal to the wall, going through the center of the ball, $P_1$. I use Hunt Crossley's method to calculate it\\\n",
    "It's direction points from the wall to the ball.\\\n",
    "2.\\\n",
    "The friction force acting on the contact point $CP_0$ is proportional to the component of the speed of $CP_0$ in the plane of the wall, the impact force and a friction factor $m_{uW}$.\\\n",
    "It is equivalent to a force through $P_1$ and to a torque acting on $A_1$, the ball 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 wall**\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 CP0 ( = contact point of ball with wall), without any further geometric considerations. Of course the friction force has opposite to the speed of CP0\\\n",
    "The friction force on CP0 is equivalent to a force on P1 (center of the ball) and a torque on A1 (ball fixed frame)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def HC_wall(N, A1, P1, r, ctau, rhodtwall, k0W):\n",
    "\n",
    "    FC = []\n",
    "    FF = []\n",
    "    TF = []\n",
    "    TT = []\n",
    "    for l, richtung in enumerate((N.z, -N.x, -N.z, N.x)):\n",
    "        \n",
    "        abstand = lW + me.dot(P1.pos_from(P0), richtung) \n",
    "        rho = r - abstand       # positive in the ball has penetrated the wall\n",
    "        CP0 = me.Point('CP0')\n",
    "        CP0.set_pos(P1, -richtung)\n",
    "        vCP0 = CP0.v2pt_theory(P1, N, A1)\n",
    "        rhodt = me.dot(vCP0, richtung)             #me.dot(P1.pos_from(P0).diff(t, N), richtung)\n",
    "        rho = sm.Max(rho, sm.S(0))\n",
    "    \n",
    "        forcec = k0W * rho**(3/2) * ( 1. + 3./2. * (1 - ctau) * rhodt / rhodtwall[l]) * (richtung *\n",
    "            sm.Heaviside(r - abstand, 0.))\n",
    "    \n",
    "        friction_force = forcec.magnitude() * muW * (-CP0.vel(N))        \n",
    "        hilfs = CP0.pos_from(P1)\n",
    "        torque = hilfs.cross(friction_force) * sm.Heaviside(r - abstand, 0.)\n",
    "        forcef = 1./me.dot(hilfs, hilfs) * torque.cross(hilfs) * sm.Heaviside(r - abstand, 0.)\n",
    "        \n",
    "        FC.append(forcec)\n",
    "        FF.append(forcef)\n",
    "        TF.append(torque)\n",
    "    \n",
    "    for i in (FC, FF, TF):\n",
    "        TT.append(i)\n",
    "    return TT    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculates the **forces** and the **torques** when the two *balls* collide.\\\n",
    "Similar to the ideas above."
   ]
  },
  {
   "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": [
    "Various **functions** needed later.\\\n",
    "The bodies must be defined here, as they are needed for the kinetic energy.\n",
    "\n",
    "- distanzw: distance from the ball to the walls\n",
    "- abstand_baelle: distance of the two balls from each other\n",
    "- rhodtwall: speed right at impact. it is calculated same as rhodt in the function HC_wall\n",
    "- rhodtmax: speed right before the impact of the two balls\n",
    "- pot_energie: potential energy\n",
    "- kin_energie: kinetic energy\n",
    "- spring_energie: energy stored in the ball during collissions with the walls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "I1 = me.inertia(A1, iXX, iYY, iZZ)                                              \n",
    "Body1 = me.RigidBody('Body1', Dmc1, A1, m, (I1, Dmc1))\n",
    "observer1 = me.Particle('observer1', m_Dmc1, mo)\n",
    "\n",
    "I2 = me.inertia(A2, iXX, iYY, iZZ)                                              \n",
    "Body2 = me.RigidBody('Body2', Dmc2, A2, m, (I2, Dmc2))\n",
    "observer2 = me.Particle('observer2', m_Dmc2, mo)\n",
    "\n",
    "BODY = [Body1, Body2, observer1, observer2]\n",
    "\n",
    "#subs_dict = {sm.Derivative(x, t): rhsx, sm.Derivative(z, t): rhsz} # needed at several place below\n",
    "subs_dict1 = {sm.Derivative(i, t): j for i, j in zip((q11, q12, q13, q21, q22, q23)\n",
    "    , (u11, u12, u13, u21, u22, u23))}\n",
    "energie_dict = {i: 0. for i in (auxx, auxy, auxz, fx, fy, fz)}\n",
    "\n",
    "distanz_list = [[lW + me.dot(ball.pos_from(P0), richtung) for richtung in (N.z, -N.x, -N.z, N.x)]\n",
    "        for ball in (Dmc1, Dmc2)]\n",
    "abstand_baelle = Dmc1.pos_from(Dmc2).magnitude()\n",
    "\n",
    "rhodtwall_list = []\n",
    "for ball, frame in zip((Dmc1, Dmc2), (A1, A2)):\n",
    "    hilfs = []\n",
    "    for richtung in (N.z, -N.x, -N.z, N.x):\n",
    "        CP0 = me.Point('CP0')\n",
    "        CP0.set_pos(ball, -richtung)\n",
    "        vCP0 = CP0.v2pt_theory(ball, N, frame)\n",
    "        rhodt = me.dot(vCP0, richtung).subs(energie_dict)             #me.dot(P1.pos_from(P0).diff(t, N), richtung)\n",
    "        hilfs.append(rhodt)\n",
    "    rhodtwall_list.append(hilfs)\n",
    "\n",
    "richtung = Dmc1.pos_from(Dmc2).normalize()\n",
    "rhodtmax_list = [hilfs0 := me.dot(Dmc2.pos_from(Dmc1).diff(t, N), richtung).subs(rhs_dict), \n",
    "                 -hilfs0]\n",
    "if term_info == True:\n",
    "    print(me.find_dynamicsymbols(rhodtmax_list[0], reference_frame=N))\n",
    "    \n",
    "    hilfs1 = set()\n",
    "    hilfs2 = set()\n",
    "    hilfs3 = set()\n",
    "    hilfs4 = set()\n",
    "    for k in range(2):\n",
    "        hilfs1 = hilfs1.union(*[me.find_dynamicsymbols(distanz_list[k][l])for l in range(4)])\n",
    "        hilfs2 = hilfs2.union(*[distanz_list[k][l].free_symbols for l in range(4)])\n",
    "        hilfs3 = hilfs3.union(*[me.find_dynamicsymbols(rhodtwall_list[k][l])for l in range(4)])\n",
    "        hilfs4 = hilfs4.union(*[rhodtwall_list[k][l].free_symbols for l in range(4)])\n",
    "    print('distanz_list DS', hilfs1)\n",
    "    print('distanz_list FS', hilfs2, '\\n')\n",
    "    print('rhodtwall_list DS', hilfs3)\n",
    "    print('rhodtwall_list FS', hilfs4, '\\n')\n",
    "\n",
    "pot_energie = (m * g * me.dot(Dmc1.pos_from(P0), N.y) + mo * g * me.dot(m_Dmc1.pos_from(P0), N.y) +\n",
    "               m * g * me.dot(Dmc2.pos_from(P0), N.y) + mo * g * me.dot(m_Dmc2.pos_from(P0), N.y)         \n",
    "    ).subs(rhs_dict).subs(energie_dict)\n",
    "kin_energie = sum([koerper.kinetic_energy(N).subs(rhs_dict).subs(energie_dict) for koerper in BODY])\n",
    "\n",
    "spring_energie = 0.\n",
    "# contribution of the walls\n",
    "for k in range(2):\n",
    "    for i in range(4):\n",
    "        rho = sm.Max(r - distanz_list[k][i], 0.)\n",
    "        spring_energie += 2./5. * k0W * rho**(5/2) * sm.Heaviside(r - distanz_list[k][i], 0.)\n",
    "# contribution of the two balls colliding        \n",
    "rho = sm.Max(2. * r - Dmc1.pos_from(Dmc2).magnitude(), 0.)\n",
    "spring_energie += (2./5. * k0 * rho**(5/2) * sm.Heaviside(2. * r - Dmc1.pos_from(Dmc2).\n",
    "                    magnitude(), 0.))\n",
    "\n",
    "if term_info == True:\n",
    "    print('pot energy DS',    me.find_dynamicsymbols(pot_energie), '\\n')\n",
    "    print('kin energy DS',    me.find_dynamicsymbols(kin_energie), '\\n')\n",
    "    print('spring energy DS', me.find_dynamicsymbols(spring_energie),'\\n')\n",
    "\n",
    "# various points\n",
    "CP1_pos = [(me.dot(CP1.pos_from(P0), uv)).subs(rhs_dict) for uv in N]\n",
    "Dmc1_pos = [(me.dot(Dmc1.pos_from(P0), uv)).subs(rhs_dict) for uv in N]# for later plotting only\n",
    "m_Dmc1_pos = [(me.dot(m_Dmc1.pos_from(P0), uv)).subs(rhs_dict) for uv in N] #later plotting only\n",
    "\n",
    "CP2_pos = [(me.dot(CP2.pos_from(P0), uv)).subs(rhs_dict) for uv in N]\n",
    "Dmc2_pos = [(me.dot(Dmc2.pos_from(P0), uv)).subs(rhs_dict) for uv in N]# for later plotting only\n",
    "m_Dmc2_pos = [(me.dot(m_Dmc2.pos_from(P0), uv)).subs(rhs_dict) for uv in N] #later plotting only"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "Set the **forces and torques** for Kane's equations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "FL = [(Dmc1, -m*g*N.y), (m_Dmc1, -mo*g*N.y), (Dmc1, fx*N.x + fy*N.y + fz*N.z)] # gravity and reaction forces\n",
    "FL.append((Dmc2, -m*g*N.y))\n",
    "FL.append((m_Dmc2, -mo*g*N.y))\n",
    "\n",
    "# collisions with the walls\n",
    "zaehler = 0\n",
    "for ball, frame in zip((Dmc1, Dmc2), (A1, A2)):\n",
    "    for i in range(4):\n",
    "        FL.append((ball,   HC_wall(N, frame, ball, r, ctau, rhodtwall[zaehler], k0W)[0][i]))  # impact force\n",
    "        FL.append((ball,   HC_wall(N, frame, ball, r, ctau, rhodtwall[zaehler], k0W)[1][i]))  # force on Dmc, due to friction\n",
    "        FL.append((frame,  HC_wall(N, frame, ball, r, ctau, rhodtwall[zaehler], k0W)[2][i]))  # torque on ball, due to friction\n",
    "    zaehler += 1\n",
    "    \n",
    "# collisions of the two balls\n",
    "FL.append((Dmc1, HC_disc(N, A2, A1, Dmc2, Dmc1, r,ctau, rhodtmax[0], k0)[0]))\n",
    "FL.append((Dmc1, HC_disc(N, A2, A1, Dmc2, Dmc1, r,ctau, rhodtmax[0], k0)[1]))\n",
    "FL.append((A1,   HC_disc(N, A2, A1, Dmc2, Dmc1, r,ctau, rhodtmax[0], k0)[2]))\n",
    "\n",
    "FL.append((Dmc2, HC_disc(N, A1, A2, Dmc1, Dmc2, r,ctau, rhodtmax[1], k0)[0]))\n",
    "FL.append((Dmc2, HC_disc(N, A1, A2, Dmc1, Dmc2, r,ctau, rhodtmax[1], k0)[1]))\n",
    "FL.append((A2,   HC_disc(N, A1, A2, Dmc1, Dmc2, r,ctau, rhodtmax[1], k0)[2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Kane's equations of motion**\n",
    "\n",
    "The formalism to set up Kane's equations.\n",
    "\n",
    "As $Dmc_1$ has a *real* speed and *virtual* speeds for the reaction forces, the reaction forces appear in the *force vector*. As they do no work, I can set them to zero. (At least I think, this is the reason they appear in the force vector) \n",
    "\n",
    "We need to numerically solve the first order differential equations for $(x_1(t), z_1(t), x_2(t), z_2(t))$, so I add the right hand sides at the bottom of the force vector. The mass matrix needs to be enlarged accordingly.\\\n",
    "Presumably I could consider them to be *speed constraints* and enter them accordingly, I did not try."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kd = [me.dot(rot11 - rot12, uv) for uv in A1] + [me.dot(rot21 - rot22, uv) for uv in A2]\n",
    "\n",
    "q = [q11, q12, q13, q21, q22, q23]\n",
    "u = [u11, u12, u13, u21, u22, u23]\n",
    "aux = [auxx, auxy, auxz]\n",
    "\n",
    "# Setting up Kane's equations\n",
    "KM = me.KanesMethod(N, q_ind=q, u_ind=u, kd_eqs=kd, u_auxiliary=aux)\n",
    "(fr, frstar) = KM.kanes_equations(BODY, FL)\n",
    "MM1 = KM.mass_matrix_full\n",
    "force1 = KM.forcing_full.subs({fx: 0., fy: 0., fz: 0.})\n",
    "\n",
    "# add the rhsx1, etc at the bottom of the force vector\n",
    "force = ((sm.Matrix.vstack(force1, sm.Matrix([rhsx1, rhsz1, rhsx2, rhsz2]))).subs(rhs_dict)).subs({i: 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 has {} operations'.format(sum([force[i].count_ops(visual=False) \n",
    "                for i in range(len(force))])), '\\n')\n",
    "\n",
    "# Enlarge MM properly\n",
    "MM2 = sm.Matrix.hstack(MM1, sm.zeros(12, 4))\n",
    "hilfs = sm.Matrix.hstack(sm.zeros(4, 12), sm.eye(4))\n",
    "MM = sm.Matrix.vstack(MM2, hilfs )\n",
    "if term_info == True:\n",
    "    print('MM DS', me.find_dynamicsymbols(MM))\n",
    "    print('MM free symbols', MM.free_symbols)\n",
    "    print('MM has {} operations'.format(sum([MM[i, j].count_ops(visual=False) \n",
    "                for i in range(MM.shape[0]) for j in range(MM.shape[1])])), '\\n')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Reaction Forces**\n",
    "\n",
    "The function to get the reaction forces is determined. This function (of course!) needs the accelerations of points, hence values of $rhs = MM^{-1} * force$ are needed. While one can get rhs symbolically, it probably HUGE. So, I calculate it numerically further down. RHS are just substitutes.\\\n",
    "I only calculate it for one ball, I do not think, anything is gained by doing it for the other ball, too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "RHS = [sm.symbols('rhs' + str(i)) for i in range(force.shape[0])]\n",
    "eingepraegt_dict = {sm.Derivative(i, t): RHS[j] for j, i in enumerate(q + u)}\n",
    "eingepraegt = (((KM.auxiliary_eqs).subs(eingepraegt_dict)).subs(rhs_dict)).subs({i: 0. for i in aux})\n",
    "if term_info == True:\n",
    "    print('eingepraegt DS', me.find_dynamicsymbols(eingepraegt))\n",
    "    print('eingepraegt free symbols', eingepraegt.free_symbols)\n",
    "    print('eingepraegt has {} operations'.format(sum([eingepraegt[i].count_ops(visual=False) \n",
    "                for i in range(len(eingepraegt))])), '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "sympy functions are converted to numpy functions, using *lambdify(....)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pL  = [m, mo, g, r, iXX, iYY, iZZ, amplitude, frequenz] + [ctau, k0, k0W, lW, mu, \n",
    "            muW] + [alpha, beta, gamma] + [rhodtwall] + [rhodtmax]\n",
    "pL1 = [r, lW, amplitude, frequenz] \n",
    "\n",
    "qL = q + u + [x1, z1, x2, z2]\n",
    "F = [fx, fy, fz]\n",
    "\n",
    "MM_lam    = sm.lambdify(qL + pL, MM, cse=True)\n",
    "force_lam = sm.lambdify(qL + pL, force, cse=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "CP1_pos_lam    = sm.lambdify(qL + pL, CP1_pos, cse=True)\n",
    "Dmc1_pos_lam   = sm.lambdify(qL + pL, Dmc1_pos, cse=True)\n",
    "m_Dmc1_pos_lam = sm.lambdify(qL + pL, m_Dmc1_pos, cse=True)\n",
    "\n",
    "CP2_pos_lam    = sm.lambdify(qL + pL, CP2_pos, cse=True)\n",
    "Dmc2_pos_lam   = sm.lambdify(qL + pL, Dmc2_pos, cse=True)\n",
    "m_Dmc2_pos_lam = sm.lambdify(qL + pL, m_Dmc2_pos, cse=True)\n",
    "\n",
    "gesamt1 = gesamt(x, z) # a function cannot be lambdified, so I have to do it like this\n",
    "gesamt_lam = sm.lambdify([x, z] + [amplitude, frequenz], gesamt1, cse=True)\n",
    "\n",
    "pot_lam       = sm.lambdify(qL + pL, pot_energie, cse=True)\n",
    "kin_lam       = sm.lambdify(qL + pL, kin_energie, cse=True)\n",
    "spring_lam    = sm.lambdify(qL + pL, spring_energie, cse=True)\n",
    "\n",
    "distanz_list_lam   = sm.lambdify(qL + pL1, distanz_list, cse=True)\n",
    "abstand_baelle_lam = sm.lambdify([x1, z1, x2, z2] + pL1, abstand_baelle, cse=True)\n",
    "\n",
    "rhodtwall_list_lam = sm.lambdify(qL + pL1, rhodtwall_list, cse=True)\n",
    "rhodtmax_list_lam  = sm.lambdify(qL + pL1, rhodtmax_list, cse=True)\n",
    "\n",
    "eingepraegt_lam = sm.lambdify(F + qL + pL + RHS, eingepraegt, cse=True)\n",
    "\n",
    "r_max_lam = sm.lambdify([x, z] + pL1, [r_max_x, r_max_z], cse=True)\n",
    "\n",
    "print('it took {:.3f} sec to establish Kanes equations'.format(time.time() - start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Numerical Integration**\n",
    "\n",
    "Input parameters / initial values of the coordinates and speeds. Whle it makes sense to name them similarly to the symbols / dynamic symbols used in setting up Kane's equations, **avoid** the **same** name: the symbols will be overwritten, with unintended consequences. The larger $frequenz_1$ and the larger *rumpel*, the samller the minimum osculatng cirlce.\n",
    "- $m_1$: mass of the ball\n",
    "- $m_{o1}$: mass of observer\n",
    "- $r_1$: radius of ball\n",
    "- $amplitude_1, frequenz_1$: define the street\n",
    "- $mu_w$: friction between ball and wall\n",
    "- $c_\\tau$: coefficient needed for the Hunt Crosssley method of impact. $c_\\tau = 1.$ corresponds to Hertz' method.\n",
    "- $\\alpha_1, \\beta_1, \\gamma_1$: location of observer relative to Dmc\n",
    "\n",
    "\n",
    "- $q_{110}, q_{120}, q_{130}, q_{210}, q_{220}, q_{230}$: initial values of the coordinates\n",
    "- $u_{110}, u_{120}, u_{130}, u_{210}, u_{220}, u_{230}$: initial values of the speeds\n",
    "- $x_{10}, z_{10}, x_{20}, z_{20}$: initial location of the contact point\n",
    "\n",
    "- intervall: running time of the integration is in [0., intervall]\n",
    "- schritte: solve_ivp returns the values evenly spaced in [0., intervall]\n",
    "\n",
    "\n",
    "An exception is raised, if the observer was put outside the ball."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "start1 = time.time()\n",
    "\n",
    "print('Arguments')\n",
    "#pL = [m, mo, g, r, iXX, iYY, iZZ, amplitude, frequenz] + [ctau, rhodtwall, k0W, lW] + [alpha, beta, gamma]\n",
    "\n",
    "print('[m, mo, g, r, iXX, iYY, iZZ, amplitude, frequenz, ctau, rhodtwall, k0W, lW, muW' +\n",
    "      'alpha, beta, gamma]', '\\n')\n",
    "\n",
    "#Input data\n",
    "#================================================================\n",
    "max_step   = 0.01\n",
    "\n",
    "m1         = 1.e0\n",
    "mo1        = 1.e-1\n",
    "r1         = 1.\n",
    "lW1        = 5.\n",
    "ctau1      = 0.9\n",
    "mu1        = 0.01\n",
    "muW1       = 0.01\n",
    "amplitude1 = 2.\n",
    "frequenz1  = 0.2\n",
    "\n",
    "alpha1, beta1, gamma1 = 0., 0.8, 0.\n",
    "\n",
    "q110, q120, q130, q210, q220, q230 = 0., 0., 0., 0., 0., 0.\n",
    "\n",
    "intervall = 12.5\n",
    "#================================================================\n",
    "# set initial angular velocities randomly\n",
    "u110, u120, u130, u210, u220, u230 = np.random.choice(np.linspace(-10, 10, 100), size=6)\n",
    "print([f'{j} = {i:.2f} rad/sec  ' for j, i in zip(('u11', 'u12', 'u13', 'u21', 'u22', 'u23'), \n",
    "                                                  (u110, u120, u130, u210, u220, u230))], '\\n')\n",
    "\n",
    "\n",
    "# find feasible starting locations for the two balls\n",
    "zaehler = 0\n",
    "while zaehler <= 100:\n",
    "    zaehler += 1\n",
    "    try:\n",
    "        x10, z10, x20, z20 = np.random.choice(np.linspace(-lW1 + 2.*r1, lW1 - 2.*r1, 100), size=4, replace=False)\n",
    "        if abstand_baelle_lam(x10, z10, x20, z20, r1, lW1, amplitude1, frequenz1) > 3. * r1:\n",
    "            raise Rausspringen()\n",
    "    except:\n",
    "        break\n",
    "\n",
    "if zaehler <= 100:\n",
    "    print(f'it took {zaehler} rounds to get valid initial conditions')\n",
    "    print(f'balls are located at ({x10:.2f}/{z10:.2f}) and at ({x20:.2f}/{z20:.2f}) respectively', '\\n')\n",
    "else:\n",
    "    raise Exception(' no good location for discs found, make lW0 larger.')\n",
    "#================================================================\n",
    "schritte = int(intervall * 300.)\n",
    "if alpha1**2 + beta1**2 + gamma1**2 >= 1.:\n",
    "    raise Exception('center of mass outside of the ball')\n",
    "    \n",
    "rhodtwall1 = [[1. for _ in range(4)] for _ in range(2)]\n",
    "rhodtmax1  = [1., -1.]\n",
    "\n",
    "#Calculate k01 and k0W1\n",
    "nu = 0.28                       # Poisson's ratio, from the internet\n",
    "EY = 2.e3                       # units: N/m^2, Young's modulus from the internet, around 2e11 for steel\n",
    "sigma = (1 - nu**2) / EY\n",
    "k01  = 4. / (3.* (sigma + sigma)) * np.sqrt(r1/2.)\n",
    "k0W1 = 4. / (3.* (sigma + sigma)) * np.sqrt(r1)\n",
    "\n",
    "\n",
    "iXXe = 2./5. * m1 * r1**2\n",
    "\n",
    "pL_vals = [m1, mo1, 9.8, r1, iXXe, iXXe, iXXe, amplitude1, frequenz1, ctau1, k01, k0W1, \n",
    "           lW1, mu1, muW1, alpha1, beta1, gamma1] + [rhodtwall1] + [rhodtmax1]\n",
    "pL1_vals = [r1, lW1, amplitude1, frequenz1]\n",
    "\n",
    "print('pL_vals', pL_vals, '\\n')\n",
    "\n",
    "y0 = [q110, q120, q130, q210, q220, q230] + [u110, u120, u130, u210, u220, u230] + [x10, z10, x20, z20]\n",
    "#================================================================\n",
    "if frequenz1 > 0.:\n",
    "#find the smallest osculating radius, given strasse, amplitude, frequenz\n",
    "    def func1(x, args):\n",
    "# just needed to get the arguments matching for minimuze\n",
    "        return np.abs(r_max_lam(*x, *args)[0])\n",
    "\n",
    "    def func2(x, args):\n",
    "# just needed to get the arguments matching for minimuze\n",
    "        return np.abs(r_max_lam(*x, *args)[1])\n",
    "\n",
    "    x0 = (0.1, 0.1)      # initial guess\n",
    "    minimal1 = minimize(func1, x0, pL1_vals)\n",
    "    minimal2 = minimize(func2, x0, pL1_vals)\n",
    "\n",
    "    minimal = min(minimal1.get('fun'), minimal2.get('fun'))\n",
    "\n",
    "    if pL_vals[3] < minimal:\n",
    "        print('selected radius = {} is less than minimally admissible radius = {:.2f}, hence o.k.'\n",
    "          .format(pL_vals[3], minimal), '\\n')\n",
    "    else:\n",
    "        print('selected radius {} is larger than admissible radius {:.2f}, hence NOT o.k.'\n",
    "          .format(pL_vals[3], minimal), '\\n')\n",
    "        raise Exception('Radius of ball is too large')\n",
    "else:\n",
    "    print('Street is flat')\n",
    "#====================================================================      \n",
    "times = np.linspace(0, intervall, schritte)\n",
    "                        \n",
    "def gradient(t, y, args):\n",
    "    \n",
    "# get the speed of the balls right before it hit a wall\n",
    "    for ball in range(2):\n",
    "        for wand in range(4):\n",
    "            if 0 < r1 -  distanz_list_lam(*y, *pL1_vals)[ball][wand] <= 0.03:\n",
    "                args[-2][ball][wand] = rhodtwall_list_lam(*y, *pL1_vals)[ball][wand]\n",
    "    if 0. < 2.*r1 - abstand_baelle_lam(y[-4], y[-3], y[-2], y[-1], *pL1_vals) <= 0.03:\n",
    "        args[-1] = rhodtmax_list_lam(*y, *pL1_vals)\n",
    "            \n",
    "    sol = np.linalg.solve(MM_lam(*y, *pL_vals), force_lam(*y, *pL_vals))\n",
    "    return np.array(sol).T[0]\n",
    "\n",
    "t_span = (0., intervall)\n",
    "resultat1 = solve_ivp(gradient, t_span, y0, t_eval=times, args=(pL_vals,), method='BDF',\n",
    "         max_step=max_step, atol=1.e-6, rtol=1.e-6)\n",
    "resultat = resultat1.y.T\n",
    "\n",
    "event_dict = {-1: 'Integration failed', 0: 'Integration finished successfully', 1: 'some termination event'}\n",
    "print(event_dict[resultat1.status], ' message was: ', resultat1.message)\n",
    "if resultat1.status == -1:\n",
    "    raise Exception()\n",
    "print('resultat shape', resultat.shape, '\\n')\n",
    "\n",
    "print(\"To numerically integrate an intervall of {} sec the routine cycled {} times and it took {:.5f} sec \".\n",
    "      format(intervall, resultat1.nfev, time.time() - start1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "Plot whatever **generalized coordinates** you want to see"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "bezeichnung = ['q11', 'q12', 'q13', 'q21', 'q22', 'q23', \n",
    "               'u11', 'u12', 'u13', 'u21', 'u22', 'u23',\n",
    "               'x1', 'z1', 'x2', 'z2']\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "for i in (6, 10, 12, 13, 14, 15):\n",
    "    ax.plot(times, resultat[:, i], label = bezeichnung[i])\n",
    "ax.set_title('Generalized coordinates')\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Draw the **hysteresis curves** of the impacts of the balls and of the balls with the walls.\\\n",
    "The red numbers on the graphs indicate the times during which the impact takes place.\\\n",
    "Only walls where impacts have taken place are shown."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "HC_kraft = []\n",
    "HC_displ = []\n",
    "HC_times = []\n",
    "zaehler  = 0\n",
    "i0 = 0\n",
    "for i in range(resultat.shape[0]):\n",
    "    abstand = 2.*r1  - abstand_baelle_lam(*[resultat[i, j] for j in range(-4, 0)], \n",
    "            *pL1_vals)\n",
    "    if abstand < 0.:\n",
    "        i0 = i+1\n",
    "\n",
    "    if abstand >= 0. and i0 == i:\n",
    "        walldt = rhodtmax_list_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL1_vals)[0]\n",
    "    if abstand >= 0.:\n",
    "        rhodt = rhodtmax_list_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL1_vals)[0]\n",
    "        kraft0 = k01 * 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 with a wall.    \n",
    "if len(HC_displ) != 0:\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 (Nm)')\n",
    "    ax.set_title(f'hysteresis curves of successive impacts of the ball_1 with ball_2, 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",
    "\n",
    "\n",
    "for ball in range(2):\n",
    "    for wall in range(4):\n",
    "\n",
    "        HC_kraft = []\n",
    "        HC_displ = []\n",
    "        HC_times = []\n",
    "        zaehler  = 0\n",
    "        i0 = 0\n",
    "        for i in range(resultat.shape[0]):\n",
    "            abstand = r1  - distanz_list_lam(*[resultat[i, j] for j in range(resultat.shape[1])], \n",
    "            *pL1_vals)[ball][wall]\n",
    "            if abstand < 0.:\n",
    "                i0 = i+1\n",
    "\n",
    "            if abstand >= 0. and i0 == i:\n",
    "                walldt = rhodtwall_list_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL1_vals)[ball][wall]\n",
    "            if abstand >= 0.:\n",
    "                rhodt = rhodtwall_list_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL1_vals)[ball][wall]\n",
    "                kraft0 = k0W1 * 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 with a wall.    \n",
    "        if len(HC_displ) != 0:\n",
    "            fig, ax = plt.subplots(figsize=(10,5))\n",
    "            ax.plot(HC_displ, HC_kraft)\n",
    "            ax.set_xlabel('penetration depth (m)')\n",
    "            ax.set_ylabel('contact force (Nm)')\n",
    "            ax.set_title(f'hysteresis curves of successive impacts of the ball_{ball+1} with wall_{wall}, ctau = {ctau1}, muW = {muW1}')\n",
    "    \n",
    "            zeitpunkte = 20\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\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the **energies of the ball**.\\\n",
    "Absent any friction, that is $mu_W = 0.$ and $c_\\tau = 0.$ the total energy should be constant. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#raise Exception()\n",
    "pot_np    = np.empty(schritte)\n",
    "kin_np    = np.empty(schritte)\n",
    "spring_np = np.empty(schritte)\n",
    "total_np = np.empty(schritte)\n",
    "for l in range(schritte):\n",
    "    pot_np[l] = pot_lam(*[resultat[l, j] for j in range(resultat.shape[1])], *pL_vals)\n",
    "    kin_np[l] = kin_lam(*[resultat[l, j] for j in range(resultat.shape[1])], *pL_vals)\n",
    "    spring_np[l] = spring_lam(*[resultat[l, j] for j in range(resultat.shape[1])], *pL_vals)\n",
    "    total_np[l] = pot_np[l] + kin_np[l] + spring_np[l]\n",
    "\n",
    "if ctau1 == 1. and muW1 == 0.:\n",
    "    fehler = ((x111:=max(total_np)) - min(total_np)) / x111 * 100.\n",
    "    print('max deviation from total energy = constant is {:.2e} % of max total energy'.format(fehler))\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "ax.plot(times, kin_np, label = 'kin energy')\n",
    "ax.plot(times, pot_np, label = 'pot energy')\n",
    "ax.plot(times, spring_np, label='spring energy')\n",
    "ax.plot(times, total_np, label = 'total energy')\n",
    "ax.set_title(f'Energy of the system, muW = {muW1}, ctau = {ctau1}')\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Reaction Forces*\n",
    "\n",
    "First the rhs is calculated numerically, then *eingepraegt = 0.* is solved numerically for the reaction forces, then they are plotted.\\\n",
    "As there are many points in time, this may take a long time to calculate for *fsolve*. To reduce the time, I only look at arounf *zeitpunkte* number of points in time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "times2 = []\n",
    "resultat2 = []\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",
    "        times2.append(times[i])\n",
    "        resultat2.append(resultat[i])\n",
    "\n",
    "schritte2 = len(times2)\n",
    "resultat2 = np.array(resultat2)\n",
    "times2 = np.array(times2)\n",
    "print('number of points considered:',len(times2))\n",
    "\n",
    "\n",
    "\n",
    "# RHS calculated numerically, too large to do it symbolically. Needed for reaction forces.\n",
    "RHS1 = np.zeros((schritte2, resultat2.shape[1]))\n",
    "for i in range(schritte2):\n",
    "    RHS1[i, :] = np.linalg.solve(MM_lam(*[resultat2[i, j]for j in range(resultat2.shape[1])], *pL_vals), \n",
    "        force_lam(*[resultat2[i, j] for j in range(resultat2.shape[1])], \n",
    "        *pL_vals)).reshape(resultat2.shape[1])\n",
    "print('RHS1 shape', RHS1.shape)\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(3)\n",
    "\n",
    "kraftx = np.zeros(schritte2)\n",
    "krafty = np.zeros(schritte2)\n",
    "kraftz = np.zeros(schritte2)\n",
    "x0 = tuple((1., 1., 1.))   # initial guess\n",
    "\n",
    "for i in range(schritte2):\n",
    "    for _ in range(2):\n",
    "        y0 = [resultat2[i, j] for j in range(resultat2.shape[1])]\n",
    "        rhs = [RHS1[i, j] for j in range(16)]\n",
    "        args = tuple(y0 + pL_vals + rhs)\n",
    "        A = fsolve(func, x0, args=args).reshape(3)\n",
    "        x0 = tuple(A)      # improved guess. Should speed up convergence of fsolve\n",
    "    kraftx[i] = A[0]\n",
    "    krafty[i] = A[1]\n",
    "    kraftz[i] = A[2]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "plt.plot(times2, kraftx, label='X force')\n",
    "plt.plot(times2, krafty, label='Y force')\n",
    "plt.plot(times2, kraftz, label='Z force')\n",
    "ax.set_title('Reaction Forces on CP1')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This shows the path of the ball, projected on the X / Z plane. The 'height' of the ball above (x, z) = (0., 0.) is indicarted by its color. \n",
    "\n",
    "HTML(...) is needed to show the animation on my iPad, no idea whether needed on other machines. it is SLOW!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "times2 = []\n",
    "resultat2 = []\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",
    "        times2.append(times[i])\n",
    "        resultat2.append(resultat[i])\n",
    "\n",
    "schritte2 = len(times2)\n",
    "resultat2 = np.array(resultat2)\n",
    "times2 = np.array(times2)\n",
    "print('number of points considered:',len(times2))\n",
    "\n",
    "\n",
    "Dmc1x  = np.empty(schritte2)\n",
    "Dmc1y  = np.empty(schritte2)\n",
    "Dmc1y1 = np.empty(schritte2)\n",
    "Dmc1z  = np.ones(schritte2)\n",
    "\n",
    "Po1x   = np.empty(schritte2)\n",
    "Po1y   = np.empty(schritte2)\n",
    "Po1z   = np.empty(schritte2)\n",
    "\n",
    "Dmc2x  = np.empty(schritte2)\n",
    "Dmc2y  = np.empty(schritte2)\n",
    "Dmc2y1 = np.empty(schritte2)\n",
    "Dmc2z  = np.ones(schritte2)\n",
    "\n",
    "Po2x   = np.empty(schritte2)\n",
    "Po2y   = np.empty(schritte2)\n",
    "Po2z   = np.empty(schritte2)\n",
    "\n",
    "for l in range(schritte2):\n",
    "    Dmc1x[l] = Dmc1_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[0]\n",
    "    Dmc1y[l] = Dmc1_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[1]\n",
    "    Dmc1z[l] = Dmc1_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[2]\n",
    "    \n",
    "    Po1x[l]  = m_Dmc1_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[0]\n",
    "    Po1y[l]  = m_Dmc1_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[1]\n",
    "    Po1z[l]  = m_Dmc1_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[2]\n",
    "    \n",
    "    Dmc2x[l] = Dmc2_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[0]\n",
    "    Dmc2y[l] = Dmc2_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[1]\n",
    "    Dmc2z[l] = Dmc2_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[2]\n",
    "    \n",
    "    Po2x[l]  = m_Dmc2_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[0]\n",
    "    Po2y[l]  = m_Dmc2_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[1]\n",
    "    Po2z[l]  = m_Dmc2_pos_lam(*[resultat2[l, j] for j in range(resultat2.shape[1])], *pL_vals)[2]\n",
    "\n",
    "for i, j in enumerate(Dmc1y):\n",
    "    Dmc1y1[i] = int(j * 100.)\n",
    "for i, j in enumerate(Dmc2y):\n",
    "    Dmc2y1[i] = int(j * 100.)\n",
    "\n",
    "ymin = min(min(Dmc1y1), min(Dmc2y1))\n",
    "ymax = max(max(Dmc1y1), max(Dmc2y1))\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",
    "farbe1 = Farbe.to_rgba(Dmc1y1[0])    # color of the starting position\n",
    "farbe2 = Farbe.to_rgba(Dmc2y1[0])    \n",
    "def animate_pendulum(times, x1, y1, z1, y11, ox, oy, oz, x21, y21, z21, y211, ox1, oy1, oz1 ):\n",
    "    \n",
    "    fig, ax = plt.subplots(figsize=(8, 8), subplot_kw={'aspect': 'equal'})\n",
    "    fig.colorbar(Farbe, label='Height of the ball above ground level \\n Factor = 100, the real height are the numbers divided by factor', shrink=0.9, ax=ax)\n",
    "    ax.axis('on')\n",
    "\n",
    "    ax.set_xlabel(' X Axis', fontsize=18)\n",
    "    ax.set_ylabel('Z Axis', fontsize=18)\n",
    "    \n",
    "    ax.set(xlim=(-lW1 - 1., lW1 + 1.), ylim=( -lW1 - 1., lW1 + 1.))\n",
    "    ax.plot([-lW1, lW1], [-lW1, -lW1], 'bo', linewidth=2, linestyle='-', markersize=0)\n",
    "    ax.text(0., -lW1-1., 'wall 0')\n",
    "    ax.plot([lW1, lW1], [-lW1, lW1], 'ro', linewidth=2, linestyle='-', markersize=0)\n",
    "    ax.text(lW1 + 0.5, 0., 'wall 1')\n",
    "    ax.plot([lW1, -lW1], [lW1, lW1], 'go', linewidth=2, linestyle='-', markersize=0)\n",
    "    ax.text(0., lW1 + .5, 'wall 2')\n",
    "    ax.plot([-lW1, -lW1], [lW1, -lW1], 'yo', linewidth=2, linestyle='-', markersize=0)\n",
    "    ax.text(-lW1-3., 0., 'wall 3')\n",
    "    \n",
    "\n",
    "    line1,  = ax.plot([], [], 'o', markersize=8, color=farbe1)       # starting point Dmc1\n",
    "    line1a, = ax.plot([], [], 'o', markersize=8, color=farbe2)       # starting point Dmc1\n",
    "    line2,  = ax.plot([], [], 'o', markersize=90/2.77*r1 /lW1 * 10.)  # moving ball1\n",
    "    line2a, = ax.plot([], [], 'o', markersize=90/2.77*r1 /lW1 * 10.)  # moving ball1\n",
    "    line3,  = ax.plot([], [], color='blue', linewidth=0.25)           # to trace the movement of Dmc1\n",
    "    line3a, = ax.plot([], [], color='red', linewidth=0.25)           # to trace the movement of Dmc2\n",
    "    line4,  = ax.plot([], [], 'o', markersize=5, color='black')       # observer on ball 1\n",
    "    line4a, = ax.plot([], [], 'o', markersize=5, color='black')       # observer on ball 2\n",
    "\n",
    "    \n",
    "    def animate(i):\n",
    "        farbe3 = Farbe.to_rgba(y11[i])                                # color of the actual point at time i\n",
    "        farbe4 = Farbe.to_rgba(y211[i])                                # color of the actual point at time i\n",
    "        \n",
    "        ax.set_title('running time {:.1f} sec, with rumpel = {}'.format(times[i], rumpel), fontsize=15)\n",
    "        \n",
    "        line1.set_data(x1[0], z1[0])\n",
    "        line1.set_color(farbe1)\n",
    "        line1a.set_data(x21[0], z21[0])\n",
    "        line1a.set_color(farbe2)\n",
    "        \n",
    "        line2.set_data(x1[i], z1[i])\n",
    "        line2.set_color(farbe3)\n",
    "        line2a.set_data(x21[i], z21[i])\n",
    "        line2a.set_color(farbe4)\n",
    "        \n",
    "        line3.set_data(x1[: i], z1[: i])\n",
    "        line3a.set_data(x21[: i], z21[: i])\n",
    "        \n",
    "        line4.set_data(ox[i], oz[i])\n",
    "        if y11[i] <= oy[i]:\n",
    "            line4.set_color('black')\n",
    "        else:\n",
    "            line4.set_color('grey')\n",
    "            \n",
    "        line4a.set_data(ox1[i], oz1[i])\n",
    "        if y211[i] <= oy1[i]:\n",
    "            line4a.set_color('black')\n",
    "        else:\n",
    "            line4a.set_color('grey')\n",
    "            \n",
    "            \n",
    "        return line1, line1a, line2, line2a, line3, line3a, line4, line4a\n",
    "\n",
    "    anim = animation.FuncAnimation(fig, animate, frames=len(times2),\n",
    "                                   interval=1000*max(times2) / len(times2),\n",
    "                                   blit=True)\n",
    "    plt.close(fig)\n",
    "    return anim\n",
    "\n",
    "anim = animate_pendulum(times2, Dmc1x, Dmc1y, Dmc1z, Dmc1y1, Po1x, Po1y, Po1z, \n",
    "                                Dmc2x, Dmc2y, Dmc2z, Dmc2y1, Po2x, Po2y, Po2z )\n",
    "start2 = time.time()\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": [
    "Plot the approximate shape of the street. I does not show the 'unevenness'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# plot approximate shape of the street\n",
    "xs = np.linspace(-lW1, lW1, 500)\n",
    "ys = np.linspace(-lW1, lW1, 500)\n",
    "X, Y = np.meshgrid(xs, ys)\n",
    "Z = gesamt_lam(X, Y, amplitude1, frequenz1)\n",
    "\n",
    "fig = plt.figure(figsize=(10,10))\n",
    "ax = fig.add_subplot(projection='3d')\n",
    "\n",
    "strasse = ax.plot_surface(X, Y, Z, cmap='plasma', linewidth=0., antialiased=True)\n",
    "ax.set_xlabel('X - Axis', fontsize = 15)\n",
    "ax.set_ylabel('Z - Axis', fontsize = 15)\n",
    "ax.set_title('Approximate shape of the street', fontsize=15)\n",
    "ax.set_zlabel('Y - Axis', fontsize=15)\n",
    "fig.colorbar(strasse, shrink=0.5, label='Height above the ground', aspect=15);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.16"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}