3D-drone-propeller-opty.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0a527636",
   "metadata": {},
   "source": [
    "I try to simulate a **drone** in 3D\\\n",
    "The body is modelled as a solid ball with center DmcD\\\n",
    "The four motors are at distance a from the body, also modelled as solid balls, with centers $P_i$\\\n",
    "The four propellers are modelled as solid discs, with centers $Pp_i$. To create lift, two propellers must rotate 'positively', two 'negatively', else the drome might rotate.\\\n",
    "Also the radius of the propeller is optimized. \n",
    "\n",
    "The objective function is either $\\int_0^{\\text{final time}} (T_1^2 + T_2^2 + T_3^2 + T_4^2) \\, dt $, where $T_i$ is the magnitude of the torque acting on the propellers i\\\n",
    "\n",
    "or $\\int_0^{\\text{final time}} \\text{kinetic energy} \\, dt $, may be selected.\n"
   ]
  },
  {
   "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",
    "to create the objective function and its gradient."
   ]
  },
  {
   "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": "090b386b",
   "metadata": {},
   "source": [
    "## Kane's equations of motion\n",
    "\n",
    "The power of the radius of the propeller used in the friction and in the lift are completely arbitrary. I just played around until I found powers, which optimize the radius inside the bound, not at either limit.\\\n",
    "Also, I have no idea, whether friction and lift depend linearly on the rotational speed of the propeller, I simply assumed without doing any checking\\\n",
    "Everything else is standard.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c665b199",
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "start0 = time.time()\n",
    "\n",
    "N, Ab = sm.symbols('N Ab', cls= me.ReferenceFrame)                              # inertial frame, frame of the drone body\n",
    "Ap1, Ap2, Ap3, Ap4 = sm.symbols('Ap1 Ap2 Ap3 Ap4', cls= me.ReferenceFrame)      # frame of the propellers\n",
    "t = me.dynamicsymbols._t\n",
    "O, DmcD, P1, P2, P3, P4 = sm.symbols('O DmcD P1 P2 P3 P4', cls= me.Point)       # origin, mass center of the drone, of the motors P1...P4\n",
    "Pp1, Pp2, Pp3, Pp4 = sm.symbols('Pp1 Pp2 Pp3 Pp4', cls= me.Point)               # mass centers of the propellers\n",
    "O.set_vel(N, 0)\n",
    "\n",
    "x, y, z   = me.dynamicsymbols('x y z')                                          # position of the mass center of the drone\n",
    "ux, uy, uz = me.dynamicsymbols('u_x, u_y, u_z')                                             \n",
    "q1, q2, q3 = me.dynamicsymbols('q_1, q_2, q_3')                                 # angle of rotation of the drone wrt N\n",
    "u1, u2, u3 = me.dynamicsymbols('u_1, u_2, u_3')         \n",
    "\n",
    "qp1, qp2, qp3, qp4 = me.dynamicsymbols('q_p1, q_p2, q_p3, q_p4')                # angle of rotation of the propellers wrt Ab\n",
    "up1, up2, up3, up4 = me.dynamicsymbols('u_p1, u_p2, u_p3, u_p4')\n",
    "\n",
    "T1, T2, T3, T4 = me.dynamicsymbols('T_1, T_2, T_3, T_4')                        # torques applied to the propellers, control input \n",
    "\n",
    "mD, mM, mP, g, a, reibung = sm.symbols('m_D,m_M, m_P, g, a, reibung')           # mass of body of ther drone, mass of the motors, mass of propellers per unit radius of it , gravity, distance of the motors from the mass center, friction coefficient\n",
    "rb, rM, rP, l1 = sm.symbols('r_b r_M r_P l_1')                                  # radius of the body, radius of the motors, radius of the propellers, location of the propellers 'above' the drone\n",
    "reibungP = sm.symbols('reibungP')                                               # friction coefficient of the propellers\n",
    "\n",
    "Ab.orient_body_fixed(N, (q3, q1, q2), '312')                                                   \n",
    "rot = Ab.ang_vel_in(N)                                                          # needed for the kinematic equations                                                       \n",
    "Ab.set_ang_vel(N, u1 * Ab.x + u2*Ab.y + u3*Ab.z)                                \n",
    "rot1 = Ab.ang_vel_in(N)                                                         # needed for the kinematic equations            \n",
    "\n",
    "Ap1.orient_axis(Ab, qp1, Ab.z)                                                  # orientation of the propellers\n",
    "Ap1.set_ang_vel(Ab, up1 * Ab.z)\n",
    "Ap2.orient_axis(Ab, qp2, Ab.z)\n",
    "Ap2.set_ang_vel(Ab, up2 * Ab.z)\n",
    "Ap3.orient_axis(Ab, qp3, Ab.z)\n",
    "Ap3.set_ang_vel(Ab, up3 * Ab.z)\n",
    "Ap4.orient_axis(Ab, qp4, Ab.z)\n",
    "Ap4.set_ang_vel(Ab, up4 * Ab.z)\n",
    "\n",
    "DmcD.set_pos(O, x*N.x + y*N.y + z*N.z)                                          # position of the mass center of the drone\n",
    "DmcD.set_vel(N, ux*N.x + uy*N.y + uz*N.z)\n",
    "\n",
    "P1.set_pos(DmcD, a * Ab.x)                                                      # position of the motors\n",
    "P2.set_pos(DmcD, -a * Ab.x)\n",
    "P3.set_pos(DmcD, a * Ab.y)\n",
    "P4.set_pos(DmcD, -a * Ab.y)\n",
    "\n",
    "P1.v2pt_theory(DmcD, N, Ab)\n",
    "P2.v2pt_theory(DmcD, N, Ab)\n",
    "P3.v2pt_theory(DmcD, N, Ab)\n",
    "P4.v2pt_theory(DmcD, N, Ab)\n",
    "\n",
    "Pp1.set_pos(P1, l1 * Ab.z)                                                      # position of the propellers\n",
    "Pp2.set_pos(P2, l1 * Ab.z)\n",
    "Pp3.set_pos(P3, l1 * Ab.z)\n",
    "Pp4.set_pos(P4, l1 * Ab.z)\n",
    "\n",
    "Pp1.v2pt_theory(P1, N, Ab)\n",
    "Pp2.v2pt_theory(P2, N, Ab)\n",
    "Pp3.v2pt_theory(P3, N, Ab)\n",
    "Pp4.v2pt_theory(P4, N, Ab)\n",
    "\n",
    "\n",
    "# Create the bodies\n",
    "# drone body and motors are modelled as solid balls of of radius rb, rM and masses mD, mM\n",
    "# the propellers are modelled as solid disks of radius rP and mass mP\n",
    "iXXd = 2./5 * mD * rb**2\n",
    "iYYd = iXXd\n",
    "iZZd = iXXd\n",
    "inertiaD = (me.inertia(Ab, iXXd, iYYd, iZZd))\n",
    "DmcDa = me.RigidBody('DmcDa', DmcD, Ab, mD, (inertiaD, DmcD))\n",
    "\n",
    "iXXm = 2./5 * mM * rM**2\n",
    "iYYm = iXXm\n",
    "iZZm = iXXm\n",
    "inertiaM = (me.inertia(Ab, iXXm, iYYm, iZZm))\n",
    "P1a = me.RigidBody('P1a', P1, Ab, mM, (inertiaM, P1))\n",
    "P2a = me.RigidBody('P2a', P2, Ab, mM, (inertiaM, P2))\n",
    "P3a = me.RigidBody('P3a', P3, Ab, mM, (inertiaM, P3))\n",
    "P4a = me.RigidBody('P4a', P4, Ab, mM, (inertiaM, P4))\n",
    "\n",
    "iXXp = 1./4. * (mP * rP) * rP**2\n",
    "iYYp = iXXp\n",
    "iZZp = 1./2. * (mP * rP) * rP**2\n",
    "inertiaP1 = (me.inertia(Ap1, iXXp, iYYp, iZZp))\n",
    "inertiaP2 = (me.inertia(Ap2, iXXp, iYYp, iZZp))\n",
    "inertiaP3 = (me.inertia(Ap3, iXXp, iYYp, iZZp))\n",
    "inertiaP4 = (me.inertia(Ap4, iXXp, iYYp, iZZp))\n",
    "Pp1a = me.RigidBody('Pp1a', Pp1, Ap1, mP, (inertiaP1, Pp1))\n",
    "Pp2a = me.RigidBody('Pp2a', Pp2, Ap2, mP, (inertiaP2, Pp2))\n",
    "Pp3a = me.RigidBody('Pp3a', Pp3, Ap3, mP, (inertiaP3, Pp3))\n",
    "Pp4a = me.RigidBody('Pp4a', Pp4, Ap4, mP, (inertiaP4, Pp4))\n",
    "\n",
    "BODY = [DmcDa, P1a, P2a, P3a, P4a, Pp1a, Pp2a, Pp3a, Pp4a]\n",
    "\n",
    "# Forces\n",
    "# I assume, that the frictional forces are proportionmal to the velocity of the drone squared, and diminish with heigtht.\n",
    "# The exponents I use for the radisu of the propellers, rD, for the fritction and for the lift are completely ARBITRARY.\n",
    "# I simply played around, until the optimum would be somewhere between the bounds.\n",
    "# I did not bother to look for 'realistic exponents, nor did I look, whether the rotational speed enters linearly, as it does here\n",
    "FL = [\n",
    "    (Ap1, T1 * Ab.z - reibungP * rP**(2/3) * up1 * Ab.z),                                                               # drive the propellers. the control variables.\n",
    "    (Ap2, T2 * Ab.z - reibungP * rP**(2/3) * up2 * Ab.z),\n",
    "    (Ap3, T3 * Ab.z - reibungP * rP**(2/3) * up3 * Ab.z),\n",
    "    (Ap4, T4 * Ab.z - reibungP * rP**(2/3) * up4 * Ab.z),\n",
    "    (Ab, -(T1 + T2 + T3 + T4) * Ab.z),                                                                      # moment of the propellers onto the drone\n",
    "    (DmcD, -mD*g*N.z - reibung * DmcD.vel(N).magnitude()*DmcD.vel(N) / (1 + z**2)), \n",
    "    (P1,   -mM*g*N.z - reibung * P1.vel(N).magnitude()*P1.vel(N) / (1 + z**2)), \n",
    "    (P2,   -mM*g*N.z - reibung * P2.vel(N).magnitude()*P2.vel(N) / (1 + z**2)),\n",
    "    (P3,   -mM*g*N.z - reibung * P3.vel(N).magnitude()*P3.vel(N) / (1 + z**2)),\n",
    "    (P4,   -mM*g*N.z - reibung * P4.vel(N).magnitude()*P4.vel(N) / (1 + z**2)),\n",
    "    (Pp1,  -mP*g*N.z - reibung * Pp1.vel(N).magnitude()*Pp1.vel(N) / (1 + z**2) + rP**2 * up1 * Ap1.z),      # propellers must run in different directions\n",
    "    (Pp2,  -mP*g*N.z - reibung * Pp2.vel(N).magnitude()*Pp2.vel(N) / (1 + z**2) - rP**2 * up2 * Ap2.z),      # otherwise the drone should rotate\n",
    "    (Pp3,  -mP*g*N.z - reibung * Pp3.vel(N).magnitude()*Pp3.vel(N) / (1 + z**2) + rP**2 * up3 * Ap3.z),\n",
    "    (Pp4,  -mP*g*N.z - reibung * Pp4.vel(N).magnitude()*Pp4.vel(N) / (1 + z**2) - rP**2 * up4 * Ap4.z),\n",
    "    ]                                                    \n",
    "\n",
    "kd = sm.Matrix([\n",
    "               ux - x.diff(t), uy - y.diff(t), uz - z.diff(t), \n",
    "               *[(rot-rot1).dot(uv) for uv in N],\n",
    "               up1 - qp1.diff(t), up2 - qp2.diff(t), up3 - qp3.diff(t), up4 - qp4.diff(t),\n",
    "               ]) \n",
    "config_constr = []\n",
    "speed_constr  = []\n",
    "q_ind = [x, y, z, q1, q2, q3, qp1, qp2, qp3, qp4]\n",
    "q_dep = []\n",
    "u_ind = [ux, uy, uz, u1, u2, u3, up1, up2, up3, up4]\n",
    "u_dep = []\n",
    "\n",
    "# Create the KanesMethod object\n",
    "KM = me.KanesMethod(\n",
    "                    N, \n",
    "                    q_ind=q_ind, \n",
    "                    q_dependent=q_dep,\n",
    "                    u_ind=u_ind, \n",
    "                    u_dependent=u_dep, \n",
    "                    kd_eqs=kd,\n",
    "                    configuration_constraints=config_constr,\n",
    "                    velocity_constraints=speed_constr,\n",
    "                    )\n",
    "(fr, frstar) = KM.kanes_equations(BODY, FL)\n",
    "EOM = kd.col_join(fr + frstar) \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')\n",
    "print(f'EOMs contain {sm.count_ops(EOM)} operations, {sm.count_ops(sm.cse(EOM))} after cse', '\\n')\n",
    "\n",
    "qL = [wert for wert in KM.q] + [wert for wert in KM.u]\n",
    "print('sequence of states:', qL)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0685c32",
   "metadata": {},
   "source": [
    "## opty\n",
    "Set up the machinery for **opty** and **solve** the optimization: Get DmcD from a predetermined point to another predetermined point, with predetermined angles $q_i$ and final speeds.\\\n",
    "For some reason, final drone speeds $u_x = u_y = 0$ are difficult, while, say, $u_x = u_y = 2.5$ pose no problems. So I iterate to $u_x = u_y \\approx 0$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc7e7050",
   "metadata": {},
   "outputs": [],
   "source": [
    "state_symbols     = tuple((wert for wert in qL))\n",
    "laenge            = len(state_symbols)\n",
    "constant_symbols  = tuple((mD, mM, mP, g, a, reibung, rb, rM, l1))\n",
    "specified_symbols = ((T1, T2, T3, T4))\n",
    "unknown_symbols   = [rP]\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",
    "kin_objective = True                # If True, the kinetic energy of the system will be minimized\n",
    "#======================================================================================================\n",
    "\n",
    "duration  = 10.                                             # The duration of the simulation.\n",
    "num_nodes = 100                                             # The number of nodes to use in the direct collocation problem.\n",
    "t0, tf    = 0., duration                                    # The initial and final times\n",
    "t_int     = duration / 2.                                   # The intermediate time\n",
    "\n",
    "initial_guess = np.ones((len(state_symbols) + len(specified_symbols)) * num_nodes + len(unknown_symbols))\n",
    "\n",
    "interval_value = duration / (num_nodes - 1)\n",
    "\n",
    "# Specify the known system parameters.\n",
    "par_map = OrderedDict()\n",
    "par_map[mD]       = 5.0                                     # Mass of the drone body\n",
    "par_map[mM]       = 0.5                                     # Mass of the motors\n",
    "par_map[mP]       = 0.1                                     # Mass of the propellers\n",
    "par_map[g]        = 9.81                                    # Acceleration due to gravity\n",
    "par_map[a]        = 3.0                                     # Distance of the mass center of the drone from P_i\n",
    "#======================================================================================================\n",
    "#======================================================================================================\n",
    "par_map[reibung]  = 0.                                      # Friction coefficient.\n",
    "par_map[reibungP] = 0.25                                    # Friction coefficient of the propellers\n",
    "#======================================================================================================\n",
    "#======================================================================================================\n",
    "par_map[rb]       = 2.0                                     # Radius of the body\n",
    "par_map[rM]       = 0.5                                     # Radius of the motors\n",
    "par_map[l1]       = 0.5                                     # Distance of the propellers from the mass center\n",
    "\n",
    "if kin_objective == False:\n",
    "    objective = sm.Integral(T1**2 + T2**2 + T3**2 + T4**2, t)   # minimal power\n",
    "else:\n",
    "    kin_energie = sum(body.kinetic_energy(N) for body in BODY)  # The kinetic energy of the system.\n",
    "    kin_energie = me.msubs(kin_energie, par_map)                # Substitute the known parameters.\n",
    "    objective = sm.Integral(kin_energie, t)                     # minimize kinetic Energy\n",
    "\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",
    "#====================================================================================\n",
    "initial_state_constraints = {\n",
    "                            x: 0,\n",
    "                            y: 0.,\n",
    "                            z: 0,\n",
    "                            q1: 0,\n",
    "                            q2: 0,\n",
    "                            q3: 0,\n",
    "                            qp1: 0,\n",
    "                            qp2: 0,\n",
    "                            qp3: 0,\n",
    "                            qp4: 0,\n",
    "                            ux: 0.,\n",
    "                            uy: 0.,\n",
    "                            uz: 0.,\n",
    "                            u1: 0.,\n",
    "                            u2: 0.,\n",
    "                            u3: 0.,\n",
    "                            up1: 0.,\n",
    "                            up2: 0.,\n",
    "                            up3: 0.,\n",
    "                            up4: 0.,\n",
    "                            } \n",
    "\n",
    "\n",
    "intermediate_state_constraints = {\n",
    "                                 x: 50.,\n",
    "                                 y: 50.,\n",
    "                                 z: 75.,\n",
    "                                 q1: 0.,\n",
    "                                 q2: 0.,\n",
    "                                 ux: 0.,   # Value has no meaning, will be set during iteration\n",
    "                                 uy: 0.,   # Value has no meaning, will be set during iteration\n",
    "                                 uz: 0.,   # Value has no meaning, will be set during iteration          \n",
    "                                 u1: 0.,\n",
    "                                 u2: 0.,                 \n",
    "                                }\n",
    "\n",
    "final_state_constraints    = {\n",
    "                            x: 100,\n",
    "                            y: 100.,\n",
    "                            z: 100.,\n",
    "                            q1: 0.,\n",
    "                            q2: 0.,\n",
    "                            ux: 0.,   # Value has no meaning, will be set during iteration\n",
    "                            uy: 0.,   # Value has no meaning, will be set during iteration\n",
    "                            uz: 0.,   # Value has no meaning, will be set during iteration\n",
    "                            u1: 0.,\n",
    "                            u2: 0.,\n",
    "                             }    \n",
    "\n",
    "\n",
    "staerke =  100.                                                 # The maximum strength of the control input\n",
    "bounds = {\n",
    "         T1: (-staerke, staerke),\n",
    "         T2: (-staerke, staerke),\n",
    "         T3: (-staerke, staerke),\n",
    "         T4: (-staerke, staerke),\n",
    "         z:  (0., np.inf),\n",
    "         q1: (-np.pi, np.pi),\n",
    "         q2: (-np.pi, np.pi),\n",
    "         rP: (0.5, 5.)\n",
    "         }\n",
    "    \n",
    "instance_constraints = (\n",
    "                       *tuple(xi.subs({t: t0}) - xi_val for xi, xi_val in initial_state_constraints.items()), \n",
    "                       *tuple(xi.subs({t: t_int}) - xi_val for xi, xi_val in intermediate_state_constraints.items()),\n",
    "                       *tuple(xi.subs({t: tf}) - xi_val for xi, xi_val in final_state_constraints.items()),\n",
    "                       )\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",
    "iterationen = 2\n",
    "teilen = iterationen / 2.5\n",
    "for i in range(iterationen): \n",
    "# for some reason, the intermediate / final states ux = uy = uz = 0 give problems. Hence I itetate towards them\n",
    "    final_state_constraints[ux] = 2.51 - (i+1)/teilen\n",
    "    final_state_constraints[uy] = 2.51 - (i+1)/teilen\n",
    "    final_state_constraints[uz] = 2.51 - (i+1)/teilen\n",
    "\n",
    "    intermediate_state_constraints[ux] = 2.51 - (i+1)/teilen\n",
    "    intermediate_state_constraints[uy] = 2.51 - (i+1)/teilen\n",
    "    intermediate_state_constraints[uz] = 2.51 - (i+1)/teilen\n",
    "\n",
    "    par_map[reibung] = 0.05  * i  # Friction coefficient. May be chaned during in each loop\n",
    "\n",
    "    instance_constraints = (\n",
    "                            *tuple(xi.subs({t: t0}) - xi_val for xi, xi_val in initial_state_constraints.items()), \n",
    "                            *tuple(xi.subs({t: t_int}) - xi_val for xi, xi_val in intermediate_state_constraints.items()),\n",
    "                            *tuple(xi.subs({t: tf}) - xi_val for xi, xi_val in final_state_constraints.items()),\n",
    "                            ) \n",
    "\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', 24000)  # 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",
    "        print(f'radius of propellers {solution[-1]:.3f} \\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 gereralized coordinates and speeds.\\\n",
    "Print, how well the **final_state_constraints** are fulfilled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9954e90e",
   "metadata": {},
   "outputs": [],
   "source": [
    "welche = -1                 # determines the solution of which iteration to use\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",
    "\n",
    "zaehler = -(len(intermediate_state_constraints) + len(final_state_constraints)) - 1\n",
    "for i in intermediate_state_constraints.keys():\n",
    "    zaehler += 1\n",
    "    print(f'error in the intermediate state constraint {i}:   {fehler[zaehler]:.3e}')\n",
    "print('\\n')\n",
    "for i in final_state_constraints.keys():\n",
    "    zaehler += 1\n",
    "    print(f'error in the final state constraint {i}:   {fehler[zaehler]:.3e}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fb77b13",
   "metadata": {},
   "source": [
    "Plot location and speed of the drone, and also the forces acting on it."
   ]
  },
  {
   "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": [
    "**animation**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "420a0087",
   "metadata": {},
   "outputs": [],
   "source": [
    "fps = 20\n",
    "\n",
    "# trace the paths of the drone body and of the propellers\n",
    "traceD  = []\n",
    "traceP1 = []\n",
    "traceP2 = []\n",
    "traceP3 = []\n",
    "traceP4 = []\n",
    "\n",
    "# Copied from Timo. Makes the animation much easier.\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",
    "qL = [*state_symbols, *specified_symbols]\n",
    "par_map[a] = 15.   # I change this for better visibility\n",
    "pL, pL_vals = [*constant_symbols], [par_map[const] for const in constant_symbols]\n",
    "\n",
    "# define final points of the force arrows\n",
    "f1, f2, f3, f4 = sm.symbols('f1 f2 f3 f4')                          # the forces of the propellers, will be calculated below\n",
    "Ff1, Ff2, Ff3, Ff4 = sm.symbols('Ff1, Ff2, Ff3, Ff4', cls = me.Point)\n",
    "Ff1.set_pos(P1, f1 * Ab.z)\n",
    "Ff2.set_pos(P2, f2 * Ab.z)\n",
    "Ff3.set_pos(P3, f3 * Ab.z)\n",
    "Ff4.set_pos(P4, f4 * Ab.z)\n",
    "\n",
    "coordinates = DmcD.pos_from(O).to_matrix(N)\n",
    "for point in [P1, P2, P3, P4, Ff1, Ff2, Ff3, Ff4]:\n",
    "    coordinates = coordinates.row_join(point.pos_from(O).to_matrix(N))\n",
    "coords_lam = sm.lambdify(qL + pL + [f1, f2, f3, f4], coordinates, cse=True)\n",
    "\n",
    "# configuration constraints seem to 'mess up ther order of the states. I need to find the indices of the states, as they are in the configuration constraints\n",
    "# to keep the order of the states in the animation\n",
    "idx = []\n",
    "for ort in initial_state_constraints.keys():\n",
    "    idx.append(qL.index(ort))\n",
    "\n",
    "# needed to give the picture the right size\n",
    "xmin, xmax = np.min(state_vals[idx[0], :]) - par_map[a], np.max(state_vals[idx[0], :]) + par_map[a]\n",
    "ymin, ymax = np.min(state_vals[idx[1], :]) - par_map[a], np.max(state_vals[idx[1], :]) + par_map[a]\n",
    "zmin, zmax = np.min(state_vals[idx[2], :]) - par_map[a], np.max(state_vals[idx[2], :]) + par_map[a]\n",
    "\n",
    "def plot_3d_plane(x_min, x_max, y_min, y_max):\n",
    "    # Create a meshgrid for x and y values\n",
    "    x = np.linspace(x_min, x_max, 100)\n",
    "    y = np.linspace(y_min, y_max, 100)\n",
    "    x, y = np.meshgrid(x, y)\n",
    "\n",
    "    # Z values are set to 0 for the plane\n",
    "    z = np.zeros_like(x)\n",
    "   \n",
    "    # Plot the 3D plane\n",
    "    ax.plot_surface(x, y, z, alpha=0.1, rstride=100, cstride=100, color='c')\n",
    "\n",
    "\n",
    "fig = plt.figure(figsize=(10, 10))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.set_xlim(xmin - 1., xmax + 1.)\n",
    "ax.set_ylim(ymin - 1., ymax + 1.)\n",
    "ax.set_zlim(zmin - 1., zmax + 1.)\n",
    "ax.set_aspect('equal')\n",
    "ebene = plot_3d_plane(xmin - 1., xmax + 1., ymin - 1., ymax + 1.)\n",
    "\n",
    "\n",
    "# Initialize an empty line\n",
    "line1, = ax.plot([], [], [], color='red', marker='o', markersize=10)            # the drone body\n",
    "line2, = ax.plot([], [], [], color='blue', marker='o', markersize=5)            # the propeller 1\n",
    "line3, = ax.plot([], [], [], color='magenta', marker='o', markersize=5)         # the propeller 2\n",
    "line4, = ax.plot([], [], [], color='orange', marker='o', markersize=5)          # the propeller 3\n",
    "line5, = ax.plot([], [], [], color='black', marker='o', markersize=5)           # the propeller 4\n",
    "line6, = ax.plot([], [], [], color='black', lw=0.5, markersize=1)               # connecting P1 to P2\n",
    "line7, = ax.plot([], [], [], color='black', lw=0.5, markersize=1)               # connecting P3 to P4\n",
    "\n",
    "line8, = ax.plot([], [], [], color='red', lw=0.5, markersize=1)                 # trace drome body\n",
    "line9, = ax.plot([], [], [], color='blue', lw=0.25, markersize=1)               # trace propeller 1\n",
    "line10, = ax.plot([], [], [], color='magenta', lw=0.25, markersize=1)           # trace propeller 2\n",
    "line11, = ax.plot([], [], [], color='orange', lw=0.25, markersize=1)            # trace propeller 3\n",
    "line12, = ax.plot([], [], [], color='black', lw=0.25, markersize=1)             # trace propeller 4\n",
    "\n",
    "\n",
    "quiver1 = ax.quiver([], [], [], [], [], [], color='r')                          # the force arrow\n",
    "quiver2 = ax.quiver([], [], [], [], [], [], color='r')                          # the force arrow\n",
    "quiver3 = ax.quiver([], [], [], [], [], [], color='r')                          # the force arrow\n",
    "quiver4 = ax.quiver([], [], [], [], [], [], color='r')                          # the force arrow\n",
    "\n",
    "ax.plot(initial_state_constraints[x], initial_state_constraints[y], initial_state_constraints[z], marker='o', markersize=5, color='red')                     # the starting point\n",
    "ax.plot(intermediate_state_constraints[x], intermediate_state_constraints[y], intermediate_state_constraints[z], marker='o', markersize=5, color='blue')     # the intermediate point\n",
    "ax.plot(final_state_constraints[x], final_state_constraints[y], final_state_constraints[z], marker='o', markersize=5, color='green')                         # the end point\n",
    "\n",
    "# Function to update the plot for each animation frame\n",
    "if kin_objective == True:\n",
    "    nachricht = 'objective function: minimize kinetic energy'\n",
    "else:\n",
    "    nachricht = 'objective function: minimize torques of the propellers'\n",
    "def update(t):\n",
    "    global quiver1, quiver2, quiver3, quiver4\n",
    "    message = (f'running time {t:.2f} sec \\n The green arrows are the forces \\n Optimal radius of the propellers is {solution[-1]:.3f} \\n {nachricht}')\n",
    "    ax.set_title(message, fontsize=15)\n",
    "\n",
    "    skala = 5.0                                                               # scaling factor for the forces, to make them look 'right'. Apparently the only way to do this with quiver in 3D\n",
    "    f11 = solution[-1]**2 * state_sol(t)[idx[-4]] / skala                     # the forces of the propellers\n",
    "    f21 = -solution[-1]**2 * state_sol(t)[idx[-3]] / skala\n",
    "    f31 = solution[-1]**2 * state_sol(t)[idx[-2]] / skala\n",
    "    f41 = -solution[-1]**2 * state_sol(t)[idx[-1]] / skala\n",
    "\n",
    "    coords = coords_lam(*state_sol(t), *input_sol(t), *pL_vals, f11, f21, f31, f41)\n",
    "    line1.set_data([coords[0, 0]], [coords[1, 0]])                          # the drone body\n",
    "    line1.set_3d_properties([coords[2, 0]])\n",
    "    line2.set_data([coords[0, 1]], [coords[1, 1]])                          # the propeller 1\n",
    "    line2.set_3d_properties([coords[2, 1]])\n",
    "    line3.set_data([coords[0, 2]], [coords[1, 2]])                          # the propeller 2\n",
    "    line3.set_3d_properties([coords[2, 2]])\n",
    "    line4.set_data([coords[0, 3]], [coords[1, 3]])                          # the propeller 3\n",
    "    line4.set_3d_properties([coords[2, 3]])\n",
    "    line5.set_data([coords[0, 4]], [coords[1, 4]])                          # the propeller 4\n",
    "    line5.set_3d_properties([coords[2, 4]])\n",
    "    \n",
    "    line6.set_data([coords[0, 1], coords[0, 2]], [coords[1, 1], coords[1, 2]])  # connecting P1 to P2\n",
    "    line7.set_data([coords[0, 3], coords[0, 4]], [coords[1, 3], coords[1, 4]])  # connecting P3 to P4\n",
    "    line6.set_3d_properties([coords[2, 1], coords[2, 2]])\n",
    "    line7.set_3d_properties([coords[2, 3], coords[2, 4]])\n",
    "\n",
    "    quiver1.remove()\n",
    "    quiver2.remove()\n",
    "    quiver3.remove()\n",
    "    quiver4.remove()\n",
    "    quiver1 = ax.quiver(coords[0, 1], coords[1, 1], coords[2, 1], coords[0, 5] - coords[0, 1], coords[1, 5] - coords[1, 1],coords[2, 5] - coords[2, 1], color='g')\n",
    "    quiver2 = ax.quiver(coords[0, 2], coords[1, 2], coords[2, 2], coords[0, 6] - coords[0, 2], coords[1, 6] - coords[1, 2],coords[2, 6] - coords[2, 2], color='g')\n",
    "    quiver3 = ax.quiver(coords[0, 3], coords[1, 3], coords[2, 3], coords[0, 7] - coords[0, 3], coords[1, 7] - coords[1, 3],coords[2, 7] - coords[2, 3], color='g')\n",
    "    quiver4 = ax.quiver(coords[0, 4], coords[1, 4], coords[2, 4], coords[0, 8] - coords[0, 4], coords[1, 8] - coords[1, 4],coords[2, 8] - coords[2, 4], color='g')\n",
    "    \n",
    "    traceD.append([coords[0, 0], coords[1, 0], coords[2, 0]])\n",
    "    line8.set_data(np.array(traceD)[:, 0:2].T)                                     # the trace of the drone body\n",
    "    line8.set_3d_properties(np.array(traceD)[:, 2])\n",
    "\n",
    "    traceP1.append([coords[0, 1], coords[1, 1], coords[2, 1]])                     # the trace of the propeller 1\n",
    "    line9.set_data(np.array(traceP1)[:, 0:2].T)\n",
    "    line9.set_3d_properties(np.array(traceP1)[:, 2])\n",
    "\n",
    "    traceP2.append([coords[0, 2], coords[1, 2], coords[2, 2]])                     # the trace of the propeller 2\n",
    "    line10.set_data(np.array(traceP2)[:, 0:2].T)\n",
    "    line10.set_3d_properties(np.array(traceP2)[:, 2])\n",
    "\n",
    "    traceP3.append([coords[0, 3], coords[1, 3], coords[2, 3]] )                    # the trace of the propeller 3\n",
    "    line11.set_data(np.array(traceP3)[:, 0:2].T)\n",
    "    line11.set_3d_properties(np.array(traceP3)[:, 2])\n",
    "\n",
    "    traceP4.append([coords[0, 4], coords[1, 4], coords[2, 4]])                     # the trace of the propeller 4\n",
    "    line12.set_data(np.array(traceP4)[:, 0:2].T)\n",
    "    line12.set_3d_properties(np.array(traceP4)[:, 2])\n",
    "\n",
    "# Set labels and legend\n",
    "ax.set_xlabel('X-axis')\n",
    "ax.set_ylabel('Y-axis')\n",
    "ax.set_zlabel('Z-axis')\n",
    "\n",
    "# Create the animation\n",
    "animation = FuncAnimation(fig, update, frames=np.arange(t0, tf, 1 / fps), interval=2.5*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()-start0:.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
}