ball-bouncing-between-two-horizontal-planes.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sm\n",
    "import sympy.physics.mechanics as me\n",
    "import time\n",
    "import numpy as np\n",
    "from scipy.integrate import solve_ivp\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mp\n",
    "%matplotlib inline\n",
    "from IPython.display import HTML\n",
    "mp.rcParams['animation.embed_limit'] = 2**128\n",
    "from matplotlib.animation import FuncAnimation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use this as a decorator, to see where the function spends its time.\\\n",
    "To see the results, put *profiler.print_stats* somewhere below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from line_profiler import LineProfiler\n",
    "\n",
    "profiler = LineProfiler()\n",
    "\n",
    "def profile(func):\n",
    "    def inner(*args, **kwargs):\n",
    "        profiler.add_function(func)\n",
    "        profiler.enable_by_count()\n",
    "        return func(*args, **kwargs)\n",
    "    return inner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An **elastic homogenious ball** of radius *r* and mass *m* is bouncing between two horizontal planes, one being at z = 0, the other one at z = H.\\\n",
    "Essentially I want to simulate the simplest example of this video: https://www.youtube.com/watch?v=g_VxOIlg7q8 \\\n",
    "An observer (a particle) of mass $m_o$ may be put anywhere inside the ball.\\\n",
    "Gravitiy points in the negative Z direction.\\\n",
    "The idea was given to me by Jason Moore.\n",
    "\n",
    "**Comment:**\\\n",
    "First I tried a *friction model*, where the force to change $\\omega$ is friction, depending on the vetical impact force and an coefficient of friction,$\\hspace{3 pt} \\mu$. I easily found parameters to stop the advancing of the ball, but I could not reverse it, like shown in the video.\\\n",
    "The I tried a *dig in model*, where the ball does not slide at all, the impact force is in the direction of the speed of the contact point. This surely made the ball revere its direction, but unrealistically so, I thought. Also some other issues\\\n",
    "So, I *mixed* the two models, like: **force on the ball = (1 - mixtur) $\\cdot$ force of friction model + mixtur $\\cdot$ force of no slip model**.\\\n",
    "With mixtur $\\approx$ 0.35 this gave reasonable results.\\\n",
    "Maybe this really describes reality to some extend, I know too little about realistic impacts. \n",
    "\n",
    "**Variables**\n",
    "\n",
    "- $N$: inertial frame\n",
    "- $O$: point fixed in N\n",
    "- $A$: frame fixed to the ball\n",
    "\n",
    "- $Dmc$: center of the ball, a *RigidBody*\n",
    "- $P_o$: the observer, a *Particle*\n",
    "- $BCP_u, BCP_o$: corresponding points fixed on the ball\n",
    "\n",
    "- $q_1, q_2, q_3$: angle of the ball w.r.t. N\n",
    "- $u_1, u_2, u_3$: angular velocities of the ball\n",
    "- $m_1, m_2, m_3$: location of the center of the ball, relative to the inertial frame N\n",
    "- $um_1, um_2, um_3$: the speeds\n",
    "\n",
    "\n",
    "- $c_{tau}$: the experimental constant needed for Hunt-Crossley\n",
    "- $m_u$: the friction constant\n",
    "- $ny_b, ny_p$: Poisson's ratio of ball and plane respectively.\n",
    "- $E_b, E_p$: Young's modulus of ball and plane respectively\n",
    "- $rhodt_{u}$: speed right before the impact with the lower plane\n",
    "- $rhodt_{o}$: speed right before the impact with the upper plane\n",
    "\n",
    "- $m, m_o, i_{XX}, i_{YY}, i_{ZZ}$: mass of ball, mass of observer, moments of inertia of the ball\n",
    "- $\\alpha, \\beta, \\gamma$: describe the location of the observer relative to the center of the ball \n",
    "- $ru_1, ...ro_3$: See explanation below.\n",
    "- $mixtur$: See explanation above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_o = time.time()\n",
    "\n",
    "t = me.dynamicsymbols._t\n",
    "\n",
    "N, A = sm.symbols('N, A', cls    = me.ReferenceFrame)\n",
    "O, Dmc, Po, BCPu, BCPo = sm.symbols('O, Dmc, Po, BCPu, BCPo', cls = me.Point)\n",
    "\n",
    "q1, q2, q3, u1, u2, u3    = me.dynamicsymbols('q1, q2, q3, u1, u2, u3')\n",
    "m1, m2, m3, um1, um2, um3 = me.dynamicsymbols('m1, m2, m3, um1, um2, um3')\n",
    "\n",
    "m, mo, g, r, alpha, beta, gamma, iXX, iYY, iZZ, H = sm.symbols('m, mo, g, r, alpha, beta, gamma, iXX, iYY, iZZ, H')\n",
    "nyb, nyp, Eb, Ep, rhodtu, rhodto, ctau, mu        = sm.symbols('nyb, nyp, Eb, Ep, rhodtu, rhodto, ctau, mu')\n",
    "ru1, ru2, ru3, ro1, ro2, ro3                      = sm.symbols('ru1, ru2, ru3, ro1, ro2, ro3')\n",
    "\n",
    "mixtur = sm.symbols('mixtur')\n",
    "\n",
    "A.orient_body_fixed(N, (q1, q2, q3), '123')\n",
    "rot = A.ang_vel_in(N)\n",
    "A.set_ang_vel(N, u1*N.x + u2*N.y + u3*N.z)\n",
    "rot1 = A.ang_vel_in(N)\n",
    "\n",
    "O.set_vel(N, 0)\n",
    "\n",
    "Dmc.set_pos(O, m1*N.x + m2*N.y + m3*N.z)\n",
    "Dmc.set_vel(N, um1*N.x + um2*N.y + um3*N.z)\n",
    "\n",
    "Po.set_pos(Dmc, r * (alpha*A.x + beta*A.y + gamma*A.z))\n",
    "Po.v2pt_theory(Dmc, N, A);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Points of impact**\\\n",
    "The potential contact points of the ball with the plane are obviously these: $BCP_u(m_1 / m_2 / -r)$ and $CP_o(m_1, m_2, r)$\n",
    "\n",
    "Their speeds (it only matters during the impacts) are: $\\begin{bmatrix} um_1 \\\\ um_2 \\\\ um2 \\end{bmatrix} + \\omega \\times \\begin{bmatrix} 0. \\\\ 0. \\\\ +/-r \\end{bmatrix}$, where $\\omega$ = angular speed of the ball, $+$ for the upper contact point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BCPu.set_pos(Dmc, -r*N.z)\n",
    "BCPo.set_pos(Dmc, r*N.z)\n",
    "\n",
    "BCPu.set_vel(N, Dmc.vel(N) + (u1*N.x + u2*N.y + u3*N.z).cross(-r*N.z))\n",
    "BCPo.set_vel(N, Dmc.vel(N) + (u1*N.x + u2*N.y + u3*N.z).cross(r*N.z))\n",
    "print('speeds of contact points:')\n",
    "BCPu.vel(N), BCPo.vel(N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Function which calculates the forces of an impact between the ball and the horizontal planes** for the *friction model*\\\n",
    "1.\\\n",
    "The impact force is on the line normal to the plane, that is in Z direction, and goes through the center of the ball, $Dmc$. I use Hunt Crossley's method to calculate it\\\n",
    "It's direction points from the plane to the ball.\\\n",
    "2.\\\n",
    "The friction force acting on the contact point $BCP_u$ is in the direction of the component of $\\dfrac{d}{dt} BCP_u$ in the plane. It's magnitude is the impact force times a friction factor $m_{uW}$.\\\n",
    "(Of course same for the upper contact point $BCP_o$)\\\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",
    "For the plane, I set $R_2 = \\infty$, so $k_0$ simplifies to $k_0 = \\frac{4}{3\\cdot(\\sigma_1 + \\sigma_2)} \\cdot \\sqrt{R_1}$, with $R_1 = r$, the radius of the ball.\\\n",
    "I am not sure, this is covered by the theory.\n",
    "\n",
    "As per the article, $n = \\frac{3}{2}$ is always to be used, I believe, Hertz derived this on theoretical grounds.\n",
    "\n",
    "*spring energy* =   $ k_0 \\cdot \\int_{0}^{\\rho} k^{3/2}\\,dk$ = $k_0 \\cdot\\frac{2}{5} \\cdot \\rho^{5/2}$\\\n",
    "I assume, the dissipated energy cannot be given in closed form, at least the article does not give one.\n",
    "\n",
    "*Note*  \n",
    "$c_\\tau = 1.$ gives **Hertz's** solution to the impact problem, also described in the article.\n",
    "\n",
    "\n",
    "**Friction when the ball hits the street**\n",
    "\n",
    "If  $\\dfrac{d}{dt} BCP_u$ ist the speed of the contact point of the ball with the plane, then $ \\left[\\dfrac{d}{dt} BCP_u -  \\left( \\dfrac{d}{dt} BCP_u \\circ N.z \\right) \\cdot N.x \\right]$ is the speed component in the plane.\n",
    "\n",
    "\n",
    "As there are numerical problems when $\\left| \\left[\\dfrac{d}{dt} BCP_u -  \\left( \\dfrac{d}{dt} BCP_u \\circ N.z \\right) \\cdot N.x \\right] \\right| \\approx 0$,  I use $\\left| \\left[\\dfrac{d}{dt} BCP_u -  \\left( \\dfrac{d}{dt} BCP_u \\circ N.z \\right) \\cdot N.x \\right] \\right| \\vee 10^{-10}$ Values smaller than $10^{-10}$ also give problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fric_lower_plane(rhodtu):\n",
    "# force acting on BCPu\n",
    "            \n",
    "    abstand = m3                                # should always be positive, the ball must not fall through the planes\n",
    "    richtung = N.z                              # pointing towards the ball\n",
    "    rho = r - abstand                           # positive in the ball has penetrated the plane\n",
    "    \n",
    "    vCP = BCPu.vel(N)                           # speed of the contact point of the ball \n",
    "    rhodt = me.dot(vCP, richtung)               # speed component in Z direction pointing from the ball to the plane\n",
    "    rho = sm.Abs(rho)\n",
    "\n",
    "# determine k0\n",
    "    wurzel = sm.sqrt(r)\n",
    "    sigma_b = (1. - nyb**2) / Eb\n",
    "    sigma_s = (1. - nyp**2) / Ep\n",
    "    k0      = 4. / (3.*(sigma_b + sigma_s)) * wurzel\n",
    "\n",
    "# Impact force on BCPu in Z direction\n",
    "# Here I assume, that the penetration is small, otherwise I would not know how to do it.\n",
    "    forcec = k0 * rho**(3/2) * ( 1. + 3./2. * (1 - ctau) * (-rhodt) / sm.Abs(rhodtu)) * richtung * sm.Heaviside(r - abstand, 0.)\n",
    "\n",
    " \n",
    " # The speed component of vCP in the plane is, of course, the speed minus its component perpendicular to the plan.    \n",
    "    vx = vCP - (me.dot(vCP, richtung)) * richtung\n",
    "    forcef = forcec.magnitude() * mu * (-vx) * sm.Heaviside(r - abstand, 0.) * 1. / sm.Max(vx.magnitude(), 1.e-10)    # Here, Heaviside probably not really needed.     \n",
    "\n",
    "# The force acting on BCPu is returned        \n",
    "    return forcec + forcef\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same comments as above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fric_upper_plane(rhodto):\n",
    "# force acting on BCPo\n",
    "            \n",
    "    abstand = H - m3                            # should always be positive, the ball must not fall through the planes\n",
    "    richtung = -N.z                             # pointing towards the ball\n",
    "    rho = r - abstand                           # positive in the ball has penetrated the plane\n",
    "    \n",
    "    vCP = BCPo.vel(N)                           # speed of the contact point of the ball \n",
    "    rhodt = me.dot(vCP, richtung)               # speed component in Z direction pointing from the ball to the plane\n",
    "    rho = sm.Abs(rho)\n",
    "\n",
    "# determine k0\n",
    "    wurzel = sm.sqrt(r)\n",
    "    sigma_b = (1. - nyb**2) / Eb\n",
    "    sigma_s = (1. - nyp**2) / Ep\n",
    "    k0      = 4. / (3.*(sigma_b + sigma_s)) * wurzel\n",
    "\n",
    "# Impact force on BCP0 in Z direction\n",
    "# Here I assume, that the penetration is small, otherwise I would not know how to do it.\n",
    "    forcec = k0 * rho**(3/2) * ( 1. + 3./2. * (1 - ctau) * (-rhodt) / sm.Abs(rhodto)) * richtung * sm.Heaviside(r - abstand, 0.)\n",
    "\n",
    " # The speed component of vCP in the plane is, of course, the speed minus its component perpendicular to the plan.    \n",
    "    vx = vCP - (me.dot(vCP, richtung)) * richtung\n",
    "    forcef = forcec.magnitude() * mu * (-vx) * sm.Heaviside(r - abstand, 0.) * 1./sm.Max(vx.magnitude(), 1.e-10)   # Here, Heaviside probably not really needed.     \n",
    "\n",
    "# The force acting on BCPu is returned  \n",
    "\n",
    "    return forcec + forcef"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here I assume, that the ball does not slide on the plane when it hits it, but 'digs into the plane' in the direction of the speed of the contact point.\\\n",
    "I do not know, whether this is covered by the Hunt-Crossley theory, which I use to calculate the impact forces.\n",
    "\n",
    "As there are numerical problems when $\\left| \\dfrac{d}{dt} BCP_u \\right| \\approx 0$,  I calculate $\\dfrac{d}{dt} BCP_u$ right before the impact during the integration, and hand $ru_1, ru_2, ru_3$ to the function, where:\\\n",
    "$\\dfrac{d}{dt} BCP_u = ru_1 \\cdot N.x + ru_2 \\cdot N.y + ru_3 \\cdot N.z$\\\n",
    "As the duration of the impact is short, I do not think, the error caused by this simplification (keeping the direction of the force constant during impact) has much effect.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def HC_lower_plane(rhodtu, ru1, ru2, ru3):\n",
    "# force acting on BCPu\n",
    "            \n",
    "    abstand = m3                                                                        # should always be positive, the ball must not fall through the planes\n",
    "    vCP = BCPu.vel(N)                                                                   # speed of the contact point on the ball\n",
    "    richtung = -(ru1*N.x + ru2*N.y + ru3*N.z)                                           # pointing towards the ball. To be consistent, I want this to be a unit vector, and always point towards the ball\n",
    "    rho = r - abstand                                                                   # positive in the ball has penetrated the plane\n",
    "    \n",
    "    rhodt = me.dot(vCP, richtung)                                                       # speed component in Z direction pointing from the ball to the plane\n",
    "    rho = sm.Abs(rho)\n",
    "\n",
    "# determine k0\n",
    "    wurzel = sm.sqrt(r)\n",
    "    sigma_b = (1. - nyb**2) / Eb\n",
    "    sigma_p = (1. - nyp**2) / Ep\n",
    "    k0      = 4. / (3.*(sigma_b + sigma_p)) * wurzel\n",
    "\n",
    "# Impact force on BCPu\n",
    "# Here I assume, that the penetration is small, otherwise I would not know how to do it.\n",
    "    forcec = k0 * rho**(3/2) * ( 1. + 3./2. * (1 - ctau) * (-rhodt) / sm.Abs(rhodtu)) * richtung * sm.Heaviside(r - abstand, 0.)\n",
    "\n",
    "# The force acting on BCPu is returned        \n",
    "    return forcec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Collision with the upper plane. Same comments as above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def HC_upper_plane(rhodto, ro1, ro2, ro3):\n",
    "# force acting on BCPo\n",
    "            \n",
    "    abstand = H - m3                                                                    # should always be positive, the ball must not fall through the planes\n",
    "    vCP = BCPo.vel(N)                                                                   # speed of the contact point on the ball\n",
    "    richtung = -(ro1*N.x + ro2*N.y + ro3*N.z)                                           # pointing towards the ball. \n",
    "    rho = r - abstand                                                                   # positive in the ball has penetrated the plane\n",
    "    \n",
    "    rhodt = me.dot(vCP, richtung)                                                       # speed component in Z direction pointing from the ball to the plane\n",
    "    rho = sm.Abs(rho)\n",
    "\n",
    "# determine k0\n",
    "    wurzel = sm.sqrt(r)\n",
    "    sigma_b = (1. - nyb**2) / Eb\n",
    "    sigma_s = (1. - nyp**2) / Ep\n",
    "    k0      = 4. / (3.*(sigma_b + sigma_s)) * wurzel\n",
    "\n",
    "# Impact force on BCPo\n",
    "# Here I assume, that the penetration is small, otherwise I would not know how to do it.\n",
    "    forcec = k0 * rho**(3/2) * ( 1. + 3./2. * (1 - ctau) * (-rhodt) / sm.Abs(rhodto)) * richtung * sm.Heaviside(r - abstand, 0.)\n",
    "\n",
    "# The force acting on BCPu is returned        \n",
    "    return forcec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Kane's equations**.\\\n",
    "Nothing special. If the frame N is used in the kinematic equations, rather than the frame A, the number of operations in force shoot up.\\\n",
    "Here I mix der forces of the 'friction model' with the forces of the 'no slip model', see the variable *mixtur* in the force terms. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "start2 = time.time()\n",
    "\n",
    "I = me.inertia(A, iXX, iYY, iZZ)\n",
    "body = me.RigidBody('body', Dmc, A, m, (I, Dmc))\n",
    "Poa  = me.Particle('Poa', Po, mo)\n",
    "BODY = [body, Poa]\n",
    "\n",
    "\n",
    "\n",
    "F1f = [(Dmc, -m*g*N.z), (Po, -mo*g*N.z)]\n",
    "F2f = [(BCPu, (sm.S(1.) - mixtur) * fric_lower_plane(rhodtu))]\n",
    "F3f = [(BCPo, (sm.S(1.) - mixtur) * fric_upper_plane(rhodto))]\n",
    "\n",
    "F1hc = [(Dmc, -m*g*N.z), (Po, -mo*g*N.z)]\n",
    "F2hc = [(BCPu, mixtur * HC_lower_plane(rhodtu, ru1, ru2, ru3))]\n",
    "F3hc = [(BCPo, mixtur * HC_upper_plane(rhodto, ro1, ro2, ro3))]\n",
    "\n",
    "FL = (F1f + F2f + F3f) + (F2hc + F3hc)                                  #Gravity needed once only\n",
    "\n",
    "q_ind = [q1, q2, q3] + [m1, m2, m3]\n",
    "u_ind = [u1, u2, u3] + [um1, um2, um3]\n",
    "\n",
    "kd = [me.dot(rot - rot1, uv) for uv in A] + [i - j.diff(t) for i, j in zip((um1, um2, um3), (m1, m2, m3))]\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",
    "force = KM.forcing_full\n",
    "\n",
    "print('force DS', me.find_dynamicsymbols(force))\n",
    "print('force free symbols', force.free_symbols)\n",
    "print(f'force has {sm.count_ops(force):,} operations, {sm.count_ops(sm.cse(force))} after cse \\n')\n",
    "\n",
    "print('MM DS', me.find_dynamicsymbols(MM))\n",
    "print('MM free symbols', MM.free_symbols)\n",
    "print(f'MM has {sm.count_ops(MM):,} operations, {sm.count_ops(sm.cse(MM))} after cse \\n')\n",
    "\n",
    "print('It took {:.3f} sec to establish Kanes equations'.format(time.time() - start2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I like to look at the *energies of the system*. This has often (just about always!) told me, I had made a mistake in the equations of motion.\\\n",
    "So, these functions are defined here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "pot_energie = m * g * me.dot(Dmc.pos_from(O), N.z) + mo * g * me.dot(Po.pos_from(O), N.z)\n",
    "kin_energie = sum([koerper.kinetic_energy(N) for koerper in BODY])\n",
    "\n",
    "rho1 = r - m3               # positive if impact with lower plane\n",
    "rho2 = r - (H - m3)         # positive if impact with upper plane\n",
    "sigma_b = (1. - nyb**2) / Eb\n",
    "sigma_p = (1. - nyp**2) / Ep\n",
    "k0      = 4. / (3.*(sigma_b + sigma_p)) * sm.sqrt(r)\n",
    "\n",
    "spring_energie = 2./5. * k0  * sm.Abs(rho1)**(5/2) * sm.Heaviside(rho1, 0.)  + 2./5. * k0 * sm.Abs(rho2)**(5/2) * sm.Heaviside(rho2, 0.)\n",
    "\n",
    "rhodtuu = me.dot(BCPu.vel(N), -N.z)      # speed of BCPu in direction of the lower plane\n",
    "rhodtoo = me.dot(BCPo.vel(N), N.z)       # analoguously dto."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert the *sympy functions* to *numpy functions* so numerical evaluations may be done. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "qL   = q_ind + u_ind\n",
    "pL   = [m, mo, g, r, alpha, beta, gamma, iXX, iYY, iZZ, H, mixtur] + [nyb, nyp, Eb, Ep, ctau, mu, rhodtu, rhodto, ru1, ru2, ru3, ro1, ro2, ro3]\n",
    "MM_lam = sm.lambdify(qL + pL, MM, cse=True)\n",
    "force_lam = sm.lambdify(qL + pL, force, cse=True)\n",
    "\n",
    "pot_lam    = sm.lambdify(qL + pL, pot_energie, cse=True)\n",
    "kin_lam    = sm.lambdify(qL + pL, kin_energie, cse=True)\n",
    "spring_lam = sm.lambdify(qL + pL, spring_energie, cse=True)\n",
    "\n",
    "rhodtu_lam = sm.lambdify(qL + pL, rhodtuu, cse=True)\n",
    "rhodto_lam = sm.lambdify(qL + pL, rhodtoo, cse=True)\n",
    "\n",
    "richtungu_lam = sm.lambdify(qL + pL, [me.dot(BCPu.vel(N).normalize(), uv) for uv in N], cse=True)\n",
    "richtungo_lam = sm.lambdify(qL + pL, [me.dot(BCPo.vel(N).normalize(), uv) for uv in N], cse=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The *parameters* and the *initial conditions* are set here. For their meaning, see above.\\\n",
    "It makes sense to use names close to the names used to set up Kane's equations, but they should **not** be the same ones. The sm.symbols and the me.dynamic symbols would be overwritten - with unintended consequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#=======================================\n",
    "# Set parameters and initial conditions.\n",
    "#=======================================\n",
    "mixtur1 = 0.375\n",
    "#=======================================\n",
    "\n",
    "m11, m21, m31 = 0., 0., 5.\n",
    "um11, um21, um31 = -5., 5., -25.\n",
    "\n",
    "q11, q21, q31 = 0., 0., 0.\n",
    "u11, u21, u31 = 0., 0., 0.\n",
    "\n",
    "mb1 = 1.\n",
    "mo1 = 0.0001\n",
    "r1 = 1.\n",
    "\n",
    "H1 = 10.\n",
    "\n",
    "nyb1  = 0.25    # Poisson\"s ratio for rubber, from the internet\n",
    "nyp1  = 0.15    # Poisson's ration for concrete, from the internet\n",
    "Eb1   = 3.e3    # Young's modulus for rubber is really about 3.e9\n",
    "Ep1   = 2.e10   # Young's modulus for steel is really about 2.e11\n",
    "\n",
    "ctau1 = 0.90    # experimental constant\n",
    "mu1   = 0.1     # coefficient of friction\n",
    "\n",
    "alpha1, beta1, gamma1 = 0.7, 0., 0.7\n",
    "\n",
    "intervall = 10.\n",
    "\n",
    "rhodtu1, rhodto1                   = 1., 1.                         # of no importance here.any value o.k.\n",
    "ru11, ru21, ru31, ro11, ro21, ro31 = 1., 1., 1., 1., 1., 1.         # dto.\n",
    "\n",
    "\n",
    "#===============================================================\n",
    "if alpha1**2 + beta1**2 + gamma1**2 >= 1.:\n",
    "    raise Exception('Observer is outside of the ball')\n",
    "    \n",
    "schritte     = int(intervall * 1000.) # this should be slightly less than nfev\n",
    "\n",
    "iXX1         = 2./5. * mb1 * r1**2      # from the internet.\n",
    "iYY1         = iXX1\n",
    "iZZ1         = iXX1\n",
    "\n",
    "# pL   = [m, mo, g, r, alpha, beta, gamma, iXX, iYY, iZZ, H] + [nyb, nyp, Eb, Ep, ctau, mu, rhodtu, rhodto, richtungu, richtungo]\n",
    "pL_vals = [mb1, mo1, 9.81, r1, alpha1, beta1, gamma1, iXX1, iYY1, iZZ1, H1, mixtur1] + [nyb1, nyp1, Eb1, Ep1, ctau1, mu1, rhodtu1, rhodto1, ru11, ru21, ru31, ro11, ro21, ro31]\n",
    "y0      = [q11, q21, q31, m11, m21, m31] + [u11, u21, u31, um11, um21, um31]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Numerical Integration**\\\n",
    "The parameters $rhodt_u, rhodt_o, ru_1, ru_2, ru_3, ro_1, ro_2, ro_3$ are available only during integration, but needed later for the Energy.\\\n",
    "So, I collect them during integration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "start2 = time.time()\n",
    "times = np.linspace(0, intervall, schritte)\n",
    "t_span = (0., intervall)\n",
    "\n",
    "impact_list = []\n",
    "zaehler = 0         # I try to reduce the number of entries into the impact_list. Not really essential.\n",
    "\n",
    "zeit_list   = []\n",
    "params_list = []\n",
    "\n",
    "#@profile\n",
    "def gradient(t, y, args):\n",
    "    global zaehler\n",
    "# set rhodtu / rhodto\n",
    "    if r1 < y[5] < r1 + 0.01:\n",
    "        zaehler += 1\n",
    "        args[-8] = rhodtu_lam(*y, *args)\n",
    "        if zaehler == 1:\n",
    "            impact_list.append([t, y[3], y[4], y[5]])\n",
    "\n",
    "        tu1, tu2, tu3 = richtungu_lam(*y, *args)\n",
    "        if tu3 < 0.:                                        # I want the speed vector only when the ball moves into the plane, not when it moves out. this would give the wrong direction.\n",
    "            args[-6], args[-5], args[-4] = tu1, tu2, tu3\n",
    "\n",
    "\n",
    "    elif r1 < H1 - y[5] < r1 + 0.01:\n",
    "        zaehler += 1\n",
    "        args[-7] = rhodto_lam(*y, *args)\n",
    "        if zaehler == 1:\n",
    "            impact_list.append([t, y[3], y[4], y[5]])\n",
    "\n",
    "        to1, to2, to3 = richtungo_lam(*y, *args)\n",
    "        if to3 > 0.:                                        # dto.\n",
    "            args[-3], args[-2], args[-1] = to1, to2, to3\n",
    "                \n",
    "    else:\n",
    "        zaehler = 0\n",
    "    \n",
    "    zeit_list.append(t)\n",
    "    params_list.append(args)\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, t_span, y0, t_eval = times, args=(pL_vals,), atol=1.e-12, rtol=1.e-12, method='Radau')\n",
    "\n",
    "resultat = resultat1.y.T\n",
    "print('resultat shape', resultat.shape, '\\n')\n",
    "event_dict = {-1: 'Integration failed', 0: 'Integration finished successfully', 1: 'some termination event'}\n",
    "print(event_dict[resultat1.status], ', message is:', resultat1.message, '\\n')\n",
    "\n",
    "print(f'To numerically integrate an intervall of {intervall} sec the routine cycled {resultat1.nfev:,} times and it took {(time.time() - start2):.3f} sec')\n",
    "#profiler.print_stats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot any **generalized coordinates or speeds** you want to see.\\\n",
    "I consider around *zeitpunkte* points in time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# reduce the number of points of time to around zeitpunkte\n",
    "times2 = []\n",
    "resultat2 = []\n",
    "index2 = []\n",
    "\n",
    "#=======================\n",
    "zeitpunkte = 1000\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",
    "schritte2 = len(times2)\n",
    "resultat2 = np.array(resultat2)\n",
    "times2 = np.array(times2)\n",
    "\n",
    "test = '$c_{tau}$'\n",
    "bezeichnung = [str(i) for i in qL]\n",
    "fig, ax = plt.subplots(figsize=(10,5))\n",
    "for i in (3, 4, 5, 6, 7):\n",
    "    ax.plot(times2, resultat2[:, i], label=bezeichnung[i])\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('units depending on coordinates selected')\n",
    "ax.set_title(f'Generalized coordinates and / or speeds \\n {test} = {ctau1}. $\\mu$ = {mu1}, mixture = {mixtur1}')\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the **energies** of the system.\\\n",
    "The parameters $rhodt_u, rhodt_o, ru_1, ru_2, ru_3, ro_1, ro_2, ro_3$ were collected at each step of the integration because they are needed here. I try to match the times when they were collected closely to the times returned at the end of the inbtegration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# I try to match the values collected in the integration closely to the times given back by the integration.\n",
    "# argumentw will hold the positions, where this is the case.\n",
    "argumente = [0]\n",
    "zeitw     = np.array(zeit_list)\n",
    "for zeit1 in times[: -1]:\n",
    "    start  = argumente[-1]\n",
    "    start1 = np.min(np.argwhere(zeitw >= zeit1))\n",
    "    argumente.append(start1)\n",
    "if len(argumente) != len(times):\n",
    "    raise Exception('Something went wrong')\n",
    "\n",
    "\n",
    "\n",
    "pot_np    = np.empty(resultat2.shape[0])\n",
    "kin_np    = np.empty(resultat2.shape[0])\n",
    "spring_np = np.empty(resultat2.shape[0])\n",
    "total_np  = np.empty(resultat2.shape[0])\n",
    "\n",
    "for i in range(schritte2):\n",
    "    pL_vals = params_list[argumente[i]]\n",
    "    pot_np[i]     = pot_lam(*resultat2[i, :], *pL_vals)\n",
    "    kin_np[i]     = kin_lam(*resultat2[i, :], *pL_vals)\n",
    "    spring_np[i]  = spring_lam(*resultat2[i, :], *pL_vals)\n",
    "    total_np[i]   = pot_np[i] + kin_np[i] + spring_np[i]\n",
    "  \n",
    "fig, ax = plt.subplots(figsize=(10,5))\n",
    "ax.plot(times2, pot_np, label='potential energy')\n",
    "ax.plot(times2, kin_np, label='kinetic energy')\n",
    "ax.plot(times2, spring_np, label='spring energy')\n",
    "ax.plot(times2, total_np, label='total energy')\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('Energy (Nm)')\n",
    "test = '$c_{tau}$'\n",
    "ax.set_title(f'Energies of the system, {test} = {ctau1}, friction $\\mu$ = {mu1}, mixture = {mixtur1}')\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Animate** the motion\\\n",
    "Substantial parts of it come from chat GPT.\\\n",
    "I consider around *zeitpunkte* points in time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The basic program was created by chat GPT.\n",
    "#=======================\n",
    "zeitpunkte = 500\n",
    "#======================\n",
    "times2 = []\n",
    "resultat2 = []\n",
    "index2 = []\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",
    "print(f'animation used {schritte2} points in time')\n",
    "resultat2 = np.array(resultat2)\n",
    "times2 = np.array(times2)\n",
    "\n",
    "# location of particle\n",
    "Po_loc = [me.dot(Po.pos_from(O), uv) for uv in N]\n",
    "Po_loc_lam = sm.lambdify(qL + pL, Po_loc, cse=True)\n",
    "\n",
    "Pox, Poy, Poz = Po_loc_lam(*[resultat2[:, j] for j in range(resultat.shape[1])], *pL_vals)\n",
    "\n",
    "# Get only one contact point for each impact. I assume, that subsequent impacts have a time difference of at least 0.1 sec. \n",
    "impact_net = []\n",
    "impact_net.append(impact_list[0])\n",
    "zeit1 = impact_net[0][0]\n",
    "for i in range(len(impact_list)):\n",
    "    if impact_list[i][0] - zeit1 < 0.1:\n",
    "        pass\n",
    "    else:\n",
    "        zeit1 = impact_list[i][0]\n",
    "        impact_net.append(impact_list[i])\n",
    "print('Number of contact points detected', len(impact_net))\n",
    "\n",
    "\n",
    "# plot the horizontal plane H\n",
    "def plot_3d_plane(x_min, x_max, y_min, y_max, z_wert, alpha1):\n",
    "    # Create a meshgrid for x and y values\n",
    "    x = np.linspace(x_min, x_max, 100)\n",
    "    y = np.linspace(y_min, y_max, 100)\n",
    "    x, y = np.meshgrid(x, y)\n",
    "\n",
    "    # Z values are set to 0 for the plane\n",
    "    z = np.ones_like(x) * z_wert\n",
    "\n",
    "    \n",
    "    # Plot the 3D plane\n",
    "    ax.plot_surface(x, y, z, alpha=alpha1, rstride=100, cstride=100, color='c') #, edgecolors='k')\n",
    "    \n",
    "# Function to create ellipsoid points\n",
    "def create_ellipsoid(a, b, c, theta, phi, delta, offset):\n",
    "    u = np.linspace(0, 2 * np.pi, 20)\n",
    "    v = np.linspace(0, np.pi, 20)\n",
    "\n",
    "    x = a * np.outer(np.cos(u), np.sin(v))\n",
    "    y = b * np.outer(np.sin(u), np.sin(v))\n",
    "    z = c * np.outer(np.ones(np.size(u)), np.cos(v))\n",
    "\n",
    "    # Rotation matrix\n",
    "    R_theta = np.array([[np.cos(theta), -np.sin(theta), 0],             # rotation around Z axis\n",
    "                       [np.sin(theta), np.cos(theta), 0],\n",
    "                       [0, 0, 1]])\n",
    "\n",
    "    R_phi = np.array([[np.cos(phi), 0, np.sin(phi)],                    # rotation around Y axis\n",
    "                     [0, 1, 0],\n",
    "                     [-np.sin(phi), 0, np.cos(phi)]])\n",
    "\n",
    "    R_delta = np.array([[1, 0, 0],                                      # rotation around X axis\n",
    "                       [0, np.cos(delta), -np.sin(delta)],\n",
    "                       [0, np.sin(delta), np.cos(delta)]])\n",
    "\n",
    "    rotated_coords = np.dot(np.dot(np.dot(np.array([x.flatten(), y.flatten(), z.flatten()]).T, R_theta), R_phi), R_delta).T\n",
    "#    rotated_coords = np.dot(np.dot(np.dot(np.array([x.flatten(), y.flatten(), z.flatten()]).T, R_phi), R_delta), R_theta).T\n",
    "    rotated_coords1 = rotated_coords.reshape(3, 20, 20)\n",
    "\n",
    "    for i in range(20):\n",
    "        for j in range(20):\n",
    "            rotated_coords1[:, i, j] = rotated_coords1[:, i, j]  + np.array([offset[0], offset[1], offset[2]])\n",
    "    return rotated_coords1\n",
    "\n",
    "maxx = np.max(resultat2[:, 3])\n",
    "minx = np.min(resultat2[:, 3])\n",
    "\n",
    "maxy = np.max(resultat2[:, 4])\n",
    "miny = np.min(resultat2[:, 4])\n",
    "\n",
    "maxx = max(maxx, maxy) + 1.\n",
    "minx = min(minx, miny) - 1\n",
    "maxy = maxx\n",
    "miny = minx\n",
    "maxz = H1 + 1.\n",
    "minz = -1.\n",
    "maxmax = max(maxx, maxz)\n",
    "minmin = min(minx, minz)\n",
    "\n",
    "# This is to asign colors of 'plasma' to the impact points.\n",
    "Test = mp.colors.Normalize(0, len(impact_net))\n",
    "Farbe = mp.cm.ScalarMappable(Test, cmap='plasma')\n",
    "\n",
    "\n",
    "# Function to update the plot in the animation\n",
    "def update(frame):\n",
    "    plt.cla()\n",
    "    ax.set_xlim([minmin, maxmax])\n",
    "    ax.set_ylim([minmin, maxmax])\n",
    "    ax.set_zlim([minmin, maxmax])\n",
    "\n",
    "    message = (f'running time {times2[frame]:.2f} sec \\n The red dot is the particle $P_O$ \\n The lighter the color of the contact point, the later in time it happened')\n",
    "    ax.set_title(message, fontsize=12)\n",
    "    ax.set_xlabel('X direction', fontsize=12)\n",
    "    ax.set_ylabel('Y direction', fontsize=12)\n",
    "    ax.set_zlabel('Z direction', fontsize=12)\n",
    "    ebeneu = plot_3d_plane(minmin, maxmax, minmin, maxmax, 0., alpha1=0.2)                              # plot lower plane\n",
    "    ebeneo = plot_3d_plane(minmin, maxmax, minmin, maxmax, H1, alpha1 = 0.1)                            # plot upper plane\n",
    "\n",
    "    ax.plot([Pox[frame]], [Poy[frame]], Poz[frame], marker='o', color='red')        # The particle\n",
    "\n",
    "    for i in range(len(impact_net)):\n",
    "        if times2[frame] > impact_net[i][0]:\n",
    "            farbe2 = Farbe.to_rgba(i)\n",
    "            ax.plot([impact_net[i][1]], [impact_net[i][2]], [impact_net[i][3]], marker='o', color=farbe2) \n",
    "            \n",
    "\n",
    "    \n",
    "\n",
    "    delta, phi, theta = resultat2[frame, 0 : 3]\n",
    "    offset = [resultat2[frame, 3], resultat2[frame, 4], resultat2[frame, 5]]\n",
    "    ellipsoid = create_ellipsoid(r1, r1, r1, delta, phi, theta, offset)\n",
    "\n",
    "    ax.plot_wireframe(ellipsoid[0], ellipsoid[1], ellipsoid[2], color='b', alpha=0.6, lw=0.25)\n",
    "\n",
    "\n",
    "# Create a Matplotlib figure and axis\n",
    "fig = plt.figure(figsize=(10, 10))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "\n",
    "ax.view_init(elev=30, azim=30, roll=0.)\n",
    "\n",
    "# Create the animation\n",
    "ani = FuncAnimation(fig, update, frames=schritte2, interval=3000*np.max(times2) / schritte2)\n",
    "\n",
    "print(f'It took {(time.time() - start_o):.3f} sec to run the program, before HTML')\n",
    "\n",
    "# Display the animation in HTML\n",
    "html_output = HTML(ani.to_jshtml())\n",
    "plt.close()  # To prevent the static plot from showing up along with the animation\n",
    "html_output"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}