librelist archives

« back to archive

Trying to parallelize scipy.intergrate.odeint calls with generated cython rhs functions

Trying to parallelize scipy.intergrate.odeint calls with generated cython rhs functions

From:
Jason Moore
Date:
2015-02-05 @ 23:52
Hi,

I'd like run ode integrations in parallel with joblib for optimization
purposes. This works fine:

import numpy as np
from scipy.integrate import odeint
from joblib import Parallel, delayed

x = np.random.random((16, 2))
t = np.linspace(0.0, 10.0, 100)
p = np.random.random(3)

def rhs(xx, tt, pp):
    return 2.0 * xx

res = Parallel(n_jobs=-1)(delayed(odeint)(rhs, x_i, t_vec, (p,)) for x_i in
x)

which is awesome! So simple. But if I use a lambda function:

rhs = lambda xx, tt, pp: 2.0 * xx

The Parallel call fails with:

TypeError: can't pickle function objects

which I get the sense from google searches that lambdas aren't supported,
but don't see that explicit anywhere.

Finally, what I really want to do is use the rhs functions generated from
PyDy [1]:

import numpy as np
from scipy.integrate import odeint
from joblib import Parallel, delayed

# from HEAD of PyDy master
from pydy.models import multi_mass_spring_damper

sys = multi_mass_spring_damper()

x = np.random.random((16, len(sys.states)))
t = np.linspace(0.0, 10.0, 100)
p = np.random.random(len(sys.constants_symbols))

# doesn't work with a lambdified function or cython function
rhs = sys.generate_ode_function(generator='lambdify')

res = Parallel(n_jobs=-1)(delayed(odeint)(rhs, x_i, t_vec, (p,)) for x_i in
x)

The rhs function causes failure in the case shown because inside PyDy,
SymPy is generating a "lambdified" function i.e. a Python lambda [2].

The rhs can also be generated as a Cython function:

rhs = sys.generate_ode_function(generator='cython')

rhs is now a function imported from a generated and compiled cython module.
In this case I get:

PicklingError: Can't pickle <function rhs at 0x7f5c6110bcf8>: it's not
found as pydy.codegen.ode_function_generators.rhs

The rhs function that is generated in PyDy is done so like:

def gen_rhs(...):
    def rhs(x, t, p):
        return call_cython_function(x, t, p)
    return rhs

So the closure could be causing issues.

Is there any way to get this to work? I'd really like it to work with the
Cythonized rhs function, but I could write a generator that generates a
normally defined Python function with def, also.

[1] http://github.com/pydy/pydy
[2] http://docs.sympy.org/dev/modules/utilities/lambdify.html

Thanks,

Jason
moorepants.info
+01 530-601-9791

Re: Trying to parallelize scipy.intergrate.odeint calls with generated cython rhs functions

From:
Jason Moore
Date:
2015-02-06 @ 00:20
May have found a solution. Maybe I'm just trying to pass too many args to
delayed.

I tried this pure multiprocessing approach and it worked:

from multiprocessing import Pool

import numpy as np
from pydy.models import multi_mass_spring_damper

sys = multi_mass_spring_damper()

x = np.random.random((16, len(sys.states)))
t = 0.0
p = np.random.random(len(sys.constants_symbols))

# doesn't work with a lambdified function or cython function
rhs = sys.generate_ode_function(generator='cython')

def rhs2(xx):
    return rhs(xx, t, p)

if __name__ == '__main__':
    p = Pool(5)
    print(p.map(rhs2, [x_i for x_i in x]))

And then the same thing with joblib:

import numpy as np
from scipy.integrate import odeint
from joblib import Parallel, delayed

from pydy.models import multi_mass_spring_damper

sys = multi_mass_spring_damper()

x = np.random.random((16, len(sys.states)))
t = 0.0
p = np.random.random(len(sys.constants_symbols))

t_vec = np.linspace(0.0, 10.0, 100)

rhs = sys.generate_ode_function(generator='cython')

def rhs2(xx):
    return rhs(xx, t, p)

def odeint2(x0):
    return odeint(rhs, x0, t_vec, args=(p,))

res = Parallel(n_jobs=-1)(delayed(rhs2)(x_i) for x_i in x)

res = Parallel(n_jobs=-1)(delayed(odeint2)(x_i) for x_i in x)

And it works. Not sure what exactly I was doing wrong though.


Jason
moorepants.info
+01 530-601-9791

On Thu, Feb 5, 2015 at 3:52 PM, Jason Moore <moorepants@gmail.com> wrote:

> Hi,
>
> I'd like run ode integrations in parallel with joblib for optimization
> purposes. This works fine:
>
> import numpy as np
> from scipy.integrate import odeint
> from joblib import Parallel, delayed
>
> x = np.random.random((16, 2))
> t = np.linspace(0.0, 10.0, 100)
> p = np.random.random(3)
>
> def rhs(xx, tt, pp):
>     return 2.0 * xx
>
> res = Parallel(n_jobs=-1)(delayed(odeint)(rhs, x_i, t_vec, (p,)) for x_i
> in x)
>
> which is awesome! So simple. But if I use a lambda function:
>
> rhs = lambda xx, tt, pp: 2.0 * xx
>
> The Parallel call fails with:
>
> TypeError: can't pickle function objects
>
> which I get the sense from google searches that lambdas aren't supported,
> but don't see that explicit anywhere.
>
> Finally, what I really want to do is use the rhs functions generated from
> PyDy [1]:
>
> import numpy as np
> from scipy.integrate import odeint
> from joblib import Parallel, delayed
>
> # from HEAD of PyDy master
> from pydy.models import multi_mass_spring_damper
>
> sys = multi_mass_spring_damper()
>
> x = np.random.random((16, len(sys.states)))
> t = np.linspace(0.0, 10.0, 100)
> p = np.random.random(len(sys.constants_symbols))
>
> # doesn't work with a lambdified function or cython function
> rhs = sys.generate_ode_function(generator='lambdify')
>
> res = Parallel(n_jobs=-1)(delayed(odeint)(rhs, x_i, t_vec, (p,)) for x_i
> in x)
>
> The rhs function causes failure in the case shown because inside PyDy,
> SymPy is generating a "lambdified" function i.e. a Python lambda [2].
>
> The rhs can also be generated as a Cython function:
>
> rhs = sys.generate_ode_function(generator='cython')
>
> rhs is now a function imported from a generated and compiled cython
> module. In this case I get:
>
> PicklingError: Can't pickle <function rhs at 0x7f5c6110bcf8>: it's not
> found as pydy.codegen.ode_function_generators.rhs
>
> The rhs function that is generated in PyDy is done so like:
>
> def gen_rhs(...):
>     def rhs(x, t, p):
>         return call_cython_function(x, t, p)
>     return rhs
>
> So the closure could be causing issues.
>
> Is there any way to get this to work? I'd really like it to work with the
> Cythonized rhs function, but I could write a generator that generates a
> normally defined Python function with def, also.
>
> [1] http://github.com/pydy/pydy
> [2] http://docs.sympy.org/dev/modules/utilities/lambdify.html
>
> Thanks,
>
> Jason
> moorepants.info
> +01 530-601-9791
>