Problems pickling objects

[TomTranter]

Describe the bug
Some objects cannot be pickled due to the following error:
AttributeError: Can't pickle local object 'primitive..f_wrapped'

To Reproduce

import pybamm
import numpy as np
import pickle

def save_obj(fname, obj):
    with open(fname, 'wb') as f:
        pickle.dump(obj, f)

pybamm.set_logging_level("INFO")

# load model
model = pybamm.lithium_ion.SPMe()

# create geometry
geometry = model.default_geometry

# load parameter values and process model and geometry
param = model.default_parameter_values
param.process_model(model)
param.process_geometry(geometry)

# set mesh
mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts)

# discretise model
disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
disc.process_model(model)

# solve model
t_eval = np.linspace(0, 0.2, 100)
solver = pybamm.KLU()
solution = solver.solve(model, t_eval)

# plot
plot = pybamm.QuickPlot(model, mesh, solution)
plot.dynamic_plot()

save_obj('test_save_param', param)

Expected behaviour
Pickle saves the object

Additional context
This type of error can occur when classes are defined within other classes or functions i.e. not the global scope. There was one example of this in pybamm with the klu solver but I think the problem may also be with autograd as the f_wrapped function is from there

[TomTranter]

Pickling geometry mesh and solution all work
param, model and disc all refer to f_wrapped
solver refers to DaeSolver.set_up in the error

[TomTranter]

CamDavidsonPilon/lifelines#188 (comment)

[TomTranter]

@tinosulzer it looks like the problem is definitely with autograd. It seems we use it everywhere though! The above comment references using sympy or casadi. I know a bit about sympy and you seem to know about casadi. What do you suggest we do?

[valentinsulzer]

Worth checking both sympy and casadi both work with pickling first, but if either one does then I'm happy to drop autograd. It shouldn't be too bad, only appears when related to differentiating functions?

Casadi would definitely work and probably be quicker, but if you're more comfortable implementing sympy and there won't be any performance loss then go for that

[TomTranter]

Is your casadi branch ready?

[valentinsulzer]

most of the casadi functionality is in master now, just a few small things to wrap up which is what I'm working on now in separate issues

Edit: just checked and you can't pickle casadi either 😦:

import pickle                                                                                                                                                                                                                                                                                                 
import casadi                                                                                                                                                                                                                                                                                                 

a = casadi.SX.sym("a")                                                                                                                                                                                                                                                                                        

with open("file","wb") as f: 
     pickle.dump(a, f) 

gives

TypeError: can't pickle SwigPyObject objects

[TomTranter]

So the only thing autograd is used for is to differentiate the interpolated functions is that right?

[valentinsulzer]

No, it's used to differentiate generic user-specified functions such as the diffusivity or OCP, interpolated functions have their own inbuilt differentiation routines. But this is the only thing it's used for

[TomTranter]

Could a solution be to create interpolated functions out of these user-specified ones? Or get the users to specify the derivative too? Sympy could possibly work. There's also a package called jax that might do it but it's linux only

[valentinsulzer]

Not a fan of getting users to specify derivatives as that could easily get out of sync, creating interpolated functions could work but requires knowing ranges of values. We could just write our own rules for all the different functions though, maybe that's the easiest, I was planning on doing this anyway at some point

[TomTranter]

Not sure what you mean by that

[valentinsulzer]

At the moment we have Symbol.diff, which applies the chain and product rules to expression trees to find derivatives. But we fall back to autograd when we haven't implemented a specific function (exp, log, tanh, etc). We could just implement all of the functions we need and specify their derivatives, and never need autograd. I've already done this for a few, but not all.

[valentinsulzer]

Have implemented this but it made simplify much slower :( will investigate ...

[TomTranter]

Sympy allows pickling of symbols but not the lambda functions

import sympy as syp
import numpy as np
import pickle
import matplotlib.pyplot as plt

sto = syp.symbols('sto')
# graphite_mcmb2528_ocp_Dualfoil
u_eq = (0.194
        + 1.5 * syp.exp(-120.0 * sto)
        + 0.0351 * syp.tanh((sto - 0.286) / 0.083)
        - 0.0045 * syp.tanh((sto - 0.849) / 0.119)
        - 0.035 * syp.tanh((sto - 0.9233) / 0.05)
        - 0.0147 * syp.tanh((sto - 0.5) / 0.034)
        - 0.102 * syp.tanh((sto - 0.194) / 0.142)
        - 0.022 * syp.tanh((sto - 0.9) / 0.0164)
        - 0.011 * syp.tanh((sto - 0.124) / 0.0226)
        + 0.0155 * syp.tanh((sto - 0.105) / 0.029)
        )
u_eq_prime = u_eq.diff(sto)
u_eq_fn = syp.lambdify(sto, expr=u_eq, modules='numpy')
u_eq_prime_fn = syp.lambdify(sto, expr=u_eq_prime, modules='numpy')
S = np.linspace(0, 1.0, 101)
plt.figure()
plt.plot(S, u_eq_fn(S))
plt.plot(S, u_eq_prime_fn(S))
sympy_objs = [u_eq, u_eq_prime, u_eq_fn, u_eq_prime_fn]

with open("sympy_symbol", "wb") as f:
    pickle.dump(u_eq, f)

with open("sympy_objs", "wb") as f:
    pickle.dump(sympy_objs, f)

[TomTranter]

Hi @tinosulzer thanks for your work on this so far. I already made a branch for this which I think @rtimms made a change in. Shall we just delete my one? issue-699-enable-pickle

[valentinsulzer]

If the changes don't clash (it doesn't look like they do, but I couldn't tell what @rtimms changed) then we can just merge the two branches

[rtimms]

i merged master into pybamm-pnm. was the new branch based off of that?

[valentinsulzer]

The problem with setup becoming slower is not such a big deal if we move to using casadi as default (possible after #685 once interpolation works). The conversion to casadi is fast, only happens at the solver stage, and is only saved within the solver object, so we wouldn't necessarily have to pickle it.

How comfortable are we generally with moving to using casadi's automatic differentiation as default, with ours only as backup? Pros: very fast, saves reinventing the wheel; cons: sunk cost, ours would probably get out of date

[rtimms]

I'm in favour of moving to casadi. It's far more sophisticated than anything we would want to implement ourselves. I think the time we spent time on our method isn't reason enough to stop us moving to something better. And for me at least, implementing stuff related to the jacobian helped me learn a lot about the codebase, debugging and python generally, so the time I invested personally isn't a sunk cost in my view.

[TomTranter]

Did we not find out that casadi can't be pickled either. I think the solver would have to be pickled for multiprocessing. Is it worth me trying to get some basic multiprocessing set up, say for parametric study so we can test these options. In general I'm in favour of not re-inventing the wheel but I don't know much about casadi and how supported it will be, regularly it will be developed. Easy it is to install on various os. It would be good to have a little tutorial from @tinosulzer about it at our next dev meeting. Thanks for your hard work on it

[valentinsulzer]

Definitely worth looking into multiprocessing :) it turns out you can pickle casadi once everything is set up, just not in the intermediate stages, but that is not a problem for us. https://web.casadi.org/blog/parfor/

[martinjrobins]

Note that there are alternative to the standard library multiprocessing library that explicitly avoid standard pickle, such as this one: https://joblib.readthedocs.io/en/latest/parallel.html. This uses cloudpickle, which can pickle lambda functions.

[TomTranter]

@martinjrobins thanks for the info. I will look into these alternatives as well but I read the link and there could be some significant slow downs potentially.