h-c-colliding-discs-friction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy.physics.mechanics as me\n",
    "import sympy as sm\n",
    "from scipy.integrate import solve_ivp\n",
    "import numpy as np\n",
    "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\n",
    "import time\n",
    "matplotlib.rcParams['animation.embed_limit'] = 2**128\n",
    "sm.init_printing(use_latex='mathjax')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is needed to exit a loop, when a feasible location of the discs within the limitations of the walls 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 four walls: $\\overline{(0, 0) (l_W, 0)}, \\overline{(l_W, 0)  (l_W, l_W)}, \\overline{(l_W, l_W) (0, l_W)}, \\overline{(0, l_W)  (0, 0)}$.\\ \n",
    "I call them $wall_0$...$wall_3$, with $wall_0$ the segment of the X - axis, and counted clock wise.\\\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 a wall through $Dmc_i$.\\\n",
    "There is speed dependent friction between discs, with coefficient $m_u$ and between walls 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",
    "**Note**  \n",
    "$c_\\tau = 1.$ gives **Hertz's** solution to the impact problem, also described in the article.\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",
    "- $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, l_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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "#==========================================\n",
    "n = 4       # 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",
    "rhodtmax = []\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    rhodtmax.append(sm.symbols('rhodtmax' + str(i) + str(j)))\n",
    "\n",
    "rhodtwall  = []\n",
    "for i in range(n):\n",
    "    rhodtwall.append([sm.symbols('rhodtwall' + '_' + str(i) + '_' + str(j)) for j in range(4)])\n",
    "        \n",
    "m0, mo, r0, lW, k0, k0W, iYY, ctau, mu, muW = sm.symbols('m0, mo, r0, lW, 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",
    "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",
    "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",
    "    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 a 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, r, ctau, rhodtwall, k0W):\n",
    "    \n",
    "    '''\n",
    "Calculates the contact force exerted by the four walls 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 wall. \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",
    "- 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 = oo.\n",
    "\n",
    "    '''\n",
    "# wall0\n",
    "    abstand = me.dot(P1.pos_from(P0), N.z)     # this is simply the Z coordinate of P1\n",
    "    richtung = N.z                             # points from wall0 to P1       \n",
    "    \n",
    "    rho = (r - abstand)       # penetration. Positive if the disc is in collision with wall0\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[0]))        \n",
    "                                ) * sm.Heaviside(r - abstand, 0.) * richtung\n",
    "    \n",
    "# wall1\n",
    "    abstand = lW - me.dot(P1.pos_from(P0), N.x) # this is the distance of P1 to wall1\n",
    "    richtung = -N.x                             # points from wall1 to P1       \n",
    "    \n",
    "    rho = (r - abstand)       # penetration. Positive if the disc is in collision with wall0\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",
    "    kraft1 = k0W * rho**(3/2) * (1. + 3./2. * (1 - ctau) * (-rhodt/sm.Abs(rhodtwall[1]))        \n",
    "                                ) * sm.Heaviside(r - abstand, 0.) * richtung \n",
    "    \n",
    "    \n",
    "# wall2\n",
    "    abstand = lW - me.dot(P1.pos_from(P0), N.z) # this is the distance of P1 to wall2\n",
    "    richtung = -N.z                             # points from wall2 to P1       \n",
    "    \n",
    "    rho = (r - abstand)       # penetration. Positive if the disc is in collision with wall0\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",
    "    kraft2 = k0W * rho**(3/2) * (1. + 3./2. * (1 - ctau) * (-rhodt/sm.Abs(rhodtwall[2]))        \n",
    "                                ) * sm.Heaviside(r - abstand, 0.) * richtung \n",
    "    \n",
    "\n",
    "# wall3\n",
    "    abstand = me.dot(P1.pos_from(P0), N.x) # this is the distance of P1 to wall3\n",
    "    richtung = N.x                             # points from wall3 to P1       \n",
    "    \n",
    "    rho = (r - abstand)       # penetration. Positive if the disc is in collision with wall0\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",
    "    kraft3 = k0W * rho**(3/2) * (1. + 3./2. * (1 - ctau) * (-rhodt/sm.Abs(rhodtwall[3]))\n",
    "        ) * sm.Heaviside(r - abstand, 0.) * richtung \n",
    "    \n",
    "    return kraft0 + kraft1 + kraft2 + kraft3   "
   ]
  },
  {
   "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",
    "    CP1, CP2 = sm.symbols('CP1, CP2', cls=me.Point)\n",
    "    \n",
    "    abstand = P2.pos_from(P1)\n",
    "    CP1.set_pos(P1, abstand / 2.)\n",
    "    CP1.v2pt_theory(P1, N, A1)\n",
    "    \n",
    "    CP2.set_pos(P2, -abstand / 2.)\n",
    "    CP2.v2pt_theory(P2, N, A2)\n",
    "    delta_v = CP2.vel(N) - CP1.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",
    "#    print(me.find_dynamicsymbols(kraft, reference_frame=N))\n",
    "    return [kraft, torque]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Friction when a disc hits a wall*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def friction_wall(N, A1, P1, 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",
    "Of course, a disc can never hit more than two wall simultaneously.\n",
    "    '''\n",
    "    \n",
    "# wall 0\n",
    "    abstand = me.dot(P1.pos_from(P0), N.z) * N.z\n",
    "    \n",
    "    CP0 = me.Point('CP0')\n",
    "    CP0.set_pos(P1, -abstand)\n",
    "    CP0.v2pt_theory(P1, N, A1)\n",
    "    force_coll = HC_wall(N, P1, r, ctau, rhodtwall, k0W).magnitude()\n",
    "    force_fric = -force_coll * mu * CP0.vel(N)\n",
    "    \n",
    "    torque0 = CP0.pos_from(P1).cross(force_fric) * sm.Heaviside(r - abstand.magnitude())\n",
    "    kraft0  = 1./me.dot(abstand, abstand) * torque0.cross(-abstand) * sm.Heaviside(r - abstand.magnitude())\n",
    "\n",
    "# wall 1\n",
    "    abstand = lW*N.x - me.dot(P1.pos_from(P0), N.x) * N.x\n",
    "    \n",
    "    CP1 = me.Point('CP1')\n",
    "    CP1.set_pos(P1, abstand)\n",
    "    CP1.v2pt_theory(P1, N, A1)\n",
    "\n",
    "    force_coll = HC_wall(N, P1, r, ctau, rhodtwall, k0W).magnitude()\n",
    "    force_fric = -force_coll * mu * CP1.vel(N)\n",
    "    \n",
    "    torque1 = CP1.pos_from(P1).cross(force_fric) * sm.Heaviside(r - abstand.magnitude())\n",
    "    kraft1  = 1./me.dot(abstand, abstand) * torque1.cross(abstand) * sm.Heaviside(r - abstand.magnitude())\n",
    "\n",
    "# wall 2\n",
    "    abstand = lW*N.z - me.dot(P1.pos_from(P0), N.z) * N.z\n",
    "    CP2 = me.Point('CP2')\n",
    "    CP2.set_pos(P1, abstand)\n",
    "    CP2.v2pt_theory(P1, N, A1)\n",
    "\n",
    "    force_coll = HC_wall(N, P1, r, ctau, rhodtwall, k0W).magnitude()\n",
    "    force_fric = -force_coll * mu * CP2.vel(N)\n",
    "    \n",
    "    torque2 = CP2.pos_from(P1).cross(force_fric) * sm.Heaviside(r - abstand.magnitude())\n",
    "    kraft2  = 1./me.dot(abstand, abstand) * torque2.cross(abstand) * sm.Heaviside(r - abstand.magnitude())\n",
    "\n",
    "    \n",
    "# wall 3\n",
    "    abstand = me.dot(P1.pos_from(P0), N.x) * N.x\n",
    "\n",
    "    CP3 = me.Point('CP3')\n",
    "    CP3.set_pos(P1, -abstand)\n",
    "    CP3.v2pt_theory(P1, N, A1)\n",
    "\n",
    "    force_coll = HC_wall(N, P1, r, ctau, rhodtwall, k0W).magnitude()\n",
    "    force_fric = -force_coll * mu * CP3.vel(N)\n",
    "    \n",
    "    torque3 = CP3.pos_from(P1).cross(force_fric) * sm.Heaviside(r - abstand.magnitude())\n",
    "    kraft3  = 1./me.dot(abstand, abstand) * torque3.cross(-abstand) * sm.Heaviside(r - abstand.magnitude())\n",
    "    \n",
    "    return [kraft0 + kraft1 + kraft2 + kraft3, torque0 + torque1 + torque2 + torque3]   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**force from the collision of discs with discs and discs with walls**\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], r0, ctau, rhodtwall[i], k0W )))\n",
    "\n",
    "#------------------------------------------------------------------------------------------------------  \n",
    "# Friction forces between discs\n",
    "zaehler = 0\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[i], 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[i], 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], r0, ctau, rhodtwall[i],\n",
    "            k0W, muW)[0]))\n",
    "    FL.append((A_list[i], friction_wall(N, A_list[i], Dmc_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",
    "    hilfs = []\n",
    "    for richtung in (N.z, -N.x, -N.z, N.x):\n",
    "        rhodt = me.dot(Dmc_list[i].vel(N), richtung)\n",
    "        hilfs.append(rhodt)\n",
    "    rhowall_list.append(hilfs)\n",
    "print('rhowall_list DS:', set.union(*[me.find_dynamicsymbols(rhowall_list[k][l]) \n",
    "        for k in range(len(rhowall_list)) for l in range(4)]), '\\n')\n",
    "print('rhowall_list free symbols:', set.union(*[rhowall_list[k][l].free_symbols \n",
    "        for k in range(len(rhowall_list)) for l in range(4)]))\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 walls   \n",
    "for i in range(n):\n",
    "# wall0\n",
    "    abstand = sm.Abs(me.dot(Dmc_list[i].pos_from(P0), N.z))     # this is simply the Z coordinate of P1\n",
    "    rho = (r0 - abstand)       # penetration. Positive if the disc is in collision with wall0\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",
    "# wall1\n",
    "    abstand = lW - sm.Abs(me.dot(Dmc_list[i].pos_from(P0), N.x))     # this is simply the Y coordinate of P1\n",
    "    rho = (r0 - abstand)       # penetration. Positive if the disc is in collision with wall0\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",
    "# wall2\n",
    "    abstand = lW - sm.Abs(me.dot(Dmc_list[i].pos_from(P0), N.z))     # this is simply the Y coordinate of P1\n",
    "    rho = (r0 - abstand)       # penetration. Positive if the disc is in collision with wall0\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",
    "# wall3\n",
    "    abstand = sm.Abs(me.dot(Dmc_list[i].pos_from(P0), N.x))     # this is simply the Y coordinate of P1\n",
    "    rho = (r0 - abstand)       # penetration. Positive if the disc is in collision with wall0\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.)"
   ]
  },
  {
   "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, lW, k0, k0W, iYY, ctau, mu, muW] + list(alpha_list) + rhodtmax + rhodtwall\n",
    "pL1 = [m0, mo, r0, lW, k0, k0W, iYY, ctau, mu, muW]\n",
    "pL2 = [m0, mo, r0, lW, k0, k0W, iYY, ctau, mu, muW] + list(alpha_list)\n",
    "\n",
    "MM_lam    = sm.lambdify(qL + pL, MM, cse=True)\n",
    "force_lam = sm.lambdify(qL + pL, force, cse=True) \n",
    "\n",
    "rhomax_list_lam  = sm.lambdify(qL + pL1, rhomax_list, cse=True)\n",
    "rhowall_list_lam = sm.lambdify(qL + pL1, rhowall_list, cse=True)\n",
    "Dmc_distanz_lam  = sm.lambdify(x_list + z_list, Dmc_distanz, cse=True)\n",
    "\n",
    "#Dmc_pos_lam = sm.lambdify(qL + pL2, Dmc_pos, cse=True)\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, 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": [
    "**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 \\to \\infty$, all discs have same radius $r_0$\\\n",
    "Now $\\displaystyle \\lim_{r_W \\to \\infty} \\sqrt{\\frac{r_0 \\cdot r_W}{r_0 + r_W}} = \\sqrt{r_0}$, and $\\sqrt{\\frac{r_0 \\cdot r_0}{r_0 + r_0}} = \\sqrt{\\frac{r_0}{2}}$\\\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$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "start1 = time.time()\n",
    "\n",
    "# Input variables\n",
    "#-------------------------------------------------------------------\n",
    "r01   = 1.0           # radius of the pendulum\n",
    "lW1   = 11.            # length of each wall\n",
    "m01   = 1.            # mass of a pendulum\n",
    "mo1   = 0.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.8 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_list1 = np.random.choice(np.linspace(2.*r01, lW1 - 2.*r01, 100), size=n, replace=False)\n",
    "        z_list1 = np.random.choice(np.linspace(2.*r01, lW1 - 2.*r01, 100), size=n, replace=False)\n",
    "        test = np.all(np.array(Dmc_distanz_lam(*x_list1, *z_list1)) - 3.*r01 > 0.)\n",
    "        x_list1 = list(x_list1)\n",
    "        z_list1 = list(z_list1)\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 lW0 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(-5., 5., 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+l+k for l in range(4)] 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",
    "sigma = (1 - nu**2) / EY\n",
    "k01  = 4. / (3.* (sigma + sigma)) * np.sqrt(r01/2.)\n",
    "k0W1 = 4. / (3.* (sigma + sigma)) * np.sqrt(r01)\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, lW1, k01, k0W1, iYY1, ctau1, mu1, muW1] + alpha_list1 + rhodtmax1 + rhodtwall1\n",
    "pL1_vals = [m01, mo1, r01, lW1, k01, k0W1, iYY1, ctau1, mu1, muW1]\n",
    "pL2_vals = [m01, mo1, r01, lW1, 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.\\\n",
    "I set *schritte* a bit larger than the number of function calls of solve_ivp, to ensure I get the collision times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "start1 = time.time()\n",
    "\n",
    "intervall = 10.                          \n",
    "max_step = 0.001\n",
    "schritte = 20000     \n",
    "times = np.linspace(0, intervall, schritte)\n",
    "\n",
    "def gradient(t, y, args):\n",
    "    \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 0. < 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",
    "# Here I find rhodtwall, the collision speed between a disc and a wall\n",
    "    for i in range(n):\n",
    "# wall 0\n",
    "        laenge = len(rhodtmax)\n",
    "        abstand = y[i+n]\n",
    "        if 0. < args[2] - abstand < 2.*max_step:\n",
    "            hilfs1 = rhowall_list_lam(*y, *pL1_vals)[i][0]\n",
    "            if np.abs(hilfs1) < 1.e-12:\n",
    "                hilfs1 = np.sign(hilfs1) * 1.e-12\n",
    "            args[10 + n + laenge + i][0] = hilfs1\n",
    "            \n",
    "# wall 1\n",
    "        laenge = len(rhodtmax)\n",
    "        abstand = args[3] - y[i]\n",
    "        hilfs1 = rhowall_list_lam(*y, *pL1_vals)[i][1]\n",
    "        if np.abs(hilfs1) < 1.e-12:\n",
    "            hilfs1 = np.sign(hilfs1) * 1.e-12\n",
    "        if 0. < args[2] - abstand < 2.*max_step:\n",
    "            args[10 + n + laenge + i][1] = hilfs1\n",
    "            \n",
    "# wall 2\n",
    "        laenge = len(rhodtmax)\n",
    "        abstand = args[3] - y[n+i]\n",
    "        if 0. < args[2] - abstand < 2.*max_step:\n",
    "            hilfs1 = rhowall_list_lam(*y, *pL1_vals)[i][2]\n",
    "            if np.abs(hilfs1) < 1.e-12:\n",
    "                hilfs1 = np.sign(hilfs1) * 1.e-12\n",
    "            args[10 + n + laenge + i][2] = hilfs1\n",
    "            \n",
    "# wall 3\n",
    "        laenge = len(rhodtmax)\n",
    "        abstand = y[i]\n",
    "        if 0. < args[2] - abstand < 2.*max_step:\n",
    "            hilfs1 = rhowall_list_lam(*y, *pL1_vals)[i][3]\n",
    "            if np.abs(hilfs1) < 1.e-12:\n",
    "                hilfs1 = np.sign(hilfs1) * 1.e-12\n",
    "            args[10 + n + laenge + i][3] = hilfs1\n",
    "            \n",
    "        \n",
    "    sol = np.linalg.solve(MM_lam(*y, *args), force_lam(*y, *args))\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='BDF' \n",
    "            , max_step=max_step, atol=1.e-6, rtol=1.e-6)\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])\n",
    "\n",
    "print('to calculate an intervall of {:.0f} sec it took {} loops and {:.3f} sec running time'.\n",
    "      format(intervall, resultat1.nfev, time.time() - start1))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*plot* whichever generalized coordinates you want to see.\\\n",
    "*schritte* is 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(5*n, 6*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.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the **hysteresis curve** of $disc_i$ when it collides with $wall_0$ and $disc_i$ when it collides with $disc_j$, $i < j$\\\n",
    "Hysteresis curves of subsequent collisions are nested, as the energy dissipated is less on every subsequent collision.\\\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": [
    "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 - resultat[i, n+l1]\n",
    "        if abstand < 0.:\n",
    "            i0 = i+1\n",
    "\n",
    "        if abstand >= 0. and i0 == i:\n",
    "            walldt = resultat[i, 4*n+l1]\n",
    "        if abstand >= 0.:\n",
    "            rhodt = resultat[i, 4*n+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",
    "    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 = 8\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\")\n",
    "    "
   ]
  },
  {
   "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."
   ]
  },
  {
   "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",
    "\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",
    "    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* is verly large, I limit to around *zeitpunkte*. Otherwise it would take a very long time to finish the animation.\\\n",
    "The location of the walls is reflected and rotated relative to my definition given above. This is a result of the orientation of the axis in matplotlib. As it does not really matter, I did not bother to 'correct' it.\\\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",
    "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",
    "farben = ['red', 'green', 'blue', 'orange', 'yellow', 'pink']\n",
    "if n > 6:\n",
    "    raise Exception('define more colors in the list farben')\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('off')\n",
    "    ax.set(xlim=(- 1., lW1 + 1.), ylim=( - 1., lW1 + 1.))\n",
    "    ax.plot([0, lW1], [0, 0], 'bo', linewidth=2, linestyle='-', markersize=0)\n",
    "    ax.text(lW1/2., -0.5, 'wall 0')\n",
    "    ax.plot([lW1, lW1], [0, lW1], 'bo', linewidth=2, linestyle='-', markersize=0)\n",
    "    ax.text(lW1 + 0.1, lW1/2.+0.1, 'wall 1')\n",
    "    ax.plot([lW1, 0], [lW1, lW1], 'bo', linewidth=2, linestyle='-', markersize=0)\n",
    "    ax.text(lW1/2. + 0.1, lW1 + 0.1, 'wall 2')\n",
    "    ax.plot([0, 0], [lW1, 0], 'bo', linewidth=2, linestyle='-', markersize=0)\n",
    "    ax.text(-1., lW1/2., 'wall 3')\n",
    "    \n",
    "    \n",
    "    LINE1 = []\n",
    "    LINE2 = []\n",
    "    LINE3 = []\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=65*11./lW1)\n",
    "        line2, = ax.plot([], [], 'o', markersize=5, color='black')\n",
    "        line3, = ax.plot([], [], '-', markersize=0, linewidth=0.3)\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",
    "        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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}