add cycle number as a key (#4798)

pybamm-team · Jan 24, 2025 · f608792 · f608792
1 parent 3d0d6ff
commit f608792
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/src/pybamm/solvers/summary_variable.py b/src/pybamm/solvers/summary_variable.py
@@ -41,7 +41,7 @@ def __init__(
         self.user_inputs = user_inputs or {}
         self.esoh_solver = esoh_solver
         self._variables = {}  # Store computed variables
-        self.cycle_number = None
+        self.cycle_number = np.array([])
 
         model = solution.all_models[0]
         self._possible_variables = model.summary_variables  # minus esoh variables
@@ -122,6 +122,8 @@ def __getitem__(self, key: str) -> float | list[float]:
         if key in self._variables:
             # return it if it exists
             return self._variables[key]
+        elif key == "Cycle number":
+            return self.cycle_number
         elif key not in self.all_variables:
             # check it's listed as a summary variable
             raise KeyError(f"Variable '{key}' is not a summary variable.")

diff --git a/tests/unit/test_solvers/test_summary_variables.py b/tests/unit/test_solvers/test_summary_variables.py
@@ -167,5 +167,7 @@ def test_get_esoh_cycle_summary_vars(self):
         sol = sim.solve()
 
         assert len(sol.summary_variables.cycles) == 10
+        assert sol.summary_variables["Cycle number"][0] == 1
+        assert sol.summary_variables["Cycle number"][9] == 10
         assert len(sol.summary_variables["x_100"]) == 10
         assert np.isclose(sol.summary_variables["x_100"][0], 0.9493)