steered-car-point-opty.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0a527636",
   "metadata": {},
   "source": [
    "I try to model a **conventional car**: The rear axle is driven, the front axle does the steering.\\\n",
    "No speed possible perpendicular to the wheels.\\\n",
    "It can drive from point A with predescribed direction of the car body to point B with another predescribed direction.\\\n",
    "The parameters I ser here is to mimick parallel parking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7cc216b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy.physics.mechanics as me\n",
    "from collections import OrderedDict\n",
    "import time\n",
    "import numpy as np\n",
    "import sympy as sm\n",
    "from scipy.interpolate import CubicSpline\n",
    "\n",
    "from opty.direct_collocation import Problem\n",
    "from opty.utils import parse_free\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "from IPython.display import HTML\n",
    "matplotlib.rcParams['animation.embed_limit'] = 2**128\n",
    "from matplotlib.animation import FuncAnimation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc31566d",
   "metadata": {},
   "source": [
    "## Timo's function\n",
    "\n",
    "To get obj, obj_grad as needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8430e7ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_objective_function(\n",
    "    objective, state_symbols, input_symbols, unknown_symbols,\n",
    "    N, node_time_interval=1.0, integration_method=\"backward euler\"):\n",
    "    \"\"\"Creates function to evaluate the objective and objective gradient.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    objective : sympy.Expr\n",
    "        The objective function to be minimized, which is a function of the\n",
    "        states and inputs. The objective function can contain non-nested\n",
    "        indefinite integrals of time, e.g. ``Integral(f(t)**2, t)``.\n",
    "    state_symbols : iterable of symbols\n",
    "        The state variables.\n",
    "    input_symbols : iterable of symbols\n",
    "        The input variables.\n",
    "    unknown_symbols : iterable of symbols\n",
    "        The unknown parameters.\n",
    "    N : int\n",
    "        Number of collocation nodes, i.e. the number of time steps.\n",
    "    node_time_interval : float\n",
    "        The value of the time interval. The default is 1.0, as this term only\n",
    "        appears in the objective function as a scaling factor.\n",
    "    integration_method : str, optional\n",
    "        The method used to integrate the system. The default is \"backward\n",
    "        euler\".\n",
    "\n",
    "    \"\"\"\n",
    "    def lambdify_function(expr, multiplication_array, take_sum):\n",
    "        if take_sum:\n",
    "            def integration_function(x):\n",
    "                return node_time_interval * np.sum(x * multiplication_array)\n",
    "        else:\n",
    "            def integration_function(x):\n",
    "                return node_time_interval * x * multiplication_array\n",
    "        return sm.lambdify(\n",
    "            (states, inputs, params), expr,\n",
    "            modules=[{int_placeholder.name: integration_function}, \"numpy\"],\n",
    "            cse=True)\n",
    "\n",
    "    def parse_expr(expr, in_integral=False):\n",
    "        if not expr.args:\n",
    "            return expr\n",
    "        if isinstance(expr, sm.Integral):\n",
    "            if in_integral:\n",
    "                raise NotImplementedError(\"Nested integrals are not supported.\")\n",
    "            if expr.limits != ((me.dynamicsymbols._t,),):\n",
    "                raise NotImplementedError(\n",
    "                    \"Only indefinite integrals of time are supported.\")\n",
    "            return int_placeholder(parse_expr(expr.function, True))\n",
    "        return expr.func(*(parse_expr(arg) for arg in expr.args))\n",
    "\n",
    "    # Parse function arguments\n",
    "    states = sm.ImmutableMatrix(state_symbols)\n",
    "    inputs = sm.ImmutableMatrix(sort_sympy(input_symbols))\n",
    "    params = sm.ImmutableMatrix(sort_sympy(unknown_symbols))\n",
    "    if states.shape[1] > 1 or inputs.shape[1] > 1 or params.shape[1] > 1:\n",
    "        raise ValueError(\n",
    "            'The state, input, and unknown symbols must be column matrices.')\n",
    "    n, q = states.shape[0], inputs.shape[0]\n",
    "    i_idx, r_idx = n * N, (n + q) * N\n",
    "\n",
    "    # Compute analytical gradient of the objective function\n",
    "    objective_grad = sm.ImmutableMatrix([objective]).jacobian(\n",
    "        states[:] + inputs[:] + params[:])\n",
    "\n",
    "    # Replace the integral with a custom function\n",
    "    int_placeholder = sm.Function(\"_IntegralFunction\")\n",
    "    objective = parse_expr(objective)\n",
    "    objective_grad = tuple(parse_expr(objective_grad))\n",
    "\n",
    "    # Replace zeros with an array of zeros, otherwise lambdify will return a\n",
    "    # scalar zero instead of an array of zeros.\n",
    "    objective_grad = tuple(\n",
    "        np.zeros(N) if grad == 0 else grad for grad in objective_grad[:n + q]\n",
    "        ) + tuple(objective_grad[n + q:])\n",
    "\n",
    "    # Define evaluation functions based on the integration method\n",
    "    if integration_method == \"backward euler\":\n",
    "        obj_expr_eval = lambdify_function(\n",
    "            objective, np.hstack((0, np.ones(N - 1))), True)\n",
    "        obj_grad_time_expr_eval = lambdify_function(\n",
    "            objective_grad[:n + q], np.hstack((0, np.ones(N - 1))), False)\n",
    "        obj_grad_param_expr_eval = lambdify_function(\n",
    "            objective_grad[n + q:], np.hstack((0, np.ones(N - 1))), True)\n",
    "        def obj(free):\n",
    "            states = free[:i_idx].reshape((n, N))\n",
    "            inputs = free[i_idx:r_idx].reshape((q, N))\n",
    "            return obj_expr_eval(states, inputs, free[r_idx:])\n",
    "\n",
    "        def obj_grad(free):\n",
    "            states = free[:i_idx].reshape((n, N))\n",
    "            inputs = free[i_idx:r_idx].reshape((q, N))\n",
    "            return np.hstack((\n",
    "                *obj_grad_time_expr_eval(states, inputs, free[r_idx:]),\n",
    "                obj_grad_param_expr_eval(states, inputs, free[r_idx:])\n",
    "             ))\n",
    "\n",
    "    elif integration_method == \"midpoint\":\n",
    "        obj_expr_eval = lambdify_function(\n",
    "            objective, np.ones(N - 1), True)\n",
    "        obj_grad_time_expr_eval = lambdify_function(\n",
    "            objective_grad[:n + q], np.hstack((0.5, np.ones(N - 2), 0.5)),\n",
    "            False)\n",
    "        obj_grad_param_expr_eval = lambdify_function(\n",
    "            objective_grad[n + q:], np.ones(N - 1), True)\n",
    "        def obj(free):\n",
    "            states = free[:i_idx].reshape((n, N))\n",
    "            states_mid = 0.5 * (states[:, :-1] + states[:, 1:])\n",
    "            inputs = free[i_idx:r_idx].reshape((q, N))\n",
    "            inputs_mid = 0.5 * (inputs[:, :-1] + inputs[:, 1:])\n",
    "            return obj_expr_eval(states_mid, inputs_mid, free[r_idx:])\n",
    "\n",
    "        def obj_grad(free):\n",
    "            states = free[:i_idx].reshape((n, N))\n",
    "            states_mid = 0.5 * (states[:, :-1] + states[:, 1:])\n",
    "            inputs = free[i_idx:r_idx].reshape((q, N))\n",
    "            inputs_mid = 0.5 * (inputs[:, :-1] + inputs[:, 1:])\n",
    "            return np.hstack((\n",
    "                *obj_grad_time_expr_eval(states, inputs, free[r_idx:]),\n",
    "                obj_grad_param_expr_eval(states_mid, inputs_mid, free[r_idx:])\n",
    "            ))\n",
    "\n",
    "    else:\n",
    "        raise NotImplementedError(\n",
    "            f\"Integration method '{integration_method}' is not implemented.\")\n",
    "\n",
    "    return obj, obj_grad\n",
    "\n",
    "\n",
    "def sort_sympy(seq):\n",
    "    \"\"\"Returns a sorted list of the symbols.\"\"\"\n",
    "    seq = list(seq)\n",
    "    try:  # symbols\n",
    "        seq.sort(key=lambda x: x.name)\n",
    "    except AttributeError:  # functions\n",
    "        seq.sort(key=lambda x: x.__class__.__name__)\n",
    "    return seq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a5b475d",
   "metadata": {},
   "source": [
    "Defines the shape of the street (strasse = German for street)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "090b386b",
   "metadata": {},
   "source": [
    "Set up **Kane's equations of motion**\\\n",
    "Note, that the holonomic constraints is added, to give $EOM = \\begin{pmatrix} \\text{kinematic constraints} \\\\ fr + fr^{\\star} \\\\ \\text{nonholonomic constraints} \\end{pmatrix} $\\\n",
    "otherwise nothing special."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c665b199",
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "\n",
    "N, A0, Ab, Af = sm.symbols('N A0 Ab Af', cls= me.ReferenceFrame)\n",
    "t = me.dynamicsymbols._t\n",
    "O, Pb, Dmc, Pf = sm.symbols('O Pb Dmc Pf', cls= me.Point)\n",
    "O.set_vel(N, 0)\n",
    "\n",
    "q0, qf  = me.dynamicsymbols('q_0 q_f')                                                      # gen. coordinates of the main body, the rear axle, the front axle\n",
    "u0, uf  = me.dynamicsymbols('u_0 u_f')                                                      # their speeds\n",
    "x, y    = me.dynamicsymbols('x y')                                                          # coordinates of the point, where the rear axle attaches to the main body\n",
    "ux, uy  = me.dynamicsymbols('u_x u_y')                                                      # their speeds\n",
    "Tf, Fb  = me.dynamicsymbols('T_f F_b')                                                      # torque at back wheel, torque at front wheel, forces on Pb, Pf. Control variables\n",
    "reibung = sm.symbols('reibung')                                                             # friction coefficient\n",
    "\n",
    "l, m0, mb, mf, iZZ0, iZZb, iZZf = sm.symbols('l m0 mb mf iZZ0, iZZb, iZZf')                 # link length, mass, inertia of skateboard\n",
    "\n",
    "A0.orient_axis(N, q0, N.z)\n",
    "A0.set_ang_vel(N, u0 * N.z)\n",
    "\n",
    "Ab.orient_axis(A0, 0, N.z)                                                                  # earlier I made the rear axle steerable too. I kept the frame because some points for plotting are defined in it.\n",
    "\n",
    "Af.orient_axis(A0, qf, N.z)\n",
    "rot = Af.ang_vel_in(N)\n",
    "Af.set_ang_vel(N, uf * N.z)\n",
    "rot1 = Af.ang_vel_in(N)\n",
    "\n",
    "Pf.set_pos(O, x * N.x + y * N.y)\n",
    "Pf.set_vel(N, ux * N.x + uy * N.y)\n",
    "\n",
    "Pb.set_pos(Pf, -l * A0.y)   \n",
    "Pb.v2pt_theory(Pf, N, A0) \n",
    "\n",
    "Dmc.set_pos(Pf, -l/2 * A0.y)\n",
    "Dmc.v2pt_theory(Pf, N, A0)\n",
    "\n",
    "vel1 = me.dot(Pb.vel(N), Ab.x) - 0                                                                                                  # no speed perpendicular to the rear wheels                                        \n",
    "vel2 = me.dot(Pf.vel(N), Af.x) - 0                                                                                                  # no speed perpendicular to the front wheels        \n",
    "\n",
    "I0 = me.inertia(A0, 0, 0, iZZ0)\n",
    "body0 = me.RigidBody('body0', Dmc, A0, m0, (I0, Dmc))\n",
    "Ib = me.inertia(Ab, 0, 0, iZZb)\n",
    "bodyb = me.RigidBody('bodyb', Pb, Ab, mb, (Ib, Pb))\n",
    "If = me.inertia(Af, 0, 0, iZZf)\n",
    "bodyf = me.RigidBody('bodyf', Pf, Af, mf, (If, Pf))\n",
    "BODY = [body0, bodyb, bodyf]\n",
    "\n",
    "\n",
    "# Forces\n",
    "FL = [(Pb, Fb * Ab.y), (Af, Tf * N.z), (Dmc, -reibung * Dmc.vel(N))]\n",
    "\n",
    "kd = sm.Matrix([ux - x.diff(t), uy - y.diff(t), u0 - q0.diff(t), me.dot(rot1- rot, N.z)])    # kinematic differential equations\n",
    "speed_constr = sm.Matrix([vel1, vel2])                                                       # speed constraints\n",
    "\n",
    "q_ind = [x, y, q0, qf]\n",
    "u_ind = [u0, uf]\n",
    "u_dep = [ux, uy]\n",
    "\n",
    "# Create the KanesMethod object\n",
    "KM = me.KanesMethod(\n",
    "                    N, q_ind=q_ind, u_ind=u_ind, \n",
    "                    kd_eqs=kd, \n",
    "                    u_dependent=u_dep, \n",
    "                    velocity_constraints=speed_constr,\n",
    "                    )\n",
    "(fr, frstar) = KM.kanes_equations(BODY, FL)\n",
    "\n",
    "EOM = kd.col_join(fr + frstar) \n",
    "EOM = EOM.col_join(speed_constr)\n",
    "\n",
    "print('EOM shape', EOM.shape, '\\n')\n",
    "print('EOM DS', me.find_dynamicsymbols(EOM), '\\n')\n",
    "print('EOM FS', EOM.free_symbols)\n",
    "print(f'EOMs contain {sm.count_ops(EOM)} operations, {sm.count_ops(sm.cse(EOM))} after cse')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0685c32",
   "metadata": {},
   "source": [
    "Set up the machinery for **opty** and **solve** the optimization.\\\n",
    "If needed, one can iterate, using the solution of iteration i as initial_guess for iteration i + 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc7e7050",
   "metadata": {},
   "outputs": [],
   "source": [
    "state_symbols     = tuple((x, y, q0, qf, ux, uy, u0, uf))\n",
    "laenge            = len(state_symbols)\n",
    "constant_symbols  = (l, m0, mb, mf, iZZ0, iZZb, iZZf, reibung)\n",
    "specified_symbols = (Fb, Tf)\n",
    "unknown_symbols   = []\n",
    "\n",
    "#======================================================================================================\n",
    "methode = 'backward euler'  # The integration method to use. backward euler or midpoint are the choices\n",
    "opty_info = True            # If True, the optimizer will print some information about the optimization process\n",
    "#======================================================================================================\n",
    "\n",
    "duration  = 15                                             # The duration of the simulation.\n",
    "num_nodes = 100                                            # The number of nodes to use in the direct collocation problem.\n",
    "t0, tf    = 0.0, duration                                  # The initial and final times\n",
    "\n",
    "initial_guess = np.ones((len(state_symbols) + len(specified_symbols)) * num_nodes) * 0.01\n",
    "\n",
    "\n",
    "interval_value = duration / (num_nodes - 1)\n",
    "\n",
    "# Specify the known system parameters.\n",
    "par_map = OrderedDict()\n",
    "par_map[m0]       = 1.0                                      # Mass of the car\n",
    "par_map[mb]       = 0.5                                      # Mass of the back wheel\n",
    "par_map[mf]       = 0.5                                      # Mass of the front wheel\n",
    "par_map[iZZ0]     = 1.                                       # Inertia of the car\n",
    "par_map[iZZb]     = 0.5                                      # Inertia of the back wheel\n",
    "par_map[iZZf]     = 0.5                                      # Inertia of the front wheel\n",
    "par_map[l]        = 3.0                                      # Length of the skateboard\n",
    "par_map[reibung]  = 0.5                                      # Friction coefficient\n",
    "\n",
    "objective = sm.Integral(Fb**2 + Tf**2, t)                    # The objective function to minimize.\n",
    "\n",
    "# Specify the objective function and it's gradient.\n",
    "obj, obj_grad = create_objective_function(\n",
    "    objective, \n",
    "    state_symbols, \n",
    "    specified_symbols, \n",
    "    unknown_symbols,\n",
    "    num_nodes, \n",
    "    interval_value, \n",
    "    methode)\n",
    "\n",
    "#====================================================================================\n",
    "initial_state_constraints = {\n",
    "                            x: 0,\n",
    "                            y: 0,\n",
    "                            q0: -np.pi/2.,                   # starting angle of the car\n",
    "                            qf: 0.5,                         # starting angle of steering wheel\n",
    "                            ux: 0.,\n",
    "                            uy: 0.,\n",
    "                            u0: 0.,\n",
    "                            uf: 0.,\n",
    "                            }\n",
    "\n",
    "final_state_constraints    = {\n",
    "                             x: 0.,\n",
    "                             ux: 0.,\n",
    "                             y : -2,                        # end position of the car\n",
    "                             uy: 0.,                  \n",
    "                             q0: -np.pi/2.,                 # end angle of the car\n",
    "                             }    \n",
    "    \n",
    "instance_constraints = tuple(xi.subs({t: t0}) - xi_val for xi, xi_val in initial_state_constraints.items()) + tuple(xi.subs({t: tf}) - xi_val for xi, xi_val in final_state_constraints.items())\n",
    "print('instance constraints:', instance_constraints, '\\n')\n",
    "grenze = 50.\n",
    "delta = np.pi/4.\n",
    "bounds = {\n",
    "        Fb: (-grenze, grenze),\n",
    "        Tf: (-grenze, grenze), \n",
    "        qf: (-np.pi/2. + delta, np.pi/2. - delta),          # restrict the steering angle to avoid locking\n",
    "        y: (-2., 0),\n",
    "        x: (-1., 1.),\n",
    "        } \n",
    "\n",
    "# Here I can iterate. This helps sometimes.\n",
    "solution_list = []                                          # list to store the solutions of each iteration\n",
    "g_list        = []                                          # list to store the values of the errors\n",
    "\n",
    "for i in range(2): \n",
    "    prob = Problem(obj, obj_grad, EOM, state_symbols, num_nodes, interval_value,\n",
    "        known_parameter_map=par_map,\n",
    "        instance_constraints=instance_constraints,\n",
    "        bounds=bounds,\n",
    "        integration_method=methode)\n",
    "\n",
    "    prob.add_option('max_iter', 120000)           # default is 3000\n",
    "\n",
    "# Find the optimal solution.\n",
    "    solution, info = prob.solve(initial_guess)\n",
    "    print(f'{i+1} - th iteration')\n",
    "    if opty_info:\n",
    "        print('message from optimizer:', info['status_msg']) \n",
    "        print('Iterations needed',len(prob.obj_value))\n",
    "        print(f\"objective value {info['obj_val']:.3e} \\n\")\n",
    "#    prob.plot_objective_value()\n",
    "    initial_guess = solution\n",
    "    solution_list.append(solution)\n",
    "    g_list.append(info['g'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87dfce60",
   "metadata": {},
   "source": [
    "Plot the **errors** in the state symbols."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9954e90e",
   "metadata": {},
   "outputs": [],
   "source": [
    "welche = 1\n",
    "fehler = g_list[welche]\n",
    "anzahl = len(state_symbols)\n",
    "fig, ax = plt.subplots(anzahl, 1, figsize=(10, 2.25*anzahl), sharex=True)\n",
    "times = np.linspace(0.0, duration, num=num_nodes)\n",
    "\n",
    "for i, j in enumerate(state_symbols):\n",
    "    if j in final_state_constraints.keys():\n",
    "        farbe1 = 'red'\n",
    "    else:\n",
    "        farbe1 = 'black'\n",
    "\n",
    "    if i < laenge:\n",
    "        farbe = 'b'\n",
    "    else:\n",
    "        farbe = 'r'\n",
    "    ax[i].plot(times, fehler[i * num_nodes:(i + 1) * num_nodes], color = farbe, lw=0.5)\n",
    "    ax[i].grid()\n",
    "    ax[i].set_ylabel(f'${str(j)}$', color=farbe1)\n",
    "ax[-1].set_xlabel('time')\n",
    "ax[0]. set_title(f'Errors in the generalized coordinates and speeds \\n coordinates/speeds in the final_state_constraints are red', color = 'b');\n",
    "print(\n",
    "     f'error in the final state constraints x:  {fehler[-5]:.3e},\\n error in the final state constrains ux: {fehler[-4]:.3e} \\n' + \n",
    "     f'error in the final state constraints y:  {fehler[-3]:.3e},\\n error in the final state constrains uy: {fehler[-2]:.3e} \\n' + \n",
    "     f'error in the final state constraints q0: {fehler[-1]:.3e}' \n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fb77b13",
   "metadata": {},
   "source": [
    "Plot generalized coordinates / speeds and forces / torques"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6f0398b",
   "metadata": {},
   "outputs": [],
   "source": [
    "solution = solution_list[welche]\n",
    "anzahl = len(state_symbols) + len(specified_symbols)\n",
    "fig, ax = plt.subplots(anzahl, 1, figsize=(10, 2.5*anzahl), sharex=True)\n",
    "times = np.linspace(0.0, duration, num=num_nodes)\n",
    "\n",
    "for i, j in enumerate(state_symbols + specified_symbols):\n",
    "    if j in final_state_constraints.keys():\n",
    "        farbe1 = 'red'\n",
    "    else:\n",
    "        farbe1 = 'black'\n",
    "\n",
    "    if i < laenge:\n",
    "        farbe = 'b'\n",
    "    else:\n",
    "        farbe = 'r'\n",
    "    ax[i].plot(times, solution[i * num_nodes:(i + 1) * num_nodes], color = farbe, lw = 0.5)\n",
    "    ax[i].grid()\n",
    "    ax[i].set_ylabel(f'${str(j)}$', color=farbe1)\n",
    "ax[-1].set_xlabel('time')\n",
    "ax[0]. set_title('Generalized coordinates and speeds', color = 'b')\n",
    "ax[laenge]. set_title('Control force(s)', color = 'r');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c384f66b",
   "metadata": {},
   "source": [
    "aminate the **car**\\\n",
    "Nothing special, just a lot of points have to be set, this is why it looks so large.\\\n",
    "The green arrows symbolize the forces which opty calculated to drive the car. They are perpendicular to the rear axle - \\\n",
    "after all, it is rear wheel drive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5546cec9",
   "metadata": {},
   "outputs": [],
   "source": [
    "fps = 30\n",
    "\n",
    "def add_point_to_data(line, x, y):\n",
    "# to trace the path of the point. Copied from Timo.\n",
    "    old_x, old_y = line.get_data()\n",
    "    line.set_data(np.append(old_x, x), np.append(old_y, y))\n",
    "\n",
    "\n",
    "state_vals, input_vals, _ = parse_free(solution, len(state_symbols), len(specified_symbols), num_nodes)\n",
    "t_arr = np.linspace(t0, tf, num_nodes)\n",
    "state_sol = CubicSpline(t_arr, state_vals.T)\n",
    "input_sol = CubicSpline(t_arr, input_vals.T)\n",
    "\n",
    "# create additional points for the axles\n",
    "Pbl, Pbr, Pfl, Pfr = sm.symbols('Pbl Pbr Pfl Pfr', cls= me.Point)\n",
    "Fbq = me.Point('Fbq')                                      # end points of the force\n",
    "la = sm.symbols('la')                                                               # length of the axles\n",
    "fb, tq = sm.symbols('f_b, t_q')                                                      # forces at the axles\n",
    "\n",
    "Pbl.set_pos(Pb, -la/2 * Ab.x)\n",
    "Pbr.set_pos(Pb, la/2 * Ab.x)\n",
    "Pfl.set_pos(Pf, -la/2 * Af.x)\n",
    "Pfr.set_pos(Pf, la/2 * Af.x)\n",
    "\n",
    "Fbq.set_pos(Pb, fb * Ab.y)\n",
    "\n",
    "coordinates = Pb.pos_from(O).to_matrix(N)\n",
    "for point in (Dmc, Pf, Pbl, Pbr, Pfl, Pfr, Fbq):\n",
    "    coordinates = coordinates.row_join(point.pos_from(O).to_matrix(N))\n",
    "\n",
    "pL, pL_vals = zip(*par_map.items())\n",
    "la1 = par_map[l] / 4.                      # length of an axle   \n",
    "coords_lam = sm.lambdify((*state_symbols, fb, tq, *pL, la), coordinates, cse=True)\n",
    "\n",
    "\n",
    "# needed to give the picture the right size.\n",
    "xmin, xmax = -5, 5.\n",
    "ymin, ymax = -5., 5.\n",
    "\n",
    "fig = plt.figure(figsize=(9, 9))\n",
    "ax  = fig.add_subplot(111)\n",
    "ax.set_xlim(xmin, xmax)\n",
    "ax.set_ylim(ymin, ymax)\n",
    "ax.set_aspect('equal')\n",
    "ax.grid()\n",
    "ax.plot(initial_state_constraints[x], initial_state_constraints[y], 'ro', markersize=10)                        # the initial position\n",
    "ax.plot(final_state_constraints[x], final_state_constraints[y], 'go', markersize=10)                            # the final position\n",
    "\n",
    "ax.axvline(bounds[x][1]+0.1, color='red', lw=0.5)\n",
    "ax.axvline(bounds[x][0] - par_map[l]-0.1, color='red', lw=0.5)\n",
    "\n",
    "# Initialize the block\n",
    "line1, = ax.plot([], [], color='orange', lw=2)                                                                  # the car\n",
    "line2, = ax.plot([], [], color='red', lw=2)                                                                     # the back axle\n",
    "line3, = ax.plot([], [], color='magenta', lw=2)                                                                 # the front axle\n",
    "line4  = ax.quiver([], [], [], [], color='green', scale=35, width=0.004, headwidth=8)                           # the force at the back axle\n",
    "line5, = ax.plot([], [], color='black', lw=0.5)                                                                 # trace the front left wheel\n",
    "line6, = ax.plot([], [], color='black', lw=0.5)                                                                 # trace the front right wheel  \n",
    "\n",
    "# Function to update the plot for each animation frame\n",
    "def update(t):\n",
    "    message = (f'running time {t:.2f} sec \\n The back axle is red, the front axle is magenta \\n The driving force is green')\n",
    "    ax.set_title(message, fontsize=12)\n",
    "\n",
    "    coords = coords_lam(*state_sol(t), *input_sol(t), *pL_vals, la1)\n",
    "    \n",
    "#   Pb, Dmc, Pf, Pbl, Pbr, Pfl, Pfr, Fbq                                      sequence of points in coords.\n",
    "    line1.set_data([coords[0, 0], coords[0, 2]], [coords[1, 0], coords[1, 2]])\n",
    "    line2.set_data([coords[0, 3], coords[0, 4]], [coords[1, 3], coords[1, 4]])\n",
    "    line3.set_data([coords[0, 5], coords[0, 6]], [coords[1, 5], coords[1, 6]])\n",
    "\n",
    "    line4.set_offsets([coords[0, 0], coords[1, 0]])\n",
    "    line4.set_UVC(coords[0, 7] - coords[0, 0] , coords[1, 7] - coords[1, 0])\n",
    "\n",
    "    add_point_to_data(line5, coords[0, 5], coords[1, 5])\n",
    "    add_point_to_data(line6, coords[0, 6], coords[1, 6])\n",
    "\n",
    "    return line1, line2, line3, line4, line5, line6,\n",
    "\n",
    "# Set labels and legend\n",
    "ax.set_xlabel('X-axis')\n",
    "ax.set_ylabel('Y-axis')\n",
    "\n",
    "# Create the animation\n",
    "animation = FuncAnimation(fig, update, frames=np.arange(t0, tf, 1 / fps), interval=fps, blit=False)\n",
    "plt.close(fig)  # Prevents the final image from being displayed directly below the animation\n",
    "# Show the plot\n",
    "display(HTML(animation.to_jshtml()))\n",
    "print(f'it took {time.time()-start:.1f} seconds to run the code.')"
   ]
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}