Mono-pod Hopper Radau.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# EOMs\n",
    "%reset\n",
    "#import libraries\n",
    "import sympy as sym\n",
    "import numpy as np\n",
    "import pickle as pickle\n",
    "\n",
    "sym.init_printing()\n",
    "from IPython.display import display #for pretty printing\n",
    "\n",
    "# Symbols --------------------------------------------------\n",
    "links = ['b','l1R','l1L','l2R','l2L']\n",
    "\n",
    "#Parameters\n",
    "m = {}\n",
    "L = {}\n",
    "I = {}\n",
    "d = {}\n",
    "g = sym.symbols('g')\n",
    "\n",
    "for l in links:\n",
    "    m.update({l: sym.symbols('m_{%s}'%l)})\n",
    "    L.update({l: sym.symbols('L_{%s}'%l)})\n",
    "    I.update({l: sym.symbols('I_{%s}'%l)})\n",
    "    d.update({l: sym.symbols('d_{%s}'%l)})\n",
    "\n",
    "#gen coordinates\n",
    "x,y,thb = sym.symbols(['x','y','\\\\theta_{b}'])\n",
    "dx,dy,dthb = sym.symbols(['\\\\dot{x}','\\\\dot{y}','\\\\dot{\\\\theta}_{b}'])\n",
    "ddx,ddy,ddthb = sym.symbols(['\\\\ddot{x}','\\\\ddot{y}','\\\\ddot{\\\\theta}_{b}'])\n",
    "\n",
    "th = {}\n",
    "dth = {}\n",
    "ddth = {}\n",
    "\n",
    "for l in links[1:]:\n",
    "    th.update({l: sym.symbols('\\\\theta_{%s}'%l)})\n",
    "    dth.update({l: sym.symbols('\\\\dot{\\\\theta}_{%s}'%l)})\n",
    "    ddth.update({l: sym.symbols('\\\\ddot{\\\\theta}_{%s}'%l)})\n",
    "\n",
    "q = sym.Matrix([y]+[th[l] for l in links[1:]])\n",
    "dq = sym.Matrix([dy]+[dth[l] for l in links[1:]])\n",
    "ddq = sym.Matrix([ddy]+[ddth[l] for l in links[1:]])\n",
    "\n",
    "# positions\n",
    "def get_R(th_in):    \n",
    "    Rz = sym.Matrix([[sym.cos(th_in),-sym.sin(th_in),0],\n",
    "                    [sym.sin(th_in),sym.cos(th_in),0],\n",
    "                    [0,0,1]])\n",
    "    \n",
    "    R = Rz\n",
    "    return R\n",
    "\n",
    "pt = {'b': sym.Matrix([[0],[y],[0]])}\n",
    "po = {'b': sym.Matrix([[0],[0],[0]])}\n",
    "\n",
    "for l in links[1:]:\n",
    "    po.update({l: sym.Matrix([[0],[0],[th[l]]])})\n",
    "\n",
    "sides = ['R','L']\n",
    "\n",
    "kneep = {}\n",
    "\n",
    "footp = {}\n",
    "footv = {}\n",
    "foota = {}\n",
    "\n",
    "for s in sides:\n",
    "    R1 = get_R(th['l1%s'%s])\n",
    "    pt.update({'l1%s'%s: pt['b'] + R1*sym.Matrix([[0],[-0.5*L['l1%s'%s]],[0]])})\n",
    "    kneep.update({s: pt['b'] + R1*sym.Matrix([[0],[-L['l1%s'%s]],[0]])})\n",
    "    \n",
    "    R2 = get_R(th['l2%s'%s])\n",
    "    pt.update({'l2%s'%s: kneep[s] + R2*sym.Matrix([[0],[-d['l2%s'%s]*L['l1%s'%s]],[0]])})\n",
    "    \n",
    "    footp.update({s: kneep[s] + R2*sym.Matrix([[0],[-L['l2%s'%s]],[0]])})\n",
    "    footv.update({s: footp[s].jacobian(q)*dq})\n",
    "    foota.update({s: footv[s].jacobian(q)*dq + footp[s].jacobian(q)*ddq})\n",
    "    \n",
    "p = {}\n",
    "v = {}\n",
    "J = {}\n",
    "for l in links:\n",
    "    p.update({l: sym.Matrix([pt[l],po[l]])})\n",
    "    v.update({l: p[l].jacobian(q)*dq})\n",
    "    J.update({l: v[l].jacobian(dq)})\n",
    "    \n",
    "# ENERGY ---------------------------------------------------------------------------------------\n",
    "T = sum([0.5*m[l]*(v[l][0]**2 + v[l][1]**2) for l in links]) + sum([0.5*I[l]*v[l][5]**2 for l in links])\n",
    "V = sum([m[l]*g*p[l][1] for l in links])\n",
    "\n",
    "# term 1\n",
    "Lg1 = sym.zeros(len(q),1)\n",
    "for i in range(len(q)):\n",
    "    dT_ddq = sym.Matrix([sym.diff(T,dq[i])]) # get partial of T in dq_i\n",
    "    Lg1[i] = dT_ddq.jacobian(q)*dq + dT_ddq.jacobian(dq)*ddq #...then get time derivative of that partial\n",
    "\n",
    "# term 3\n",
    "Lg3 = sym.Matrix([T]).jacobian(q).transpose() # partial of T in q\n",
    "\n",
    "# term 4\n",
    "Lg4 = sym.Matrix([V]).jacobian(q).transpose() # partial of U in q\n",
    "\n",
    "#GEN FORCES-----------\n",
    "# constraint forces\n",
    "dj = sym.symbols('d_j')\n",
    "FcxR, FcyR, FcxL, FcyL = sym.symbols(['F_{cxR}','F_{cyR}','F_{cxL}','F_{cyL}'])\n",
    "\n",
    "FcL = sym.Matrix([[FcxL],[FcyL],[0]])\n",
    "FcL_r = footp['L']\n",
    "FcL_Q = FcL_r.jacobian(q).transpose()*FcL\n",
    "\n",
    "FcR = sym.Matrix([[FcxR],[FcyR],[0]])\n",
    "FcR_r = kneep['R'] + get_R(th['l2R'])*sym.Matrix([[0],[-dj],[0]])\n",
    "FcR_Q = FcR_r.jacobian(q).transpose()*FcR\n",
    "\n",
    "jointp = {'L': FcL_r, 'R': FcR_r}\n",
    "jointv = {}\n",
    "jointa = {}\n",
    "for s in sides:\n",
    "    jointv.update({s: jointp[s].jacobian(q)*dq})\n",
    "    jointa.update({s: jointv[s].jacobian(q)*dq + jointv[s].jacobian(dq)*ddq})\n",
    "# ground reaction force\n",
    "Gx, Gy = sym.symbols(['G_x','G_y'])\n",
    "\n",
    "G = sym.Matrix([[Gx],[Gy],[0]])\n",
    "G_r = footp['R']\n",
    "G_Q = G_r.jacobian(q).transpose()*G\n",
    "\n",
    "# actuators\n",
    "tau = {}\n",
    "for s in sides:\n",
    "    tau.update({s: sym.symbols('\\\\tau_{%s}'%s)})\n",
    "\n",
    "tau_Q = sym.Matrix([[0],[tau['R']],[tau['L']],[0],[0]])\n",
    "\n",
    "Q = FcL_Q + FcR_Q + G_Q + tau_Q\n",
    "\n",
    "# AND combine!\n",
    "EOM = Lg1 - Lg3 + Lg4 - Q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LAMBDIFY\n",
    "from pyomo.environ import*\n",
    "from pyomo.opt import SolverFactory\n",
    "from pyomo.opt import SolverStatus, TerminationCondition\n",
    "\n",
    "func_map = {'sin':sin, 'cos':cos} \n",
    "\n",
    "sym_list = [g] + [m[l] for l in links] +\\\n",
    "           [L[l] for l in links] +\\\n",
    "           [I[l] for l in links] +\\\n",
    "           [d[l] for l in links] + [dj] +\\\n",
    "           [y] + [th[l] for l in links[1:]] +\\\n",
    "           [dy] + [dth[l] for l in links[1:]] +\\\n",
    "           [ddy] + [ddth[l] for l in links[1:]] +\\\n",
    "           [[tau['%sR'%p],tau['%sL'%p] for p in planes] +\\\n",
    "            Gx,Gy,FcxR, FcyR, FcxL, FcyL\n",
    "\n",
    "\n",
    "\n",
    "DOFs = ['y'] + ['th_%s'%l for l in links[1:]]\n",
    "\n",
    "for dof_i,dof in enumerate(DOFs):\n",
    "    lamb_EOM.update({dof: sym.lambdify(sym_list,EOM[dof_i],modules = [func_map])})\n",
    "    \n",
    "TDOFs = ['x','y']\n",
    "lamb_footp = {}\n",
    "lamb_footv = {}\n",
    "for dof_i,dof in enumerate(TDOFs):\n",
    "    lamb_footp.update({dof: sym.lambdify(sym_list,footp['R'][dof_i],modules = [func_map])})\n",
    "    lamb_footv.update({dof: sym.lambdify(sym_list,footv['R'][dof_i],modules = [func_map])})\n",
    "\n",
    "lamb_jointp = {}\n",
    "lamb_jointv = {}\n",
    "lamb_jointa = {}\n",
    "for s in sides:\n",
    "    for dof_i,dof in enumerate(TDOFs):\n",
    "        lamb_jointp.update({(s,dof): sym.lambdify(sym_list,jointp[s][dof_i],modules = [func_map])})\n",
    "        lamb_jointv.update({(s,dof): sym.lambdify(sym_list,jointv[s][dof_i],modules = [func_map])})\n",
    "        lamb_jointa.update({(s,dof): sym.lambdify(sym_list,jointa[s][dof_i],modules = [func_map])})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PARAMETERS:\n",
    "mass = {'b': 2.8 + 2.542,\n",
    "       'l1R': 0.079,\n",
    "       'l1L': 0.079,\n",
    "       'l2R': 0.101,\n",
    "       'l2L': 0.100}\n",
    "\n",
    "length = {'b': 0.250,\n",
    "         'l1R': 0.1745,\n",
    "         'l1L': 0.1745,\n",
    "         'l2R': 0.289,\n",
    "         'l2L': 0.250}\n",
    "\n",
    "inertia = {'b': 0,\n",
    "          'l1R': 157713/(1000*1000**2),\n",
    "          'l1L': 157713/(1000*1000**2),\n",
    "          'l2R': 927418/(1000*1000**2),\n",
    "          'l2L': 817690/(1000*1000**2)}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P = 2\n",
    "\n",
    "def get_var_list(m,n,p):\n",
    "    var_list = [m.g]+\\\n",
    "                [m.m[l] for l in links]+\\\n",
    "                [m.l[l] for l in links]+\\\n",
    "                [m.I[l] for l in links]+\\\n",
    "                [m.dCOM[l] for l in links]+[m.dj]+\\\n",
    "                [m.q[n,p,dof] for dof in DOFs]+\\\n",
    "                [m.dq[n,p,dof] for dof in DOFs]+\\\n",
    "                [m.ddq[n,p,dof] for dof in DOFs]+\\\n",
    "                [m.tau[n,s] for s in sides]+\\\n",
    "                [m.GRF[n,p,dof] for dof in TDOFs]+\\\n",
    "                [m.Fc[n,p,s,dof] for s in sides for dof in TDOFs]\n",
    "    return var_list\n",
    "\n",
    "def create_model(N,hm,mu):\n",
    "    m = ConcreteModel()\n",
    "\n",
    "    # SETS-----------------------------------------------------------------------------------------------------------------------\n",
    "    m.N = RangeSet(N) \n",
    "    m.P = Set(initialize = range(P+1))\n",
    "    m.L = Set(initialize = links)\n",
    "    m.DOF = Set(initialize = DOFs) \n",
    "    m.TDOF = Set(initialize = TDOFs) \n",
    "    m.S = Set(initialize = sides)\n",
    "    m.p = Set(intialize = planes)\n",
    "    \n",
    "    # PARAMETERS-----------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "    m.g = Param(initialize = 9.81)\n",
    "    m.m = Param(m.L, initialize = mass)\n",
    "    m.l = Param(m.L, initialize = length)\n",
    "    m.I = Param(m.L, initialize = inertia)\n",
    "    m.dCOM = Param(m.L, initialize = 0.5)\n",
    "    m.dj = Param(initialize = length['l2L_F'])\n",
    "    \n",
    "    # VARIABLES -----------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "    # generalized coordinates\n",
    "    m.q = Var(m.N, m.P, m.DOF) # position\n",
    "    m.dq = Var(m.N, m.P, m.DOF) # velocity\n",
    "    m.ddq = Var(m.N, m.P, m.DOF) # acceleration\n",
    "        \n",
    "    for n in range(1,N+1):\n",
    "        for p in range(P+1):\n",
    "            m.q[n,p,'y'].setlb(0)\n",
    "            m.q[n,p,'th_l1R'].setlb(0)\n",
    "            m.q[n,p,'th_l1R'].setub(np.pi)\n",
    "            m.q[n,p,'th_l1L'].setlb(-np.pi)\n",
    "            m.q[n,p,'th_l1L'].setub(0)\n",
    "    # forces\n",
    "    m.tau = Var(m.N, m.S)\n",
    "    m.Fc = Var(m.N, m.P, m.S, m.TDOF)\n",
    "    m.GRF = Var(m.N, m.P, m.TDOF)\n",
    "    \n",
    "    for n in range(1,N+1):\n",
    "        for p in range(P+1):\n",
    "            m.GRF[n,p,'y'].setlb(0)\n",
    "\n",
    "    #TIME AND INTEGRATION ---------------------------------------------------------------------------------------------------\n",
    "\n",
    "    # timestep\n",
    "    m.hm = Param(initialize = hm, mutable = True) # master timestep\n",
    "    m.h = Var(m.N, bounds = (0.8,1.2))\n",
    "\n",
    "    # Integration constraints \n",
    "    radau_mat = [[0.416666666666667, -0.0833333333333333],\n",
    "                 [0.75, 0.25]]\n",
    "\n",
    "    def integrate_p(m,n,p,dof):\n",
    "        if p == 0:\n",
    "            if n == 1: return Constraint.Skip\n",
    "            else:\n",
    "                return m.q[n,0,dof] == m.q[n-1,P,dof]\n",
    "        else:\n",
    "            vel = sum([radau_mat[p-1][pp-1]*m.dq[n,pp,dof] for pp in range(1,P+1)])\n",
    "            return m.q[n,p,dof] == m.q[n,0,dof] + m.h[n]*m.hm*vel\n",
    "    m.integrate_p = Constraint(m.N, m.P, m.DOF, rule = integrate_p)\n",
    "\n",
    "    def integrate_v(m,n,p,dof):\n",
    "        if p == 0:\n",
    "            if n == 1: return Constraint.Skip\n",
    "            else:\n",
    "                return m.dq[n,0,dof] == m.dq[n-1,P,dof]\n",
    "        else:\n",
    "            acc = sum([radau_mat[p-1][pp-1]*m.ddq[n,pp,dof] for pp in range(1,P+1)])\n",
    "            return m.dq[n,p,dof] == m.dq[n,0,dof] + m.h[n]*m.hm*acc\n",
    "    m.integrate_v = Constraint(m.N, m.P, m.DOF, rule = integrate_v)\n",
    "    \n",
    "    # dynamics --------------------------------------------------------------------------------------    \n",
    "    \n",
    "    def dynamics(m,n,p,dof):\n",
    "        if p == 0: return Constraint.Skip\n",
    "        var_list = get_var_list(m,n,p)\n",
    "        return lamb_EOM[dof](*var_list) == 0\n",
    "    m.dynamics = Constraint(m.N, m.P, m.DOF, rule = dynamics)\n",
    "    \n",
    "    # motor model --------------------------------------------------------------------------------------\n",
    "    signs = ['ps','ng']\n",
    "    m.sgn = Set(initialize = signs)\n",
    "\n",
    "    max_rpm = 80*48/2\n",
    "    max_dth = max_rpm*2*np.pi/60\n",
    "    max_tau = 8.92*2\n",
    "    \n",
    "    def speed_bounds(omega):\n",
    "            ub =(-(max_tau)/(max_dth))*omega + max_tau\n",
    "            lb = (-(max_tau)/(max_dth))*omega - max_tau\n",
    "            return(lb,ub)\n",
    "            \n",
    "    def motor_torque_speed_curve(m,n,s,sgn):\n",
    "        lb,ub = speed_bounds(m.dq[n,'th_l1R'])\n",
    "        if sgn == 'ps':\n",
    "            for s in sides:\n",
    "                return m.tau[n,'%s'%s] <= ub\n",
    "        else:\n",
    "            for s in sides:\n",
    "                return m.tau[n,'%s'%s] >= lb\n",
    "        return Constraint.Skip\n",
    "    m.motor_torque_speed_curve = Constraint(m.N, m.S, m.sgn, rule = motor_torque_speed_curve)\n",
    "        \n",
    "#     for n in range(1,N+1):\n",
    "#         for s in sides:\n",
    "#             m.tau[n,s].setlb(-3*6.5)\n",
    "#             m.tau[n,s].setub(3*6.5)\n",
    "#             m.dq[n,p,'th_l1%s'%s].setlb(-max_dth/3)\n",
    "#             m.dq[n,p,'th_l1%s'%s].setub(max_dth/3)\n",
    "# #             m.power[n,p,s].setlb(-1200)\n",
    "# #             m.power[n,p,s].setub(1200)\n",
    "    \n",
    "    # connection -------------------------------------------------------------------------------------\n",
    "    m.jointp = Var(m.N, m.S, m.TDOF)\n",
    "    def get_jointp(m,n,s,dof):\n",
    "        var_list = get_var_list(m,n,0)\n",
    "        return m.jointp[n,s,dof] == lamb_jointp[s,dof](*var_list)\n",
    "    m.get_jointp = Constraint(m.N, m.S, m.TDOF, rule = get_jointp)\n",
    "    \n",
    "#     m.jointa = Var(m.N, m.P, m.S, m.TDOF)\n",
    "#     def get_jointa(m,n,p,s,dof):\n",
    "#         if p == 0: return Constraint.Skip\n",
    "#         var_list = get_var_list(m,n,p)\n",
    "#         return m.jointa[n,p,s,dof] == lamb_jointa[s,dof](*var_list)\n",
    "#     m.get_jointa = Constraint(m.N, m.P, m.S, m.TDOF, rule = get_jointa)\n",
    "    \n",
    "    def constraint_forces(m,n,p,dof):\n",
    "        if p == 0: return Constraint.Skip\n",
    "        return m.Fc[n,p,'L',dof] == -m.Fc[n,p,'R',dof]\n",
    "    m.constraint_forces = Constraint(m.N, m.P, m.TDOF, rule = constraint_forces)\n",
    "    \n",
    "    def connect_p(m,n,dof):\n",
    "#         if n > 1: return Constraint.Skip\n",
    "        return m.jointp[n,'R',dof] == m.jointp[n,'L',dof]\n",
    "    m.connect_p = Constraint(m.N, m.TDOF, rule = connect_p)\n",
    "    \n",
    "#     def connect_a(m,n,p,dof):\n",
    "#         if p == 0: return Constraint.Skip\n",
    "#         return m.jointa[n,p,'R',dof] == m.jointa[n,p,'L',dof]\n",
    "#     m.connect_a = Constraint(m.N, m.P, m.TDOF, rule = connect_a)\n",
    "    \n",
    "    # GROUND ----------------------------------------------------------------------------------------------------------\n",
    "    \n",
    "    # impact\n",
    "    m.footp = Var(m.N, m.P, m.TDOF, bounds = (0,None))\n",
    "    def get_footp(m,n,p,dof):\n",
    "        if dof == 'x': return Constraint.Skip\n",
    "        if p == 0 and n > 1:\n",
    "            return m.footp[n,p,dof] == m.footp[n-1,P,dof]\n",
    "        else:\n",
    "            var_list = get_var_list(m,n,p)\n",
    "            return m.footp[n,p,dof] == lamb_footp[dof](*var_list)\n",
    "    m.get_footp = Constraint(m.N, m.P, m.TDOF, rule = get_footp)\n",
    "    \n",
    "    m.footv = Var(m.N, m.P, m.TDOF)\n",
    "    def get_footv(m,n,p,dof):\n",
    "        if dof == 'y': return Constraint.Skip\n",
    "        if p == 0: return Constraint.Skip\n",
    "        else:\n",
    "            var_list = get_var_list(m,n,p)\n",
    "            return m.footv[n,p,dof] == lamb_footv[dof](*var_list)\n",
    "    m.get_footv = Constraint(m.N, m.P, m.TDOF, rule = get_footv)\n",
    "    \n",
    "    m.contact_penalty = Var(m.N, m.P, bounds = (0,None))\n",
    "    def contact_complementarity(m,n,p):\n",
    "        if p == 0: return Constraint.Skip\n",
    "        if n < N:\n",
    "#             A = sum([m.footy[n,pp] for pp in range(p,P+1)]) + sum([m.footy[n+1,pp] for pp in range(P+1)])\n",
    "            A = sum([m.footp[n+1,pp,'y'] for pp in range(P+1)])\n",
    "        else:\n",
    "            A = sum([m.footp[n,pp,'y'] for pp in range(p,P+1)])\n",
    "        B = m.GRF[n,p,'y']\n",
    "        return A*B <= m.contact_penalty[n,p]\n",
    "    m.contact_complementarity = Constraint(m.N, m.P, rule = contact_complementarity)\n",
    "    \n",
    "    # friction\n",
    "    m.mu = Param(initialize = mu)\n",
    "    \n",
    "    GDOFs = ['x']\n",
    "    m.GDOF = Set(initialize = GDOFs)\n",
    "        \n",
    "    dres = 2\n",
    "    directions = RangeSet(dres)\n",
    "    m.D = Set(initialize = directions) # directions\n",
    "    def get_d(m,dof,d):\n",
    "        th = np.linspace(0,2*np.pi,dres+1)\n",
    "        if dof == 'x':\n",
    "            return np.cos(th[d-1])\n",
    "        if dof == 'z':\n",
    "            return np.sin(th[d-1])\n",
    "    m.d = Param(m.GDOF, m.D, initialize = get_d) # unit vectors\n",
    "    \n",
    "    m.alpha = Var(m.N, m.P, m.D, bounds = (0,1)) # activation in each direction\n",
    "    m.F = Var(m.N, m.P, bounds = (0,None)) # magnitude of friction force\n",
    "    m.gamma = Var(m.N, m.P, bounds = (0,None))\n",
    "    m.dir = Var(m.N, m.P, m.D, bounds = (0,None))\n",
    "\n",
    "    def def_gamma(m,n,p,d):\n",
    "        if p == 0: return Constraint.Skip\n",
    "        dx = m.footv[n,p,'x']\n",
    "        return m.gamma[n,p] + sum([m.d[dof,d]*dx for dof in GDOFs]) == m.dir[n,p,d]\n",
    "    m.def_gamma = Constraint(m.N, m.P, m.D, rule = def_gamma)\n",
    "    \n",
    "    def friction(m,n,p): # insert velocity-dependent friction law here\n",
    "        if p == 0: return Constraint.Skip\n",
    "        dx = m.footv[n,p,'x']\n",
    "        fm = m.mu\n",
    "        return m.F[n,p] == fm*m.GRF[n,p,'y']\n",
    "    m.friction = Constraint(m.N, m.P, rule = friction)\n",
    "\n",
    "    def get_GRF(m,n,p,dof):\n",
    "        if p == 0: return Constraint.Skip\n",
    "        return m.GRF[n,p,dof] == m.F[n,p]*sum([m.alpha[n,p,d]*m.d[dof,d] for d in directions])\n",
    "    m.get_GRF = Constraint(m.N, m.P, m.GDOF, rule = get_GRF)\n",
    "\n",
    "    # complementarity\n",
    "    m.direction_penalty = Var(m.N, m.P, m.D, bounds = (0.0,None))\n",
    "    m.magnitude_penalty = Var(m.N, m.P, bounds = (0,None))\n",
    "\n",
    "    def direction_complementarity(m,n,p,d):\n",
    "        if p == 0: return Constraint.Skip\n",
    "        A = m.alpha[n,p,d]\n",
    "        B = m.dir[n,p,d]\n",
    "        return  A*B <= m.direction_penalty[n,p,d]\n",
    "    m.direction_complementarity = Constraint(m.N, m.P, m.D, rule = direction_complementarity)\n",
    "\n",
    "    def magnitude_complementarity(m,n,p):\n",
    "        if p == 0: return Constraint.Skip\n",
    "        A = m.gamma[n,p]\n",
    "        B = 1-sum([m.alpha[n,p,d] for d in directions])\n",
    "        return A*B <= m.magnitude_penalty[n,p]\n",
    "    m.magnitude_complementarity = Constraint(m.N, m.P, rule = magnitude_complementarity)\n",
    "    \n",
    "    def MinPenalty(m):\n",
    "        return sum([m.contact_penalty[n,p] + m.magnitude_penalty[n,p] for n in range(1,N+1) for p in range(P+1)]) +\\\n",
    "                sum([m.direction_penalty[n,p,d] for n in range(1,N+1) for p in range(P+1) for d in range(1,dres+1)])\n",
    "    m.MinPenalty = Objective(rule = MinPenalty)\n",
    "    m.MinPenalty.deactivate()\n",
    "    def ss_p(m,dof):\n",
    "#         if dof != 'y': return Constraint.Skip\n",
    "        return m.q[1,0,dof] == m.q[N,0,dof]\n",
    "    m.ss_p = Constraint(m.DOF, rule = ss_p)\n",
    "    \n",
    "    def ss_v(m,dof):\n",
    "        return m.dq[1,0,dof] == m.dq[N,0,dof]\n",
    "    m.ss_v = Constraint(m.DOF, rule = ss_v)\n",
    "    \n",
    "    def MaxHeight(m):\n",
    "        return -m.q[1,0,'y']\n",
    "    m.MaxHeight = Objective(rule = MaxHeight)\n",
    "    #m.MaxHeight.deactivate()\n",
    "    \n",
    "    return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_opt(): # initializes solver\n",
    "    opt = SolverFactory('ipopt',executable = '/home/immanuel/Solver/build/bin/ipopt')\n",
    "    opt.options[\"linear_solver\"] = 'ma86'\n",
    "    opt.options[\"print_level\"] = 5 # prints a log with each iteration (you want to this - it's the only way to see progress.)\n",
    "    opt.options[\"max_iter\"] = 30000 # maximum number of iterations\n",
    "    opt.options[\"max_cpu_time\"] = 3600 # maximum cpu time in seconds\n",
    "    opt.options[\"Tol\"] = 1e-6 # the tolerance for feasibility. Considers constraints satisfied when they're within this margin.\n",
    "    \n",
    "    opt.options[\"OF_acceptable_obj_change_tol\"] = 1e-4\n",
    "    opt.options[\"OF_ma86_scaling\"] = 'none'\n",
    "    \n",
    "    return opt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "N = 100+1\n",
    "hm = 0.01\n",
    "mu = 1.0\n",
    "m = create_model(N,hm,mu)\n",
    "# # HIGH DROP -----------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "m.footp[1,0,'y'].fix(0)\n",
    "m.q[1,0,'y'].setlb(0.2)\n",
    "#m.q[1,'x'].fix(0)\n",
    "#m.q[1,'th_b'].fix(0)\n",
    "#m.footp[1,'y'].fix(0)\n",
    "#rest pose\n",
    "m.dq[1,0,:].fix(0)\n",
    "m.h[1].fix(1)\n",
    "\n",
    "#final\n",
    "#m.q[N,'x'].setlb(1)\n",
    "#m.q[N,'y'].setlb(0.9)\n",
    "#m.footp[N,'x'].setlb(0)\n",
    "#m.footp[N,'y'].fix(0)\n",
    "#m.q[N,'th_b'].fix(0)\n",
    "#m.dq[N,:].fix(0)\n",
    "#set initial pose\n",
    "\n",
    "\n",
    "m.footp[0.5*(N-1),0,'y'].setlb(0.3)\n",
    "#m.q[0.5*(N-1),'y'].setlb(0.6)\n",
    "#m.footp[0.5*(N-1),'x'].fix(0)\n",
    "# m.q[0.5*]\n",
    "for n in range(1,N+1):\n",
    "    m.footp[n,,:,'x'].fix(0)\n",
    "# m.footp[0.5*(N-1),'x'].setub(0.01)\n",
    "\n",
    "opt = init_opt()\n",
    "\n",
    "results = opt.solve(m,tee = True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "m.q[1,0,'y'].unfix()\n",
    "m.MinPenalty.deactivate()\n",
    "m.MaxHeight.activate()\n",
    "\n",
    "thr = 1e-4\n",
    "for n in range(1,N+1):\n",
    "    for p in range(P+1):\n",
    "        m.contact_penalty[n,p].setub(thr)\n",
    "        m.magnitude_penalty[n,p].setub(thr)\n",
    "        for d in range(1,3):\n",
    "            m.direction_penalty[n,p,d].setub(thr)\n",
    "\n",
    "results = opt.solve(m, tee = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#animate it\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as ani\n",
    "from IPython.display import HTML\n",
    "import numpy as np\n",
    "from matplotlib.animation import FuncAnimation, FFMpegWriter\n",
    "%matplotlib inline\n",
    "\n",
    "fig1, ax1 = plt.subplots(1,1) #create axes\n",
    "ax1.set_aspect('equal')\n",
    "\n",
    "ymax = np.max([m.q[n,p,'y'].value for n in range(1,N+1) for p in range(P+1)])\n",
    "\n",
    "def plot_model(n,m,ax): #update function for animation\n",
    "    ax.clear()\n",
    "    ax.set_xlim([-1, 1])\n",
    "    ax.set_ylim([-0.2, ymax + 0.2])\n",
    "    ax.plot([-10,10],[0,0],linewidth=3,color='black')\n",
    "    p = 0\n",
    "    \n",
    "    y = m.q[n,p,'y'].value\n",
    "    ax.scatter([0],[y],s=50,color='black',zorder=50)\n",
    "    \n",
    "    cols = ['xkcd:bright pink','xkcd:bright blue']\n",
    "    for si, s in enumerate(sides):\n",
    "        l1 = m.l['l1%s'%s]\n",
    "        th1 = m.q[n,p,'th_l1%s'%s].value\n",
    "        knee_x = l1*np.sin(th1)\n",
    "        knee_y = y-l1*np.cos(th1)\n",
    "        ax.plot([0,knee_x],[y,knee_y],color=cols[si])\n",
    "        \n",
    "        l2 = m.l['l2%s'%s]\n",
    "        th2 = m.q[n,p,'th_l2%s'%s].value\n",
    "        foot_x = knee_x + l2*np.sin(th2)\n",
    "        foot_y = knee_y - l2*np.cos(th2)\n",
    "        ax.plot([knee_x,foot_x],[knee_y,foot_y],color=cols[si])\n",
    "        \n",
    "       \n",
    "update = lambda n: plot_model(n,m,ax1) #lambdify update function\n",
    "\n",
    "animate = ani.FuncAnimation(fig1,update,range(1,N+1),interval = 100,repeat=True)\n",
    "writer = FFMpegWriter(fps=1.0/0.1, metadata=dict(artist='Me'))\n",
    "animate.save('hop_animation_higher_order2.mp4', writer=writer)\n",
    "\n",
    "HTML(animate.to_html5_video()) #you need to convert the animation to HTML5 to embed it in the notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotfig2, \n",
    "fig2, ax2 = plt.subplots(1,1)\n",
    "time = [0]\n",
    "for n in range(1,N+1): time.append(time[n-1] + hm*m.h[n].value)\n",
    "time.pop()\n",
    "dt = [m.h[n].value for n in range(1,N+1)]\n",
    "m.h[1].value = 1.0 # left out of integrationtime_sum = hm*sum(model.h[n].value for n in nodes)\n",
    "#print(f'Duration: {time_sum}s')\n",
    "#ax2.plot(time, [m.GRF[n,1,'y'].value for n in m.N])\n",
    "ax2.plot(time, [m.q[n,1,'th_l1L'].value for n in m.N])\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "e554e1ece3fdd66bc6d1dd782a6f7ee29b41f45152844efbf0ce8f22e8a6c640"
  },
  "kernelspec": {
   "display_name": "Python 3.8.8 64-bit ('base': conda)",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}