Repository: pybamm-team/PyBaMM (python)

Issue:
[Bug]: IDAKLU SOLVER output_variables long simulation and state issues
### PyBaMM Version

24.11.2

### Python Version

3.12.4

### Describe the bug

IDAKLU SOLVER output_variables long simulation and state issues:

Hi! I have been running ageing simulations with CASADI for a while and wanted to use IDAKLU for CC ageing or calendar ageing as it is faster. With the new option for output_variables, the memory allocation is much smaller (x10), good for parallel or threaded computing. 

Now, two issues arise. Importantly, none of these errors occur with Casadi or IDAKLU (when output_variables are not specified)


**So basically using last_state in looped simulations and using .cycles to retrieve cycle specif data does not work too well with IDAKLU when output_variables is active :)**


**1)** When using cycles and selecting output_variables, one cannot use the previous way of working:
 for example like this: simulation_IDA3_checkup.solution.cycles[-1]['Throughput capacity [A.h]'].entries
but instead must go with a work around like this: 
    **capacity_full_checkup = simulation_IDA3_checkup.solution.get_data_dict()['Throughput capacity [A.h]']
    cycles = simulation_IDA3_checkup.solution.get_data_dict()['Cycle']
    capacity_discharge_section= [d for d, c in zip(capacity_full_checkup, cycles) if c == 2]**   
I find this work around inconvenient and would prefer that what works with CASADI and IDAKLU (without output_Variables selected) works also with output_variables selected

**2)** when using cycling and for example using last states, I get an error with starting_solution
example:

    simulation_IDA3_ageing_cylic.solve(starting_solution=last_state)    
    last_state = simulation_IDA3_ageing_cylic.solution.last_state
    simulation_IDA3_checkup.solve(starting_solution=last_state)

this will give me an error if simulation_IDA3_ageing_cylic = pybamm.Simulation(... , solver=solver_IDAKLU2 , ....)
but it will not give me an error if  simulation_IDA3_ageing_cylic = pybamm.Simulation(... , solver=solver_IDAKLU1 , ....)
solver_IDAKLU1 = pybamm.IDAKLUSolver(rtol=0.00001, atol=1e-07, root_method='casadi', root_tol=1e-07, extrap_tol=None)
solver_IDAKLU2 = pybamm.IDAKLUSolver(rtol=0.00001, atol=1e-07, root_method='casadi', root_tol=1e-07, extrap_tol=None, output_variables=['Time [s]', 'Current [A]', "Voltage [V]" ], options=None)

for this last example I have a work around which is to use **starting_solution = simulation_IDA3_ageing_calendar.solution**, so using the whole solution instead of the last_state (which is I think more efficient and has a smaller memory footprint)

for 2) this means passing on a simulation object which, iteration by iteration, will grow in size. Using save_cycles in the simulation options only leads to more errors which I will not address here. all in all, I can 

Ideally, I would like to be able to activate output_variable in the IDAKLU solver, while retaining the capability of calling and using:

- starting_solution=last_state (in .solve step)
- simulation.solution.cycles[-1]['XXX'].entries (after .solve )
- save_cycles = XXX (in simulation step)

Am I getting this wrong? I.e., recongizning that it might not work just as if I swapped CASADI, are there other workarounds possibile? Let me know :) Thanks for the help 

Alessio Lodge, TNO - Imperial

### Steps to Reproduce

# Code is above. Use any experiment and run with this:


model = pybamm.lithium_ion.DFN()
solver_CasadiSolverSafe = pybamm.CasadiSolver(dt_max = 100,mode = "safe", rtol=1e-7 ,atol=1e-7)
solver_IDAKLU1 = pybamm.IDAKLUSolver(rtol=0.00001, atol=1e-07, root_method='casadi', root_tol=1e-07, extrap_tol=None)
solver_IDAKLU2 = pybamm.IDAKLUSolver(rtol=0.00001, atol=1e-07, root_method='casadi', root_tol=1e-07, extrap_tol=None, output_variables=['Time [s]', 'Current [A]', "Voltage [V]" ], options=None)
solver_IDAKLU3 =  pybamm.IDAKLUSolver(rtol=0.00001, atol=1e-07, root_method='casadi', root_tol=1e-07, extrap_tol=None, output_variables=['Time [s]', 'Current [A]', "Voltage [V]", 'Throughput capacity [A.h]', 'Discharge capacity [A.h]'], options=  {"dt_max": 1000.0})

# FOR CASADI OR IDAKLU WITHOUT OUTPUT VARS
simulation_IDA3_checkup.solution.cycles[-1]['Throughput capacity [A.h]'].entries


# FOR  IDAKLU WITH OUTPUT VARS **1)**
    capacity_full_checkup = simulation_IDA3_checkup.solution.get_data_dict()['Throughput capacity [A.h]']
    cycles = simulation_IDA3_checkup.solution.get_data_dict()['Cycle']
    capacity_discharge_section= [d for d, c in zip(capacity_full_checkup, cycles) if c == 2]   
#I find this work around inconvenient and would prefer that what works with CASADI and IDAKLU (without output_Variables selected) works also with #output_variables selected


# FOR  IDAKLU WITH OUTPUT VARS **2)**
starting_solution = simulation_IDA3_ageing_calendar.solution #


### Relevant log output

```shell
for 1) (when I make the error happen, without work around)

  File "c:\Users\XXX\Desktop\Journal pape\Code\Python_ageing_code_working_in_IDAKLU.py", line 81, in get_EOL_cap
    print(simulation_IDA3_checkup.solution.cycles[-1]['Throughput capacity [A.h]'].entries - min(simulation_IDA3_checkup.solution.cycles[-1]['Throughput capacity [A.h]'].entries)) #DOESNT WORK WITH IDAKLU, WHY???                                                                                                                                                                             WITH IDAKLU, WHY???
          ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\lodgeaa\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\solution.py", line 710, in __getitem__
    self.update(key)
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\solution.py", line 590, in update
    self._update_variable(variable)
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\solution.py", line 608, in _update_variable
    raise KeyError(
KeyError: "Cannot process variable 'Throughput capacity [A.h]' as it was not part of the solve. Please re-run the solve with `output_variables` set to include this variable."


for 2) (when I make the error happen, without work around)

Traceback (most recent call last):
  File "c:\Users\XXX\Desktop\Journal pape\Code\python_ageing_code_working_in_IDAKLU_with_output_variables", line 173, in <module>
    simulation_IDA3_checkup.solve(starting_solution=last_state)
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\simulation.py", line 864, in solve
    self._solution = self._solution + cycle_solution
                     ~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\solution.py", line 994, in __add__
    v: self._variables[v]._update(other._variables[v], new_sol)
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\processed_variable_computed.py", line 475, in _update
    new_var = self.__class__(bv, bvc, bvd, new_sol)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\processed_variable_computed.py", line 88, in __init__
    self.initialise_0D()
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\processed_variable_computed.py", line 202, in initialise_0D
    self._xr_data_array = xr.DataArray(entries, coords=[("t", self.t_pts)])
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\xarray\core\dataarray.py", line 456, in __init__
    coords, dims = _infer_coords_and_dims(data.shape, coords, dims)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\xarray\core\dataarray.py", line 195, in _infer_coords_and_dims
    _check_coords_dims(shape, new_coords, dims_tuple)
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\xarray\core\dataarray.py", line 129, in _check_coords_dims
    raise ValueError(
ValueError: conflicting sizes for dimension 't': length 143 on the data but length 144 on coordinate 't'
```

Pull request description:
Create a base processed variable class
# Description

This PR addresses an issue encountered when using `first_state` or `last_state` solution objects with the `output_variables` option in the IDAKLU solver - for example, when supplying an starting solution to .solve() within an experiment.

When adding a `first_state` or `last_state` solution to a new solution computed with the `output_variables` option, a conflict arises due to the differing types of processed variables:
* first_state/last_state solutions dynamically generate `ProcessedVariable` objects on demand.
* New solutions (e.g. from .solve()) use `ComputedProcessedVariable` objects.
These two classes currently share no functionality and cannot interoperate. Consequently, attempting to combine them to form a new solution results in failure.

This PR introduces a new abstract base class for both `ProcessedVariable` and `ComputedProcessedVariable`, providing a shared `_update()` method. This enables consistent combination of these objects when merging solutions. The concatenated object is always a `ComputedProcessedVariable`, since converting to a `ProcessedVariable` is not viable - it requires access to the full state vector, which is unavailable when `output_variables` is used.

Fixes #4743 

# Important checks:

Please confirm the following before marking the PR as ready for review:
- No style issues: `nox -s pre-commit`
- All tests pass: `nox -s tests`
- The documentation builds: `nox -s doctests`
- Code is commented for hard-to-understand areas
- Tests added that prove fix is effective or that feature works


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2ea515708..a64926a57 100644
--- a/CHANGELOG.md
+++ b/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))
diff --git a/src/pybamm/solvers/base_processed_variable.py b/src/pybamm/solvers/base_processed_variable.py
new file mode 100644
index 000000000..389847b48
--- /dev/null
+++ b/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)
diff --git a/src/pybamm/solvers/processed_variable.py b/src/pybamm/solvers/processed_variable.py
index 0685b9a52..e1b19be32 100644
--- a/src/pybamm/solvers/processed_variable.py
+++ b/src/pybamm/solvers/processed_variable.py
@@ -7,8 +7,10 @@ from pybammsolvers import idaklu
 
 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 @@ class ProcessedVariable:
     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__(
diff --git a/src/pybamm/solvers/processed_variable_computed.py b/src/pybamm/solvers/processed_variable_computed.py
index dd6e09ba4..b3a47b434 100644
--- a/src/pybamm/solvers/processed_variable_computed.py
+++ b/src/pybamm/solvers/processed_variable_computed.py
@@ -10,8 +10,10 @@ from scipy.integrate import cumulative_trapezoid
 
 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 @@ class ProcessedVariableComputed:
 
         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 @@ class ProcessedVariableComputed:
             return {}
         return self._sensitivities
 
-    def _update(
+    def _concat(
         self, other: pybamm.ProcessedVariableComputed, new_sol: pybamm.Solution
     ) -> pybamm.ProcessedVariableComputed:
         """
@@ -716,4 +721,9 @@ class ProcessedVariableComputed:
 
         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
diff --git a/src/pybamm/solvers/solution.py b/src/pybamm/solvers/solution.py
index 01f614ef9..b958932f0 100644
--- a/src/pybamm/solvers/solution.py
+++ b/src/pybamm/solvers/solution.py
@@ -801,17 +801,10 @@ class Solution:
         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
 


Test patch:
diff --git a/tests/integration/test_solvers/test_idaklu.py b/tests/integration/test_solvers/test_idaklu.py
index 14e079c4c..cf4dfafc8 100644
--- a/tests/integration/test_solvers/test_idaklu.py
+++ b/tests/integration/test_solvers/test_idaklu.py
@@ -209,3 +209,36 @@ class TestIDAKLUSolver:
             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
+        )
diff --git a/tests/unit/test_solvers/test_processed_variable.py b/tests/unit/test_solvers/test_processed_variable.py
index d14bb8988..934c0c63a 100644
--- a/tests/unit/test_solvers/test_processed_variable.py
+++ b/tests/unit/test_solvers/test_processed_variable.py
@@ -1486,3 +1486,155 @@ class TestProcessedVariable:
 
         # 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)
+        )
diff --git a/tests/unit/test_solvers/test_processed_variable_computed.py b/tests/unit/test_solvers/test_processed_variable_computed.py
index 709870393..ae49e4ff4 100644
--- a/tests/unit/test_solvers/test_processed_variable_computed.py
+++ b/tests/unit/test_solvers/test_processed_variable_computed.py
@@ -107,7 +107,7 @@ class TestProcessedVariableComputed:
         )
 
         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 @@ class TestProcessedVariableComputed:
         )
 
         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
diff --git a/tests/unit/test_solvers/test_solution.py b/tests/unit/test_solvers/test_solution.py
index 2057e5092..907a8eacb 100644
--- a/tests/unit/test_solvers/test_solution.py
+++ b/tests/unit/test_solvers/test_solution.py
@@ -195,15 +195,6 @@ class TestSolution:
         # 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)]