Chapter-24-Part-11-Rolling-Disk.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4b06bc6e",
   "metadata": {},
   "source": [
    "## Rolling Disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e18de4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy import *\n",
    "import sympy.physics.mechanics as me\n",
    "me.init_vprinting()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcb0e079",
   "metadata": {},
   "outputs": [],
   "source": [
    "psi, theta, phi, xc, yc = coords = me.dynamicsymbols(\"psi, theta, phi, x_c, y_c\")\n",
    "psid, thetad, phid, xcd, ycd = me.dynamicsymbols(\"psi, theta, phi, x_c, y_c\", 1)\n",
    "omega1, omega2, omega3, uc, vc = speeds = me.dynamicsymbols(\"omega1, omega2, omega3, u_c, v_c\")\n",
    "m, r, g = symbols(\"m, r, g\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c2ee4f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "A, B, C, D = symbols(\"A, B, C, D\", cls=me.ReferenceFrame)\n",
    "B.orient_axis(A, A.z, psi)\n",
    "C.orient_axis(B, B.x, theta)\n",
    "D.orient_axis(C, C.y, phi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "646c647e-a7fd-43a6-97e7-24b8d2643976",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = symbols(\"t\")\n",
    "\n",
    "O = me.Point(\"O\")\n",
    "Cp = O.locatenew(\"Cp\", xc * A.x + yc * A.y)\n",
    "G = Cp.locatenew(\"G\", r * C.z)\n",
    "Q = G.locatenew(\"G\", -r * C.z)\n",
    "\n",
    "O.set_vel(A, 0)\n",
    "Cp.set_vel(B, 0)\n",
    "Cp.set_vel(A, Cp.pos_from(O).diff(t, A))\n",
    "G.set_vel(C, 0)\n",
    "G.set_vel(D, 0)\n",
    "Q.v2pt_theory(G, A, D)\n",
    "Q.a2pt_theory(G, A, D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe484474-0676-4c81-b419-0e4f5dc886c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# angular velocity of the disk on A\n",
    "A_w_D = D.ang_vel_in(A).express(D)\n",
    "A_w_D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76cbf293",
   "metadata": {},
   "outputs": [],
   "source": [
    "rot_kin = A_w_D - (omega1 * D.x + omega2 * D.y + omega3 * D.z)\n",
    "rot_kin = list(rot_kin.to_matrix(D))\n",
    "transl_kin = [xcd - uc, ycd - vc]\n",
    "kdes = transl_kin + rot_kin\n",
    "display(*kdes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d71e83-4b2f-4774-9b76-044c6419146b",
   "metadata": {},
   "outputs": [],
   "source": [
    "v_zero = Q.vel(A)\n",
    "f_v = [v_zero.dot(A.x), v_zero.dot(A.y)]\n",
    "display(*f_v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc6789ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "disc_inertia = m * r ** 2 / 4 * me.inertia(D, 1, 2, 1)\n",
    "disc = me.RigidBody(\"Disc\", G, D, m, (disc_inertia, G))\n",
    "bodies = [disc]\n",
    "loads = [(G, -m * g * A.z)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "198c0ee0",
   "metadata": {},
   "outputs": [],
   "source": [
    "kane = me.KanesMethod(\n",
    "    A,\n",
    "    q_ind=[psi, theta, phi, xc, yc], \n",
    "    u_ind=[omega1, omega2, omega3],\n",
    "    u_dependent=[uc, vc],\n",
    "    kd_eqs=kdes,\n",
    "    velocity_constraints=f_v\n",
    ")\n",
    "(fr, frstar) = kane.kanes_equations(bodies, loads)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f1bd707",
   "metadata": {},
   "source": [
    "### Numerical integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8db2b99d-abc7-4245-a577-5a09efc8fd07",
   "metadata": {},
   "outputs": [],
   "source": [
    "kane.kindiffdict()[phid]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28837334",
   "metadata": {},
   "outputs": [],
   "source": [
    "sd = {\n",
    "    g: 9.81, r: 0.5, m: 5,\n",
    "    # initial condition for positions\n",
    "    psi: 0, theta: pi / 20, phi: 0.9 * pi / 2, xc: 0, yc: 0,\n",
    "    # initial condition for independent velocities\n",
    "    omega1: 0.5, omega2: 2, omega3: 0.5\n",
    "}\n",
    "kindiffdict = {k: v.subs(sd).n() for k, v in kane.kindiffdict().items()}\n",
    "kindiffdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1519423",
   "metadata": {},
   "outputs": [],
   "source": [
    "uc_vc = solve(Matrix(f_v).subs(kindiffdict).subs(sd), [uc, vc])\n",
    "uc_vc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c94d2893",
   "metadata": {},
   "outputs": [],
   "source": [
    "sd.update(uc_vc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd803036",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydy.system import System\n",
    "\n",
    "sys = System(kane)\n",
    "sys.constants = {m: sd[m], r: sd[r], g: sd[g]}\n",
    "for s in [m, r, g]:\n",
    "    sd.pop(s)\n",
    "sys.initial_conditions = sd\n",
    "\n",
    "fps = 60\n",
    "t0, tf = 0, 15\n",
    "n = int(fps * (tf - t0))\n",
    "sys.times = np.linspace(t0, tf, n)\n",
    "results = sys.integrate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27122d3d-a751-40f3-a28b-bc288ddf2161",
   "metadata": {},
   "source": [
    "### Extraction of Information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a328fe0-b281-4ede-acab-23d4a4879344",
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a5c2e52-36df-4d88-bfee-4a976c7c7a47",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, ax = plt.subplots(2, 1, layout=\"constrained\", sharex=True)\n",
    "ax[0].plot(sys.times, results[:, : len(coords)])\n",
    "ax[0].legend([\"$%s$\" % latex(c) for c in coords])\n",
    "ax[0].set_ylabel(\"Generalized Coordinates\")\n",
    "ax[1].plot(sys.times, results[:, len(coords) :])\n",
    "ax[1].legend([\"$%s$\" % latex(c) for c in speeds])\n",
    "ax[1].set_xlabel(\"Time [s]\")\n",
    "ax[1].set_ylabel(\"Generalized Speeds\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "543971f2-29fc-42b6-ad5c-5e3a089ed68a",
   "metadata": {},
   "source": [
    "#### Positions of points $C_{p}$ and $G$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbcfd173-6c55-4b31-86c9-59e123aa226e",
   "metadata": {},
   "outputs": [],
   "source": [
    "OG = G.pos_from(O).to_matrix(A)\n",
    "G_func = lambdify(coords + [r], OG)\n",
    "G_coords = G_func(*results[:, : len(coords)].T, sys.constants[r])\n",
    "G_coords = G_coords.squeeze().T\n",
    "\n",
    "fig, ax = plt.subplots(2, 1, figsize=(5, 8), height_ratios=[1, 0.25], layout=\"constrained\")\n",
    "ax[0].plot(results[:, 3], results[:, 4], label=\"$C_{p}$\")\n",
    "ax[0].plot(G_coords[:, 0], G_coords[:, 1], label=\"$G$\")\n",
    "ax[0].set_aspect(\"equal\")\n",
    "ax[0].set_xlabel(\"x [m]\")\n",
    "ax[0].set_ylabel(\"y [m]\")\n",
    "ax[0].set_title(\"Position on the horizontal plane\")\n",
    "ax[0].legend(loc=\"upper left\")\n",
    "\n",
    "ax[1].plot(sys.times, G_coords[:, 2])\n",
    "ax[1].set_xlabel(\"Time [s]\")\n",
    "ax[1].set_ylabel(\"$z_{G}$ [m]\")\n",
    "ax[1].set_title(\"Height of center of mass\")\n",
    "ax[1].set_ylim(0, 0.55)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "299c2af0-d8d3-4426-87ca-5c8bceb4f99f",
   "metadata": {},
   "source": [
    "#### Time derivatives of the generalized coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c00ea868-736e-4d9a-a202-4413e32473c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy.physics.mechanics import vlatex\n",
    "\n",
    "arguments = coords + speeds\n",
    "num_functions = {\n",
    "    k: lambdify(arguments, v, cse=True) for k, v in kane.kindiffdict().items()\n",
    "}\n",
    "\n",
    "fig, ax = plt.subplots(layout=\"constrained\")\n",
    "for k, v in num_functions.items():\n",
    "    ax.plot(sys.times, v(*results.T), label=\"$%s$\" % vlatex(k))\n",
    "ax.legend()\n",
    "ax.set_xlabel(\"Time [s]\")\n",
    "ax.set_title(\"Time derivative of the generalized coordinates\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea73c6d1-0be9-4ade-b7c6-1d333c8c3956",
   "metadata": {},
   "source": [
    "#### Acceleration of Q in A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a7fd46d-74cc-4f0d-bc46-1ed911e455f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "a_Q = Q.acc(A).to_matrix(A)\n",
    "a_Q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd26cad0-27de-4712-bc3f-5c78c89d5183",
   "metadata": {},
   "outputs": [],
   "source": [
    "me.find_dynamicsymbols(a_Q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b026af85-08f2-4984-bda8-6d9d90493f5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "a_Q = a_Q.subs(kane.kindiffdict()).doit()\n",
    "a_Q = a_Q.subs(kane.kindiffdict()).doit()\n",
    "a_Q = a_Q.simplify()\n",
    "a_Q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96481f1a-90c3-427e-8d9f-1b1fa876c319",
   "metadata": {},
   "outputs": [],
   "source": [
    "me.find_dynamicsymbols(a_Q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d38c35e-61d5-4353-90a3-ea66319d9a90",
   "metadata": {},
   "outputs": [],
   "source": [
    "arguments = coords + speeds + list(sys.constants.keys())\n",
    "mm_func = lambdify(arguments, kane.mass_matrix_full, cse=True)\n",
    "fm_func = lambdify(arguments, kane.forcing_full, cse=True)\n",
    "\n",
    "states_dot = np.zeros_like(results)\n",
    "for i in range(results.shape[0]):\n",
    "    arguments_values = list(results[i, :]) + list(sys.constants.values())\n",
    "    states_dot[i, :] = np.linalg.solve(\n",
    "        mm_func(*arguments_values),\n",
    "        fm_func(*arguments_values)\n",
    "    ).flatten()\n",
    "\n",
    "new_arguments = sys.states + [f.diff(t) for f in sys.states]  + list(sys.constants.keys())\n",
    "a_Q_func = lambdify(new_arguments, a_Q, cse=False)\n",
    "\n",
    "a_Q_res = np.zeros((results.shape[0], 3))\n",
    "for i in range(results.shape[0]):\n",
    "    values = list(results[i, :]) + list(states_dot[i, :]) + list(sys.constants.values())\n",
    "    a_Q_res[i, :] = a_Q_func(*values).flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c768f16-ac53-4e70-b9da-a17730c6ef68",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(sys.times, a_Q_res[:, 0], label=\"$a_{Q, x}$\")\n",
    "ax.plot(sys.times, a_Q_res[:, 1], label=\"$a_{Q, y}$\")\n",
    "ax.plot(sys.times, a_Q_res[:, 2], label=\"$a_{Q, z}$\")\n",
    "ax.legend()\n",
    "ax.set_xlabel(\"Time [s]\")\n",
    "ax.set_ylabel(r\"Acceleration $\\left[m / s^{2}\\right]$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb8fd8b4",
   "metadata": {},
   "source": [
    "### Animation with PyDy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d523fe15",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydy.viz import Scene, Cylinder, VisualizationFrame, Sphere, Plane\n",
    "\n",
    "disc_geom = Cylinder(r / 10, r, name=\"disc\", color=\"orange\")\n",
    "plane_geom = Plane(50, 50, name=\"ground plane\", color=\"white\")\n",
    "center_of_mass_geom = Sphere(r / 10, color=\"black\", name=\"G\")\n",
    "contact_point_geom = Sphere(r / 10, color=\"blue\", name=\"CP\")\n",
    "\n",
    "# to make it easier to visualize the rolling, let's create\n",
    "# two rods, one aligned with the positive x-direction of\n",
    "# the disk, the other aligned with the positive z-direction.\n",
    "cylx = Cylinder(r, r / 15, name=\"cylx\", color=\"red\")\n",
    "cylz = Cylinder(r, r / 15, name=\"cylz\", color=\"green\")\n",
    "cylx_frame, cylz_frame = symbols(\"R1, R2\", cls=me.ReferenceFrame)\n",
    "cylx_frame.orient_axis(D, D.z, -pi/2)\n",
    "cylz_frame.orient_axis(D, D.x, -pi/2)\n",
    "\n",
    "rod1 = me.RigidBody('rod1', frame=cylx_frame)\n",
    "rod2 = me.RigidBody('rod2', frame=cylz_frame)\n",
    "rod1.masscenter.set_pos(disc.masscenter, r/2 * D.x)\n",
    "rod2.masscenter.set_pos(disc.masscenter, r/2 * D.z)\n",
    "\n",
    "vf_disc = VisualizationFrame(\"c\", D, G, disc_geom)\n",
    "vf_center_of_mass = VisualizationFrame(\"s1\", D, G, center_of_mass_geom)\n",
    "vf_contact_point = VisualizationFrame(\"s1\", A, Cp, contact_point_geom)\n",
    "vf_rod1 = VisualizationFrame(\"vf r1\", rod1.frame, rod1.masscenter, cylx)\n",
    "vf_rod2 = VisualizationFrame(\"vf r2\", rod2.frame, rod2.masscenter, cylz)\n",
    "vf_plane = VisualizationFrame(\"asd\", A, O, plane_geom)\n",
    "\n",
    "# because PyDy's default orientation is y-axis up, we need to\n",
    "# perform a rotation in order to get z-axis up\n",
    "N = me.ReferenceFrame(\"N\")\n",
    "N.orient_axis(A, A.x, pi/2)\n",
    "scene = Scene(\n",
    "    N, O,\n",
    "    vf_disc, vf_center_of_mass, vf_contact_point,\n",
    "    vf_plane, vf_rod1, vf_rod2,\n",
    "    system=sys)\n",
    "scene.display_jupyter(axes_arrow_length=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c282f09",
   "metadata": {},
   "source": [
    "### Animation with Matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30d34cfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import HTML\n",
    "\n",
    "mv, rv, gv = [float(sys.constants[t]) for t in [m, r, g]]\n",
    "psiv, thetav, phiv, xcv, ycv = results[0, :5]\n",
    "\n",
    "OG = lambdify([psi, theta, phi, xc, yc, r], G.pos_from(O).express(A).to_matrix(A))\n",
    "\n",
    "# circle perimeter\n",
    "angle = np.linspace(0, 2*np.pi, 100)\n",
    "xd = rv * np.cos(angle)\n",
    "zd = rv * np.sin(angle)\n",
    "yd = np.zeros_like(zd)\n",
    "disk_coords = np.stack([xd, yd, zd])\n",
    "\n",
    "# disk surface\n",
    "rr = np.linspace(0, rv, 2)\n",
    "rr, aa = np.meshgrid(rr, angle)\n",
    "xds = rr * np.cos(aa)\n",
    "zds = rr * np.sin(aa)\n",
    "yds = np.zeros_like(zds)\n",
    "shape = xds.shape\n",
    "disk_surf_coords = np.stack([xds.flatten(), yds.flatten(), zds.flatten()])\n",
    "\n",
    "# rotation matrix from ground frame to disk frame\n",
    "R_AD = A.dcm(D)\n",
    "R_AD = lambdify([psi, theta, phi], R_AD)\n",
    "\n",
    "new_disk_coords = (R_AD(psiv, thetav, phiv) @ disk_coords) + OG(psiv, thetav, phiv, xcv, ycv, rv)\n",
    "new_disk_surf_coords = (R_AD(psiv, thetav, phiv) @ disk_surf_coords) + OG(psiv, thetav, phiv, xcv, ycv, rv)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(projection=\"3d\")\n",
    "circle, = ax.plot(new_disk_coords[0, :], new_disk_coords[1, :], new_disk_coords[2, :], color=\"tab:blue\")\n",
    "# reference point fixed on the surface of the disk\n",
    "ref_point = ax.scatter(new_disk_coords[0, 0], new_disk_coords[1, 0], new_disk_coords[2, 0], color=\"r\")\n",
    "trajectory, = ax.plot(results[0, 3], results[0, 4], 0, color=\"tab:orange\")\n",
    "disk_surface = ax.plot_surface(\n",
    "    new_disk_surf_coords[0, :].reshape(shape),\n",
    "    new_disk_surf_coords[1, :].reshape(shape),\n",
    "    new_disk_surf_coords[2, :].reshape(shape),\n",
    "    alpha=0.5, color=\"tab:blue\"\n",
    ")\n",
    "l = 3\n",
    "ax.axis([xcv-l, xcv+l, ycv-l, ycv+l, -1, 3])\n",
    "ax.set_zlim(0, rv * 2.5)\n",
    "ax.set_zticks([0, rv * 2.5])\n",
    "ax.set_aspect(\"equal\")\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d25a1b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "matplotlib.rcParams[\"animation.embed_limit\"] = 200\n",
    "\n",
    "def update(idx):\n",
    "    global disk_surface\n",
    "    psiv, thetav, phiv, xcv, ycv = results[idx, :5]\n",
    "    # update disc position and orientation\n",
    "    new_disk_coords = (R_AD(psiv, thetav, phiv) @ disk_coords) + OG(psiv, thetav, phiv, xcv, ycv, rv)\n",
    "    circle.set_data_3d([new_disk_coords[0, :], new_disk_coords[1, :], new_disk_coords[2, :]])\n",
    "    # update trajectory\n",
    "    cur_x, cur_y, cur_z = trajectory.get_data_3d()\n",
    "    cur_x = np.append(cur_x, xcv)\n",
    "    cur_y = np.append(cur_y, ycv)\n",
    "    cur_z = np.append(cur_z, 0)\n",
    "    trajectory.set_data_3d([cur_x, cur_y, cur_z])\n",
    "    # scatter needs at least 2 points for update to work\n",
    "    points = np.stack([new_disk_coords[:, 0], new_disk_coords[:, 0]])\n",
    "    ref_point._offsets3d = (points[:, 0], points[:, 1], points[:, 2])\n",
    "    # disk surface\n",
    "    new_disk_surf_coords = (R_AD(psiv, thetav, phiv) @ disk_surf_coords) + OG(psiv, thetav, phiv, xcv, ycv, rv)\n",
    "    disk_surface.remove()\n",
    "    disk_surface = ax.plot_surface(\n",
    "        new_disk_surf_coords[0, :].reshape(shape),\n",
    "        new_disk_surf_coords[1, :].reshape(shape),\n",
    "        new_disk_surf_coords[2, :].reshape(shape),\n",
    "        alpha=0.5, color=\"tab:blue\"\n",
    "    )\n",
    "    ax.set_title(\"t = {:.2f} s\".format(sys.times[idx]))\n",
    "\n",
    "ani = FuncAnimation(fig, update, frames=len(sys.times))\n",
    "HTML(ani.to_jshtml(fps=fps))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "213396e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "ani.save(\"rolling-disk.mp4\", fps=60)"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}