pybamm-team#3530 add classes to handle termination

pybamm/__init__.py

@@ -228,6 +228,7 @@
 #
 from .experiment.experiment import Experiment
 from . import experiment
+from .experiment import step
 
 
 #

pybamm/experiment/experiment.py

@@ -3,7 +3,7 @@
 #
 
 import pybamm
-from pybamm.step._steps_util import (
+from .step._steps_util import (
     _convert_time_to_seconds,
     _convert_temperature_to_kelvin,
 )

pybamm/step/__init__.py → pybamm/experiment/step/__init__.py

@@ -1,2 +1,3 @@
 from .steps import *
 from .steps import _Step
+from .step_termination import *

pybamm/step/_steps_util.py → pybamm/experiment/step/_steps_util.py

@@ -4,6 +4,7 @@
 import pybamm
 import numpy as np
 from datetime import datetime
+from .step_termination import read_termination
 
 _examples = """
 
@@ -136,8 +137,10 @@ def __init__(
             termination = [termination]
         self.termination = []
         for term in termination:
-            typ, value = _convert_electric(term)
-            self.termination.append({"type": typ, "value": value})
+            if isinstance(term, str):
+                term = _convert_electric(term)
+            term = read_termination(term)
+            self.termination.append(term)
 
         self.temperature = _convert_temperature_to_kelvin(temperature)
 
@@ -193,10 +196,7 @@ def to_dict(self):
         }
 
     def __eq__(self, other):
-        return (
-            isinstance(other, _Step)
-            and self.hash_args == other.hash_args
-        )
+        return isinstance(other, _Step) and self.hash_args == other.hash_args
 
     def __hash__(self):
         return hash(self.basic_repr())

pybamm/experiment/step/step_termination.py

@@ -0,0 +1,78 @@
+import pybamm
+import numpy as np
+
+
+class BaseTermination:
+    def __init__(self, value):
+        self.value = value
+
+    def get_event(self, variables, step_value):
+        raise NotImplementedError
+
+
+class CrateTermination(BaseTermination):
+    def get_event(self, variables, step_value):
+        event = pybamm.Event(
+            "C-rate cut-off [A] [experiment]",
+            abs(variables["C-rate"]) - self.value,
+        )
+        return event
+
+
+class CurrentTermination(BaseTermination):
+    def get_event(self, variables, step_value):
+        event = pybamm.Event(
+            "Current cut-off [A] [experiment]",
+            abs(variables["Current [A]"]) - self.value,
+        )
+        return event
+
+
+class VoltageTermination(BaseTermination):
+    def get_event(self, variables, step_value):
+        # The voltage event should be positive at the start of charge/
+        # discharge. We use the sign of the current or power input to
+        # figure out whether the voltage event is greater than the starting
+        # voltage (charge) or less (discharge) and set the sign of the
+        # event accordingly
+        if isinstance(step_value, pybamm.Symbol):
+            inpt = {"start time": 0}
+            init_curr = step_value.evaluate(t=0, inputs=inpt).flatten()[0]
+        else:
+            init_curr = step_value
+        sign = np.sign(init_curr)
+        if sign > 0:
+            name = "Discharge"
+        else:
+            name = "Charge"
+        if sign != 0:
+            # Event should be positive at initial conditions for both
+            # charge and discharge
+            event = pybamm.Event(
+                f"{name} voltage cut-off [V] [experiment]",
+                sign * (variables["Battery voltage [V]"] - self.value),
+            )
+            return event
+
+
+class CustomTermination(BaseTermination):
+    def __init__(self, name, event_function):
+        self.name = name
+        self.event_function = event_function
+
+    def get_event(self, variables, step_value):
+        return pybamm.Event(self.name, self.event_function(variables))
+
+
+def read_termination(termination):
+    if isinstance(termination, tuple):
+        typ, value = termination
+    else:
+        return termination
+
+    termination_class = {
+        "current": CurrentTermination,
+        "voltage": VoltageTermination,
+        "C-rate": CrateTermination,
+    }[typ]
+    return termination_class(value)

pybamm/simulation.py

@@ -192,16 +192,6 @@ def set_up_and_parameterise_experiment(self):
                 op_conds.type = "current"
                 op_conds.value = op_conds.value * capacity
 
-            # Update terminations
-            termination = op_conds.termination
-            for term in termination:
-                term_type = term["type"]
-                if term_type == "C-rate":
-                    # Change type to current
-                    term["type"] = "current"
-                    # Scale C-rate with capacity to obtain current
-                    term["value"] = term["value"] * capacity
-
             # Add time to the experiment times
             dt = op_conds.duration
             if dt is None:
@@ -294,46 +284,9 @@ def set_up_and_parameterise_model_for_experiment(self):
 
     def update_new_model_events(self, new_model, op):
         for term in op.termination:
-            if term["type"] == "current":
-                new_model.events.append(
-                    pybamm.Event(
-                        "Current cut-off [A] [experiment]",
-                        abs(new_model.variables["Current [A]"]) - term["value"],
-                    )
-                )
-
-            # add voltage events to the model
-            if term["type"] == "voltage":
-                # The voltage event should be positive at the start of charge/
-                # discharge. We use the sign of the current or power input to
-                # figure out whether the voltage event is greater than the starting
-                # voltage (charge) or less (discharge) and set the sign of the
-                # event accordingly
-                if isinstance(op.value, pybamm.Interpolant) or isinstance(
-                    op.value, pybamm.Multiplication
-                ):
-                    inpt = {"start time": 0}
-                    init_curr = op.value.evaluate(t=0, inputs=inpt).flatten()[0]
-                    sign = np.sign(init_curr)
-                else:
-                    sign = np.sign(op.value)
-                if sign > 0:
-                    name = "Discharge"
-                else:
-                    name = "Charge"
-                if sign != 0:
-                    # Event should be positive at initial conditions for both
-                    # charge and discharge
-                    new_model.events.append(
-                        pybamm.Event(
-                            f"{name} voltage cut-off [V] [experiment]",
-                            sign
-                            * (
-                                new_model.variables["Battery voltage [V]"]
-                                - term["value"]
-                            ),
-                        )
-                    )
+            event = term.get_event(new_model.variables, op.value)
+            if event is not None:
+                new_model.events.append(event)
 
         # Keep the min and max voltages as safeguards but add some tolerances
         # so that they are not triggered before the voltage limits in the