Repository: tinosulzer/PyBaMM

Issue:
Allow Concatenation of Variables as a key in the model
This will make writing the model down easier, but requires some adaptations:
- Set y_slices for the children of the concatenation, not the concatenation itself
- Take gradient and divergence appropriately
- Reading initial and boundary conditions *should* work ok as-is

### Example

```python3
cn = Variable("conc", ["negative electrode"])
cs = Variable("conc", ["separator"])
cp = Variable("conc", ["positive electrode"])
c = Concatenation(cn, cs, cp)
jn = butler_volmer(cn)
jp = butler_volmer(cp)
j = Concatenation(jn, 0, jp)
model.rhs = {c: div(grad(c)) + sj} # avoids having to write this equation down 3 times
```

Discussion:

Patch:
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index de5a745e..bb0b8388 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -79,10 +79,19 @@ class Discretisation(object):
         variables : iterable of :class:`pybamm.Variables`
             The variables for which to set slices
         """
-        y_slices = {variable.id: None for variable in variables}
+        # Unpack symbols in variables that are concatenations of variables
+        unpacked_variables = []
+        for symbol in variables:
+            if isinstance(symbol, pybamm.Concatenation):
+                unpacked_variables.extend([var for var in symbol.children])
+            else:
+                unpacked_variables.append(symbol)
+        # Set up y_slices
+        y_slices = {variable.id: None for variable in unpacked_variables}
         start = 0
         end = 0
-        for variable in variables:
+        # Iterate through unpacked variables, adding appropriate slices to y_slices
+        for variable in unpacked_variables:
             # If domain is empty then variable has size 1
             if variable.domain == []:
                 end += 1
@@ -333,7 +342,15 @@ class Discretisation(object):
             Discretised right-hand side equations
 
         """
-        ids = {v.id for v in var_eqn_dict.keys()}
+        # Unpack symbols in variables that are concatenations of variables
+        unpacked_variables = []
+        for symbol in var_eqn_dict.keys():
+            if isinstance(symbol, pybamm.Concatenation):
+                unpacked_variables.extend([var for var in symbol.children])
+            else:
+                unpacked_variables.append(symbol)
+        # Check keys from the given var_eqn_dict against self._y_slices
+        ids = {v.id for v in unpacked_variables}
         if ids != self._y_slices.keys():
             given_variable_names = [v.name for v in var_eqn_dict.keys()]
             raise pybamm.ModelError(
@@ -342,7 +359,7 @@ class Discretisation(object):
             )
 
         equations = list(var_eqn_dict.values())
-        slices = [self._y_slices[var.id] for var in var_eqn_dict.keys()]
+        slices = [self._y_slices[var.id] for var in unpacked_variables]
 
         # sort equations according to slices
         sorted_equations = [eq for _, eq in sorted(zip(slices, equations))]
diff --git a/pybamm/models/base_model.py b/pybamm/models/base_model.py
index 7c484f00..60d0df17 100644
--- a/pybamm/models/base_model.py
+++ b/pybamm/models/base_model.py
@@ -232,13 +232,20 @@ class BaseModel(object):
         vars_in_algebraic_keys = set()
         vars_in_eqns = set()
         # Get all variables ids from rhs and algebraic keys and equations
+        # For equations we look through the whole expression tree.
+        # "Variables" can be Concatenations so we also have to look in the whole
+        # expression tree
         for var, eqn in self.rhs.items():
-            vars_in_rhs_keys.add(var.id)
+            vars_in_rhs_keys.update(
+                [x.id for x in var.pre_order() if isinstance(x, pybamm.Variable)]
+            )
             vars_in_eqns.update(
                 [x.id for x in eqn.pre_order() if isinstance(x, pybamm.Variable)]
             )
         for var, eqn in self.algebraic.items():
-            vars_in_algebraic_keys.add(var.id)
+            vars_in_algebraic_keys.update(
+                [x.id for x in var.pre_order() if isinstance(x, pybamm.Variable)]
+            )
             vars_in_eqns.update(
                 [x.id for x in eqn.pre_order() if isinstance(x, pybamm.Variable)]
             )


Test patch:
diff --git a/tests/test_discretisations/test_discretisation.py b/tests/test_discretisations/test_discretisation.py
index 7327f60b..01995f97 100644
--- a/tests/test_discretisations/test_discretisation.py
+++ b/tests/test_discretisations/test_discretisation.py
@@ -71,6 +71,29 @@ class TestDiscretise(unittest.TestCase):
         np.testing.assert_array_equal(y[disc._y_slices[d.id]], d_true)
         np.testing.assert_array_equal(y[disc._y_slices[jn.id]], jn_true)
 
+        # Variables with a concatenation
+        js = pybamm.Variable("js", domain=["separator"])
+        jp = pybamm.Variable("jp", domain=["positive electrode"])
+        j = pybamm.Concatenation(jn, js, jp)
+        variables = [c, d, j]
+        disc.set_variable_slices(variables)
+        self.assertEqual(
+            disc._y_slices,
+            {
+                c.id: slice(0, 100),
+                d.id: slice(100, 200),
+                jn.id: slice(200, 240),
+                js.id: slice(240, 265),
+                jp.id: slice(265, 300),
+            },
+        )
+        d_true = 4 * combined_submesh.nodes
+        jn_true = mesh["negative electrode"].nodes ** 3
+        y = np.concatenate([c_true, d_true, jn_true])
+        np.testing.assert_array_equal(y[disc._y_slices[c.id]], c_true)
+        np.testing.assert_array_equal(y[disc._y_slices[d.id]], d_true)
+        np.testing.assert_array_equal(y[disc._y_slices[jn.id]], jn_true)
+
     def test_process_symbol_base(self):
         # create discretisation
         mesh = get_mesh_for_testing()
@@ -437,6 +460,36 @@ class TestDiscretise(unittest.TestCase):
         with self.assertRaises(pybamm.ModelError):
             disc.process_model(model)
 
+    def test_process_model_concatenation(self):
+        # concatenation of variables as the key
+        cn = pybamm.Variable("c", domain=["negative electrode"])
+        cs = pybamm.Variable("c", domain=["separator"])
+        cp = pybamm.Variable("c", domain=["positive electrode"])
+        c = pybamm.Concatenation(cn, cs, cp)
+        N = pybamm.grad(c)
+        model = pybamm.BaseModel()
+        model.rhs = {c: pybamm.div(N)}
+        model.initial_conditions = {c: pybamm.Scalar(3)}
+
+        model.boundary_conditions = {N: {"left": 0, "right": 0}}
+        model.check_well_posedness()
+
+        # create discretisation
+        disc = get_discretisation_for_testing()
+        mesh = disc.mesh
+
+        combined_submesh = mesh.combine_submeshes(
+            "negative electrode", "separator", "positive electrode"
+        )
+
+        disc.process_model(model)
+        y0 = model.concatenated_initial_conditions
+        np.testing.assert_array_equal(y0, 3 * np.ones_like(combined_submesh.nodes))
+
+        # grad and div are identity operators here
+        np.testing.assert_array_equal(y0, model.concatenated_rhs.evaluate(None, y0))
+        model.check_well_posedness()
+
     def test_broadcast(self):
         whole_cell = ["negative electrode", "separator", "positive electrode"]