Repository: pybamm-team/PyBaMM

Issue:
Pickling parameter module
Another pickling issue is coming from `standard_parameters_lithium_ion` being a module and therefore not pickle-able. This can easily be [converted into a class](https://stackoverflow.com/questions/16717577/programmatically-turn-module-set-of-functions-into-a-python-class) to fix this issue

Discussion:

Patch:
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
index 3f0b6ffa1..70b47c057 100644
--- a/.github/PULL_REQUEST_TEMPLATE.md
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -6,7 +6,7 @@ Fixes # (issue)
 
 ## Type of change
 
-Please add a line in the relevant section of [CHANGELOG.md](https://github.com/pybamm-team/PyBaMM/blob/master/CHANGELOG.md) to document the change (include PR #) - note reverse order of PR #s.
+Please add a line in the relevant section of [CHANGELOG.md](https://github.com/pybamm-team/PyBaMM/blob/master/CHANGELOG.md) to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.
 
 - [ ] New feature (non-breaking change which adds functionality)
 - [ ] Optimization (back-end change that speeds up the code)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d33b72437..e414a1965 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,7 @@
 -   Simplified the interface for setting current functions ([#723](https://github.com/pybamm-team/PyBaMM/pull/723))
 -   Added Heaviside operator ([#723](https://github.com/pybamm-team/PyBaMM/pull/723))
 -   Added some "Getting Started" documentation ([#703](https://github.com/pybamm-team/PyBaMM/pull/703))
--   Added Simulation class ([#693](https://github.com/pybamm-team/PyBaMM/pull/693))
+-   Added Simulation class ([#693](https://github.com/pybamm-team/PyBaMM/pull/693)) with load/save functionality ([#732](https://github.com/pybamm-team/PyBaMM/pull/732))
 -   Added interface to CasADi solver ([#687](https://github.com/pybamm-team/PyBaMM/pull/687), [#691](https://github.com/pybamm-team/PyBaMM/pull/691), [#714](https://github.com/pybamm-team/PyBaMM/pull/714)). This makes the SUNDIALS DAE solvers (Scikits and KLU) truly optional (though IDA KLU is recommended for solving the DFN).
 -   Added option to use CasADi's Algorithmic Differentiation framework to calculate Jacobians ([#687](https://github.com/pybamm-team/PyBaMM/pull/687))
 -   Added method to evaluate parameters more easily ([#669](https://github.com/pybamm-team/PyBaMM/pull/669))
@@ -36,6 +36,13 @@
 -   Added warning if `ProcessedVariable` is called outside its interpolation range ([#681](https://github.com/pybamm-team/PyBaMM/pull/681))
 -   Improved the way `ProcessedVariable` objects are created in higher dimensions ([#581](https://github.com/pybamm-team/PyBaMM/pull/581))
 
+## Breaking changes
+
+-   The parameters "Bruggeman coefficient" must now be specified separately as "Bruggeman coefficient (electrolyte)" and "Bruggeman coefficient (electrode)"
+-   The current classes (`GetConstantCurrent`, `GetUserCurrent` and `GetUserData`) have now been removed. Please refer to the [`change-input-current` notebook](https://github.com/pybamm-team/PyBaMM/blob/master/examples/notebooks/change-input-current.ipynb) for information on how to specify an input current
+-   Parameter functions must now use pybamm functions instead of numpy functions (e.g. `pybamm.exp` instead of `numpy.exp`), as these are then used to construct the expression tree directly. Generally, pybamm syntax follows numpy syntax; please get in touch if a function you need is missing.
+
+
 # [v0.1.0](https://github.com/pybamm-team/PyBaMM/tree/v0.1.0) - 2019-10-08
 
 This is the first official version of PyBaMM.
diff --git a/pybamm/__init__.py b/pybamm/__init__.py
index 4e3d07864..690be6e08 100644
--- a/pybamm/__init__.py
+++ b/pybamm/__init__.py
@@ -246,7 +246,7 @@ from .solvers.idaklu_solver import IDAKLUSolver, have_idaklu
 from .processed_variable import post_process_variables, ProcessedVariable
 from .quick_plot import QuickPlot, ax_min, ax_max
 
-from .simulation import Simulation
+from .simulation import Simulation, load_sim
 
 #
 # Remove any imported modules, so we don't expose them as part of pybamm
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index d48f09c04..708fbeb50 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -137,13 +137,8 @@ class Discretisation(object):
             # since they point to the same object
             model_disc = model
         else:
-            # create a model of the same class as the original model
-            model_disc = model.__class__(model.options)
-            model_disc.name = model.name
-            model_disc.options = model.options
-            model_disc.use_jacobian = model.use_jacobian
-            model_disc.use_simplify = model.use_simplify
-            model_disc.convert_to_format = model.convert_to_format
+            # create an empty copy of the original model
+            model_disc = model.new_copy()
 
         model_disc.bcs = self.bcs
 
diff --git a/pybamm/expression_tree/binary_operators.py b/pybamm/expression_tree/binary_operators.py
index 9a78bfd91..cec1dc1f0 100644
--- a/pybamm/expression_tree/binary_operators.py
+++ b/pybamm/expression_tree/binary_operators.py
@@ -227,7 +227,9 @@ class Power(BinaryOperator):
 
     def _binary_evaluate(self, left, right):
         """ See :meth:`pybamm.BinaryOperator._binary_evaluate()`. """
-        return left ** right
+        # don't raise RuntimeWarning for NaNs
+        with np.errstate(invalid="ignore"):
+            return left ** right
 
     def _binary_simplify(self, left, right):
         """ See :meth:`pybamm.BinaryOperator._binary_simplify()`. """
diff --git a/pybamm/models/base_model.py b/pybamm/models/base_model.py
index 3956fce42..ecf51e994 100644
--- a/pybamm/models/base_model.py
+++ b/pybamm/models/base_model.py
@@ -1,12 +1,22 @@
 #
 # Base model class
 #
-import pybamm
-
+import inspect
 import numbers
+import pybamm
 import warnings
 
 
+class ParamClass:
+    """Class for converting a module of parameters into a class. For pickling."""
+
+    def __init__(self, methods):
+        for k, v in methods.__dict__.items():
+            # don't save module attributes (e.g. pybamm, numpy)
+            if not (k.startswith("__") or inspect.ismodule(v)):
+                self.__dict__[k] = v
+
+
 class BaseModel(object):
     """Base model class for other models to extend.
 
@@ -263,8 +273,14 @@ class BaseModel(object):
         self._jacobian_algebraic = jacobian_algebraic
 
     @property
-    def set_of_parameters(self):
-        return self._set_of_parameters
+    def param(self):
+        return self._param
+
+    @param.setter
+    def param(self, values):
+        # convert module into a class
+        # (StackOverflow: https://tinyurl.com/yk3euon3)
+        self._param = ParamClass(values)
 
     @property
     def options(self):
@@ -277,6 +293,16 @@ class BaseModel(object):
     def __getitem__(self, key):
         return self.rhs[key]
 
+    def new_copy(self, options=None):
+        "Create an empty copy with identical options, or new options if specified"
+        options = options or self.options
+        new_model = self.__class__(options)
+        new_model.name = self.name
+        new_model.use_jacobian = self.use_jacobian
+        new_model.use_simplify = self.use_simplify
+        new_model.convert_to_format = self.convert_to_format
+        return new_model
+
     def update(self, *submodels):
         """
         Update model to add new physics from submodels
diff --git a/pybamm/parameters/parameter_values.py b/pybamm/parameters/parameter_values.py
index ce81251fa..8566ece72 100644
--- a/pybamm/parameters/parameter_values.py
+++ b/pybamm/parameters/parameter_values.py
@@ -271,12 +271,7 @@ class ParameterValues(dict):
             model = unprocessed_model
         else:
             # create a blank model of the same class
-            model = unprocessed_model.__class__(unprocessed_model.options)
-            model.name = unprocessed_model.name
-            model.options = unprocessed_model.options
-            model.use_jacobian = unprocessed_model.use_jacobian
-            model.use_simplify = unprocessed_model.use_simplify
-            model.convert_to_format = unprocessed_model.convert_to_format
+            model = unprocessed_model.new_copy()
 
         if len(unprocessed_model.rhs) == 0 and len(unprocessed_model.algebraic) == 0:
             raise pybamm.ModelError("Cannot process parameters for empty model")
diff --git a/pybamm/simulation.py b/pybamm/simulation.py
index ae8ae094e..1a75fc563 100644
--- a/pybamm/simulation.py
+++ b/pybamm/simulation.py
@@ -1,3 +1,7 @@
+#
+# Simulation class
+#
+import pickle
 import pybamm
 import numpy as np
 import copy
@@ -71,7 +75,7 @@ class Simulation:
         """
         A method to reset a simulation back to its unprocessed state.
         """
-        self.model = self._model_class(self._model_options)
+        self.model = self.model.new_copy(self._model_options)
         self.geometry = copy.deepcopy(self._unprocessed_geometry)
         self._model_with_set_params = None
         self._built_model = None
@@ -321,3 +325,23 @@ class Simulation:
             or spatial_methods
         ):
             self.reset()
+
+    def save(self, filename):
+        """Save simulation using pickle"""
+        if self.model.convert_to_format == "python":
+            # We currently cannot save models in the 'python'
+            raise NotImplementedError(
+                """
+                Cannot save simulation if model format is python.
+                Set model.convert_to_format = 'casadi' instead.
+                """
+            )
+        with open(filename, "wb") as f:
+            pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
+
+
+def load_sim(filename):
+    """Load a saved simulation"""
+    with open(filename, "rb") as f:
+        sim = pickle.load(f)
+    return sim
diff --git a/pybamm/solvers/casadi_solver.py b/pybamm/solvers/casadi_solver.py
index c14f49802..b88ceed73 100644
--- a/pybamm/solvers/casadi_solver.py
+++ b/pybamm/solvers/casadi_solver.py
@@ -194,7 +194,11 @@ class CasadiSolver(pybamm.DaeSolver):
         solve_start_time = timer.time()
         pybamm.logger.debug("Calling DAE solver")
         solution = self.integrate_casadi(
-            self.casadi_problem, self.y0, t_eval, mass_matrix=model.mass_matrix.entries
+            self.casadi_rhs,
+            self.casadi_algebraic,
+            self.y0,
+            t_eval,
+            mass_matrix=model.mass_matrix.entries,
         )
         solve_time = timer.time() - solve_start_time
 
@@ -203,7 +207,7 @@ class CasadiSolver(pybamm.DaeSolver):
 
         return solution, solve_time, termination
 
-    def integrate_casadi(self, problem, y0, t_eval, mass_matrix=None):
+    def integrate_casadi(self, rhs, algebraic, y0, t_eval, mass_matrix=None):
         """
         Solve a DAE model defined by residuals with initial conditions y0.
 
@@ -233,11 +237,24 @@ class CasadiSolver(pybamm.DaeSolver):
             options["calc_ic"] = True
 
         # set up and solve
+        t = casadi.MX.sym("t")
+        y_diff = casadi.MX.sym("y_diff", rhs(0, y0).shape[0])
+        if algebraic is None:
+            problem = {"t": t, "x": y_diff, "ode": rhs(t, y_diff)}
+        else:
+            y_alg = casadi.MX.sym("y_alg", algebraic(0, y0).shape[0])
+            y = casadi.vertcat(y_diff, y_alg)
+            problem = {
+                "t": t,
+                "x": y_diff,
+                "z": y_alg,
+                "ode": rhs(t, y),
+                "alg": algebraic(t, y),
+            }
         integrator = casadi.integrator("F", self.method, problem, options)
         try:
             # Try solving
-            len_rhs = problem["x"].size()[0]
-            y0_diff, y0_alg = np.split(y0, [len_rhs])
+            y0_diff, y0_alg = np.split(y0, [y_diff.shape[0]])
             sol = integrator(x0=y0_diff, z0=y0_alg)
             y_values = np.concatenate([sol["xf"].full(), sol["zf"].full()])
             return pybamm.Solution(t_eval, y_values, None, None, "final time")
diff --git a/pybamm/solvers/dae_solver.py b/pybamm/solvers/dae_solver.py
index ae4316228..7ea8283b7 100644
--- a/pybamm/solvers/dae_solver.py
+++ b/pybamm/solvers/dae_solver.py
@@ -146,8 +146,7 @@ class DaeSolver(pybamm.BaseSolver):
                 jac_algebraic = pybamm.EvaluatorPython(jac_algebraic)
                 jac = pybamm.EvaluatorPython(jac)
 
-            def jacobian(t, y):
-                return jac.evaluate(t, y, known_evals={})[0]
+            jacobian = Jacobian(jac.evaluate)
 
             def jacobian_alg(t, y):
                 return jac_algebraic.evaluate(t, y, known_evals={})[0]
@@ -167,11 +166,8 @@ class DaeSolver(pybamm.BaseSolver):
             }
 
         # Calculate consistent initial conditions for the algebraic equations
-        def rhs(t, y):
-            return concatenated_rhs.evaluate(t, y, known_evals={})[0][:, 0]
-
-        def algebraic(t, y):
-            return concatenated_algebraic.evaluate(t, y, known_evals={})[0][:, 0]
+        rhs = Rhs(concatenated_rhs.evaluate)
+        algebraic = Algebraic(concatenated_algebraic.evaluate)
 
         if len(model.algebraic) > 0:
             y0 = self.calculate_consistent_initial_conditions(
@@ -184,38 +180,19 @@ class DaeSolver(pybamm.BaseSolver):
             # can use DAE solver to solve ODE model
             y0 = model.concatenated_initial_conditions[:, 0]
 
-        # Create functions to evaluate residuals
-        def residuals(t, y, ydot):
-            pybamm.logger.debug(
-                "Evaluating residuals for {} at t={}".format(model.name, t)
-            )
-            y = y[:, np.newaxis]
-            rhs_eval, known_evals = concatenated_rhs.evaluate(t, y, known_evals={})
-            # reuse known_evals
-            alg_eval = concatenated_algebraic.evaluate(t, y, known_evals=known_evals)[0]
-            # turn into 1D arrays
-            rhs_eval = rhs_eval[:, 0]
-            alg_eval = alg_eval[:, 0]
-            return (
-                np.concatenate((rhs_eval, alg_eval)) - model.mass_matrix.entries @ ydot
-            )
-
         # Create event-dependent function to evaluate events
-        def event_fun(event):
-            def eval_event(t, y):
-                return event.evaluate(t, y)
-
-            return eval_event
-
-        event_funs = [event_fun(event) for event in events.values()]
+        def get_event_class(event):
+            return EvalEvent(event.evaluate)
 
         # Add the solver attributes
         self.y0 = y0
         self.rhs = rhs
         self.algebraic = algebraic
-        self.residuals = residuals
+        self.residuals = Residuals(
+            model, concatenated_rhs.evaluate, concatenated_algebraic.evaluate
+        )
         self.events = events
-        self.event_funs = event_funs
+        self.event_funs = [get_event_class(event) for event in events.values()]
         self.jacobian = jacobian
 
     def set_up_casadi(self, model):
@@ -269,8 +246,7 @@ class DaeSolver(pybamm.BaseSolver):
                 "jacobian", [t_casadi, y_casadi], [casadi_jac_alg]
             )
 
-            def jacobian(t, y):
-                return casadi_jac_fn(t, y)
+            jacobian = JacobianCasadi(casadi_jac_fn)
 
             def jacobian_alg(t, y):
                 return casadi_jac_alg_fn(t, y)
@@ -280,11 +256,8 @@ class DaeSolver(pybamm.BaseSolver):
             jacobian_alg = None
 
         # Calculate consistent initial conditions for the algebraic equations
-        def rhs(t, y):
-            return concatenated_rhs_fn(t, y).full()[:, 0]
-
-        def algebraic(t, y):
-            return concatenated_algebraic_fn(t, y).full()[:, 0]
+        rhs = RhsCasadi(concatenated_rhs_fn)
+        algebraic = AlgebraicCasadi(concatenated_algebraic_fn)
 
         if len(model.algebraic) > 0:
 
@@ -298,25 +271,10 @@ class DaeSolver(pybamm.BaseSolver):
             # can use DAE solver to solve ODE model
             y0 = model.concatenated_initial_conditions[:, 0]
 
-        # Create functions to evaluate residuals
-
-        def residuals(t, y, ydot):
-            pybamm.logger.debug(
-                "Evaluating residuals for {} at t={}".format(model.name, t)
-            )
-            states_eval = all_states_fn(t, y).full()[:, 0]
-            return states_eval - model.mass_matrix.entries @ ydot
-
         # Create event-dependent function to evaluate events
-        def event_fun(event):
+        def get_event_class(event):
             casadi_event_fn = casadi.Function("event", [t_casadi, y_casadi], [event])
-
-            def eval_event(t, y):
-                return casadi_event_fn(t, y)
-
-            return eval_event
-
-        event_funs = [event_fun(event) for event in casadi_events.values()]
+            return EvalEvent(casadi_event_fn)
 
         # Add the solver attributes
         # Note: these are the (possibly) converted to python version rhs, algebraic
@@ -324,19 +282,14 @@ class DaeSolver(pybamm.BaseSolver):
         self.y0 = y0
         self.rhs = rhs
         self.algebraic = algebraic
-        self.residuals = residuals
+        self.residuals = ResidualsCasadi(model, all_states_fn)
         self.events = model.events
-        self.event_funs = event_funs
+        self.event_funs = [get_event_class(event) for event in casadi_events.values()]
         self.jacobian = jacobian
 
-        # Create CasADi problem for the CasADi solver
-        self.casadi_problem = {
-            "t": t_casadi,
-            "x": y_diff,
-            "z": y_alg,
-            "ode": concatenated_rhs,
-            "alg": concatenated_algebraic,
-        }
+        # Save CasADi functions for the CasADi solver
+        self.casadi_rhs = concatenated_rhs_fn
+        self.casadi_algebraic = concatenated_algebraic_fn
 
     def calculate_consistent_initial_conditions(
         self, rhs, algebraic, y0_guess, jac=None
@@ -456,3 +409,103 @@ class DaeSolver(pybamm.BaseSolver):
             A function that takes in t, y and ydot and returns the Jacobian
         """
         raise NotImplementedError
+
+
+class Rhs:
+    "Returns information about rhs at time t and state y"
+
+    def __init__(self, concatenated_rhs_fn):
+        self.concatenated_rhs_fn = concatenated_rhs_fn
+
+    def __call__(self, t, y):
+        return self.concatenated_rhs_fn(t, y, known_evals={})[0][:, 0]
+
+
+class RhsCasadi(Rhs):
+    "Returns information about rhs at time t and state y, with CasADi"
+
+    def __call__(self, t, y):
+        return self.concatenated_rhs_fn(t, y).full()[:, 0]
+
+
+class Algebraic:
+    "Returns information about algebraic equations at time t and state y"
+
+    def __init__(self, concatenated_algebraic_fn):
+        self.concatenated_algebraic_fn = concatenated_algebraic_fn
+
+    def __call__(self, t, y):
+        return self.concatenated_algebraic_fn(t, y, known_evals={})[0][:, 0]
+
+
+class AlgebraicCasadi(Algebraic):
+    "Returns information about algebraic equations at time t and state y, with CasADi"
+
+    def __call__(self, t, y):
+        return self.concatenated_algebraic_fn(t, y).full()[:, 0]
+
+
+class Residuals:
+    "Returns information about residuals at time t and state y"
+
+    def __init__(self, model, concatenated_rhs_fn, concatenated_algebraic_fn):
+        self.model = model
+        self.concatenated_rhs_fn = concatenated_rhs_fn
+        self.concatenated_algebraic_fn = concatenated_algebraic_fn
+        self.mass_matrix = model.mass_matrix.entries
+
+    def __call__(self, t, y, ydot):
+        pybamm.logger.debug(
+            "Evaluating residuals for {} at t={}".format(self.model.name, t)
+        )
+        y = y[:, np.newaxis]
+        rhs_eval, known_evals = self.concatenated_rhs_fn(t, y, known_evals={})
+        # reuse known_evals
+        alg_eval = self.concatenated_algebraic_fn(t, y, known_evals=known_evals)[0]
+        # turn into 1D arrays
+        rhs_eval = rhs_eval[:, 0]
+        alg_eval = alg_eval[:, 0]
+        return np.concatenate([rhs_eval, alg_eval]) - self.mass_matrix @ ydot
+
+
+class ResidualsCasadi(Residuals):
+    "Returns information about residuals at time t and state y, with CasADi"
+
+    def __init__(self, model, all_states_fn):
+        self.model = model
+        self.all_states_fn = all_states_fn
+        self.mass_matrix = model.mass_matrix.entries
+
+    def __call__(self, t, y, ydot):
+        pybamm.logger.debug(
+            "Evaluating residuals for {} at t={}".format(self.model.name, t)
+        )
+        states_eval = self.all_states_fn(t, y).full()[:, 0]
+        return states_eval - self.mass_matrix @ ydot
+
+
+class EvalEvent:
+    "Returns information about events at time t and state y"
+
+    def __init__(self, event_fn):
+        self.event_fn = event_fn
+
+    def __call__(self, t, y):
+        return self.event_fn(t, y)
+
+
+class Jacobian:
+    "Returns information about the jacobian at time t and state y"
+
+    def __init__(self, jac_fn):
+        self.jac_fn = jac_fn
+
+    def __call__(self, t, y):
+        return self.jac_fn(t, y, known_evals={})[0]
+
+
+class JacobianCasadi(Jacobian):
+    "Returns information about the jacobian at time t and state y, with CasADi"
+
+    def __call__(self, t, y):
+        return self.jac_fn(t, y)
diff --git a/pybamm/solvers/ode_solver.py b/pybamm/solvers/ode_solver.py
index 168bc0479..6feba0d98 100644
--- a/pybamm/solvers/ode_solver.py
+++ b/pybamm/solvers/ode_solver.py
@@ -118,28 +118,13 @@ class OdeSolver(pybamm.BaseSolver):
                 name: pybamm.EvaluatorPython(event) for name, event in events.items()
             }
 
-        # Create function to evaluate rhs
-        def dydt(t, y):
-            pybamm.logger.debug("Evaluating RHS for {} at t={}".format(model.name, t))
-            y = y[:, np.newaxis]
-            dy = concatenated_rhs.evaluate(t, y, known_evals={})[0]
-            return dy[:, 0]
-
         # Create event-dependent function to evaluate events
-        def event_fun(event):
-            def eval_event(t, y):
-                return event.evaluate(t, y)
-
-            return eval_event
-
-        event_funs = [event_fun(event) for event in events.values()]
+        def get_event_class(event):
+            return EvalEvent(event.evaluate)
 
         # Create function to evaluate jacobian
         if jac_rhs is not None:
-
-            def jacobian(t, y):
-                return jac_rhs.evaluate(t, y, known_evals={})[0]
-
+            jacobian = Jacobian(jac_rhs.evaluate)
         else:
             jacobian = None
 
@@ -147,9 +132,9 @@ class OdeSolver(pybamm.BaseSolver):
         # Note: these are the (possibly) converted to python version rhs, algebraic
         # etc. The expression tree versions of these are attributes of the model
         self.y0 = y0
-        self.dydt = dydt
+        self.dydt = Dydt(model, concatenated_rhs.evaluate)
         self.events = events
-        self.event_funs = event_funs
+        self.event_funs = [get_event_class(event) for event in events.values()]
         self.jacobian = jacobian
 
     def set_up_casadi(self, model):
@@ -191,21 +176,10 @@ class OdeSolver(pybamm.BaseSolver):
             "rhs", [t_casadi, y_casadi], [concatenated_rhs]
         )
 
-        def dydt(t, y):
-            pybamm.logger.debug("Evaluating RHS for {} at t={}".format(model.name, t))
-            dy = concatenated_rhs_fn(t, y).full()
-            return dy[:, 0]
-
         # Create event-dependent function to evaluate events
-        def event_fun(event):
+        def get_event_class(event):
             casadi_event_fn = casadi.Function("event", [t_casadi, y_casadi], [event])
-
-            def eval_event(t, y):
-                return casadi_event_fn(t, y)
-
-            return eval_event
-
-        event_funs = [event_fun(event) for event in casadi_events.values()]
+            return EvalEvent(casadi_event_fn)
 
         # Create function to evaluate jacobian
         if model.use_jacobian:
@@ -214,18 +188,16 @@ class OdeSolver(pybamm.BaseSolver):
             casadi_jac_fn = casadi.Function(
                 "jacobian", [t_casadi, y_casadi], [casadi_jac]
             )
-
-            def jacobian(t, y):
-                return casadi_jac_fn(t, y)
+            jacobian = JacobianCasadi(casadi_jac_fn)
 
         else:
             jacobian = None
 
         # Add the solver attributes
         self.y0 = y0
-        self.dydt = dydt
+        self.dydt = DydtCasadi(model, concatenated_rhs_fn)
         self.events = model.events
-        self.event_funs = event_funs
+        self.event_funs = [get_event_class(event) for event in casadi_events.values()]
         self.jacobian = jacobian
 
     def integrate(
@@ -251,3 +223,56 @@ class OdeSolver(pybamm.BaseSolver):
             A function that takes in t and y and returns the Jacobian
         """
         raise NotImplementedError
+
+
+# Set up caller classes outside of the solver object to allow pickling
+
+
+class Dydt:
+    "Returns information about time derivatives at time t and state y"
+
+    def __init__(self, model, concatenated_rhs_fn):
+        self.model = model
+        self.concatenated_rhs_fn = concatenated_rhs_fn
+
+    def __call__(self, t, y):
+        pybamm.logger.debug("Evaluating RHS for {} at t={}".format(self.model.name, t))
+        y = y[:, np.newaxis]
+        dy = self.concatenated_rhs_fn(t, y, known_evals={})[0]
+        return dy[:, 0]
+
+
+class DydtCasadi(Dydt):
+    "Returns information about time derivatives at time t and state y, with CasADi"
+
+    def __call__(self, t, y):
+        pybamm.logger.debug("Evaluating RHS for {} at t={}".format(self.model.name, t))
+        dy = self.concatenated_rhs_fn(t, y).full()
+        return dy[:, 0]
+
+
+class EvalEvent:
+    "Returns information about events at time t and state y"
+
+    def __init__(self, event_fn):
+        self.event_fn = event_fn
+
+    def __call__(self, t, y):
+        return self.event_fn(t, y)
+
+
+class Jacobian:
+    "Returns information about the jacobian at time t and state y"
+
+    def __init__(self, jac_fn):
+        self.jac_fn = jac_fn
+
+    def __call__(self, t, y):
+        return self.jac_fn(t, y, known_evals={})[0]
+
+
+class JacobianCasadi(Jacobian):
+    "Returns information about the jacobian at time t and state y, with CasADi"
+
+    def __call__(self, t, y):
+        return self.jac_fn(t, y)


Test patch:
diff --git a/tests/unit/test_simulation.py b/tests/unit/test_simulation.py
index 81bf3992b..25cd6c7bf 100644
--- a/tests/unit/test_simulation.py
+++ b/tests/unit/test_simulation.py
@@ -130,6 +130,69 @@ class TestSimulation(unittest.TestCase):
         sim.specs(spatial_methods=spatial_methods)
         sim.build()
 
+    def test_save_load(self):
+        model = pybamm.lead_acid.LOQS()
+        model.use_jacobian = True
+        sim = pybamm.Simulation(model)
+
+        sim.save("test.pickle")
+        sim_load = pybamm.load_sim("test.pickle")
+        self.assertEqual(sim.model.name, sim_load.model.name)
+
+        # save after solving
+        sim.solve()
+        sim.save("test.pickle")
+        sim_load = pybamm.load_sim("test.pickle")
+        self.assertEqual(sim.model.name, sim_load.model.name)
+
+        # with python formats
+        model.convert_to_format = None
+        sim = pybamm.Simulation(model)
+        sim.solve()
+        sim.save("test.pickle")
+        model.convert_to_format = "python"
+        sim = pybamm.Simulation(model)
+        sim.solve()
+        with self.assertRaisesRegex(
+            NotImplementedError, "Cannot save simulation if model format is python"
+        ):
+            sim.save("test.pickle")
+
+    def test_save_load_dae(self):
+        model = pybamm.lead_acid.LOQS({"surface form": "algebraic"})
+        model.use_jacobian = True
+        sim = pybamm.Simulation(model)
+
+        # save after solving
+        sim.solve()
+        sim.save("test.pickle")
+        sim_load = pybamm.load_sim("test.pickle")
+        self.assertEqual(sim.model.name, sim_load.model.name)
+
+        # with python format
+        model.convert_to_format = None
+        sim = pybamm.Simulation(model)
+        sim.solve()
+        sim.save("test.pickle")
+
+        # with Casadi solver
+        sim = pybamm.Simulation(model, solver=pybamm.CasadiSolver())
+        sim.solve()
+        sim.save("test.pickle")
+        sim_load = pybamm.load_sim("test.pickle")
+        self.assertEqual(sim.model.name, sim_load.model.name)
+
+    @unittest.skipIf(not pybamm.have_idaklu(), "idaklu solver is not installed")
+    def test_save_load_klu(self):
+        model = pybamm.lead_acid.LOQS({"surface form": "algebraic"})
+        model.use_jacobian = True
+        # with KLU solver
+        sim = pybamm.Simulation(model, solver=pybamm.IDAKLUSolver())
+        sim.solve()
+        sim.save("test.pickle")
+        sim_load = pybamm.load_sim("test.pickle")
+        self.assertEqual(sim.model.name, sim_load.model.name)
+
     def test_set_defaults(self):
         model = pybamm.lithium_ion.SPM()
 
diff --git a/tests/unit/test_solvers/test_casadi_solver.py b/tests/unit/test_solvers/test_casadi_solver.py
index fb0e2938a..13d66ff0d 100644
--- a/tests/unit/test_solvers/test_casadi_solver.py
+++ b/tests/unit/test_solvers/test_casadi_solver.py
@@ -14,13 +14,14 @@ class TestCasadiSolver(unittest.TestCase):
         # Constant
         solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8, method="idas")
 
+        t = casadi.MX.sym("t")
         y = casadi.MX.sym("y")
         constant_growth = casadi.MX(0.5)
-        problem = {"x": y, "ode": constant_growth}
+        rhs = casadi.Function("rhs", [t, y], [constant_growth])
 
         y0 = np.array([0])
         t_eval = np.linspace(0, 1, 100)
-        solution = solver.integrate_casadi(problem, y0, t_eval)
+        solution = solver.integrate_casadi(rhs, None, y0, t_eval)
         np.testing.assert_array_equal(solution.t, t_eval)
         np.testing.assert_allclose(0.5 * solution.t, solution.y[0])
 
@@ -28,11 +29,11 @@ class TestCasadiSolver(unittest.TestCase):
         solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8, method="cvodes")
 
         exponential_decay = -0.1 * y
-        problem = {"x": y, "ode": exponential_decay}
+        rhs = casadi.Function("rhs", [t, y], [exponential_decay])
 
         y0 = np.array([1])
         t_eval = np.linspace(0, 1, 100)
-        solution = solver.integrate_casadi(problem, y0, t_eval)
+        solution = solver.integrate_casadi(rhs, None, y0, t_eval)
         np.testing.assert_allclose(solution.y[0], np.exp(-0.1 * solution.t))
         self.assertEqual(solution.termination, "final time")
 
@@ -40,16 +41,17 @@ class TestCasadiSolver(unittest.TestCase):
         # Turn off warnings to ignore sqrt error
         warnings.simplefilter("ignore")
 
+        t = casadi.MX.sym("t")
         y = casadi.MX.sym("y")
         sqrt_decay = -np.sqrt(y)
 
         y0 = np.array([1])
         t_eval = np.linspace(0, 3, 100)
         solver = pybamm.CasadiSolver()
-        problem = {"x": y, "ode": sqrt_decay}
+        rhs = casadi.Function("rhs", [t, y], [sqrt_decay])
         # Expect solver to fail when y goes negative
         with self.assertRaises(pybamm.SolverError):
-            solver.integrate_casadi(problem, y0, t_eval)
+            solver.integrate_casadi(rhs, None, y0, t_eval)
 
         # Set up as a model and solve
         # Create model