Create a base processed variable class (pybamm-team#5076)

* Creates an abstract base class for variables and can convert a processed to computed variable.

* sort imports

* Propegate sensitivities in PVC

* style: pre-commit fixes

* working version with some tests

* Use entries instead, add integration test

* style: pre-commit fixes

* move stub solution

* skip testing line & add note

* Edit changelog

* make _update public

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Martin Robinson <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- Creates `BaseProcessedVariable` to enable object combination when adding solutions together ([#5076](https://github.com/pybamm-team/PyBaMM/pull/5076))
 - Added a `Constant` symbol for named constants. This is a subclass of `Scalar` and is used to represent named constants such as the gas constant. This avoids constants being simplified out when constructing expressions. ([#5070](https://github.com/pybamm-team/PyBaMM/pull/5070))
 - Generalise `pybamm.DiscreteTimeSum` to allow it to be embedded in other expressions ([#5044](https://github.com/pybamm-team/PyBaMM/pull/5044))
 - Adds `all` key-value pair to `output_variables` sensitivity dictionaries, accessible through `solution[var].sensitivities['all']`. Aligns shape with conventional solution sensitivities object. ([#5067](https://github.com/pybamm-team/PyBaMM/pull/5067))

src/pybamm/solvers/base_processed_variable.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import abc
+
+import pybamm
+
+
+class BaseProcessedVariable(abc.ABC):
+    """
+    Shared API for both 'on-demand' and 'computed' processed variables.
+    """
+
+    @abc.abstractmethod
+    def __call__(self, **coords): ...
+
+    @abc.abstractmethod
+    def as_computed(self) -> pybamm.ProcessedVariableComputed: ...
+
+    def update(
+        self,
+        other: BaseProcessedVariable,
+        new_sol: pybamm.Solution,
+    ) -> pybamm.ProcessedVariableComputed:
+        """
+        Return a *new* CPV that is `self` followed by `other`.
+        Works no matter which concrete types we start with.
+        """
+        return self.as_computed()._concat(other.as_computed(), new_sol)

src/pybamm/solvers/processed_variable.py

@@ -7,8 +7,10 @@
 
 import pybamm
 
+from .base_processed_variable import BaseProcessedVariable
 
-class ProcessedVariable:
+
+class ProcessedVariable(BaseProcessedVariable):
     """
     An object that can be evaluated at arbitrary (scalars or vectors) t and x, and
     returns the (interpolated) value of the base variable at that t and x.
@@ -494,6 +496,58 @@ def _is_discrete_time_method(self):
     def hermite_interpolation(self):
         return self.all_yps is not None
 
+    def as_computed(self):
+        """
+        Allows a ProcessedVariable to be converted to a ComputedProcessedVariable for
+        use together, e.g. when using a last state solution with a new simulation running
+        with output variables in the solver.
+        """
+
+        def _stub_solution(self):
+            """
+            Return a lightweight object that looks like the parts of
+            `pybamm.Solution` required by ProcessedVariableComputed, but without
+            keeping the full state vector in memory.
+            """
+
+            class StubSolution:
+                def __init__(self, ts, ys, inputs, inputs_casadi, sensitivities, t_pts):
+                    self.all_ts = ts
+                    self.all_ys = ys
+                    self.all_inputs = inputs
+                    self.all_inputs_casadi = inputs_casadi
+                    self.sensitivities = sensitivities
+                    self.t = t_pts
+
+            return StubSolution(
+                self.all_ts,
+                self.all_ys,
+                self.all_inputs,
+                self.all_inputs_casadi,
+                self.sensitivities,
+                self.t_pts,
+            )
+
+        entries = self.entries  # shape: (..., n_t)
+
+        # Move time to axis 0, then flatten spatial dims per timestep
+        reshaped = np.moveaxis(entries, -1, 0)  # shape: (n_t, ...)
+        base_data = [reshaped.reshape(reshaped.shape[0], -1)]  # (n_t, n_vars)
+
+        cpv = pybamm.ProcessedVariableComputed(
+            self.base_variables,
+            self.base_variables_casadi,
+            base_data,
+            _stub_solution(self),
+        )
+
+        # add sensitivities if they exist
+        if self.sensitivities:
+            # TODO: test once #5058 is fixed
+            cpv._sensitivities = self.sensitivities  # pragma: no cover
+
+        return cpv
+
 
 class ProcessedVariable0D(ProcessedVariable):
     def __init__(

src/pybamm/solvers/processed_variable_computed.py

@@ -10,8 +10,10 @@
 
 import pybamm
 
+from .base_processed_variable import BaseProcessedVariable
 
-class ProcessedVariableComputed:
+
+class ProcessedVariableComputed(BaseProcessedVariable):
     """
     An object that can be evaluated at arbitrary (scalars or vectors) t and x, and
     returns the (interpolated) value of the base variable at that t and x.
@@ -133,6 +135,9 @@ def __init__(
 
         raise NotImplementedError(f"Shape not recognized for {base_variables[0]}")
 
+    def as_computed(self) -> ProcessedVariableComputed:
+        return self
+
     def add_sensitivity(self, param, data):
         # unroll from sparse representation into n-d matrix
         # then flatten for consistency with full state-vector
@@ -694,7 +699,7 @@ def sensitivities(self):
             return {}
         return self._sensitivities
 
-    def _update(
+    def _concat(
         self, other: pybamm.ProcessedVariableComputed, new_sol: pybamm.Solution
     ) -> pybamm.ProcessedVariableComputed:
         """
@@ -716,4 +721,9 @@ def _update(
 
         new_var = self.__class__(bv, bvc, bvd, new_sol)
 
+        new_var._sensitivities = {
+            k: np.concatenate((self.sensitivities[k], other.sensitivities[k]))
+            for k in self.sensitivities.keys()
+        }
+
         return new_var

src/pybamm/solvers/solution.py

@@ -801,17 +801,10 @@ def __add__(self, other):
         new_sol._sub_solutions = self.sub_solutions + other.sub_solutions
 
         # update variables which were derived at the solver stage
-        if other._variables and all(
-            isinstance(v, pybamm.ProcessedVariableComputed)
-            for v in other._variables.values()
-        ):
-            if not self._variables:
-                new_sol._variables = other._variables.copy()
-            else:
-                new_sol._variables = {
-                    v: self._variables[v]._update(other._variables[v], new_sol)
-                    for v in self._variables.keys()
-                }
+        if any([self.variables_returned, other.variables_returned]):
+            vars = {*self._variables.keys(), *other._variables.keys()}
+            new_sol._variables = {v: self[v].update(other[v], new_sol) for v in vars}
+            new_sol.variables_returned = True
 
         return new_sol
 

tests/integration/test_solvers/test_idaklu.py

@@ -209,3 +209,36 @@ def test_with_experiments(self):
             sols[0].cycles[-1]["Current [A]"].data,
             sols[1].cycles[-1]["Current [A]"].data,
         )
+
+    def test_outvars_with_experiments_multi_simulation(self):
+        model = pybamm.lithium_ion.SPM()
+
+        experiment = pybamm.Experiment(
+            [
+                "Discharge at C/2 for 10 minutes",
+                "Rest for 10 minutes",
+            ]
+        )
+
+        solver = pybamm.IDAKLUSolver(
+            output_variables=[
+                "Discharge capacity [A.h]",  # 0D variables
+                "Time [s]",
+                "Current [A]",
+                "Voltage [V]",
+                "Pressure [Pa]",  # 1D variable
+                "Positive particle effective diffusivity [m2.s-1]",  # 2D variable
+            ]
+        )
+
+        sim = pybamm.Simulation(model, experiment=experiment, solver=solver)
+        solution = sim.solve()
+
+        sim2 = pybamm.Simulation(model, experiment=experiment, solver=solver)
+
+        new_sol1 = sim2.solve(starting_solution=solution.copy())
+        new_sol2 = sim2.solve(starting_solution=solution.copy().last_state)
+
+        np.testing.assert_array_equal(
+            new_sol1["Voltage [V]"].entries[63:], new_sol2["Voltage [V]"].entries
+        )

tests/unit/test_solvers/test_processed_variable.py

@@ -1486,3 +1486,155 @@ def spl(t):
 
         # Check that the unsorted and sorted arrays are the same
         assert np.all(y_unsorted == y_sorted[idxs_unsort])
+
+    def test_as_computed_0D(self):
+        # 0D
+        t = pybamm.t
+        y = pybamm.StateVector(slice(0, 1))
+        var = t * y
+        model = pybamm.BaseModel()
+        t_sol = np.linspace(0, 1)
+        y_sol = np.array([np.linspace(0, 5)])
+        yp_sol = self._get_yps(y_sol, False)
+        var_casadi = to_casadi(var, y_sol)
+        processed_var = pybamm.process_variable(
+            "test",
+            [var],
+            [var_casadi],
+            self._sol_default(t_sol, y_sol, yp_sol, model),
+        )
+        computed_var = processed_var.as_computed()
+
+        np.testing.assert_array_equal(computed_var.entries, t_sol * y_sol[0])
+
+    def test_as_computed_1D(self):
+        var = pybamm.Variable("var", domain=["negative electrode", "separator"])
+        x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"])
+
+        disc = tests.get_discretisation_for_testing()
+        disc.set_variable_slices([var])
+        x_sol = disc.process_symbol(x).entries[:, 0]
+        var_sol = disc.process_symbol(var)
+        t_sol = np.linspace(0, 1)
+        y_sol = x_sol[:, np.newaxis] * (5 * t_sol)
+        yp_sol = self._get_yps(y_sol, False, values=5)
+
+        var_casadi = to_casadi(var_sol, y_sol)
+        processed_var = pybamm.process_variable(
+            "test",
+            [var_sol],
+            [var_casadi],
+            self._sol_default(t_sol, y_sol, yp_sol),
+        )
+
+        computed_var = processed_var.as_computed()
+
+        # 2 vectors
+        np.testing.assert_allclose(
+            computed_var(t_sol, x_sol), y_sol, rtol=1e-7, atol=1e-6
+        )
+        # 1 vector, 1 scalar
+        np.testing.assert_allclose(
+            computed_var(0.5, x_sol), 2.5 * x_sol, rtol=1e-7, atol=1e-6
+        )
+        np.testing.assert_allclose(
+            computed_var(t_sol, x_sol[-1]),
+            x_sol[-1] * np.linspace(0, 5),
+            rtol=1e-7,
+            atol=1e-6,
+        )
+        # 2 scalars
+        np.testing.assert_allclose(
+            computed_var(0.5, x_sol[-1]), 2.5 * x_sol[-1], rtol=1e-7, atol=1e-6
+        )
+
+    def test_as_computed_2D(self):
+        var = pybamm.Variable(
+            "var",
+            domain=["negative particle"],
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
+        x = pybamm.SpatialVariable("x", domain=["negative electrode"])
+        r = pybamm.SpatialVariable(
+            "r",
+            domain=["negative particle"],
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
+
+        disc = tests.get_p2d_discretisation_for_testing()
+        disc.set_variable_slices([var])
+        x_sol = disc.process_symbol(x).entries[:, 0]
+        r_sol = disc.process_symbol(r).entries[:, 0]
+        # Keep only the first iteration of entries
+        r_sol = r_sol[: len(r_sol) // len(x_sol)]
+        var_sol = disc.process_symbol(var)
+        t_sol = np.linspace(0, 1)
+        y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis] * np.linspace(0, 5)
+        yp_sol = self._get_yps(y_sol, False)
+
+        var_casadi = to_casadi(var_sol, y_sol)
+        processed_var = pybamm.process_variable(
+            "test",
+            [var_sol],
+            [var_casadi],
+            self._sol_default(t_sol, y_sol, yp_sol),
+        )
+
+        computed_var = processed_var.as_computed()
+        # 3 vectors
+        np.testing.assert_array_equal(
+            computed_var(t_sol, x_sol, r_sol).shape, (10, 40, 50)
+        )
+        np.testing.assert_allclose(
+            computed_var(t_sol, x_sol, r_sol),
+            np.reshape(y_sol, [len(r_sol), len(x_sol), len(t_sol)]),
+            rtol=1e-7,
+            atol=1e-6,
+        )
+        # 2 vectors, 1 scalar
+        np.testing.assert_array_equal(computed_var(0.5, x_sol, r_sol).shape, (10, 40))
+        np.testing.assert_array_equal(computed_var(t_sol, 0.2, r_sol).shape, (10, 50))
+        np.testing.assert_array_equal(computed_var(t_sol, x_sol, 0.5).shape, (40, 50))
+        # 1 vectors, 2 scalar
+        np.testing.assert_array_equal(computed_var(0.5, 0.2, r_sol).shape, (10,))
+        np.testing.assert_array_equal(computed_var(0.5, x_sol, 0.5).shape, (40,))
+        np.testing.assert_array_equal(computed_var(t_sol, 0.2, 0.5).shape, (50,))
+        # 3 scalars
+        np.testing.assert_array_equal(computed_var(0.2, 0.2, 0.2).shape, ())
+
+    def test_as_computed_3D(self):
+        disc = tests.get_size_distribution_disc_for_testing(xpts=5, rpts=6, Rpts=7)
+        var = pybamm.Variable(
+            "var",
+            domain=["negative particle"],
+            auxiliary_domains={
+                "secondary": ["negative particle size"],
+                "tertiary": ["negative electrode"],
+            },
+        )
+        x_sol = disc.mesh["negative electrode"].nodes
+        disc.set_variable_slices([var])
+
+        Nx = len(x_sol)
+        R_sol = disc.mesh["negative particle size"].nodes
+        r_sol = disc.mesh["negative particle"].nodes
+        var_sol = disc.process_symbol(var)
+        t_sol = np.linspace(0, 1)
+        y_sol = np.ones(len(x_sol) * len(R_sol) * len(r_sol))[:, np.newaxis] * t_sol
+        yp_sol = self._get_yps(y_sol, False)
+
+        var_casadi = to_casadi(var_sol, y_sol)
+        geometry_options = {"options": {"particle size": "distribution"}}
+        model = tests.get_base_model_with_battery_geometry(**geometry_options)
+        processed_var = pybamm.process_variable(
+            "test",
+            [var_sol],
+            [var_casadi],
+            self._sol_default(t_sol, y_sol, yp_sol, model),
+        )
+        computed_var = processed_var.as_computed()
+
+        # 4 vectors
+        np.testing.assert_array_equal(
+            computed_var(t=t_sol, x=x_sol, R=R_sol, r=r_sol).shape, (6, 7, Nx, 50)
+        )

tests/unit/test_solvers/test_processed_variable_computed.py

@@ -107,7 +107,7 @@ def test_processed_variable_0D(self):
         )
 
         comb_sol = sol + sol_2
-        comb_var = processed_var._update(processed_var2, comb_sol)
+        comb_var = processed_var.update(processed_var2, comb_sol)
         np.testing.assert_array_equal(comb_var.entries, np.append(y_sol, y_sol2))
 
     # check empty sensitivity works
@@ -275,7 +275,7 @@ def test_processed_variable_1D_update(self):
         )
 
         comb_sol = sol1 + sol2
-        comb_var = processed_var1._update(var_2, comb_sol)
+        comb_var = processed_var1.update(var_2, comb_sol)
 
         # Ordering from idaklu with output_variables set is different to
         # the full solver

tests/unit/test_solvers/test_solution.py

@@ -195,15 +195,6 @@ def test_add_solutions_with_computed_variables(self):
         # check solution still tagged as 'variables_returned'
         assert sol_sum.variables_returned is True
 
-        # add a solution with computed variable to an empty solution
-        empty_sol = pybamm.Solution(
-            sol1.all_ts, sol1["2u"].base_variables_data, model, {u: 0, v: 1}
-        )
-
-        sol4 = empty_sol + sol2
-        assert sol4["2u"] == sol2["2u"]
-        assert sol4.variables_returned is True
-
     def test_copy(self):
         # Set up first solution
         t1 = [np.linspace(0, 1), np.linspace(1, 2, 5)]