Repository: tinosulzer/PyBaMM

Issue:
discretisation of different parameter values in different domains
Different parameters can have different values in different domains. This can be represented in a model by a concatentation of Parameters, which will change to a concat of Scalars. The discretisation class need to change this to a concatenation of `Vectors`.

Discussion:

Patch:
diff --git a/pybamm/discretisations/base_discretisation.py b/pybamm/discretisations/base_discretisation.py
index f262cd54..c62e4b39 100644
--- a/pybamm/discretisations/base_discretisation.py
+++ b/pybamm/discretisations/base_discretisation.py
@@ -48,7 +48,8 @@ class BaseDiscretisation(object):
         )
 
         # Discretise right-hand sides, passing domain from variable
-        model.rhs = self.process_rhs(model.rhs, model.boundary_conditions, y_slices)
+        model.rhs = self.process_rhs(
+            model.rhs, model.boundary_conditions, y_slices)
 
     def get_variable_slices(self, variables):
         """Set the slicing for variables.
@@ -90,10 +91,12 @@ class BaseDiscretisation(object):
 
         """
         for variable, equation in initial_conditions.items():
-            discretised_ic = self.process_symbol(equation, variable.domain).evaluate()
+            discretised_ic = self.process_symbol(
+                equation, variable.domain).evaluate()
 
             if isinstance(discretised_ic, numbers.Number):
-                discretised_ic = discretised_ic * self.vector_of_ones(variable.domain)
+                discretised_ic = discretised_ic * \
+                    self.vector_of_ones(variable.domain)
             else:
                 raise NotImplementedError(
                     "Currently only accepts scalar initial conditions"
@@ -164,7 +167,8 @@ class BaseDiscretisation(object):
 
         elif isinstance(symbol, pybamm.BinaryOperator):
             left, right = symbol.children
-            new_left = self.process_symbol(left, domain, y_slices, boundary_conditions)
+            new_left = self.process_symbol(
+                left, domain, y_slices, boundary_conditions)
             new_right = self.process_symbol(
                 right, domain, y_slices, boundary_conditions
             )
@@ -182,6 +186,15 @@ class BaseDiscretisation(object):
             )
             return pybamm.StateVector(y_slices[symbol.id])
 
+        elif isinstance(symbol, pybamm.Concatenation):
+            # only know how to discretise a concatenation of scalars...
+            all_scalars = all(isinstance(child, pybamm.Scalar)
+                              for child in symbol.children)
+            if all_scalars:
+                return self.scalar_to_vector(symbol.children)
+            else:
+                raise NotImplementedError
+
         else:
             # hack to copy the symbol but without a parent
             # (building tree from bottom up)
@@ -227,6 +240,49 @@ class BaseDiscretisation(object):
         """
         raise NotImplementedError
 
+    def scalar_to_vector(self, scalars, domain=None):
+        """
+        Convert a :class:`Scalar` (or a list of :class:`Scalar`) to a uniform
+        Vector of size given by mesh. The size of the Vector is based on the
+        domains associated with the Scalar, and the size of the mesh.
+
+        Parameters
+        ----------
+        scalar: :class:`Scalar` or list of :class:`Scalar`
+            The scalar values to assign to the vector. A list can be given for
+            different values for different domains
+
+        domain : list, optional
+            If given, overrides the domains given in scalar.domain
+
+        """
+        # convert scalar to list if a single one
+        if isinstance(scalars, pybamm.Scalar):
+            scalars = [scalars]
+
+        # work out total size of vector
+        subvector_sizes = []
+        for s in scalars:
+            if domain is None:
+                this_domain = s.domain
+            else:
+                this_domain = domain
+            subvector_size = sum([self.mesh[dom].npts for dom in this_domain])
+            subvector_sizes.append(subvector_size)
+
+        # allocate vector
+        vector = np.empty(sum(subvector_sizes), dtype=float)
+
+        # assign slices of vector associated with concatenated scalars
+        start = 0
+        end = 0
+        for s, size in zip(scalars, subvector_sizes):
+            end += size
+            vector[start:end] = s.value
+            start = end
+
+        return pybamm.Vector(vector)
+
     def vector_of_ones(self, domain):
         """
         Returns a Vector of ones of the size given by the mesh.
@@ -265,7 +321,8 @@ class MatrixVectorDiscretisation(BaseDiscretisation):
         # Check that boundary condition keys are hashes (ids)
         for key in boundary_conditions.keys():
             assert isinstance(key, int), TypeError(
-                "boundary condition keys should be hashes, not {}".format(type(key))
+                "boundary condition keys should be hashes, not {}".format(
+                    type(key))
             )
         # Discretise symbol
         discretised_symbol = self.process_symbol(
@@ -296,7 +353,8 @@ class MatrixVectorDiscretisation(BaseDiscretisation):
         # Check that boundary condition keys are hashes (ids)
         for key in boundary_conditions.keys():
             assert isinstance(key, int), TypeError(
-                "boundary condition keys should be hashes, not {}".format(type(key))
+                "boundary condition keys should be hashes, not {}".format(
+                    type(key))
             )
         # Discretise symbol
         discretised_symbol = self.process_symbol(
diff --git a/pybamm/expression_tree/scalar.py b/pybamm/expression_tree/scalar.py
index 56563ea2..96d78e41 100644
--- a/pybamm/expression_tree/scalar.py
+++ b/pybamm/expression_tree/scalar.py
@@ -6,7 +6,7 @@ from __future__ import print_function, unicode_literals
 import pybamm
 
 
-class Scalar(pybamm.Symbol):
+class Scalar(pybamm.Domain, pybamm.Symbol):
     """A node in the expression tree representing a scalar value
 
     **Extends:** :class:`Symbol`
@@ -19,10 +19,12 @@ class Scalar(pybamm.Symbol):
     name : str, optional
         the name of the node. Defaulted to ``str(value)``
         if not provided
+    domain : iterable of str, optional
+        list of domains the parameter is valid over, defaults to empty list
 
     """
 
-    def __init__(self, value, name=None):
+    def __init__(self, value, name=None, domain=[]):
         """
 
         """
@@ -30,7 +32,7 @@ class Scalar(pybamm.Symbol):
         if name is None:
             name = str(value)
 
-        super().__init__(name)
+        super().__init__(name, domain=domain)
         self.value = value
 
     @property
diff --git a/pybamm/parameters/parameter_values.py b/pybamm/parameters/parameter_values.py
index 0121d935..2773fd32 100644
--- a/pybamm/parameters/parameter_values.py
+++ b/pybamm/parameters/parameter_values.py
@@ -117,7 +117,7 @@ class ParameterValues(dict):
         """
         if isinstance(symbol, pybamm.Parameter):
             value = self.get_parameter_value(symbol)
-            return pybamm.Scalar(value)
+            return pybamm.Scalar(value, domain=symbol.domain)
 
         elif isinstance(symbol, pybamm.BinaryOperator):
             left, right = symbol.children


Test patch:
diff --git a/tests/test_discretisations/test_base_discretisations.py b/tests/test_discretisations/test_base_discretisations.py
index e4378d8e..5ebfd2d5 100644
--- a/tests/test_discretisations/test_base_discretisations.py
+++ b/tests/test_discretisations/test_base_discretisations.py
@@ -11,7 +11,10 @@ class MeshForTesting(pybamm.BaseMesh):
     def __init__(self):
         super().__init__(None)
         self["whole cell"] = self.submeshclass(np.linspace(0, 1, 100))
-        self["negative electrode"] = self.submeshclass(self["whole cell"].nodes[:40])
+        self["negative electrode"] = self.submeshclass(
+            self["whole cell"].nodes[:40])
+        self["positive electrode"] = self.submeshclass(
+            self["whole cell"].nodes[40:])
 
 
 class DiscretisationForTesting(pybamm.MatrixVectorDiscretisation):
@@ -56,7 +59,8 @@ class TestDiscretise(unittest.TestCase):
         y_slices = disc.get_variable_slices(variables)
         self.assertEqual(
             y_slices,
-            {c.id: slice(0, 100), d.id: slice(100, 200), jn.id: slice(200, 240)},
+            {c.id: slice(0, 100), d.id: slice(
+                100, 200), jn.id: slice(200, 240)},
         )
         d_true = 4 * mesh["whole cell"].nodes
         jn_true = mesh["negative electrode"].nodes ** 3
@@ -119,13 +123,17 @@ class TestDiscretise(unittest.TestCase):
         self.assertTrue(isinstance(exp_disc, pybamm.Division))
         # left side
         self.assertTrue(isinstance(exp_disc.children[0], pybamm.Multiplication))
-        self.assertTrue(isinstance(exp_disc.children[0].children[0], pybamm.Scalar))
-        self.assertTrue(isinstance(exp_disc.children[0].children[1], pybamm.Addition))
+        self.assertTrue(isinstance(
+            exp_disc.children[0].children[0], pybamm.Scalar))
+        self.assertTrue(isinstance(
+            exp_disc.children[0].children[1], pybamm.Addition))
         self.assertTrue(
-            isinstance(exp_disc.children[0].children[1].children[0], pybamm.Parameter)
+            isinstance(
+                exp_disc.children[0].children[1].children[0], pybamm.Parameter)
         )
         self.assertTrue(
-            isinstance(exp_disc.children[0].children[1].children[1], pybamm.StateVector)
+            isinstance(
+                exp_disc.children[0].children[1].children[1], pybamm.StateVector)
         )
         self.assertEqual(
             exp_disc.children[0].children[1].children[1].y_slice, y_slices[var2.id]
@@ -136,15 +144,18 @@ class TestDiscretise(unittest.TestCase):
             isinstance(exp_disc.children[1].children[0], pybamm.Subtraction)
         )
         self.assertTrue(
-            isinstance(exp_disc.children[1].children[0].children[0], pybamm.StateVector)
+            isinstance(
+                exp_disc.children[1].children[0].children[0], pybamm.StateVector)
         )
         self.assertEqual(
             exp_disc.children[1].children[0].children[0].y_slice, y_slices[var1.id]
         )
         self.assertTrue(
-            isinstance(exp_disc.children[1].children[0].children[1], pybamm.Parameter)
+            isinstance(
+                exp_disc.children[1].children[0].children[1], pybamm.Parameter)
         )
-        self.assertTrue(isinstance(exp_disc.children[1].children[1], pybamm.Scalar))
+        self.assertTrue(isinstance(
+            exp_disc.children[1].children[1], pybamm.Scalar))
 
     def test_discretise_spatial_operator(self):
         mesh = MeshForTesting()
@@ -156,7 +167,8 @@ class TestDiscretise(unittest.TestCase):
 
             self.assertTrue(isinstance(eqn_disc, pybamm.Multiplication))
             self.assertTrue(isinstance(eqn_disc.children[0], pybamm.Matrix))
-            self.assertTrue(isinstance(eqn_disc.children[1], pybamm.StateVector))
+            self.assertTrue(isinstance(
+                eqn_disc.children[1], pybamm.StateVector))
 
             y = mesh["whole cell"].nodes ** 2
             var_disc = disc.process_symbol(var, None, y_slices)
@@ -184,7 +196,8 @@ class TestDiscretise(unittest.TestCase):
         mesh = MeshForTesting()
         disc = DiscretisationForTesting(mesh)
         y0 = disc.process_initial_conditions(initial_conditions)
-        np.testing.assert_array_equal(y0, 3 * np.ones_like(mesh["whole cell"].nodes))
+        np.testing.assert_array_equal(
+            y0, 3 * np.ones_like(mesh["whole cell"].nodes))
 
         # two equations
         T = pybamm.Variable("T", domain=["negative electrode"])
@@ -227,10 +240,12 @@ class TestDiscretise(unittest.TestCase):
         y_slices = disc.get_variable_slices(rhs.keys())
         concatenated_rhs = disc.process_rhs(rhs, boundary_conditions, y_slices)
         np.testing.assert_array_equal(
-            y[y_slices[c.id]], concatenated_rhs.evaluate(None, y)[y_slices[c.id]]
+            y[y_slices[c.id]], concatenated_rhs.evaluate(None, y)[
+                y_slices[c.id]]
         )
         np.testing.assert_array_equal(
-            y[y_slices[T.id]], concatenated_rhs.evaluate(None, y)[y_slices[T.id]]
+            y[y_slices[T.id]], concatenated_rhs.evaluate(None, y)[
+                y_slices[T.id]]
         )
 
     def test_process_model(self):
@@ -246,7 +261,8 @@ class TestDiscretise(unittest.TestCase):
 
         disc.process_model(model)
         y0 = model.initial_conditions
-        np.testing.assert_array_equal(y0, 3 * np.ones_like(mesh["whole cell"].nodes))
+        np.testing.assert_array_equal(
+            y0, 3 * np.ones_like(mesh["whole cell"].nodes))
         np.testing.assert_array_equal(y0, model.rhs.evaluate(None, y0))
 
         # two equations
@@ -270,11 +286,43 @@ class TestDiscretise(unittest.TestCase):
         )
         np.testing.assert_array_equal(y0, model.rhs.evaluate(None, y0))
 
+    def test_scalar_to_vector(self):
+        a = pybamm.Scalar(5)
+        mesh = MeshForTesting()
+        disc = pybamm.BaseDiscretisation(mesh)
+        a_vec = disc.scalar_to_vector(a, ["whole cell"])
+        expected_vector = 5 * np.ones_like(mesh["whole cell"].nodes)
+        np.testing.assert_allclose(a_vec.evaluate(), expected_vector)
+
+        a = pybamm.Scalar(5, domain=["whole cell"])
+        a_vec = disc.scalar_to_vector(a)
+        np.testing.assert_allclose(a_vec.evaluate(), expected_vector)
+
+        a_vec = disc.scalar_to_vector(a, ["whole cell", "negative electrode"])
+        expected_vector = np.concatenate([
+            5 * np.ones_like(mesh["whole cell"].nodes),
+            5 * np.ones_like(mesh["negative electrode"].nodes),
+        ])
+        np.testing.assert_allclose(a_vec.evaluate(), expected_vector)
+
+        a = pybamm.Scalar(5, domain=["whole cell", "negative electrode"])
+        a_vec = disc.scalar_to_vector(a)
+        np.testing.assert_allclose(a_vec.evaluate(), expected_vector)
+
+        a = pybamm.Scalar(5, domain=["whole cell"])
+        b = pybamm.Scalar(4, domain=["negative electrode"])
+        a_vec = disc.scalar_to_vector([a, b])
+        expected_vector = np.concatenate([
+            5 * np.ones_like(mesh["whole cell"].nodes),
+            4 * np.ones_like(mesh["negative electrode"].nodes),
+        ])
+        np.testing.assert_allclose(a_vec.evaluate(), expected_vector)
+
     def test_vector_of_ones(self):
         mesh = MeshForTesting()
         disc = pybamm.BaseDiscretisation(mesh)
         vec = disc.vector_of_ones(["whole cell"])
-        self.assertEqual(vec.evaluate(None)[0], 1)
+        self.assertEqual(vec.evaluate()[0], 1)
         self.assertEqual(vec.shape, mesh["whole cell"].nodes.shape)
 
     def test_concatenation(self):
@@ -285,6 +333,28 @@ class TestDiscretise(unittest.TestCase):
         conc = disc.concatenate(a, b, c)
         self.assertTrue(isinstance(conc, pybamm.Concatenation))
 
+    def test_concatenation_of_scalars(self):
+        a = pybamm.Scalar(5, domain=["negative electrode"])
+        b = pybamm.Scalar(4, domain=["positive electrode"])
+        mesh = MeshForTesting()
+        disc = pybamm.BaseDiscretisation(mesh)
+        var = pybamm.Variable("var", domain=["whole cell"])
+        y_slices = disc.get_variable_slices([var])
+
+        eqn = pybamm.Concatenation(a, b)
+        eqn_disc = disc.process_symbol(eqn, var.domain, y_slices, {})
+        self.assertTrue(isinstance(eqn_disc, pybamm.Vector))
+        expected_vector = np.concatenate([
+            5 * np.ones_like(mesh["negative electrode"].nodes),
+            4 * np.ones_like(mesh["positive electrode"].nodes),
+        ])
+        np.testing.assert_allclose(eqn_disc.evaluate(), expected_vector)
+
+        # should only be able to concatentate scalars
+        eqn = pybamm.Concatenation(a, var)
+        with self.assertRaises(NotImplementedError):
+            eqn_disc = disc.process_symbol(eqn, var.domain, y_slices, {})
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")