Repository: pybamm-team/PyBaMM

Issue:
Add handling gradients of subdomain variables
Currently if we try to take the gradient of `phi_e_n` we get a symbol of the wrong size because `phi_e_n` is not in the boundary conditions (and shouldn't be for checking purposes) instead `phi_e` is in the boundary conditions. 

Discussion:

Patch:
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index 14ebffa96..b6ff10ea6 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -86,6 +86,7 @@ class Discretisation(object):
         self.set_variable_slices(variables)
 
         # set boundary conditions (only need key ids for boundary_conditions)
+        model = self.set_internal_boundary_conditions(model)
         self.bcs = self.process_boundary_conditions(model)
 
         # set up inplace vs not inplace
@@ -97,6 +98,8 @@ class Discretisation(object):
             # create a blank model so that original model is unchanged
             model_disc = pybamm.BaseModel()
 
+        model_disc.bcs = self.bcs
+
         # Process initial condtions
         ics, concat_ics = self.process_initial_conditions(model)
         model_disc.initial_conditions = ics
@@ -167,6 +170,68 @@ class Discretisation(object):
         # reset discretised_symbols
         self._discretised_symbols = {}
 
+    def set_internal_boundary_conditions(self, model):
+        """
+        A method to set the internal boundary conditions for the submodel.
+        These are required to properly calculate the gradient.
+        Note: this method modifies the state of self.boundary_conditions.
+        """
+
+        def boundary_gradient(left_orphan, right_orphan):
+
+            left_domain = left_orphan.domain[0]
+            right_domain = right_orphan.domain[0]
+
+            left_submesh = self._spatial_methods[left_domain].mesh[left_domain][0]
+            right_submesh = self._spatial_methods[right_domain].mesh[right_domain][0]
+
+            x_left = left_submesh.nodes[-1]
+            x_right = right_submesh.nodes[0]
+
+            dy = pybamm.Index(right_orphan, 0) - pybamm.Index(left_orphan, -1)
+            dx = x_right - x_left
+
+            return dy / dx
+
+        bc_key_ids = [key.id for key in list(model.boundary_conditions.keys())]
+
+        internal_bcs = {}
+        for var in model.boundary_conditions.keys():
+            if isinstance(var, pybamm.Concatenation):
+                children = var.children
+
+                first_child = children[0]
+                first_orphan = first_child.new_copy()
+                next_child = children[1]
+                next_orphan = next_child.new_copy()
+
+                lbc = model.boundary_conditions[var]["left"]
+                rbc = (boundary_gradient(first_orphan, next_orphan), "Neumann")
+
+                if first_child.id not in bc_key_ids:
+                    internal_bcs.update({first_child: {"left": lbc, "right": rbc}})
+
+                for i, _ in enumerate(children[1:-1]):
+                    current_child = next_child
+                    current_orphan = next_orphan
+                    next_child = children[i + 2]
+                    next_orphan = next_child.new_copy()
+
+                    lbc = rbc
+                    rbc = (boundary_gradient(current_orphan, next_orphan), "Neumann")
+                    if current_child.id not in bc_key_ids:
+                        internal_bcs.update(
+                            {current_child: {"left": lbc, "right": rbc}}
+                        )
+
+                lbc = rbc
+                rbc = model.boundary_conditions[var]["right"]
+                if children[-1].id not in bc_key_ids:
+                    internal_bcs.update({children[-1]: {"left": lbc, "right": rbc}})
+
+        model.boundary_conditions.update(internal_bcs)
+        return model
+
     def process_initial_conditions(self, model):
         """Discretise model initial_conditions.
 
diff --git a/pybamm/expression_tree/unary_operators.py b/pybamm/expression_tree/unary_operators.py
index 97e383213..ef4536c24 100644
--- a/pybamm/expression_tree/unary_operators.py
+++ b/pybamm/expression_tree/unary_operators.py
@@ -144,7 +144,11 @@ class Index(UnaryOperator):
 
     def __init__(self, child, index, name=None):
         self.index = index
-        if isinstance(index, int):
+        if index == -1:
+            self.slice = slice(index, None)
+            if name is None:
+                name = "Index[-1]"
+        elif isinstance(index, int):
             self.slice = slice(index, index + 1)
             if name is None:
                 name = "Index[" + str(index) + "]"
@@ -158,11 +162,17 @@ class Index(UnaryOperator):
         else:
             raise TypeError("index must be integer or slice")
 
-        if self.slice.stop > child.size:
+        if self.slice in (slice(0, 1), slice(-1, None)):
+            pass
+        elif self.slice.stop > child.size:
             raise ValueError("slice size exceeds child size")
 
         super().__init__(name, child)
 
+        # no domain for integer value
+        if isinstance(index, int):
+            self.domain = []
+
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
 
diff --git a/pybamm/parameters/parameter_values.py b/pybamm/parameters/parameter_values.py
index b02b90924..d3aa235ea 100644
--- a/pybamm/parameters/parameter_values.py
+++ b/pybamm/parameters/parameter_values.py
@@ -273,6 +273,8 @@ class ParameterValues(dict):
             elif isinstance(symbol, pybamm.BoundaryOperator):
                 # BoundaryValue or BoundaryFlux
                 new_symbol = symbol.__class__(new_child, symbol.side)
+            elif isinstance(symbol, pybamm.Index):
+                new_symbol = symbol.__class__(new_child, symbol.index)
             else:
                 new_symbol = symbol.__class__(new_child)
             # ensure domain remains the same


Test patch:
diff --git a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py
index 736b991c1..fcd9218ac 100644
--- a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py
+++ b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py
@@ -50,6 +50,75 @@ class TestDFN(unittest.TestCase):
         modeltest = tests.StandardModelTest(model, var_pts=var_pts)
         modeltest.test_all()
 
+    def test_gradient_of_subdomain_var(self):
+        # a hacky test that will be replaced by a more general one in
+        # standard output tests when #481 is completed
+
+        options = {"thermal": None}
+        model = pybamm.lithium_ion.DFN(options)
+        c_e = model.variables["Electrolyte concentration"]
+        c_e_n, c_e_s, c_e_p = c_e.orphans
+        phi_e = model.variables["Electrolyte potential"]
+        phi_e_n, phi_e_s, phi_e_p = phi_e.orphans
+        model.variables.update(
+            {
+                "Test grad c_e": pybamm.grad(c_e),
+                "Test grad c_e_n": pybamm.grad(c_e_n),
+                "Test grad c_e_s": pybamm.grad(c_e_s),
+                "Test grad c_e_p": pybamm.grad(c_e_p),
+                "Test grad phi_e": pybamm.grad(phi_e),
+                "Test grad phi_e_n": pybamm.grad(phi_e_n),
+                "Test grad phi_e_s": pybamm.grad(phi_e_s),
+                "Test grad phi_e_p": pybamm.grad(phi_e_p),
+            }
+        )
+        geometry = model.default_geometry
+        param = model.default_parameter_values
+        param.process_model(model)
+        param.process_geometry(geometry)
+        mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts)
+        disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
+        disc.process_model(model)
+
+        t_eval = np.linspace(0, 2, 100)
+        solution = model.default_solver.solve(model, t_eval)
+        t = solution.t
+        x_n = disc.mesh["negative electrode"][0].nodes
+        x_s = disc.mesh["separator"][0].nodes
+        x_p = disc.mesh["positive electrode"][0].nodes
+        whole_cell = ["negative electrode", "separator", "positive electrode"]
+        x = disc.mesh.combine_submeshes(*whole_cell)[0].nodes
+
+        variables = pybamm.post_process_variables(
+            model.variables, solution.t, solution.y, mesh=disc.mesh
+        )
+
+        grad_c_e_n = variables["Test grad c_e_n"]
+        grad_c_e_s = variables["Test grad c_e_s"]
+        grad_c_e_p = variables["Test grad c_e_p"]
+        grad_c_e = variables["Test grad c_e"]
+
+        grad_c_e_combined = np.concatenate(
+            (grad_c_e_n(t, x_n), grad_c_e_s(t, x_s), grad_c_e_p(t, x_p)), axis=0
+        )
+
+        np.testing.assert_array_almost_equal(
+            grad_c_e(t, x), grad_c_e_combined, decimal=6
+        )
+
+        grad_phi_e_n = variables["Test grad phi_e_n"]
+        grad_phi_e_s = variables["Test grad phi_e_s"]
+        grad_phi_e_p = variables["Test grad phi_e_p"]
+        grad_phi_e = variables["Test grad phi_e"]
+
+        grad_phi_e_combined = np.concatenate(
+            (grad_phi_e_n(t, x_n), grad_phi_e_s(t, x_s), grad_phi_e_p(t, x_p)), axis=0
+        )
+
+        np.testing.assert_array_almost_equal(
+            grad_phi_e(t, x), grad_phi_e_combined, decimal=6
+        )
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/unit/test_discretisations/test_discretisation.py b/tests/unit/test_discretisations/test_discretisation.py
index 964a34010..7357c6df8 100644
--- a/tests/unit/test_discretisations/test_discretisation.py
+++ b/tests/unit/test_discretisations/test_discretisation.py
@@ -44,6 +44,23 @@ class TestDiscretise(unittest.TestCase):
         disc = pybamm.Discretisation(None, None)
         self.assertEqual(disc._spatial_methods, {})
 
+    def test_add_internal_boundary_conditions(self):
+        model = pybamm.BaseModel()
+        c_e = pybamm.standard_variables.c_e
+        lbc = (pybamm.Scalar(0), "Neumann")
+        rbc = (pybamm.Scalar(0), "Neumann")
+        model.boundary_conditions = {c_e: {"left": lbc, "right": rbc}}
+
+        mesh = get_mesh_for_testing()
+        spatial_methods = {
+            "macroscale": pybamm.SpatialMethod,
+        }
+        disc = pybamm.Discretisation(mesh, spatial_methods)
+        model = disc.set_internal_boundary_conditions(model)
+
+        for child in c_e.children:
+            self.assertTrue(child in model.boundary_conditions.keys())
+
     def test_discretise_slicing(self):
         # create discretisation
         mesh = get_mesh_for_testing()