Chapter-24-Part-13.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4c9dce40-a9fb-42e9-b4e2-cedecddc4d38",
   "metadata": {},
   "source": [
    "## Reaction Wheel Pendulum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "233012d9-fa7f-4876-b9ff-a3351f49d8ca",
   "metadata": {},
   "source": [
    "### Equation of Motion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9418fa3c-7092-4a1a-8d20-add2730bed8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy import *\n",
    "from sympy.physics.mechanics import *\n",
    "init_vprinting()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2f4c853-8b59-428d-9063-e8915ecc910b",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha, theta = dynamicsymbols(\"alpha, theta\")\n",
    "alphad, thetad = dynamicsymbols(\"alpha, theta\", 1)\n",
    "u1, u2 = dynamicsymbols(\"u1, u2\")\n",
    "Tm = dynamicsymbols(\"T_m\")\n",
    "Lg, L, g, m_pend, m_disk = symbols(\"L_G, L, g, m_p, m_d\")\n",
    "\n",
    "N, A, B = symbols(\"N, A, B\", cls=ReferenceFrame)\n",
    "A.orient_axis(N, N.z, theta)\n",
    "B.orient_axis(A, A.z, alpha)\n",
    "\n",
    "# O: pin joint of the pendulum\n",
    "# G: center of mass of the rod\n",
    "# P: pin joint of the wheel\n",
    "O, P, G = symbols(\"O, P, G\", cls=Point)\n",
    "G.set_pos(O, -Lg * A.y)\n",
    "P.set_pos(O, -L * A.y)\n",
    "\n",
    "O.set_vel(N, 0)\n",
    "O.set_vel(A, 0)\n",
    "G.v2pt_theory(O, N, A)\n",
    "G.a2pt_theory(O, N, A)\n",
    "P.v2pt_theory(O, N, A)\n",
    "P.a2pt_theory(O, N, A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2897c7c3-d199-40a1-8db0-58882674178b",
   "metadata": {},
   "outputs": [],
   "source": [
    "Ixxp, Iyyp, Izzp = symbols(\"I_xx__p I_yy__p I_zz__p\")\n",
    "Ixxd, Iyyd, Izzd = symbols(\"I_xx__d I_yy__d I_zz__d\")\n",
    "# centrail inertias of the rod and disk\n",
    "inertia_pend = inertia(A, Ixxp, Iyyp, Izzp)\n",
    "inertia_disk = inertia(B, Ixxd, Iyyd, Izzd)\n",
    "\n",
    "bodies = [\n",
    "    RigidBody(\"pend\", G, A, m_pend, (inertia_pend, G)),\n",
    "    RigidBody(\"disk\", P, B, m_disk, (inertia_disk, P))\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd1a9c8a-a5ac-4448-ae1a-8c356bbda8b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "loads = [\n",
    "    (G, -m_pend * g * N.y),\n",
    "    (P, -m_disk * g * N.y),\n",
    "    (B, Tm * N.z),\n",
    "    (A, -Tm * N.z)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b71ab3e6-b5c2-4e22-8686-f5aa81d2780b",
   "metadata": {},
   "outputs": [],
   "source": [
    "kd_eqs = [\n",
    "    thetad - u1,\n",
    "    alphad - u2,\n",
    "]\n",
    "kane = KanesMethod(\n",
    "    N,\n",
    "    q_ind=[theta, alpha],\n",
    "    u_ind=[u1, u2],\n",
    "    kd_eqs=kd_eqs\n",
    ")\n",
    "\n",
    "fr, frstar = kane.kanes_equations(bodies, loads)\n",
    "fr + frstar"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19669276-d980-475d-ba49-3190fe0b6a9a",
   "metadata": {},
   "source": [
    "### Optimization with Opty - Fixed duration simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de6679c4-aeb2-444c-b3b1-20527d41846e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import opty\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02a5c484-9af0-4f40-b6cd-791993da51a9",
   "metadata": {},
   "source": [
    "#### Prepare the EoM "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27def060-f46b-463f-b804-27330d2299e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "eom = Matrix.col_join(Matrix(kd_eqs), fr + frstar)\n",
    "eom"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aecd0ccb-f9ac-4f87-bf76-b2e22fb127a2",
   "metadata": {},
   "source": [
    "#### Identify the known parameters, the control inputs and the state variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aca28ac6-f063-41af-8901-d15590330078",
   "metadata": {},
   "outputs": [],
   "source": [
    "Lg_val, L_val, m_pend_val, m_disk_val, r_val = 1.25, 2, 1, 1, 0.5\n",
    "known_parameter_map = {\n",
    "    g: 9.81,\n",
    "    Lg: Lg_val,\n",
    "    L: L_val,\n",
    "    m_pend: m_pend_val,\n",
    "    m_disk: m_disk_val,\n",
    "    Izzp: m_pend_val * L_val**2 / 12,\n",
    "    Izzd: m_disk_val * r_val**2 / 2\n",
    "}\n",
    "known_parameter_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9557168b-a35b-4cae-9341-e084bc81cbb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "unknown_parameter = []\n",
    "unknown_input_trajectories = [\n",
    "    Tm\n",
    "]\n",
    "state_symbols = [\n",
    "    theta, alpha, u1, u2\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d4d5723-12b0-472c-b2cd-158df9ad6c98",
   "metadata": {},
   "source": [
    "#### Decide between a fixed duration solution or variable duration solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3c6f6be-2578-458e-ba32-54ac35c7e243",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_nodes = 500\n",
    "t0, tf = 0, 10\n",
    "times = np.linspace(t0, tf, num_nodes)\n",
    "h = (tf - t0) / (num_nodes - 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3403917b-d343-411d-bff4-480167e1d802",
   "metadata": {},
   "source": [
    "#### Create the objective function and its gradient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "581cb2b4-aa63-4702-9b46-59d16b6d8052",
   "metadata": {},
   "outputs": [],
   "source": [
    "opty.create_objective_function?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a335a284-d769-4728-8225-fd0b37596a03",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = symbols(\"t\")\n",
    "objective = Integral(Tm**2, t)\n",
    "objective_func, objective_jacobian_func = opty.create_objective_function(\n",
    "    objective,\n",
    "    state_symbols,\n",
    "    unknown_input_trajectories,\n",
    "    unknown_parameter,\n",
    "    num_nodes,\n",
    "    h\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de7902f2-4951-4e87-81d8-10d4eadad4fa",
   "metadata": {},
   "source": [
    "#### Create the necessary constraints on the states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe64d4e5-f151-439b-9421-55009359e97b",
   "metadata": {},
   "outputs": [],
   "source": [
    "theta_0 = 0\n",
    "alpha_0 = 0\n",
    "instance_constraints = [\n",
    "    theta.subs(t, 0) - theta_0,\n",
    "    theta.subs(t, tf) - pi,\n",
    "    alpha.subs(t, 0) - alpha_0,\n",
    "    u1.subs(t, 0),\n",
    "    u1.subs(t, tf),\n",
    "    u2.subs(t, 0),\n",
    "    u2.subs(t, tf)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0335eb5-4e47-4a3f-b85a-fa55d4406eab",
   "metadata": {},
   "source": [
    "#### Create the necessary bounds to unknown parameters, unknown trajectories and states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19fc209c-6bfb-4e95-9b5e-838f23df038e",
   "metadata": {},
   "outputs": [],
   "source": [
    "bounds = {\n",
    "    Tm: (-10, 10),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a37373a7-99dc-4623-b950-bd9c077ee369",
   "metadata": {},
   "source": [
    "#### Setup the optimization problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0448105e-5073-4d26-b199-1c3884cd8d0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "opty.Problem?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3d506db-7aa0-4d65-b26e-343555c36893",
   "metadata": {},
   "outputs": [],
   "source": [
    "problem = opty.Problem(\n",
    "    objective_func,\n",
    "    objective_jacobian_func,\n",
    "    eom,\n",
    "    state_symbols,\n",
    "    num_nodes,\n",
    "    h,\n",
    "    known_parameter_map=known_parameter_map,\n",
    "    instance_constraints=instance_constraints,\n",
    "    bounds=bounds,\n",
    "    time_symbol=t,\n",
    "    integration_method=\"midpoint\"\n",
    ")\n",
    "# optional step\n",
    "problem.add_option('max_iter', 3000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a845e91f-e3c1-40c8-aeb0-f6f72d16cc32",
   "metadata": {},
   "source": [
    "#### Create initial conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cac679c3-af60-4a4e-8560-c64be3adc5e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_guess = np.random.randn(problem.num_free)\n",
    "problem.plot_trajectories(initial_guess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac475623-01eb-4ba0-a365-30d18636a597",
   "metadata": {},
   "outputs": [],
   "source": [
    "problem.num_free"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "074170d7-44e1-4dd2-adc3-dcf2ad1c2596",
   "metadata": {},
   "source": [
    "#### Run the optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87793e8b-4c27-4e0c-afc5-f7a2f30ccb86",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "solution, info = problem.solve(initial_guess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b364edb3-6b2f-420f-9194-da7a73d82c0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(info['status_msg'])\n",
    "print(info['obj_val'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c66e7c97-5944-46cd-8d44-e0288c23ca99",
   "metadata": {},
   "source": [
    "#### Analyze the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78659005-6c38-4477-a18c-c35249438899",
   "metadata": {},
   "outputs": [],
   "source": [
    "problem.plot_trajectories(solution)\n",
    "problem.plot_objective_value()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "596a777c-3a5c-46b3-a581-ff08563430d0",
   "metadata": {},
   "source": [
    "### Animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8d8d849-a653-4b37-a344-5dc132849a42",
   "metadata": {},
   "outputs": [],
   "source": [
    "opty.utils.parse_free?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14383318-f93d-4e1e-b8be-43a3cb865d5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "state_traj, input_traj, constants = opty.utils.parse_free(\n",
    "    solution, problem.collocator.num_states,\n",
    "    problem.collocator.num_unknown_input_trajectories,\n",
    "    problem.collocator.num_collocation_nodes,\n",
    "    variable_duration=problem.collocator._variable_duration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dab224db-edf5-4c13-aee3-3de1222947a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import HTML\n",
    "\n",
    "xP = lambdify([L, theta], P.pos_from(O) & N.x)\n",
    "yP = lambdify([L, theta], P.pos_from(O) & N.y)\n",
    "rot = lambdify([alpha, theta], N.dcm(B))\n",
    "points = np.array([[r_val, 0, 0], [0, 0, 0], [0, r_val, 0]])\n",
    "points_on_wheel = np.array([[0, r_val, 0]])\n",
    "\n",
    "def create_animation(state_traj):\n",
    "    fig, ax = plt.subplots()\n",
    "    \n",
    "    idx0 = 0\n",
    "    th0 = state_traj[0, idx0]\n",
    "    al0 = state_traj[1, idx0]\n",
    "    xP0 = xP(L_val, th0)\n",
    "    yP0 = yP(L_val, th0)\n",
    "    phi = np.linspace(0, 2*np.pi, 50, endpoint=True)\n",
    "    x_wheel = r_val * np.cos(phi)\n",
    "    y_wheel = r_val * np.sin(phi)\n",
    "    points_on_wheel_transformed = (rot(al0, th0) @ points_on_wheel.T).T\n",
    "    \n",
    "    pendulum, = ax.plot([0, xP0], [0, yP0])\n",
    "    wheel, = ax.plot(xP0 + x_wheel, yP0 + y_wheel)\n",
    "    points_on_wheel_handler = ax.scatter(xP0 + points_on_wheel_transformed[:, 0], yP0 + points_on_wheel_transformed[:, 1])\n",
    "    \n",
    "    ax.set_aspect(\"equal\")\n",
    "    w = (L_val + r_val) * 1.25\n",
    "    ax.axis([-w, w, -w, w])\n",
    "    \n",
    "    def update(idx):\n",
    "        theta = state_traj[0, idx]\n",
    "        alpha = state_traj[1, idx]\n",
    "        xP_cur = xP(L_val, theta)\n",
    "        yP_cur = yP(L_val, theta)\n",
    "        points_on_wheel_transformed = (rot(alpha, theta) @ points_on_wheel.T).T\n",
    "        \n",
    "        pendulum.set_data([0, xP_cur], [0, yP_cur])\n",
    "        wheel.set_data(xP_cur + x_wheel, yP_cur + y_wheel)\n",
    "        points_on_wheel_handler.set_offsets(np.c_[xP_cur + points_on_wheel_transformed[:, 0], yP_cur + points_on_wheel_transformed[:, 1]])\n",
    "        ax.set_title(\"t = {:.2f} s\".format(times[idx]))\n",
    "    \n",
    "    # update(0)\n",
    "    plt.close(fig)\n",
    "    return FuncAnimation(fig, func=update, frames=len(times))\n",
    "\n",
    "ani = create_animation(state_traj)\n",
    "HTML(ani.to_jshtml(fps=30))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed5324df-91ab-499e-bfd4-694022959b1c",
   "metadata": {},
   "source": [
    "### Variable duration simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dc50282-00dd-4dee-b002-cb628c229ad9",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_nodes = 500\n",
    "h = symbols(\"t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef798d8e-d709-4f51-b1d1-b2c0f23d2087",
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective_func(free):\n",
    "    \"\"\"Minimize the sum of the squares of the control torque.\"\"\"\n",
    "    T, h = free[4*num_nodes:-1], free[-1]\n",
    "    return h * np.sum(T**2)\n",
    "\n",
    "def objective_jacobian_func(free):\n",
    "    T, h = free[4*num_nodes:-1], free[-1]\n",
    "    grad = np.zeros_like(free)\n",
    "    grad[4*num_nodes:-1] = 2 * h * T\n",
    "    grad[-1] = np.sum(T**2)\n",
    "    return grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09392c5a-8ffa-47a4-80b2-811e8f86bd01",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf = (num_nodes - 1) * h\n",
    "theta_0 = 0\n",
    "alpha_0 = 0\n",
    "instance_constraints = [\n",
    "    theta.subs(t, 0) - theta_0,\n",
    "    theta.subs(t, tf) - pi,\n",
    "    alpha.subs(t, 0) - alpha_0,\n",
    "    u1.subs(t, 0),\n",
    "    u1.subs(t, tf),\n",
    "    u2.subs(t, 0),\n",
    "    u2.subs(t, tf)\n",
    "]\n",
    "\n",
    "bounds = {\n",
    "    Tm: (-10, 10),\n",
    "    h: (0, 0.05)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a8f994a-0c7e-4ff2-903d-2d17fbb4bc1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "problem = opty.Problem(\n",
    "    objective_func,\n",
    "    objective_jacobian_func,\n",
    "    eom,\n",
    "    state_symbols,\n",
    "    num_nodes,\n",
    "    h,\n",
    "    known_parameter_map=known_parameter_map,\n",
    "    instance_constraints=instance_constraints,\n",
    "    bounds=bounds,\n",
    "    time_symbol=t\n",
    ")\n",
    "# optional step\n",
    "problem.add_option('max_iter', 3000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dd3a9e4-cfb1-4629-88df-60a188914717",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "max_attempts = 10\n",
    "i = 0\n",
    "converged = False\n",
    "logs = {}\n",
    "\n",
    "while (i < max_attempts) and (not converged):\n",
    "    initial_guess = np.random.randn(problem.num_free)\n",
    "    initial_guess[-1] = 1 / num_nodes\n",
    "    solution, info = problem.solve(initial_guess)\n",
    "    converged = not info[\"status\"]\n",
    "    logs[i] = info.copy()\n",
    "    i += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb3cfc09-ec2f-45b2-99f9-cd736ab8396c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(info[\"status_msg\"])\n",
    "print(info[\"obj_val\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1103776-a2e2-4c26-8573-c057360af02c",
   "metadata": {},
   "outputs": [],
   "source": [
    "problem.plot_trajectories(solution)\n",
    "problem.plot_objective_value()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "066748c3-254b-49d9-b310-8ca05c58d1b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "state_traj, input_traj, constants, h_val = opty.utils.parse_free(\n",
    "    solution, problem.collocator.num_states,\n",
    "    problem.collocator.num_unknown_input_trajectories,\n",
    "    problem.collocator.num_collocation_nodes,\n",
    "    variable_duration=problem.collocator._variable_duration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fde9f7f-3a2d-4d30-a4aa-62ae1b0580f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ani = create_animation(state_traj)\n",
    "HTML(ani.to_jshtml(fps=30))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b392d3c-6335-4bd1-9fd7-3010aaa2f7c4",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}