Repository: pybamm-team/PyBaMM
Pull request #2931 (2023-05-08 17:08:14)

save step and cycle number from a solution
# Description

I keep doing this manually and it might be useful here instead

## 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.

- [ ] New feature (non-breaking change which adds functionality)
- [ ] Optimization (back-end change that speeds up the code)
- [ ] Bug fix (non-breaking change which fixes an issue)

# Key checklist:

- [ ] No style issues: `$ pre-commit run` (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)
- [ ] All tests pass: `$ python run-tests.py --all`
- [ ] The documentation builds: `$ python run-tests.py --doctest`

You can run unit and doctests together at once, using `$ python run-tests.py --quick`.

## Further checks:

- [ ] Code is commented, particularly in 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 ec99475d4..84d68222a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # [Unreleased](https://github.com/pybamm-team/PyBaMM/)
 
+## Features
+
+- If a solution contains cycles and steps, the cycle number and step number are now saved when `solution.save_data()` is called ([#2931](https://github.com/pybamm-team/PyBaMM/pull/2931))
+
 ## Optimizations
 
 - Update Jax (0.4.8) and JaxLib (0.4.7) compatibility ([#2927](https://github.com/pybamm-team/PyBaMM/pull/2927))
diff --git a/pybamm/solvers/solution.py b/pybamm/solvers/solution.py
index 1ac1e7fd4..278355f7f 100644
--- a/pybamm/solvers/solution.py
+++ b/pybamm/solvers/solution.py
@@ -535,6 +535,71 @@ class Solution(object):
         with open(filename, "wb") as f:
             pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
 
+    def get_data_dict(self, variables=None, short_names=None, cycles_and_steps=True):
+        """
+        Construct a (standard python) dictionary of the solution data containing the
+        variables in `variables`. If `variables` is None then all variables are
+        returned. Any variable names in short_names are replaced with the corresponding
+        short name.
+
+        If the solution has cycles, then the cycle numbers and step numbers are also
+        returned in the dictionary.
+
+        Parameters
+        ----------
+        variables : list, optional
+            List of variables to return. If None, returns all variables in solution.data
+        short_names : dict, optional
+            Dictionary of shortened names to use when saving.
+        cycles_and_steps : bool, optional
+            Whether to include the cycle numbers and step numbers in the dictionary
+
+        Returns
+        -------
+        dict
+            A dictionary of the solution data
+        """
+        if variables is None:
+            # variables not explicitly provided -> save all variables that have been
+            # computed
+            data_long_names = self.data
+        else:
+            if isinstance(variables, str):
+                variables = [variables]
+            # otherwise, save only the variables specified
+            data_long_names = {}
+            for name in variables:
+                data_long_names[name] = self[name].data
+        if len(data_long_names) == 0:
+            raise ValueError(
+                """
+                Solution does not have any data. Please provide a list of variables
+                to save.
+                """
+            )
+
+        # Use any short names if provided
+        data_short_names = {}
+        short_names = short_names or {}
+        for name, var in data_long_names.items():
+            name = short_names.get(name, name)  # return name if no short name
+            data_short_names[name] = var
+
+        # Save cycle number and step number if the solution has them
+        if cycles_and_steps and len(self.cycles) > 0:
+            data_short_names["Cycle"] = np.array([])
+            data_short_names["Step"] = np.array([])
+            for i, cycle in enumerate(self.cycles):
+                data_short_names["Cycle"] = np.concatenate(
+                    [data_short_names["Cycle"], i * np.ones_like(cycle.t)]
+                )
+                for j, step in enumerate(cycle.steps):
+                    data_short_names["Step"] = np.concatenate(
+                        [data_short_names["Step"], j * np.ones_like(step.t)]
+                    )
+
+        return data_short_names
+
     def save_data(
         self, filename=None, variables=None, to_format="pickle", short_names=None
     ):
@@ -565,46 +630,19 @@ class Solution(object):
         -------
         data : str, optional
             str if 'csv' or 'json' is chosen and filename is None, otherwise None
-
-
         """
-        if variables is None:
-            # variables not explicitly provided -> save all variables that have been
-            # computed
-            data = self.data
-        else:
-            # otherwise, save only the variables specified
-            data = {}
-            for name in variables:
-                data[name] = self[name].data
-        if len(data) == 0:
-            raise ValueError(
-                """
-                Solution does not have any data. Please provide a list of variables
-                to save.
-                """
-            )
-
-        # Use any short names if provided
-        data_short_names = {}
-        short_names = short_names or {}
-        for name, var in data.items():
-            # change to short name if it exists
-            if name in short_names:
-                data_short_names[short_names[name]] = var
-            else:
-                data_short_names[name] = var
+        data = self.get_data_dict(variables=variables, short_names=short_names)
 
         if to_format == "pickle":
             if filename is None:
                 raise ValueError("pickle format must be written to a file")
             with open(filename, "wb") as f:
-                pickle.dump(data_short_names, f, pickle.HIGHEST_PROTOCOL)
+                pickle.dump(data, f, pickle.HIGHEST_PROTOCOL)
         elif to_format == "matlab":
             if filename is None:
                 raise ValueError("matlab format must be written to a file")
             # Check all the variable names only contain a-z, A-Z or _ or numbers
-            for name in data_short_names.keys():
+            for name in data.keys():
                 # Check the string only contains the following ASCII:
                 # a-z (97-122)
                 # A-Z (65-90)
@@ -627,23 +665,23 @@ class Solution(object):
                             "['Electrolyte concentration'], to_format='matlab, "
                             "short_names={'Electrolyte concentration': 'c_e'})"
                         )
-            savemat(filename, data_short_names)
+            savemat(filename, data)
         elif to_format == "csv":
-            for name, var in data_short_names.items():
+            for name, var in data.items():
                 if var.ndim >= 2:
                     raise ValueError(
                         "only 0D variables can be saved to csv, but '{}' is {}D".format(
                             name, var.ndim - 1
                         )
                     )
-            df = pd.DataFrame(data_short_names)
+            df = pd.DataFrame(data)
             return df.to_csv(filename, index=False)
         elif to_format == "json":
             if filename is None:
-                return json.dumps(data_short_names, cls=NumpyEncoder)
+                return json.dumps(data, cls=NumpyEncoder)
             else:
                 with open(filename, "w") as outfile:
-                    json.dump(data_short_names, outfile, cls=NumpyEncoder)
+                    json.dump(data, outfile, cls=NumpyEncoder)
         else:
             raise ValueError("format '{}' not recognised".format(to_format))
 


Test patch:
diff --git a/tests/unit/test_solvers/test_solution.py b/tests/unit/test_solvers/test_solution.py
index 478bf94d8..2ef01d743 100644
--- a/tests/unit/test_solvers/test_solution.py
+++ b/tests/unit/test_solvers/test_solution.py
@@ -206,8 +206,6 @@ class TestSolution(TestCase):
         model.variables["c"] = c
         model.variables["2c"] = 2 * c
 
-        disc = pybamm.Discretisation()
-        disc.process_model(model)
         solution = pybamm.ScipySolver().solve(model, np.linspace(0, 1))
 
         # test create a new processed variable
@@ -230,8 +228,6 @@ class TestSolution(TestCase):
         model.variables["c"] = c
         model.variables["2c"] = 2 * c
 
-        disc = pybamm.Discretisation()
-        disc.process_model(model)
         solution = pybamm.ScipySolver().solve(model, np.linspace(0, 1))
 
         solution.plot(["c", "2c"], testing=True)
@@ -327,6 +323,27 @@ class TestSolution(TestCase):
         np.testing.assert_array_equal(solution["c"].entries, solution_load["c"].entries)
         np.testing.assert_array_equal(solution["d"].entries, solution_load["d"].entries)
 
+    def test_get_data_cycles_steps(self):
+        model = pybamm.BaseModel()
+        c = pybamm.Variable("c")
+        model.rhs = {c: -c}
+        model.initial_conditions = {c: 1}
+        model.variables["c"] = c
+
+        solver = pybamm.ScipySolver()
+        sol1 = solver.solve(model, np.linspace(0, 1))
+        sol2 = solver.solve(model, np.linspace(1, 2))
+
+        sol = sol1 + sol2
+        sol.cycles = [sol]
+        sol.cycles[0].steps = [sol1, sol2]
+
+        data = sol.get_data_dict("c")
+        np.testing.assert_array_equal(data["Cycle"], 0)
+        np.testing.assert_array_equal(
+            data["Step"], np.concatenate([np.zeros(50), np.ones(50)])
+        )
+
     def test_solution_evals_with_inputs(self):
         model = pybamm.lithium_ion.SPM()
         geometry = model.default_geometry