Repository: pybamm-team/PyBaMM (python)

Issue:
[Bug]: pybamm.DiscreteTimeSum sensitivities not being summed over time
### PyBaMM Version

develop

### Python Version

3.10

### Describe the bug

When calculating the sensitivities for the `pybamm.DiscreteTimeSum`, the sum over time isn't being performed, giving a sensitivity vector of the same dimension as the number of time points. This is unsuprising since this is not currently being tested either, so tests and a fix are needed

### Steps to Reproduce

create a model with a variable using `pybamm.DiscreteTimeSum`, then calculate the sensitivities and check the dimension of the result

### Relevant log output

```shell

```

Pull request description:
bug: fix sensitivity calculation for discrete time sum
# Description

Fixes #5007

## 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 #)

# 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 ca36038a9..06e1858e4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@
 
 - Fixed a bug in the `QuickPlot` which would return empty values for 1D variables at the beginning and end of a timespan. ([#4991](https://github.com/pybamm-team/PyBaMM/pull/4991))
 - Fixed a bug in the `Exponential1DSubMesh` where the mesh was not being created correctly for non-zero minimum values. ([#4989](https://github.com/pybamm-team/PyBaMM/pull/4989))
+- Fixed sensitivity calculation for `pybamm.DiscreteTimeSum`. ([#5007](https://github.com/pybamm-team/PyBaMM/pull/5008))
 
 ## Breaking changes
 
diff --git a/src/pybamm/solvers/processed_variable.py b/src/pybamm/solvers/processed_variable.py
index b49a5fec6..82cb251bd 100644
--- a/src/pybamm/solvers/processed_variable.py
+++ b/src/pybamm/solvers/processed_variable.py
@@ -15,6 +15,8 @@ class ProcessedVariable:
 
     Parameters
     ----------
+    name : str
+        The name of the variable
     base_variables : list of :class:`pybamm.Symbol`
         A list of base variables with a method `evaluate(t,y)`, each entry of which
         returns the value of that variable for that particular sub-solution.
@@ -34,11 +36,13 @@ class ProcessedVariable:
 
     def __init__(
         self,
+        name: str,
         base_variables,
         base_variables_casadi,
         solution,
         time_integral: Optional[pybamm.ProcessedVariableTimeIntegral] = None,
     ):
+        self._name = name
         self.base_variables = base_variables
         self.base_variables_casadi = base_variables_casadi
 
@@ -458,6 +462,26 @@ class ProcessedVariable:
 
             # Compute sensitivity
             S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval
+
+            # post fix for discrete time integral won't give correct result
+            # if ts are not equal to the discrete times. Raise error
+            # in this case
+            if (
+                self.time_integral is not None
+                and self.time_integral.method == "discrete"
+            ):
+                if not (
+                    len(ts) == len(self.time_integral.discrete_times)
+                    and np.allclose(ts, self.time_integral.discrete_times, atol=1e-10)
+                ):
+                    raise pybamm.SolverError(
+                        f'Processing discrete-time-sum variable "{self._name}": solution times '
+                        "and discrete times of the time integral are not equal. Set 't_interp=discrete_sum_times' to "
+                        f"ensure the correct times are used.\nSolution times: {ts}\nDiscrete Sum times: {self.time_integral.discrete_times}"
+                    )
+
+            S_var = self._observe_postfix(S_var, ts)
+
             all_S_var.append(S_var)
 
         S_var = casadi.vertcat(*all_S_var)
@@ -481,6 +505,7 @@ class ProcessedVariable:
 class ProcessedVariable0D(ProcessedVariable):
     def __init__(
         self,
+        name: str,
         base_variables,
         base_variables_casadi,
         solution,
@@ -488,6 +513,7 @@ class ProcessedVariable0D(ProcessedVariable):
     ):
         self.dimensions = 0
         super().__init__(
+            name,
             base_variables,
             base_variables_casadi,
             solution,
@@ -526,6 +552,8 @@ class ProcessedVariable1D(ProcessedVariable):
 
     Parameters
     ----------
+    variable : str
+        The name of the variable
     base_variables : list of :class:`pybamm.Symbol`
         A list of base variables with a method `evaluate(t,y)`, each entry of which
         returns the value of that variable for that particular sub-solution.
@@ -543,6 +571,7 @@ class ProcessedVariable1D(ProcessedVariable):
 
     def __init__(
         self,
+        name: str,
         base_variables,
         base_variables_casadi,
         solution,
@@ -550,6 +579,7 @@ class ProcessedVariable1D(ProcessedVariable):
     ):
         self.dimensions = 1
         super().__init__(
+            name,
             base_variables,
             base_variables_casadi,
             solution,
@@ -607,6 +637,8 @@ class ProcessedVariable2D(ProcessedVariable):
 
     Parameters
     ----------
+    variable : str
+        The name of the variable
     base_variables : list of :class:`pybamm.Symbol`
         A list of base variables with a method `evaluate(t,y)`, each entry of which
         returns the value of that variable for that particular sub-solution.
@@ -624,6 +656,7 @@ class ProcessedVariable2D(ProcessedVariable):
 
     def __init__(
         self,
+        name: str,
         base_variables,
         base_variables_casadi,
         solution,
@@ -631,6 +664,7 @@ class ProcessedVariable2D(ProcessedVariable):
     ):
         self.dimensions = 2
         super().__init__(
+            name,
             base_variables,
             base_variables_casadi,
             solution,
@@ -749,6 +783,8 @@ class ProcessedVariable2DSciKitFEM(ProcessedVariable2D):
 
     Parameters
     ----------
+    variable : str
+        The name of the variable
     base_variables : list of :class:`pybamm.Symbol`
         A list of base variables with a method `evaluate(t,y)`, each entry of which
         returns the value of that variable for that particular sub-solution.
@@ -766,6 +802,7 @@ class ProcessedVariable2DSciKitFEM(ProcessedVariable2D):
 
     def __init__(
         self,
+        name: str,
         base_variables,
         base_variables_casadi,
         solution,
@@ -773,6 +810,7 @@ class ProcessedVariable2DSciKitFEM(ProcessedVariable2D):
     ):
         self.dimensions = 2
         super(ProcessedVariable2D, self).__init__(
+            name,
             base_variables,
             base_variables_casadi,
             solution,
@@ -814,6 +852,8 @@ class ProcessedVariable3D(ProcessedVariable):
 
     Parameters
     ----------
+    variable : str
+        The name of the variable
     base_variables : list of :class:`pybamm.Symbol`
         A list of base variables with a method `evaluate(t,y)`, each entry of which
         returns the value of that variable for that particular sub-solution.
@@ -831,6 +871,7 @@ class ProcessedVariable3D(ProcessedVariable):
 
     def __init__(
         self,
+        name: str,
         base_variables,
         base_variables_casadi,
         solution,
@@ -838,6 +879,7 @@ class ProcessedVariable3D(ProcessedVariable):
     ):
         self.dimensions = 3
         super().__init__(
+            name,
             base_variables,
             base_variables_casadi,
             solution,
@@ -1005,6 +1047,8 @@ class ProcessedVariable3DSciKitFEM(ProcessedVariable3D):
 
     Parameters
     ----------
+    variable : str
+        The name of the variable
     base_variables : list of :class:`pybamm.Symbol`
         A list of base variables with a method `evaluate(t,y)`, each entry of which
         returns the value of that variable for that particular sub-solution.
@@ -1022,6 +1066,7 @@ class ProcessedVariable3DSciKitFEM(ProcessedVariable3D):
 
     def __init__(
         self,
+        name: str,
         base_variables,
         base_variables_casadi,
         solution,
@@ -1029,6 +1074,7 @@ class ProcessedVariable3DSciKitFEM(ProcessedVariable3D):
     ):
         self.dimensions = 3
         super(ProcessedVariable3D, self).__init__(
+            name,
             base_variables,
             base_variables_casadi,
             solution,
@@ -1079,7 +1125,7 @@ class ProcessedVariable3DSciKitFEM(ProcessedVariable3D):
         return entries
 
 
-def process_variable(base_variables, *args, **kwargs):
+def process_variable(name: str, base_variables, *args, **kwargs):
     mesh = base_variables[0].mesh
     domain = base_variables[0].domain
 
@@ -1093,22 +1139,22 @@ def process_variable(base_variables, *args, **kwargs):
         and "current collector" in domain
         and isinstance(mesh, pybamm.ScikitSubMesh2D)
     ):
-        return ProcessedVariable2DSciKitFEM(base_variables, *args, **kwargs)
+        return ProcessedVariable2DSciKitFEM(name, base_variables, *args, **kwargs)
     if hasattr(base_variables[0], "secondary_mesh"):
         if "current collector" in base_variables[0].domains["secondary"] and isinstance(
             base_variables[0].secondary_mesh, pybamm.ScikitSubMesh2D
         ):
-            return ProcessedVariable3DSciKitFEM(base_variables, *args, **kwargs)
+            return ProcessedVariable3DSciKitFEM(name, base_variables, *args, **kwargs)
 
     # check variable shape
     if len(base_eval_shape) == 0 or base_eval_shape[0] == 1:
-        return ProcessedVariable0D(base_variables, *args, **kwargs)
+        return ProcessedVariable0D(name, base_variables, *args, **kwargs)
 
     n = mesh.npts
     base_shape = base_eval_shape[0]
     # Try some shapes that could make the variable a 1D variable
     if base_shape in [n, n + 1]:
-        return ProcessedVariable1D(base_variables, *args, **kwargs)
+        return ProcessedVariable1D(name, base_variables, *args, **kwargs)
 
     # Try some shapes that could make the variable a 2D variable
     first_dim_nodes = mesh.nodes
@@ -1118,7 +1164,7 @@ def process_variable(base_variables, *args, **kwargs):
         len(first_dim_nodes),
         len(first_dim_edges),
     ]:
-        return ProcessedVariable2D(base_variables, *args, **kwargs)
+        return ProcessedVariable2D(name, base_variables, *args, **kwargs)
 
     # Try some shapes that could make the variable a 3D variable
     tertiary_pts = base_variables[0].tertiary_mesh.nodes
@@ -1126,7 +1172,7 @@ def process_variable(base_variables, *args, **kwargs):
         len(first_dim_nodes),
         len(first_dim_edges),
     ]:
-        return ProcessedVariable3D(base_variables, *args, **kwargs)
+        return ProcessedVariable3D(name, base_variables, *args, **kwargs)
 
     raise NotImplementedError(f"Shape not recognized for {base_variables[0]}")
 
diff --git a/src/pybamm/solvers/solution.py b/src/pybamm/solvers/solution.py
index 98ec1b346..90ab2ed6b 100644
--- a/src/pybamm/solvers/solution.py
+++ b/src/pybamm/solvers/solution.py
@@ -468,7 +468,7 @@ class Solution:
                 model._variables_casadi[variable] = var_casadi
             vars_casadi.append(var_casadi)
         var = pybamm.process_variable(
-            vars_pybamm, vars_casadi, self, time_integral=time_integral
+            variable, vars_pybamm, vars_casadi, self, time_integral=time_integral
         )
 
         self._variables[variable] = var


Test patch:
diff --git a/tests/unit/test_solvers/test_processed_variable.py b/tests/unit/test_solvers/test_processed_variable.py
index fa165f438..51d12c563 100644
--- a/tests/unit/test_solvers/test_processed_variable.py
+++ b/tests/unit/test_solvers/test_processed_variable.py
@@ -79,6 +79,7 @@ class TestProcessedVariable:
         var_casadi = to_casadi(var_sol, y_sol)
         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),
@@ -102,6 +103,7 @@ class TestProcessedVariable:
         yp_sol = self._get_yps(y_sol, hermite_interp)
         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),
@@ -116,6 +118,7 @@ class TestProcessedVariable:
         sol = self._sol_default(t_sol, y_sol, yp_sol, model)
         var_casadi = to_casadi(var, y_sol)
         processed_var = pybamm.process_variable(
+            "test",
             [var],
             [var_casadi],
             sol,
@@ -162,6 +165,7 @@ class TestProcessedVariable:
         model = pybamm.BaseModel()
         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),
@@ -187,6 +191,7 @@ class TestProcessedVariable:
         model = pybamm.BaseModel()
         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),
@@ -209,6 +214,7 @@ class TestProcessedVariable:
         yp_sol = self._get_yps(y_sol, hermite_interp)
         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, pybamm.BaseModel()),
@@ -227,6 +233,7 @@ class TestProcessedVariable:
         inputs = {"a": np.array([1.0])}
         var_casadi = to_casadi(var, y_sol, inputs=inputs)
         processed_var = pybamm.process_variable(
+            "test",
             [var],
             [var_casadi],
             pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), inputs),
@@ -255,6 +262,7 @@ class TestProcessedVariable:
 
         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),
@@ -265,6 +273,7 @@ class TestProcessedVariable:
         )
         eqn_casadi = to_casadi(eqn_sol, y_sol)
         processed_eqn = pybamm.process_variable(
+            "test",
             [eqn_sol],
             [eqn_casadi],
             self._sol_default(t_sol, y_sol, yp_sol),
@@ -287,6 +296,7 @@ class TestProcessedVariable:
         x_s_edge.mesh = disc.mesh["separator"]
         x_s_casadi = to_casadi(x_s_edge, y_sol)
         processed_x_s_edge = pybamm.process_variable(
+            "test",
             [x_s_edge],
             [x_s_casadi],
             self._sol_default(t_sol, y_sol, yp_sol),
@@ -303,6 +313,7 @@ class TestProcessedVariable:
         yp_sol = self._get_yps(y_sol, hermite_interp, values=0)
         eqn_casadi = to_casadi(eqn_sol, y_sol)
         processed_eqn2 = pybamm.process_variable(
+            "test",
             [eqn_sol],
             [eqn_casadi],
             self._sol_default(t_sol, y_sol, yp_sol),
@@ -342,7 +353,7 @@ class TestProcessedVariable:
         c = pybamm.StateVector(slice(0, var_pts[x]), domain=["SEI layer"])
         c.mesh = mesh["SEI layer"]
         c_casadi = to_casadi(c, y_sol)
-        pybamm.process_variable([c], [c_casadi], solution)
+        pybamm.process_variable("test", [c], [c_casadi], solution)
 
     @pytest.mark.parametrize("hermite_interp", _hermite_args)
     def test_processed_variable_2D_x_r(self, hermite_interp):
@@ -465,6 +476,7 @@ class TestProcessedVariable:
         x_s_edge.secondary_mesh = disc.mesh["current collector"]
         x_s_casadi = to_casadi(x_s_edge, y_sol)
         processed_x_s_edge = pybamm.process_variable(
+            "test",
             [x_s_edge],
             [x_s_casadi],
             self._sol_default(t_sol, y_sol, yp_sol),
@@ -500,6 +512,7 @@ class TestProcessedVariable:
 
         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),
@@ -525,6 +538,7 @@ class TestProcessedVariable:
 
         var_casadi = to_casadi(var_sol, u_sol)
         processed_var = pybamm.process_variable(
+            "test",
             [var_sol],
             [var_casadi],
             self._sol_default(t_sol, u_sol, yp_sol),
@@ -552,6 +566,7 @@ class TestProcessedVariable:
             options={"dimensionality": 2}
         )
         processed_var = pybamm.process_variable(
+            "test",
             [var_sol],
             [var_casadi],
             pybamm.Solution(t_sol, u_sol, model, {}, all_yps=yp_sol),
@@ -573,6 +588,7 @@ class TestProcessedVariable:
         yp_sol = self._get_yps(y_sol, hermite_interp, values=5)
         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),
@@ -585,6 +601,7 @@ class TestProcessedVariable:
 
         eqn_casadi = to_casadi(eqn, y_sol)
         processed_eqn = pybamm.process_variable(
+            "test",
             [eqn],
             [eqn_casadi],
             self._sol_default(t_sol, y_sol, yp_sol),
@@ -610,6 +627,7 @@ class TestProcessedVariable:
         yp_sol = self._get_yps(y_sol, hermite_interp)
         eqn_casadi = to_casadi(eqn, y_sol)
         processed_var = pybamm.process_variable(
+            "test",
             [eqn],
             [eqn_casadi],
             self._sol_default(t_sol, y_sol, yp_sol, pybamm.BaseModel()),
@@ -635,6 +653,7 @@ class TestProcessedVariable:
 
         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),
@@ -660,6 +679,7 @@ class TestProcessedVariable:
         )
         eqn_casadi = to_casadi(eqn_sol, y_sol)
         processed_eqn = pybamm.process_variable(
+            "test",
             [eqn_sol],
             [eqn_casadi],
             self._sol_default(t_sol, y_sol, yp_sol),
@@ -682,6 +702,7 @@ class TestProcessedVariable:
         x_casadi = to_casadi(x_disc, y_sol)
 
         processed_x = pybamm.process_variable(
+            "test",
             [x_disc],
             [x_casadi],
             self._sol_default(t_sol, y_sol, yp_sol),
@@ -697,6 +718,7 @@ class TestProcessedVariable:
         r_n.mesh = disc.mesh["negative particle"]
         r_n_casadi = to_casadi(r_n, y_sol)
         processed_r_n = pybamm.process_variable(
+            "test",
             [r_n],
             [r_n_casadi],
             self._sol_default(t_sol, y_sol, yp_sol),
@@ -717,6 +739,7 @@ class TestProcessedVariable:
             options={"particle size": "distribution"}
         )
         processed_R_n = pybamm.process_variable(
+            "test",
             [R_n],
             [R_n_casadi],
             pybamm.Solution(t_sol, y_sol, model, {}),
@@ -743,6 +766,7 @@ class TestProcessedVariable:
 
         eqn_casadi = to_casadi(eqn_sol, y_sol)
         processed_var = pybamm.process_variable(
+            "test",
             [eqn_sol],
             [eqn_casadi],
             self._sol_default(t_sol, y_sol, yp_sol),
@@ -798,6 +822,7 @@ class TestProcessedVariable:
         )
         with pytest.raises(NotImplementedError, match="Spatial variable name"):
             pybamm.process_variable(
+                "test",
                 [var_sol],
                 [var_casadi],
                 pybamm.Solution(t_sol, y_sol, model, {}, all_yps=yp_sol),
@@ -830,6 +855,7 @@ class TestProcessedVariable:
 
         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),
@@ -879,6 +905,7 @@ class TestProcessedVariable:
 
         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),
@@ -915,6 +942,7 @@ class TestProcessedVariable:
 
         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),
@@ -947,6 +975,7 @@ class TestProcessedVariable:
 
         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),
@@ -988,6 +1017,7 @@ class TestProcessedVariable:
 
         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),
@@ -1013,6 +1043,7 @@ class TestProcessedVariable:
 
         var_casadi = to_casadi(var_sol, u_sol)
         processed_var = pybamm.process_variable(
+            "test",
             [var_sol],
             [var_casadi],
             self._sol_default(t_sol, u_sol, yp_sol),
@@ -1060,6 +1091,7 @@ class TestProcessedVariable:
 
         var_casadi = to_casadi(var_sol, u_sol)
         processed_var = pybamm.process_variable(
+            "test",
             [var_sol],
             [var_casadi],
             self._sol_default(t_sol, u_sol, yp_sol),
@@ -1124,6 +1156,7 @@ class TestProcessedVariable:
             }
         )
         processed_var = pybamm.process_variable(
+            "test",
             [var_sol],
             [var_casadi],
             pybamm.Solution(t_sol, y_sol, model, {}, all_yps=yp_sol),
@@ -1192,6 +1225,7 @@ class TestProcessedVariable:
         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),
@@ -1269,6 +1303,7 @@ class TestProcessedVariable:
         var_casadi = to_casadi(var_sol, u_sol)
 
         processed_var = pybamm.process_variable(
+            "test",
             [var_sol],
             [var_casadi],
             self._sol_default(t_sol, u_sol, up_sol),
@@ -1324,6 +1359,7 @@ class TestProcessedVariable:
         y_sol = np.array([np.linspace(0, 5)])
         var_casadi = to_casadi(var, y_sol)
         processed_var = pybamm.process_variable(
+            "test",
             [var],
             [var_casadi],
             pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}),
@@ -1375,6 +1411,7 @@ class TestProcessedVariable:
         sol = sol1 + sol2
         var_casadi = to_casadi(var, sol.all_ys[0])
         processed_var = pybamm.process_variable(
+            "test",
             [var] * len(sol.all_ts),
             [var_casadi] * len(sol.all_ts),
             sol,
diff --git a/tests/unit/test_solvers/test_processed_variable_computed.py b/tests/unit/test_solvers/test_processed_variable_computed.py
index 652cf659d..d05123bb4 100644
--- a/tests/unit/test_solvers/test_processed_variable_computed.py
+++ b/tests/unit/test_solvers/test_processed_variable_computed.py
@@ -396,6 +396,7 @@ class TestProcessedVariableComputed:
         x_s_edge.secondary_mesh = disc.mesh["current collector"]
         x_s_casadi = to_casadi(x_s_edge, y_sol)
         processed_x_s_edge = pybamm.process_variable(
+            "test",
             [x_s_edge],
             [x_s_casadi],
             pybamm.Solution(
diff --git a/tests/unit/test_solvers/test_solution.py b/tests/unit/test_solvers/test_solution.py
index c100d5856..95600a796 100644
--- a/tests/unit/test_solvers/test_solution.py
+++ b/tests/unit/test_solvers/test_solution.py
@@ -526,12 +526,63 @@ class TestSolution:
         model.rhs = {c: -a * c}
         model.initial_conditions = {c: 1}
         model.variables["data_comparison"] = data_comparison
+        model.variables["data"] = data
+        model.variables["c"] = c
 
         solver = solver_class()
         for a in [0.5, 1.0, 2.0]:
-            sol = solver.solve(model, t_eval=t_eval, inputs={"a": a})
+            sol = solver.solve(model, t_eval=t_eval, t_interp=t_interp, inputs={"a": a})
             y_sol = np.exp(-a * data_times)
             expected = np.sum((y_sol - data_values) ** 2)
             np.testing.assert_allclose(
                 sol["data_comparison"](), expected, rtol=1e-3, atol=1e-2
             )
+
+            # sensitivity calculation only supported for IDAKLUSolver
+            if solver_class == pybamm.IDAKLUSolver:
+                sol = solver.solve(
+                    model,
+                    t_eval=t_eval,
+                    t_interp=t_interp,
+                    inputs={"a": a},
+                    calculate_sensitivities=True,
+                )
+                y_sol = np.exp(-a * data_times)
+                dy_sol_da = -data_times * y_sol
+
+                np.testing.assert_allclose(
+                    sol["data"].sensitivities["a"].full().flatten(),
+                    np.zeros_like(data_times),
+                    rtol=1e-3,
+                    atol=1e-2,
+                )
+                np.testing.assert_allclose(
+                    sol["c"].data,
+                    y_sol,
+                    rtol=1e-3,
+                    atol=1e-2,
+                )
+                np.testing.assert_allclose(
+                    sol["c"].sensitivities["a"].full().flatten(),
+                    dy_sol_da,
+                    rtol=1e-3,
+                    atol=1e-2,
+                )
+                np.testing.assert_allclose(
+                    sol["data_comparison"].sensitivities["a"],
+                    np.sum(2 * (y_sol - data_values) * dy_sol_da),
+                    rtol=1e-3,
+                    atol=1e-2,
+                )
+
+                # should raise error if t_interp is not equal to data_times
+                with pytest.raises(
+                    pybamm.SolverError,
+                    match="solution times and discrete times of the time integral are not equal",
+                ):
+                    solver.solve(
+                        model,
+                        t_eval=t_eval,
+                        inputs={"a": a},
+                        calculate_sensitivities=True,
+                    )["data_comparison"].sensitivities["a"]