Add functions of time in experiment step

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- Added functionality to pass in arbitrary functions of time as the argument for a (`pybamm.step`). ([#4222](https://github.com/pybamm-team/PyBaMM/pull/4222))
 - Added new parameters `"f{pref]Initial inner SEI on cracks thickness [m]"` and `"f{pref]Initial outer SEI on cracks thickness [m]"`, instead of hardcoding these to `L_inner_0 / 10000` and `L_outer_0 / 10000`. ([#4168](https://github.com/pybamm-team/PyBaMM/pull/4168))
 - Added `pybamm.DataLoader` class to fetch data files from [pybamm-data](https://github.com/pybamm-team/pybamm-data/releases/tag/v1.0.0) and store it under local cache. ([#4098](https://github.com/pybamm-team/PyBaMM/pull/4098))
 - Added `time` as an option for `Experiment.termination`. Now allows solving up to a user-specified time while also allowing different cycles and steps in an experiment to be handled normally. ([#4073](https://github.com/pybamm-team/PyBaMM/pull/4073))

pybamm/experiment/step/base_step.py

@@ -37,7 +37,7 @@
     ----------
     value : float
         The value of the step, corresponding to the type of step. Can be a number, a
-        2-tuple (for cccv_ode), or a 2-column array (for drive cycles)
+        2-tuple (for cccv_ode), a 2-column array (for drive cycles), or a 1-argument function of t
     duration : float, optional
         The duration of the step in seconds.
     termination : str or list, optional
@@ -73,6 +73,7 @@
     ):
         # Check if drive cycle
         self.is_drive_cycle = isinstance(value, np.ndarray)
+        is_python_function = callable(value)
         if self.is_drive_cycle:
             if value.ndim != 2 or value.shape[1] != 2:
                 raise ValueError(
@@ -83,6 +84,21 @@
             t = value[:, 0]
             if t[0] != 0:
                 raise ValueError("Drive cycle must start at t=0")
+        elif is_python_function:
+            t0 = 0
+            # Check if the function is only a function of t
+            try:
+                value_t0 = value(t0)
+            except TypeError:
+                raise TypeError(
+                    "Input function must have only 1 positional argument for time"
+                ) from None
+
+            # Check if the value at t0 is feasible
+            if not (np.isfinite(value_t0) and np.isscalar(value_t0)):
+                raise ValueError(
+                    f"Input function must return a real number output at t = {t0}"
+                )
 
         # Set duration
         if duration is None:
@@ -135,6 +151,10 @@
                 name="Drive Cycle",
             )
             self.period = np.diff(t).min()
+        elif is_python_function:
+            t = pybamm.t - pybamm.InputParameter("start time")
+            self.value = value(t)
+            self.period = _convert_time_to_seconds(period)
         else:
             self.value = value
             self.period = _convert_time_to_seconds(period)

pybamm/experiment/step/steps.py

@@ -426,6 +426,9 @@ def value_based_charge_or_discharge(step_value):
     elif isinstance(step_value, pybamm.Symbol):
         inpt = {"start time": 0}
         init_curr = step_value.evaluate(t=0, inputs=inpt).flatten()[0]
+    elif callable(step_value):
+        inpt = {"start time": 0}
+        init_curr = step_value(pybamm.t).evaluate(t=0, inputs=inpt).flatten()[0]
     else:
         init_curr = step_value
     sign = np.sign(init_curr)

tests/unit/test_simulation.py

@@ -365,6 +365,33 @@ def test_step_with_inputs(self):
             sim.solution.all_inputs[1]["Current function [A]"], 2
         )
 
+    def test_time_varying_input_function(self):
+        tf = 20.0
+
+        def oscillating(t):
+            return 3.6 + 0.1 * np.sin(2 * np.pi * t / tf)
+
+        model = pybamm.lithium_ion.SPM()
+
+        operating_modes = {
+            "Current [A]": pybamm.step.current,
+            "C-rate": pybamm.step.c_rate,
+            "Voltage [V]": pybamm.step.voltage,
+            "Power [W]": pybamm.step.power,
+        }
+        for name in operating_modes:
+            step = operating_modes[name](oscillating, duration=tf / 2)
+            experiment = pybamm.Experiment([step, step], period=f"{tf / 100} seconds")
+
+            solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8)
+            sim = pybamm.Simulation(model, experiment=experiment, solver=solver)
+            sim.solve()
+            for sol in sim.solution.sub_solutions:
+                t0 = sol.t[0]
+                np.testing.assert_array_almost_equal(
+                    sol[name].entries, np.array(oscillating(sol.t - t0))
+                )
+
     def test_save_load(self):
         with TemporaryDirectory() as dir_name:
             test_name = os.path.join(dir_name, "tests.pickle")