h-c-colliding-ellipse-real-1-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 minimize, root\n",
    "import numpy as np\n",
    "\n",
    "from numba import njit\n",
    "\n",
    "from itertools import permutations\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from matplotlib import animation\n",
    "from matplotlib import patches\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 ellipses**, named $Dmc_0....Dmc_{n-1}$ with semi axes $a, b$ 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",
    "An observer, a particle of mass $m_o$ may be attached anywhere within each disc.\n",
    "\n",
    "**General Comments**\\\n",
    "While the collision of the ellipse with the wall posed no major problems, this was not so with the ellipses colliding into each other.\n",
    "1. The best way to integrate was to use no *method* in solv_ivp, but select a very small *max_step* instead\n",
    "2. I tried two ways to get the distance between the ellipse:\n",
    "- minimize $|{}^{CPhe_i} \\bar r^{CPhe_j}(\\epsilon_i, \\epsilon_j)|$ directly, using scipy's minimize function\n",
    "- calculate $\\dfrac{d}{d\\epsilon_i} |{}^{CPhe_i} \\bar r^{CPhe_j}(\\epsilon_i, \\epsilon_j)|$ and $\\dfrac{d}{d\\epsilon_j} |{}^{CPhe_i} \\bar r^{CPhe_j}(\\epsilon_i, \\epsilon_j)|$, and solve for $\\epsilon_i, \\epsilon_j$.\n",
    " \n",
    "The second, the $\\nabla$ - method worked much better.\n",
    "\n",
    "3. If Young' modulus, $EY_e, EY_w$ are small, say, $10^4$, the total energy with friction = 0 and $c_{tau} = 1.$ are not constant. I think, the reason is this:\\\n",
    "In the spring constant of the Hunt-Crossley method there is a term $\\sqrt{\\dfrac{R_1 \\cdot R_2 }{R_1 + R_2}}$, where the $R_i$ are the osculating radii of the colliding bodies.\\\n",
    "My ellipses continue to rotate while penetrating each other, so this term varies during a collision. As the integration is discrete, this may be the reason.\\\n",
    "If $EY_e, EY_w$ are large, say $10^7$, the total energy is pretty constant, as now the penetration time is very short, and therefore the ellipse do not rotate much.\n",
    "\n",
    "\n",
    "**Note about the force during the collisions**\n",
    "\n",
    " **Hunt Crossley's method**\n",
    " \n",
    "My reference is this article, given to me by JM\\\n",
    "https://www.sciencedirect.com/science/article/pii/S0094114X23000782 \\\n",
    "\n",
    " \n",
    "This is with dissipation during the collision, the general force is given in (63) as\\\n",
    "$f_n = k_0 \\cdot \\rho + \\chi \\cdot \\dot \\rho$, with $k_0$ as above, $\\rho$ the penetration, and $\\dot\\rho$ the speed of the penetration.\\\n",
    "In the article it is stated, that $n = \\frac{3}{2}$ is a good choice, it is derived in Hertz' approach. Of course, $\\rho, \\dot\\rho$ must be the signed magnitudes of the respective vectors.\n",
    "\n",
    "A more realistic force is given in (64) as:\\\n",
    "$f_n = k_0 \\cdot \\rho^n + \\chi \\cdot \\rho^n\\cdot \\dot \\rho$, as this avoids discontinuity at the moment of impact.\n",
    "\n",
    "**Hunt and Crossley** give this value for $\\chi$, see table 1:\n",
    "\n",
    "$\\chi = \\dfrac{3}{2} \\cdot(1 - c_\\tau) \\cdot \\dfrac{k_0}{\\dot \\rho^{(-)}}$, \n",
    "where $c_\\tau = \\dfrac{v_1^{(+)} - v_2^{(+)}}{v_1^{(-)} - v_2^{(-)}}$, where $v_i^{(-)}, v_i^{(+)}$ are the speeds of $body_i$, before and after the collosion, see (45), $\\dot\\rho^{(-)}$ is the speed right at the time the impact starts. $c_\\tau$ is an experimental factor, apparently around 0.8 for steel.\n",
    "\n",
    "Using (64), this results in their expression for the force:\n",
    "\n",
    "$f_n = k_0 \\cdot \\rho^n \\left[1 + \\dfrac{3}{2} \\cdot(1 - c_\\tau) \\cdot \\dfrac{\\dot\\rho}{\\dot\\rho^{(-)}}\\right]$\n",
    "\n",
    "with $k_0 = \\frac{4}{3\\cdot(\\sigma_1 + \\sigma_2)} \\cdot \\sqrt{\\frac{R_1 \\cdot R_2}{R_1 + R_2}}$, where $\\sigma_i = \\frac{1 - \\nu_i^2}{E_i}$, with $\\nu_i$ = Poisson's ratio, $E_i$ = Young\"s modulus, $R_1, R_2$ the radii of the colliding bodies, $\\rho$ the penetration depth. All is near equations (54) and (61) of this article.\n",
    "\n",
    "As per the article, $n = \\frac{3}{2}$ is always to be used.\n",
    "\n",
    "*spring energy* =   $ k_0 \\cdot \\int_{0}^{\\rho} k^{3/2}\\,dk$ = $k_0 \\cdot\\frac{2}{5} \\cdot \\rho^{5/2}$\\\n",
    "I assume, the dissipated energy cannot be given in closed form, at least the article does not give one.\n",
    "\n",
    "**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 ellipse's' point of view, the wall is concave. I model this by taking $R_2 = -R_W$.\\\n",
    "As $max(a, b)| < |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 ellipses\n",
    "- $q_0...q_{n-1}$: generalized coordinates for the ellipse\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 ellipse\n",
    "- $N$: frame of inertia\n",
    "- $P_0$: point fixed in $N$\n",
    "- $A_i$: body fixed frame of the i-th ellipse\n",
    "\n",
    "- $m_i$: mass of the i-th ellipse\n",
    "- $Dmc_i$: center of the i-th ellipse\n",
    "- $Po_i$: observer (particle) on i-th ellipse\n",
    "- $\\alpha_i. \\beta_i$: distance of observer on i-th ellipse\n",
    "- $a, b, R_W$: semi axes of the ellipses, radius of the wall\n",
    "- $i_{YY_i}$: moment of ineratia of the i-th ellipse\n",
    "\n",
    "\n",
    "- $reibung$: coefficient of friction between ellipses / between ellipse and wall.\n",
    "- $nu_e, nu_w$: Poison's coefficients of the ellipses / of the wall\n",
    "- $EY_e, EY_w$: dto for Young's moduli\n",
    "- $c_\\tau$: the experimental constant needed for Hunt-Crossley\n",
    "- $rhodt_{max}$: the collision speed between two ellipses, to be determined during integration, needed for Hunt_Crossley\n",
    "- $rhodt_{wall}$: the collision speeds when $disc_i$ hits a wall\n",
    "- $CPh_i$: contact point of $ellipse_i$ with the wall\n",
    "- $CPhs_i$: point on the ellipse which has had contact with the wall. $|{}^{CPh_i} r^{CPhs_i}|$ is the penetration depth.\n",
    "\n",
    "- $CPhe_i$: potential contact points of $ellipse_i$ with $ellipse_j$. Penetration depth is $|{}^{CPhe_i} r^{CPhe_j}|$\n",
    "- $l_{list}, le_{list}$: lists holding the penetration depth of $ellipse_i$ with the wall / penetration depth of $ellipse_i$ and $ellipse_j$.\n",
    "- $epsilon_{list}, epsilone_{list}$: lists holding the angles of the different contact points as discribed in the body fixed frames $A_i$. \n",
    "\n",
    "With *friktion = True* the friction between ellipses / ellipses and wall will be considered. For some reason, this increases the time for integration a lot.\\\n",
    "With *elli_ghost = True* the ellipses may move through each other without colliding. This shows the impacts with the wall only.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start0 = time.time()\n",
    "#==========================================\n",
    "n = 2         # n > 1\n",
    "friktion   = True\n",
    "elli_ghost = False\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",
    "CPh_list  = list(sm.symbols(f'CPh:{n}', cls=me.Point))\n",
    "CPhx_list  = list(sm.symbols(f'CPhx:{n}'))\n",
    "CPhz_list  = list(sm.symbols(f'CPhz:{n}'))\n",
    "\n",
    "CPhs_list  = list(sm.symbols(f'CPhs:{n}', cls=me.Point)) \n",
    "CPhsx_list  = list(sm.symbols(f'CPhsx:{n}'))\n",
    "CPhsz_list  = list(sm.symbols(f'CPhsz:{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",
    "alpha_list   = list(sm.symbols(f'alpha:{n}'))\n",
    "beta_list    = list(sm.symbols(f'beta:{n}'))\n",
    "\n",
    "epsilon_list  = list(sm.symbols(f'epsilon:{n}'))\n",
    "l_list        = list(sm.symbols(f'l:{n}'))\n",
    "\n",
    "rhodtmax      = [sm.symbols(f'rhodtmax{i}{j}') for i, j in permutations(range(n), r=2)]\n",
    "le_list       = [sm.symbols(f'le{i}{j}') for i, j in permutations(range(n), r=2)]\n",
    "epsilone_list = [sm.symbols(f'epsilone{i}{j}, epsilone{j}{i}') for i, j in permutations(range(n), r=2)]\n",
    "rhodtwall     = list(sm.symbols(f'rhodtwall:{n}'))\n",
    "\n",
    "richtung_list = [sm.symbols(f'richtungx{i}, richtungz{j}')for i, j in permutations(range(n), r=2)]\n",
    "\n",
    "t = me.dynamicsymbols._t\n",
    "\n",
    "m0, mo, iYY, a, b, RW, nue, nuw, EYe, EYw, ctau, reibung = sm.symbols('m0, mo, iYY, a, b, RW, nue, nuw, EYe, EYw, ctau, reibung')\n",
    "\n",
    "N  = me.ReferenceFrame('N')\n",
    "P0 = me.Point('P0')\n",
    "P0.set_vel(N, 0)\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], a * alpha_list[i] * A_list[i].x + b * beta_list[i] * A_list[i].z)\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\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Find the point where the ellipse hits the wall**\\\n",
    "I simply look for the points on the ellipse and the wall, where their distance is the closest.\\\n",
    "I use *scipy.minimize* to do it numerically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lang = sm.symbols('lang')\n",
    "CPh  = [sm.symbols('CPh' + str(i), cls=me.Point) for i in range(n)]\n",
    "CPhs = [sm.symbols('CPhs' + str(i), cls=me.Point) for i in range(n)]    \n",
    "\n",
    "def CPhxe(epsilon):\n",
    "    return a * sm.cos(epsilon)\n",
    "def CPhze(epsilon):\n",
    "    return b * sm.sin(epsilon)\n",
    "\n",
    "CPh_list  = []\n",
    "CPhs_list = []                  # just needed for the plot of the initial conditions\n",
    "# define CPh\n",
    "for i in range(n):\n",
    "    CPh[i].set_pos(Dmc_list[i], CPhxe(epsilon_list[i])*A_list[i].x + CPhze(epsilon_list[i])*A_list[i].z )\n",
    "\n",
    "    nhat_hilfs = CPh[i].pos_from (P0).normalize()\n",
    "    CPhs[i].set_pos(CPh[i], sm.Abs(lang)*nhat_hilfs)\n",
    "\n",
    "    CPh_list.append([me.dot(CPh[i].pos_from(P0), uv) for uv in (N.x, N.z)])\n",
    "    CPhs_list.append([me.dot(CPhs[i].pos_from(P0), uv) for uv in (N.x, N.z)])\n",
    "\n",
    "def distanzCPhiCPhj(i, epsilon):\n",
    "    CPh[i].set_pos(Dmc_list[i], CPhxe(epsilon)*A_list[i].x + CPhze(epsilon)*A_list[i].z)\n",
    "    return RW - CPh[i].pos_from(P0).magnitude() # if this is < 0., there is a collision.\n",
    "\n",
    "# this function will be minimized during integration to get the distance from ellipse to wall\n",
    "abstand = [distanzCPhiCPhj(i, epsilon_list[i]) for i in range(n)]\n",
    "min_distanzCPhiCPhj_lam = [njit(sm.lambdify([epsilon_list[i]] + q_list + x_list + z_list + [a, b, RW], abstand[i], cse=True)) for i in range(n)]\n",
    "\n",
    "\n",
    "# needed for plotting initial conditions only.\n",
    "CPh_list_lam  = [sm.lambdify(q_list + x_list + z_list + [a, b, RW,  epsilon_list[i]], CPh_list[i], cse=True) for i in range(n)]\n",
    "CPhs_list_lam = [sm.lambdify(q_list + x_list + z_list + [a, b, RW, lang, epsilon_list[i]],  CPhs_list[i], cse=True) for i in range(n)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Find the potential collision points of any two ellipses**\\\n",
    "I assume, that no more than two ellipses collide at exactly thew same time. Seems a reasonable assumption to me.\\\n",
    "In order to get the potential contact points, say, $CPhe_i, CPhe_j$, I simply try to find the minimum distance between any two points on the circumferences on the respective ellipses.\\\n",
    "I do this in two ways:\n",
    "- minimize $|{}^{CPhe_i} \\bar r^{CPhe_j}(\\epsilon_i, \\epsilon_j)|$ directly, using scipy's minimize function\n",
    "- calculate $\\dfrac{d}{d\\epsilon_i} |{}^{CPhe_i} \\bar r^{CPhe_j}(\\epsilon_i, \\epsilon_j)|$ and $\\dfrac{d}{d\\epsilon_j} |{}^{CPhe_i} \\bar r^{CPhe_j}(\\epsilon_i, \\epsilon_j)|$, and solve for $\\epsilon_i, \\epsilon_j$. This is only a sufficient condition for a minimum, it could also give a maximum. With the right initial guess, scipy's **root** function should give the minimum. During integration, this seems to be faster than the first option. \n",
    "\n",
    "\n",
    "If $| {}^{CPhe_i} \\bar r^{CPhe_j} | \\approx 0.$ it becomes numerically critical to get the direction ${}^{CPhe_i} \\bar r^{CPhe_j}$. So, I allow the possibility to fix the direction before the distance becomes too small. See also the comment in the numerical integration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def CPhxe(epsilon):\n",
    "    return a * sm.cos(epsilon)\n",
    "def CPhze(epsilon):\n",
    "    return b * sm.sin(epsilon)\n",
    "\n",
    "def vorzeichen(i, j, epsilon1, epsilon2):\n",
    "# the idea is this: I calculate the triangle abc, with a := Dmc_i - CPh_i, b := CPhe_j, and get c using\n",
    "# the cosine theorem: c^2 = a^2 + b^2 - 2 * a * b * cos(gamma)\n",
    "# while c < Dmc_i.pos_from(Dmc_j).magnitude, the ellipses are separated.\n",
    "    Pi = me.Point('Pi')\n",
    "    Pj = me.Point('Pj')\n",
    "    Pi.set_pos(Dmc_list[i], CPhxe(epsilon1)*A_list[i].x + CPhze(epsilon1)*A_list[i].z)\n",
    "    Pj.set_pos(Dmc_list[j], CPhxe(epsilon2)*A_list[j].x + CPhze(epsilon2)*A_list[j].z)\n",
    "\n",
    "    rr = Dmc_list[i].pos_from(Dmc_list[j]).magnitude()\n",
    "    r1 = Dmc_list[i].pos_from(Pi)\n",
    "    r2 = Dmc_list[j].pos_from(Pj)\n",
    "    gamma_cos = (me.dot(r1.normalize(), r2.normalize()))\n",
    "    r1 = r1.magnitude()\n",
    "    r2 = r2.magnitude() \n",
    "    r3 = sm.sqrt(r1**2 + r2**2 - 2. * r1 * r2 * gamma_cos)\n",
    "    hilfs1 = rr - r3\n",
    "    hilfs2 = sm.Piecewise((-1., hilfs1 <= 0.), (1., hilfs1 > 0.))\n",
    "    return hilfs2  # -1., if the ellipses have penetrated\n",
    "\n",
    "\n",
    "def distanzCPheiCPhej(i, j, epsilon1, epsilon2):\n",
    "    P1, P2 = sm.symbols('P1, P2', cls=me.Point)\n",
    "    P1.set_pos(Dmc_list[i], CPhxe(epsilon1)*A_list[i].x + CPhze(epsilon1)*A_list[i].z)\n",
    "    P2.set_pos(Dmc_list[j], CPhxe(epsilon2)*A_list[j].x + CPhze(epsilon2)*A_list[j].z)\n",
    "    vektor = P2.pos_from(P1)\n",
    "    return vektor.magnitude() * vorzeichen(i, j, epsilon1, epsilon2)    \n",
    "\n",
    "def richtungCPheiCPhej(i, j, epsilon1, epsilon2):\n",
    "    P1, P2 = sm.symbols('P1, P2', cls=me.Point)\n",
    "    P1.set_pos(Dmc_list[i], CPhxe(epsilon1)*A_list[i].x + CPhze(epsilon1)*A_list[i].z)\n",
    "    P2.set_pos(Dmc_list[j], CPhxe(epsilon2)*A_list[j].x + CPhze(epsilon2)*A_list[j].z)\n",
    "    vektor = P2.pos_from(P1).normalize()\n",
    "    return [me.dot(vektor, uv) for uv in (N.x, N.z) ]# ((A_list[i].x + A_list[j].x), (A_list[i].z + A_list[j].z))]\n",
    "\n",
    "# This function will be minimized numerically during integratrion to get the distance from ellipse_i to ellipse_j\n",
    "min_distanzCPheiCPhej_lam  = []\n",
    "min_distanzCPheiCPhej_lam1 = []\n",
    "\n",
    "zaehler = -1\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    zaehler += 1\n",
    "    abstand = distanzCPheiCPhej(i, j, epsilone_list[zaehler][0], epsilone_list[zaehler][1])\n",
    "    abstanddeidej = [abstand.diff(epsilone_list[zaehler][0]), abstand.diff(epsilone_list[zaehler][1])]\n",
    "    \n",
    "    min_distanzCPheiCPhej_lam.append((sm.lambdify([epsilone_list[zaehler][0], epsilone_list[zaehler][1]] + q_list + x_list + z_list + [a, b], abstand, cse=True)))\n",
    "    min_distanzCPheiCPhej_lam1.append((sm.lambdify([epsilone_list[zaehler][0], epsilone_list[zaehler][1]] + q_list + x_list + z_list + [a, b], abstanddeidej, cse=True)))\n",
    "\n",
    "\n",
    "richtung_lam = []\n",
    "zaehler = -1\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    zaehler +=1\n",
    "    richtungij = richtungCPheiCPhej(i, j, epsilone_list[zaehler][0], epsilone_list[zaehler][1])\n",
    "    richtung_lam.append((sm.lambdify([epsilone_list[zaehler][0], epsilone_list[zaehler][1]] + q_list + x_list + z_list + [a, b], richtungij, cse=True)))\n",
    "\n",
    "\n",
    "# just needed for plotting initial situation\n",
    "epsilonei = sm.symbols('epsilonei')\n",
    "CPhe = list(sm.symbols('CPhe' + str(i), cls=me.Point) for i in range(n))\n",
    "\n",
    "CPhe_list = []\n",
    "for i in range(n):\n",
    "    CPhe[i].set_pos(Dmc_list[i], CPhxe(epsilonei)*A_list[i].x + CPhze(epsilonei)*A_list[i].z )\n",
    "    CPhe_list.append([me.dot(CPhe[i].pos_from(P0), uv) for uv in (N.x, N.z)])\n",
    "CPhe_list_lam = sm.lambdify(q_list + x_list + z_list + [a, b, epsilonei], CPhe_list, cse=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Collision force between ellipse and the wall**\\\n",
    "From the ellipse's point of view, the wall is concave. Hence I use $R_2 = -R_W. \\space$ I do not know, whether this is 'covered' by the theory.\\\n",
    "I add a speed dependent frictional force between ellipse and wall. It acts in the line of the tangent at the collision point, directed opposite to the speed component in that direction. It is proportional to the magnitude of this speed component and to the magnitude of the impact force. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# needed in the functions below. \n",
    "rhodt_dict = {sm.Derivative(i, t): j for i, j in zip(q_list + x_list + z_list, u_list + ux_list + uz_list)}\n",
    "\n",
    "def HC_wall(i, N, A, Dmc, CPh, epsilon, a, b, RW, nue, nuw, EYe, EYw, ctau, reibung, rhodtwall, l):\n",
    "    \n",
    "# curvature of the ellipse at the point (a*cos(epsilon) / b*sin(epsilon)) from the internet\n",
    "    kappa1 = (a * b) / (sm.sqrt((a*sm.sin(epsilon))**2 + (b*sm.cos(epsilon))**2))**3\n",
    "\n",
    "    R1 = 1. / kappa1\n",
    "    R2 = -RW\n",
    "    sigmae = (1. - nue**2) / EYe\n",
    "    sigmaw = (1. - nuw**2) / EYw\n",
    "    k0 = 4./3. * 1./(sigmae + sigmaw) * sm.sqrt(R1*R2 / (R1 + R2))\n",
    "    nhat = CPh.pos_from(P0).normalize()#.subs({epsilonwi: epsilon})\n",
    "    rhodt = me.dot(CPh.pos_from(P0).diff(t, N), nhat).subs(rhodt_dict)#.subs({epsilonwi: epsilon})\n",
    "    rho   = sm.Abs(l) * sm.Heaviside(-l, sm.S(0))\n",
    "#    print('rhodt DS', me.find_dynamicsymbols(rhodt))\n",
    "#    print('rhodt FS', rhodt.free_symbols)\n",
    "    fHC_betrag = k0 * rho**(3/2) * (1. + 3./2. * (1. - ctau) * (rhodt) / sm.Abs(rhodtwall)) \n",
    "    fHC = fHC_betrag * (-nhat) * sm.Heaviside(-l, sm.S(0))  # force is acting on CPh, hence the minus sign.\n",
    "\n",
    "#    print('fHC DS', me.find_dynamicsymbols(fHC, reference_frame=N))\n",
    "#    print('fHC FS', fHC.free_symbols(reference_frame=N))\n",
    "\n",
    "# friction force on CPh\n",
    "    that = nhat.cross(A_list[i].y)\n",
    "    vCPh = (me.dot(CPh.pos_from(P0).diff(t, N), that)).subs(rhodt_dict)\n",
    "    F_friction = fHC.magnitude() * reibung * vCPh * (-that) * sm.Heaviside(-l, sm.S(0))\n",
    "\n",
    "#    print('F_friction DS', me.find_dynamicsymbols(F_friction, reference_frame=N))\n",
    "#    print('F_friction FS', F_friction.free_symbols(reference_frame=N))\n",
    "\n",
    "    if friktion == True:\n",
    "        return fHC + F_friction\n",
    "    else:\n",
    "        return fHC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Collision between any two ellipses**\\\n",
    "I use Hunt-Crossley's method to model it.\\\n",
    "I add a speed dependent frictional force between the ellipses. It acts in the line of the tangent at the collision point, directed opposite to the speed component in that direction. It is proportional to the magnitude of this speed component and to the magnitude of the impact force. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def HC_ellipse(i, j, epsiloni, epsilonj, l, rhodtellipse, richtungx, richtungz):\n",
    "# this calculates the force of ellipse_i on ellipse_j during their collision.\n",
    "# i, j are the respective ellipses\n",
    "# epsilone list of angles\n",
    "# l is the distance between ellipse_i and ellipse_j, negative during penetration\n",
    "# rhodtellipse is the collision speed right before impact.\n",
    "    \n",
    "# curvature of the ellipse at the point (a*cos(delta) / b*sin(delta)) from the internet\n",
    "    kappa1 = (a * b) / (sm.sqrt((a*sm.sin(epsiloni))**2 + (b*sm.cos(epsiloni))**2))**3\n",
    "    kappa2 = (a * b) / (sm.sqrt((a*sm.sin(epsilonj))**2 + (b*sm.cos(epsilonj))**2))**3\n",
    "\n",
    "    R1 = 1. / kappa1\n",
    "    R2 = 1. / kappa2\n",
    "    sigmae = (1. - nue**2) / EYe\n",
    "    k0 = 4./3. * 1./(sigmae + sigmae) * sm.sqrt(R1*R2 / (R1 + R2))\n",
    "\n",
    "    P1, P2 = sm.symbols('P1, P2', cls=me.Point)\n",
    "    P1.set_pos(Dmc_list[i], CPhxe(epsiloni)*A_list[i].x + CPhze(epsiloni)*A_list[i].z)\n",
    "    P2.set_pos(Dmc_list[j], CPhxe(epsilonj)*A_list[j].x + CPhze(epsilonj)*A_list[j].z)\n",
    "    \n",
    "    \n",
    "    vei = P1.pos_from(P0).diff(t, N)\n",
    "    vej = P2.pos_from(P0).diff(t, N)\n",
    "# only the speed in direction of the collision is important here\n",
    "    richtung = richtungx*N.x + richtungz*N.z #richtungx*(A_list[i].x + A_list[j].x) + richtungz * (A_list[i].z + A_list[j].z)           # DURING a collision, this points basically from ellipse_i to ellipse_j\n",
    "    rhodt = me.dot(vei - vej, richtung).subs(rhodt_dict)\n",
    "    rho   = sm.Abs(l) * sm.Heaviside(-l, sm.S(0))\n",
    "#    print('rhodt DS', me.find_dynamicsymbols(rhodt))\n",
    "#    print('rhodt FS', rhodt.free_symbols)\n",
    "\n",
    "    fHC_betrag = k0 * rho**(3/2) * (1. + 3./2. * (1. - ctau) * (rhodt) / sm.Abs(rhodtellipse))\n",
    "    fHC = fHC_betrag * richtung * sm.Heaviside(-l, sm.S(0))\n",
    "\n",
    "# friction force on CPhej\n",
    "    that = richtung.cross(A_list[i].y)\n",
    "    vCPhej = (me.dot(vei - vej, that)).subs(rhodt_dict)\n",
    "    F_friction = fHC.magnitude() * (-reibung) * vCPhej * (-that) * sm.Heaviside(-l, sm.S(0))\n",
    "\n",
    "#    print('F_friction DS', me.find_dynamicsymbols(F_friction, reference_frame=N))\n",
    "#    print('F_friction FS', F_friction.free_symbols(reference_frame=N))\n",
    "\n",
    "#    print('fHC DS', me.find_dynamicsymbols(fHC, reference_frame=N))\n",
    "#    print('fHC FS', fHC.free_symbols(reference_frame=N))\n",
    "\n",
    "    if friktion == True:\n",
    "        return fHC + F_friction\n",
    "    else:\n",
    "        return fHC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the **force** acting on the system\n",
    "\n",
    "i do not do this the most economical way, as I do not consider that $\\bar f_{ellipse_i \\space on \\space ellipse_j} = -\\bar f_{ellipse_j \\space on \\space ellipse_i}$\\\n",
    "But it makes the 'book keeping' easier. I assume (I do not know!) that cse = True removes most of this inefficiency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#def HC_wall(i, N, A, Dmc, CPh, qepsilon, a, b, RW, nue, nuw, EYe, EYw, ctau, rhodtwall, l))\n",
    "CPhej = list(sm.symbols(f'CPhej:{n}', cls =me.Point))\n",
    "FL_wall = []\n",
    "for i in range(n):\n",
    "    FL_wall.append((CPh[i], HC_wall(i, N, A_list[i], Dmc_list[i], CPh[i], epsilon_list[i], a, b, RW, nue, nuw, EYe, EYw, ctau, reibung, rhodtwall[i], l_list[i])))\n",
    "    \n",
    "#def HC_ellipse(i, j, epsiloni, epsilonj, l, rhodtellipse, richtungx, richtungz):\n",
    "FL_ellipse = []\n",
    "zaehler = -1\n",
    "\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    zaehler += 1\n",
    "    CPhej[j].set_pos(Dmc_list[j], CPhxe(epsilone_list[zaehler][1])*A_list[j].x + CPhze(epsilone_list[zaehler][1])*A_list[j].z )\n",
    "    CPhej[j].set_vel(N, (CPhej[j].pos_from(P0).diff(t, N)))\n",
    "    FL_ellipse.append((CPhej[j], HC_ellipse(i, j, epsilone_list[zaehler][0], epsilone_list[zaehler][1], le_list[zaehler], rhodtmax[zaehler], richtung_list[zaehler][0], richtung_list[zaehler][1])))\n",
    "\n",
    "if elli_ghost == True:\n",
    "    FL = FL_wall\n",
    "else:\n",
    "    FL = FL_wall + FL_ellipse\n",
    "    \n",
    "print('FL FS', set.union(*[FL[j][1].free_symbols(reference_frame=N) for j in range(len(FL))]))\n",
    "print('FL DS', set.union(*[me.find_dynamicsymbols(FL[j][1], reference_frame=N) for j in range(len(FL))]))"
   ]
  },
  {
   "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 = q_list + x_list + z_list\n",
    "u_ind = u_list + ux_list + uz_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(f'MM contains {sm.count_ops(MM):,} operations, {sm.count_ops(sm.cse(MM)):,} after cse', '\\n')\n",
    "\n",
    "force = KM.forcing_full\n",
    "print('force DS', me.find_dynamicsymbols(force))\n",
    "print('force free symbols', force.free_symbols)\n",
    "print(f'force contains {sm.count_ops(force):,} operations, {sm.count_ops(sm.cse(force)):,} after cse')"
   ]
  },
  {
   "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 $ellipse_j$ and $ellipse_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 $ellipse_i$ and the wall.\n",
    "- *Po_pos*: Holds the locations of each observer.\n",
    "- *Dmc_distanz*: Holds the distance between the centers of $ellipse_j$ and $ellipse_i$, $0 \\le i, j \\le n-1$, $i \\neq j$. Needed for initial conditions.\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(q_list + x_list + z_list, u_list + ux_list + uz_list)}\n",
    "\n",
    "rhodtwall_list = []\n",
    "for i in range(n):\n",
    "    richtung = CPh[i].pos_from(P0).normalize()\n",
    "    rhodt = me.dot(CPh[i].pos_from(P0).diff(t, N), richtung).subs(derivative_dict)\n",
    "    rhodtwall_list.append(rhodt)\n",
    "print('rhodtwall_list DS:', set.union(*[me.find_dynamicsymbols(rhodtwall_list[k]) for k in range(len(rhodtwall_list))]))\n",
    "print('rhodtwall_list free symbols:', set.union(*[rhodtwall_list[k].free_symbols for k in range(len(rhodtwall_list))]), '\\n')\n",
    "\n",
    "# collision speed between ellipse_i and ellipse_j right brfore impact.\n",
    "# it is the speed with which ellipse_i hits ellipse_j\n",
    "zaehler = -1\n",
    "rhodtmax_list = []\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    zaehler += 1\n",
    "    vei = CPhe[i].pos_from(P0).diff(t, N).subs(derivative_dict).subs({epsilonei: epsilone_list[zaehler][0]})\n",
    "    vej = CPhe[j].pos_from(P0).diff(t, N).subs(derivative_dict).subs({epsilonei: epsilone_list[zaehler][1]})\n",
    "# only the speed in direction of the collision is important here\n",
    "    richtung = richtung_list[zaehler][0] * N.x + richtung_list[zaehler][1] * N.z\n",
    "    rhodtmax_list.append(me.dot(vei - vej, richtung).subs(derivative_dict))\n",
    "print('rhodtmax_list DS:', set.union(*[me.find_dynamicsymbols(rhodtmax_list[k]) for k in range(len(rhodtmax_list))]))\n",
    "print('rhodtmax_list free symbols:', set.union(*[rhodtmax_list[k].free_symbols for k in range(len(rhodtmax_list))]), '\\n')\n",
    "\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",
    "kin_energie = sum([body.kinetic_energy(N) for body in BODY])\n",
    "print('kinetic energy FS', kin_energie.free_symbols)\n",
    "print('kinetic energy DS', me.find_dynamicsymbols(kin_energie, reference_frame=N), '\\n')\n",
    "\n",
    "spring_energie = 0.\n",
    "# 1. collisions of discs\n",
    "if elli_ghost == False:\n",
    "    zaehler = -1\n",
    "    for i, j in permutations(range(n), r=2):\n",
    "        zaehler += 1\n",
    "        kappa1 = (a * b) / (sm.sqrt((a*sm.sin(epsilone_list[zaehler][0]))**2 + (b*sm.cos(epsilone_list[zaehler][0]))**2))**3\n",
    "        kappa2 = (a * b) / (sm.sqrt((a*sm.sin(epsilone_list[zaehler][1]))**2 + (b*sm.cos(epsilone_list[zaehler][1]))**2))**3\n",
    "\n",
    "        R1 = 1. / kappa1\n",
    "        R2 = 1. / kappa2\n",
    "        sigmae = (1. - nue**2) / EYe\n",
    "        k0 = 4./3. * 1./(sigmae + sigmae) * sm.sqrt(R1*R2 / (R1 + R2))\n",
    "\n",
    "        rho = sm.Abs(le_list[zaehler])\n",
    "        rho = rho**(5/2)\n",
    "        spring_energie += (k0 * 2./5. * rho * sm.Heaviside(-le_list[zaehler], 0.))  * 0.5\n",
    "else:\n",
    "    pass   \n",
    "      \n",
    "# 2. Collision of discs with the wall   \n",
    "for i in range(n): \n",
    "    kappa1 = (a * b) / (sm.sqrt((a*sm.sin(epsilon_list[i]))**2 + (b*sm.cos(epsilon_list[i]))**2))**3\n",
    "    R1 = 1. / kappa1\n",
    "    R2 = -RW\n",
    "    sigmae = (1. - nue**2) / EYe\n",
    "    sigmaw = (1. - nuw**2) / EYw\n",
    "    k0 = 4./3. * 1./(sigmae + sigmaw) * sm.sqrt(R1*R2 / (R1 + R2))\n",
    "    rho = sm.Abs(l_list[i])**(5/2)\n",
    "    spring_energie += k0 * 2./5. * rho * sm.Heaviside(-l_list[i], 0.)\n",
    "print('spring energy FS', spring_energie.free_symbols)\n",
    "print('spring energy DS', me.find_dynamicsymbols(spring_energie, reference_frame=N), '\\n')\n",
    "\n",
    "# Needed only for the initial conditions\n",
    "Dmc_distanz = [Dmc_list[i].pos_from(Dmc_list[j]).magnitude() for i, j in permutations(range(n), r=2)]\n",
    "print('Dmc_distanz DS:', set.union(*[me.find_dynamicsymbols(Dmc_distanz[k]) for k in range(len(Dmc_distanz))]))\n",
    "print('Dmc_distanz free symbols:', set.union(*[Dmc_distanz[k].free_symbols for k in range(len(Dmc_distanz))]), '\\n')"
   ]
  },
  {
   "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, iYY, a, b, RW, nue, nuw, EYe, EYw, ctau, reibung] + l_list + le_list + epsilon_list + epsilone_list + alpha_list + beta_list + rhodtwall + rhodtmax + richtung_list\n",
    "pL1 = [m0, mo, iYY, a, b, RW, nue, nuw, EYe, EYw, ctau, reibung] + l_list + le_list + epsilon_list + epsilone_list\n",
    "pL2 = [m0, mo, iYY, a, b, RW, nue, nuw, EYe, EYw, ctau, reibung] + alpha_list + beta_list\n",
    "pL3 = [m0, mo, iYY, a, b, RW, nue, nuw, EYe, EYw, ctau, reibung] + l_list + epsilon_list + le_list + epsilone_list\n",
    "pL4 = [m0, mo, iYY, a, b, RW, nue, nuw, EYe, EYw, ctau, reibung] + l_list + epsilon_list\n",
    "\n",
    "MM_lam    = sm.lambdify(qL + pL, MM, cse=True)\n",
    "force_lam = sm.lambdify(qL + pL, force, cse = True) \n",
    "\n",
    "rhodtwall_lam     = (sm.lambdify(qL + pL, rhodtwall_list, cse=True))\n",
    "rhodtmax_list_lam = (sm.lambdify(qL + pL, rhodtmax_list, cse=True))\n",
    "\n",
    "Po_pos_lam  = sm.lambdify(q_list + x_list + z_list + alpha_list + beta_list + [a, b], Po_pos, cse=True)\n",
    "\n",
    "kin_lam    = sm.lambdify(qL + pL2, kin_energie, cse=True)\n",
    "spring_lam = sm.lambdify(qL + pL3, spring_energie, cse=True)\n",
    "spring_lam1 = sm.lambdify(qL + pL4, spring_energie, cse=True)\n",
    "\n",
    "Dmc_distanz_lam = sm.lambdify(x_list + z_list, Dmc_distanz, cse=True)\n",
    "\n",
    "print(f'it took {time.time()-start0:.3f} sec to establish Kanes equation for {n} bodies')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Set initial conditions and parameters**\n",
    "\n",
    "1.\\\n",
    "The ellipses are randomly placed within the wall, 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",
    "Doing this, I consider the ellipses to be discs, with $r_0 = max(a_1, b_1)$. This is 'on the safe side', so to speak.\\\n",
    "2.\\\n",
    "Assign random generalized speeds to each ellipse, 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "start1 = time.time()\n",
    "\n",
    "# Input variables\n",
    "#-------------------------------------------------------------------\n",
    "m01      = 1.                                                                          # mass of a pendulum\n",
    "mo1      = 1.                                                                          # mass of the observer\n",
    "a1       = 2.                                                                           # semi axis of ellipses\n",
    "b1       = 1.                                                                           # dto.\n",
    "RW1      = 7.                                                                           # radius of wall\n",
    "nue1     = 0.28                                                                         # Poisson's number of ellipse\n",
    "nuw1     =  0.28                                                                        # dto. for wall\n",
    "EYe1     = 1.e7                                                                         # Young's modulus for ellipse\n",
    "EYw1     = 1.e7                                                                         # dto. for wall\n",
    "ctau1    = 0.95                                                                          # experimerntal constant needed for H-C's method \n",
    "reibung1 = 0.5                                                                           # friction between ellipses / between ellipses and the wall\n",
    "\n",
    "np.random.seed(123456789)\n",
    "\n",
    "q_list1 = [*np.random.choice(np.linspace(-np.pi, np.pi, 100), size=n)]              # initial angle of the ellipses\n",
    "alpha_list1 = [0.5 for _ in range(len(alpha_list))]                                 # location of observer\n",
    "beta_list1  = [0.5 for _ in range(len(beta_list))]                                  # dto.\n",
    "\n",
    "ux_list1 = list(2. * np.random.choice(np.linspace(-5., 5., 100), size=n))                # initial speed of center of the ellipse in x direction\n",
    "uz_list1 = list(2. * np.random.choice(np.linspace(-5., 5., 100), size=n))                # initial speed of center of the ellipse in Z direction\n",
    "u_list1  = list(1. * np.random.choice(np.linspace(-5., 5., 100), size=len(q_list)))    # initial rotationmal speed\n",
    "\n",
    "rhodtwall1 = [1. + k for k in range(n)]                                             # initial walues of no consequence, anything will do.\n",
    "rhodtmax1  = [1. for _ in range(n*(n-1))]                                           # dto.\n",
    "richtung_list1 = [(1., 0.) for _ in range(n*(n-1))]\n",
    "\n",
    "solver_dict = {1: 'Powell', 2: 'Nelder-Mead', 3: 'L-BFGS-B', 4: 'BFGS', 5: 'COBYLA', 6:'SLSQP', 7: 'TNC'}\n",
    "methode = solver_dict[1]\n",
    "option_dict = None#{'xtol': 1.e-6, 'ftol': 1.e-6}\n",
    "\n",
    "#-------------------------------------------------------------------\n",
    "iYY1  = 0.25 * m01 * (a1**2 + b1**2)                                                # Moment of inertia of the ellipse.\n",
    "\n",
    "# 1. randomly place the ellipses as described above\n",
    "zaehler = 0\n",
    "r01     = max(a1, b1)\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",
    "else:\n",
    "    raise Exception(' no good location for ellipses found, make RW1 larger, or try again.')\n",
    "\n",
    "# make a plot of the initial situation\n",
    "fig, ax = plt.subplots(figsize=(10, 10))\n",
    "ax.set_aspect('equal')\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",
    "ax.plot(0., 0., marker='x', color='black')\n",
    "\n",
    "# This is to asign colors of 'plasma' to the ellipses.\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",
    "# just intermediate storage\n",
    "l_hilfs       = []\n",
    "epsilon_hilfs = []\n",
    "TEST1         = []\n",
    "\n",
    "# Just to adapt the arguments so they fit minimize\n",
    "def funcw(x0, args):\n",
    "    return min_distanzCPhiCPhj_lam[kk](*x0, *args)\n",
    "\n",
    "def funce(x0, args):\n",
    "    return min_distanzCPheiCPhej_lam[zaehler](*x0, *args)\n",
    "\n",
    "\n",
    "# get the closest distance between ellipse and wall\n",
    "for kk in range(n):\n",
    "    x0 = 1.\n",
    "    args1 = q_list1 + x_list1 + z_list1 + [a1, b1, RW1]\n",
    "    epsilon_min = minimize(funcw, x0, args1, method = methode, options = option_dict)\n",
    "    min_eps = (epsilon_min.x % (2.*np.pi))\n",
    "    ll_min = funcw(min_eps, args1)\n",
    "    TEST1.append((ll_min, kk, min_eps[0]))\n",
    "    l_hilfs.append(ll_min)\n",
    "    epsilon_hilfs.append(min_eps[0])\n",
    "\n",
    "\n",
    "    elli = patches.Ellipse((x_list1[kk], z_list1[kk]), width=2.*a1, height=2.*b1, angle=-np.rad2deg(q_list1[kk]), zorder=1, fill=True, color=farben[kk], ec='black')\n",
    "    ax.add_patch(elli)\n",
    "    weite = 10.\n",
    "    ax.plot(x_list1[kk], z_list1[kk], color='yellow', marker='o', markersize=2)\n",
    "    ax.plot(Po_pos_lam(*q_list1, *x_list1, *z_list1, *alpha_list1, *beta_list1, a1, b1)[kk][0], \n",
    "            Po_pos_lam(*q_list1, *x_list1, *z_list1, *alpha_list1, *beta_list1, a1, b1)[kk][1], color='white', marker='o', markersize=5)\n",
    "    ax.set_title('possible contact points' )\n",
    "\n",
    "for kk in range(len(TEST1)):\n",
    "    koerper1 = TEST1[kk][1]\n",
    "    epsilon1 = TEST1[kk][2]\n",
    "    laenge = TEST1[kk][0]\n",
    "    x11 =  CPh_list_lam[koerper1](*q_list1, *x_list1, *z_list1, a1, b1, RW1, epsilon1)[0]\n",
    "    x12 = CPhs_list_lam[koerper1](*q_list1, *x_list1, *z_list1, a1, b1, RW1, laenge, epsilon1)[0]\n",
    "    z11 =  CPh_list_lam[koerper1](*q_list1, *x_list1, *z_list1, a1, b1, RW1, epsilon1)[1]\n",
    "    z12 = CPhs_list_lam[koerper1](*q_list1, *x_list1, *z_list1, a1, b1, RW1, laenge, epsilon1)[1]\n",
    "    \n",
    "    ax.plot([x11, x12], [z11, z12], color=farben[koerper1], linestyle='-')\n",
    "\n",
    "    \n",
    "# find possible collision points between the ellipses\n",
    "zaehler        = -1\n",
    "TEST3          = []\n",
    "le_hilfs       = []\n",
    "epsilone_hilfs = []\n",
    "\n",
    "for i, j in permutations(range(n), r=2):\n",
    "    zaehler += 1\n",
    "    x0 = (1., 1.)                                                       # initial guess\n",
    "    args1 = q_list1 + x_list1 + z_list1 + [a1, b1]\n",
    "    epsilon_min = minimize(funce, x0, args1, method = methode, options = option_dict)\n",
    "    min_eps = epsilon_min.x % (2.*np.pi)\n",
    "    ll_min = funce(min_eps, args1)\n",
    "    TEST3.append([ll_min, i, j, min_eps[0], min_eps[1], zaehler])\n",
    "    le_hilfs.append(ll_min)\n",
    "    epsilone_hilfs.append((min_eps[0], min_eps[1]))\n",
    "\n",
    "for kk in range(len(TEST3)):\n",
    "    koerper1 = TEST3[kk][1]\n",
    "    koerper2 = TEST3[kk][2]\n",
    "    epsilon1 = TEST3[kk][3]\n",
    "    epsilon2 = TEST3[kk][4]\n",
    "    x11 = CPhe_list_lam(*q_list1, *x_list1, *z_list1, a1, b1, epsilon1)[koerper1][0]\n",
    "    x12 = CPhe_list_lam(*q_list1, *x_list1, *z_list1, a1, b1, epsilon2)[koerper2][0]\n",
    "    z11 = CPhe_list_lam(*q_list1, *x_list1, *z_list1, a1, b1, epsilon1)[koerper1][1]\n",
    "    z12 = CPhe_list_lam(*q_list1, *x_list1, *z_list1, a1, b1, epsilon2)[koerper2][1]\n",
    "    \n",
    "    ax.plot([x11, x12], [z11, z12], color=farben[koerper1], linestyle='dotted')\n",
    "\n",
    "print(f'it took {(time.time() - start1):.3f} sec to get initial conditions', '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "**Numerical integration**\\\n",
    "Slow, as the distances between the ellipse and between ellipses and the wall must be found numerically.\\\n",
    "With *option1 = True* the 'direct method' to minimize $|{}^{CPhe_i} \\bar r^{CPhe_j}(\\epsilon_i, \\epsilon_j)|$ is used, else $ \\nabla{|{}^{CPhe_i} \\bar r^{CPhe_j}(\\epsilon_i, \\epsilon_j)|} = 0$ is used.\n",
    "\n",
    "I collect the values $\\epsilon_{min_i}, \\epsilon_{min_j}$ and $l_{min} = |{}^{CPhe_i} \\bar r^{CPhe_j}(\\epsilon_{min_i}, \\epsilon_{min_j})|$ during integration, for later use.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#=============================\n",
    "option1 = False\n",
    "\n",
    "intervall = 7.5                      \n",
    "schritte = 600  \n",
    "max_step = 0.0001\n",
    "#============================\n",
    "\n",
    "# Just to adapt the arguments so they fit minimize\n",
    "def funcw1(x0, args):\n",
    "    kk = args[-1]\n",
    "    args1 = [args[i] for i in range(len(args)-1)]\n",
    "    return min_distanzCPhiCPhj_lam[kk](*x0, *args1)\n",
    "\n",
    "def funce1(x0, args):\n",
    "    zaehler = args[-1]\n",
    "    args1 = [args[i] for i in range(len(args)-1)]\n",
    "    return min_distanzCPheiCPhej_lam[zaehler](*x0, *args1)\n",
    "\n",
    "def funce1didj(x0, args):\n",
    "    zaehler = args[-1]\n",
    "    args1 = [args[i] for i in range(len(args)-1)]\n",
    "    return min_distanzCPheiCPhej_lam1[zaehler](*x0, *args1)\n",
    "\n",
    "l_list1 = l_hilfs\n",
    "epsilon_list1 = epsilon_hilfs\n",
    "le_list1 = le_hilfs\n",
    "epsilone_list1 = epsilone_hilfs\n",
    "\n",
    "y0       = q_list1 + x_list1 + z_list1 + u_list1 + ux_list1 + uz_list1\n",
    "pL_vals  = [m01, mo1, iYY1, a1, b1, RW1, nue1, nuw1, EYe1, EYw1, ctau1, reibung1] + l_list1 + le_list1 + epsilon_list1 + epsilone_list1 + alpha_list1 + beta_list1 + rhodtwall1 + rhodtmax1 + richtung_list1\n",
    "pL1_vals = [m01, mo1, iYY1, a1, b1, RW1, nue1, nuw1, EYe1, EYw1, ctau1, reibung1] + l_list1 + le_list1 + epsilon_list1 + epsilone_list1\n",
    "pL2_vals = [m01, mo1, iYY1, a1, b1, RW1, nue1, nuw1, EYe1, EYw1, ctau1, reibung1] + alpha_list1 + beta_list1\n",
    "\n",
    "print(pL_vals)\n",
    "\n",
    "start1 = time.time()\n",
    "times = np.linspace(0, intervall, int(schritte*intervall))\n",
    "\n",
    "LAENGEe  = []\n",
    "zeit     = []\n",
    "ellipsee = []\n",
    "\n",
    "def gradient(t, y, args):\n",
    "    \n",
    "# find the distance of the ellipses fro9m the wall. and also find rhodtwall\n",
    "    for kk in range(n):\n",
    "        x0 = args[12 + n + n*(n-1) + kk]\n",
    "\n",
    "        args1 = [y[i] for i in range(3*n)] + [a1, b1, RW1] + [kk]\n",
    "        epsilon_min = minimize(funcw1, x0, args1, method=methode, options=option_dict)\n",
    "        min_eps = epsilon_min.x % (2.*np.pi)\n",
    "        ll_min = funcw1(min_eps, args1)\n",
    "        args[12 + kk] = ll_min\n",
    "        args[12 + n + n*(n-1) + kk] = min_eps[0]\n",
    " \n",
    "        if 0. <= ll_min <= 0.1:\n",
    "            args[12 + 4*n + 2*n*(n-1) + kk] = rhodtwall_lam(*y, *args)[kk]\n",
    "    \n",
    "\n",
    "# find possible collision points between the ellipses, and also find rhodtmax\n",
    "    if elli_ghost == False:\n",
    "        zaehler        = -1\n",
    "        LAENGEe1 = []\n",
    "        ellipsee1 = []\n",
    "#        for i, j in permutations(range(n), r=2):\n",
    "        for _ in range(n*(n-1)):\n",
    "\n",
    "            zaehler += 1\n",
    "            x0 = args[12 + n + n*(n-1) + n + zaehler]                                          # initial guess\n",
    "            args1 = [y[ij] for ij in range(3*n)] + [a1, b1] + [zaehler]   \n",
    "\n",
    "            if option1 == True:          \n",
    "                epsilon_min = minimize(funce1, x0, args1, method=methode, options=option_dict)\n",
    "            else:\n",
    "                epsilon_min = root(funce1didj, x0, args1)#, method='broyden1')\n",
    "\n",
    "            min_eps = epsilon_min.x % (2.*np.pi)\n",
    "            ll_min = funce1(min_eps, args1)\n",
    "\n",
    "            LAENGEe1.append(ll_min)\n",
    "            ellipsee1.append((min_eps[0], min_eps[1]))\n",
    "\n",
    "            args[12 + n + zaehler] = ll_min\n",
    "            args[12 + n + n*(n-1) + n + zaehler] = (min_eps[0], min_eps[1])\n",
    "\n",
    "            if 0. <= ll_min <= 0.01:\n",
    "                rhodteiej = rhodtmax_list_lam(*y, *args)[zaehler]\n",
    "                args[12 + 5*n + 2*n*(n-1) + zaehler] = rhodteiej #rhodtmax_list_lam(*y, *args)[zaehler]\n",
    "        \n",
    "# If the distance between the contact points becomes very small, calculating the direction CPhe_i.pos_from(CPhej) seems to become numerically difficult. This allows me to fix the direction once \n",
    "# the distance is small. Of course, mechanically speaking not correct.\n",
    "            if ll_min > 1.e-5:\n",
    "                args[12 + 5*n + 3*n*(n-1) + zaehler] = richtung_lam[zaehler](args[12 + n + n*(n-1) + n + zaehler][0], args[12 + n + n*(n-1) + n + zaehler][1], *[y[i] for i in range(3*n)], a1, b1)\n",
    "# After penetration, the direction has to be reversed to ensure that the contact is the force of ellipse_i on ellipse_j \n",
    "            elif ll_min < -1.e-5:\n",
    "                args[12 + 5*n + 3*n*(n-1) + zaehler] = (-richtung_lam[zaehler](args[12 + n + n*(n-1) + n + zaehler][0], args[12 + n + n*(n-1) + n + zaehler][1], *[y[i] for i in range(3*n)], a1, b1)[0], \n",
    "                                                        -richtung_lam[zaehler](args[12 + n + n*(n-1) + n + zaehler][0], args[12 + n + n*(n-1) + n + zaehler][1], *[y[i] for i in range(3*n)], a1, b1)[1])\n",
    "            else:\n",
    "                pass\n",
    "        LAENGEe.append(LAENGEe1)\n",
    "        zeit.append(t)\n",
    "        ellipsee.append(ellipsee1)\n",
    "    else:\n",
    "        pass\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, max_step=max_step)#, atol=1.e-4, rtol=1.e-4)# method='Radau')#, 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], '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('shape of resultat', resultat.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot any **coordinates** you want to see"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N2 = 500\n",
    "N1 = max(1, 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 = (['q' + str(i) for i in range(n)] +\n",
    "               ['x' + str(i) for i in range(n)] +\n",
    "               ['z' + str(i) for i in range(n)] + \n",
    "               ['u' + 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",
    "              \n",
    "              ) \n",
    "fig, ax = plt.subplots(figsize=(10,5))\n",
    "for i in range(1*n, 3*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 **energies** of the system.\n",
    "\n",
    "If $c_{tau} < 1.$ or $reibung \\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Just to adapt the arguments so they fit minimize\n",
    "def funcw(x0, args):\n",
    "    return min_distanzCPhiCPhj_lam[kk](*x0, *args)\n",
    "\n",
    "def funce(x0, args):\n",
    "    return min_distanzCPheiCPhej_lam[zaehler](*x0, *args)\n",
    "\n",
    "def funcedidj(x0, args):\n",
    "    return min_distanzCPheiCPhej_lam1[zaehler](*x0, *args)\n",
    "\n",
    "schritte  = resultat.shape[0]\n",
    "kin_np    = np.empty(schritte)\n",
    "spring_np = np.empty(schritte)\n",
    "total_np  = np.empty(schritte)\n",
    "\n",
    "# fid the distance of the ellipses from the wall\n",
    "x0 = 1.\n",
    "laengew1 = []\n",
    "elliw1   = []\n",
    "for i in range(schritte):\n",
    "    laengew2 = []\n",
    "    elliw2   = []\n",
    "    for kk in range(n):\n",
    "        args1 = [resultat[i, j] for j in range(3*n)] + [a1, b1, RW1]\n",
    "        epsilon_min = minimize(funcw, x0, args1, method=methode, options=option_dict)\n",
    "        min_eps = epsilon_min.x % (2.*np.pi)\n",
    "        ll_min = funcw(min_eps, args1)\n",
    "        x0 = min_eps[0]\n",
    "        laengew2.append(ll_min)\n",
    "        elliw2.append(min_eps[0])\n",
    "    laengew1.append(laengew2)    \n",
    "    elliw1.append(elliw2)\n",
    "\n",
    "\n",
    "# find the distances between ellipse\n",
    "# I try to match the values collected in the integration closely to the times given back by the integration.\n",
    "# argumente will hold the positions, where this is the case.\n",
    "if elli_ghost == False:\n",
    "    argumente = [0]\n",
    "    zeit = np.array(zeit)\n",
    "    for zeit1 in times[: -1]:\n",
    "        start = argumente[-1]\n",
    "        start1 = np.min(np.argwhere(zeit >= zeit1))\n",
    "        argumente.append(start1)\n",
    "    if len(argumente) != len(times):\n",
    "        raise Exception('Something went wrong')\n",
    "\n",
    "    laengee1 = []\n",
    "    ellie1   = []\n",
    "\n",
    "    for index in argumente:\n",
    "        laengee1.append(LAENGEe[index])\n",
    "        ellie1.append(ellipsee[index])\n",
    "else:\n",
    "    pass\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",
    "    if elli_ghost == False:\n",
    "        spring_np[i] = spring_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *[pL_vals[k] for k in range(12)], *laengew1[i], *elliw1[i], *laengee1[i], *ellie1[i]) \n",
    "    else:\n",
    "        spring_np[i] = spring_lam1(*[resultat[i, j] for j in range(resultat.shape[1])], *[pL_vals[k] for k in range(12)], *laengew1[i], *elliw1[i])\n",
    "\n",
    "    total_np[i]  = spring_np[i] + kin_np[i]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "#ax.plot(times, kin_np, label = 'kinetic energy')\n",
    "for i, j in zip((kin_np, spring_np), ('kinetic energy', 'spring energy', 'total energy')):\n",
    "    ax.plot(times[: schritte], i, label=j)\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('Energy (Nm)')\n",
    "ax.set_title(f'A: Energies of the system with {n} ellipses, with ctau = {ctau1} and friction = {reibung1}')\n",
    "ax.legend();\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "#ax.plot(times, kin_np, label = 'kinetic energy')\n",
    "for i, j in zip((kin_np, spring_np, total_np), ('kinetic energy', 'spring energy', 'total energy')):\n",
    "    ax.plot(times[: schritte], i, label=j)\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('Energy (Nm)')\n",
    "ax.set_title(f'B. Energies of the system with {n} ellipses, with ctau = {ctau1} and friction = {reibung1}')\n",
    "ax.legend();\n",
    "\n",
    "koerper = 0\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "laengew1 = np.array(laengew1)\n",
    "ax.plot(times, laengew1[:, koerper])\n",
    "ax.set_title(f'distance of ellipse_{koerper} from wall')\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('distance (m)')\n",
    "\n",
    "laengew2 = [min(laengew1[jj, koerper], 0.) for jj in range(len(laengew1))]\n",
    "fig, ax = plt.subplots(figsize=(10,5))\n",
    "ax.plot(times, laengew2)\n",
    "ax.set_title(f'penetration of of ellipse_{koerper} into the wall')\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('distance (m)')\n",
    "\n",
    "if elli_ghost == False:\n",
    "    fig, ax = plt.subplots(figsize=(10, 5))\n",
    "    laengee1 = np.array(laengee1)\n",
    "    ax.plot(times, laengee1[:, 0])\n",
    "    ax.set_title('distance of ellipse_0 from ellipse_1')\n",
    "    ax.set_xlabel('time (sec)')\n",
    "    ax.set_ylabel('distance (m)')\n",
    "\n",
    "\n",
    "    laengee2 = [min(laengee1[jj, 0], 0.0) for jj in range(len(laengee1))]\n",
    "    fig, ax = plt.subplots(figsize=(10,5))\n",
    "    ax.plot(times, laengee2)\n",
    "    ax.set_title('penetration of ellipse0 into ellipse1')\n",
    "    ax.set_xlabel('time (sec)')\n",
    "    ax.set_ylabel('distance (m)')\n",
    "    print(np.min(laengee1[:, 0]))\n",
    "else:\n",
    "    pass"
   ]
  },
  {
   "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."
   ]
  },
  {
   "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(resultat.shape[0]/zeitpunkte))\n",
    "\n",
    "for i in range(resultat.shape[0]):\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, 2*n)] for i in range(schritte2)])                  # X - coordinates of the centers of the ellipses\n",
    "Dmc_Z = np.array([[resultat2[i, j] for j in range(2*n, 3*n)] for i in range(schritte2)])                # Z - coordinates of the centers of the ellipses\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(int(resultat.shape[1]/2.))], *alpha_list1, *beta_list1, a1, b1 )[l][0] for l in range(n) ]\n",
    "    Po_Z[i] = [Po_pos_lam(*[resultat2[i, j] for j in range(int(resultat.shape[1]/2.))], *alpha_list1, *beta_list1, a1, b1 )[l][1] for l in range(n) ]\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",
    "    LINE4 = []\n",
    "\n",
    "    for i in range(n):\n",
    "        x1 = resultat2[0, n+i]\n",
    "        z1 = resultat2[0, 2*n+i]\n",
    "        elli = patches.Ellipse((x1, z1), width=2.*a1, height=2.*b1, angle=-np.rad2deg(resultat2[0, i]), zorder=1, fill=True, color=farben[i], ec='black') # the ellipses\n",
    "        line1 = ax.add_patch(elli)\n",
    "        line2, = ax.plot([], [], 'o', markersize=5, color='white')      # the observers\n",
    "        line3, = ax.plot([], [], '-', markersize=0, linewidth=0.3)      # tracing the centers of the ellipses\n",
    "        line4, = ax.plot([], [], 'o', markersize=5, color='yellow')      # the centers of the ellipses\n",
    "     \n",
    "        LINE1.append(line1)\n",
    "        LINE2.append(line2)\n",
    "        LINE3.append(line3)\n",
    "        LINE4.append(line4)\n",
    "\n",
    "    def animate(i):\n",
    "        ax.set_title(f'System with {n} bodies, running time {times2[i]:.2f} sec' + '\\n' +\n",
    "                     f' ctau = {ctau1}, friction = {reibung1}', fontsize=12)\n",
    "        for j in range(n):\n",
    "            LINE1[j].set_center((resultat2[i, n+j], resultat2[i, 2*n+j]))\n",
    "            LINE1[j].set_angle(-np.rad2deg(resultat2[i, j]))\n",
    "            LINE1[j].set_color(farben[j])\n",
    "\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",
    "            LINE4[j].set_data(([resultat2[i, n+j]], [resultat2[i, 2*n+j]]))\n",
    "            \n",
    "            \n",
    "        return LINE1 + LINE2 + LINE3 + LINE4\n",
    "\n",
    "    anim = animation.FuncAnimation(fig, animate, frames=schritte2,\n",
    "                                   interval=2000*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() - start0):.3f} sec to run the program, BEFORE HTML')\n",
    "HTML(anim.to_jshtml())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "math",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}