Merge branch 'develop' into telemetry

valentinsulzer · valentinsulzer · commit 215252b45265 · 2024-11-07T12:56:25.000-05:00
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -4,7 +4,7 @@ ci:
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.7.1"
+    rev: "v0.7.2"
     hooks:
       - id: ruff
         args: [--fix, --show-fixes]
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -44,7 +44,7 @@ You now have everything you need to start making changes!
 
 ### B. Writing your code
 
-6. PyBaMM is developed in [Python](https://www.python.org)), and makes heavy use of [NumPy](https://numpy.org/) (see also [NumPy for MatLab users](https://numpy.org/doc/stable/user/numpy-for-matlab-users.html) and [Python for R users](https://www.rebeccabarter.com/blog/2023-09-11-from_r_to_python)).
+6. PyBaMM is developed in [Python](https://www.python.org), and makes heavy use of [NumPy](https://numpy.org/).
 7. Make sure to follow our [coding style guidelines](#coding-style-guidelines).
 8. Commit your changes to your branch with [useful, descriptive commit messages](https://chris.beams.io/posts/git-commit/): Remember these are
    publicly visible and should still make sense a few months ahead in time.
@@ -116,8 +116,8 @@ PyBaMM provides a utility function `import_optional_dependency`, to check for th
 
 Optional dependencies should never be imported at the module level, but always inside methods. For example:
 
-```
-def use_pybtex(x,y,z):
+```python
+def use_pybtex(x, y, z):
     pybtex = import_optional_dependency("pybtex")
     ...
 ```
diff --git a/docs/source/examples/index.rst b/docs/source/examples/index.rst
@@ -87,6 +87,7 @@ The notebooks are organised into subfolders, and can be viewed in the galleries
     notebooks/parameterization/change-input-current.ipynb
     notebooks/parameterization/parameter-values.ipynb
     notebooks/parameterization/parameterization.ipynb
+    notebooks/parameterization/sensitivities_and_data_fitting.ipynb
 
 .. nbgallery::
     :caption: Simulations and Experiments
diff --git a/docs/source/examples/notebooks/parameterization/sensitivities_and_data_fitting.ipynb b/docs/source/examples/notebooks/parameterization/sensitivities_and_data_fitting.ipynb
diff --git a/src/pybamm/__init__.py b/src/pybamm/__init__.py
@@ -39,6 +39,7 @@
 from .expression_tree.parameter import Parameter, FunctionParameter
 from .expression_tree.scalar import Scalar
 from .expression_tree.variable import *
+from .expression_tree.coupled_variable import *
 from .expression_tree.independent_variable import *
 from .expression_tree.independent_variable import t
 from .expression_tree.vector import Vector
diff --git a/src/pybamm/discretisations/discretisation.py b/src/pybamm/discretisations/discretisation.py
@@ -500,8 +500,8 @@ def check_tab_conditions(self, symbol, bcs):
 
         if domain != "current collector":
             raise pybamm.ModelError(
-                f"""Boundary conditions can only be applied on the tabs in the domain
-            'current collector', but {symbol} has domain {domain}"""
+                "Boundary conditions can only be applied on the tabs in the domain "
+                f"'current collector', but {symbol} has domain {domain}"
             )
         # Replace keys with "left" and "right" as appropriate for 1D meshes
         if isinstance(mesh, pybamm.SubMesh1D):
@@ -893,11 +893,9 @@ def _process_symbol(self, symbol):
                 y_slices = self.y_slices[symbol]
             except KeyError as error:
                 raise pybamm.ModelError(
-                    f"""
-                    No key set for variable '{symbol.name}'. Make sure it is included in either
-                    model.rhs or model.algebraic in an unmodified form
-                    (e.g. not Broadcasted)
-                    """
+                    f"No key set for variable '{symbol.name}'. Make sure it is included in either "
+                    "model.rhs or model.algebraic in an unmodified form "
+                    "(e.g. not Broadcasted)"
                 ) from error
             # Add symbol's reference and multiply by the symbol's scale
             # so that the state vector is of order 1
@@ -938,6 +936,11 @@ def _process_symbol(self, symbol):
             if symbol._expected_size is None:
                 symbol._expected_size = expected_size
             return symbol.create_copy()
+
+        elif isinstance(symbol, pybamm.CoupledVariable):
+            new_symbol = self.process_symbol(symbol.children[0])
+            return new_symbol
+
         else:
             # Backup option: return the object
             return symbol
diff --git a/src/pybamm/expression_tree/averages.py b/src/pybamm/expression_tree/averages.py
@@ -251,8 +251,8 @@ def z_average(symbol: pybamm.Symbol) -> pybamm.Symbol:
     # Symbol must have domain [] or ["current collector"]
     if symbol.domain not in [[], ["current collector"]]:
         raise pybamm.DomainError(
-            f"""z-average only implemented in the 'current collector' domain,
-            but symbol has domains {symbol.domain}"""
+            "z-average only implemented in the 'current collector' domain, "
+            f"but symbol has domains {symbol.domain}"
         )
     # If symbol doesn't have a domain, its average value is itself
     if symbol.domain == []:
@@ -285,8 +285,8 @@ def yz_average(symbol: pybamm.Symbol) -> pybamm.Symbol:
     # Symbol must have domain [] or ["current collector"]
     if symbol.domain not in [[], ["current collector"]]:
         raise pybamm.DomainError(
-            f"""y-z-average only implemented in the 'current collector' domain,
-            but symbol has domains {symbol.domain}"""
+            "y-z-average only implemented in the 'current collector' domain, "
+            f"but symbol has domains {symbol.domain}"
         )
     # If symbol doesn't have a domain, its average value is itself
     if symbol.domain == []:
diff --git a/src/pybamm/expression_tree/binary_operators.py b/src/pybamm/expression_tree/binary_operators.py
@@ -36,7 +36,7 @@ def _preprocess_binary(
     # Check both left and right are pybamm Symbols
     if not (isinstance(left, pybamm.Symbol) and isinstance(right, pybamm.Symbol)):
         raise NotImplementedError(
-            f"""BinaryOperator not implemented for symbols of type {type(left)} and {type(right)}"""
+            f"BinaryOperator not implemented for symbols of type {type(left)} and {type(right)}"
         )
 
     # Do some broadcasting in special cases, to avoid having to do this manually
@@ -389,9 +389,9 @@ def _binary_jac(self, left_jac, right_jac):
             return left @ right_jac
         else:
             raise NotImplementedError(
-                f"""jac of 'MatrixMultiplication' is only
-             implemented for left of type 'pybamm.Array',
-             not {left.__class__}"""
+                f"jac of 'MatrixMultiplication' is only "
+                "implemented for left of type 'pybamm.Array', "
+                f"not {left.__class__}"
             )
 
     def _binary_evaluate(self, left, right):
@@ -1541,8 +1541,8 @@ def source(
 
     if left.domain != ["current collector"] or right.domain != ["current collector"]:
         raise pybamm.DomainError(
-            f"""'source' only implemented in the 'current collector' domain,
-            but symbols have domains {left.domain} and {right.domain}"""
+            "'source' only implemented in the 'current collector' domain, "
+            f"but symbols have domains {left.domain} and {right.domain}"
         )
     if boundary:
         return pybamm.BoundaryMass(right) @ left
diff --git a/src/pybamm/expression_tree/coupled_variable.py b/src/pybamm/expression_tree/coupled_variable.py
@@ -0,0 +1,55 @@
+import pybamm
+
+from pybamm.type_definitions import DomainType
+
+
+class CoupledVariable(pybamm.Symbol):
+    """
+    A node in the expression tree representing a variable whose equation is set by a different model or submodel.
+
+
+    Parameters
+    ----------
+    name : str
+        name of the node
+    domain : iterable of str
+        list of domains that this coupled variable is valid over
+    """
+
+    def __init__(
+        self,
+        name: str,
+        domain: DomainType = None,
+    ) -> None:
+        super().__init__(name, domain=domain)
+
+    def _evaluate_for_shape(self):
+        """
+        Returns the scalar 'NaN' to represent the shape of a parameter.
+        See :meth:`pybamm.Symbol.evaluate_for_shape()`
+        """
+        return pybamm.evaluate_for_shape_using_domain(self.domains)
+
+    def create_copy(self):
+        """Creates a new copy of the coupled variable."""
+        new_coupled_variable = CoupledVariable(self.name, self.domain)
+        return new_coupled_variable
+
+    @property
+    def children(self):
+        return self._children
+
+    @children.setter
+    def children(self, expr):
+        self._children = expr
+
+    def set_coupled_variable(self, symbol, expr):
+        """Sets the children of the coupled variable to the expression passed in expr. If the symbol is not the coupled variable, then it searches the children of the symbol for the coupled variable. The coupled variable will be replaced by its first child (symbol.children[0], which should be expr) in the discretisation step."""
+        if self == symbol:
+            symbol.children = [
+                expr,
+            ]
+        else:
+            for child in symbol.children:
+                self.set_coupled_variable(child, expr)
+        symbol.set_id()
diff --git a/src/pybamm/expression_tree/operations/convert_to_casadi.py b/src/pybamm/expression_tree/operations/convert_to_casadi.py
@@ -7,6 +7,7 @@
 import casadi
 import numpy as np
 from scipy import special
+from scipy import interpolate
 
 
 class CasadiConverter:
@@ -165,6 +166,18 @@ def _convert(self, symbol, t, y, y_dot, inputs):
                             # for some reason, pybamm.Interpolant always returns a column vector, so match that
                             test = test.T
                         return test
+                    elif solver == "bspline":
+                        bspline = interpolate.make_interp_spline(
+                            symbol.x[0], symbol.y, k=3
+                        )
+                        knots = [bspline.t]
+                        coeffs = bspline.c.flatten()
+                        degree = [bspline.k]
+                        m = len(coeffs) // len(symbol.x[0])
+                        f = casadi.Function.bspline(
+                            symbol.name, knots, coeffs, degree, m
+                        )
+                        return f(converted_children[0])
                     else:
                         return casadi.interpolant(
                             "LUT", solver, symbol.x, symbol.y.flatten()
@@ -176,6 +189,20 @@ def _convert(self, symbol, t, y, y_dot, inputs):
                             symbol.y.ravel(order="F"),
                             converted_children,
                         )
+                    elif solver == "bspline" and len(converted_children) == 2:
+                        bspline = interpolate.RectBivariateSpline(
+                            symbol.x[0], symbol.x[1], symbol.y
+                        )
+                        [tx, ty, c] = bspline.tck
+                        [kx, ky] = bspline.degrees
+                        knots = [tx, ty]
+                        coeffs = c
+                        degree = [kx, ky]
+                        m = 1
+                        f = casadi.Function.bspline(
+                            symbol.name, knots, coeffs, degree, m
+                        )
+                        return f(casadi.hcat(converted_children).T).T
                     else:
                         LUT = casadi.interpolant(
                             "LUT", solver, symbol.x, symbol.y.ravel(order="F")
@@ -231,8 +258,6 @@ def _convert(self, symbol, t, y, y_dot, inputs):
 
         else:
             raise TypeError(
-                f"""
-                Cannot convert symbol of type '{type(symbol)}' to CasADi. Symbols must all be
-                'linear algebra' at this stage.
-                """
+                f"Cannot convert symbol of type '{type(symbol)}' to CasADi. Symbols must all be "
+                "'linear algebra' at this stage."
             )
diff --git a/src/pybamm/expression_tree/unary_operators.py b/src/pybamm/expression_tree/unary_operators.py
@@ -992,8 +992,8 @@ def __init__(self, name, child, side):
         if side in ["negative tab", "positive tab"]:
             if child.domain[0] != "current collector":
                 raise pybamm.ModelError(
-                    f"""Can only take boundary value on the tabs in the domain
-                'current collector', but {child} has domain {child.domain[0]}"""
+                    "Can only take boundary value on the tabs in the domain "
+                    f"'current collector', but {child} has domain {child.domain[0]}"
                 )
         self.side = side
         # boundary value of a child takes the primary domain from secondary domain
diff --git a/src/pybamm/expression_tree/vector.py b/src/pybamm/expression_tree/vector.py
@@ -29,9 +29,7 @@ def __init__(
             entries = entries[:, np.newaxis]
         if entries.shape[1] != 1:
             raise ValueError(
-                f"""
-                Entries must have 1 dimension or be column vector, not have shape {entries.shape}
-                """
+                f"Entries must have 1 dimension or be column vector, not have shape {entries.shape}"
             )
         if name is None:
             name = f"Column vector of length {entries.shape[0]!s}"
diff --git a/src/pybamm/meshes/one_dimensional_submeshes.py b/src/pybamm/meshes/one_dimensional_submeshes.py
@@ -297,8 +297,8 @@ def __init__(self, lims, npts, edges=None):
 
         if (npts + 1) != len(edges):
             raise pybamm.GeometryError(
-                f"""User-suppled edges has should have length (npts + 1) but has length
-                {len(edges)}.Number of points (npts) for domain {spatial_var.domain} is {npts}.""".replace(
+                "User-suppled edges has should have length (npts + 1) but has length "
+                f"{len(edges)}.Number of points (npts) for domain {spatial_var.domain} is {npts}.".replace(
                     "\n                ", " "
                 )
             )
diff --git a/src/pybamm/meshes/scikit_fem_submeshes.py b/src/pybamm/meshes/scikit_fem_submeshes.py
@@ -92,8 +92,8 @@ def read_lims(self, lims):
         # check coordinate system agrees
         if spatial_vars[0].coord_sys != spatial_vars[1].coord_sys:
             raise pybamm.DomainError(
-                f"""spatial variables should have the same coordinate system,
-                but have coordinate systems {spatial_vars[0].coord_sys} and {spatial_vars[1].coord_sys}"""
+                "spatial variables should have the same coordinate system, "
+                f"but have coordinate systems {spatial_vars[0].coord_sys} and {spatial_vars[1].coord_sys}"
             )
         return spatial_vars, tabs
 
@@ -360,9 +360,9 @@ def __init__(self, lims, npts, y_edges=None, z_edges=None):
             # check that npts equals number of user-supplied edges
             if npts[var.name] != len(edges[var.name]):
                 raise pybamm.GeometryError(
-                    f"""User-suppled edges has should have length npts but has length {len(edges[var.name])}.
-                     Number of points (npts) for variable {var.name} in
-                     domain {var.domain} is {npts[var.name]}."""
+                    f"User-supplied edges has should have length npts but has length {len(edges[var.name])}. "
+                    f"Number of points (npts) for variable {var.name} in "
+                    f"domain {var.domain} is {npts[var.name]}."
                 )
 
             # check end points of edges agree with spatial_lims
diff --git a/src/pybamm/models/base_model.py b/src/pybamm/models/base_model.py
@@ -56,6 +56,7 @@ def __init__(self, name="Unnamed model"):
         self._boundary_conditions = {}
         self._variables_by_submodel = {}
         self._variables = pybamm.FuzzyDict({})
+        self._coupled_variables = {}
         self._summary_variables = []
         self._events = []
         self._concatenated_rhs = None
@@ -182,6 +183,31 @@ def boundary_conditions(self):
     def boundary_conditions(self, boundary_conditions):
         self._boundary_conditions = BoundaryConditionsDict(boundary_conditions)
 
+    @property
+    def coupled_variables(self):
+        """Returns a dictionary mapping strings to expressions representing variables needed by the model but whose equations were set by other models."""
+        return self._coupled_variables
+
+    @coupled_variables.setter
+    def coupled_variables(self, coupled_variables):
+        for name, var in coupled_variables.items():
+            if (
+                isinstance(var, pybamm.CoupledVariable)
+                and var.name != name
+                # Exception if the variable is also there under its own name
+                and not (
+                    var.name in coupled_variables and coupled_variables[var.name] == var
+                )
+            ):
+                raise ValueError(
+                    f"Coupled variable with name '{var.name}' is in coupled variables dictionary with "
+                    f"name '{name}'. Names must match."
+                )
+        self._coupled_variables = coupled_variables
+
+    def list_coupled_variables(self):
+        return list(self._coupled_variables.keys())
+
     @property
     def variables(self):
         """Returns a dictionary mapping strings to expressions representing the model's useful variables."""
@@ -1084,7 +1110,7 @@ def check_ics_bcs(self):
         for var in self.rhs.keys():
             if var not in self.initial_conditions.keys():
                 raise pybamm.ModelError(
-                    f"""no initial condition given for variable '{var}'"""
+                    f"no initial condition given for variable '{var}'"
                 )
 
     def check_variables(self):
@@ -1106,11 +1132,9 @@ def check_variables(self):
         for var in all_vars:
             if var not in vars_in_keys:
                 raise pybamm.ModelError(
-                    f"""
-                    No key set for variable '{var}'. Make sure it is included in either
-                    model.rhs or model.algebraic, in an unmodified form
-                    (e.g. not Broadcasted)
-                    """
+                    f"No key set for variable '{var}'. Make sure it is included in either "
+                    "model.rhs or model.algebraic, in an unmodified form "
+                    "(e.g. not Broadcasted)"
                 )
 
     def check_no_repeated_keys(self):
@@ -1524,9 +1548,7 @@ def check_and_convert_bcs(self, boundary_conditions):
                 # Check types
                 if bc[1] not in ["Dirichlet", "Neumann"]:
                     raise pybamm.ModelError(
-                        f"""
-                        boundary condition types must be Dirichlet or Neumann, not '{bc[1]}'
-                        """
+                        f"boundary condition types must be Dirichlet or Neumann, not '{bc[1]}'"
                     )
 
         return boundary_conditions
diff --git a/src/pybamm/models/submodels/thermal/base_thermal.py b/src/pybamm/models/submodels/thermal/base_thermal.py
@@ -143,8 +143,12 @@ def _get_standard_coupled_variables(self, variables):
             phase_names = ["primary ", "secondary "]
 
         if self.options.electrode_types["negative"] == "planar":
-            Q_rxn_n = pybamm.FullBroadcast(
-                0, ["negative electrode"], "current collector"
+            i_n = variables["Lithium metal total interfacial current density [A.m-2]"]
+            eta_r_n = variables["Lithium metal interface reaction overpotential [V]"]
+            Q_rxn_n = pybamm.PrimaryBroadcast(
+                i_n * eta_r_n / self.param.n.L,
+                ["negative electrode"],
+                "current collector",
             )
             Q_rev_n = pybamm.FullBroadcast(
                 0, ["negative electrode"], "current collector"
diff --git a/src/pybamm/parameters/parameter_values.py b/src/pybamm/parameters/parameter_values.py
diff --git a/src/pybamm/simulation.py b/src/pybamm/simulation.py
diff --git a/src/pybamm/solvers/casadi_algebraic_solver.py b/src/pybamm/solvers/casadi_algebraic_solver.py
diff --git a/tests/unit/test_expression_tree/test_coupled_variable.py b/tests/unit/test_expression_tree/test_coupled_variable.py
diff --git a/tests/unit/test_expression_tree/test_symbol.py b/tests/unit/test_expression_tree/test_symbol.py
