PyBaMM-GEM 0.1.0__py3-none-any.whl

pybamm_gem/__init__.py

@@ -0,0 +1,77 @@
+"""
+Copyright (c) 2026 Mohammed Asheruddin N. All rights reserved.
+
+PyBaMM-GEM: A PyBaMM-compatible extension package for advanced battery models,
+research models, and simulation tools.
+"""
+from __future__ import annotations
+
+from importlib.metadata import version
+
+import pybamm
+from pybamm_gem.entry_point import Model, models, parameter_sets
+
+__version__ = version("PyBaMM-GEM")
+
+
+def ParameterValues(name: str) -> pybamm.ParameterValues:
+    """
+    Return a PyBaMM ParameterValues object from a GEM parameter set.
+
+    Examples
+    --------
+    >>> import pybamm_gem
+    >>> params = pybamm_gem.ParameterValues("Imperial2026")
+    """
+    values = dict(parameter_sets[name])
+    values.pop("chemistry", None)
+    values.pop("notcite", None)
+    return pybamm.ParameterValues(values)
+
+
+def about() -> None:
+    """
+    Display information about the installed PyBaMM-GEM package.
+    """
+    print(
+    f"""
+════════════════════════════════════════════════════════════
+
+PyBaMM-GEM {__version__}
+
+Extending PyBaMM {pybamm.__version__} with advanced battery
+models, submodels, and parameter sets.
+
+https://github.com/mohammedasher/PyBaMM-GEM
+
+════════════════════════════════════════════════════════════
+""".strip()
+)
+
+def version() -> dict[str, str]:
+    """
+    Return version information for PyBaMM-GEM and the compatible PyBaMM installation.
+
+    Returns
+    -------
+    dict
+        Dictionary containing GEM and PyBaMM version information.
+    """
+    return {
+        "gem": __version__,
+        "pybamm": pybamm.__version__,
+    }
+
+__all__ = [
+    "__version__",
+    "pybamm",
+    "parameter_sets",
+    "ParameterValues",
+    "Model",
+    "models",
+    "about",
+    "version",
+]
+
+# Display package banner on import
+about()

pybamm_gem/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

pybamm_gem/entry_point.py

@@ -0,0 +1,154 @@
+"""
+This code is adopted from the PyBaMM project under the BSD-3-Clause
+
+Copyright (c) 2018-2024, the PyBaMM team.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+
+
+import importlib.metadata
+import sys
+import textwrap
+from collections.abc import Mapping
+from typing import Callable
+
+class EntryPoint(Mapping):
+    """
+    Dict-like interface for accessing parameter sets and models through entry points in copier template.
+    Access via :py:data:`pybamm_gem.parameter_sets` for parameter_sets
+    Access via :py:data:`pybamm_gem.Model` for Models
+
+    Examples
+    --------
+    Listing available parameter sets:
+        >>> import pybamm_gem
+        >>> list(pybamm_gem.parameter_sets)
+        ['Chen2020', ...]
+        >>> list(pybamm_gem.models)
+        ['SPM', ...]
+
+    Get the docstring for a parameter set/model:
+
+
+        >>> print(pybamm_gem.parameter_sets.get_docstring("Ai2020"))
+        <BLANKLINE>
+        Parameters for the Enertech cell (Ai2020), from the papers :footcite:t:`Ai2019`,
+        :footcite:t:`rieger2016new` and references therein.
+        ...
+
+        >>> print(pybamm_gem.models.get_docstring("SPM"))
+        <BLANKLINE>
+        Single Particle Model (SPM) model of a lithium-ion battery, from :footcite:t:`Marquis2019`. This class differs from the :class:`pybamm.lithium_ion.SPM` model class in that it shows the whole model in a single class. This comes at the cost of flexibility in combining different physical effects, and in general the main SPM class should be used instead.
+        ...
+        See also: :ref:`adding-parameter-sets`
+    """
+
+    _instances = 0
+    def __init__(self, group):
+        """Dict of entry points for parameter sets or models, lazily load entry points as"""
+        if not hasattr(self, 'initialized'):    # Ensure __init__ is called once per instance
+            self.initialized = True
+            EntryPoint._instances += 1
+            self._all_entries = dict()
+            self.group = group
+            for entry_point in self.get_entries(self.group):
+                self._all_entries[entry_point.name] = entry_point
+
+    @staticmethod
+    def get_entries(group_name):
+        """Wrapper for the importlib version logic"""
+        if sys.version_info < (3, 10):  # pragma: no cover
+            return importlib.metadata.entry_points()[group_name]
+        else:
+            return importlib.metadata.entry_points(group=group_name)
+
+    def __new__(cls, group):
+        """Ensure only two instances of entry points exist, one for parameter sets and the other for models"""
+        if EntryPoint._instances < 2:
+            cls.instance = super().__new__(cls)
+        return cls.instance
+
+    def __getitem__(self, key) -> dict:
+        return self._load_entry_point(key)()
+
+    def _load_entry_point(self, key) -> Callable:
+        """Check that ``key`` is a registered ``parameter_sets`` or ``models` ,
+        and return the entry point for the parameter set/model, loading it needed."""
+        if key not in self._all_entries:
+            raise KeyError(f"Unknown parameter set or model: {key}")
+        ps = self._all_entries[key]
+        try:
+            ps = self._all_entries[key] = ps.load()
+        except AttributeError:
+            pass
+        return ps
+
+    def __iter__(self):
+        return self._all_entries.__iter__()
+
+    def __len__(self) -> int:
+        return len(self._all_entries)
+
+    def get_docstring(self, key):
+        """Return the docstring for the ``key`` parameter set or model"""
+        return textwrap.dedent(self._load_entry_point(key).__doc__)
+
+    def __getattribute__(self, name):
+        try:
+            return super().__getattribute__(name)
+        except AttributeError as error:
+              raise error
+
+#: Singleton Instance of :class:ParameterSets """
+parameter_sets = EntryPoint(group="parameter_sets")
+
+#: Singleton Instance of :class:ModelEntryPoints"""
+models = EntryPoint(group="models")
+
+def Model(model:str):
+    """
+    Returns the loaded model object
+
+    Parameters
+    ----------
+    model : str
+        The model name or author name of the model mentioned at the model entry point.
+    Returns
+    -------
+    pybamm.model
+        Model object of the initialised model.
+    Examples
+    --------
+    Listing available models:
+        >>> import pybamm_gem
+        >>> list(pybamm_gem.models)
+        ['SPM', ...]
+        >>> pybamm_gem.Model('Author/Year')
+        <pybamm_gem.models.input.SPM.SPM object>
+    """
+    return models[model]

pybamm_gem/models/input/ExampleReservoir.py

@@ -0,0 +1,200 @@
+"""
+This code is adopted from the PyBaMM project under the BSD-3-Clause
+
+Copyright (c) 2018-2024, the PyBaMM team.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+
+
+#
+# Basic Reservoir Model
+#
+import pybamm
+
+class BasicReservoir(pybamm.lithium_ion.BaseModel):
+    """Reservoir model of a lithium-ion battery, from
+    :footcite:t:`Marquis2019`.
+
+    Parameters
+    ----------
+    name : str, optional
+        The name of the model.
+    """
+
+    def __init__(self, name="Reservoir Model"):
+        super().__init__({}, name)
+        pybamm.citations.register("Marquis2019")
+        # `param` is a class containing all the relevant parameters and functions for
+        # this model. These are purely symbolic at this stage, and will be set by the
+        # `ParameterValues` class when the model is processed.
+        param = self.param
+
+        ######################
+        # Variables
+        ######################
+        # Variables that depend on time only are created without a domain
+        Q = pybamm.Variable("Discharge capacity [A.h]")
+        # Variables that vary spatially are created with a domain
+        sto_n = pybamm.Variable(
+            "Average negative particle stoichiometry",
+            domain="current collector",
+            bounds=(0, 1),
+        )
+        sto_p = pybamm.Variable(
+            "Average positive particle stoichiometry",
+            domain="current collector",
+            bounds=(0, 1),
+        )
+
+        # Constant temperature
+        T = param.T_init
+
+        ######################
+        # Other set-up
+        ######################
+
+        # Current density
+        i_cell = param.current_density_with_time
+        a_n = 3 * param.n.prim.epsilon_s_av / param.n.prim.R_typ
+        a_p = 3 * param.p.prim.epsilon_s_av / param.p.prim.R_typ
+        j_n = i_cell / (param.n.L * a_n)
+        j_p = -i_cell / (param.p.L * a_p)
+
+        ######################
+        # State of Charge
+        ######################
+        I = param.current_with_time
+        # The `rhs` dictionary contains differential equations, with the key being the
+        # variable in the d/dt
+        self.rhs[Q] = I / 3600
+        # Initial conditions must be provided for the ODEs
+        self.initial_conditions[Q] = pybamm.Scalar(0)
+
+        ######################
+        # Particles
+        ######################
+
+        self.rhs[sto_n] = -i_cell / (
+            param.n.L * param.n.prim.epsilon_s_av * param.n.prim.c_max * param.F
+        )
+        self.rhs[sto_p] = i_cell / (
+            param.p.L * param.p.prim.epsilon_s_av * param.p.prim.c_max * param.F
+        )
+        # c_n_init and c_p_init are functions of r and x, but for the reservoir model
+        # we take the x-averaged and r-averaged value since there are no x-dependence
+        # nor r-dependencein the particles
+        self.initial_conditions[sto_n] = (
+            pybamm.x_average(pybamm.r_average(param.n.prim.c_init)) / param.n.prim.c_max
+        )
+        self.initial_conditions[sto_p] = (
+            pybamm.x_average(pybamm.r_average(param.p.prim.c_init)) / param.p.prim.c_max
+        )
+
+        self.events += [
+            pybamm.Event(
+                "Minimum negative particle surface stoichiometry",
+                sto_n - 0.01,
+            ),
+            pybamm.Event(
+                "Maximum negative particle surface stoichiometry",
+                (1 - 0.01) - sto_n,
+            ),
+            pybamm.Event(
+                "Minimum positive particle surface stoichiometry",
+                sto_p - 0.01,
+            ),
+            pybamm.Event(
+                "Maximum positive particle surface stoichiometry",
+                (1 - 0.01) - sto_p,
+            ),
+        ]
+
+        # Note that the reservoir model does not have any algebraic equations, so the
+        # `algebraic` dictionary remains empty
+
+        ######################
+        # (Some) variables
+        ######################
+        # Interfacial reactions
+        RT_F = param.R * T / param.F
+        j0_n = param.n.prim.j0(param.c_e_init_av, sto_n * param.n.prim.c_max, T)
+        j0_p = param.p.prim.j0(param.c_e_init_av, sto_p * param.p.prim.c_max, T)
+        eta_n = (2 / param.n.prim.ne) * RT_F * pybamm.arcsinh(j_n / (2 * j0_n))
+        eta_p = (2 / param.p.prim.ne) * RT_F * pybamm.arcsinh(j_p / (2 * j0_p))
+        phi_s_n = 0
+        phi_e = -eta_n - param.n.prim.U(sto_n, T)
+        phi_s_p = eta_p + phi_e + param.p.prim.U(sto_p, T)
+        V = phi_s_p
+        num_cells = pybamm.Parameter(
+            "Number of cells connected in series to make a battery"
+        )
+        c_s_n = sto_n * param.n.prim.c_max
+        c_s_p = sto_p * param.p.prim.c_max
+
+        whole_cell = ["negative electrode", "separator", "positive electrode"]
+        # The `variables` dictionary contains all variables that might be useful for
+        # visualising the solution of the model
+        # Primary broadcasts are used to broadcast scalar quantities across a domain
+        # into a vector of the right shape, for multiplying with other vectors
+        self.variables = {
+            "Time [s]": pybamm.t,
+            "Discharge capacity [A.h]": Q,
+            "X-averaged negative particle concentration [mol.m-3]": pybamm.PrimaryBroadcast(
+                c_s_n, "negative particle"
+            ),
+            "Negative particle surface "
+            "concentration [mol.m-3]": pybamm.PrimaryBroadcast(
+                c_s_n, "negative electrode"
+            ),
+            "Electrolyte concentration [mol.m-3]": pybamm.PrimaryBroadcast(
+                param.c_e_init_av, whole_cell
+            ),
+            "X-averaged positive particle concentration [mol.m-3]": pybamm.PrimaryBroadcast(
+                c_s_p, "positive particle"
+            ),
+            "Positive particle surface "
+            "concentration [mol.m-3]": pybamm.PrimaryBroadcast(
+                c_s_p, "positive electrode"
+            ),
+            "Current [A]": I,
+            "Current variable [A]": I,  # for compatibility with pybamm.Experiment
+            "Negative electrode potential [V]": pybamm.PrimaryBroadcast(
+                phi_s_n, "negative electrode"
+            ),
+            "Electrolyte potential [V]": pybamm.PrimaryBroadcast(phi_e, whole_cell),
+            "Positive electrode potential [V]": pybamm.PrimaryBroadcast(
+                phi_s_p, "positive electrode"
+            ),
+            "Voltage [V]": V,
+            "Battery voltage [V]": V * num_cells,
+        }
+        # Events specify points at which a solution should terminate
+        self.events += [
+            pybamm.Event("Minimum voltage [V]", V - param.voltage_low_cut),
+            pybamm.Event("Maximum voltage [V]", param.voltage_high_cut - V),
+        ]

pybamm_gem/models/input/ExampleSPM.py

@@ -0,0 +1,215 @@
+"""
+This code is adopted from the PyBaMM project under the BSD-3-Clause
+
+Copyright (c) 2018-2024, the PyBaMM team.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+
+
+#
+# Basic Single Particle Model (SPM)
+#
+import pybamm
+
+
+class SPM(pybamm.lithium_ion.BaseModel):
+    """Single Particle Model (SPM) model of a lithium-ion battery, from
+    :footcite:t:`Marquis2019`.
+
+    This class differs from the :class:`pybamm.lithium_ion.SPM` model class in that it
+    shows the whole model in a single class. This comes at the cost of flexibility in
+    combining different physical effects, and in general the main SPM class should be
+    used instead.
+
+    Parameters
+    ----------
+    name : str, optional
+        The name of the model.
+    """
+
+    def __init__(self, name="Single Particle Model"):
+        super().__init__({}, name)
+        pybamm.citations.register("Marquis2019")
+        # `param` is a class containing all the relevant parameters and functions for
+        # this model. These are purely symbolic at this stage, and will be set by the
+        # `ParameterValues` class when the model is processed.
+        param = self.param
+
+        ######################
+        # Variables
+        ######################
+        # Variables that depend on time only are created without a domain
+        Q = pybamm.Variable("Discharge capacity [A.h]")
+        # Variables that vary spatially are created with a domain
+        c_s_n = pybamm.Variable(
+            "X-averaged negative particle concentration [mol.m-3]",
+            domain="negative particle",
+        )
+        c_s_p = pybamm.Variable(
+            "X-averaged positive particle concentration [mol.m-3]",
+            domain="positive particle",
+        )
+
+        # Constant temperature
+        T = param.T_init
+
+        ######################
+        # Other set-up
+        ######################
+
+        # Current density
+        i_cell = param.current_density_with_time
+        a_n = 3 * param.n.prim.epsilon_s_av / param.n.prim.R_typ
+        a_p = 3 * param.p.prim.epsilon_s_av / param.p.prim.R_typ
+        j_n = i_cell / (param.n.L * a_n)
+        j_p = -i_cell / (param.p.L * a_p)
+
+        ######################
+        # State of Charge
+        ######################
+        I = param.current_with_time
+        # The `rhs` dictionary contains differential equations, with the key being the
+        # variable in the d/dt
+        self.rhs[Q] = I / 3600
+        # Initial conditions must be provided for the ODEs
+        self.initial_conditions[Q] = pybamm.Scalar(0)
+
+        ######################
+        # Particles
+        ######################
+
+        # The div and grad operators will be converted to the appropriate matrix
+        # multiplication at the discretisation stage
+        N_s_n = -param.n.prim.D(c_s_n, T) * pybamm.grad(c_s_n)
+        N_s_p = -param.p.prim.D(c_s_p, T) * pybamm.grad(c_s_p)
+        self.rhs[c_s_n] = -pybamm.div(N_s_n)
+        self.rhs[c_s_p] = -pybamm.div(N_s_p)
+        # Surf takes the surface value of a variable, i.e. its boundary value on the
+        # right side. This is also accessible via `boundary_value(x, "right")`, with
+        # "left" providing the boundary value of the left side
+        c_s_surf_n = pybamm.surf(c_s_n)
+        c_s_surf_p = pybamm.surf(c_s_p)
+        # Boundary conditions must be provided for equations with spatial derivatives
+        self.boundary_conditions[c_s_n] = {
+            "left": (pybamm.Scalar(0), "Neumann"),
+            "right": (
+                -j_n / (param.F * pybamm.surf(param.n.prim.D(c_s_n, T))),
+                "Neumann",
+            ),
+        }
+        self.boundary_conditions[c_s_p] = {
+            "left": (pybamm.Scalar(0), "Neumann"),
+            "right": (
+                -j_p / (param.F * pybamm.surf(param.p.prim.D(c_s_p, T))),
+                "Neumann",
+            ),
+        }
+        # c_n_init and c_p_init are functions of r and x, but for the SPM we
+        # take the x-averaged value since there is no x-dependence in the particles
+        self.initial_conditions[c_s_n] = pybamm.x_average(param.n.prim.c_init)
+        self.initial_conditions[c_s_p] = pybamm.x_average(param.p.prim.c_init)
+        # Events specify points at which a solution should terminate
+        sto_surf_n = c_s_surf_n / param.n.prim.c_max
+        sto_surf_p = c_s_surf_p / param.p.prim.c_max
+        self.events += [
+            pybamm.Event(
+                "Minimum negative particle surface stoichiometry",
+                pybamm.min(sto_surf_n) - 0.01,
+            ),
+            pybamm.Event(
+                "Maximum negative particle surface stoichiometry",
+                (1 - 0.01) - pybamm.max(sto_surf_n),
+            ),
+            pybamm.Event(
+                "Minimum positive particle surface stoichiometry",
+                pybamm.min(sto_surf_p) - 0.01,
+            ),
+            pybamm.Event(
+                "Maximum positive particle surface stoichiometry",
+                (1 - 0.01) - pybamm.max(sto_surf_p),
+            ),
+        ]
+
+        # Note that the SPM does not have any algebraic equations, so the `algebraic`
+        # dictionary remains empty
+
+        ######################
+        # (Some) variables
+        ######################
+        # Interfacial reactions
+        RT_F = param.R * T / param.F
+        j0_n = param.n.prim.j0(param.c_e_init_av, c_s_surf_n, T)
+        j0_p = param.p.prim.j0(param.c_e_init_av, c_s_surf_p, T)
+        eta_n = (2 / param.n.prim.ne) * RT_F * pybamm.arcsinh(j_n / (2 * j0_n))
+        eta_p = (2 / param.p.prim.ne) * RT_F * pybamm.arcsinh(j_p / (2 * j0_p))
+        phi_s_n = 0
+        phi_e = -eta_n - param.n.prim.U(sto_surf_n, T)
+        phi_s_p = eta_p + phi_e + param.p.prim.U(sto_surf_p, T)
+        V = phi_s_p
+        num_cells = pybamm.Parameter(
+            "Number of cells connected in series to make a battery"
+        )
+
+        whole_cell = ["negative electrode", "separator", "positive electrode"]
+        # The `variables` dictionary contains all variables that might be useful for
+        # visualising the solution of the model
+        # Primary broadcasts are used to broadcast scalar quantities across a domain
+        # into a vector of the right shape, for multiplying with other vectors
+        self.variables = {
+            "Time [s]": pybamm.t,
+            "Discharge capacity [A.h]": Q,
+            "X-averaged negative particle concentration [mol.m-3]": c_s_n,
+            "Negative particle surface "
+            "concentration [mol.m-3]": pybamm.PrimaryBroadcast(
+                c_s_surf_n, "negative electrode"
+            ),
+            "Electrolyte concentration [mol.m-3]": pybamm.PrimaryBroadcast(
+                param.c_e_init_av, whole_cell
+            ),
+            "X-averaged positive particle concentration [mol.m-3]": c_s_p,
+            "Positive particle surface "
+            "concentration [mol.m-3]": pybamm.PrimaryBroadcast(
+                c_s_surf_p, "positive electrode"
+            ),
+            "Current [A]": I,
+            "Current variable [A]": I,  # for compatibility with pybamm.Experiment
+            "Negative electrode potential [V]": pybamm.PrimaryBroadcast(
+                phi_s_n, "negative electrode"
+            ),
+            "Electrolyte potential [V]": pybamm.PrimaryBroadcast(phi_e, whole_cell),
+            "Positive electrode potential [V]": pybamm.PrimaryBroadcast(
+                phi_s_p, "positive electrode"
+            ),
+            "Voltage [V]": V,
+            "Battery voltage [V]": V * num_cells,
+        }
+        # Events specify points at which a solution should terminate
+        self.events += [
+            pybamm.Event("Minimum voltage [V]", V - param.voltage_low_cut),
+            pybamm.Event("Maximum voltage [V]", param.voltage_high_cut - V),
+        ]