Issue 3530 custom termination

CHANGELOG.md

@@ -2,10 +2,12 @@
 
 ## Features
 
+- Added method to get QuickPlot axes by variable ([#3596](https://github.com/pybamm-team/PyBaMM/pull/3596))
+- Added custom experiment terminations ([#3596](https://github.com/pybamm-team/PyBaMM/pull/3596))
+- Mechanical parameters are now a function of stoichiometry and temperature ([#3576](https://github.com/pybamm-team/PyBaMM/pull/3576))
 - Added a new unary operator, `EvaluateAt`, that evaluates a spatial variable at a given position ([#3573](https://github.com/pybamm-team/PyBaMM/pull/3573))
 - Added a method, `insert_reference_electrode`, to `pybamm.lithium_ion.BaseModel` that insert a reference electrode to measure the electrolyte potential at a given position in space and adds new variables that mimic a 3E cell setup. ([#3573](https://github.com/pybamm-team/PyBaMM/pull/3573))
 - Serialisation added so models can be written to/read from JSON ([#3397](https://github.com/pybamm-team/PyBaMM/pull/3397))
-- Mechanical parameters are now a function of stoichiometry and temperature ([#3576](https://github.com/pybamm-team/PyBaMM/pull/3576))
 
 ## Bug fixes
 

docs/source/api/experiment/experiment_steps.rst

@@ -18,3 +18,28 @@ directly:
 
 .. autoclass:: pybamm.step._Step
     :members:
+
+Step terminations
+-----------------
+
+Standard step termination events are implemented by the following classes, which are
+called when the termination is specified by a specific string. These classes can be
+either be called directly or via the string format specified in the class docstring
+
+.. autoclass:: pybamm.step.CrateTermination
+    :members:
+
+.. autoclass:: pybamm.step.CurrentTermination
+    :members:
+
+.. autoclass:: pybamm.step.VoltageTermination
+    :members:
+
+The following classes can be used to define custom terminations for an experiment
+step:
+
+.. autoclass:: pybamm.step.BaseTermination
+    :members:
+
+.. autoclass:: pybamm.step.CustomTermination
+    :members:

docs/source/api/plotting/quick_plot.rst

@@ -7,3 +7,6 @@ Quick Plot
   :members:
 
 .. autofunction:: pybamm.dynamic_plot
+
+.. autoclass:: pybamm.QuickPlotAxes
+  :members:

docs/source/examples/notebooks/getting_started/tutorial-1-how-to-run-a-model.ipynb

@@ -205,7 +205,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.0"
+   "version": "3.11.3"
   }
  },
  "nbformat": 4,

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,157 @@
+import pybamm
+import numpy as np
+
+
+class BaseTermination:
+    """
+    Base class for a termination event for an experiment step. To create a custom
+    termination, a class must implement `get_event` to return a :class:`pybamm.Event`
+    corresponding to the desired termination. In most cases the class
+    :class:`pybamm.step.CustomTermination` can be used to assist with this.
+
+    Parameters
+    ----------
+    value : float
+        The value at which the event is triggered
+    """
+
+    def __init__(self, value):
+        self.value = value
+
+    def get_event(self, variables, step_value):
+        """
+        Return a :class:`pybamm.Event` object corresponding to the termination event
+
+        Parameters
+        ----------
+        variables : dict
+            Dictionary of model variables, to be used for selecting the variable(s) that
+            determine the event
+        step_value : float or :class:`pybamm.Symbol`
+            Value of the step for which this is a termination event, to be used in some
+            cases to determine the sign of the event.
+        """
+        raise NotImplementedError
+
+
+class CrateTermination(BaseTermination):
+    """
+    Termination based on C-rate, created when a string termination of the C-rate type
+    (e.g. "C/10") is provided
+    """
+
+    def get_event(self, variables, step_value):
+        """
+        See :meth:`BaseTermination.get_event`
+        """
+        event = pybamm.Event(
+            "C-rate cut-off [experiment]",
+            abs(variables["C-rate"]) - self.value,
+        )
+        return event
+
+
+class CurrentTermination(BaseTermination):
+    """
+    Termination based on current, created when a string termination of the current type
+    (e.g. "1A") is provided
+    """
+
+    def get_event(self, variables, step_value):
+        """
+        See :meth:`BaseTermination.get_event`
+        """
+        event = pybamm.Event(
+            "Current cut-off [A] [experiment]",
+            abs(variables["Current [A]"]) - self.value,
+        )
+        return event
+
+
+class VoltageTermination(BaseTermination):
+    """
+    Termination based on voltage, created when a string termination of the voltage type
+    (e.g. "4.2V") is provided
+    """
+
+    def get_event(self, variables, step_value):
+        """
+        See :meth:`BaseTermination.get_event`
+        """
+        # 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):
+    """
+    Define a custom termination event using a function. This can be used to create an
+    event based on any variable in the model.
+
+    Parameters
+    ----------
+    name : str
+        Name of the event
+    event_function : callable
+        A function that takes in a dictionary of variables and evaluates the event
+        value. Must be positive before the event is triggered and zero when the
+        event is triggered.
+
+    Example
+    -------
+    Add a cut-off based on negative electrode stoichiometry. The event will trigger
+    when the negative electrode stoichiometry reaches 10%.
+
+    >>> def neg_stoich_cutoff(variables):
+    >>>    return variables["Negative electrode stoichiometry"] - 0.1
+
+    >>> neg_stoich_termination = pybamm.step.CustomTermination(
+    >>>    name="Negative stoichiometry cut-off", event_function=neg_stoich_cutoff
+    >>> )
+    """
+
+    def __init__(self, name, event_function):
+        if not name.endswith(" [experiment]"):
+            name += " [experiment]"
+        self.name = name
+        self.event_function = event_function
+
+    def get_event(self, variables, step_value):
+        """
+        See :meth:`BaseTermination.get_event`
+        """
+        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/plotting/quick_plot.py

@@ -486,14 +486,14 @@ def plot(self, t, dynamic=False):
         self.plots = {}
         self.time_lines = {}
         self.colorbars = {}
-        self.axes = []
+        self.axes = QuickPlotAxes()
 
         # initialize empty handles, to be created only if the appropriate plots are made
         solution_handles = []
 
         for k, (key, variable_lists) in enumerate(self.variables.items()):
             ax = self.fig.add_subplot(self.gridspec[k])
-            self.axes.append(ax)
+            self.axes.add(key, ax)
             x_min, x_max, y_min, y_max = self.axis_limits[key]
             ax.set_xlim(x_min, x_max)
             if y_min is not None and y_max is not None:
@@ -803,3 +803,40 @@ def create_gif(self, number_of_images=80, duration=0.1, output_filename="plot.gi
         # remove the generated images
         for image in images:
             os.remove(image)
+
+
+class QuickPlotAxes:
+    """
+    Class to store axes for the QuickPlot
+    """
+
+    _by_variable = {}
+    _axes = []
+
+    def add(self, keys, axis):
+        """
+        Add axis
+
+        Parameters
+        ----------
+        keys : iter
+            Iterable of keys of variables being plotted on the axis
+        axis : matplotlib Axis object
+            The axis object
+        """
+        self._axes.append(axis)
+        for k in keys:
+            self._by_variable[k] = axis
+
+    def __getitem__(self, index):
+        """
+        Get axis by index
+        """
+        return self._axes[index]
+
+    @property
+    def by_variable(self, key):
+        """
+        Get axis by variable name
+        """
+        return self._by_variable[key]

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