gistfile1.py

from casadi import *

##### lets integrate something 100 timesteps with 3 different methods #####
# this is whate we want to integrate
def dynamics_fun(x):
      xdot = # some complicated shit
      return xdot
    
## 1: inline everything, standard approach to AD afaik
def rk4(x0, h):
    k1 = dynamics_fun(x0)
    k2 = dynamics_fun(x0 + h/2*k1)
    k3 = dynamics_fun(x0 + h/2*k2)
    k4 = dynamics_fun(x0 + h/2*k3)
    return x0 + h/6*(k1 + 2*k2 + 2*k3 + k4)

x = SX.sym('x',4,1);
h = SX.sym('h');
xf = x
for k in range(100):
    xf = rk4(xf,h)

final_fun = SXFunction([x,h],[xf])


## 2: make the single step integrator its own "atomic" operation
def rk4(x0,h):
    k1 = dynamics_fun(x0)
    k2 = dynamics_fun(x0 + h/2*k1)
    k3 = dynamics_fun(x0 + h/2*k2)
    k4 = dynamics_fun(x0 + h/2*k3)
    return x0 + h/6*(k1 + 2*k2 + 2*k3 + k4)
x = SX.sym('x',4,1)
h = SX.sym('h')
y = rk4(x,h)
integrator = SXFunction([x,h],[y]) # an atomic integrator

z = MX.sym('z',4,1)
zh = MX.sym('zh')
xf = z
for k in range(100):
    xf = integrator.call([xf,zh]) # each timestep is just a call to this integrator node

final_fun = MXFunction([z,zh],[xf])



## 3: make integrator fun and dynamics fun atomic
x = SX.sym('x',4,1)
xdot = dynamics_fun(x)
dfun = SXFunction([x],[xdot]) # one atomic dynamics evaluation

def rk4(x0,h):
    k1 = dfun.call([x0])
    k2 = dfun.call([x0 + h/2*k1])
    k3 = dfun.call([x0 + h/2*k2])
    k4 = dfun.call([x0 + h/2*k3])
    return x0 + h/6*(k1 + 2*k2 + 2*k3 + k4)
x = MX.sym('x',4,1)
h = MX.sym('h')
integrator_fun = MXFunction([x,h],[rk4(x,h)]) # one atomic integration step which itnernally makes atomic dynamics calls 

z = MX.sym('x',4,1);
zh = MX.sym('h');
xf = z
for k in range(100):
    xf = integrator_fun.call([xf,zh])

final_fun = MXFunction([z,zh],[xf])