Improve memory use in experiments (#3261)

* #3081 add test

* #3081 change __hash__ and __eq__ to only check certain arguments

* #3081 remove unused variable

* #3081 redefine step processing to avoid processing redundant models

* #3081 add test

* #3081 fix failing test

* #3081 update CHANGELOG

* #3081 add non-battery PDE test

* #3081 fix failing notebooks

* Revert "#3081 fix failing notebooks"

This reverts commit 6f0b6a8.

* Revert "#3081 add non-battery PDE test"

This reverts commit 64cd12c.

* #3081 update docstring

CHANGELOG.md

@@ -23,12 +23,15 @@
 - Error generated when invalid parameter values are passed ([#3132](https://github.com/pybamm-team/PyBaMM/pull/3132))
 - Parameters in `Prada2013` have been updated to better match those given in the paper, which is a 2.3 Ah cell, instead of the mix-and-match with the 1.1 Ah cell from Lain2019 ([#3096](https://github.com/pybamm-team/PyBaMM/pull/3096))
 
+## Optimizations
 
+- Improved how steps are processed in simulations to reduce memory usage ([#3261](https://github.com/pybamm-team/PyBaMM/pull/3261))
 
 ## Breaking changes
 
 - The class `pybamm.thermal.OneDimensionalX` has been moved to `pybamm.thermal.pouch_cell.OneDimensionalX` to reflect the fact that the model formulation implicitly assumes a pouch cell geometry ([#3257](https://github.com/pybamm-team/PyBaMM/pull/3257))
 - The "lumped" thermal option now always used the parameters "Cell cooling surface area [m2]", "Cell volume [m3]" and "Total heat transfer coefficient [W.m-2.K-1]" to compute the cell cooling regardless of the chosen "cell geometry" option. The user must now specify the correct values for these parameters instead of them being calculated based on e.g. a pouch cell. An `OptionWarning` is raised to let users know to update their parameters ([#3257](https://github.com/pybamm-team/PyBaMM/pull/3257))
+- Numpy functions now work with PyBaMM symbols (e.g. `np.exp(pybamm.Symbol("a"))` returns `pybamm.Exp(pybamm.Symbol("a"))`). This means that parameter functions can be specified using numpy functions instead of pybamm functions. Additionally, combining numpy arrays with pybamm objects now works (the numpy array is converted to a pybamm array) ([#3205](https://github.com/pybamm-team/PyBaMM/pull/3205))
 - Added option to use an empirical hysteresis model for the diffusivity and exchange-current density ([#3194](https://github.com/pybamm-team/PyBaMM/pull/3194))
 - Double-layer capacity can now be provided as a function of temperature ([#3174](https://github.com/pybamm-team/PyBaMM/pull/3174))
 - `pybamm_install_jax` is deprecated. It is now replaced with `pip install pybamm[jax]` ([#3163](https://github.com/pybamm-team/PyBaMM/pull/3163))

pybamm/experiment/experiment.py

@@ -68,11 +68,6 @@ def __init__(
             termination,
         )
 
-        self.datetime_formats = [
-            "Day %j %H:%M:%S",
-            "%Y-%m-%d %H:%M:%S",
-        ]
-
         operating_conditions_cycles = []
         for cycle in operating_conditions:
             # Check types and convert to list
@@ -89,21 +84,36 @@ def __init__(
 
         # Convert strings to pybamm.step._Step objects
         # We only do this once per unique step, do avoid unnecessary conversions
-        unique_steps_unprocessed = set(operating_conditions_steps_unprocessed)
+        # Assign experiment period and temperature if not specified in step
+        self.period = _convert_time_to_seconds(period)
+        self.temperature = _convert_temperature_to_kelvin(temperature)
+
         processed_steps = {}
-        for step in unique_steps_unprocessed:
-            if isinstance(step, str):
-                processed_steps[step] = pybamm.step.string(step)
+        for step in operating_conditions_steps_unprocessed:
+            if repr(step) in processed_steps:
+                continue
+            elif isinstance(step, str):
+                processed_step = pybamm.step.string(step)
             elif isinstance(step, pybamm.step._Step):
-                processed_steps[step] = step
+                processed_step = step
+
+            if processed_step.period is None:
+                processed_step.period = self.period
+            if processed_step.temperature is None:
+                processed_step.temperature = self.temperature
+
+            processed_steps[repr(step)] = processed_step
+
+        self.operating_conditions_steps = [
+            processed_steps[repr(step)]
+            for step in operating_conditions_steps_unprocessed
+        ]
 
         # Save the processed unique steps and the processed operating conditions
         # for every step
         self.unique_steps = set(processed_steps.values())
-        self.operating_conditions_steps = [
-            processed_steps[step] for step in operating_conditions_steps_unprocessed
-        ]
 
+        # Allocate experiment global variables
         self.initial_start_time = self.operating_conditions_steps[0].start_time
 
         if (
@@ -118,15 +128,6 @@ def __init__(
         self.termination_string = termination
         self.termination = self.read_termination(termination)
 
-        # Modify steps with period and temperature in place
-        self.period = _convert_time_to_seconds(period)
-        self.temperature = _convert_temperature_to_kelvin(temperature)
-        for step in self.unique_steps:
-            if step.period is None:
-                step.period = self.period
-            if step.temperature is None:
-                step.temperature = self.temperature
-
     def __str__(self):
         return str(self.operating_conditions_cycles)
 

pybamm/simulation.py

@@ -255,7 +255,7 @@ def set_up_and_parameterise_model_for_experiment(self):
             parameterised_model = new_parameter_values.process_model(
                 new_model, inplace=False
             )
-            self.experiment_unique_steps_to_model[repr(op)] = parameterised_model
+            self.experiment_unique_steps_to_model[op.basic_repr()] = parameterised_model
 
         # Set up rest model if experiment has start times
         if self.experiment.initial_start_time:
@@ -778,8 +778,8 @@ def solve(
                     else:
                         dt = op_conds.duration
                     op_conds_str = str(op_conds)
-                    model = self.op_conds_to_built_models[repr(op_conds)]
-                    solver = self.op_conds_to_built_solvers[repr(op_conds)]
+                    model = self.op_conds_to_built_models[op_conds.basic_repr()]
+                    solver = self.op_conds_to_built_solvers[op_conds.basic_repr()]
 
                     logs["step number"] = (step_num, cycle_length)
                     logs["step operating conditions"] = op_conds_str

pybamm/solvers/solution.py

@@ -443,7 +443,7 @@ def initial_start_time(self):
 
     @initial_start_time.setter
     def initial_start_time(self, value):
-        """Updates the reason for termination"""
+        """Updates the initial start time of the experiment"""
         self._initial_start_time = value
 
     def set_summary_variables(self, all_summary_variables):

pybamm/step/_steps_util.py

@@ -71,22 +71,25 @@ def __init__(
     ):
         self.type = typ
 
-        # Record all the args for repr
-        self.args = f"{typ}, {value}"
+        # Record all the args for repr and hash
+        self.repr_args = f"{typ}, {value}"
+        self.hash_args = f"{typ}, {value}"
         if duration:
-            self.args += f", duration={duration}"
+            self.repr_args += f", duration={duration}"
         if termination:
-            self.args += f", termination={termination}"
+            self.repr_args += f", termination={termination}"
+            self.hash_args += f", termination={termination}"
         if period:
-            self.args += f", period={period}"
+            self.repr_args += f", period={period}"
         if temperature:
-            self.args += f", temperature={temperature}"
+            self.repr_args += f", temperature={temperature}"
+            self.hash_args += f", temperature={temperature}"
         if tags:
-            self.args += f", tags={tags}"
+            self.repr_args += f", tags={tags}"
         if start_time:
-            self.args += f", start_time={start_time}"
+            self.repr_args += f", start_time={start_time}"
         if description:
-            self.args += f", description={description}"
+            self.repr_args += f", description={description}"
 
         # Check if drive cycle
         self.is_drive_cycle = isinstance(value, np.ndarray)
@@ -158,7 +161,15 @@ def __str__(self):
             return repr(self)
 
     def __repr__(self):
-        return f"_Step({self.args})"
+        return f"_Step({self.repr_args})"
+
+    def basic_repr(self):
+        """
+        Return a basic representation of the step, only with type, value, termination
+        and temperature, which are the variables involved in processing the model. Also
+        used for hashing.
+        """
+        return f"_Step({self.hash_args})"
 
     def to_dict(self):
         """
@@ -184,13 +195,11 @@ def to_dict(self):
     def __eq__(self, other):
         return (
             isinstance(other, _Step)
-            and self.__repr__() == other.__repr__()
-            and self.next_start_time == other.next_start_time
-            and self.end_time == other.end_time
+            and self.hash_args == other.hash_args
         )
 
     def __hash__(self):
-        return hash(repr(self))
+        return hash(self.basic_repr())
 
     @property
     def unit(self):

tests/unit/test_experiments/test_simulation_with_experiment.py

@@ -38,8 +38,12 @@ def test_set_up(self):
             [3600, 3 / Crate * 3600, 24 * 3600, 24 * 3600],
         )
 
-        model_I = sim.experiment_unique_steps_to_model[repr(op_conds[1])]  # CC charge
-        model_V = sim.experiment_unique_steps_to_model[repr(op_conds[2])]  # CV hold
+        model_I = sim.experiment_unique_steps_to_model[
+            op_conds[1].basic_repr()
+        ]  # CC charge
+        model_V = sim.experiment_unique_steps_to_model[
+            op_conds[2].basic_repr()
+        ]  # CV hold
         self.assertIn(
             "Current cut-off [A] [experiment]",
             [event.name for event in model_V.events],
@@ -729,6 +733,38 @@ def test_experiment_start_time_starting_solution(self):
         # test that the final time is correct (i.e. starting solution correctly set)
         self.assertEqual(new_solution["Time [s]"].entries[-1], 5400)
 
+    def test_experiment_start_time_identical_steps(self):
+        # Test that if we have the same step twice, with different start times,
+        # they get processed only once
+        model = pybamm.lithium_ion.SPM()
+
+        experiment = pybamm.Experiment(
+            [
+                pybamm.step.string(
+                    "Discharge at C/2 for 10 minutes",
+                    start_time=datetime(2023, 1, 1, 8, 0, 0),
+                ),
+                pybamm.step.string("Discharge at C/3 for 10 minutes"),
+                pybamm.step.string(
+                    "Discharge at C/2 for 10 minutes",
+                    start_time=datetime(2023, 1, 1, 10, 0, 0),
+                ),
+                pybamm.step.string("Discharge at C/3 for 10 minutes"),
+            ]
+        )
+
+        sim = pybamm.Simulation(model, experiment=experiment)
+        sim.solve(calc_esoh=False)
+
+        # Check that there are 4 steps
+        self.assertEqual(len(experiment.operating_conditions_steps), 4)
+
+        # Check that there are only 2 unique steps
+        self.assertEqual(len(sim.experiment.unique_steps), 2)
+
+        # Check that there are only 3 built models (unique steps + padding rest)
+        self.assertEqual(len(sim.op_conds_to_built_models), 3)
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")