add operator to experiment terminations (#4770)

add operator to experiment terminations

CHANGELOG.md

@@ -1,5 +1,8 @@
 # [Unreleased](https://github.com/pybamm-team/PyBaMM/)
 
+## Features
+- Added Operators to current and voltage termination events. ([#4770](https://github.com/pybamm-team/PyBaMM/pull/4770))
+
 # [v25.1.0](https://github.com/pybamm-team/PyBaMM/tree/v25.1.0) - 2025-01-14
 
 ## Features

src/pybamm/experiment/step/__init__.py

@@ -1,5 +1,5 @@
 from .steps import *
 from .base_step import BaseStep, BaseStepExplicit, BaseStepImplicit
-from .step_termination import *
+from .step_termination import BaseTermination, CurrentTermination, VoltageTermination, CustomTermination, CrateTermination, _read_termination
 
 __all__ = ['base_step', 'step_termination', 'steps']

src/pybamm/experiment/step/base_step.py

@@ -67,8 +67,14 @@ def __init__(
         tags=None,
         start_time=None,
         description=None,
-        direction=None,
+        direction: str | None = None,
     ):
+        potential_directions = ["charge", "discharge", "rest", None]
+        if direction not in potential_directions:
+            raise ValueError(
+                f"Invalid direction: {direction}. Must be one of {potential_directions}"
+            )
+        self.input_duration = duration
         self.input_duration = duration
         self.input_value = value
         # Check if drive cycle
@@ -386,11 +392,11 @@ def value_based_charge_or_discharge(self):
             init_curr = self.value
         sign = np.sign(init_curr)
         if sign == 0:
-            return "Rest"
+            return "rest"
         elif sign > 0:
-            return "Discharge"
+            return "discharge"
         else:
-            return "Charge"
+            return "charge"
 
     def record_tags(
         self,

src/pybamm/experiment/step/step_termination.py

@@ -14,8 +14,11 @@ class BaseTermination:
         The value at which the event is triggered
     """
 
-    def __init__(self, value):
+    def __init__(self, value, operator=None):
         self.value = value
+        if operator not in ["<", ">", None]:
+            raise ValueError(f"Invalid operator: {operator}")
+        self.operator = operator
 
     def get_event(self, variables, step):
         """
@@ -67,9 +70,19 @@ def get_event(self, variables, step):
         """
         See :meth:`BaseTermination.get_event`
         """
+        operator = self.operator
+        if operator == ">":
+            expr = self.value - variables["Current [A]"]
+            event_string = f"Current [A] > {self.value} [A] [experiment]"
+        elif operator == "<":
+            expr = variables["Current [A]"] - self.value
+            event_string = f"Current [A] < {self.value} [A] [experiment]"
+        else:
+            expr = abs(variables["Current [A]"]) - self.value
+            event_string = f"abs(Current [A]) < {self.value} [A] [experiment]"
         event = pybamm.Event(
-            "Current cut-off [A] [experiment]",
-            abs(variables["Current [A]"]) - self.value,
+            event_string,
+            expr,
         )
         return event
 
@@ -89,24 +102,48 @@ def get_event(self, variables, step):
         # figure out whether the voltage event is greater than the starting
         # voltage (charge) or less (discharge) and set the sign of the
         # event accordingly
-        direction = step.direction.capitalize()
-        if direction == "Charge":
+        operator = self.operator
+        if operator is None:
+            direction = step.direction
+            if direction == "charge":
+                operator = ">"
+            elif direction == "discharge":
+                operator = "<"
+            else:
+                # No event for rest steps
+                return None
+
+        if operator == ">":
             sign = -1
-        elif direction == "Discharge":
+        else:
+            # operator can only be "<" or ">"
             sign = 1
-        elif direction == "Rest":
-            # No event for rest steps
-            return None
 
         # Event should be positive at initial conditions for both
         # charge and discharge
         event = pybamm.Event(
-            f"{direction} voltage cut-off [V] [experiment]",
+            f"Voltage {operator} {self.value} [V] [experiment]",
             sign * (variables["Battery voltage [V]"] - self.value),
         )
         return event
 
 
+class Voltage:
+    def __gt__(self, value):
+        return VoltageTermination(value, operator=">")
+
+    def __lt__(self, value):
+        return VoltageTermination(value, operator="<")
+
+
+class Current:
+    def __gt__(self, value):
+        return CurrentTermination(value, operator=">")
+
+    def __lt__(self, value):
+        return CurrentTermination(value, operator="<")
+
+
 class CustomTermination(BaseTermination):
     """
     Define a custom termination event using a function. This can be used to create an

tests/unit/test_experiments/test_experiment_steps.py

@@ -307,3 +307,42 @@ def custom_step_voltage(variables):
 
         with pytest.raises(ValueError, match="control must be"):
             pybamm.step.CustomStepImplicit(custom_step_voltage, control="bla")
+
+    def test_bad_direction(self):
+        with pytest.raises(ValueError, match="Invalid direction"):
+            pybamm.step.Voltage(4.1, direction="foo")
+
+    def test_steps_with_operators(self):
+        # voltage
+        step = pybamm.step.voltage(1, duration=3600)
+        termination_lt_4_1 = pybamm.step.VoltageTermination(4.1, operator="<")
+        termination_gt_4_1 = pybamm.step.VoltageTermination(4.1, operator=">")
+        variables = {"Battery voltage [V]": 4.2}
+        event_lt_4_1 = termination_lt_4_1.get_event(variables, step)
+        np.testing.assert_allclose(event_lt_4_1.expression, 4.2 - 4.1)
+        event_gt_4_1 = termination_gt_4_1.get_event(variables, step)
+        np.testing.assert_allclose(event_gt_4_1.expression, 4.1 - 4.2)
+
+        # current
+        step = pybamm.step.current(1, duration=3600)
+        termination_lt_0_05 = pybamm.step.CurrentTermination(0.05, operator="<")
+        termination_gt_0_05 = pybamm.step.CurrentTermination(0.05, operator=">")
+        variables = {"Current [A]": 0.06}
+        event_lt_0_05 = termination_lt_0_05.get_event(variables, step)
+        np.testing.assert_allclose(event_lt_0_05.expression, 0.06 - 0.05)
+        event_gt_0_05 = termination_gt_0_05.get_event(variables, step)
+        np.testing.assert_allclose(event_gt_0_05.expression, 0.05 - 0.06)
+
+        # error
+        with pytest.raises(ValueError, match="Invalid operator"):
+            pybamm.step.CurrentTermination(0.05, operator="=")
+
+        # operator overloading
+        termination_lt_0_05_oo = pybamm.step.step_termination.Current() < 0.05
+        termination_gt_0_05_oo = pybamm.step.step_termination.Current() > 0.05
+        assert termination_lt_0_05_oo == termination_gt_0_05
+        assert termination_gt_0_05_oo == termination_gt_0_05
+        termination_lt_4_1_oo = pybamm.step.step_termination.Voltage() < 4.1
+        termination_gt_4_1_oo = pybamm.step.step_termination.Voltage() > 4.1
+        assert termination_lt_4_1_oo == termination_gt_4_1
+        assert termination_gt_4_1_oo == termination_gt_4_1

tests/unit/test_experiments/test_simulation_with_experiment.py

@@ -33,10 +33,10 @@ def test_set_up(self):
             steps[1].basic_repr()
         ]  # CC charge
         model_V = sim.experiment_unique_steps_to_model[steps[2].basic_repr()]  # CV hold
-        assert "Current cut-off [A] [experiment]" in [
+        assert "abs(Current [A]) < 0.05 [A] [experiment]" in [
             event.name for event in model_V.events
         ]
-        assert "Charge voltage cut-off [V] [experiment]" in [
+        assert "Voltage > 4.1 [V] [experiment]" in [
             event.name for event in model_I.events
         ]