Repository: pybamm-team/PyBaMM (python)

Issue:
Make IDAKLU output variables compatible with Experiments
This was initially reported as a bug (see report below) but it turns out it is not an implemented feature yet. It seems the main issue is that when a new experiment step is initialised with the previous step solution, that one is empty so the relevant variables are not available. I believe that amending the `last_state` so it stores the whole state vector even when output variables are used should fix this.

### PyBaMM Version

develop

### Python Version

3.11.6

### Describe the bug

I am running a model with the `IDAKLUSolver` and some `output_variables`. These include, amongst others, the `"Discharge capacity [A.h]"`, however PyBaMM complains that that variable is not found.

### Steps to Reproduce

```python3
import pybamm

pybamm.set_logging_level("NOTICE")

model = pybamm.lithium_ion.DFN()

output_variables = ["Discharge capacity [A.h]", "Time [s]", "Current [A]", "Voltage [V]"]
solver = pybamm.IDAKLUSolver(output_variables=output_variables)

experiment = pybamm.Experiment(
    [
        (
            "Charge at 1C until 4.2 V",
            "Hold at 4.2 V until C/50",
            "Rest for 1 hour",
        )
    ]
)

sim = pybamm.Simulation(
    model,
    experiment=experiment,
    solver=solver,
)

sol = sim.solve()
```

### Relevant log output

```shell
Traceback (most recent call last):
  File "/home/brosaplanella/PyBaMM/src/pybamm/models/base_model.py", line 808, in set_initial_conditions_from
    final_state = solution[var.name]
                  ~~~~~~~~^^^^^^^^^^
  File "/home/brosaplanella/PyBaMM/src/pybamm/solvers/solution.py", line 547, in __getitem__
    self.update(key)
  File "/home/brosaplanella/PyBaMM/src/pybamm/solvers/solution.py", line 481, in update
    raise KeyError(
KeyError: "Cannot process variable 'Discharge capacity [A.h]' as it was not part of the solve. Please re-run the solve with `output_variables` set to include this variable."

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/brosaplanella/battery_modelling/degradation-data/examples/mwe.py", line 29, in <module>
    sol = sim.solve()
          ^^^^^^^^^^^
  File "/home/brosaplanella/PyBaMM/src/pybamm/simulation.py", line 702, in solve
    step_solution = solver.step(
                    ^^^^^^^^^^^^
  File "/home/brosaplanella/PyBaMM/src/pybamm/solvers/base_solver.py", line 1222, in step
    _, concatenated_initial_conditions = model.set_initial_conditions_from(
                                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/brosaplanella/PyBaMM/src/pybamm/models/base_model.py", line 810, in set_initial_conditions_from
    raise pybamm.ModelError(
pybamm.expression_tree.exceptions.ModelError: To update a model from a solution, each variable in model.initial_conditions must appear in the solution with the same key as the variable name. In the solution provided, 'Cannot process variable 'Discharge capacity [A.h]' as it was not part of the solve. Please re-run the solve with `output_variables` set to include this variable.' was not found.
```


Pull request:
Allow using IDAKLU(output_variables=...) with Experiments
# Description

Allows users to use the `output_variables` option with the IDAKLU solver within an experiment.

Fixes #4386

## Type of change

Please add a line in the relevant section of [CHANGELOG.md](https://github.com/pybamm-team/PyBaMM/blob/develop/CHANGELOG.md) to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

- [x] New feature (non-breaking change which adds functionality)

# Key checklist:

- [x] No style issues: `$ pre-commit run` (or `$ nox -s pre-commit`) (see [CONTRIBUTING.md](https://github.com/pybamm-team/PyBaMM/blob/develop/CONTRIBUTING.md#installing-and-using-pre-commit) for how to set this up to run automatically when committing locally, in just two lines of code)
- [x] All tests pass: `$ python run-tests.py --all` (or `$ nox -s tests`)
- [x] The documentation builds: `$ python run-tests.py --doctest` (or `$ nox -s doctests`)

You can run integration tests, unit tests, and doctests together at once, using `$ python run-tests.py --quick` (or `$ nox -s quick`).

## Further checks:

- [x] Code is commented, particularly in hard-to-understand areas
- [x] Tests added that prove fix is effective or that feature works


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4fbc22902..9167391c1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Features
 
+- Adds support to `pybamm.Experiment` for the `output_variables` option in the `IDAKLUSolver`. ([#4534](https://github.com/pybamm-team/PyBaMM/pull/4534))
 - Adds an option "voltage as a state" that can be "false" (default) or "true". If "true" adds an explicit algebraic equation for the voltage. ([#4507](https://github.com/pybamm-team/PyBaMM/pull/4507))
 - Improved `QuickPlot` accuracy for simulations with Hermite interpolation. ([#4483](https://github.com/pybamm-team/PyBaMM/pull/4483))
 - Added Hermite interpolation to the (`IDAKLUSolver`) that improves the accuracy and performance of post-processing variables. ([#4464](https://github.com/pybamm-team/PyBaMM/pull/4464))
diff --git a/src/pybamm/solvers/base_solver.py b/src/pybamm/solvers/base_solver.py
index a7fc79fe3..19aab6540 100644
--- a/src/pybamm/solvers/base_solver.py
+++ b/src/pybamm/solvers/base_solver.py
@@ -1452,6 +1452,7 @@ class BaseSolver:
                 solution.t_event,
                 solution.y_event,
                 solution.termination,
+                variables_returned=solution.variables_returned,
             )
             event_sol.solve_time = 0
             event_sol.integration_time = 0
diff --git a/src/pybamm/solvers/idaklu_solver.py b/src/pybamm/solvers/idaklu_solver.py
index 80eaffebf..484c1ed9b 100644
--- a/src/pybamm/solvers/idaklu_solver.py
+++ b/src/pybamm/solvers/idaklu_solver.py
@@ -863,6 +863,7 @@ class IDAKLUSolver(pybamm.BaseSolver):
             termination,
             all_sensitivities=yS_out,
             all_yps=yp,
+            variables_returned=bool(save_outputs_only),
         )
 
         newsol.integration_time = integration_time
diff --git a/src/pybamm/solvers/processed_variable_computed.py b/src/pybamm/solvers/processed_variable_computed.py
index 4f0cccc8c..befe6314b 100644
--- a/src/pybamm/solvers/processed_variable_computed.py
+++ b/src/pybamm/solvers/processed_variable_computed.py
@@ -1,6 +1,7 @@
 #
-# Processed Variable class
+# Processed Variable Computed class
 #
+from __future__ import annotations
 import casadi
 import numpy as np
 import pybamm
@@ -450,3 +451,27 @@ class ProcessedVariableComputed:
         if len(self.all_inputs[0]) == 0:
             return {}
         return self._sensitivities
+
+    def _update(
+        self, other: pybamm.ProcessedVariableComputed, new_sol: pybamm.Solution
+    ) -> pybamm.ProcessedVariableComputed:
+        """
+        Returns a new ProcessedVariableComputed object that is the result of appending
+        the data from other to this object. Used exclusively in running experiments, to
+        append the data from one cycle to the next.
+
+        Parameters
+        ----------
+        other : :class:`pybamm.ProcessedVariableComputed`
+            The other ProcessedVariableComputed object to append to this one
+        new_sol : :class:`pybamm.Solution`
+            The new solution object to be used to create the processed variables
+        """
+
+        bv = self.base_variables + other.base_variables
+        bvc = self.base_variables_casadi + other.base_variables_casadi
+        bvd = self.base_variables_data + other.base_variables_data
+
+        new_var = self.__class__(bv, bvc, bvd, new_sol)
+
+        return new_var
diff --git a/src/pybamm/solvers/solution.py b/src/pybamm/solvers/solution.py
index 7410c2161..256d596fd 100644
--- a/src/pybamm/solvers/solution.py
+++ b/src/pybamm/solvers/solution.py
@@ -62,6 +62,10 @@ class Solution:
         True if sensitivities included as the solution of the explicit forwards
         equations.  False if no sensitivities included/wanted. Dict if sensitivities are
         provided as a dict of {parameter: [sensitivities]} pairs.
+    variables_returned: bool
+        Bool to indicate if `all_ys` contains the full state vector, or is empty because
+        only requested variables have been returned. True if `output_variables` is used
+        with a solver, otherwise False.
 
     """
 
@@ -76,6 +80,7 @@ class Solution:
         termination="final time",
         all_sensitivities=False,
         all_yps=None,
+        variables_returned=False,
         check_solution=True,
     ):
         if not isinstance(all_ts, list):
@@ -93,6 +98,8 @@ class Solution:
             all_yps = [all_yps]
         self._all_yps = all_yps
 
+        self.variables_returned = variables_returned
+
         # Set up inputs
         if not isinstance(all_inputs, list):
             all_inputs_copy = dict(all_inputs)
@@ -460,9 +467,15 @@ class Solution:
         else:
             all_yps = self.all_yps[0][:, :1]
 
+        if not self.variables_returned:
+            all_ys = self.all_ys[0][:, :1]
+        else:
+            # Get first state from initial conditions as all_ys is empty
+            all_ys = self.all_models[0].y0full.reshape(-1, 1)
+
         new_sol = Solution(
             self.all_ts[0][:1],
-            self.all_ys[0][:, :1],
+            all_ys,
             self.all_models[:1],
             self.all_inputs[:1],
             None,
@@ -500,9 +513,15 @@ class Solution:
         else:
             all_yps = self.all_yps[-1][:, -1:]
 
+        if not self.variables_returned:
+            all_ys = self.all_ys[-1][:, -1:]
+        else:
+            # Get last state from y_event as all_ys is empty
+            all_ys = self.y_event.reshape(len(self.y_event), 1)
+
         new_sol = Solution(
             self.all_ts[-1][-1:],
-            self.all_ys[-1][:, -1:],
+            all_ys,
             self.all_models[-1:],
             self.all_inputs[-1:],
             self.t_event,
@@ -580,15 +599,11 @@ class Solution:
         # Iterate through all models, some may be in the list several times and
         # therefore only get set up once
         vars_casadi = []
-        for i, (model, ts, ys, inputs, var_pybamm) in enumerate(
-            zip(self.all_models, self.all_ts, self.all_ys, self.all_inputs, vars_pybamm)
+        for i, (model, ys, inputs, var_pybamm) in enumerate(
+            zip(self.all_models, self.all_ys, self.all_inputs, vars_pybamm)
         ):
-            if (
-                ys.size == 0
-                and var_pybamm.has_symbol_of_classes(
-                    pybamm.expression_tree.state_vector.StateVector
-                )
-                and not ts.size == 0
+            if self.variables_returned and var_pybamm.has_symbol_of_classes(
+                pybamm.expression_tree.state_vector.StateVector
             ):
                 raise KeyError(
                     f"Cannot process variable '{variable}' as it was not part of the "
@@ -682,7 +697,7 @@ class Solution:
 
         Returns
         -------
-        :class:`pybamm.ProcessedVariable`
+        :class:`pybamm.ProcessedVariable` or :class:`pybamm.ProcessedVariableComputed`
             A variable that can be evaluated at any time or spatial point. The
             underlying data for this variable is available in its attribute ".data"
         """
@@ -950,6 +965,7 @@ class Solution:
             other.termination,
             all_sensitivities=all_sensitivities,
             all_yps=all_yps,
+            variables_returned=other.variables_returned,
         )
 
         new_sol.closest_event_idx = other.closest_event_idx
@@ -966,6 +982,19 @@ class Solution:
         # Set sub_solutions
         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()
+                }
+
         return new_sol
 
     def __radd__(self, other):
@@ -983,6 +1012,7 @@ class Solution:
             self.termination,
             self._all_sensitivities,
             self.all_yps,
+            self.variables_returned,
         )
         new_sol._all_inputs_casadi = self.all_inputs_casadi
         new_sol._sub_solutions = self.sub_solutions
@@ -992,6 +1022,13 @@ class Solution:
         new_sol.integration_time = self.integration_time
         new_sol.set_up_time = self.set_up_time
 
+        # copy over variables which were derived at the solver stage
+        if self._variables and all(
+            isinstance(v, pybamm.ProcessedVariableComputed)
+            for v in self._variables.values()
+        ):
+            new_sol._variables = self._variables.copy()
+
         return new_sol
 
     def plot_voltage_components(
@@ -1094,6 +1131,7 @@ def make_cycle_solution(
         sum_sols.termination,
         sum_sols._all_sensitivities,
         sum_sols.all_yps,
+        sum_sols.variables_returned,
     )
     cycle_solution._all_inputs_casadi = sum_sols.all_inputs_casadi
     cycle_solution._sub_solutions = sum_sols.sub_solutions


Test patch:
diff --git a/tests/integration/test_solvers/test_idaklu.py b/tests/integration/test_solvers/test_idaklu.py
index 88faa80dd..3ee96a9cc 100644
--- a/tests/integration/test_solvers/test_idaklu.py
+++ b/tests/integration/test_solvers/test_idaklu.py
@@ -152,3 +152,55 @@ class TestIDAKLUSolver:
             # test that y[1:3] = to true solution
             true_solution = b_value * sol.t
             np.testing.assert_array_almost_equal(sol.y[1:3], true_solution)
+
+    def test_with_experiments(self):
+        summary_vars = []
+        sols = []
+        for out_vars in [True, False]:
+            model = pybamm.lithium_ion.SPM()
+
+            if out_vars:
+                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
+                ]
+            else:
+                output_variables = None
+
+            solver = pybamm.IDAKLUSolver(output_variables=output_variables)
+
+            experiment = pybamm.Experiment(
+                [
+                    (
+                        "Charge at 1C until 4.2 V",
+                        "Hold at 4.2 V until C/50",
+                        "Rest for 1 hour",
+                    )
+                ]
+            )
+
+            sim = pybamm.Simulation(
+                model,
+                experiment=experiment,
+                solver=solver,
+            )
+
+            sol = sim.solve()
+            sols.append(sol)
+            summary_vars.append(sol.summary_variables)
+
+        # check computed variables are propegated sucessfully
+        np.testing.assert_array_equal(
+            sols[0]["Pressure [Pa]"].data, sols[1]["Pressure [Pa]"].data
+        )
+        np.testing.assert_array_almost_equal(
+            sols[0]["Voltage [V]"].data, sols[1]["Voltage [V]"].data
+        )
+
+        # check summary variables are the same if output variables are specified
+        for var in summary_vars[0].keys():
+            assert summary_vars[0][var] == summary_vars[1][var]
diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py
index 39918d73a..e4d6559e7 100644
--- a/tests/unit/test_solvers/test_idaklu_solver.py
+++ b/tests/unit/test_solvers/test_idaklu_solver.py
@@ -939,6 +939,9 @@ class TestIDAKLUSolver:
             with pytest.raises(KeyError):
                 sol[varname].data
 
+        # Check Solution is marked
+        assert sol.variables_returned is True
+
         # Mock a 1D current collector and initialise (none in the model)
         sol["x_s [m]"].domain = ["current collector"]
         sol["x_s [m]"].entries
diff --git a/tests/unit/test_solvers/test_processed_variable_computed.py b/tests/unit/test_solvers/test_processed_variable_computed.py
index 0fa46b441..f28c053fd 100644
--- a/tests/unit/test_solvers/test_processed_variable_computed.py
+++ b/tests/unit/test_solvers/test_processed_variable_computed.py
@@ -1,7 +1,7 @@
 #
 # Tests for the Processed Variable Computed class
 #
-# This class forms a container for variables (and sensitivities) calculted
+# This class forms a container for variables (and sensitivities) calculated
 #  by the idaklu solver, and does not possesses any capability to calculate
 #  values itself since it does not have access to the full state vector
 #
@@ -76,11 +76,12 @@ class TestProcessedVariableComputed:
         t_sol = np.array([0])
         y_sol = np.array([1])[:, np.newaxis]
         var_casadi = to_casadi(var, y_sol)
+        sol = pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {})
         processed_var = pybamm.ProcessedVariableComputed(
             [var],
             [var_casadi],
             [y_sol],
-            pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}),
+            sol,
         )
         # Assert that the processed variable is the same as the solution
         np.testing.assert_array_equal(processed_var.entries, y_sol[0])
@@ -94,6 +95,22 @@ class TestProcessedVariableComputed:
         processed_var.cumtrapz_ic = 1
         processed_var.entries
 
+        # check _update
+        t_sol2 = np.array([1])
+        y_sol2 = np.array([2])[:, np.newaxis]
+        var_casadi = to_casadi(var, y_sol2)
+        sol_2 = pybamm.Solution(t_sol2, y_sol2, pybamm.BaseModel(), {})
+        processed_var2 = pybamm.ProcessedVariableComputed(
+            [var],
+            [var_casadi],
+            [y_sol2],
+            sol_2,
+        )
+
+        comb_sol = sol + sol_2
+        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
     def test_processed_variable_0D_no_sensitivity(self):
         # without space
@@ -217,6 +234,60 @@ class TestProcessedVariableComputed:
         c_casadi = to_casadi(c, y_sol)
         pybamm.ProcessedVariableComputed([c], [c_casadi], [y_sol], solution)
 
+    def test_processed_variable_1D_update(self):
+        # variable 1
+        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_sol1 = disc.process_symbol(x).entries[:, 0]
+        var_sol1 = disc.process_symbol(var)
+        t_sol1 = np.linspace(0, 1)
+        y_sol1 = np.ones_like(x_sol1)[:, np.newaxis] * np.linspace(0, 5)
+
+        var_casadi1 = to_casadi(var_sol1, y_sol1)
+        sol1 = pybamm.Solution(t_sol1, y_sol1, pybamm.BaseModel(), {})
+        processed_var1 = pybamm.ProcessedVariableComputed(
+            [var_sol1],
+            [var_casadi1],
+            [y_sol1],
+            sol1,
+        )
+
+        # variable 2 -------------------
+        var2 = pybamm.Variable("var2", domain=["negative electrode", "separator"])
+        z = pybamm.SpatialVariable("z", domain=["negative electrode", "separator"])
+
+        disc = tests.get_discretisation_for_testing()
+        disc.set_variable_slices([var2])
+        z_sol2 = disc.process_symbol(z).entries[:, 0]
+        var_sol2 = disc.process_symbol(var2)
+        t_sol2 = np.linspace(2, 3)
+        y_sol2 = np.ones_like(z_sol2)[:, np.newaxis] * np.linspace(5, 1)
+
+        var_casadi2 = to_casadi(var_sol2, y_sol2)
+        sol2 = pybamm.Solution(t_sol2, y_sol2, pybamm.BaseModel(), {})
+        var_2 = pybamm.ProcessedVariableComputed(
+            [var_sol2],
+            [var_casadi2],
+            [y_sol2],
+            sol2,
+        )
+
+        comb_sol = sol1 + sol2
+        comb_var = processed_var1._update(var_2, comb_sol)
+
+        # Ordering from idaklu with output_variables set is different to
+        # the full solver
+        y_sol1 = y_sol1.reshape((y_sol1.shape[1], y_sol1.shape[0])).transpose()
+        y_sol2 = y_sol2.reshape((y_sol2.shape[1], y_sol2.shape[0])).transpose()
+
+        np.testing.assert_array_equal(
+            comb_var.entries, np.concatenate((y_sol1, y_sol2), axis=1)
+        )
+        np.testing.assert_array_equal(comb_var.entries, comb_var.data)
+
     def test_processed_variable_2D_x_r(self):
         var = pybamm.Variable(
             "var",
diff --git a/tests/unit/test_solvers/test_solution.py b/tests/unit/test_solvers/test_solution.py
index 2fb25f79a..1ee652d7a 100644
--- a/tests/unit/test_solvers/test_solution.py
+++ b/tests/unit/test_solvers/test_solution.py
@@ -174,6 +174,50 @@ class TestSolution:
         ):
             sol_sum.y
 
+    @pytest.mark.skipif(
+        not pybamm.has_idaklu(), reason="idaklu solver is not installed"
+    )
+    def test_add_solutions_with_computed_variables(self):
+        model = pybamm.BaseModel()
+        u = pybamm.Variable("u")
+        v = pybamm.Variable("v")
+        model.rhs = {u: 1 * v}
+        model.algebraic = {v: 1 - v}
+        model.initial_conditions = {u: 0, v: 1}
+        model.variables = {"2u": 2 * u}
+
+        disc = pybamm.Discretisation()
+        disc.process_model(model)
+
+        # Set up first solution
+        t1 = np.linspace(0, 1, 50)
+        solver = pybamm.IDAKLUSolver(output_variables=["2u"])
+
+        sol1 = solver.solve(model, t1)
+
+        # second solution
+        t2 = np.linspace(2, 3, 50)
+        sol2 = solver.solve(model, t2)
+
+        sol_sum = sol1 + sol2
+
+        # check varaibles concat appropriately
+        assert sol_sum["2u"].data[0] == sol1["2u"].data[0]
+        assert sol_sum["2u"].data[-1] == sol2["2u"].data[-1]
+        # Check functions still work
+        sol_sum["2u"].unroll()
+        # 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)]
@@ -194,6 +238,34 @@ class TestSolution:
         assert sol_copy.solve_time == sol1.solve_time
         assert sol_copy.integration_time == sol1.integration_time
 
+    @pytest.mark.skipif(
+        not pybamm.has_idaklu(), reason="idaklu solver is not installed"
+    )
+    def test_copy_with_computed_variables(self):
+        model = pybamm.BaseModel()
+        u = pybamm.Variable("u")
+        v = pybamm.Variable("v")
+        model.rhs = {u: 1 * v}
+        model.algebraic = {v: 1 - v}
+        model.initial_conditions = {u: 0, v: 1}
+        model.variables = {"2u": 2 * u}
+
+        disc = pybamm.Discretisation()
+        disc.process_model(model)
+
+        # Set up first solution
+        t1 = np.linspace(0, 1, 50)
+        solver = pybamm.IDAKLUSolver(output_variables=["2u"])
+
+        sol1 = solver.solve(model, t1)
+
+        sol2 = sol1.copy()
+
+        assert (
+            sol1._variables[k] == sol2._variables[k] for k in sol1._variables.keys()
+        )
+        assert sol2.variables_returned is True
+
     def test_last_state(self):
         # Set up first solution
         t1 = [np.linspace(0, 1), np.linspace(1, 2, 5)]
@@ -214,6 +286,39 @@ class TestSolution:
         assert sol_last_state.solve_time == 0
         assert sol_last_state.integration_time == 0
 
+    @pytest.mark.skipif(
+        not pybamm.has_idaklu(), reason="idaklu solver is not installed"
+    )
+    def test_first_last_state_empty_y(self):
+        # check that first and last state work when y is empty
+        # due to only variables being returned (required for experiments)
+        model = pybamm.BaseModel()
+        u = pybamm.Variable("u")
+        v = pybamm.Variable("v")
+        model.rhs = {u: 1 * v}
+        model.algebraic = {v: 1 - v}
+        model.initial_conditions = {u: 0, v: 1}
+        model.variables = {"2u": 2 * u, "4u": 4 * u}
+        model._summary_variables = {"4u": model.variables["4u"]}
+
+        disc = pybamm.Discretisation()
+        disc.process_model(model)
+
+        # Set up first solution
+        t1 = np.linspace(0, 1, 50)
+        solver = pybamm.IDAKLUSolver(output_variables=["2u"])
+
+        sol1 = solver.solve(model, t1)
+
+        np.testing.assert_array_equal(
+            sol1.first_state.all_ys[0], np.array([[0.0], [1.0]])
+        )
+        # check summay variables not in the solve can be evaluated at the final timestep
+        # via 'last_state
+        np.testing.assert_array_almost_equal(
+            sol1.last_state["4u"].entries, np.array([4.0])
+        )
+
     def test_cycles(self):
         model = pybamm.lithium_ion.SPM()
         experiment = pybamm.Experiment(