chain-link.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sm\n",
    "import sympy.physics.mechanics as me\n",
    "import time\n",
    "import numpy as np\n",
    "from scipy.integrate import solve_ivp\n",
    "from scipy.optimize import fsolve, minimize\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mp\n",
    "from matplotlib import patches\n",
    "%matplotlib inline\n",
    "from IPython.display import HTML\n",
    "mp.rcParams['animation.embed_limit'] = 2**128\n",
    "from matplotlib import animation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**A chain**\n",
    "\n",
    "I want to simulate a chain, which I model as a simple open 2D n-link pendulum, where each link is modelled a a thin rod.\\\n",
    "I saw this some time ago as an example of some other modelling library, and I want to see, how this will work with sympy mechanics.\\\n",
    "\n",
    "I noticed the following:\n",
    "- If I define the rotation of frame $A[i]$ relative to frame $A[i-1]$ the number of operations shoot up: for $n = 12$ the **mass matrix has 25 mio operations**. If I define the rotation of $A[i]$ relative to the inertial frame $O$, the numbers become very reasonable, for $n = 50$, around 40,000 operations.\n",
    "- Setting up Kane's equations seems to take quite long, given the low numbers of operations in the mass matrix and the force vector. (This is by my experience, I have run very many examples, but certainly not 'scientific'.)\n",
    "- Also the *lambdification* runs quite long given the small number of operations, again by my experience. With a system, where the mass matrix has 40,000 operations, the force vector has 50,000 operations and the rection force vector has 1,700 operations setting up Kane's equations normally takes a few seconds, here it takes 2 min! \n",
    "\n",
    "**Symbols**\n",
    "\n",
    "- $O$: Reference frame\n",
    "- $PO$: point fixed in $O$\n",
    "- $A[i]$: body fixed frame of link $i$ with $0 \\leq n < n$\n",
    "- $Dmc[i]$: center of gravitiy of link $i$\n",
    "- $P[i]$: point, where frame $A[i]$ joins frame $A[i+1]$\n",
    "- $l$: length of the pendulum, that is each link has length = $\\dfrac{l}{n}$\n",
    "\n",
    "- $m$: mass of each link\n",
    "- $iZZ$: moment of inertial of each link around $A[i].z$, relative to $Dmc[i]$\n",
    "- $reibung$: speed dependent friction in each joint.\n",
    "- $auxx, auxy, fx, fy$: virtual speeds and reaction forces on point $O$\n",
    "- $rhs$: 'place holders' for the $rhs = MM^{-1} \\cdot force$ to be evaluated numerically later\n",
    "\n",
    "- $q[i]$: generalized coordinate of frame $A[i]$ relative to the inertial frame $O$.\n",
    "- $u[i]$: angular speed dto."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#==================\n",
    "n = 50\n",
    "term_info = False\n",
    "#==================\n",
    "start = time.time()\n",
    "\n",
    "m, g, iZZ, l, reibung = sm.symbols('m, g, iZZ, l, reibung')\n",
    "q   = me.dynamicsymbols(f'q:{n}')\n",
    "u   = me.dynamicsymbols(f'u:{n}')\n",
    "\n",
    "auxx, auxy, fx, fy = me.dynamicsymbols('auxx, auxy, fx, fy')\n",
    "t = me.dynamicsymbols._t\n",
    "\n",
    "A   = sm.symbols(f'A:{n}', cls=me.ReferenceFrame)\n",
    "Dmc = sm.symbols(f'Dmc:{n}', cls=me.Point)\n",
    "P   = sm.symbols(f'P:{n}', cls=me.Point)\n",
    "rhs = list(sm.symbols(f'rhs:{n}'))\n",
    "\n",
    "O   = me.ReferenceFrame('O')\n",
    "PO  = me.Point('PO')\n",
    "PO.set_vel(O, auxx*O.x + auxy*O.y)\n",
    "\n",
    "l1 = l/n\n",
    "l2 = l/(2 * n)\n",
    "\n",
    "A[0].orient_axis(O, q[0], O.z)\n",
    "A[0].set_ang_vel(O, u[0] * O.z)\n",
    "\n",
    "Dmc[0].set_pos(PO, l2*A[0].x)\n",
    "Dmc[0].v2pt_theory(PO, O, A[0])\n",
    "P[0].set_pos(PO, l1*A[0].x)\n",
    "P[0].v2pt_theory(PO, O, A[0])\n",
    "\n",
    "for i in range(1, n):\n",
    "    A[i].orient_axis(O, q[i], O.z)         #(A[i-1], q[i], A[i-1].z)\n",
    "    A[i].set_ang_vel(O, u[i] * O.z)    #(A[i-1], u[i] * A[i-1].z)\n",
    "\n",
    "    Dmc[i].set_pos(P[i-1], l2*A[i].x)\n",
    "    Dmc[i].v2pt_theory(P[i-1], O, A[i])\n",
    "    P[i].set_pos(P[i-1], l1*A[i].x)\n",
    "    P[i].v2pt_theory(P[i-1], O, A[i])\n",
    "\n",
    "BODY = []\n",
    "for i in range(n):\n",
    "    inertia = me.inertia(A[i], 0., 0., iZZ)\n",
    "    body    = me.RigidBody('body' + str(i), Dmc[i], A[i], m, (inertia, Dmc[i]))\n",
    "    BODY.append(body)\n",
    "\n",
    "FL1    = [(Dmc[i], -m*g*O.y) for i in range(n)] + [(PO, fx*O.x + fy*O.y)]\n",
    "Torque = [(A[i], - u[i] * reibung * A[i].z) for i in range(n)]\n",
    "FL     = FL1 + Torque\n",
    "\n",
    "kd     = [u[i] - q[i].diff(t) for i in range(n)]\n",
    "aux    = [auxx, auxy]\n",
    "\n",
    "KM = me.KanesMethod(O, q_ind=q, u_ind=u, kd_eqs=kd, u_auxiliary=aux)\n",
    "(fr, frstar) = KM.kanes_equations(BODY, FL)\n",
    "\n",
    "MM = KM.mass_matrix_full\n",
    "if term_info == True:\n",
    "    print('MM DS', me.find_dynamicsymbols(MM))\n",
    "    print('MM free symbols', MM.free_symbols)\n",
    "    print('MM contains {} operations'.format(sum([MM[i, j].count_ops(visual=False) for i in range(MM.shape[0]) for j in range(MM.shape[1])])), '\\n')\n",
    "\n",
    "force = KM.forcing_full\n",
    "if term_info == True:\n",
    "    print('force DS', me.find_dynamicsymbols(force))\n",
    "    print('force free symbols', force.free_symbols)\n",
    "    print('force contains {} operations'.format(sum([force[i].count_ops(visual=False) for i in range(force.shape[0])])), '\\n')\n",
    "\n",
    "eingepraegt_dict = {sm.Derivative(i, t): j for i, j in zip(u, rhs)} \n",
    "eingepraegt      = KM.auxiliary_eqs.subs(eingepraegt_dict)\n",
    "if term_info == True:\n",
    "    print('eingepraegt DS', me.find_dynamicsymbols(eingepraegt))\n",
    "    print('eingepraegt free symbols', eingepraegt.free_symbols)\n",
    "    print('eingepraegt has {} operations'.format(sum([eingepraegt[i].count_ops(visual=False) for i in range(len(eingepraegt))])), '\\n')\n",
    "print(f'it took {time.time() - start:.3f} sec to set up Kanes equations for {n} links')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Functions for the kinetic and the potential energies.\\\n",
    "Always useful to detect mistakes.\\\n",
    "orte is needed for the animation only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "kin_energie = sum([koerper.kinetic_energy(O) for koerper in BODY]).subs({i: 0. for i in aux})\n",
    "pot_energie = sum([m*g*me.dot(koerper.pos_from(PO), O.y) for koerper in Dmc])\n",
    "\n",
    "orte        = [[me.dot(P[i].pos_from(PO), uv) for uv in (O.x, O.y)] for i in range(n) ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lambdification.\\\n",
    "As mentioned above, it takes quite long to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start3 = time.time()\n",
    "qL = q + u\n",
    "pL = [m, g, l, iZZ, reibung]\n",
    "F  = [fx, fy]\n",
    "\n",
    "MM_lam          = sm.lambdify(qL + pL, MM, cse=True)\n",
    "force_lam       = sm.lambdify(qL + pL, force, cse=True)\n",
    "eingepraegt_lam = sm.lambdify(F + qL + pL + rhs, eingepraegt, cse=True)\n",
    "\n",
    "kin_lam = sm.lambdify(qL + pL, kin_energie, cse=True)\n",
    "pot_lam = sm.lambdify(qL + pL, pot_energie, cse=True)\n",
    "ort_lam = sm.lambdify(qL + pL, orte, cse=True)\n",
    "\n",
    "print(f'it took {time.time()-start3:.3f} sec to do the lambdification')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Numerical Integration**\n",
    "\n",
    "Nothing special here.\\\n",
    "method='Radau' in solve_ivp gives a more constant total energy if $reibung = 0.$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start2 = time.time()\n",
    "#==========================\n",
    "# Input variables\n",
    "#==========================\n",
    "m1       = 1.\n",
    "g1       = 9.8\n",
    "l1       = 20\n",
    "reibung1 = 0.\n",
    "\n",
    "q1      = [3.*np.pi/2. + np.pi * i / n for i in range(1, n+1)]\n",
    "u1      = [0. for _ in range(n)]\n",
    "\n",
    "intervall = 10.\n",
    "punkte    = 50\n",
    "#==========================\n",
    "\n",
    "schritte = int(intervall * punkte)\n",
    "times    = np.linspace(0., intervall, schritte)\n",
    "t_span   = (0., intervall)\n",
    "\n",
    "iZZ1 = 1./12. * m1 * (l1/n)**2        # from the internet\n",
    "\n",
    "pL_vals = [m1, g1, l1, iZZ1, reibung1]\n",
    "y0      =[*q1, *u1]\n",
    "\n",
    "\n",
    "def gradient(t, y, args):\n",
    "    sol = np.linalg.solve(MM_lam(*y, *args), force_lam(*y, *args))\n",
    "    return np.array(sol).T[0]\n",
    "        \n",
    "resultat1 = solve_ivp(gradient, t_span, y0, t_eval = times, args=(pL_vals,), method='Radau')\n",
    "#                , atol=1.e-3, rtol=1.e-3)\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(\"To numerically integrate an intervall of {} sec the routine cycled {} times and it took {:.3f} sec \"\n",
    "      .format(intervall, resultat1.nfev, time.time() - start2), '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the energies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kin_np   = np.empty(schritte)\n",
    "pot_np   = np.empty(schritte)\n",
    "total_np = np.empty(schritte)\n",
    "\n",
    "for i in range(schritte):\n",
    "    kin_np[i] = kin_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)\n",
    "    pot_np[i] = pot_lam(*[resultat[i, j] for j in range(resultat.shape[1])], *pL_vals)\n",
    "    total_np[i] = kin_np[i] + pot_np[i]\n",
    "\n",
    "if reibung1 == 0.:\n",
    "    max_total = np.max(np.abs(total_np))\n",
    "    min_total = np.min(np.abs(total_np))\n",
    "    delta = max_total - min_total\n",
    "    print(f'max deviation of total energy from zero is {delta/max_total * 100:.3e} % of max. total energy')\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "for i, j in zip((kin_np, pot_np, total_np), ('kinetic energy', 'potential energy', 'total energy')):\n",
    "    ax.plot(times, i, label=j)\n",
    "ax.set_title(\"Energies of the system\")\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('energy (Nm)')\n",
    "ax.legend();\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Reaction force on $PO$**\n",
    "\n",
    "First I solve the equation $rhs = MM^{-1} \\cdot force$ numerically.\\\n",
    "Then I solve $eingepraegt = 0$ for $fx, fy$ numerically.\\\n",
    "I iterate over fsolve twice, this sometimes improves the accuracy.\\\n",
    "I iterate over the whole process twice: Sometimes there is a message that convergence is slow. If I run it again, usually this message does not come again. I do not know, why.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RHS calculated numerically, too large to do it symbolically. Needed for reaction forces.\n",
    "\n",
    "# this is only needed because I basically copied it from some other program, where resultat2, etc \n",
    "# was different from resultat.\n",
    "resultat2 = resultat\n",
    "schritte2 = schritte\n",
    "times2    = times\n",
    "\n",
    "RHS1 = np.empty((schritte2, resultat2.shape[1]))\n",
    "for i in range(schritte2):\n",
    "    RHS1[i, :] = np.linalg.solve(MM_lam(*[resultat2[i, j]for j in range(resultat2.shape[1])], *pL_vals), \n",
    "        force_lam(*[resultat2[i, j] for j in range(resultat2.shape[1])], \n",
    "        *pL_vals)).reshape(resultat2.shape[1])\n",
    "print('RHS1 shape', RHS1.shape)\n",
    "\n",
    "#calculate implied forces numerically\n",
    "def func (x, *args):\n",
    "# just serves to 'modify' the arguments for fsolve.\n",
    "    return eingepraegt_lam(*x, *args).reshape(2)\n",
    "\n",
    "for _ in range(2):\n",
    "    kraftx = np.empty(schritte2)\n",
    "    krafty = np.empty(schritte2)\n",
    "    x0 = tuple((1., 1.))   # initial guess\n",
    "\n",
    "    for i in range(schritte2):\n",
    "        for _ in range(2):\n",
    "            y0 = [resultat2[i, j] for j in range(resultat2.shape[1])]\n",
    "            rhs = [RHS1[i, j] for j in range(n, 2*n)]\n",
    "            args = tuple(y0 + pL_vals + rhs)\n",
    "            AAA = fsolve(func, x0, args=args).reshape(2)\n",
    "            x0 = tuple(AAA)      # improved guess. Should speed up convergence of fsolve\n",
    "        kraftx[i] = AAA[0]\n",
    "        krafty[i] = AAA[1]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 5))\n",
    "plt.plot(times2, kraftx, label='X force')\n",
    "plt.plot(times2, krafty, label='Y force')\n",
    "ax.set_title('Reaction Forces on suspension point PO')\n",
    "ax.set_xlabel('time (sec)')\n",
    "ax.set_ylabel('force (N)')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Animation**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# get Carthesian coordinates\n",
    "x_coords = np.empty((len(times), n+1))\n",
    "y_coords = np.empty((len(times), n+1))\n",
    "\n",
    "for j in range(schritte):\n",
    "    x_coords[j] = [0.] + [ort_lam(*[resultat[j, k] for k in range(resultat.shape[1])], *pL_vals)[k11][0] \n",
    "                         for k11 in range(n) ]\n",
    "    y_coords[j] = [0.] + [ort_lam(*[resultat[j, k] for k in range(resultat.shape[1])], *pL_vals)[k11][1] \n",
    "                         for k11 in range(n)]\n",
    "\n",
    "max_x = max([abs(x_coords[i, j])for i in range(schritte) for j in range(n+1)])\n",
    "max_y = max([abs(y_coords[i, j])for i in range(schritte) for j in range(n+1)])\n",
    "max_xy = max(max_x, max_y)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 7))\n",
    "ax.axis('on')\n",
    "ax.set(xlim=(-max_xy-1., max_xy+1.), ylim=(-max_xy-1., max_xy+1.))\n",
    "ax.plot(0., 0.4, marker='v', markersize=15, color='red')\n",
    "\n",
    "Lines = []\n",
    "for i in range(n):     # sets the points. Red points are fixed\n",
    "    line, = ax.plot([], [], 'o-', lw=0.5, color='blue', markersize=0.)\n",
    "    Lines.append(line)\n",
    "\n",
    "line, = ax.plot([], [], 'o-', lw=0.5, color='blue', markersize=0) #connects the points\n",
    "Lines.append(line)\n",
    "   \n",
    "def animate_pendulum(times, x, y):    \n",
    "    def animate(i):\n",
    "        ax.set_title('Running time is {:.2f} sec'.format(i / schritte * intervall), fontsize=12)\n",
    "        for j in range(n+1):\n",
    "            Lines[j].set_data(x[i, j], y[i, j])\n",
    "        Lines[-1].set_data(x[i], y[i])\n",
    "        return Lines\n",
    "\n",
    "    anim = animation.FuncAnimation(fig, animate, frames=len(times),\n",
    "                                   interval=1000*times.max() / len(times),\n",
    "                                   blit=True) #, init_func=init)\n",
    "    plt.close(fig)\n",
    "    return anim\n",
    "\n",
    "anim = animate_pendulum(times, x_coords, y_coords)\n",
    "print('to run the program before HTML(..) took {:.3f} sec'.format(time.time() - start))\n",
    "HTML(anim.to_jshtml())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}