parallel_map_eval.py

import casadi as ca
import pylab as pl
from time import time

# Run with CasADi 2.4.3

OPTNUMDIR = True

if OPTNUMDIR:

    ca.CasadiOptions.setOptimizedNumDir(1)


# Problem setup

# Aim: add up two dense inverted matrices, then invert the sum,
# and afterwards compute the trace

N = 200

x = ca.MX.sym("x", N)

a = pl.random(N)
C = ca.mul([a, x.T])

b = pl.random(N)
D = ca.mul([b, x.T])

x_init = pl.random(N)

f_p1 = ca.solve(C, ca.MX.eye(N), "csparse")
f_p2 = ca.solve(D, ca.MX.eye(N), "csparse")

fcn_p1 = ca.MXFunction("fcn_p1", [x], [f_p1])
fcn_p2 = ca.MXFunction("fcn_p2", [x], [f_p2])

# Parallel evaluation

f_switch = ca.Switch("f_switch", [fcn_p1, fcn_p2], fcn_p1)

[f_switch_mapped] = f_switch.map( \
    [ca.DMatrix(range(2)).T, ca.repmat(x,1,2)],"openmp")
f_mapped_eval = ca.MXFunction("f_mapped_eval", \
    [x], [f_switch_mapped[:,:N] + f_switch_mapped[:,N:]])

f_parallel = ca.trace(ca.solve(f_mapped_eval([x])[0], ca.MX.eye(N), "csparse"))
fcn_parallel = ca.MXFunction("fcn_parallel", [x], [f_parallel])

# Serial evaluation

f_serial = ca.trace(ca.solve(f_p1 + f_p2, ca.MX.eye(N), "csparse"))
fcn_serial = ca.MXFunction("fcn_serial", [x], [f_serial])

# Problem initialization

t_start = time()

nlp_parallel = ca.MXFunction("nlp_parallel", \
    ca.nlpIn(x = x), ca.nlpOut(f = f_parallel))
nlpsol_parallel = ca.NlpSolver("nlpsol_parallel", \
    "ipopt", nlp_parallel)

t_stop = time()

print "Duration init parallel: " + str(t_stop - t_start) + " s"


t_start = time()

nlp_serial = ca.MXFunction("nlp_serial", \
    ca.nlpIn(x = x), ca.nlpOut(f = f_serial))
nlpsol_serial = ca.NlpSolver("nlpsol_serial", \
    "ipopt", nlp_serial)

t_stop = time()

print "Duration init serial: " + str(t_stop - t_start) + " s"

# Problem evaluation

print "nlp_parallel(x_init) = " + str(nlp_parallel(x = x_init)["f"])
print "nlp_serial(x_init) = " + str(nlp_serial(x = x_init)["f"])