Added InputParameter support in PyBamm experiments

src/pybamm/experiment/experiment.py

@@ -42,13 +42,17 @@ def __init__(
         self,
         operating_conditions: list[str | tuple[str] | BaseStep],
         period: str | None = None,
+        input_parameters: dict[str, any] | None = None,
         temperature: float | None = None,
         termination: list[str] | None = None,
     ):
+        self.input_parameters = input_parameters if input_parameters is not None else {}
+
         # Save arguments for copying
         self.args = (
             operating_conditions,
             period,
+            input_parameters,
             temperature,
             termination,
         )
@@ -70,7 +74,7 @@ def __init__(
         self.temperature = _convert_temperature_to_kelvin(temperature)
 
         processed_steps = self.process_steps(
-            steps_unprocessed, self.period, self.temperature
+            steps_unprocessed, self.period, self.temperature, self.input_parameters
         )
 
         self.steps = [processed_steps[repr(step)] for step in steps_unprocessed]
@@ -93,7 +97,7 @@ def __init__(
         self.termination = self.read_termination(termination)
 
     @staticmethod
-    def process_steps(unprocessed_steps, period, temp):
+    def process_steps(unprocessed_steps, period, temp, input_parameters=None):
         processed_steps = {}
         for step in unprocessed_steps:
             if repr(step) in processed_steps:
@@ -103,7 +107,7 @@ def process_steps(unprocessed_steps, period, temp):
             elif isinstance(step, pybamm.step.BaseStep):
                 # Copy the step to avoid modifying the original with the period and
                 # temperature and any other changes
-                processed_step = step.copy()
+                processed_step = step.copy(input_parameters=input_parameters)
             else:
                 raise TypeError("Operating conditions must be a Step object or string.")
 

src/pybamm/experiment/step/base_step.py

@@ -68,6 +68,7 @@ def __init__(
         start_time=None,
         description=None,
         direction: str | None = None,
+        input_parameters: dict[str, any] | None = None,
     ):
         potential_directions = ["charge", "discharge", "rest", None]
         if direction not in potential_directions:
@@ -77,6 +78,7 @@ def __init__(
         self.input_duration = duration
         self.input_duration = duration
         self.input_value = value
+        self.input_parameters = input_parameters if input_parameters is not None else {}
         # Check if drive cycle
         is_drive_cycle = isinstance(value, np.ndarray)
         is_python_function = callable(value)
@@ -194,7 +196,7 @@ def __init__(
         self.next_start_time = None
         self.end_time = None
 
-    def copy(self):
+    def copy(self, input_parameters: dict[str, any] | None = None):
         """
         Return a copy of the step.
 
@@ -203,6 +205,12 @@ def copy(self):
         :class:`pybamm.Step`
             A copy of the step.
         """
+
+        new_input_parameters = (
+            input_parameters
+            if input_parameters is not None
+            else getattr(self, "input_parameters", {})
+        )
         return self.__class__(
             self.input_value,
             duration=self.input_duration,
@@ -213,6 +221,7 @@ def copy(self):
             start_time=self.start_time,
             description=self.description,
             direction=self.direction,
+            input_parameters=new_input_parameters,
         )
 
     def __str__(self):
@@ -387,7 +396,18 @@ def value_based_charge_or_discharge(self):
         """
         if isinstance(self.value, pybamm.Symbol):
             inpt = {"start time": 0}
-            init_curr = self.value.evaluate(t=0, inputs=inpt).flatten()[0]
+            merged_inputs = {}
+            if hasattr(self, "input_parameters"):
+                merged_inputs.update(self.input_parameters)
+            merged_inputs.update(inpt)
+            if (
+                isinstance(self.value, pybamm.InputParameter)
+                and self.value.name not in merged_inputs
+            ):
+                merged_inputs[self.value.name] = self.input_parameters.get(
+                    self.value.name
+                )
+            init_curr = self.value.evaluate(t=0, inputs=merged_inputs).flatten()[0]
         else:
             init_curr = self.value
         sign = np.sign(init_curr)

src/pybamm/experiment/step/steps.py

@@ -320,8 +320,13 @@ def __init__(self, current_value_function, **kwargs):
     def current_value(self, variables):
         return self.current_value_function(variables)
 
-    def copy(self):
-        return CustomStepExplicit(self.current_value_function, **self.kwargs)
+    def copy(self, input_parameters: dict[str, any] | None = None):
+        new_kwargs = dict(self.kwargs)
+        if input_parameters is not None:
+            new_kwargs["input_parameters"] = input_parameters
+        elif "input_parameters" not in new_kwargs:
+            new_kwargs["input_parameters"] = {}
+        return CustomStepExplicit(self.current_value_function, **new_kwargs)
 
 
 class CustomStepImplicit(BaseStepImplicit):
@@ -414,7 +419,10 @@ def get_submodel(self, model):
             model.param, self.current_rhs_function, model.options, control=self.control
         )
 
-    def copy(self):
-        return CustomStepImplicit(
-            self.current_rhs_function, self.control, **self.kwargs
-        )
+    def copy(self, input_parameters: dict[str, any] | None = None):
+        new_kwargs = dict(self.kwargs)
+        if input_parameters is not None:
+            new_kwargs["input_parameters"] = input_parameters
+        elif "input_parameters" not in new_kwargs:
+            new_kwargs["input_parameters"] = {}
+        return CustomStepImplicit(self.current_rhs_function, self.control, **new_kwargs)

tests/unit/test_experiments/test_experiment.py

@@ -214,3 +214,37 @@ def test_set_next_start_time(self):
 
         # TODO: once #3176 is completed, the test should pass for
         # operating_conditions_steps (or equivalent) as well
+
+    def test_experiment_input_parameter(self):
+        current_step = pybamm.step.current(
+            1, input_parameters=pybamm.InputParameter("I_app"), termination="2.5V"
+        )
+
+        experiment = pybamm.Experiment(
+            operating_conditions=[current_step],
+            input_parameters={"I_app": 1},
+            period="1 minute",
+            temperature=25.00,
+            termination="2.5 V",
+        )
+
+        processed_steps = experiment.process_steps(
+            [current_step], experiment.period, 298.15, experiment.input_parameters
+        )
+
+        for step in processed_steps.values():
+            assert hasattr(step, "input_parameters")
+            assert "I_app" in step.input_parameters
+            assert step.input_parameters["I_app"] == 1
+
+    def test_experiment_copy_preserves_input_parameters(self):
+        current_step = pybamm.step.current(
+            1, input_parameters=pybamm.InputParameter("I_app"), termination="2.5 V"
+        )
+        experiment = pybamm.Experiment(
+            input_parameters={"I_app": 1},
+            termination="2.5 V",
+            operating_conditions=[current_step],
+        )
+        experiment_copy = experiment.copy()
+        assert experiment_copy.input_parameters == experiment.input_parameters