Five-Link-Biped-demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d25b1885",
   "metadata": {},
   "source": [
    "# Five-Link biped walking loop problem: interactive demonstration\n",
    "\n",
    "Hello and welcome. This is a Jupyter Notebook, a kind of document that can alternate between static content, like text and images, and executable cells of code.\n",
    "\n",
    "This document ilustrates the Five-link biped walking loop test case of the paper: \"Collocation Methods for Second Order Systems\", submitted to RSS 2022.\n",
    "\n",
    "In order to run the cells of code, you can select the cell and clic on the small \"play\" button in the bar above or press shift+enter. Alternatively, you can select the option \"run\" -> \"run all cells\" in order to run all the code in order. Beware that some cells can take several minutes!\n",
    "\n",
    "All of the code used in this example is open-source and free to use.\n",
    "\n",
    "[SymPy](https://www.sympy.org/en/index.html) is used for Symbolic formulation and manipulation of the problem.\n",
    "\n",
    "[Numpy](https://numpy.org/) is used for numerical arrays and operations.\n",
    "\n",
    "[CasADI](https://web.casadi.org/) is used for optimization.\n",
    "\n",
    "[Chords](https://github.com/AunSiro/optibot) is the name of the package where we are compiling our code. We aim to produce a toolbox for Optimal Control Problems, focused on robotics, including a high level, readable and clean interface between the prior three packages."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f047156",
   "metadata": {},
   "source": [
    "## Package imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1f6cd9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cf02b82",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy import (symbols, simplify)\n",
    "from sympy.physics.mechanics import dynamicsymbols, init_vprinting\n",
    "from sympy.physics.mechanics import Lagrangian, ReferenceFrame, Point, Particle,inertia, RigidBody, angular_momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf1ec76d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from chords.symbolic import lagrange, diff_to_symb, SimpLagrangesMethod\n",
    "from chords.numpy import unpack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "072a034c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#SymPy vector-like latex rendering inizialization:\n",
    "\n",
    "init_vprinting()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adaeac40",
   "metadata": {},
   "source": [
    "## Symbolic Problem Modelling\n",
    "\n",
    "The first step is to model our problem taking advantage of the high level object syntax of the mechanics module in SymPy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48219687",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating symbols and dynamic symbols\n",
    "\n",
    "m0, m1, m2, m3, m4, l0, l1, l2, l3, l4, t, g = symbols('m_0:5 l_0:5 t g')\n",
    "I0, I1, I2, I3, I4, d0, d1, d2, d3, d4 = symbols('I_0:5 d_0:5')\n",
    "q0, q1, q2, q3, q4 = dynamicsymbols('q_0:5')\n",
    "m0, m1, m2, m3, m4, l0, l1, l2, l3, l4, t, g, I0, I1, I2, I3, I4, d0, d1, d2, d3, d4, q0, q1, q2, q3, q4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "572b34fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Definition of the physics system\n",
    "\n",
    "N_in = ReferenceFrame('N')\n",
    "P0 = Point('P0')\n",
    "P0.set_vel(N_in, 0)\n",
    "\n",
    "N0 = N_in.orientnew('N0', 'Axis', [q0, N_in.z])\n",
    "P1 = P0.locatenew('P1', l0 * N0.y)\n",
    "P1.set_vel(N_in, P1.pos_from(P0).dt(N_in))\n",
    "CM0 = P0.locatenew('CM0', (l0-d0) * N0.y)\n",
    "CM0.set_vel(N_in, CM0.pos_from(P0).dt(N_in))\n",
    "I_0 = inertia(N0, 0, 0, I0)\n",
    "body0 = RigidBody('Stance_Tibia', CM0, N0, m0, (I_0,CM0))\n",
    "body0.potential_energy = m0 * g * CM0.pos_from(P0).dot(N_in.y)\n",
    "\n",
    "\n",
    "N1 = N_in.orientnew('N1', 'Axis', [q1, N_in.z])\n",
    "P2 = P1.locatenew('P2', l1 * N1.y)\n",
    "P2.set_vel(N_in, P2.pos_from(P0).dt(N_in))\n",
    "CM1 = P1.locatenew('CM1', (l1-d1) * N1.y)\n",
    "CM1.set_vel(N_in, CM1.pos_from(P0).dt(N_in))\n",
    "I_1 = inertia(N1, 0, 0, I1)\n",
    "body1 = RigidBody('Stance_Femur', CM1, N1, m1, (I_1,CM1))\n",
    "body1.potential_energy = m1 * g * CM1.pos_from(P0).dot(N_in.y)\n",
    "\n",
    "\n",
    "N2 = N_in.orientnew('N2', 'Axis', [q2, N_in.z])\n",
    "P3 = P2.locatenew('P3', l2 * N2.y)\n",
    "P3.set_vel(N_in, P3.pos_from(P0).dt(N_in))\n",
    "CM2 = P2.locatenew('CM2', d2 * N2.y)\n",
    "CM2.set_vel(N_in, CM2.pos_from(P0).dt(N_in))\n",
    "I_2 = inertia(N2, 0, 0, I2)\n",
    "body2 = RigidBody('Torso', CM2, N2, m2, (I_2,CM2))\n",
    "body2.potential_energy = m2 * g * CM2.pos_from(P0).dot(N_in.y)\n",
    "\n",
    "\n",
    "N3 = N_in.orientnew('N3', 'Axis', [q3, N_in.z])\n",
    "P4 = P2.locatenew('P4', -l3 * N3.y)\n",
    "P4.set_vel(N_in, P4.pos_from(P0).dt(N_in))\n",
    "CM3 = P2.locatenew('CM3', -d3 * N3.y)\n",
    "CM3.set_vel(N_in, CM3.pos_from(P0).dt(N_in))\n",
    "I_3 = inertia(N3, 0, 0, I3)\n",
    "body3 = RigidBody('Swing_Femur', CM3, N3, m3, (I_3,CM3))\n",
    "body3.potential_energy = m3 * g * CM3.pos_from(P0).dot(N_in.y)\n",
    "\n",
    "\n",
    "N4 = N_in.orientnew('N4', 'Axis', [q4, N_in.z])\n",
    "P5 = P4.locatenew('P5', -l4 * N4.y)\n",
    "P5.set_vel(N_in, P5.pos_from(P0).dt(N_in))\n",
    "CM4 = P4.locatenew('CM4', -d4 * N4.y)\n",
    "CM4.set_vel(N_in, CM4.pos_from(P0).dt(N_in))\n",
    "I_4 = inertia(N4, 0, 0, I4)\n",
    "body4 = RigidBody('Swing_Tibia', CM4, N4, m4, (I_4,CM4))\n",
    "body4.potential_energy = m4 * g * CM4.pos_from(P0).dot(N_in.y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95980f54",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Computing the Lagrangian\n",
    "\n",
    "Lag_simp = Lagrangian(N_in, body0, body1, body2, body3, body4)\n",
    "Lag_simp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e7fe237",
   "metadata": {},
   "outputs": [],
   "source": [
    "from chords.symbolic import ImplicitLagrangesMethod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f931cee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defining the control forces and external actions, and applying them to our system\n",
    "\n",
    "u0, u1, u2, u3, u4 = symbols('u_:5')\n",
    "FL = [\n",
    "    (N0, (u0-u1) * N_in.z),\n",
    "    (N1, (u1-u2) * N_in.z),\n",
    "    (N2, (u2-u3) * N_in.z),\n",
    "    (N3, (u3-u4) * N_in.z),\n",
    "    (N4, u4 * N_in.z)\n",
    "]\n",
    "LM_small = ImplicitLagrangesMethod(Lag_simp, [q0, q1, q2, q3, q4], forcelist=FL, frame=N_in)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fe1aa6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generating the dynamic equations\n",
    "\n",
    "LM_small.form_lagranges_equations()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2735e6b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "impl_x = LM_small.implicit_dynamics_x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c1e2b7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "impl_q = LM_small.implicit_dynamics_q"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4723d4d1",
   "metadata": {},
   "source": [
    "### Generating auxiliar functions\n",
    "\n",
    "Later in the problem we need some expressions derived from the problem. Here we will generate them as symbolic expressions, and then convert them to numerical functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba987f28",
   "metadata": {},
   "outputs": [],
   "source": [
    "import casadi as cas\n",
    "from sympy import lambdify\n",
    "from chords.casadi import implicit_dynamic_q_to_casadi_function, implicit_dynamic_x_to_casadi_function, sympy2casadi\n",
    "from chords.symbolic import find_arguments, diff_to_symb_expr\n",
    "from sympy.physics.mechanics import kinetic_energy, potential_energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7253d31",
   "metadata": {},
   "outputs": [],
   "source": [
    "imp_dyn_x_f_cas = implicit_dynamic_x_to_casadi_function(impl_x, list(dynamicsymbols('x_0:10')), verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94932d92",
   "metadata": {},
   "outputs": [],
   "source": [
    "imp_dyn_q_f_cas = implicit_dynamic_q_to_casadi_function(impl_q, list(LM_small.q), verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2721faa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "imp_dyn_x_f_cas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62c52025",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "feet_x = P5.pos_from(P0).dot(N_in.x)\n",
    "feet_x = diff_to_symb_expr(feet_x)\n",
    "feet_x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99ab5029",
   "metadata": {},
   "outputs": [],
   "source": [
    "feet_y = P5.pos_from(P0).dot(N_in.y)\n",
    "feet_y = diff_to_symb_expr(feet_y)\n",
    "feet_y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "828f68d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "feet_y_vel = P5.vel(N_in).dot(N_in.y)   #pos_from(P0).dot(N_in.y)\n",
    "feet_y_vel = diff_to_symb_expr(feet_y_vel)\n",
    "feet_y_vel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce528447",
   "metadata": {},
   "outputs": [],
   "source": [
    "cm_pos = m0*CM0.pos_from(P0)\n",
    "cm_pos += m1*CM1.pos_from(P0)\n",
    "cm_pos += m2*CM2.pos_from(P0)\n",
    "cm_pos += m3*CM3.pos_from(P0)\n",
    "cm_pos += m4*CM4.pos_from(P0)\n",
    "cm_pos = cm_pos/(m0+m1+m2+m3+m4)\n",
    "\n",
    "sys_CM = P0.locatenew('Sys_CM', cm_pos)\n",
    "sys_CM_x = simplify(sys_CM.pos_from(P0).dot(N_in.x))\n",
    "sys_CM_y = simplify(sys_CM.pos_from(P0).dot(N_in.y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc652dcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "sym_x = dynamicsymbols('q_0:5')\n",
    "sym_x = sym_x + [ii.diff() for ii in sym_x]\n",
    "sym_x = [diff_to_symb(ii) for ii in sym_x]\n",
    "sym_params = list(symbols('I_0:5 d_0:5 g l_0:2 l_3 m_0:5'))\n",
    "sym_add_params = [symbols('l_4'),]\n",
    "sym_vars = sym_x + sym_params + sym_add_params\n",
    "\n",
    "print(len(sym_vars), sym_vars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec81e447",
   "metadata": {},
   "outputs": [],
   "source": [
    "cas_x_args = cas.MX.sym(\"x\", len(sym_x))\n",
    "cas_params = cas.MX.sym(\"p\", len(sym_params))\n",
    "cas_add_params = cas.MX.sym(\"p_add\", len(sym_add_params))\n",
    "cas_all_vars = [cas_x_args[ii] for ii in range(len(sym_x))]\n",
    "cas_all_vars += [cas_params[ii] for ii in range(len(sym_params))]\n",
    "cas_all_vars += [cas_add_params[ii] for ii in range(len(sym_add_params))]\n",
    "print(len(cas_all_vars), cas_all_vars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5322ec96",
   "metadata": {},
   "outputs": [],
   "source": [
    "_cas_expr_temp_x = sympy2casadi(feet_x, sym_vars, cas_all_vars)\n",
    "feet_x_cas = cas.Function(\n",
    "        \"Feet_x\",\n",
    "        [cas_x_args, cas_params, cas_add_params],\n",
    "        [_cas_expr_temp_x,],\n",
    "        [\"x\", \"params\", \"additional_params\"],\n",
    "        [\"feet_x_position\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78ab102b",
   "metadata": {},
   "outputs": [],
   "source": [
    "_cas_expr_temp_y = sympy2casadi(feet_y, sym_vars, cas_all_vars)\n",
    "feet_y_cas = cas.Function(\n",
    "        \"Feet_y\",\n",
    "        [cas_x_args, cas_params, cas_add_params],\n",
    "        [_cas_expr_temp_y,],\n",
    "        [\"x\", \"params\", \"additional_params\"],\n",
    "        [\"feet_y_position\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f0cfbd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "_cas_expr_temp_y_vel = sympy2casadi(feet_y_vel, sym_vars, cas_all_vars)\n",
    "feet_y_vel_cas = cas.Function(\n",
    "        \"Feet_y_vel\",\n",
    "        [cas_x_args, cas_params, cas_add_params],\n",
    "        [_cas_expr_temp_y_vel,],\n",
    "        [\"x\", \"params\", \"additional_params\"],\n",
    "        [\"feet_y_speed\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "271543e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simetric_cond_casadi(n = 5):\n",
    "    x1 = cas.MX.sym('x_1', 2*n)\n",
    "    x2 = cas.MX.sym('x_2', 2*n)\n",
    "    cond = [x1[ii] - x2[n-1-ii] for ii in range(n)]\n",
    "    cas_funcs = cas.horzcat(*cond)\n",
    "    return cas.Function(\n",
    "        \"Sim_cond\",\n",
    "        [x1, x2],\n",
    "        [cas_funcs,],\n",
    "        [\"x_1\", \"x2\"],\n",
    "        [\"residue\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dc41926",
   "metadata": {},
   "outputs": [],
   "source": [
    "simetric_5_links = simetric_cond_casadi(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e20a65f1",
   "metadata": {},
   "source": [
    "Creating and simplifying symbolically the expressions of the heel impact may require some time, but it alows for a faster problem formulation later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a8b731f",
   "metadata": {},
   "outputs": [],
   "source": [
    "bodies = [body0, body1, body2, body3, body4]\n",
    "points_right = [P0, P1, P2, P2, P4]\n",
    "points_left = [P5, P4, P2, P2, P1]\n",
    "subs_key = list(zip(dynamicsymbols('q_0:5'),dynamicsymbols('q_p_0:5')))\n",
    "\n",
    "impact_eqs = []\n",
    "for ii in range(5):\n",
    "    print('calculating eq', ii)\n",
    "    print('\\tleft side')\n",
    "    left_side = angular_momentum(points_left[ii], N_in, *bodies[:5-ii]).dot(N_in.z)\n",
    "    left_side = simplify(left_side)\n",
    "    print('\\tright side')\n",
    "    right_side = angular_momentum(points_right[ii], N_in, *bodies[ii:]).dot(N_in.z)\n",
    "    right_side = simplify(right_side).subs(subs_key)\n",
    "    impact_eqs.append(left_side-right_side)\n",
    "#impact_eqs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a75f77f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def impact_cond_casadi(eqs, x1_sym, x2_sym, sym_params, sym_add_params):\n",
    "    x1_sym = [diff_to_symb(ii) for ii in x1_sym]\n",
    "    x2_sym = [diff_to_symb(ii) for ii in x2_sym]\n",
    "    eqs = [diff_to_symb_expr(ii) for ii in eqs]\n",
    "\n",
    "    all_vars = x1_sym + x2_sym + sym_params + sym_add_params\n",
    "    n = len(x1_sym)\n",
    "    cas_x1 = cas.MX.sym('x_1', n)\n",
    "    cas_x2 = cas.MX.sym('x_2', n)\n",
    "    cas_params = cas.MX.sym(\"p\", len(sym_params))\n",
    "    cas_add_params = cas.MX.sym(\"p_add\", len(sym_add_params))\n",
    "    cas_all_vars = [cas_x1[ii] for ii in range(n)]\n",
    "    cas_all_vars += [cas_x2[ii] for ii in range(n)]\n",
    "    cas_all_vars += [cas_params[ii] for ii in range(len(sym_params))]\n",
    "    cas_all_vars += [cas_add_params[ii] for ii in range(len(sym_add_params))]\n",
    "    \n",
    "    cas_funcs = []\n",
    "    for function in eqs:\n",
    "        cas_funcs.append(sympy2casadi(function, all_vars, cas_all_vars))\n",
    "    cas_funcs = cas.horzcat(*cas_funcs)\n",
    "    return cas.Function(\n",
    "        \"Sim_cond\",\n",
    "        [cas_x1, cas_x2, cas_params, cas_add_params],\n",
    "        [cas_funcs,],\n",
    "        [\"x_1\", \"x2\", 'params', 'additional_params'],\n",
    "        [\"residue\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac93f537",
   "metadata": {},
   "outputs": [],
   "source": [
    "sym_x = dynamicsymbols('q_0:5')\n",
    "sym_x = sym_x + [ii.diff() for ii in sym_x]\n",
    "subs_key = list(zip(dynamicsymbols('q_0:5'),dynamicsymbols('q_p_0:5')))\n",
    "sym_x_2 = [ii.subs(subs_key) for ii in sym_x]\n",
    "impact_cond_cas_f = impact_cond_casadi(impact_eqs, sym_x,  sym_x_2, sym_params, sym_add_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c582149f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sys_cm_np = lambdify([sym_x, sym_params], [sys_CM_x, sys_CM_y],'numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8edb4b5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ang_mom_p0 = angular_momentum(P0, N_in, *bodies).dot(N_in.z)\n",
    "ang_mom_p0_np = lambdify([sym_x, sym_params], ang_mom_p0,'numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8321b26d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ang_mom_p5 = angular_momentum(P5, N_in, *bodies).dot(N_in.z)\n",
    "ang_mom_p5_np = lambdify([sym_x, sym_params, sym_add_params], ang_mom_p5,'numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0c745d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "P5_static = P5.locatenew('P5_static', 0 * N_in.y)\n",
    "P5_static.set_vel(N_in, 0 * N_in.y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ca1b8cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "ang_mom_p5_static = angular_momentum(P5_static, N_in, *bodies).dot(N_in.z)\n",
    "ang_mom_p5_static_np = lambdify([sym_x, sym_params, sym_add_params], ang_mom_p5_static,'numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7739c8d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "angular_momentum(P0, N_in, bodies[0]).dot(N_in.z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "265729fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_energy = potential_energy(*bodies) + kinetic_energy(N_in, *bodies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03bd4e11",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_energy_np = lambdify([sym_x, sym_params], system_energy,'numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12ac30f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "mass_matrix_np = lambdify([sym_x, sym_params], LM_small.mass_matrix,'numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6678011",
   "metadata": {},
   "outputs": [],
   "source": [
    "sym_u = symbols('u_:5')\n",
    "F_impl_np = lambdify([sym_x, sym_u, sym_params], LM_small.forcing,'numpy')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8ef0a60",
   "metadata": {},
   "source": [
    "### Scheme definitions\n",
    "\n",
    "Each scheme is defined here as a function that must be equal to zero at each interval.\n",
    "Note that functions that contain \"mod\" in the name are those we define as \"second order\",\n",
    "and use separate conditions for q and v.\n",
    "\n",
    "Note that we will operate this problem wihout combining the scheme equations $F(q_k, q_{k+1}, q'_k, q'_{k+1} q''_k, q''_{k+1}, SchemeParams) = 0$ and the dynamics equations $H(q, q', q'', u, params) = 0$ imposed at the collocation points. This approach allows us to solve this problem without inverting the mass matrix.\n",
    "\n",
    "If you wish to define your own schemes, do it here.\n",
    "\n",
    "Be careful to respect the function structure:\n",
    "\n",
    "    restriction(x, x_n, a, a_n, dt, scheme_params) = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce78cc29",
   "metadata": {},
   "outputs": [],
   "source": [
    "from chords.piecewise import index_div\n",
    "from copy import copy\n",
    "\n",
    "def euler_accel_restr(x, x_n, a, a_n, dt, scheme_params):\n",
    "    first_ind, last_ind = index_div(x)\n",
    "    x_d = copy(x)\n",
    "    x_d[first_ind] = x[last_ind]\n",
    "    x_d[last_ind] = a\n",
    "    return x_n - (x + dt * x_d)\n",
    "\n",
    "\n",
    "def trapz_accel_restr(x, x_n, a, a_n, dt, scheme_params):\n",
    "    first_ind, last_ind, all_but_first_ind, all_but_last_ind = index_div(x)\n",
    "    x_d = copy(x)\n",
    "    x_d[first_ind] = x[last_ind]\n",
    "    x_d[last_ind] = a\n",
    "    x_d_n = copy(x)\n",
    "    x_d_n[first_ind] = x_n[last_ind]\n",
    "    x_d_n[last_ind] = a_n\n",
    "    return x_n - (x + dt / 2 * (x_d + x_d_n))\n",
    "\n",
    "\n",
    "def trapz_mod_accel_restr(x, x_n, a, a_n, dt, scheme_params):\n",
    "    res = copy(x)\n",
    "    first_ind, last_ind, all_but_first_ind, all_but_last_ind = index_div(x)\n",
    "    res[last_ind] = x[last_ind] + dt / 2 * (a + a_n)\n",
    "    res[first_ind] = x[first_ind] + dt * x[last_ind] + dt ** 2 / 6 * (a_n + 2 * a)\n",
    "    return x_n - res\n",
    "\n",
    "\n",
    "def hs_half_x(x, x_n, x_d, x_d_n, dt):\n",
    "    x_c = (x + x_n) / 2 + dt / 8 * (x_d - x_d_n)\n",
    "    return x_c\n",
    "\n",
    "\n",
    "def hs_accel_restr(x, x_n, a, a_n, dt, scheme_params):\n",
    "    a_c = scheme_params\n",
    "    first_ind, last_ind, all_but_first_ind, all_but_last_ind = index_div(x)\n",
    "    x_d = copy(x)\n",
    "    x_d[first_ind] = x[last_ind]\n",
    "    x_d[last_ind] = a\n",
    "\n",
    "    x_d_n = copy(x)\n",
    "    x_d_n[first_ind] = x_n[last_ind]\n",
    "    x_d_n[last_ind] = a_n\n",
    "\n",
    "    x_c = hs_half_x(x, x_n, x_d, x_d_n, dt)\n",
    "    x_d_c = copy(x)\n",
    "    x_d_c[first_ind] = x_c[last_ind]\n",
    "    x_d_c[last_ind] = a_c\n",
    "    return x + dt / 6 * (x_d + 4 * x_d_c + x_d_n) - x_n\n",
    "\n",
    "\n",
    "def hs_mod_half_x(x, x_n, a, a_n, dt):\n",
    "    x_c = copy(x)\n",
    "    first_ind, last_ind, all_but_first_ind, all_but_last_ind = index_div(x)\n",
    "    q = x[first_ind]\n",
    "    v = x[last_ind]\n",
    "    q_n = x_n[first_ind]\n",
    "    v_n = x_n[last_ind]\n",
    "    q_c = q + dt / 32 * (13 * v + 3 * v_n) + dt**2 / 192 * (11 * a - 5 * a_n)\n",
    "    v_c = (v + v_n) / 2 + dt / 8 * (a - a_n)\n",
    "    x_c[first_ind] = q_c\n",
    "    x_c[last_ind] = v_c\n",
    "    return x_c\n",
    "\n",
    "\n",
    "def hs_mod_accel_restr(x, x_n, a, a_n, dt, scheme_params):\n",
    "    a_c = scheme_params\n",
    "    res = copy(x)\n",
    "    first_ind, last_ind, all_but_first_ind, all_but_last_ind = index_div(x)\n",
    "    q = x[first_ind]\n",
    "    v = x[last_ind]\n",
    "    res[last_ind] = v + dt / 6 * (a + 4 * a_c + a_n)\n",
    "    res[first_ind] = q + dt * v + dt ** 2 / 6 * (a + 2 * a_c)\n",
    "    return x_n - res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0f69168",
   "metadata": {},
   "source": [
    "### Casadi optimization\n",
    "\n",
    "We have generated the system equations symbolicaly. Now, we translate them to CasADi objects in order to perform the optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89045c9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from chords.casadi import accelrestriction2casadi\n",
    "#from chords.schemes import (euler_accel_restr, trapz_accel_restr, trapz_mod_accel_restr,\n",
    "#                             hs_mod_accel_restr, hs_accel_restr, hs_half_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed677222",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Numerical values of the paramenters\n",
    "\n",
    "I_0_n, I_1_n, I_2_n, I_3_n, I_4_n = 0.93, 1.08, 2.22, 1.08, 0.93\n",
    "d_0_n, d_1_n, d_2_n, d_3_n, d_4_n = 0.128, 0.163, 0.2, 0.163, 0.128\n",
    "g_n = 9.81\n",
    "l_0_n, l_1_n, l_2_n, l_3_n, l_4_n = 0.4, 0.4, 0.625, 0.4, 0.4\n",
    "m_0_n, m_1_n, m_2_n, m_3_n, m_4_n = 3.2, 6.8, 20, 6.8, 3.2\n",
    "params = [\n",
    "    I_0_n, I_1_n, I_2_n, I_3_n, I_4_n,\n",
    "    d_0_n, d_1_n, d_2_n, d_3_n, d_4_n,\n",
    "    g_n,\n",
    "    l_0_n, l_1_n, l_3_n,\n",
    "    m_0_n, m_1_n, m_2_n, m_3_n, m_4_n\n",
    "]\n",
    "additional_params = [l_4_n,]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a91b399",
   "metadata": {},
   "outputs": [],
   "source": [
    "opti = cas.Opti()\n",
    "p_opts = {}#{\"expand\":True,'ipopt.print_level':0, 'print_time':0}\n",
    "s_opts = {}#{\"max_iter\": 10000, 'tol': 1e-26}#, 'linear_solver' : \"MA27\"}\n",
    "opti.solver(\"ipopt\",p_opts,\n",
    "                    s_opts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a495fea",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 25\n",
    "X = opti.variable(N+1,10)\n",
    "X_dot = opti.variable(N+1,10)\n",
    "U = opti.variable(N+1,5)\n",
    "U_c = opti.variable(N,5)\n",
    "X_c = opti.variable(N,10)\n",
    "X_dot_c = opti.variable(N,10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fb416e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "T = opti.parameter()\n",
    "u_m = opti.parameter()\n",
    "Params_opti = opti.parameter(len(params))\n",
    "Add_params_opti = opti.parameter(len(additional_params))\n",
    "D = opti.parameter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa2d65a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Definition of the cost function\n",
    "\n",
    "#cost = cas.sum2((cas.sum1(U[:,:]**2)+cas.sum1(U[1:-1,:]**2))/N)\n",
    "cost = cas.sum2((4*cas.sum1(U_c[:,:]**2) + cas.sum1(U[:,:]**2)+cas.sum1(U[1:-1,:]**2))/(3*N))\n",
    "#cost = cas.sum2(cas.sum1(U**2))\n",
    "opti.minimize(cost)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51fcabcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Periodic gait constraint:\n",
    "opti.subject_to(simetric_5_links(X[0,:], X[-1,:]) == 0)\n",
    "opti.subject_to(impact_cond_cas_f(X[-1,:], X[0,:], Params_opti, Add_params_opti) == 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bb8a74d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Step size constraint:\n",
    "opti.subject_to(feet_x_cas(X[-1,:], Params_opti, Add_params_opti) == D)\n",
    "opti.subject_to(feet_y_cas(X[-1,:], Params_opti, Add_params_opti) == 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b3c88be",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Small Feet Conditions:\n",
    "opti.subject_to(U[:,0] == 0)\n",
    "opti.subject_to(U_c[:,0] == 0)\n",
    "opti.subject_to(feet_y_vel_cas(X[0,:], Params_opti, Add_params_opti)>0)\n",
    "opti.subject_to(feet_y_vel_cas(X[-1,:], Params_opti, Add_params_opti)<0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66514d84",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Feet over ground Restrictions:\n",
    "for ii in range(1,N):\n",
    "    opti.subject_to(feet_y_cas(X[ii,:], Params_opti, Add_params_opti) > 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3051db54",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Dynamics Constraints:\n",
    "for ii in range(N+1):\n",
    "    opti.subject_to(imp_dyn_x_f_cas(X[ii,:], X_dot[ii,:], U[ii,:], [], Params_opti) == 0)\n",
    "for ii in range(N):\n",
    "    opti.subject_to(X_c[ii,:] == hs_mod_half_x(X[ii,:], X[ii+1,:], X_dot[ii,5:], X_dot[ii+1,5:], T/N))\n",
    "    opti.subject_to(imp_dyn_x_f_cas(X_c[ii,:], X_dot_c[ii,:], U_c[ii,:], [], Params_opti) == 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2e51940",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#Scheme Constraints\n",
    "#cas_accel_restr = accelrestriction2casadi(trapz_mod_accel_restr, 5)\n",
    "cas_accel_restr = accelrestriction2casadi(hs_mod_accel_restr, 5, 5)\n",
    "for ii in range(N):\n",
    "    opti.subject_to(cas_accel_restr(X[ii,:], X[ii+1,:], X_dot[ii, 5:], X_dot[ii+1, 5:],T/N, X_dot_c[ii,5:]) == 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "372760e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "opti.set_value(T, 0.7)#0.7\n",
    "opti.set_value(D, 0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a160b61d",
   "metadata": {},
   "outputs": [],
   "source": [
    "opti.set_value(Params_opti, params)\n",
    "opti.set_value(Add_params_opti, additional_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f0b9de3",
   "metadata": {},
   "outputs": [],
   "source": [
    "q_0_guess = np.array([-0.3, 0.7, 0, -0.5, -0.6])\n",
    "q_1_guess = q_0_guess[::-1]\n",
    "s_arr = np.linspace(0, 1, N+1)\n",
    "q_guess = np.expand_dims(q_0_guess,0)+ np.expand_dims(s_arr,1)*np.expand_dims((q_1_guess - q_0_guess),0)\n",
    "\n",
    "q_dot_guess = (q_1_guess - q_0_guess) * np.ones([N+1,1])/opti.value(T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97e731fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "opti.set_initial(X[:,:5], q_guess)\n",
    "opti.set_initial(X[:,5:], q_dot_guess)\n",
    "\n",
    "opti.set_initial(X_c[:,:5], (q_guess[:-1,:]+q_guess[1:,:])/2)\n",
    "opti.set_initial(X_c[:,5:], q_dot_guess[:-1,:])\n",
    "\n",
    "opti.set_initial(X_dot[:,:5], q_dot_guess)\n",
    "opti.set_initial(X_dot[:,5:], 0)\n",
    "\n",
    "opti.set_initial(X_dot_c[:,:5], q_dot_guess[:-1,:])\n",
    "opti.set_initial(X_dot_c[:,5:], 0)\n",
    "\n",
    "opti.set_initial(U, 0)\n",
    "opti.set_initial(U_c, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24c36bae",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "sol = opti.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9a994d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "U_sol = sol.value(U)\n",
    "U_c_sol = sol.value(U_c)\n",
    "X_sol = sol.value(X)\n",
    "X_c_sol = sol.value(X_c)\n",
    "X_dot_sol = sol.value(X_dot)\n",
    "X_dot_c_sol = sol.value(X_dot_c)\n",
    "T_sol = sol.value(T)\n",
    "T_sol_arr = np.linspace(0, T_sol, N+1)\n",
    "T_c_arr = (T_sol_arr[:-1]+T_sol_arr[1:])/2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff18ba9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=[14,10])\n",
    "labels= ['stance anckle', 'stance knee', 'stance hip', 'swing hip', 'swing knee']\n",
    "for ii in range(5):\n",
    "    plt.plot(T_sol_arr,U_sol[:,ii], marker = 'o', label = labels[ii] + ' u_k')\n",
    "    plt.plot(T_c_arr,U_c_sol[:,ii], 'o', label = labels[ii] + ' u_c')\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "plt.title('u(t)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41478613",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=[14,10])\n",
    "labels= ['stance tibia', 'stance femur', 'torso', 'swing femur', 'swing tibia']\n",
    "for ii in range(5):\n",
    "    plt.plot(T_sol_arr, X_sol[:,ii], marker = 'o', label = labels[ii] + ' q_k')\n",
    "    plt.plot(T_c_arr,X_c_sol[:,ii], 'o', label = labels[ii] + ' q_c')\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "plt.title('q(t)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b133395",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chain_to_draw(x,params):\n",
    "    [\n",
    "    I_0_n, I_1_n, I_2_n, I_3_n, I_4_n,\n",
    "    d_0_n, d_1_n, d_2_n, d_3_n, d_4_n,\n",
    "    g_n,\n",
    "    l_0_n, l_1_n, l_3_n,\n",
    "    m_0_n, m_1_n, m_2_n, m_3_n, m_4_n\n",
    "    ] = params\n",
    "    points_x = [0, ]\n",
    "    points_y = [0, ]\n",
    "    points_x.append(points_x[-1] - l_0_n*np.sin(x[0]))\n",
    "    points_x.append(points_x[-1] - l_1_n*np.sin(x[1]))\n",
    "    points_x.append(points_x[-1] - l_2_n*np.sin(x[2]))\n",
    "    points_x.append(points_x[-2])\n",
    "    points_x.append(points_x[-1] + l_3_n*np.sin(x[3]))\n",
    "    points_x.append(points_x[-1] + l_4_n*np.sin(x[4]))\n",
    "    \n",
    "    \n",
    "    points_y.append(points_y[-1] + l_0_n*np.cos(x[0]))\n",
    "    points_y.append(points_y[-1] + l_1_n*np.cos(x[1]))\n",
    "    points_y.append(points_y[-1] + l_2_n*np.cos(x[2]))\n",
    "    points_y.append(points_y[-2])\n",
    "    points_y.append(points_y[-1] - l_3_n*np.cos(x[3]))\n",
    "    points_y.append(points_y[-1] - l_4_n*np.cos(x[4]))\n",
    "    \n",
    "    return points_x, points_y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce54d361",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "points_x, points_y = chain_to_draw(X_sol[0], params)\n",
    "plt.figure(figsize=[15,15])\n",
    "plt.grid()\n",
    "for ii in range(0, N, 1):\n",
    "    points_x, points_y = chain_to_draw(X_sol[ii], params)\n",
    "    plt.plot(points_x, points_y, lw=1,  color = plt.cm.viridis(ii/N))\n",
    "    \n",
    "\n",
    "plt.gca().set_aspect('equal')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cac3d5f4",
   "metadata": {},
   "source": [
    "total_mass = m_0_n + m_1_n + m_2_n + m_3_n + m_4_n\n",
    "ang_mom_arr = [ang_mom_p0_np(X_sol[ii,:],params) for ii in range(N+1)]\n",
    "ang_mom_swing_foot_arr = [ang_mom_p5_np(X_sol[ii,:],params, additional_params) for ii in range(N+1)]\n",
    "ang_mom_swing_foot_static_arr = [ang_mom_p5_static_np(X_sol[ii,:],params, additional_params) for ii in range(N+1)]\n",
    "cm_torque_arr = [total_mass * -g_n * sys_cm_np(X_sol[ii,:], params)[0] for ii in range(N+1)]\n",
    "ang_mom_arr_deriv = np.gradient(ang_mom_arr, T_sol_arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c13008e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from chords.piecewise import interpolated_array, interpolated_array_derivative\n",
    "from chords.analysis import dynamic_error_implicit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13728f2a",
   "metadata": {},
   "source": [
    "## Sistematic comparative of schemes for different values of N\n",
    "\n",
    "Now let's solve the problem with different methods.\n",
    "\n",
    "### Caution!\n",
    "\n",
    "Executing the next cell may require some time!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14b8d250",
   "metadata": {},
   "outputs": [],
   "source": [
    "def q_init(N):\n",
    "    q_0_guess = np.array([-0.3, 0.7, 0, -0.5, -0.6])\n",
    "    q_1_guess = q_0_guess[::-1]\n",
    "    s_arr = np.linspace(0, 1, N+1)\n",
    "    q_guess = np.expand_dims(q_0_guess,0)+ np.expand_dims(s_arr,1)*np.expand_dims((q_1_guess - q_0_guess),0)\n",
    "    q_dot_guess = (q_1_guess - q_0_guess) * np.ones([N+1,1])/opti.value(T)\n",
    "    return q_guess, q_dot_guess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a045db4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "def chrono_solve(opti, solve_repetitions):\n",
    "    cput0 = time.time()\n",
    "    for ii in range(solve_repetitions):\n",
    "        sol = opti.solve()\n",
    "    cput1 = time.time()\n",
    "    cpudt = (cput1-cput0)/solve_repetitions\n",
    "    return sol, cpudt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "017cb9c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def casadi_biped(N = 25, scheme = \"trapz\", solve_repetitions = 1, t_end = 0.7, step_length = 0.5):\n",
    "    opti = cas.Opti()\n",
    "    p_opts = {\"expand\":True,'ipopt.print_level':0, 'print_time':0}\n",
    "    s_opts = {\"max_iter\": 10000, 'tol': 1e-26}#, 'linear_solver' : \"MA27\"}\n",
    "    opti.solver(\"ipopt\",p_opts,\n",
    "                        s_opts)\n",
    "    \n",
    "    restr_schemes = {\n",
    "        'trapz': trapz_accel_restr,\n",
    "        'trapz_mod' : trapz_mod_accel_restr,\n",
    "        'hs': hs_accel_restr,\n",
    "        'hs_mod': hs_mod_accel_restr,\n",
    "        'hs_parab': hs_accel_restr,\n",
    "        'hs_mod_parab': hs_mod_accel_restr\n",
    "    }\n",
    "    \n",
    "    f_restr = restr_schemes[scheme]\n",
    "\n",
    "    X = opti.variable(N+1,10)\n",
    "    X_dot = opti.variable(N+1,10)\n",
    "    U = opti.variable(N+1,5)\n",
    "    if 'hs' in scheme:\n",
    "        U_c = opti.variable(N,5)\n",
    "        X_c = opti.variable(N,10)\n",
    "        X_dot_c = opti.variable(N,10)\n",
    "\n",
    "    T = opti.parameter()\n",
    "    u_m = opti.parameter()\n",
    "    Params_opti = opti.parameter(len(params))\n",
    "    Add_params_opti = opti.parameter(len(additional_params))\n",
    "    D = opti.parameter()\n",
    "\n",
    "    # Cost\n",
    "    if 'parab' in scheme:\n",
    "        cost = cas.sum2((4*cas.sum1(U_c[:,:]**2) + cas.sum1(U[:,:]**2)+cas.sum1(U[1:-1,:]**2))/(3*N))\n",
    "    else:\n",
    "        cost = cas.sum2((cas.sum1(U[:,:]**2)+cas.sum1(U[1:-1,:]**2))/N)\n",
    "    #cost = cas.sum2(cas.sum1(U**2))\n",
    "    opti.minimize(cost)\n",
    "\n",
    "    #Periodic gait constraint:\n",
    "    opti.subject_to(simetric_5_links(X[0,:], X[-1,:]) == 0)\n",
    "    opti.subject_to(impact_cond_cas_f(X[-1,:], X[0,:], Params_opti, Add_params_opti) == 0)\n",
    "\n",
    "    #Step size constraint:\n",
    "    opti.subject_to(feet_x_cas(X[-1,:], Params_opti, Add_params_opti) == D)\n",
    "    opti.subject_to(feet_y_cas(X[-1,:], Params_opti, Add_params_opti) == 0)\n",
    "\n",
    "    #Small Feet Conditions:\n",
    "    opti.subject_to(U[:,0] == 0)\n",
    "    opti.subject_to(feet_y_vel_cas(X[0,:], Params_opti, Add_params_opti)>0)\n",
    "    opti.subject_to(feet_y_vel_cas(X[-1,:], Params_opti, Add_params_opti)<0)\n",
    "    if 'hs' in scheme:\n",
    "        opti.subject_to(U_c[:,0] == 0)\n",
    "\n",
    "    #Feet over ground Restrictions:\n",
    "    for ii in range(1,N):\n",
    "        opti.subject_to(feet_y_cas(X[ii,:], Params_opti, Add_params_opti) > 0)\n",
    "\n",
    "    #Dynamics Constraints:\n",
    "    for ii in range(N+1):\n",
    "        opti.subject_to(imp_dyn_x_f_cas(X[ii,:], X_dot[ii,:], U[ii,:], [], Params_opti) == 0)\n",
    "    if 'hs' in scheme:\n",
    "        for ii in range(N):\n",
    "            opti.subject_to(X_c[ii,:] == hs_half_x(X[ii,:], X[ii+1,:], X_dot[ii,:], X_dot[ii+1,:], T/N))\n",
    "            opti.subject_to(imp_dyn_x_f_cas(X_c[ii,:], X_dot_c[ii,:], U_c[ii,:], [], Params_opti) == 0)\n",
    "        if 'parab' not in scheme:\n",
    "            for ii in range(N):\n",
    "                opti.subject_to(U_c[ii,:] == (U[ii,:]+U[ii+1,:])/2)\n",
    "\n",
    "    #Scheme Constraints\n",
    "    if 'hs' in scheme:\n",
    "        cas_accel_restr = accelrestriction2casadi(f_restr, 5, 5)\n",
    "        for ii in range(N):\n",
    "            opti.subject_to(cas_accel_restr(X[ii,:], X[ii+1,:], X_dot[ii, 5:], X_dot[ii+1, 5:],T/N, X_dot_c[ii,5:]) == 0)\n",
    "    else:\n",
    "        cas_accel_restr = accelrestriction2casadi(f_restr, 5)\n",
    "        for ii in range(N):\n",
    "            opti.subject_to(cas_accel_restr(X[ii,:], X[ii+1,:], X_dot[ii, 5:], X_dot[ii+1, 5:],T/N, []) == 0)\n",
    "    \n",
    "\n",
    "    opti.set_value(T, t_end)#0.7\n",
    "    opti.set_value(D, step_length)#0.5\n",
    "\n",
    "    opti.set_value(Params_opti, params)\n",
    "    opti.set_value(Add_params_opti, additional_params)\n",
    "    \n",
    "    q_guess, q_dot_guess = q_init(N)\n",
    "    opti.set_initial(X[:,:5], q_guess)\n",
    "    opti.set_initial(X[:,5:], q_dot_guess)\n",
    "    \n",
    "    opti.set_initial(X_dot[:,:5], q_dot_guess)\n",
    "    opti.set_initial(X_dot[:,5:], 0)\n",
    "    opti.set_initial(U, 0)\n",
    "    \n",
    "    if 'hs' in scheme:\n",
    "        opti.set_initial(X_c[:,:5], (q_guess[:-1,:]+q_guess[1:,:])/2)\n",
    "        opti.set_initial(X_c[:,5:], q_dot_guess[:-1,:])\n",
    "\n",
    "        opti.set_initial(X_dot_c[:,:5], q_dot_guess[:-1,:])\n",
    "        opti.set_initial(X_dot_c[:,5:], 0)\n",
    "        opti.set_initial(U_c, 0)\n",
    "\n",
    "    sol, cpudt = chrono_solve(opti, solve_repetitions)\n",
    "\n",
    "    U_sol = sol.value(U)\n",
    "    X_sol = sol.value(X)\n",
    "    X_dot_sol = sol.value(X_dot)\n",
    "    T_sol = sol.value(T)\n",
    "    T_sol_arr = np.linspace(0, T_sol, N+1)\n",
    "    T_c_arr = (T_sol_arr[:-1]+T_sol_arr[1:])/2\n",
    "    cost_sol = sol.value(cost)\n",
    "    if 'hs' in scheme:\n",
    "        U_c_sol = sol.value(U_c)\n",
    "        X_c_sol = sol.value(X_c)\n",
    "        X_dot_c_sol = sol.value(X_dot_c)\n",
    "    else:\n",
    "        U_c_sol = None\n",
    "        X_c_sol = None\n",
    "        X_dot_c_sol = None\n",
    "        \n",
    "    return{\n",
    "        'u':U_sol,\n",
    "        'x':X_sol,\n",
    "        'x_dot':X_dot_sol,\n",
    "        't':T_sol,\n",
    "        't_array':T_sol_arr,\n",
    "        't_c_array': T_c_arr,\n",
    "        'cpudt':cpudt,\n",
    "        'u_c':U_c_sol,\n",
    "        'x_c':X_c_sol,\n",
    "        'x_dot_c':X_dot_c_sol,\n",
    "        'cost':cost_sol\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "422b213d",
   "metadata": {},
   "outputs": [],
   "source": [
    "schemes = ['hs_parab', 'hs_mod_parab','trapz', 'trapz_mod']\n",
    "solve_repetitions = 3\n",
    "N_arr = [20, 25, 30, 40, 50, 60]\n",
    "results = {}\n",
    "\n",
    "for scheme in schemes:\n",
    "    key = scheme\n",
    "    print('Problem:', key)\n",
    "    results[key] = {'N_arr':N_arr}\n",
    "    for N in N_arr:\n",
    "        print(f'\\tN = {N}')\n",
    "        results[key][N] = casadi_biped(\n",
    "            N = N,\n",
    "            scheme = scheme,\n",
    "            solve_repetitions = solve_repetitions,\n",
    "            t_end = 0.7,\n",
    "            step_length = 0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d9c453b",
   "metadata": {},
   "source": [
    "### Calculating dynamic errors for each case\n",
    "\n",
    "Caution! May take several seconds to run!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96962668",
   "metadata": {},
   "outputs": [],
   "source": [
    "schemes = ['hs_parab', 'hs_mod_parab','trapz', 'trapz_mod']\n",
    "\n",
    "n_graph = 2000 # A higher number here will provide more exact results but take longer to run\n",
    "t_arr = np.linspace(0,0.7,n_graph)\n",
    "for scheme in schemes:\n",
    "    key = scheme\n",
    "    if 'parab' in scheme:\n",
    "        u_scheme = 'parab'\n",
    "    else:\n",
    "        u_scheme = 'lin'\n",
    "    print('Problem:', key)\n",
    "    N_arr = results[key]['N_arr']\n",
    "    for N in N_arr:\n",
    "        print(f'\\tN = {N}')\n",
    "        dynamic_errors = dynamic_error_implicit(\n",
    "            x_arr=results[key][N]['x'],\n",
    "            u_arr=results[key][N]['u'],\n",
    "            t_end=results[key][N]['t'],\n",
    "            params = params,\n",
    "            F_impl = F_impl_np,\n",
    "            M = mass_matrix_np,\n",
    "            scheme = scheme, \n",
    "            u_scheme = u_scheme,\n",
    "            scheme_params={'u_c':results[key][N]['u_c'],\n",
    "                          'x_dot_c': results[key][N]['x_dot_c'],\n",
    "                          'x_c': results[key][N]['x_c']},\n",
    "            n_interp= n_graph)\n",
    "        \n",
    "        dyn_err_q = dynamic_errors[0][0]\n",
    "        dyn_err_v = dynamic_errors[0][1]\n",
    "        dyn_err_2 = dynamic_errors[1][0]\n",
    "            \n",
    "        results[key][N]['dyn_err_q'] = dyn_err_q\n",
    "        results[key][N]['dyn_err_v'] = dyn_err_v\n",
    "        results[key][N]['dyn_err_2'] = dyn_err_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0384aee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot settings\n",
    "plt.rcParams.update({'font.size': 15})\n",
    "oct_fig_size = [10,6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be19a701",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "schemes = ['hs_parab','hs_mod_parab', 'trapz', 'trapz_mod']\n",
    "titles = ['Hermite Simpson','2nd order Hermite Simpson', 'Trapezoidal', '2nd order Trapezoidal']\n",
    "colors = ['b', 'orange', 'g', 'r', 'purple']\n",
    "n_int = len(t_arr)\n",
    "N = 25\n",
    "interv_n = (N * t_arr)/results[scheme][N]['t']\n",
    "for kk in range(len(schemes)):\n",
    "    scheme = schemes[kk]\n",
    "    plt.figure(figsize=[14,8])\n",
    "    for ii in range(5):\n",
    "        cut_p = 0\n",
    "        for ll in range(1,N+1):\n",
    "            jj = np.searchsorted(interv_n, ll)\n",
    "            plt.plot(t_arr[cut_p:jj],results[scheme][N]['dyn_err_q'][cut_p:jj,ii], '-', c = colors[ii], label = f'$q_{ii+1}$' if cut_p == 0 else None)\n",
    "            cut_p = jj\n",
    "    plt.plot(np.linspace(0,results[scheme][N]['t'],N+1), np.zeros(N+1), 'ok')\n",
    "    plt.legend()\n",
    "    plt.grid()\n",
    "    if kk == 1:\n",
    "        plt.ylim([-0.00001, 0.00001])\n",
    "    elif kk == 3:\n",
    "        plt.ylim([-0.001, 0.001])\n",
    "    plt.title(r'First order dynamic error $\\varepsilon^{[1]}_{q_i}$,'+f' {titles[kk]} scheme')\n",
    "    plt.xlabel('Time(s)')\n",
    "    plt.ylabel('Dynamic error $(rad/s)$')\n",
    "    plt.tight_layout(pad = 0.0)\n",
    "    sch_type = titles[kk].replace(' ','_')\n",
    "    \n",
    "    # If you are running the notebook locally and want to save the plots,\n",
    "    # uncomment the next line\n",
    "    #plt.savefig(f'5_link_First_Order_Dynamic_Error_{sch_type}_scheme.eps', format='eps')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d90c9d5",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "schemes = ['hs_parab','hs_mod_parab', 'trapz', 'trapz_mod']\n",
    "titles = ['Hermite Simpson','2nd order Hermite Simpson', 'Trapezoidal', '2nd order Trapezoidal']\n",
    "colors = ['b', 'orange', 'g', 'r', 'purple']\n",
    "n_int = len(t_arr)\n",
    "N = 25\n",
    "interv_n = (N * t_arr)/results[scheme][N]['t']\n",
    "for kk in range(len(schemes)):\n",
    "    scheme = schemes[kk]\n",
    "    plt.figure(figsize=[14,8])\n",
    "    for ii in range(5):\n",
    "        cut_p = 0\n",
    "        for ll in range(1,N+1):\n",
    "            jj = np.searchsorted(interv_n, ll)\n",
    "            plt.plot(t_arr[cut_p:jj],results[scheme][N]['dyn_err_2'][cut_p:jj,ii], '-', c = colors[ii], label = f'$q_{ii+1}$' if cut_p == 0 else None)\n",
    "            cut_p = jj\n",
    "    plt.plot(results[scheme][N]['t_array'], np.zeros(N+1), 'ok', label = 'knot & collocation points')\n",
    "    if 'hs' in scheme:\n",
    "        plt.plot(results[scheme][N]['t_c_array'], np.zeros(N), 'ow', markeredgecolor='b', label = 'collocation points')\n",
    "        plt.ylim([-0.08, 0.08])\n",
    "    else:\n",
    "        plt.ylim([-1.75, 1.75])\n",
    "    plt.legend()\n",
    "    plt.grid()\n",
    "    plt.title(r'Second order dynamic error $\\varepsilon^{{[2]}}_{{q_i}}$,'+f' {titles[kk]} scheme')\n",
    "    plt.xlabel('Time(s)')\n",
    "    plt.ylabel('Dynamic error $(rad/s^2)$')\n",
    "    plt.tight_layout(pad = 0.0)\n",
    "    sch_type = titles[kk].replace(' ','_')\n",
    "    \n",
    "    # If you are running the notebook locally and want to save the plots,\n",
    "    # uncomment the next line\n",
    "    #plt.savefig(f'5_link_Second_Order_Dynamic_Error_{sch_type}_scheme.eps', format='eps')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a48d8285",
   "metadata": {},
   "outputs": [],
   "source": [
    "def arr_mod(x):\n",
    "    x_1 = np.sum(x*x, axis=1)\n",
    "    return np.sqrt(x_1)\n",
    "def arr_sum(x):\n",
    "    return np.sum(np.abs(x), axis = 1)\n",
    "def arr_max(x):\n",
    "    return np.max(np.abs(x), axis = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d0bc6dc",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "schemes = ['hs_mod_parab','hs_parab']#, 'trapz', 'trapz_mod']\n",
    "titles = ['2nd order Hermite Simpson','Hermite Simpson']#, 'Trapezoidal', 'Modified Trapezoidal']\n",
    "colors = ['b', 'orange', 'g', 'r', 'purple']\n",
    "funcs = [arr_sum,]#arr_mod,  arr_max\n",
    "#func_tittles = ['Module of', 'Sum of absolute', 'Maximum of absolute']\n",
    "y_max_list = [0.12, 0.2, 0.09]\n",
    "n_int = len(t_arr)\n",
    "N = 25\n",
    "interv_n = (N * t_arr)/results[scheme][N]['t']\n",
    "for ii in range(1):\n",
    "    plt.figure(figsize=oct_fig_size)\n",
    "    for kk in [1,0]:\n",
    "        scheme = schemes[kk]\n",
    "        cut_p = 0\n",
    "        for ll in range(1,N+1):\n",
    "            jj = np.searchsorted(interv_n, ll)\n",
    "            y_plot = funcs[ii](results[scheme][N]['dyn_err_2'])\n",
    "            plt.plot(t_arr[cut_p:jj],y_plot[cut_p:jj], '-', c = f'C{kk}', label = titles[kk] if cut_p == 0 else None)\n",
    "            cut_p = jj\n",
    "    plt.plot(results[scheme][N]['t_array'], np.zeros(N+1), 'ok', label = 'knot & collocation points')\n",
    "    plt.plot(results[scheme][N]['t_c_array'], np.zeros(N), 'ow', markeredgecolor='k', label = 'collocation points')\n",
    "    plt.legend()\n",
    "    plt.grid()\n",
    "    plt.ylim([-0.01,y_max_list[ii]])\n",
    "    plt.title(r'Second order dynamic error $\\varepsilon^{[2]}$,'+f' N = {N}')\n",
    "    plt.xlabel('Time(s)')\n",
    "    plt.ylabel('Dynamic error $(rad/s^2)$')\n",
    "    plt.tight_layout(pad = 0.0)\n",
    "    \n",
    "    # If you are running the notebook locally and want to save the plots,\n",
    "    # uncomment the next line\n",
    "    #plt.savefig(f'5_link_HS_N{N}_second_order_dynamic_error.eps', format='eps')\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cbf4ba0",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "schemes = ['trapz', 'trapz_mod']\n",
    "titles = ['Trapezoidal', '2nd order Trapezoidal']\n",
    "funcs = [arr_sum,]#arr_mod,  arr_max\n",
    "y_max_list = [0.12, 0.2, 0.09]\n",
    "n_int = len(t_arr)\n",
    "N = 50\n",
    "interv_n = (N * t_arr)/results[scheme][N]['t']\n",
    "for ii in range(1):\n",
    "    plt.figure(figsize=oct_fig_size)\n",
    "    for kk in range(2):\n",
    "        scheme = schemes[kk]\n",
    "        cut_p = 0\n",
    "        for ll in range(1,N+1):\n",
    "            jj = np.searchsorted(interv_n, ll)\n",
    "            y_plot = funcs[ii](results[scheme][N]['dyn_err_2'])\n",
    "            plt.plot(t_arr[cut_p:jj],y_plot[cut_p:jj], '-', c = f'C{kk+2}', label = titles[kk] if cut_p == 0 else None)\n",
    "            cut_p = jj\n",
    "    plt.plot(results[scheme][N]['t_array'], np.zeros(N+1), 'ok', label = 'knot & collocation points')\n",
    "    plt.legend()\n",
    "    plt.grid()\n",
    "    plt.title(r'Second order dynamic error $\\varepsilon^{[2]}$,'+f' N = {N}')\n",
    "    plt.xlabel('Time(s)')\n",
    "    plt.ylabel('Dynamic error $(rad/s^2)$')\n",
    "    plt.tight_layout(pad = 0.0)\n",
    "    \n",
    "    # If you are running the notebook locally and want to save the plots,\n",
    "    # uncomment the next line\n",
    "    #plt.savefig(f'5_link_Trapezoidal_N{N}_second_order_dynamic_error.eps', format='eps')\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80ef2fcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def total_state_error(t_arr, dyn_err):\n",
    "    errors = np.trapz(np.abs(dyn_err), t_arr, axis=0)\n",
    "    return errors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ed89ae9",
   "metadata": {},
   "outputs": [],
   "source": [
    "schemes = ['hs_parab', 'hs_mod_parab','trapz', 'trapz_mod']\n",
    "N_arr = [10,15,20,25,30,40,50,75,100,150]\n",
    "\n",
    "t_arr = np.linspace(0,0.7,n_graph)\n",
    "for scheme in schemes:\n",
    "    key = scheme\n",
    "    print('Problem:', key)\n",
    "    N_arr = results[key]['N_arr']\n",
    "    for N in N_arr:\n",
    "        print(f'\\tN = {N}')\n",
    "        for letter in 'qv2':\n",
    "            results[key][N][f'integ_dyn_err_{letter}']= total_state_error(t_arr, results[scheme][N][f'dyn_err_{letter}'])\n",
    "            results[key][N][f'module_dyn_err_{letter}']= np.sqrt(np.dot(results[key][N][f'integ_dyn_err_{letter}'], results[key][N][f'integ_dyn_err_{letter}']))\n",
    "            results[key][N][f'sum_dyn_err_{letter}'] = np.sum(results[key][N][f'integ_dyn_err_{letter}'])\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ae21888",
   "metadata": {},
   "outputs": [],
   "source": [
    "for scheme in schemes:\n",
    "    key = scheme\n",
    "    print('Problem:', key)\n",
    "    N_arr = results[key]['N_arr']\n",
    "    \n",
    "    for letter in 'qv2':\n",
    "        list_mod = []\n",
    "        list_sum = []\n",
    "        for N in N_arr:\n",
    "            #print(f'\\tN = {N}')\n",
    "            list_mod.append(results[key][N][f'module_dyn_err_{letter}'])\n",
    "            list_sum.append(results[key][N][f'sum_dyn_err_{letter}'])\n",
    "        results[key][f'module_dyn_err_{letter}_array'] = np.array(list_mod)\n",
    "        results[key][f'sum_dyn_err_{letter}_array'] = np.array(list_sum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3ab743e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For each scheme, the number of collocation points can be obtained\n",
    "for scheme in results.keys():\n",
    "    if 'hs' in scheme:\n",
    "        n_coll = np.array(results[scheme]['N_arr'])*2-1\n",
    "        results[scheme]['N_coll_arr'] = n_coll\n",
    "    else:\n",
    "        results[scheme]['N_coll_arr'] = results[scheme]['N_arr']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad11764a",
   "metadata": {},
   "outputs": [],
   "source": [
    "schemes = ['hs_mod_parab','hs_parab', 'trapz', 'trapz_mod']\n",
    "titles = ['2nd order Hermite Simpson', 'Hermite Simpson','Trapezoidal', '2nd order Trapezoidal']\n",
    "plt.figure(figsize=oct_fig_size)\n",
    "for ii in [2,3,1,0]:\n",
    "    key = schemes[ii]\n",
    "    plt.plot(results[key]['N_arr'], results[key][f'sum_dyn_err_2_array'], marker = 'o', c = f'C{ii}',label = titles[ii])\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "plt.yscale('log')\n",
    "plt.title('Second order dynamic error $E^{[2]}$')\n",
    "plt.xlabel('Number of intervals')\n",
    "plt.ylabel('Dynamic error ($rad/s$)')\n",
    "plt.tight_layout(pad = 0.0)\n",
    "\n",
    "# If you are running the notebook locally and want to save the plots,\n",
    "# uncomment the next line\n",
    "#plt.savefig(f'5_link_Sum_second_order_dynamic_error_vs_interval_number.eps', format='eps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "503b9b89",
   "metadata": {},
   "outputs": [],
   "source": [
    "for scheme in schemes:\n",
    "    key = scheme\n",
    "    print('Problem:', key)\n",
    "    N_arr = results[key]['N_arr']\n",
    "    list_cpudt = []\n",
    "    for N in N_arr:\n",
    "        #print(f'\\tN = {N}')\n",
    "        list_cpudt.append(results[key][N]['cpudt'])\n",
    "    results[key]['cpudt_array'] = np.array(list_cpudt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fc2d3aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "schemes = ['hs_mod_parab','hs_parab', 'trapz', 'trapz_mod']\n",
    "titles = ['2nd order Hermite Simpson', 'Hermite Simpson','Trapezoidal', '2nd order Trapezoidal']\n",
    "plt.figure(figsize=oct_fig_size)\n",
    "for ii in [2,3,1,0]:\n",
    "    key = schemes[ii]\n",
    "    plt.plot(results[key]['N_arr'], results[key][f'cpudt_array'], marker = 'o', c = f'C{ii}', label = titles[ii])\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "#plt.yscale('log')\n",
    "plt.title('Optimization time')\n",
    "plt.xlabel('Number of intervals')\n",
    "plt.ylabel('Time (s)')\n",
    "plt.tight_layout(pad = 0.0)\n",
    "\n",
    "# If you are running the notebook locally and want to save the plots,\n",
    "# uncomment the next line\n",
    "#plt.savefig(f'5_link_optimization_vs_interval_number.eps', format='eps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12222ba2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we print the data shown in Table III of the paper\n",
    "\n",
    "for scheme in ['trapz', 'trapz_mod', 'hs_parab', 'hs_mod_parab']:\n",
    "    key = scheme\n",
    "    N = 25 if 'hs' in scheme else 50\n",
    "    print('scheme:', scheme, 'N:', N,'\\n\\ttime:', results[key][N][f'cpudt'],\n",
    "              '\\n\\tErr 1:', results[key][N]['sum_dyn_err_q'], '\\n\\tErr 2:', results[key][N]['sum_dyn_err_2'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ebe7ee9",
   "metadata": {},
   "source": [
    "## Animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24c3fa78",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import animation, rc\n",
    "import matplotlib.patches as patches\n",
    "from matplotlib.transforms import Affine2D\n",
    "from IPython.display import HTML\n",
    "import matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6916e60b",
   "metadata": {},
   "outputs": [],
   "source": [
    "matplotlib.rcParams['animation.embed_limit'] = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e2cda42",
   "metadata": {},
   "outputs": [],
   "source": [
    "def body_tray(X, params):\n",
    "    res = []\n",
    "    for ii in range(X.shape[0]):\n",
    "        res.append(list(chain_to_draw(X[ii,:], params)))\n",
    "    return np.array(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf7dde8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def loop_body_tray(X, params):\n",
    "    point_tray = body_tray(X, params)\n",
    "    point_tray_loop = np.append(\n",
    "        point_tray,\n",
    "        np.expand_dims(\n",
    "            np.array(list(chain_to_draw(X[0,[4,3,2,1,0,5,6,7,8,9]],params)))\n",
    "            ,0),\n",
    "        0)\n",
    "    return point_tray_loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0350d366",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mod_sum(iterable, start):\n",
    "    for element in iterable:\n",
    "        start += element\n",
    "    return start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcc65a06",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_anim(X, U, params, n_loops = 1):\n",
    "    [\n",
    "    I_0_n, I_1_n, I_2_n, I_3_n, I_4_n,\n",
    "    d_0_n, d_1_n, d_2_n, d_3_n, d_4_n,\n",
    "    g_n,\n",
    "    l_0_n, l_1_n, l_3_n,\n",
    "    m_0_n, m_1_n, m_2_n, m_3_n, m_4_n\n",
    "    ] = params\n",
    "    \n",
    "    N = X.shape[0]\n",
    "    fig, ax = plt.subplots()\n",
    "    draw_width = 14\n",
    "    draw_height = 14\n",
    "    \n",
    "    fig.set_dpi(72)\n",
    "    fig.set_size_inches([draw_width,draw_height])\n",
    "    ax.set_xlim(( -1, 1))\n",
    "    ax.set_ylim(( -0.2, 1.8))\n",
    "    \n",
    "    body, = ax.plot([], [], lw=4, ms = 12, marker = 'o')\n",
    "    trail, = ax.plot([], [], lw=1, color = 'k')\n",
    "    old_trail, = ax.plot([], [], lw=1, color = 'k')\n",
    "    next_trail, = ax.plot([], [], lw=1, color = 'k')\n",
    "    \n",
    "    \n",
    "    point_tray = body_tray(X, params)\n",
    "    point_tray_loop = loop_body_tray(X, params)\n",
    "    #sys_cm_point, = ax.plot([], [], 'go', ms=12)\n",
    "    #line_sys_cm, = ax.plot([], [], 'k:', lw=1)\n",
    "    \n",
    "    print_vars = [X[:,ii] for ii in range(5)]+[np.linspace(0, N-1, N, dtype=int)]\n",
    "    print_var_names = [f'q_{ii}' for ii in range(5)]+['step']\n",
    "    texts = []\n",
    "    ii = 0.8\n",
    "    for arr in print_vars:\n",
    "        texts.append(ax.text(-0.8, ii, \"\", fontsize = 12))\n",
    "        ii -= 0.2\n",
    "    \n",
    "    ax.grid()\n",
    "    \n",
    "    def init():\n",
    "        body.set_data([], [])\n",
    "        trail.set_data(point_tray_loop[0,0,-1], point_tray_loop[0,1,-1])\n",
    "        old_trail.set_data(point_tray_loop[:,0,-1]-0.5, point_tray_loop[:,1,-1])\n",
    "        #next_trail.set_data(point_tray_loop[:,0,-1]+0.5, point_tray_loop[:,1,-1])\n",
    "        #sys_cm_point.set_data([], [])\n",
    "        #line_sys_cm.set_data([], [])\n",
    "        return (body,)\n",
    "    \n",
    "    def animate(i):\n",
    "        margin_x = -0.25 + i * 0.5/N\n",
    "        trail.set_data(point_tray_loop[0:i+1,0,-1], point_tray_loop[0:i+1,1,-1])\n",
    "        #sys_cm_coords = sys_cm_np(X[i,:], params)\n",
    "        #sys_cm_point.set_data(sys_cm_coords)\n",
    "        #line_sys_cm.set_data([0, sys_cm_coords[0]], [0, sys_cm_coords[1]])\n",
    "        \n",
    "        ax.set_xlim(( -1+ margin_x, 1+ margin_x))\n",
    "        points_x, points_y = point_tray[i,:,:]\n",
    "        body.set_data(points_x, points_y) \n",
    "        \n",
    "        for ii in range(len(texts)):\n",
    "            text = texts[ii]\n",
    "            name = print_var_names[ii]\n",
    "            arr = print_vars[ii]\n",
    "            text.set_position((-0.9 + margin_x, 1.7 - 0.05*ii))\n",
    "            if name == 'step':\n",
    "                text.set_text(\"$step$ = \" + str(arr[i]))\n",
    "            else:\n",
    "                text.set_text(\"$\" + name + \"$ = %.3f\" % arr[i])\n",
    "        return (body,)\n",
    "    iterable_frames = mod_sum([[jj for jj in range(N)]for kk in range(n_loops)], start = [])\n",
    "    anim = animation.FuncAnimation(fig, animate, init_func=init,\n",
    "                               frames=iterable_frames, interval=20, \n",
    "                               blit=True)\n",
    "    return anim\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2460218",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "anim = create_anim(results['hs_mod_parab'][25]['x'][:-1,:],results['hs_mod_parab'][25]['u'], params, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b8bdaf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7548391e",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = r\"biped_animation.mp4\" \n",
    "writervideo = animation.FFMpegWriter(fps=25//0.7) \n",
    "# If you are running the notebook locally and want to save the animation,\n",
    "# uncomment the next line\n",
    "#anim.save(f, writer=writervideo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f84a647b",
   "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}