- Trying to parallelize scipy.intergrate.odeint calls with generated cython rhs functions by Jason Moore

- 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

- 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 >