h-c-colliding-discs-friction-circle.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy.physics.mechanics as me\n",
    "import sympy as sm\n",
    "from scipy.integrate import solve_ivp\n",
    "from scipy.optimize import fsolve\n",
    "import numpy as np\n",
    "from itertools import permutations\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from matplotlib import animation\n",
    "from IPython.display import HTML\n",
    "import matplotlib as mp\n",
    "import time\n",
    "import matplotlib\n",
    "matplotlib.rcParams['animation.embed_limit'] = 2**126"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is needed to exit a loop, when a feasible initial location of the discs within the limitations of the wall was found."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Rausspringen(Exception):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**n discs**, named $Dmc_0....Dmc_{n-1}$ with radius $r_0$ and mass $m_0$ are sliding on the frictionless horizontal X/Z plane.\\\n",
    "Their space is limited by a circular wall of radius $R_W$ and center at the origin. \n",
    "\n",
    "The collision force is always on the line $\\overline{Dmc_i, Dmc_j}$, $0 \\le i, j \\le n-1$, $i \\neq j$, or on a line perpendicular to the tangent of the circular wall at the collision point. Of course, this line will go through Dmc[i]\\\n",
    "There is speed dependent friction between discs, with coefficient $m_u$ and between wall and discs, with coefficient $m_{uW}$\n",
    "\n",
    "An observer, a particle of mass $m_o$ may be attached anywhere within each disc.\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",
    "**Notes**\n",
    "\n",
    "1.\\\n",
    "$c_\\tau = 1.$ gives **Hertz's** solution to the impact problem, also described in the article.\n",
    "\n",
    "2.\\\n",
    "From the discs' point of view, the wall is concave. I model this by taking $R_2 = -R_W$.\\\n",
    "As $|R_1| < |R_2|$ this should give no issues. **I do not know, whether this approach is theoretically correct.**\n",
    "\n",
    "\n",
    "**Variables**\n",
    "\n",
    "- $n$: number of discs\n",
    "- $q_0...q_{n-1}$: generalized coordinates for the discs\n",
    "- $u_0...u_{n-1}$: the angular speeds\n",
    "- $x_i, z_i$: the coordinates, in the inertial frame $N$, of the center of the i-th disc\n",
    "- $ux_i, uz_i$ their speeds\n",
    "- $N$: frame of inertia\n",
    "- $P_0$: point fixed in $N$\n",
    "- $A_i$: body fixed frame of the i-th disc\n",
    "\n",
    "- $m_i$: mass of the i-th disc\n",
    "- $Dmc_i$: center of the i-th disc\n",
    "- $Po_i$: observer (particle) on i-th disc\n",
    "- $\\alpha_i$: distance of observer on i-th disc\n",
    "- $m_W$: mass of the wall.\n",
    "- $r_0, R_W, k_0$: radius of the discs, length of the wall, modulus of elasticity of the pendulum body\n",
    "- $k_{0W}$: modulus of elasticity of the collision disc with a wall\n",
    "- $i_{YY_i}$: moment of ineratia of the i-th disc\n",
    "\n",
    "\n",
    "- $m_u$: coefficient of friction between discs / between disc and wall.\n",
    "- $c_\\tau$: the experimental constant needed for Hunt-Crossley\n",
    "- $rhodt_{max}$: the collission speed, to be determined during integration, needed for Hunt_Crossley\n",
    "- $rhodt_{wall}$: the collision speeds when $disc_i$ hits a wall\n",
    "- $CP_i$: contact point of $disc_i$ with the wall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "#==========================================\n",
    "n = 3       # n > 1\n",
    "#==========================================\n",
    "\n",
    "if isinstance(n, int) == False or n < 2:\n",
    "    raise Exception('n must be an integer larger than 1')\n",
    "    \n",
    "q_list  = me.dynamicsymbols(f'q:{n}')\n",
    "u_list  = me.dynamicsymbols(f'u:{n}')\n",
    "x_list  = me.dynamicsymbols(f'x:{n}')\n",
    "z_list  = me.dynamicsymbols(f'z:{n}')\n",
    "ux_list = me.dynamicsymbols(f'ux:{n}')\n",
    "uz_list = me.dynamicsymbols(f'uz:{n}')\n",
    "\n",
    "CPx_list = list(sm.symbols(f'CPx:{n}'))\n",
    "CPz_list = list(sm.symbols(f'CPz:{n}'))\n",
    "CP_list  = list(sm.symbols(f'CP:{n}', cls=me.Point))\n",
    "\n",
    "\n",
    "rhodtmax = []\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    rhodtmax.append(sm.symbols('rhodtmax' + str(i) + str(j)))\n",
    "\n",
    "rhodtwall  = list(sm.symbols(f'rhodtwall:{n}'))\n",
    "\n",
    "m0, mo, r0, RW, k0, k0W, iYY, ctau, mu, muW = sm.symbols('m0, mo, r0, RW, k0, k0W, iYY, ctau, mu, muW')\n",
    "\n",
    "N  = me.ReferenceFrame('N')\n",
    "P0 = me.Point('P0')\n",
    "P0.set_vel(N, 0)\n",
    "\n",
    "t = me.dynamicsymbols._t\n",
    "\n",
    "\n",
    "\n",
    "A_list     = sm.symbols(f'A:{n}', cls=me.ReferenceFrame)\n",
    "Dmc_list   = sm.symbols(f'Dmc:{n}', cls=me.Point)\n",
    "Po_list    = sm.symbols(f'Po:{n}', cls=me.Point)\n",
    "CP_list    = list(sm.symbols(f'CP:{n}', cls=me.Point))\n",
    "alpha_list = sm.symbols(f'alpha:{n}')\n",
    "\n",
    "Body1 = []\n",
    "Body2 = []\n",
    "for i in range(n):\n",
    "    A_list[i].orient_axis(N, q_list[i], N.y)\n",
    "    A_list[i].set_ang_vel(N, u_list[i] * N.y)\n",
    "    \n",
    "    Dmc_list[i].set_pos(P0, x_list[i]*N.x + z_list[i]*N.z)\n",
    "    Dmc_list[i].set_vel(N, ux_list[i]*N.x + uz_list[i]*N.z)\n",
    "    \n",
    "    Po_list[i].set_pos(Dmc_list[i], r0 * alpha_list[i] * A_list[i].x)\n",
    "    Po_list[i].v2pt_theory(Dmc_list[i], N, A_list[i])\n",
    "    \n",
    "    CP_list[i].set_pos(P0, CPx_list[i]*N.x + CPz_list[i]*N.z)\n",
    "    \n",
    "    I = me.inertia(A_list[i], 0, iYY, 0)                                              \n",
    "    body = me.RigidBody('body' + str(i), Dmc_list[i], A_list[i], m0, (I, Dmc_list[i]))\n",
    "    teil = me.Particle('teil' + str(i), Po_list[i], mo)\n",
    "    Body1.append(body)\n",
    "    Body2.append(teil)\n",
    "BODY = Body1 + Body2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Colliding force between discs**\n",
    "\n",
    "This function returns the force, as given by Hunt-Crossley, which $P_1$ excerts on $P_2$.\\\n",
    "The relevant speed is only the speed component in $\\overline{P_1 P_2}$ direction.\\\n",
    "$sm.Heaviside(2. \\cdot r - abstand, 0)$ ensures, that there is a force only during the collision. \n",
    "\n",
    "I get $\\dot\\rho^{(-)}$ during integration. The way I defined it, it should always be negative, just like $\\dot\\rho$ in the first phase of the penetration. For whatever reason, this is not always so.\\\n",
    "Using $\\dfrac{-\\dot\\rho}{|\\dot\\rho^{(-)}|}$ gives the right results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def HC_disc(N, P1, P2, r, ctau, rhodtmax, k0):\n",
    "    '''\n",
    "Calculates the contact force exerted by P1 on P2, according to the Hunt-Crossley formula given above.\n",
    "I assume, that the contact force always acts along the line P1/P2. I think, this is a fair assymption if\n",
    "the colliding balls are homogenious.\n",
    "\n",
    "The variables in the list are\n",
    "\n",
    "- N is an me.ReferenceFrame, the inertial frame\n",
    "- P1, P2 are me.Point objects. They are the centers of two me.RigidBody objects, here assumed to be two\n",
    "  ball each of radius r\n",
    "- radius of the ball\n",
    "- ctau, the experimental constant needed\n",
    "- rhodtmax, the relative speeds of P1 to P2, right at the impact time, measured in N. \n",
    "  This has to be calculated numerically during the integration, I found no ther way.\n",
    "- k0 is the force constant\n",
    "\n",
    "\n",
    "    '''\n",
    "    vektorP1P2 = P2.pos_from(P1)\n",
    "    abstand = vektorP1P2.magnitude() \n",
    "    richtung = vektorP1P2.normalize()        \n",
    "    \n",
    "    rho = (2.*r - abstand)/2.      # penetration. Positive if the two balls are in collision\n",
    "    geschw = vektorP1P2.diff(t, N)\n",
    "    rhodt = me.dot(geschw, richtung)\n",
    "    rho = sm.Max(rho, sm.S(0))   # if rho < 0., rho**(3/2) will give problems\n",
    "    \n",
    "    kraft = k0 * rho**(3/2) * (1. + 3./2. * (1 - ctau) * (-rhodt/sm.Abs(rhodtmax))\n",
    "        ) * sm.Heaviside(2. * r - abstand, 0.) * richtung\n",
    "    \n",
    "    return kraft   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Collision force between disc and the wall**\\\n",
    "As far as the penetration depth is concerned, this is like two discs colliding: the penetration depth of the wall into the disc is the same as the disc into the wall. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def HC_wall(N, P1, CP, r, ctau, rhodtwall, k0W):\n",
    "    \n",
    "    '''\n",
    "Calculates the contact force exerted by the wall on P1, according to the Hunt-Crossley formula \n",
    "given above. I assume, that the contact force always acts in the direction normal to the tangent of the\n",
    "wall at the collision point CP\n",
    "(Of course, a disc can maximally be in simultaneous contact with two walls)\n",
    "\n",
    "The variables in the list are\n",
    "\n",
    "- N is an me.ReferenceFrame, the inertial frame\n",
    "- P1, is a me.Point object, the center of the disc\n",
    "- CP is the point where the disc collides with the wall.\n",
    "- radius of the disc\n",
    "- ctau, the experimental constant needed\n",
    "- rhodtmax, a list containing the speed right before P1 hits wall_i \n",
    "  This has to be calculated numerically during the integration, I found no ther way.\n",
    "- k0W is the force constant. This is different from k0 as the wall is like a disc with R = -RW.\n",
    "\n",
    "    '''\n",
    "\n",
    "    vektor   = P1.pos_from(CP)   \n",
    "    abstand  = vektor.magnitude()\n",
    "    richtung = vektor.normalize()\n",
    "    \n",
    "    rho = (r - abstand)       # penetration. Positive if the disc is in collision with the wall\n",
    "    rhodt = me.dot(P1.vel(N), richtung)\n",
    "    rho = sm.Max(rho, sm.S(0))   # if rho < 0., rho**(3/2) will give problems\n",
    "    \n",
    "    kraft0 = k0W * rho**(3/2) * (1. + 3./2. * (1 - ctau) * (-rhodt/sm.Abs(rhodtwall))        \n",
    "                                ) * sm.Heaviside(r - abstand, 0.) * richtung\n",
    "    \n",
    "    return kraft0   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Friction when a disc hits another disc*\n",
    "\n",
    "This website:\n",
    "https://math.stackexchange.com/questions/2195047/solve-the-vector-cross-product-equation\n",
    "\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 P1, without any further geometric considerations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def friction_disc(N, A1, A2, P1, P2, r, ctau, rhodtmax, k0, mu):\n",
    "    '''\n",
    "when two discs collide, their surface speed in general will be different.\n",
    "There is a force caused by friction. Here I calculate the force acting on the contact point CP1 of P1 in\n",
    "direction of the tangent on disc1, whose center is P1, proportional to\n",
    "* v(CP2) - v(CP1)\n",
    "* the magnitude of the collision force which P2 excerts on P1\n",
    "* the coefficient of friction mu\n",
    "The force on CP is equivalent to a force on Dmc and a torque on A\n",
    "    '''\n",
    "    CPa, CPb = sm.symbols('CPa, CPb', cls=me.Point)\n",
    "    \n",
    "    abstand = P2.pos_from(P1)\n",
    "    CPa.set_pos(P1, abstand / 2.)\n",
    "    CPa.v2pt_theory(P1, N, A1)\n",
    "    \n",
    "    CPb.set_pos(P2, -abstand / 2.)\n",
    "    CPb.v2pt_theory(P2, N, A2)\n",
    "    delta_v = CPb.vel(N) - CPa.vel(N)\n",
    "\n",
    "    force_coll = HC_disc(N, P2, P1, r, ctau, rhodtmax, k0).magnitude()\n",
    "    force_fric = (force_coll * mu * delta_v)\n",
    "\n",
    "    torque = 0.5 * abstand.cross(force_fric) * sm.Heaviside(2. * r - abstand.magnitude())\n",
    "    \n",
    "    kraft  = 1./(0.25 * me.dot(abstand, abstand)) * torque.cross(abstand) * sm.Heaviside( 2. * \n",
    "            r - abstand.magnitude())\n",
    "\n",
    "    return [kraft, torque]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Friction when a disc hits the wall*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def friction_wall(N, A1, P1, CP, r, ctau, rhodtwall, k0W, mu):\n",
    "    '''\n",
    "When a disc collides with a wall, there will be a force due to friction.\n",
    "I calculate the force due to friction on the contact point CP1.\n",
    "This force is parallel to the wall, and proportional to\n",
    "* the speed of CP1\n",
    "* the magnitude of the collision force\n",
    "* the coefficient of friction mu\n",
    "    '''\n",
    "    \n",
    "    vektor = P1.pos_from(CP)\n",
    "    abstand = vektor.normalize() \n",
    "    \n",
    "    CPa = me.Point('CPa')\n",
    "    CPa.set_pos(P1, -abstand)\n",
    "    CPa.v2pt_theory(P1, N, A1)\n",
    "    force_coll = HC_wall(N, P1, CP, r, ctau, rhodtwall, k0W).magnitude()\n",
    "    force_fric = -force_coll * mu * CPa.vel(N)\n",
    "    \n",
    "    torque0 = CPa.pos_from(P1).cross(force_fric) * sm.Heaviside(r - vektor.magnitude())\n",
    "    kraft0  = 1./me.dot(abstand, abstand) * torque0.cross(-abstand) * sm.Heaviside(r - vektor.magnitude())\n",
    "\n",
    "    \n",
    "    return [kraft0, torque0]   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**force from the collision of discs with discs and discs with the wall**\n",
    "\n",
    "With *permutations(range(n), r=2)* I get all possible forces of $disc_j$ on $disc_i$, $0 \\le i, j \\le n-1$, $i \\neq j$\\\n",
    "This is not the most efficient way to do it, as for example $F^{disc_i \\space \\backslash \\space disc_j} = -F^{disc_j \\space \\backslash \\space disc_i}$, but this way seems easier to keep it all 'straight'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "FL = []\n",
    "\n",
    "#------------------------------------------------------------------------------------------------------\n",
    "# Hunt - Crossley type collision forces between discs and between disc and wall\n",
    "zaehler = -1\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    zaehler += 1\n",
    "    FL.append((Dmc_list[i], HC_disc(N, Dmc_list[j], Dmc_list[i], r0, ctau, rhodtmax[zaehler], k0)))\n",
    "\n",
    "for i in range(n):\n",
    "    FL.append((Dmc_list[i], HC_wall(N, Dmc_list[i], CP_list[i], r0, ctau, rhodtwall[i], k0W )))\n",
    "\n",
    "#------------------------------------------------------------------------------------------------------  \n",
    "# Friction forces between discs\n",
    "zaehler = -1\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    zaehler += 1\n",
    "    FL.append((Dmc_list[i], friction_disc(N, A_list[i], A_list[j], Dmc_list[i], Dmc_list[j], r0, ctau,\n",
    "        rhodtmax[zaehler], k0, mu)[0]))\n",
    "    FL.append((A_list[i], friction_disc(N, A_list[i], A_list[j], Dmc_list[i], Dmc_list[j], r0, ctau,\n",
    "        rhodtmax[zaehler], k0, mu)[1]))\n",
    "    \n",
    "#------------------------------------------------------------------------------------------------------\n",
    "# Friction forces between disc and wall\n",
    "for i in range(n):\n",
    "    FL.append((Dmc_list[i], friction_wall(N, A_list[i], Dmc_list[i], CP_list[i], r0, ctau, rhodtwall[i],\n",
    "            k0W, muW)[0]))\n",
    "    FL.append((A_list[i], friction_wall(N, A_list[i], Dmc_list[i], CP_list[i], r0, ctau, rhodtwall[i],\n",
    "            k0W, muW)[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Kane's equations** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "kd = [i - sm.Derivative(j, t) for i, j in zip(u_list + ux_list + uz_list, q_list + x_list + z_list)]\n",
    "\n",
    "q_ind = x_list + z_list + q_list\n",
    "u_ind = ux_list + uz_list + u_list\n",
    "\n",
    "KM = me.KanesMethod(N, q_ind=q_ind, u_ind=u_ind, kd_eqs=kd)\n",
    "(fr, frstar) = KM.kanes_equations(BODY, FL)\n",
    "\n",
    "MM = KM.mass_matrix_full\n",
    "print('MM DS', me.find_dynamicsymbols(MM))\n",
    "print('MM 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\n",
    "print('force DS', me.find_dynamicsymbols(force), '\\n')\n",
    "print('force free symbols', force.free_symbols, '\\n')\n",
    "print('force contains {} operations'.format(sum([force[i].count_ops(visual=False) for i in range(force.shape[0])])), '\\n')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here various **functions** are defined, which are needed later.\n",
    "\n",
    "- *rhomax_list*: It is used during integration to calculate the speeds just before impact between $disc_j$ and $disc_i$, $0 \\le i, j \\le n-1$, $i \\neq j$\n",
    "- *rhowall_list*: It is used during integration to calculate the speeds just before impact between $disc_i$ and the walls.\n",
    "- *Dmc_pos*: Holds the locations of the centers of the discs. Only for plotting.\n",
    "- *Po_pos*: Holds the locations of each pbserver. Dto.\n",
    "- *Dmc_distanz*: Holds the distance between $disc_j$ and $disc_i$, $0 \\le i, j \\le n-1$, $i \\neq j$. Needed during integration\n",
    "- *kinetic_energie*: calculates the kinetic energy of the bodies and particles.\n",
    "- *spring_energie*: calculates the spring energy of the colliding bodies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "derivative_dict = {sm.Derivative(i, t): j for i, j in zip(x_list + z_list, ux_list + uz_list)}\n",
    "\n",
    "rhomax_list = []\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    vektor = Dmc_list[i].pos_from(Dmc_list[j])\n",
    "    richtung = vektor.normalize()\n",
    "    geschw = vektor.diff(t, N).subs(derivative_dict)\n",
    "    rhodt = me.dot(geschw, richtung)\n",
    "    rhomax_list.append(rhodt)\n",
    "print('rhomax_list DS:', set.union(*[me.find_dynamicsymbols(rhomax_list[k]) \n",
    "        for k in range(len(rhomax_list))]), '\\n')\n",
    "print('rhomax_list free symbols:', set.union(*[rhomax_list[k].free_symbols \n",
    "        for k in range(len(rhomax_list))]), '\\n')\n",
    "\n",
    "\n",
    "rhowall_list = []\n",
    "for i in range(n):\n",
    "    richtung = CP_list[i].pos_from(Dmc_list[i]).normalize()\n",
    "    rhodt = me.dot(Dmc_list[i].vel(N), richtung)\n",
    "    rhowall_list.append(rhodt)\n",
    "print('rhowall_list DS:', set.union(*[me.find_dynamicsymbols(rhowall_list[k]) \n",
    "        for k in range(len(rhowall_list))], '\\n'))\n",
    "print('rhowall_list free symbols:', set.union(*[rhowall_list[k].free_symbols \n",
    "        for k in range(len(rhowall_list))]))\n",
    "\n",
    "Po_pos  = [[me.dot(Po_list[i].pos_from(P0), uv) for uv in (N.x, N.z)] for i in range(n)]\n",
    "\n",
    "Dmc_distanz = [Dmc_list[i].pos_from(Dmc_list[j]).magnitude() for i, j in permutations(range(n), r=2)]\n",
    "\n",
    "kin_energie = sum([body.kinetic_energy(N) for body in BODY])\n",
    "\n",
    "spring_energie = 0.\n",
    "# 1. collisions of discs\n",
    "for i in range(n-1):\n",
    "    for j in range(i+1, n):\n",
    "        distanz = Dmc_list[i].pos_from(Dmc_list[j]).magnitude()\n",
    "        rho = sm.Max((2.*r0 - distanz)/2., sm.S(0.))\n",
    "        rho = rho**(5/2)\n",
    "        spring_energie += 2. * k0 * 2./5. * rho * sm.Heaviside(2.*r0 - distanz, 0.)        \n",
    "# 2. Collision of discs with the wall   \n",
    "for i in range(n):\n",
    "    abstand = Dmc_list[i].pos_from(CP_list[i]).magnitude()     # this is simply the Z coordinate of P1\n",
    "    rho = (r0 - abstand)       # penetration. Positive if the disc is in collision with the wall\n",
    "    rho = sm.Max(rho, sm.S(0))   # if rho < 0., rho**(3/2) will give problems\n",
    "    rho = rho**(5/2)\n",
    "    spring_energie += k0W * 2./5. * rho * sm.Heaviside(r0 - abstand, 0.)\n",
    "\n",
    "# Distance from Dmc[i] to collission point[i]\n",
    "#abstandDmcCP = []\n",
    "#for i in range(n):\n",
    "#    abstandDmcCP.append(CP_list[i].pos_from(Dmc_list[i]).magnitude())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Lambdification*\\\n",
    "The sympy functions are converted to numpy functions for numerical calculations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "qL  = q_ind + u_ind\n",
    "pL  = [m0, mo, r0, RW, k0, k0W, iYY, ctau, mu, muW] + list(alpha_list) + rhodtmax + rhodtwall\n",
    "pL1 = [m0, mo, r0, RW, k0, k0W, iYY, ctau, mu, muW]\n",
    "pL2 = [m0, mo, r0, RW, k0, k0W, iYY, ctau, mu, muW] + list(alpha_list)\n",
    "print(CPx_list, CPz_list)\n",
    "MM_lam    = sm.lambdify(qL + pL, MM, cse=True)\n",
    "force_lam = sm.lambdify(qL + pL + CPx_list + CPz_list, force, cse=True) \n",
    "\n",
    "rhomax_list_lam  = sm.lambdify(qL + pL1, rhomax_list, cse=True)\n",
    "rhowall_list_lam = sm.lambdify(qL + pL1 + CPx_list + CPz_list, rhowall_list, cse=True)\n",
    "Dmc_distanz_lam  = sm.lambdify(x_list + z_list, Dmc_distanz, cse=True)\n",
    "\n",
    "#abstandDmcCP_lam = []\n",
    "#for i in range(n):    \n",
    "#    abstand   = sm.lambdify(qL + pL + CPx_list + CPz_list, abstandDmcCP[i], cse=True)\n",
    "#    abstandDmcCP_lam.append(abstand)\n",
    "    \n",
    "Po_pos_lam  = sm.lambdify(qL + pL2, Po_pos, cse=True)\n",
    "\n",
    "kin_lam    = sm.lambdify(qL + pL2, kin_energie, cse=True)\n",
    "spring_lam = sm.lambdify(qL + pL2 + CPx_list + CPz_list, spring_energie, cse=True)\n",
    "\n",
    "print(f'it took {time.time()-start:.3f} sec to establish Kanes equation for {n} bodies')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Determine the contact points of the discs with the wall**\n",
    "\n",
    "Given a triangle with sides $a, b, c$ and opposing angles $\\alpha, \\beta, \\gamma$ the *law of the sinus* is:\n",
    "$\\dfrac{a}{sin(\\alpha)} = \\dfrac{b}{sin(\\beta)} = \\dfrac{c}{sin(\\gamma)} \\space \\space$ from the internet.\\\n",
    "The *law of cosine* is $c^2 = a^2+b^2 - 2 \\cdot a \\cdot b \\cdot sin(\\gamma) \\space \\space$ also from the internet\\\n",
    "In this case:\n",
    "- $b = | {}^{Dmc}r^{P0} |$\n",
    "- $a = R_W - r_0$\n",
    "- $\\alpha = angle(b, c)$\n",
    "\n",
    "The direction of $b$ is known, see above and call it $\\hat b$, and the direction of $c$ is: $\\hat c = \\hat v_{Dmc}$\\\n",
    "$cos(\\alpha) = \\hat c \\cdot \\hat b$\\\n",
    "$sin(\\alpha) = \\sqrt{1-cos(\\alpha)^2}$\\\n",
    "From the sinus law I get: $sin(\\beta) = \\dfrac{b}{a} \\cdot sin(\\alpha)$\\\n",
    "Hence $\\gamma = \\pi - \\alpha - \\beta$\\\n",
    "So, the sought after $c$ is: $c = \\sqrt{a^2 + b^2 - 2 \\cdot a \\cdot b \\cdot cos(\\gamma)}$ \n",
    "\n",
    "Now the location of $Dmc$ at contact is known, call it $\\overline{Dmc}$\\\n",
    "To get the contact point $CP$, all I have to do is: $CP = \\overline{Dmc} + r_0 \\cdot \\hat a$\n",
    "\n",
    "The aproach above works as long as the discs do not touch the wall.\\\n",
    "However, when the disc is at the wall, it is easy to find the contact point:\\\n",
    "$CP = R_W \\cdot {}^{Dmc}\\hat r^{P_0}$\\\n",
    "I check during the numerical integration whether a disc is at the wall."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aCP = list(sm.symbols(f'aCP:{n}'))\n",
    "bCP = list(sm.symbols(f'aCP:{n}'))\n",
    "cCP = list(sm.symbols(f'cCP:{n}'))\n",
    "\n",
    "alphaCP = list(sm.symbols(f'alphaCP:{n}'))\n",
    "betaCP  = list(sm.symbols(f'betaCP:{n}'))\n",
    "gammaCP = list(sm.symbols(f'gammaCP:{n}'))\n",
    "\n",
    "alphacos = list(sm.symbols(f'alphacos:{n}'))\n",
    "alphasin  = list(sm.symbols(f'alphasin:{n}'))\n",
    "betasin  = list(sm.symbols(f'betasin:{n}'))\n",
    "\n",
    "DmcCP   = list(sm.symbols(f'DmcCP:{n}', cls=me.Point))\n",
    "CPh     = list(sm.symbols(f'CPh:{n}', cls=me.Point))\n",
    "CPhh    = list(sm.symbols(f'CPhh:{n}', cls=me.Point))\n",
    "\n",
    "abstandDcmCP  = ['x' for _ in range(n)]\n",
    "\n",
    "CP_ort   = [['x', 'x'] for _ in range(n)]\n",
    "CPh_ort  = [['x', 'x'] for _ in range(n)]\n",
    "\n",
    "subs_dict1 = {sm.Derivative(i, t): j for i, j in zip(x_list + z_list, ux_list + uz_list)}\n",
    "for i in range(n):\n",
    "    bCP[i] = P0.pos_from(Dmc_list[i]).magnitude()\n",
    "    aCP[i] = RW - r0\n",
    "    \n",
    "    vDmch = ((Dmc_list[i].pos_from(P0).diff(t, N)).normalize()).subs(subs_dict1)\n",
    "    \n",
    "    alphacos[i] = me.dot(P0.pos_from(Dmc_list[i]).normalize(), vDmch)\n",
    "    alphasin[i] = sm.sqrt(sm.Abs(1. - alphacos[i]**2))\n",
    "    betasin[i]  = sm.Piecewise((bCP[i] / aCP[i] * alphasin[i], bCP[i] <= aCP[i]), (0., True))\n",
    "    alphaCP[i]  = sm.Piecewise((sm.asin(alphasin[i]), alphacos[i] >= 0. ),\n",
    "                               (-sm.asin(alphasin[i]) + sm.pi, True) )\n",
    "    betaCP[i]   = sm.asin(betasin[i])\n",
    "    gammaCP[i]  = sm.pi - alphaCP[i] - betaCP[i]\n",
    "    \n",
    "    cCP[i] = sm.sqrt(sm.Abs(aCP[i]**2 + bCP[i]**2 - 2.*aCP[i]*bCP[i]*sm.cos(gammaCP[i])))\n",
    "\n",
    "    DmcCP[i].set_pos(Dmc_list[i], cCP[i] * vDmch)\n",
    "    CPh[i].set_pos(DmcCP[i], r0 * DmcCP[i].pos_from(P0).normalize()) \n",
    "    CPhh[i].set_pos(P0, RW * Dmc_list[i].pos_from(P0).normalize())\n",
    "\n",
    "    abstandDcmCP[i] = Dmc_list[i].pos_from(CPhh[i]).magnitude()\n",
    "    CP_ort[i]   = [me.dot(CPh[i].pos_from(P0), uv) for uv in (N.x, N.z)]\n",
    "    CPh_ort[i]  = [me.dot(CPhh[i].pos_from(P0), uv) for uv in (N.x, N.z)]\n",
    "    \n",
    "abstandDmcCP_lam   = sm.lambdify(qL + pL1, abstandDcmCP, cse=True)\n",
    "CP_ort_lam    = sm.lambdify(qL + pL1, CP_ort, cse=True)\n",
    "CPh_ort_lam   = sm.lambdify(qL + pL1, CPh_ort, cse=True)\n",
    "bCP_lam       = sm.lambdify(qL + pL1, bCP, cse=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Set initial conditions and parameters**\n",
    "\n",
    "1.\\\n",
    "The discs are randomly placed within the walls, such that they have a distance of at least $r_0$ from the walls, and they have a distance of at least $r_0$ from one another. If this cannot be found after 200 trials an exception is raised. As soon as a good placement is found, the loop is left\\\n",
    "2.\\\n",
    "Assign random linear speeds to each disc, in the range [-5., 5.] for each component.\\\n",
    "3.\\\n",
    "Assign arbitray non zero values to rhodtmax and rhowall. They will be overwritten during the integration and (hopefully) filled with the correct values.\\\n",
    "4.\\\n",
    "Calculate $k_0$ and $k_{0W}$. I model the wall as a disc with $r_W = -R_W$, all discs have same radius $r_0$\\\n",
    "$\\sqrt{\\frac{r_0 \\cdot r_0}{r_0 + r_0}} = \\sqrt{\\frac{r_0}{2}} \\space \\space$ for collisions between discs\n",
    "I make the assumption, that the material constants of discs and walls are the same.\\\n",
    "If Young's modulus is large, the integration runs 'forever'. I use $2\\cdot 10^3$.\n",
    "\n",
    "**NOTE**: If I set Young's modulus, $E_Y$ much larger than $10^2$ the numerical integration runs very long, or even aborts. I guess, the system is too stiff then."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "start1 = time.time()\n",
    "\n",
    "# Input variables\n",
    "#-------------------------------------------------------------------\n",
    "r01   = 1.0           # radius of the pendulum\n",
    "RW1   = 5.            # length of each wall\n",
    "m01   = 1.            # mass of a pendulum\n",
    "mo1   = 1.           # mass of the observer\n",
    "\n",
    "mu1   = 0.1           # friction between discs\n",
    "muW1  = 0.1           # friction between disc and wall\n",
    "ctau1 = 0.9          # given in the article\n",
    "\n",
    "# initial conditions for the rotation of each disc\n",
    "q_list1 = [0. for _ in range(len(q_list))]\n",
    "\n",
    "# location of the observers\n",
    "alpha_list1 = [0.99 for _ in range(len(alpha_list))]\n",
    "#-------------------------------------------------------------------\n",
    "#np.random.seed(11)\n",
    "# 1. randomly place the discs as described above\n",
    "zaehler = 0\n",
    "while zaehler <= 200:\n",
    "    zaehler += 1\n",
    "    try:\n",
    "        x_listen = []\n",
    "        z_listen = []\n",
    "        for i in range(n):\n",
    "            x_listen.append(np.random.choice(np.linspace(-RW1 + 2.*r01, RW1 - 2.*r01, 100)))\n",
    "            z_listen.append(np.random.choice(np.linspace(-np.sqrt(RW1**2 - x_listen[-1]**2) + 2.*r01, \n",
    "                    np.sqrt(RW1**2 - x_listen[-1]**2) - 2.*r01, 100)))\n",
    "        test = np.all(np.array(Dmc_distanz_lam(*x_listen, *z_listen)) - 3.*r01 > 0.)\n",
    "        x_list1 = x_listen\n",
    "        z_list1 = z_listen\n",
    "        \n",
    "        if test == True:\n",
    "            raise Rausspringen\n",
    "    except:\n",
    "        break\n",
    "\n",
    "if zaehler <= 200:\n",
    "    print(f'it took {zaehler} rounds to get valid initial conditions')\n",
    "    print('distance between discs is', [f'{np.array(Dmc_distanz_lam(*x_list1, *z_list1))[l] - 2.*r01:.2f}' \n",
    "                for l in range(len(Dmc_distanz))])\n",
    "else:\n",
    "    raise Exception(' no good location for discs found, make RW1 larger, or try again.')\n",
    "\n",
    "# 2. Assign random linear and angular speeds to each disc\n",
    "ux_list1 = np.random.choice(np.linspace(-5., 5., 100), size=n)\n",
    "uz_list1 = np.random.choice(np.linspace(-5., 5., 100), size=n)\n",
    "u_list1  = np.random.choice(np.linspace(-15., 15., 100), size=len(q_list))\n",
    "\n",
    "ux_list1 = list(ux_list1)\n",
    "uz_list1 = list(uz_list1)\n",
    "u_list1  = list(u_list1)\n",
    "\n",
    "# 3. Assign non-zero values to rhodtmax and rhodtwall\n",
    "rhodtmax1  = [1. + k for k in range(len(rhodtmax))]\n",
    "rhodtwall1 = [1. + k for k in range(n)]\n",
    "\n",
    "# 4. Calculate k01 and k0W1\n",
    "nu = 0.28                       # Poisson's ratio for steel, from the internet\n",
    "EY = 2.e3                       # units: N/m^2, Young's modulus from the internet, around 2*10^11 for steel\n",
    "RW11 = -RW1\n",
    "sigma = (1 - nu**2) / EY\n",
    "k01  = 4. / (3.* (sigma + sigma)) * np.sqrt(r01/2.)\n",
    "k0W1 = 4. / (3.* (sigma + sigma)) * np.sqrt((r01*RW11)/(r01 + RW11))\n",
    "\n",
    "iYY1 = 0.25 * m01 * r01**2      # from the internet\n",
    "\n",
    "y0       = x_list1 + z_list1 + q_list1 + ux_list1 + uz_list1 + u_list1 \n",
    "pL_vals  = [m01, mo1, r01, RW1, k01, k0W1, iYY1, ctau1, mu1, muW1] + alpha_list1 + rhodtmax1 + rhodtwall1\n",
    "pL1_vals = [m01, mo1, r01, RW1, k01, k0W1, iYY1, ctau1, mu1, muW1]\n",
    "pL2_vals = [m01, mo1, r01, RW1, k01, k0W1, iYY1, ctau1, mu1, muW1] + list(alpha_list1)\n",
    "print('initial conditions are:', [f'{pL_vals[l]}' for l in range(len(pL_vals))])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "**Numerical integration**\n",
    "\n",
    "$\\dot \\rho^{(-)}$ has to be found numerically during integration, at least I could not think of any other way.\\\n",
    "If $\\dot \\rho^{(-)} \\approx 0.$ I set it to $10^{-12}$ to avoid numerical issues. \\\n",
    "I use *max_step* in solve_ivp to ensure it does not miss the time close to a collision. Hoever, this drastically increases the time it take to integrate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#--------------------------------------\n",
    "# I set these values here only, so I can integrate the same initial conditions with different \n",
    "# parameters for the integration routine.\n",
    "intervall = 10.                          \n",
    "max_step = 1.\n",
    "schritte = 20000  \n",
    "#--------------------------------------\n",
    "\n",
    "start1 = time.time()\n",
    "times = np.linspace(0, intervall, schritte)\n",
    "\n",
    "def gradient(t, y, args):\n",
    "    global X0, zaehler, zaehlerm\n",
    "# here I find rhodtmax, the collision speed of two discs just before the impact\n",
    "    zaehler = -1\n",
    "    for i, j in permutations(range(n), r=2):\n",
    "        zaehler += 1\n",
    "        if Dmc_distanz_lam(*[y[l] for l in range(2*n)])[zaehler] - 2.*args[2] < 2. * max_step:\n",
    "            hilfs1 = rhomax_list_lam(*y, *pL1_vals)[zaehler]\n",
    "            if np.abs(hilfs1) < 1.e-12:\n",
    "                hilfs1 = np.sign(hilfs1) * 1.e-12\n",
    "            args[10 + n + zaehler] = hilfs1\n",
    "\n",
    "#=======================================================================            \n",
    "# Here I find rhodtwall, the collision speed between a disc and the wall\n",
    "# To do so, I first have to find the contact points.\n",
    "    CPx_list1 = []\n",
    "    CPz_list1 = []\n",
    "    args1 = [args[j] for j in range(10)]\n",
    "    for i in range(n):\n",
    "        if bCP_lam(*y, *args1)[i] < RW1 - r01:\n",
    "            CPx_list1.append(CP_ort_lam(*y, *args1)[i][0])\n",
    "            CPz_list1.append(CP_ort_lam(*y, *args1)[i][1])\n",
    "        else:\n",
    "            CPx_list1.append(CPh_ort_lam(*y, *args1)[i][0])\n",
    "            CPz_list1.append(CPh_ort_lam(*y, *args1)[i][1])\n",
    "    \n",
    "    laenge = len(rhodtmax)\n",
    "    for i in range(n):\n",
    "        abstand = abstandDmcCP_lam(*y, *args1)[i]\n",
    "        if abstand - args[2] < 2.*max_step:\n",
    "            hilfs1 = rhowall_list_lam(*y, *pL1_vals, *CPx_list1, *CPz_list1)[i]\n",
    "\n",
    "            if np.abs(hilfs1) < 1.e-12:\n",
    "                hilfs1 = np.sign(hilfs1) * 1.e-12\n",
    "            args[10 + n + laenge + i] = hilfs1\n",
    "        \n",
    "    sol = np.linalg.solve(MM_lam(*y, *args), force_lam(*y, *args, *CPx_list1, *CPz_list1))\n",
    "    return np.array(sol).T[0]\n",
    "\n",
    "    \n",
    "resultat1 = solve_ivp(gradient, (0., float(intervall)), y0, args=(pL_vals,), t_eval=times, method='Radau' \n",
    "            , atol=1.e-6, rtol=1.e-6, max_step=max_step)\n",
    "resultat = resultat1.y.T\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",
    "print('to calculate an intervall of {:.2f} sec it took {} loops and {:.3f} sec running time'.\n",
    "      format(intervall, resultat1.nfev, time.time() - start1))\n",
    "print(resultat.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This shows, how close the contact points are to the vall. Ideally this should be zero.\\\n",
    "$CP_X, CP_Z$ are used further down, when plotting the energies, so this must be run before the energies are run.\\\n",
    "For some reason, this takes a long time to run.\\\n",
    "In my trial runs, the errors of all $CP_i$ were (almost?) identical, I do not know, why."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "CP_X = np.empty((schritte, n))\n",
    "CP_Z = np.empty((schritte, n))\n",
    "\n",
    "for i in range(schritte):\n",
    "    for kk in range(n):\n",
    "        if bCP_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL1_vals)[kk] < RW1 - r01:\n",
    "            CP_X[i] = CP_ort_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL1_vals)[kk][0]\n",
    "            CP_Z[i] = CP_ort_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL1_vals)[kk][1]\n",
    "        else:\n",
    "            CP_X[i] = CPh_ort_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL1_vals)[kk][0]\n",
    "            CP_Z[i] = CPh_ort_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL1_vals)[kk][1]\n",
    "\n",
    "CPRR = []\n",
    "for kk in range(n):\n",
    "    CPR = []\n",
    "    for i in range(schritte):\n",
    "        CPR.append(np.sqrt(np.abs(RW1**2 - CP_X[i][kk]**2 - CP_Z[i][kk]**2)))\n",
    "    CPRR.append(CPR)\n",
    "\n",
    "\n",
    "fig, axes = plt.subplots(n, 1, sharex=True)\n",
    "fig.set_size_inches((10.0, 2.0*n))\n",
    "\n",
    "for kk in range(n):\n",
    "    if kk == 0:\n",
    "        axes[kk].set_title(\"Deviation of contact points from ideal position\")\n",
    "    axes[kk].plot(times, CPRR[kk], label='CP_'+str(kk))\n",
    "    axes[kk].set_ylabel('deviation (m)')\n",
    "    if kk == n-1:\n",
    "        axes[kk].set_xlabel('time( sec)')\n",
    "    axes[kk].legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*plot* whichever generalized coordinates you want to see.\\\n",
    "*schritte* may be very large, to catch the impacts, this is not needed her. Hence I reduce the number of points to be plotted to around $N_2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "N2 = 500\n",
    "N1 = int(resultat.shape[0] / N2)\n",
    "times1 = []\n",
    "resultat1 = []\n",
    "for i in range(resultat.shape[0]):\n",
    "    if i % N1 == 0:\n",
    "        times1.append(times[i])\n",
    "        resultat1.append(resultat[i])\n",
    "resultat1 = np.array(resultat1)\n",
    "times1 = np.array(times1)\n",
    "\n",
    "bezeichnung = (['x' + str(i) for i in range(n)] +\n",
    "               ['z' + str(i) for i in range(n)] +\n",
    "               ['q' + str(i) for i in range(n)] + \n",
    "               ['ux' + str(i) for i in range(n)] +\n",
    "               ['uz' + str(i) for i in range(n)] +\n",
    "               ['u' + str(i) for i in range(n)]  \n",
    "              \n",
    "              ) \n",
    "fig, ax = plt.subplots(figsize=(10,5))\n",
    "for i in range(0*n, 2*n):\n",
    "    label = 'gen. coord. ' + str(i)\n",
    "    ax.plot(times1, resultat1[:, i], label=bezeichnung[i])\n",
    "ax.set_title('Generalized coordinates')\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('units of whichever coordinates are chosen')\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the **hysteresis curve** of $disc_i$ when it collides with the wall and $disc_i$ when it collides with $disc_j$, $i < j$\\\n",
    "The $i_0$ is needed to get the (approximately) correct $\\dot \\rho^{(-)}$, the speed at the impact.\\\n",
    "I only plot, if an impact did take place.\\\n",
    "The black numbers on the curves give the time of the impact."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "vDmc       = [Dmc_list[i].vel(N).magnitude() for i in range(n)]\n",
    "vorzeichen = [sm.sign(me.dot(Dmc_list[i].pos_from(P0).normalize(), Dmc_list[i].vel(N).normalize())) \n",
    "                for i in range(n)]\n",
    "vDmc_lam = sm.lambdify(qL + pL2, vDmc, cse=True)\n",
    "vorzeichen_lam = sm.lambdify(qL + pL2, vorzeichen, cse=True)\n",
    "\n",
    "for l1 in range(n):\n",
    "    HC_kraft = []\n",
    "    HC_displ = []\n",
    "    HC_times = []\n",
    "    zaehler  = 0\n",
    "\n",
    "    i0 = 0\n",
    "    for i in range(resultat.shape[0]):\n",
    "        abstand = r01 - abstandDmcCP_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL1_vals)[l1]\n",
    "        if abstand < 0.:\n",
    "            i0 = i+1\n",
    "\n",
    "        if abstand >= 0. and i0 == i:\n",
    "            walldt = vDmc_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL2_vals)[l1]\n",
    "        if abstand >= 0.:\n",
    "            rhodt = vorzeichen_lam(*[resultat[i, j] for j in range(resultat.shape[1])], \n",
    "                *pL2_vals)[l1] * vDmc_lam(*[resultat[i, j] for j in range(resultat.shape[1])], \n",
    "                *pL2_vals)[l1] \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",
    "    if len(HC_displ) > 0:\n",
    "        HC_displ = np.array(HC_displ)\n",
    "        HC_kraft = np.array(HC_kraft)\n",
    "\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 (N)')\n",
    "        ax.set_title(f'hysteresis curves of successive impacts of dics_{l1} with wall_0, ctau = {ctau1}' + \n",
    "                     f'mu = {mu1}, 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=\"black\")\n",
    "\n",
    "\n",
    "# plot the hysteresis curve of disc i colliding with disc j, for i < j only,\n",
    "zaehler1 = -1\n",
    "for l1, l2 in permutations(range(n), r=2):\n",
    "    zaehler1 += 1\n",
    "    HC_kraft = []\n",
    "    HC_displ = []\n",
    "    HC_times = []\n",
    "    zaehler  = 0\n",
    "\n",
    "    i0 = 0\n",
    "    for i in range(resultat.shape[0]):\n",
    "        abstand = 2.*r01 - Dmc_distanz_lam(*[resultat[i, j] for j in range(2*n)])[zaehler1]\n",
    "        if abstand < 0.:\n",
    "            i0 = i+1\n",
    "\n",
    "        if abstand >= 0. and i0 == i:\n",
    "            walldt = rhomax_list_lam(*[resultat[i, j] for j in range(resultat.shape[1])], \n",
    "                    *pL1_vals)[zaehler1]\n",
    "        if abstand >= 0.:\n",
    "            rhodt  = rhomax_list_lam(*[resultat[i, j] for j in range(resultat.shape[1])], \n",
    "                    *pL1_vals)[zaehler1]\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",
    "    if len(HC_displ) > 0 and l1 < l2:\n",
    "        HC_displ = np.array(HC_displ)\n",
    "        HC_kraft = np.array(HC_kraft)\n",
    "\n",
    "        fig, ax = plt.subplots(figsize=(10,5))\n",
    "        ax.plot(HC_displ, HC_kraft, color='red')\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 dics_{l1} with disc_{l2}' +\n",
    "                     f' ctau = {ctau1}, mu = {mu1}, muW = {muW1}');\n",
    "    \n",
    "        zeitpunkte = 12\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=\"black\")    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the **energies** of the system.\n",
    "\n",
    "If $c_{tau} < 1.$, or $m_u \\neq 0.$ or $m_{uW} \\neq 0.$ the total energy should drop monotonically. This is due to the Hunt-Crossley prescription of the forces during a collision, due to friction during the collisions respectively.\\\n",
    "Unless the step size is very small, impacts may not be shown. However, this increases the running time a lot!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "kin_np    = np.empty(schritte)\n",
    "spring_np = np.empty(schritte)\n",
    "total_np  = np.empty(schritte)\n",
    "\n",
    "for i in range(schritte):\n",
    "    kin_np[i]    = kin_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL2_vals)\n",
    "    spring_np[i] = spring_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL2_vals, \n",
    "        *[CP_X[i, j] for j in range(n)], *[CP_Z[i, j] for j in range(n)])\n",
    "    total_np[i]  = spring_np[i] + kin_np[i]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "for i, j in zip((kin_np, spring_np, total_np), ('kinetic energy', 'spring energy', 'total energy')):\n",
    "    ax.plot(times, i, label=j)\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('Energy (Nm)')\n",
    "ax.set_title(f'Energies of the system with {n} balls, with ctau = {ctau1}, mu = {mu1} and muW = {muW1}')\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Animation*\n",
    "\n",
    "As the number of points in time, given as *schritte* may be verly large, I limit to around *zeitpunkte*. Otherwise it would take a very long time to finish the animation.\n",
    "\n",
    "The size of the discs, and the depth of penetration is not to scale. I would not know, how to do this."
   ]
  },
  {
   "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",
    "Dmc_X = np.array([[resultat2[i, j] for j in range(n)] for i in range(schritte2)])\n",
    "Dmc_Z = np.array([[resultat2[i, j] for j in range(n, 2*n)] for i in range(schritte2)])\n",
    "\n",
    "Po_X = np.empty((schritte2, n))\n",
    "Po_Z = np.empty((schritte2, n))\n",
    "CP_X = np.empty((schritte2, n))\n",
    "CP_Z = np.empty((schritte2, n))\n",
    "\n",
    "for i in range(schritte2):\n",
    "    Po_X[i] = [Po_pos_lam(*[resultat2[i, j] for j in range(resultat.shape[1])], *pL2_vals)[l][0] \n",
    "            for l in range(n) ]\n",
    "    Po_Z[i] = [Po_pos_lam(*[resultat2[i, j] for j in range(resultat.shape[1])], *pL2_vals)[l][1] \n",
    "            for l in range(n)]\n",
    "\n",
    "for i in range(schritte2):\n",
    "    for kk in range(n):\n",
    "        if bCP_lam(*[resultat2[i, j] for j in range(resultat2.shape[1])], *pL1_vals)[kk] < RW1 - r01:\n",
    "            CP_X[i] = CP_ort_lam(*[resultat2[i, j] for j in range(resultat.shape[1])], *pL1_vals)[kk][0]\n",
    "            CP_Z[i] = CP_ort_lam(*[resultat2[i, j] for j in range(resultat.shape[1])], *pL1_vals)[kk][1]\n",
    "        else:\n",
    "            CP_X[i] = CPh_ort_lam(*[resultat2[i, j] for j in range(resultat.shape[1])], *pL1_vals)[kk][0]\n",
    "            CP_Z[i] = CPh_ort_lam(*[resultat2[i, j] for j in range(resultat.shape[1])], *pL1_vals)[kk][1]\n",
    "        \n",
    "# This is to asign colors of 'plasma' to the discs.\n",
    "Test = mp.colors.Normalize(0, n)\n",
    "Farbe = mp.cm.ScalarMappable(Test, cmap='plasma')\n",
    "farben = [Farbe.to_rgba(l) for l in range(n)]    # color of the starting position\n",
    "    \n",
    "def animate_pendulum(times2, Dmc_x, Dmc_Z, Po_X, Po_Z):\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(8, 8))\n",
    "    ax.axis('on')\n",
    "    theta = np.linspace(0., 2.*np.pi, 200)\n",
    "    aa = RW1 * np.sin(theta)\n",
    "    bb = RW1 * np.cos(theta)\n",
    "    ax.plot(aa, bb, linewidth=2)\n",
    "    \n",
    "    LINE1 = []\n",
    "    LINE2 = []\n",
    "    LINE3 = []\n",
    "\n",
    "    for i in range(n):\n",
    "# picking the 'right' radius of the discs I do by trial and error. I did not try to get a formula\n",
    "        line1, = ax.plot([], [], 'o', markersize=400./RW1)\n",
    "        line2, = ax.plot([], [], 'o', markersize=5, color='black')\n",
    "        line3, = ax.plot([], [], '-', markersize=0, linewidth=0.3)\n",
    "     \n",
    "        LINE1.append(line1)\n",
    "        LINE2.append(line2)\n",
    "        LINE3.append(line3)\n",
    "\n",
    "    def animate(i):\n",
    "        ax.set_title(f'System with {n} bodies, running time {i/schritte2 * intervall:.2f} sec' + '\\n' +\n",
    "                     f' ctau = {ctau1}, mu = {mu1}, muW = {muW1}', fontsize=12)\n",
    "        for j in range(n):\n",
    "            LINE1[j].set_data(Dmc_X[i, j], Dmc_Z[i, j])\n",
    "            LINE1[j].set_color(farben[j])\n",
    "            LINE2[j].set_data(Po_X[i, j], Po_Z[i, j])\n",
    "            LINE3[j].set_data(Dmc_X[:i, j], Dmc_Z[:i, j])\n",
    "            LINE3[j].set_color(farben[j])\n",
    "            \n",
    "            x_werte = (Dmc_X[i, j], CP_X[i, j])\n",
    "            z_werte = (Dmc_Z[i, j], CP_Z[i, j])\n",
    "           \n",
    "        return LINE1 + LINE2 + LINE3\n",
    "\n",
    "    anim = animation.FuncAnimation(fig, animate, frames=schritte2,\n",
    "                                   interval=1000*times2.max() / schritte2,\n",
    "                                   blit=True)\n",
    "    plt.close(fig)\n",
    "    return anim\n",
    "\n",
    "anim = animate_pendulum(times2, Dmc_X, Dmc_Z, Po_X, Po_Z)\n",
    "print(f'it took {time.time() - start :.3f} sec to run the program, BEFORE HTML')\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}