pseudo MX optimizations

[jgillis]

pseudo MX optimizations

[jgillis]

Proof of concept for option 4:

from casadi import *

h = MX.sym("h")
v = MX.sym("v")
m = MX.sym("m")


f = Function('f',[h],[h**2])
cg = CodeGenerator('f')
cg.add(f)
cg.add(f.forward(1))
cg.generate()

f = external('f',Importer('f.c','shell'))


g = Function('g',[h],[cos(h)])
cg = CodeGenerator('g')
cg.add(g)
cg.add(g.forward(1))
cg.generate()

g = external('g',Importer('g.c','shell'))

a = MX.sym("a")
b = MX.sym("b")

u = 9

p = [a,b]

y = [v, (u-a*v**2)/m-9.81, -b*u**2+g(f(h)*m+sqrt(v))]

x = [h,v,m]


def my_jacobian(y,x,p):


	[y_lifted,w,wdef] = extract(y,{"lift_shared":False,"lift_calls":True})

	sp = jacobian_sparsity(vcat(wdef),vcat(w))
	nnz = [jacobian(wdef[r],w[c]) for r,c in zip(sp.row(),sp.get_col())]

	J = Function("J",[vcat(x),vcat(w),vcat(p)],[vcat(nnz)],{"always_inline":True})
	
	jac = MX.sym("jac",sp)
	
	jacnnz = MX.sym("jacnnz",len(nnz))
	jac = sparsity_cast(jacnnz,sp)
	
	f = Function('f',[vcat(x),vcat(w),vcat(p),jacnnz],[jacobian(vcat(y_lifted),vcat(w)) @ solve(DM.eye(jac.shape[0])- jac,jacobian(vcat(wdef),vcat(x))) +jacobian(vcat(y_lifted),vcat(x))])
	f = f.expand('f',{"cse":True,"always_inline":True})
	
	[w_expr,_] = substitute_inplace(w,wdef,y_lifted)
	
	return f(vcat(x),vcat(w_expr),vcat(p),J(vcat(x),vcat(w_expr),vcat(p)))


ref = Function('ref',[vcat(x),vcat(p)],[jacobian(vcat(y),vcat(x))])
f = Function('f',[vcat(x),vcat(p)],[my_jacobian(y,x,p)])

DM.rng(1)

x_ = DM.rand(3)
p_ = DM.rand(2)

print(ref(x_,p_))
print(f(x_,p_))