Allow sim to be pickled with output variables (#5225)

* Bugfix #5224 add output variable pkl attributes to _setup
Allows pickling solver when output variables are set

* Add to changelog

---------

Co-authored-by: Martin Robinson <[email protected]>

Author: pipliggins

CHANGELOG.md

@@ -11,6 +11,7 @@
 
 ## Bug fixes
 
+- Fixed a bug where simulations using output variables in `IDAKLUSolver` couldn't be pickled ([#5225](https://github.com/pybamm-team/PyBaMM/pull/5225))
 - Added explicit warning in installation docs about unmaintained Conda recipe due to pybammsolvers split (Fixes #5155). See pull request [#5206](https://github.com/pybamm-team/PyBaMM/pull/5206)
 - Fixed a bug where time-based Heaviside or modulo discontinuities could trigger out-of-bounds errors in time arrays. ([#5205](https://github.com/pybamm-team/PyBaMM/pull/5205))
 - Fixed a bug using a time-varying input with heaviside or modulo functions using the `IDAKLUSolver`. ([#4994](https://github.com/pybamm-team/PyBaMM/pull/4994))

src/pybamm/solvers/idaklu_solver.py

@@ -445,8 +445,11 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False):
             "output_variables": self.output_variables,
             "var_fcns": self.computed_var_fcns,
             "var_idaklu_fcns": self.var_idaklu_fcns,
+            "var_idaklu_fcns_pkl": self.var_idaklu_fcns_pkl,
             "dvar_dy_idaklu_fcns": self.dvar_dy_idaklu_fcns,
+            "dvar_dy_idaklu_fcns_pkl": self.dvar_dy_idaklu_fcns_pkl,
             "dvar_dp_idaklu_fcns": self.dvar_dp_idaklu_fcns,
+            "dvar_dp_idaklu_fcns_pkl": self.dvar_dp_idaklu_fcns_pkl,
         }
 
         solver = self._setup["solver_function"](
@@ -483,6 +486,8 @@ def __getstate__(self):
         if not hasattr(self, "_setup"):
             return self.__dict__
 
+        self.var_idaklu_fcns = []
+
         for key in [
             "solver",
             "solver_function",
@@ -492,6 +497,9 @@ def __getstate__(self):
             "mass_action",
             "sensfn",
             "rootfn",
+            "var_idaklu_fcns",
+            "dvar_dy_idaklu_fcns",
+            "dvar_dp_idaklu_fcns",
         ]:
             del self._setup[key]
         return self.__dict__
@@ -503,6 +511,10 @@ def __setstate__(self, d):
         if not hasattr(self, "_setup"):
             return
 
+        self.var_idaklu_fcns = [
+            idaklu.generate_function(f) for f in self.var_idaklu_fcns_pkl
+        ]
+
         for key in [
             "rhs_algebraic",
             "jac_times_cjmass",
@@ -513,6 +525,15 @@ def __setstate__(self, d):
         ]:
             self._setup[key] = idaklu.generate_function(self._setup[key + "_pkl"])
 
+        for key in [
+            "var_idaklu_fcns",
+            "dvar_dy_idaklu_fcns",
+            "dvar_dp_idaklu_fcns",
+        ]:
+            self._setup[key] = [
+                idaklu.generate_function(f) for f in self._setup[key + "_pkl"]
+            ]
+
         self._setup["solver_function"] = idaklu.create_casadi_solver_group
 
         self._setup["solver"] = self._setup["solver_function"](

tests/unit/test_simulation.py

@@ -736,6 +736,16 @@ def test_save_load_model(self):
 
         os.remove("sim_save.json")
 
+    def test_save_load_outvars(self, tmp_path):
+        filename = str(tmp_path / "test.pkl")
+        model = pybamm.lithium_ion.SPM()
+        solver = pybamm.IDAKLUSolver(output_variables=["Voltage [V]"])
+        sim = pybamm.Simulation(model, solver=solver)
+        sim.solve([0, 600])
+        sim.save(filename)
+        pkl_obj = pybamm.load_sim(filename)
+        assert list(pkl_obj.solver.output_variables) == ["Voltage [V]"]
+
     def test_plot(self):
         sim = pybamm.Simulation(pybamm.lithium_ion.SPM())