Repository: pybamm-team/PyBaMM

Issue:
`IndefiniteIntegral` node to edge or node to node
For some of the lead-acid models we need an indefinite integral where the integrand evaluates on nodes and the integral evaluates on either nodes or edges

Discussion:

Patch:
diff --git a/docs/tutorials/add-spatial-method.rst b/docs/tutorials/add-spatial-method.rst
index a6f256c6e..2d6ad71c3 100644
--- a/docs/tutorials/add-spatial-method.rst
+++ b/docs/tutorials/add-spatial-method.rst
@@ -63,6 +63,7 @@ Optionally, a new spatial method can also overwrite the default behaviour for th
 - :meth:`pybamm.SpatialMethod.broadcast`
 - :meth:`pybamm.SpatialMethod.mass_matrix`
 - :meth:`pybamm.SpatialMethod.process_binary_operators`
+- :meth:`pybamm.SpatialMethod.concatenation`
 
 For an example of an existing spatial method implementation, see the Finite Volume
 `API docs <https://pybamm.readthedocs.io/en/latest/source/spatial_methods/finite_volume.html>`_
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index ce330a1cf..efb60253c 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -428,7 +428,7 @@ class Discretisation(object):
 
         elif isinstance(symbol, pybamm.Concatenation):
             new_children = [self.process_symbol(child) for child in symbol.children]
-            new_symbol = pybamm.DomainConcatenation(new_children, self.mesh)
+            new_symbol = spatial_method.concatenation(new_children)
 
             return new_symbol
 
diff --git a/pybamm/spatial_methods/finite_volume.py b/pybamm/spatial_methods/finite_volume.py
index 5a4773685..4cfb67e82 100644
--- a/pybamm/spatial_methods/finite_volume.py
+++ b/pybamm/spatial_methods/finite_volume.py
@@ -189,9 +189,7 @@ class FiniteVolume(pybamm.SpatialMethod):
         return pybamm.Matrix(matrix)
 
     def integral(self, domain, symbol, discretised_symbol):
-        """Vector-vector dot product to implement the integral operator.
-        See :meth:`pybamm.BaseDiscretisation.integral`
-        """
+        """Vector-vector dot product to implement the integral operator. """
         # Calculate integration vector
         integration_vector = self.definite_integral_vector(domain)
 
@@ -240,18 +238,14 @@ class FiniteVolume(pybamm.SpatialMethod):
         return pybamm.Matrix(vector[np.newaxis, :])
 
     def indefinite_integral(self, domain, symbol, discretised_symbol):
-        """Implementation of the indefinite integral operator. The
-        input discretised symbol must be defined on the internal mesh edges.
-        See :meth:`pybamm.BaseDiscretisation.indefinite_integral`
-        """
-
-        if not symbol.evaluates_on_edges():
-            raise pybamm.ModelError(
-                "Symbol to be integrated must be valid on the mesh edges"
-            )
+        """Implementation of the indefinite integral operator. """
 
-        # Calculate integration matrix
-        integration_matrix = self.indefinite_integral_matrix(domain)
+        # Different integral matrix depending on whether the integrand evaluates on
+        # edges or nodes
+        if symbol.evaluates_on_edges():
+            integration_matrix = self.indefinite_integral_matrix_edges(domain)
+        else:
+            integration_matrix = self.indefinite_integral_matrix_nodes(domain)
 
         # Don't need to check for spherical domains as spherical polars
         # only change the diveregence (symbols here have grad and no div)
@@ -262,9 +256,10 @@ class FiniteVolume(pybamm.SpatialMethod):
         self.test_shape(out)
         return out
 
-    def indefinite_integral_matrix(self, domain):
+    def indefinite_integral_matrix_edges(self, domain):
         """
-        Matrix for finite-volume implementation of the indefinite integral
+        Matrix for finite-volume implementation of the indefinite integral where the
+        integrand is evaluated on mesh edges
 
         .. math::
             F(x) = \\int_0^x\\!f(u)\\,du
@@ -295,8 +290,8 @@ class FiniteVolume(pybamm.SpatialMethod):
 
         Returns
         -------
-        :class:`pybamm.Vector`
-            The finite volume integral vector for the domain
+        :class:`pybamm.Matrix`
+            The finite volume integral matrix for the domain
         """
 
         # Create appropriate submesh by combining submeshes in domain
@@ -323,6 +318,41 @@ class FiniteVolume(pybamm.SpatialMethod):
 
         return pybamm.Matrix(matrix)
 
+    def indefinite_integral_matrix_nodes(self, domain):
+        """
+        Matrix for finite-volume implementation of the indefinite integral where the
+        integrand is evaluated on mesh nodes.
+        This is just a straightforward cumulative sum of the integrand
+
+        Parameters
+        ----------
+        domain : list
+            The domain(s) of integration
+
+        Returns
+        -------
+        :class:`pybamm.Matrix`
+            The finite volume integral matrix for the domain
+        """
+
+        # Create appropriate submesh by combining submeshes in domain
+        submesh_list = self.mesh.combine_submeshes(*domain)
+        submesh = submesh_list[0]
+        n = submesh.npts
+        sec_pts = len(submesh_list)
+
+        du_n = submesh.d_edges
+        du_entries = [du_n] * (n)
+        offset = -np.arange(1, n + 1, 1)
+        sub_matrix = diags(du_entries, offset, shape=(n + 1, n))
+        # Convert to csr_matrix so that we can take the index (row-slicing), which is
+        # not supported by the default kron format
+        # Note that this makes column-slicing inefficient, but this should not be an
+        # issue
+        matrix = csr_matrix(kron(eye(sec_pts), sub_matrix))
+
+        return pybamm.Matrix(matrix)
+
     def add_ghost_nodes(self, symbol, discretised_symbol, bcs):
         """
         Add ghost nodes to a symbol.
@@ -522,20 +552,64 @@ class FiniteVolume(pybamm.SpatialMethod):
         # no need to do any averaging
         if left_evaluates_on_edges == right_evaluates_on_edges:
             pass
-        # If only left child evaluates on edges, compute diffusivity for right child
+        # If only left child evaluates on edges, map right child onto edges
         elif left_evaluates_on_edges and not right_evaluates_on_edges:
-            disc_right = self.compute_diffusivity(disc_right)
-        # If only right child evaluates on edges, compute diffusivity for left child
+            disc_right = self.node_to_edge(disc_right)
+        # If only right child evaluates on edges, map left child onto edges
         elif right_evaluates_on_edges and not left_evaluates_on_edges:
-            disc_left = self.compute_diffusivity(disc_left)
+            disc_left = self.node_to_edge(disc_left)
         # Return new binary operator with appropriate class
         out = bin_op.__class__(disc_left, disc_right)
         self.test_shape(out)
         return out
 
-    def compute_diffusivity(self, discretised_symbol):
+    def concatenation(self, disc_children):
+        """Discrete concatenation, taking `edge_to_node` for children that evaluate on
+        edges.
+        See :meth:`pybamm.SpatialMethod.concatenation`
+        """
+        for idx, child in enumerate(disc_children):
+            n_nodes = sum(
+                len(mesh.nodes) for mesh in self.mesh.combine_submeshes(*child.domain)
+            )
+            n_edges = sum(
+                len(mesh.edges) for mesh in self.mesh.combine_submeshes(*child.domain)
+            )
+            child_size = child.size
+            if child_size != n_nodes:
+                # Average any children that evaluate on the edges (size n_edges) to
+                # evaluate on nodes instead, so that concatenation works properly
+                if child_size == n_edges:
+                    disc_children[idx] = self.edge_to_node(child)
+                else:
+                    raise pybamm.ShapeError(
+                        """
+                        child must have size n_nodes (number of nodes in the mesh)
+                        or n_edges (number of edges in the mesh)
+                        """
+                    )
+        return pybamm.DomainConcatenation(disc_children, self.mesh)
+
+    def edge_to_node(self, discretised_symbol):
         """
-        Compute the diffusivity at cell edges, based on the diffusivity at cell nodes.
+        Convert a discretised symbol evaluated on the cell edges to a discretised symbol
+        evaluated on the cell nodes.
+        See :meth:`pybamm.FiniteVolume.shift`
+        """
+        return self.shift(discretised_symbol, "edge to node")
+
+    def node_to_edge(self, discretised_symbol):
+        """
+        Convert a discretised symbol evaluated on the cell nodes to a discretised symbol
+        evaluated on the cell edges.
+        See :meth:`pybamm.FiniteVolume.shift`
+        """
+        return self.shift(discretised_symbol, "node to edge")
+
+    def shift(self, discretised_symbol, shift_key):
+        """
+        Convert a discretised symbol evaluated at edges/nodes, to a discretised symbol
+        evaluated at nodes/edges.
         For now we just take the arithemtic mean, though it may be better to take the
         harmonic mean based on [1].
 
@@ -546,18 +620,22 @@ class FiniteVolume(pybamm.SpatialMethod):
         ----------
         discretised_symbol : :class:`pybamm.Symbol`
             Symbol to be averaged. When evaluated, this symbol returns either a scalar
-            or an array of shape (n,), where n is the number of points in the mesh for
-            the symbol's domain (n = self.mesh[symbol.domain].npts)
+            or an array of shape (n,) or (n+1,), where n is the number of points in the
+            mesh for the symbol's domain (n = self.mesh[symbol.domain].npts)
+        shift_key : str
+            Whether to shift from nodes to edges ("node to edge"), or from edges to
+            nodes ("edge to node")
 
         Returns
         -------
-        :class:`pybamm.Function`
+        :class:`pybamm.Symbol`
             Averaged symbol. When evaluated, this returns either a scalar or an array of
-            shape (n+1,) as appropriate.
+            shape (n+1,) (if `shift_key = "node to edge"`) or (n,) (if
+            `shift_key = "edge to node"`)
         """
 
         def arithmetic_mean(array):
-            """Calculate the arithemetic mean of an array using matrix multiplication"""
+            """Calculate the arithmetic mean of an array using matrix multiplication"""
             # Create appropriate submesh by combining submeshes in domain
             submesh_list = self.mesh.combine_submeshes(*array.domain)
 
@@ -567,13 +645,21 @@ class FiniteVolume(pybamm.SpatialMethod):
             # Create 1D matrix using submesh
             n = submesh.npts
 
-            sub_matrix_left = csr_matrix(([1.5, -0.5], ([0, 0], [0, 1])), shape=(1, n))
-            sub_matrix_center = diags([0.5, 0.5], [0, 1], shape=(n - 1, n))
-            sub_matrix_right = csr_matrix(
-                ([-0.5, 1.5], ([0, 0], [n - 2, n - 1])), shape=(1, n)
-            )
-            sub_matrix = vstack([sub_matrix_left, sub_matrix_center, sub_matrix_right])
-
+            if shift_key == "node to edge":
+                sub_matrix_left = csr_matrix(
+                    ([1.5, -0.5], ([0, 0], [0, 1])), shape=(1, n)
+                )
+                sub_matrix_center = diags([0.5, 0.5], [0, 1], shape=(n - 1, n))
+                sub_matrix_right = csr_matrix(
+                    ([-0.5, 1.5], ([0, 0], [n - 2, n - 1])), shape=(1, n)
+                )
+                sub_matrix = vstack(
+                    [sub_matrix_left, sub_matrix_center, sub_matrix_right]
+                )
+            elif shift_key == "edge to node":
+                sub_matrix = diags([0.5, 0.5], [0, 1], shape=(n, n + 1))
+            else:
+                raise ValueError("shift key '{}' not recognised".format(shift_key))
             # Second dimension length
             second_dim_len = len(submesh_list)
 
diff --git a/pybamm/spatial_methods/spatial_method.py b/pybamm/spatial_methods/spatial_method.py
index 3702b6058..98667f564 100644
--- a/pybamm/spatial_methods/spatial_method.py
+++ b/pybamm/spatial_methods/spatial_method.py
@@ -8,8 +8,8 @@ from scipy.sparse import eye, kron, coo_matrix
 
 class SpatialMethod:
     """
-    A general spatial methods class, with default (trivial) behaviour for broadcast,
-    mass_matrix and compute_diffusivity.
+    A general spatial methods class, with default (trivial) behaviour for some spatial
+    operations.
     All spatial methods will follow the general form of SpatialMethod in
     that they contain a method for broadcasting variables onto a mesh,
     a gradient operator, and a diverence operator.
@@ -267,6 +267,21 @@ class SpatialMethod:
         """
         return bin_op.__class__(disc_left, disc_right)
 
+    def concatenation(self, disc_children):
+        """Discrete concatenation object.
+
+        Parameters
+        ----------
+        disc_children : list
+            List of discretised children
+
+        Returns
+        -------
+        :class:`pybamm.DomainConcatenation`
+            Concatenation of the discretised children
+        """
+        return pybamm.DomainConcatenation(disc_children, self.mesh)
+
     @staticmethod
     def test_shape(symbol):
         """


Test patch:
diff --git a/tests/integration/test_models/standard_output_tests.py b/tests/integration/test_models/standard_output_tests.py
index 8c9768207..5bc1b04f8 100644
--- a/tests/integration/test_models/standard_output_tests.py
+++ b/tests/integration/test_models/standard_output_tests.py
@@ -455,6 +455,8 @@ class CurrentTests(BaseOutputTest):
 
         self.i_s_n = self.get_var("Negative electrode current density")
         self.i_s_p = self.get_var("Positive electrode current density")
+        self.i_s = self.get_var("Electrode current density")
+        self.i_e = self.get_var("Electrolyte current density")
 
     def test_interfacial_current_average(self):
         """Test that average of the interfacial current density is equal to the true
@@ -469,8 +471,20 @@ class CurrentTests(BaseOutputTest):
     def test_conservation(self):
         """Test sum of electrode and electrolyte current densities give the applied
         current density"""
+        t, x_n, x_s, x_p = self.t, self.x_n, self.x_s, self.x_p
 
-        # TODO: add a total function
+        current_param = pybamm.electrical_parameters.current_with_time
+        i_cell = self.param.process_symbol(current_param).evaluate(t=t)
+        for x in [x_n, x_s, x_p]:
+            np.testing.assert_array_almost_equal(
+                self.i_s(t, x) + self.i_e(t, x), i_cell, decimal=3
+            )
+        np.testing.assert_array_almost_equal(
+            self.i_s(t, x_n), self.i_s_n(t, x_n), decimal=3
+        )
+        np.testing.assert_array_almost_equal(
+            self.i_s(t, x_p), self.i_s_p(t, x_p), decimal=3
+        )
 
     def test_current_density_boundaries(self):
         """Test the boundary values of the current densities"""
diff --git a/tests/unit/test_spatial_methods/test_finite_volume.py b/tests/unit/test_spatial_methods/test_finite_volume.py
index e8a465380..1ba54c823 100644
--- a/tests/unit/test_spatial_methods/test_finite_volume.py
+++ b/tests/unit/test_spatial_methods/test_finite_volume.py
@@ -14,23 +14,48 @@ import unittest
 
 
 class TestFiniteVolume(unittest.TestCase):
-    def test_compute_diffusivity(self):
+    def test_node_to_edge_to_node(self):
         # create discretisation
         mesh = get_mesh_for_testing()
         fin_vol = pybamm.FiniteVolume(mesh)
         n = mesh["negative electrode"][0].npts
 
+        # node to edge
         c = pybamm.Vector(np.ones(n), domain=["negative electrode"])
-        diffusivity_c = fin_vol.compute_diffusivity(c)
+        diffusivity_c = fin_vol.node_to_edge(c)
         np.testing.assert_array_equal(diffusivity_c.evaluate(), np.ones((n + 1, 1)))
 
-        d = pybamm.StateVector(slice(0, n), domain=["negative electrode"])
-        y_test = np.ones(n)
-        diffusivity_d = fin_vol.compute_diffusivity(d)
+        # edge to node
+        d = pybamm.StateVector(slice(0, n + 1), domain=["negative electrode"])
+        y_test = np.ones(n + 1)
+        diffusivity_d = fin_vol.edge_to_node(d)
         np.testing.assert_array_equal(
-            diffusivity_d.evaluate(None, y_test), np.ones((n + 1, 1))
+            diffusivity_d.evaluate(None, y_test), np.ones((n, 1))
         )
 
+        # bad shift key
+        with self.assertRaisesRegex(ValueError, "shift key"):
+            fin_vol.shift(c, "bad shift key")
+
+    def test_concatenation(self):
+        mesh = get_mesh_for_testing()
+        fin_vol = pybamm.FiniteVolume(mesh)
+        whole_cell = ["negative electrode", "separator", "positive electrode"]
+        edges = [pybamm.Vector(mesh[dom][0].edges, domain=dom) for dom in whole_cell]
+        # Concatenation of edges should get averaged to nodes first, using edge_to_node
+        v_disc = fin_vol.concatenation(edges)
+        np.testing.assert_array_equal(
+            v_disc.evaluate()[:, 0], mesh.combine_submeshes(*whole_cell)[0].nodes
+        )
+
+        # test for bad shape
+        edges = [
+            pybamm.Vector(np.ones(mesh[dom][0].npts + 2), domain=dom)
+            for dom in whole_cell
+        ]
+        with self.assertRaisesRegex(pybamm.ShapeError, "child must have size n_nodes"):
+            fin_vol.concatenation(edges)
+
     def test_extrapolate_left_right(self):
         # create discretisation
         mesh = get_mesh_for_testing()
@@ -1016,29 +1041,39 @@ class TestFiniteVolume(unittest.TestCase):
         np.testing.assert_array_almost_equal(c_exact, c_approx, decimal=3)
         self.assertEqual(left_boundary_value_disc.evaluate(y=c_exact), 0)
 
-        # ------------
-        # check raises error for variales not on mesh edges
-        phi = pybamm.Variable("phi", domain=["separator", "positive electrode"])
-        no_grad_or_div = phi
-        x = pybamm.SpatialVariable("x", ["separator", "positive electrode"])
-        int_grad_phi = pybamm.IndefiniteIntegral(no_grad_or_div, x)
-        disc.set_variable_slices([phi])
+    def test_indefinite_integral_on_nodes(self):
+        mesh = get_mesh_for_testing()
+        spatial_methods = {"macroscale": pybamm.FiniteVolume}
+        disc = pybamm.Discretisation(mesh, spatial_methods)
 
-        with self.assertRaisesRegex(pybamm.ModelError, "integrated"):
-            disc.process_symbol(int_grad_phi)
+        # input a phi, take grad, then integrate to recover phi approximation
+        # (need to test this way as check evaluated on edges using if has grad
+        # and no div)
+        phi = pybamm.Variable("phi", domain=["negative electrode", "separator"])
 
-        grad_and_div = pybamm.div(pybamm.grad(phi))
-        int_grad_phi = pybamm.IndefiniteIntegral(grad_and_div, x)
+        x = pybamm.SpatialVariable("x", ["negative electrode", "separator"])
+        int_phi = pybamm.IndefiniteIntegral(phi, x)
         disc.set_variable_slices([phi])
-        # Add boundary conditions so that the integrand has the right shape
-        disc._bcs = {
-            phi.id: {
-                "left": (pybamm.Scalar(0), "Dirichlet"),
-                "right": (pybamm.Scalar(0), "Neumann"),
-            }
-        }
-        with self.assertRaisesRegex(pybamm.ModelError, "integrated"):
-            disc.process_symbol(int_grad_phi)
+        # Set boundary conditions (required for shape but don't matter)
+        int_phi_disc = disc.process_symbol(int_phi)
+
+        combined_submesh = mesh.combine_submeshes("negative electrode", "separator")
+
+        # constant case
+        phi_exact = np.ones((combined_submesh[0].npts, 1))
+        int_phi_exact = combined_submesh[0].edges
+        int_phi_approx = int_phi_disc.evaluate(None, phi_exact)[:, 0]
+        np.testing.assert_array_equal(int_phi_exact, int_phi_approx)
+        # linear case
+        phi_exact = combined_submesh[0].nodes[:, np.newaxis]
+        int_phi_exact = combined_submesh[0].edges[:, np.newaxis] ** 2 / 2
+        int_phi_approx = int_phi_disc.evaluate(None, phi_exact)
+        np.testing.assert_array_almost_equal(int_phi_exact, int_phi_approx)
+        # cos case
+        phi_exact = np.cos(combined_submesh[0].nodes[:, np.newaxis])
+        int_phi_exact = np.sin(combined_submesh[0].edges[:, np.newaxis])
+        int_phi_approx = int_phi_disc.evaluate(None, phi_exact)
+        np.testing.assert_array_almost_equal(int_phi_exact, int_phi_approx, decimal=5)
 
     def test_discretise_spatial_variable(self):
         # create discretisation