Repository: tinosulzer/PyBaMM

Issue:
for algebraic equations: change from list to dict
dict key: `pybamm.Variable`
dict value: algebraic equation associated with that variable

Discussion:

Patch:
diff --git a/pybamm/discretisations/base_discretisation.py b/pybamm/discretisations/base_discretisation.py
index 55f6263a..fb66ed60 100644
--- a/pybamm/discretisations/base_discretisation.py
+++ b/pybamm/discretisations/base_discretisation.py
@@ -56,7 +56,7 @@ class BaseDiscretisation(object):
         self._mesh = self.mesh_type(geometry, self.submesh_types, self.submesh_pts)
 
         # Set the y split for variables
-        variables = self.get_all_variables(model)
+        variables = list(model.rhs.keys()) + list(model.algebraic.keys())
         y_slices = self.get_variable_slices(variables)
 
         # Process initial condtions
@@ -75,46 +75,6 @@ class BaseDiscretisation(object):
         # Check that resulting model makes sense
         self.check_model(model)
 
-    def get_all_variables(self, model):
-        """get all variables in a model
-
-        This returns all the variables in the model in a list. Given a model with n
-        differential equations in model.rhs, the returned list will have the
-        corresponding variables for these equations (in order) in the first n elements.
-        The remaining elements of the list will have the remaining algebraic variables
-        in no specific order
-
-        Parameters
-        ----------
-        model : :class:`pybamm.BaseModel` (or subclass)
-            Model to dicretise.
-
-
-        Returns
-        -------
-        variables: list of :class:`pybamm.Variable`
-            All the variables in the model
-        """
-        variables = list(model.rhs.keys())
-        algebraic_variables = {}
-        for equation in model.algebraic:
-            # find all variables in the expression
-            eq_variables = {
-                variable.id: pybamm.Variable(variable.name, variable.domain)
-                for variable in equation.pre_order()
-                if isinstance(variable, pybamm.Variable)
-            }
-
-            # add to the total set
-            algebraic_variables.update(eq_variables)
-
-        # remove rhs variables from algebraic variables (if they exist)
-        for variable in variables:
-            algebraic_variables.pop(variable.id, None)
-
-        # return both rhs and algebraic variables
-        return variables + list(algebraic_variables.values())
-
     def get_variable_slices(self, variables):
         """Set the slicing for variables.
 
@@ -128,7 +88,6 @@ class BaseDiscretisation(object):
         y_slices : dict of {variable id: slice}
             The slices to take when solving (assigning chunks of y to each vector)
         """
-
         y_slices = {variable.id: None for variable in variables}
         start = 0
         end = 0
@@ -196,10 +155,10 @@ class BaseDiscretisation(object):
         model.concatenated_rhs = self.concatenate(*model.rhs.values())
 
         # Discretise and concatenate algebraic equations
-        model.algebraic = self.process_list(
+        model.algebraic = self.process_dict(
             model.algebraic, y_slices, model.boundary_conditions
         )
-        model.concatenated_algebraic = self.concatenate(*model.algebraic)
+        model.concatenated_algebraic = self.concatenate(*model.algebraic.values())
 
     def process_dict(self, var_eqn_dict, y_slices=None, boundary_conditions={}):
         """Discretise a dictionary of {variable: equation}, broadcasting if necessary
@@ -238,36 +197,6 @@ class BaseDiscretisation(object):
 
         return var_eqn_dict
 
-    def process_list(self, var_eqn_list, y_slices, boundary_conditions):
-        """Discretise a list of [equation]
-        (typically model.algebraic).
-
-        Parameters
-        ----------
-        var_eqn_list:list
-            Equations to dicretise
-        y_slices : dict of {variable id: slice}
-            The slices to assign to StateVectors when discretising
-        boundary_conditions : dict
-            Boundary conditions ({symbol.id: {"left": left bc, "right": right bc}} dict)
-            associated with var_eqn_dict to dicretise
-
-        Returns
-        -------
-        var_eqn_dict : dict
-            Discretised right-hand side equations
-
-        """
-        boundary_conditions = {
-            key.id: value for key, value in boundary_conditions.items()
-        }
-        for i, equation in enumerate(var_eqn_list):
-            var_eqn_list[i] = self.process_symbol(
-                equation, y_slices, boundary_conditions
-            )
-
-        return var_eqn_list
-
     def process_symbol(self, symbol, y_slices=None, boundary_conditions={}):
         """Discretise operators in model equations.
 
diff --git a/pybamm/models/base_model.py b/pybamm/models/base_model.py
index b7488a6d..d79435a2 100644
--- a/pybamm/models/base_model.py
+++ b/pybamm/models/base_model.py
@@ -18,7 +18,10 @@ class BaseModel(object):
         A dictionary that maps expressions (variables) to expressions that represent
         the rhs
     algebraic: dict
-        A list of algebraic expressions that are assumed to equate to zero
+        A dictionary that maps expressions (variables) to expressions that represent
+        the algebraic equations. The algebraic expressions are assumed to equate
+        to zero. Note that all the variables in the model must exist in the keys of
+        `rhs` or `algebraic`.
     initial_conditions: dict
         A dictionary that maps expressions (variables) to expressions that represent
         the initial conditions for the state variables y
@@ -37,7 +40,7 @@ class BaseModel(object):
     def __init__(self):
         # Initialise empty model
         self._rhs = {}
-        self._algebraic = []
+        self._algebraic = {}
         self._initial_conditions = {}
         self._initial_conditions_ydot = {}
         self._boundary_conditions = {}
@@ -88,7 +91,7 @@ class BaseModel(object):
             ]
         ):
             raise pybamm.DomainError(
-                """variable and equation in {} must have the same domain""".format(name)
+                "variable and equation in '{}' must have the same domain".format(name)
             )
 
         return dict
@@ -107,7 +110,7 @@ class BaseModel(object):
 
     @algebraic.setter
     def algebraic(self, algebraic):
-        self._algebraic = algebraic
+        self._algebraic = self._set_dict(algebraic, "algebraic")
 
     @property
     def initial_conditions(self):
@@ -195,33 +198,46 @@ class BaseModel(object):
     def check_well_posedness(self):
         """
         Check that the model is well-posed by executing the following tests:
-        - All the variables that appear in the rhs and algebraic equations appear in
-        the rhs keys, with leeway for exactly n variables where n is the number of
-        algebraic equations. Overdetermined if more equations than variables,
-        underdetermined if more variables than equations.
+        - Model is not over- or underdetermined, by comparing keys and equations in rhs
+        and algebraic. Overdetermined if more equations than variables, underdetermined
+        if more variables than equations.
         - There is an initial condition in self.initial_conditions for each
         variable/equation pair in self.rhs
         - There are appropriate boundary conditions in self.boundary_conditions for each
-        variable/equation pair in self.rhs
+        variable/equation pair in self.rhs and self.algebraic
         """
         # Equations (differential and algebraic)
         # Get all the variables from differential and algebraic equations
-        variable_ids_in_eqns = set()
-        for eqn in list(self.rhs.values()) + self.algebraic:
-            variable_ids_in_eqns.update(
+        vars_in_rhs_keys = set()
+        vars_in_algebraic_keys = set()
+        vars_in_eqns = set()
+        # Get all variables ids from rhs and algebraic keys and equations
+        for var, eqn in self.rhs.items():
+            vars_in_rhs_keys.add(var.id)
+            vars_in_eqns.update(
                 [x.id for x in eqn.pre_order() if isinstance(x, pybamm.Variable)]
             )
-        # Get all variables ids from rhs keys
-        variable_ids_in_keys = set([var.id for var in self.rhs.keys()])
-        # Compare eqns and keys
-        unaccounted_variables = variable_ids_in_eqns.difference(variable_ids_in_keys)
-        # Count how many variables are not accounted for by algebraic equations
-        n_extra_variables = len(unaccounted_variables) - len(self.algebraic)
-        # Make sure n_extra_variables == 0
-        if n_extra_variables > 0:
-            raise pybamm.ModelError("model is underdetermined")
-        elif n_extra_variables < 0:
-            raise pybamm.ModelError("model is overdetermined")
+        for var, eqn in self.algebraic.items():
+            vars_in_algebraic_keys.add(var.id)
+            vars_in_eqns.update(
+                [x.id for x in eqn.pre_order() if isinstance(x, pybamm.Variable)]
+            )
+        # If any keys are repeated between rhs and algebraic then the model is
+        # overdetermined
+        if not set(vars_in_rhs_keys).isdisjoint(vars_in_algebraic_keys):
+            raise pybamm.ModelError("model is overdetermined (repeated keys)")
+        # If any algebraic keys don't appear in the eqns then the model is
+        # overdetermined (but rhs keys can be absent from the eqns, e.g. dcdt = -1 is
+        # fine)
+        extra_algebraic_keys = vars_in_algebraic_keys.difference(vars_in_eqns)
+        if extra_algebraic_keys:
+            raise pybamm.ModelError("model is overdetermined (extra algebraic keys)")
+        # If any variables in the equations don't appear in the keys then the model is
+        # underdetermined
+        vars_in_keys = vars_in_rhs_keys.union(vars_in_algebraic_keys)
+        extra_variables = vars_in_eqns.difference(vars_in_keys)
+        if extra_variables:
+            raise pybamm.ModelError("model is underdetermined (too many variables)")
 
         # Initial conditions
         for var in self.rhs.keys():
@@ -231,7 +247,7 @@ class BaseModel(object):
                 )
 
         # Boundary conditions
-        for var, eqn in self.rhs.items():
+        for var, eqn in {**self.rhs, **self.algebraic}.items():
             if eqn.has_spatial_derivatives():
                 # Variable must be in at least one expression in the boundary condition
                 # keys (to account for both Dirichlet and Neumann boundary conditions)


Test patch:
diff --git a/tests/test_discretisations/test_base_discretisations.py b/tests/test_discretisations/test_base_discretisations.py
index d5b61fd3..d24f0c94 100644
--- a/tests/test_discretisations/test_base_discretisations.py
+++ b/tests/test_discretisations/test_base_discretisations.py
@@ -38,42 +38,6 @@ class TestDiscretise(unittest.TestCase):
         with self.assertRaises(pybamm.ModelError):
             result = disc._concatenate_init(initial_conditions, y_slices)
 
-    def test_find_all_variables(self):
-        # create discretisation
-        defaults = shared.TestDefaults1DMacro()
-        disc = shared.DiscretisationForTesting(
-            defaults.mesh_type, defaults.submesh_pts, defaults.submesh_types
-        )
-        disc.mesh_geometry(defaults.geometry)
-
-        whole_cell = ["negative electrode", "separator", "positive electrode"]
-
-        a = pybamm.Variable("a", domain=whole_cell)
-        b = pybamm.Variable("b", domain=whole_cell)
-        c = pybamm.Variable("c", domain=whole_cell)
-        model = pybamm.BaseModel()
-        model.rhs = {c: pybamm.Scalar(1)}
-
-        variables = disc.get_all_variables(model)
-        self.assertListEqual([v.id for v in variables], [v.id for v in [c]])
-
-        model.rhs = {c: pybamm.Scalar(1), b: pybamm.Scalar(1)}
-
-        variables = disc.get_all_variables(model)
-        self.assertListEqual([v.id for v in variables], [v.id for v in [c, b]])
-
-        model.rhs = {c: pybamm.Scalar(1), b: pybamm.Scalar(1)}
-        model.algebraic = [a - c]
-
-        variables = disc.get_all_variables(model)
-        self.assertListEqual([v.id for v in variables], [v.id for v in [c, b, a]])
-
-        model.rhs = {}
-        model.algebraic = [a - c]
-
-        variables = disc.get_all_variables(model)
-        self.assertSetEqual({v.id for v in variables}, {v.id for v in [a, c]})
-
     def test_discretise_slicing(self):
         # create discretisation
         defaults = shared.TestDefaults1DMacro()
@@ -424,7 +388,7 @@ class TestDiscretise(unittest.TestCase):
         N = pybamm.grad(c)
         model = pybamm.BaseModel()
         model.rhs = {c: pybamm.div(N)}
-        model.algebraic = [d - 2 * c]
+        model.algebraic = {d: d - 2 * c}
         model.initial_conditions = {d: pybamm.Scalar(6), c: pybamm.Scalar(3)}
         model.initial_conditions_ydot = {d: pybamm.Scalar(2), c: pybamm.Scalar(1)}
 
@@ -477,7 +441,7 @@ class TestDiscretise(unittest.TestCase):
         # test that not enough initial conditions for ydot raises an error
         model = pybamm.BaseModel()
         model.rhs = {c: pybamm.div(N)}
-        model.algebraic = [d - 2 * c]
+        model.algebraic = {d: d - 2 * c}
         model.initial_conditions = {d: pybamm.Scalar(6), c: pybamm.Scalar(3)}
         model.initial_conditions_ydot = {c: pybamm.Scalar(1)}
 
diff --git a/tests/test_models/test_base_model.py b/tests/test_models/test_base_model.py
index c175476f..439784bc 100644
--- a/tests/test_models/test_base_model.py
+++ b/tests/test_models/test_base_model.py
@@ -38,7 +38,7 @@ class TestBaseModel(unittest.TestCase):
 
     def test_algebraic_set_get(self):
         model = pybamm.BaseModel()
-        algebraic = [pybamm.Symbol("c") - pybamm.Symbol("a")]
+        algebraic = {pybamm.Symbol("b"): pybamm.Symbol("c") - pybamm.Symbol("a")}
         model.algebraic = algebraic
         self.assertEqual(algebraic, model.algebraic)
 
@@ -173,25 +173,31 @@ class TestBaseModel(unittest.TestCase):
 
         # Well-posed DAE model
         e = pybamm.Variable("e", domain=whole_cell)
-        model.algebraic = [e - c - d]
+        model.algebraic = {e: e - c - d}
         model.check_well_posedness()
 
         # Underdetermined model - not enough differential equations
         model.rhs = {c: 5 * pybamm.div(pybamm.grad(d)) - 1}
-        model.algebraic = [e - c - d]
+        model.algebraic = {e: e - c - d}
         with self.assertRaisesRegex(pybamm.ModelError, "underdetermined"):
             model.check_well_posedness()
 
         # Underdetermined model - not enough algebraic equations
-        model.rhs = {c: 5 * pybamm.div(pybamm.grad(d)) - 1, d: -c - e}
-        model.algebraic = []
+        model.algebraic = {}
         with self.assertRaisesRegex(pybamm.ModelError, "underdetermined"):
             model.check_well_posedness()
 
-        # Overdetermined model - too many algebraic equations
-        # Model cannot be overdetermined by too many differential equations, as rhs
-        # keys are all unique variables
-        model.algebraic = [c - d, e + d]
+        # Overdetermined model - repeated keys
+        model.algebraic = {c: c - d, d: e + d}
+        with self.assertRaisesRegex(pybamm.ModelError, "overdetermined"):
+            model.check_well_posedness()
+        # Overdetermined model - extra keys in algebraic
+        model.rhs = {c: 5 * pybamm.div(pybamm.grad(d)) - 1, d: -d}
+        model.algebraic = {e: c - d}
+        with self.assertRaisesRegex(pybamm.ModelError, "overdetermined"):
+            model.check_well_posedness()
+        model.rhs = {c: 1, d: -1}
+        model.algebraic = {e: c - d}
         with self.assertRaisesRegex(pybamm.ModelError, "overdetermined"):
             model.check_well_posedness()
 
@@ -229,6 +235,20 @@ class TestBaseModel(unittest.TestCase):
         with self.assertRaisesRegex(pybamm.ModelError, "boundary condition"):
             model.check_well_posedness()
 
+        # Algebraic well-posed model
+        whole_cell = ["negative electrode", "separator", "positive electrode"]
+        model = pybamm.BaseModel()
+        model.algebraic = {c: 5 * pybamm.div(pybamm.grad(c)) - 1}
+        model.boundary_conditions = {2 * c: {"left": 0, "right": 0}}
+        model.check_well_posedness()
+        model.boundary_conditions = {pybamm.grad(c): {"left": 0, "right": 0}}
+        model.check_well_posedness()
+
+        # Algebraic model with bad boundary conditions
+        model.boundary_conditions = {d: {"left": 0, "right": 0}}
+        with self.assertRaisesRegex(pybamm.ModelError, "boundary condition"):
+            model.check_well_posedness()
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/test_solvers/test_scikits_solvers.py b/tests/test_solvers/test_scikits_solvers.py
index e529ba9f..59ab7d47 100644
--- a/tests/test_solvers/test_scikits_solvers.py
+++ b/tests/test_solvers/test_scikits_solvers.py
@@ -90,7 +90,7 @@ class TestScikitsSolver(unittest.TestCase):
         var1 = pybamm.Variable("var1", domain=whole_cell)
         var2 = pybamm.Variable("var2", domain=whole_cell)
         model.rhs = {var1: 0.1 * var1}
-        model.algebraic = [2 * var1 - var2]
+        model.algebraic = {var2: 2 * var1 - var2}
         model.initial_conditions = {var1: 1, var2: 2}
         model.initial_conditions_ydot = {var1: 0.1, var2: 0.2}
         disc = model.default_discretisation