Peter Stahlecker's example that has a large number of operations

Frame_test_n_pendulum.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of operations in rhs[-2] is 13235\n"
     ]
    }
   ],
   "source": [
    "#============================\n",
    "# n body pendulum.\n",
    "# it seems, that if I rotate the bodies using the 'Body' method is MUCH slower that using two intermediate frames with the 'Axis' method\n",
    "#============================\n",
    "\n",
    "from sympy.physics.mechanics import *\n",
    "import sympy\n",
    "from sympy import Dummy, lambdify, symbols, sin, cos, simplify, Matrix, solve, count_ops\n",
    "from scipy.integrate import odeint, solve_ivp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "from IPython.display import HTML\n",
    "import matplotlib\n",
    "import time\n",
    "matplotlib.rcParams['animation.embed_limit'] = 2**128\n",
    "\n",
    "#Konstanten======\n",
    "n = 1                                 # number of pendulum bodies\n",
    "\n",
    "\n",
    "#----------------------------------------------------------------------------------------------------------------------------------\n",
    "start = time.time()\n",
    "\n",
    "# Generalized coordinates and velocities\n",
    "q = []\n",
    "qd = []\n",
    "u = []\n",
    "ud = []\n",
    "for i in range(1, n+1):                                     #Notice: data of body i in list location i-1\n",
    "    for j in ('x', 'y', 'z'):\n",
    "        q.append(dynamicsymbols('q'+str(i) + str(j)))          #Rotation of Frame i relative to inertial frame N\n",
    "        u.append(dynamicsymbols('u'+str(i) + str(j)))          \n",
    "        qd.append(dynamicsymbols('q'+str(i) + str(j), 1))      \n",
    "        ud.append(dynamicsymbols('u'+str(i) + str(j), 1))\n",
    "\n",
    "ux, uy, uz, f1, f2, f3 = dynamicsymbols('ux, uy, uz, f1, f2, f3')     # needed to calculate the forces at the fixed point of the pendulum\n",
    "\n",
    "g, l, m, reibung, t = symbols('g, l, m, reibung, t')                  # Reibung is German for friction\n",
    "r = symbols('r')                                                \n",
    "iXX, iYY, iZZ = symbols('iXX, iYY, iZZ')                        \n",
    "\n",
    "#Reference frame, reference point\n",
    "N = ReferenceFrame('N')\n",
    "P0 = Point('P0')\n",
    "P01 = P0.locatenew('P01', 0)\n",
    "P01.set_vel(N, ux*N.x + uy*N.y + uz*N.z)\n",
    "\n",
    "#Frame points, bodies for each Pendulum, numbered 1, 2, ..., n\n",
    "rot = [0]\n",
    "rot1 = [0]\n",
    "A = [N]\n",
    "P = [P01]\n",
    "DMC = [0]       #center of grvity of each body relative to Pi\n",
    "BODY = [Particle('P01a', P01, m)]\n",
    "\n",
    "#two versions to rotate each body: intermediate frames or all in once. Using intermediate frames run much faster here\n",
    "\n",
    "# version one: intermediate frames\n",
    "for i in range(1, n+1):\n",
    "    Lx = N.orientnew('Lx', 'Axis', [q[3*i-2], N.y])\n",
    "    Ly = Lx.orientnew('Ly', 'axis', [q[3*i-3], Lx.x])\n",
    "    Ai = Ly.orientnew('A' + str(i), 'Axis', [q[3*i-1], Ly.z])\n",
    "\n",
    "# version two: all in once\n",
    "#Ai = N.orientnew('A' + str(i), 'Body', (q[3*i-2], q[3*i-3], q[3*i-1]), '213')\n",
    "    \n",
    "    rot.append(Ai.ang_vel_in(N))\n",
    "    Ai.set_ang_vel(N, u[3*i-3] * N.x + u[3*i-2] * N.y + u[3*i-1] * N.z )\n",
    "    rot1.append(Ai.ang_vel_in(N))\n",
    "    A.append(Ai)\n",
    "\n",
    "    Pi = P[i-1].locatenew('P' + str(i), l * A[i].y)\n",
    "    Pi.v2pt_theory(P[i-1], N, A[i])\n",
    "    P.append(Pi)\n",
    "    Dmci = P[i-1].locatenew('Dmc' + str(i), r * A[i].y)         #Schwerpunktslage in A[i]\n",
    "    Dmci.v2pt_theory(P[i-1], N, A[i])\n",
    "    DMC.append(Dmci)\n",
    "\n",
    "    I = inertia(A[i], iXX, iYY, iZZ)\n",
    "    Bodyi = RigidBody('Body' + str(i), Dmci, A[i], m, (I, Dmci))\n",
    "    BODY.append(Bodyi)   \n",
    "    \n",
    "#Kinematic differential equations\n",
    "kd = []               \n",
    "for i in range(1, n+1):\n",
    "    for uv in N:\n",
    "        kd.append(dot(rot[i] - rot1[i], uv))\n",
    "\n",
    "FL = [(P01, -m*g*N.y + f1*N.x + f2*N.y + f3*N.z)]       # to calculate forces in the fixed point, of no interest to myissue\n",
    "\n",
    "for i in range(1, n+1):\n",
    "    kraft_auf_punkt = (DMC[i], -m*g*N.y)\n",
    "    torque_on_body = (A[i], -reibung * (dot(A[i].ang_vel_in(N),A[i].x)*A[i].x + dot(A[i].ang_vel_in(N),A[i].y)*A[i].y + dot(A[i].ang_vel_in(N),A[i].z)*A[i].z))\n",
    "    FL.append(kraft_auf_punkt)\n",
    "    FL.append(torque_on_body)\n",
    "\n",
    "aux = [ux, uy, uz]\n",
    "\n",
    "# equations of motion\n",
    "KM = KanesMethod(N, q_ind=q, u_ind=u, u_auxiliary=aux, kd_eqs=kd)\n",
    "(fr, frstar) = KM.kanes_equations(BODY, FL)\n",
    "\n",
    "MM = KM.mass_matrix_full\n",
    "force = KM.forcing_full\n",
    "rhs = KM.rhs()\n",
    "AAA = rhs[-2].count_ops(visual=False)      #calculate the number of operation in one element of rhs\n",
    "print('number of operations in rhs[-2] is {}'.format(AAA))\n",
    "\n",
    "\n",
    "# forces on the fixed point\n",
    "ersetze = [i.diff(t) for i in u]\n",
    "subs_dict = {}\n",
    "zz = 3 * n\n",
    "for i in range(3*n):\n",
    "    subs_dict.update({ersetze[i]: rhs[-zz]})    # replace accelerations in ersetze with the respective element of rhs\n",
    "    zz -= 1\n",
    "eingepraegt = KM.auxiliary_eqs                        \n",
    "\n",
    "eingepraegt = solve(eingepraegt, (f1, f2, f3))        \n",
    "schluessel = list(eingepraegt.keys())                 # to simply adress the keys in that dictionary\n",
    "\n",
    "eingepraegt_x = eingepraegt[schluessel[0]].subs(subs_dict)\n",
    "eingepraegt_y = eingepraegt[schluessel[1]].subs(subs_dict)\n",
    "eingepraegt_z = eingepraegt[schluessel[2]].subs(subs_dict)\n",
    "\n",
    "# Lambdification\n",
    "qL = q + u\n",
    "pL = [g, m, l, r, iXX, iYY, iZZ, reibung]                                          \n",
    "\n",
    "MM_lam = lambdify(qL + pL, MM)\n",
    "force_lam = lambdify(qL + pL, force)\n",
    "eingepraegt_x_lam = lambdify(qL + pL, eingepraegt_x)\n",
    "eingepraegt_y_lam = lambdify(qL + pL, eingepraegt_y)\n",
    "eingepraegt_z_lam = lambdify(qL + pL, eingepraegt_z)\n",
    "\n",
    "\n",
    "print(\"it took {:.3f} sec to establish Kane's equations\".format(time.time() - start))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Numerisches Integrieren\n",
    "laenge = 25/n\n",
    "start = time.time()\n",
    "\n",
    "intervall = 25\n",
    "schritte = 500\n",
    "times = np.linspace(0, intervall, schritte)\n",
    "\n",
    "#Anfangsbedingungen\n",
    "#qL = [q0x, q0z] + q + [u0x, u0z] + u\n",
    "#pL = [g, m, l, r, iXX, iYY, iZZ, reibung]                 #r ist Abstand vom Schwerpunkt zum Aufhängepunkt                          \n",
    "pL_vals = [9.8, 1, 10, 5, 12, 24, 12., 10.]                   #reibung ist Dämpfung in den Gelenken\n",
    "\n",
    "y0 = [0, 0, 0] * int(len(q)/3) + [0., 3., 1.] * int(len(u)/3)\n",
    "                                            \n",
    "\n",
    "zaehler = 0\n",
    "def gradient(y, t, args):                           #ich denke, das erzwingt, dass y, t an den richtigen Plätzen für odeint sind.\n",
    "    global zaehler\n",
    "    zaehler += 1\n",
    "    vals = np.concatenate((y, args))\n",
    "    sol = np.linalg.solve(MM_lam(*vals), force_lam(*vals))\n",
    "    return np.array(sol).T[0]\n",
    "    \n",
    "resultat = odeint(gradient, y0, times, args=(pL_vals,))\n",
    "print(resultat.shape)\n",
    "print('to calculate an intervall of {} sec it took {} loops and {:.3f} sec running time'.format(intervall, zaehler, time.time() - start))\n",
    "#==================================================================================\n",
    "#drucken der Koordinaten\n",
    "fig, ax = plt.subplots(figsize=(12, 6))\n",
    "for i in range(6*n):\n",
    "    ax.plot(times, resultat[:, i], label='Koordinate {}'.format(i+1))\n",
    "ax.legend();\n",
    "\n",
    "\n",
    "#Drucken der eingeprägten Kräfte\n",
    "kraft_x = np.zeros(schritte)\n",
    "kraft_y = np.zeros(schritte)\n",
    "kraft_z = np.zeros(schritte)\n",
    "for i in range(schritte):\n",
    "    kraft_x[i] = eingepraegt_x_lam(*[resultat[i, j] for j in range(6*n)], *pL_vals)\n",
    "    kraft_y[i] = eingepraegt_y_lam(*[resultat[i, j] for j in range(6*n)], *pL_vals)\n",
    "    kraft_z[i] = eingepraegt_z_lam(*[resultat[i, j] for j in range(6*n)], *pL_vals)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(12,6))\n",
    "plt.plot(times, kraft_x, label='X force')\n",
    "plt.plot(times, kraft_y, label='Y force')\n",
    "plt.plot(times, kraft_z, label='Z force')\n",
    "plt.legend();\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.6+"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}