Chapter-24-Part-10-Spinning-Top.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "473a8b86-2eb1-47dc-b8d2-3d52de83befd",
   "metadata": {},
   "source": [
    "## Spinning Top\n",
    "\n",
    "Notes: \n",
    "* the code below might be slightly different than what is presented on the book. Here, many functions were created in order to promote code reusability throughout this notebook.\n",
    "* the [SymPy Plotting Backends module](https://sympy-plot-backends.readthedocs.io/) and [K3D Jupyter](https://github.com/K3D-tools/K3D-jupyter) are required to visualize the results of the numerical integration in a 3D space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e2adf22-5b82-463d-9e98-9ee1ff5d6d18",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy import *\n",
    "from sympy.physics.mechanics import *\n",
    "init_vprinting()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f67aab36-4a33-4540-b19f-22a2409d7b32",
   "metadata": {},
   "outputs": [],
   "source": [
    "phi, theta, psi = dynamicsymbols(\"phi, theta, psi\")\n",
    "phid, thetad, psid = dynamicsymbols(\"phi, theta, psi\", 1)\n",
    "omega_x, omega_y, omega_z = dynamicsymbols(\"omega_x omega_y omega_z\")\n",
    "A, C, d, m, g = symbols(\"A, C, d, m, g\")\n",
    "\n",
    "N, B = symbols(\"N, B\", cls=ReferenceFrame)\n",
    "B.orient_body_fixed(N, [phi, theta, psi], 313)\n",
    "\n",
    "O, G = symbols(\"O, G\", cls=Point)\n",
    "G.set_pos(O, d * B.z)\n",
    "\n",
    "O.set_vel(N, 0)\n",
    "G.v2pt_theory(O, N, B)\n",
    "G.a2pt_theory(O, N, B)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1af5d51e-776e-4ec8-be2f-11b052695c61",
   "metadata": {},
   "source": [
    "## First choice of generalized speeds and KDEs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77903455-e4d4-46c4-a977-010ad5c9e358",
   "metadata": {},
   "outputs": [],
   "source": [
    "B_w_N = B.ang_vel_in(N).express(B)\n",
    "B_w_N"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f937ee23-2f22-48c5-b6f0-c2a9a2b300e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_body = omega_x * B.x + omega_y * B.y + omega_z * B.z\n",
    "rot_kin = B_w_N - w_body\n",
    "rot_kin = list(rot_kin.to_matrix(B))\n",
    "display(*rot_kin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8de0bf7f-32a9-40a5-82bc-1ff2c0d61d0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_inertia = inertia(B, A, A, C)\n",
    "top = RigidBody(\"Top\", G, B, m, (top_inertia, G))\n",
    "bodies = [top]\n",
    "loads = [(G, -m * g * N.z)]\n",
    "\n",
    "kane1 = KanesMethod(\n",
    "    N,\n",
    "    q_ind=[psi, theta, phi],\n",
    "    u_ind=[omega_x, omega_y, omega_z],\n",
    "    kd_eqs=rot_kin\n",
    ")\n",
    "fr, frstar = kane1.kanes_equations(bodies, loads)\n",
    "\n",
    "eom1 = fr + frstar\n",
    "eom1.simplify()\n",
    "eom1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "998437b2-e9b0-4c37-963a-573aa23d6c3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "angle_dots = {k: v.simplify() for k, v in kane1.kindiffdict().items()}\n",
    "angle_dots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13a87624-6cf2-4b02-aa88-7c779ed0c6c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "A_func = lambda m, h, r: 3*m * (h**2 + 4*r**2) / 80\n",
    "C_func =lambda m, h, r: 3*m*r**2 / 10\n",
    "\n",
    "r_val = 2.5e-02\n",
    "h_val = r_val * 4\n",
    "d_val = 3 * h_val / 4\n",
    "m_val = 0.5\n",
    "\n",
    "A_val = A_func(m_val, h_val, r_val)\n",
    "C_val = C_func(m_val, h_val, r_val)\n",
    "\n",
    "constants = {\n",
    "    m: m_val, g: 9.81, d: d_val,\n",
    "    A: A_val, C: C_val\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db786416-4fad-4628-9c00-a5c189a7919a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_func_to_integrate(kane, states, constants):\n",
    "    mm = kane.mass_matrix_full\n",
    "    fm = kane.forcing_full\n",
    "    mm.simplify()\n",
    "    fm.simplify()\n",
    "\n",
    "    display(mm, fm)\n",
    "    \n",
    "    arguments = states + list(constants.keys())\n",
    "    mm_func = lambdify(arguments, mm, cse=True)\n",
    "    fm_func = lambdify(arguments, fm, cse=True)\n",
    "    \n",
    "    def func_to_integrate(t, x, args):\n",
    "        values = np.concatenate((x, args))\n",
    "        solution = np.linalg.solve(\n",
    "            mm_func(*values),\n",
    "            fm_func(*values)\n",
    "        )\n",
    "        return solution.flatten()\n",
    "        \n",
    "    return func_to_integrate\n",
    "\n",
    "states = list(kane1.q) + list(kane1.u)\n",
    "func_to_integrate = create_func_to_integrate(kane1, states, constants)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9417a0c2-0243-43aa-b467-fa307e617b16",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.integrate import solve_ivp\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "def perform_integration(\n",
    "    func_to_integrate, t0, tf, states, initial_conditions,\n",
    "    constants, rtol=1e-03, atol=1e-06, method=\"RK45\", n=None\n",
    "):\n",
    "    t_span = [t0, tf]\n",
    "    t_eval = np.linspace(t0, tf, n) if n else n\n",
    "    x0 = [float(initial_conditions[k]) for k in states]\n",
    "\n",
    "    start_time = time.time()\n",
    "    results = solve_ivp(\n",
    "        func_to_integrate, t_span, x0,\n",
    "        args=(list(constants.values()), ),\n",
    "        method=method, t_eval=t_eval,\n",
    "        rtol=rtol, atol=atol\n",
    "    )\n",
    "    end_time = time.time()\n",
    "    print(\"Integration took about: %s seconds.\" % (end_time - start_time))\n",
    "    return results\n",
    "\n",
    "initial_conditions = {\n",
    "    phi: 0, theta: pi/10, psi: 0,\n",
    "    omega_x: 0, omega_y: 0, omega_z: 1000\n",
    "}\n",
    "results1 = perform_integration(\n",
    "    func_to_integrate, 0, 2,\n",
    "    states, initial_conditions, constants\n",
    ")\n",
    "results1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ad86993-63f8-4b65-b3cf-36aae1c6bc17",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def plot_states(results, states, angle_dots):\n",
    "    angle_dots_func = {k: lambdify(states, v) for k, v in angle_dots.items()}\n",
    "    angle_dots_res = {k: v(*results.y) for k, v in angle_dots_func.items()}\n",
    "    \n",
    "    fig, ax = plt.subplots(3, 3, sharex=True, layout=\"constrained\", figsize=(10, 5))\n",
    "    ax[0, 0].plot(results.t, results.y.T[:, 0])\n",
    "    ax[0, 0].set_ylabel(\"$%s$\" % vlatex(states[0]))\n",
    "    ax[0, 0].set_title(\"Angles [rad]\")\n",
    "    ax[1, 0].plot(results.t, results.y.T[:, 1])\n",
    "    ax[1, 0].set_ylabel(\"$%s$\" % vlatex(states[1]))\n",
    "    ax[2, 0].plot(results.t, results.y.T[:, 2])\n",
    "    ax[2, 0].set_ylabel(\"$%s$\" % vlatex(states[2]))\n",
    "    ax[2, 0].set_xlabel(\"Time [s]\")\n",
    "    ax[0, 1].plot(results.t, results.y.T[:, 3])\n",
    "    ax[0, 1].set_ylabel(\"$%s$\" % vlatex(states[3]))\n",
    "    ax[0, 1].set_title(\"Velocities [rad/s]\")\n",
    "    ax[1, 1].plot(results.t, results.y.T[:, 4])\n",
    "    ax[1, 1].set_ylabel(\"$%s$\" % vlatex(states[4]))\n",
    "    ax[2, 1].plot(results.t, results.y.T[:, 5])\n",
    "    ax[2, 1].set_ylabel(\"$%s$\" % vlatex(states[5]))\n",
    "    ax[2, 1].set_xlabel(\"Time [s]\")\n",
    "\n",
    "    t = symbols(\"t\")\n",
    "    for i, k in enumerate(states[:3]):\n",
    "        ax[i, 2].plot(results.t, angle_dots_res[k.diff(t)])\n",
    "        ax[i, 2].set_ylabel(\"$%s$\" % vlatex(k))\n",
    "    plt.show()\n",
    "\n",
    "plot_states(results1, states, angle_dots)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abcbeaf5-a9f7-40f3-b020-45001085c65f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compuse_free_end_position(states, results):\n",
    "    pos = B.z.express(N)\n",
    "    pos_components = [pos & vec for vec in N]\n",
    "    pos_components_funcs = [lambdify(states, comp) for comp in pos_components]\n",
    "    pos_components_res = [f(*results.y) for f in pos_components_funcs]\n",
    "    return pos_components_res\n",
    "\n",
    "pos_components_res = compuse_free_end_position(states, results1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d8c3a12-633d-4b63-8503-4ebc4ce689d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import k3d\n",
    "from spb.utils import get_vertices_indices\n",
    "\n",
    "def plot_free_end_trajectory_3d(pos_components_res, results):\n",
    "\n",
    "    theta_num, phi_num = np.mgrid[0:2*np.pi:50j, 0:np.pi:50j]\n",
    "    x = np.cos(theta_num) * np.sin(phi_num)\n",
    "    y = np.sin(theta_num) * np.sin(phi_num)\n",
    "    z = np.cos(phi_num)\n",
    "    vertices, indices = get_vertices_indices(x, y, z)\n",
    "    vertices = vertices.astype(np.float32)\n",
    "    \n",
    "    fig = k3d.plot(grid_visible=False)\n",
    "    mesh1 = k3d.mesh(vertices, indices, side=\"double\", opacity=0.2)\n",
    "    mesh2 = k3d.mesh(vertices, indices, side=\"double\", opacity=0.025, color=0, wireframe=True)\n",
    "    line_coordinates = np.vstack(pos_components_res).T.astype(np.float32)\n",
    "    path = k3d.line(\n",
    "        line_coordinates,\n",
    "        shader=\"simple\",\n",
    "        attribute=results.t\n",
    "    )\n",
    "    fig += mesh1\n",
    "    fig += mesh2\n",
    "    fig += path\n",
    "    fig.display()\n",
    "\n",
    "plot_free_end_trajectory_3d(pos_components_res, results1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "159b6d12-b3d6-4cc3-8e20-148d660a1992",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.collections import LineCollection\n",
    "\n",
    "# from Matplotlib's documentation:\n",
    "# https://matplotlib.org/stable/gallery/lines_bars_and_markers/multicolored_line.html\n",
    "def colored_line(x, y, c, ax, **lc_kwargs):\n",
    "    \"\"\"\n",
    "    Plot a line with a color specified along the line by a third value.\n",
    "\n",
    "    It does this by creating a collection of line segments. Each line segment is\n",
    "    made up of two straight lines each connecting the current (x, y) point to the\n",
    "    midpoints of the lines connecting the current point with its two neighbors.\n",
    "    This creates a smooth line with no gaps between the line segments.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x, y : array-like\n",
    "        The horizontal and vertical coordinates of the data points.\n",
    "    c : array-like\n",
    "        The color values, which should be the same size as x and y.\n",
    "    ax : Axes\n",
    "        Axis object on which to plot the colored line.\n",
    "    **lc_kwargs\n",
    "        Any additional arguments to pass to matplotlib.collections.LineCollection\n",
    "        constructor. This should not include the array keyword argument because\n",
    "        that is set to the color argument. If provided, it will be overridden.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    matplotlib.collections.LineCollection\n",
    "        The generated line collection representing the colored line.\n",
    "    \"\"\"\n",
    "    if \"array\" in lc_kwargs:\n",
    "        warnings.warn('The provided \"array\" keyword argument will be overridden')\n",
    "\n",
    "    # Default the capstyle to butt so that the line segments smoothly line up\n",
    "    default_kwargs = {\"capstyle\": \"butt\"}\n",
    "    default_kwargs.update(lc_kwargs)\n",
    "\n",
    "    # Compute the midpoints of the line segments. Include the first and last points\n",
    "    # twice so we don't need any special syntax later to handle them.\n",
    "    x = np.asarray(x)\n",
    "    y = np.asarray(y)\n",
    "    x_midpts = np.hstack((x[0], 0.5 * (x[1:] + x[:-1]), x[-1]))\n",
    "    y_midpts = np.hstack((y[0], 0.5 * (y[1:] + y[:-1]), y[-1]))\n",
    "\n",
    "    # Determine the start, middle, and end coordinate pair of each line segment.\n",
    "    # Use the reshape to add an extra dimension so each pair of points is in its\n",
    "    # own list. Then concatenate them to create:\n",
    "    # [\n",
    "    #   [(x1_start, y1_start), (x1_mid, y1_mid), (x1_end, y1_end)],\n",
    "    #   [(x2_start, y2_start), (x2_mid, y2_mid), (x2_end, y2_end)],\n",
    "    #   ...\n",
    "    # ]\n",
    "    coord_start = np.column_stack((x_midpts[:-1], y_midpts[:-1]))[:, np.newaxis, :]\n",
    "    coord_mid = np.column_stack((x, y))[:, np.newaxis, :]\n",
    "    coord_end = np.column_stack((x_midpts[1:], y_midpts[1:]))[:, np.newaxis, :]\n",
    "    segments = np.concatenate((coord_start, coord_mid, coord_end), axis=1)\n",
    "\n",
    "    lc = LineCollection(segments, **default_kwargs)\n",
    "    lc.set_array(c)  # set the colors of each segment\n",
    "\n",
    "    return ax.add_collection(lc)\n",
    "\n",
    "def cart2sph(x, y, z):\n",
    "    # convert cartesian coordinates to spherical coordinates\n",
    "    xy = np.sqrt(x**2 + y**2)\n",
    "    r = np.sqrt(x**2 + y**2 + z**2)\n",
    "    theta = np.arctan2(y, x) \n",
    "    phi = np.arctan2(xy, z) \n",
    "    return r, theta, phi\n",
    "\n",
    "def plot_free_end_trajectory_2d(pos_components_res, results):\n",
    "    # t: theta (longitude), p: phi for spherical to cylindrical projection.\n",
    "    # Not to be confused with theta and phi used in the generation of EoM\n",
    "    _, t, p = cart2sph(*pos_components_res)\n",
    "    t, p = [np.degrees(e) for e in [t, p]]\n",
    "    t_with_nan = t.copy()\n",
    "    \n",
    "    # t ranges from -180deg to +180deg. As t cross over +180deg and start\n",
    "    # the loop again at -180deg, Matplotlib will draw an horizontal line,\n",
    "    # which is not desired. Let's remove it by inserting a NaN value\n",
    "    # where this discontinuity happens.\n",
    "    for i in range(1, len(t)):\n",
    "        if abs(t[i] - t[i-1]) > 90:\n",
    "            t_with_nan[i] = np.nan\n",
    "    \n",
    "    fig, ax = plt.subplots(figsize=(10, 5))\n",
    "    cl = colored_line(t_with_nan, 90 - p, results.t, ax, cmap=\"plasma\")\n",
    "    fig.colorbar(cl, label=\"Time [s]\")\n",
    "    ax.set_xlabel(r\"$\\theta$ [deg]\")\n",
    "    ax.set_ylabel(r\"$\\phi$ [deg]\")\n",
    "    ax.set_xlim(-180, 180)\n",
    "    ax.set_ylim(55, 80)\n",
    "    ax.set_title(\"Drift \")\n",
    "    plt.show()\n",
    "\n",
    "plot_free_end_trajectory_2d(pos_components_res, results1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dccbe41-e54f-49da-b1ba-ef50b5911707",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_and_plot_energies(kane, states, constants, results):\n",
    "    vG = G.vel(N)\n",
    "    T = (B_w_N & top_inertia & B_w_N) / 2 + m * (vG & vG) / 2\n",
    "    T = T.subs(kane.kindiffdict()).simplify().doit()\n",
    "    U = m * g * (G.pos_from(O) & N.z)\n",
    "    display(T, U)\n",
    "\n",
    "    arguments = states + list(constants.keys())\n",
    "    T_func = lambdify(arguments, T)\n",
    "    U_func = lambdify(arguments, U)\n",
    "    T_res = T_func(*results.y, *list(constants.values()))\n",
    "    U_res = U_func(*results.y, *list(constants.values()))\n",
    "    \n",
    "    fig, ax = plt.subplots(3, 1, sharex=True)\n",
    "    ax[0].plot(results.t, T_res)\n",
    "    ax[0].set_ylabel(\"T [J]\")\n",
    "    ax[1].plot(results.t, U_res)\n",
    "    ax[1].set_ylabel(\"U [J]\")\n",
    "    ax[2].plot(results.t, T_res + U_res)\n",
    "    ax[2].set_xlabel(\"Times [s]\")\n",
    "    ax[2].set_ylabel(\"T+U [J]\")\n",
    "    fig.suptitle(\"Check for Energy Conservation\")\n",
    "    plt.show()\n",
    "\n",
    "    return T_res, U_res\n",
    "\n",
    "compute_and_plot_energies(kane1, states, constants, results1);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "250ebaca-2b30-4c7b-9a51-9adb959d5b7b",
   "metadata": {},
   "source": [
    "### Improving accuracy of the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ba1a6c6-4eda-4d81-8e39-f19f81334036",
   "metadata": {},
   "outputs": [],
   "source": [
    "results2 = perform_integration(\n",
    "    func_to_integrate, 0, 2,\n",
    "    states, initial_conditions, constants,\n",
    "    rtol=1e-06, atol=1e-09\n",
    ")\n",
    "pos_components_res2 = compuse_free_end_position(states, results2)\n",
    "plot_free_end_trajectory_3d(pos_components_res2, results2)\n",
    "plot_free_end_trajectory_2d(pos_components_res2, results2)\n",
    "T_res, U_res = compute_and_plot_energies(kane1, states, constants, results2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d723850-1bac-4d5a-8940-197721d1557e",
   "metadata": {},
   "source": [
    "## Second choice of generalized speeds and KDEs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1ca62b3-e4d7-4590-bf05-47a63cfc5d07",
   "metadata": {},
   "outputs": [],
   "source": [
    "omega_1, omega_2, omega_3 = dynamicsymbols(\"omega_1, omega_2, omega_3\")\n",
    "kd_eqs = [\n",
    "    omega_1 - phid, omega_2 - thetad, omega_3 - psid\n",
    "]\n",
    "\n",
    "kane3 = KanesMethod(\n",
    "    N,\n",
    "    q_ind=[psi, theta, phi],\n",
    "    u_ind=[omega_1, omega_2, omega_3],\n",
    "    kd_eqs=kd_eqs\n",
    ")\n",
    "fr, frstar = kane3.kanes_equations(bodies, loads)\n",
    "\n",
    "eom3 = fr + frstar\n",
    "eom3.simplify()\n",
    "eom3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df955ff9-3617-474d-b73f-da1020dc5386",
   "metadata": {},
   "outputs": [],
   "source": [
    "states3 = list(kane3.q) + list(kane3.u)\n",
    "func_to_integrate3 = create_func_to_integrate(kane3, states3, constants)\n",
    "\n",
    "initial_conditions3 = {\n",
    "    phi: 0, theta: pi/10, psi: 0,\n",
    "    omega_1: 0, omega_2: 0, omega_3: 1000\n",
    "}\n",
    "\n",
    "results3 = perform_integration(\n",
    "    func_to_integrate3, 0, 2,\n",
    "    states3, initial_conditions3, constants,\n",
    "    method=\"RK45\", n=2000,\n",
    "    # rtol=1e-06, atol=1e-09\n",
    ")\n",
    "pos_components_res3 = compuse_free_end_position(states3, results3)\n",
    "plot_free_end_trajectory_3d(pos_components_res3, results3)\n",
    "plot_free_end_trajectory_2d(pos_components_res3, results3)\n",
    "T_res3, U_res3 = compute_and_plot_energies(kane3, states3, constants, results3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbfb4014-c512-47b1-9fea-50d87bbabc3a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}