From 48475a4d1c0d5395fe71a771c4202545b1f5610a Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Fri, 20 Dec 2024 14:12:20 +0000 Subject: [PATCH 1/3] feat: auto-add state variables to output variables --- src/pybamm/discretisations/discretisation.py | 80 ++++++++++--------- .../test_discretisation.py | 46 ++++++----- 2 files changed, 72 insertions(+), 54 deletions(-) diff --git a/src/pybamm/discretisations/discretisation.py b/src/pybamm/discretisations/discretisation.py index 3d9579ff9c..6cf454b7ca 100644 --- a/src/pybamm/discretisations/discretisation.py +++ b/src/pybamm/discretisations/discretisation.py @@ -202,7 +202,10 @@ def process_model(self, model, inplace=True): # Note that we **do not** discretise the keys of model.rhs, # model.initial_conditions and model.boundary_conditions pybamm.logger.verbose(f"Discretise variables for {model.name}") - model_disc.variables = self.process_dict(model.variables) + pre_processed_variables = self._pre_process_variables( + model.variables, model_disc.initial_conditions + ) + model_disc.variables = self.process_dict(pre_processed_variables) # Process parabolic and elliptic equations pybamm.logger.verbose(f"Discretise model equations for {model.name}") @@ -657,6 +660,46 @@ def create_mass_matrix(self, model): return mass_matrix, mass_matrix_inv + def _pre_process_variables( + self, + variables: dict[str, pybamm.Symbol], + initial_conditions: dict[pybamm.Variable, pybamm.Symbol], + ): + """ + Pre-process variables before discretisation. This involves: + - ensuring that all the state variables are included in the variables, + any missing are added + + Parameters + ---------- + variables : dict + Dictionary of variables to pre-process + initial_conditions : dict + Dictionary of initial conditions + + Returns + ------- + dict + Pre-processed variables (copy of input variables with any missing state) + + Raises + ------ + :class:`pybamm.ModelError` + If any state variable names are already included but with + incorrect expressions + """ + new_variables = {k: v for k, v in variables.items()} + for var in initial_conditions.keys(): + if var.name not in new_variables: + new_variables[var.name] = var + else: + if new_variables[var.name] != var: + raise pybamm.ModelError( + f"Variable '{var.name}' should have expression " + f"'{var}', but has expression '{new_variables[var.name]}'" + ) + return new_variables + def process_dict(self, var_eqn_dict, ics=False): """Discretise a dictionary of {variable: equation}, broadcasting if necessary (can be model.rhs, model.algebraic, model.initial_conditions or @@ -1008,7 +1051,6 @@ def _concatenate_in_order(self, var_eqn_dict, check_complete=False, sparse=False def check_model(self, model): """Perform some basic checks to make sure the discretised model makes sense.""" self.check_initial_conditions(model) - self.check_variables(model) def check_initial_conditions(self, model): # Check initial conditions are a numpy array @@ -1049,40 +1091,6 @@ def check_initial_conditions(self, model): f"{model.algebraic[var].shape} and initial_conditions.shape = {model.initial_conditions[var].shape} for variable '{var}'." ) - def check_variables(self, model): - """ - Check variables in variable list against rhs. - Be lenient with size check if the variable in model.variables is broadcasted, or - a concatenation - (if broadcasted, variable is a multiplication with a vector of ones) - """ - for rhs_var in model.rhs.keys(): - if rhs_var.name in model.variables.keys(): - var = model.variables[rhs_var.name] - - different_shapes = not np.array_equal( - model.rhs[rhs_var].shape, var.shape - ) - - not_concatenation = not isinstance(var, pybamm.Concatenation) - - not_mult_by_one_vec = not ( - isinstance( - var, (pybamm.Multiplication, pybamm.MatrixMultiplication) - ) - and ( - pybamm.is_matrix_one(var.left) - or pybamm.is_matrix_one(var.right) - ) - ) - - if different_shapes and not_concatenation and not_mult_by_one_vec: - raise pybamm.ModelError( - "variable and its eqn must have the same shape after " - "discretisation but variable.shape = " - f"{var.shape} and rhs.shape = {model.rhs[rhs_var].shape} for variable '{var}'. " - ) - def is_variable_independent(self, var, all_vars_in_eqns): pybamm.logger.verbose("Removing independent blocks.") if not isinstance(var, pybamm.Variable): diff --git a/tests/unit/test_discretisations/test_discretisation.py b/tests/unit/test_discretisations/test_discretisation.py index 78b6cc0d59..d9c0d595fe 100644 --- a/tests/unit/test_discretisations/test_discretisation.py +++ b/tests/unit/test_discretisations/test_discretisation.py @@ -998,9 +998,9 @@ def test_concatenation_2D(self): assert expr.children[2].evaluate(0, y).shape == (105, 1) def test_exceptions(self): - c_n = pybamm.Variable("c", domain=["negative electrode"]) + c_n = pybamm.Variable("c_n", domain=["negative electrode"]) N_n = pybamm.grad(c_n) - c_s = pybamm.Variable("c", domain=["separator"]) + c_s = pybamm.Variable("c_s", domain=["separator"]) N_s = pybamm.grad(c_s) model = pybamm.BaseModel() model.rhs = {c_n: pybamm.div(N_n), c_s: pybamm.div(N_s)} @@ -1011,22 +1011,6 @@ def test_exceptions(self): } disc = get_discretisation_for_testing() - - # check raises error if different sized key and output var - model.variables = {c_n.name: c_s} - with pytest.raises(pybamm.ModelError, match="variable and its eqn"): - disc.process_model(model) - - # check doesn't raise if concatenation - model.variables = {c_n.name: pybamm.concatenation(2 * c_n, 3 * c_s)} - disc.process_model(model, inplace=False) - - # check doesn't raise if broadcast - model.variables = { - c_n.name: pybamm.PrimaryBroadcast( - pybamm.InputParameter("a"), ["negative electrode"] - ) - } disc.process_model(model) # Check setting up a 0D spatial method with 1D mesh raises error @@ -1239,3 +1223,29 @@ def test_independent_rhs_with_event(self): disc = pybamm.Discretisation(remove_independent_variables_from_rhs=True) disc.process_model(model) assert len(model.rhs) == 3 + + def test_pre_process_variables(self): + a = pybamm.Variable("a") + b = pybamm.Variable("b") + model = pybamm.BaseModel() + model.rhs = {a: b, b: a} + model.initial_conditions = { + a: pybamm.Scalar(0), + b: pybamm.Scalar(1), + } + model.variables = { + "a": a, # correct + # b missing + } + disc = pybamm.Discretisation() + disc_model = disc.process_model(model, inplace=False) + assert list(disc_model.variables.keys()) == ["a", "b"] + + model.variables = { + "a": a, + "b": 2 * a, + } + with pytest.raises( + pybamm.ModelError, match="Variable 'b' should have expression" + ): + disc.process_model(model, inplace=False) From 37b3e62322d2b254e71dc6c352fe22f6ffe9a12a Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Fri, 20 Dec 2024 14:14:17 +0000 Subject: [PATCH 2/3] add changelog --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8322884eda..43011dd1e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ - Made composite electrode model compatible with particle size distribution ([#4687](https://github.com/pybamm-team/PyBaMM/pull/4687)) - Added `Symbol.post_order()` method to return an iterable that steps through the tree in post-order fashion. ([#4684](https://github.com/pybamm-team/PyBaMM/pull/4684)) - Added two more submodels (options) for the SEI: Lars von Kolzenberg (2020) model and Tunneling Limit model ([#4394](https://github.com/pybamm-team/PyBaMM/pull/4394)) - +- Automatically add state variables of the model to the output variables if they are not already present ([#4700](https://github.com/pybamm-team/PyBaMM/pull/4700)) ## Breaking changes From f1100b665f780922ed731ae849c5af85f1265d88 Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Fri, 20 Dec 2024 15:54:08 +0000 Subject: [PATCH 3/3] fix lead acid model tests --- src/pybamm/discretisations/discretisation.py | 9 ++-- src/pybamm/expression_tree/concatenations.py | 42 ++++++++++--------- .../oxygen_diffusion/full_oxygen_diffusion.py | 6 ++- 3 files changed, 33 insertions(+), 24 deletions(-) diff --git a/src/pybamm/discretisations/discretisation.py b/src/pybamm/discretisations/discretisation.py index 6cf454b7ca..1a70f1dd56 100644 --- a/src/pybamm/discretisations/discretisation.py +++ b/src/pybamm/discretisations/discretisation.py @@ -193,6 +193,11 @@ def process_model(self, model, inplace=True): model_disc.bcs = self.bcs + # pre-process variables so that all state variables are included + pre_processed_variables = self._pre_process_variables( + model.variables, model.initial_conditions + ) + pybamm.logger.verbose(f"Discretise initial conditions for {model.name}") ics, concat_ics = self.process_initial_conditions(model) model_disc.initial_conditions = ics @@ -202,9 +207,7 @@ def process_model(self, model, inplace=True): # Note that we **do not** discretise the keys of model.rhs, # model.initial_conditions and model.boundary_conditions pybamm.logger.verbose(f"Discretise variables for {model.name}") - pre_processed_variables = self._pre_process_variables( - model.variables, model_disc.initial_conditions - ) + model_disc.variables = self.process_dict(pre_processed_variables) # Process parabolic and elliptic equations diff --git a/src/pybamm/expression_tree/concatenations.py b/src/pybamm/expression_tree/concatenations.py index a1905f9080..62adac5265 100644 --- a/src/pybamm/expression_tree/concatenations.py +++ b/src/pybamm/expression_tree/concatenations.py @@ -4,6 +4,7 @@ from __future__ import annotations import copy from collections import defaultdict +from typing import Optional import numpy as np import sympy @@ -146,9 +147,9 @@ def _concatenation_new_copy(self, children, perform_simplifications: bool = True children before creating the new copy. """ if perform_simplifications: - return concatenation(*children) + return concatenation(*children, name=self.name) else: - return self.__class__(*children) + return self.__class__(*children, name=self.name) def _concatenation_jac(self, children_jacs): """Calculate the Jacobian of a concatenation.""" @@ -468,17 +469,18 @@ def _concatenation_new_copy(self, children, perform_simplifications=True): class ConcatenationVariable(Concatenation): """A Variable representing a concatenation of variables.""" - def __init__(self, *children): - # Name is the intersection of the children names (should usually make sense - # if the children have been named consistently) - name = intersect(children[0].name, children[1].name) - for child in children[2:]: - name = intersect(name, child.name) - if len(name) == 0: - name = None - # name is unchanged if its length is 1 - elif len(name) > 1: - name = name[0].capitalize() + name[1:] + def __init__(self, *children, name: Optional[str] = None): + if name is None: + # Name is the intersection of the children names (should usually make sense + # if the children have been named consistently) + name = intersect(children[0].name, children[1].name) + for child in children[2:]: + name = intersect(name, child.name) + if len(name) == 0: + name = None + # name is unchanged if its length is 1 + elif len(name) > 1: + name = name[0].capitalize() + name[1:] if len(children) > 0: if all(child.scale == children[0].scale for child in children): @@ -523,7 +525,7 @@ def intersect(s1: str, s2: str): return intersect.lstrip().rstrip() -def simplified_concatenation(*children): +def simplified_concatenation(*children, name: Optional[str] = None): """Perform simplifications on a concatenation.""" # remove children that are None children = list(filter(lambda x: x is not None, children)) @@ -534,29 +536,29 @@ def simplified_concatenation(*children): elif len(children) == 1: return children[0] elif all(isinstance(child, pybamm.Variable) for child in children): - return pybamm.ConcatenationVariable(*children) + return pybamm.ConcatenationVariable(*children, name=name) else: # Create Concatenation to easily read domains - concat = Concatenation(*children) + concat = Concatenation(*children, name=name) if all( isinstance(child, pybamm.Broadcast) and child.child == children[0].child for child in children ): unique_child = children[0].orphans[0] if isinstance(children[0], pybamm.PrimaryBroadcast): - return pybamm.PrimaryBroadcast(unique_child, concat.domain) + return pybamm.PrimaryBroadcast(unique_child, concat.domain, name=name) else: return pybamm.FullBroadcast( - unique_child, broadcast_domains=concat.domains + unique_child, broadcast_domains=concat.domains, name=name ) else: return concat -def concatenation(*children): +def concatenation(*children, name: Optional[str] = None): """Helper function to create concatenations.""" # TODO: add option to turn off simplifications - return simplified_concatenation(*children) + return simplified_concatenation(*children, name=name) def simplified_numpy_concatenation(*children): diff --git a/src/pybamm/models/submodels/oxygen_diffusion/full_oxygen_diffusion.py b/src/pybamm/models/submodels/oxygen_diffusion/full_oxygen_diffusion.py index 7ecad6fa4c..4066dcfed7 100644 --- a/src/pybamm/models/submodels/oxygen_diffusion/full_oxygen_diffusion.py +++ b/src/pybamm/models/submodels/oxygen_diffusion/full_oxygen_diffusion.py @@ -36,7 +36,11 @@ def get_fundamental_variables(self): domain="positive electrode", auxiliary_domains={"secondary": "current collector"}, ) - c_ox_s_p = pybamm.concatenation(c_ox_s, c_ox_p) + c_ox_s_p = pybamm.concatenation( + c_ox_s, + c_ox_p, + name="Separator and positive electrode oxygen concentration [mol.m-3]", + ) variables = { "Separator and positive electrode oxygen concentration [mol.m-3]": c_ox_s_p }