pseudo MX optimizations

[jgillis]

In some contexts, MX is used as SX (everything is scalar) with a few shortcomings lifted (being able to embed Functions).

For these situations, there are important optimizations missing.

Consider

x = X.sym("x")
y = X.sym("y")

r = sqrt(x+y)

X = vertcat(x,y)
E = vertcat(sqrt(x+y),sin(x))

J = jacobian(E,X)

For X=SX, we have

@1=sqrt((x+y)), @2=(1./(@1+@1)), 
[[@2, @2], 
 [cos(x), 00]]

For Y=MX, we have J:

@1=vertsplit(ones(2x1,1nz)){0}, @2=(2.*sqrt((x+y))), ((zeros(2x2,3nz)[:4:2] = vertcat((@1/@2), (cos(x)*@1)))[1] = (vertsplit(ones(2x1,1nz)){1}/@2))'

This may conceivable be represented as, J_terse

@1=(1./(2.*sqrt((x+y)))), sparsity_cast(vertcat(@1, @1, cos(x))

Some related ideas:

• Refactor Jacobian calculation to support large-scale problems #1055
• Constant folding in MX #172
• simplify concat(x[:n],x[n:]) #992

One could:

1. Not implement this
2. Implement generic simplifications that bring J into the shorter form J_terse (simplifications on-the-fly + simplifications in an MXFunction context that depend on the refcount), possibly displaying beneficial side-effects in other use cases
3. Supply a helper routine that performs the specific simplification of J to J_terse
4. Implement a new jacobian command (or option in existing command) to output form J_terse immediately without first doing J

[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_))