Repository: tinosulzer/PyBaMM

Issue:
allow concatenation of variables and scalars
From @tinosulzer  in #77:

It would useful to allow concatenation of variables and scalars, to deal with a case like the following:
Interfacial current G is some function in the positive and negative electrodes, and zero in the separator. In the model, this would probably be coded as

Gn = cn ** 2
Gs = Scalar(0)
Gp = cp ** 3
G = Concatenation(Gn, Gs, Gp)
Then G should get discretised to a concatenation of an expression tree with a StateVector, representing cn, for the negative electrode, a Vector of zeros with the right shape, and an expression tree StateVector, representing cp, for the positive electrode.

Discussion:

Patch:
diff --git a/docs/source/expression_tree/concatenations.rst b/docs/source/expression_tree/concatenations.rst
new file mode 100644
index 00000000..3083003a
--- /dev/null
+++ b/docs/source/expression_tree/concatenations.rst
@@ -0,0 +1,13 @@
+Concatenations
+==============
+
+.. autoclass:: pybamm.Concatenation
+  :members:
+
+.. autoclass:: pybamm.NumpyConcatenation
+  :members:
+
+.. autoclass:: pybamm.DomainConcatenation
+  :members:
+
+.
diff --git a/docs/source/expression_tree/index.rst b/docs/source/expression_tree/index.rst
index fb267a58..8b63c475 100644
--- a/docs/source/expression_tree/index.rst
+++ b/docs/source/expression_tree/index.rst
@@ -12,3 +12,4 @@ Expression Tree
   scalar
   matrix
   vector
+  concatenations
diff --git a/pybamm/__init__.py b/pybamm/__init__.py
index 59856547..c1eacaec 100644
--- a/pybamm/__init__.py
+++ b/pybamm/__init__.py
@@ -58,7 +58,11 @@ from .expression_tree.binary_operators import (
     Multiplication,
     Division,
 )
-from .expression_tree.concatenations import Concatenation, NumpyConcatenation
+from .expression_tree.concatenations import (
+    Concatenation,
+    NumpyConcatenation,
+    DomainConcatenation,
+)
 from .expression_tree.array import Array
 from .expression_tree.matrix import Matrix
 from .expression_tree.parameter import Parameter
diff --git a/pybamm/discretisations/base_discretisation.py b/pybamm/discretisations/base_discretisation.py
index a19743e4..c27fddfa 100644
--- a/pybamm/discretisations/base_discretisation.py
+++ b/pybamm/discretisations/base_discretisation.py
@@ -188,14 +188,16 @@ 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
+            new_children = [
+                self.process_symbol(child, y_slices, boundary_conditions)
+                for child in symbol.children
+            ]
+            new_symbol = pybamm.DomainConcatenation(new_children, self.mesh)
+
+            if new_symbol.is_constant():
+                return pybamm.Vector(new_symbol.evaluate())
+
+            return new_symbol
 
         else:
             # hack to copy the symbol but without a parent
@@ -238,49 +240,6 @@ 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.
diff --git a/pybamm/expression_tree/array.py b/pybamm/expression_tree/array.py
index b4950393..162d1095 100644
--- a/pybamm/expression_tree/array.py
+++ b/pybamm/expression_tree/array.py
@@ -10,19 +10,23 @@ class Array(pybamm.Symbol):
     """node in the expression tree that holds an tensor type variable
     (e.g. :class:`numpy.array`)
 
-    Arguements:
-    ``entries``
+    Parameters
+    ----------
+
+    entries : numpy.array
         the array associated with the node
-    ``name``
-        the name of the node. Optional, defaults to ``str(entries)`` if not provided
+    name : str, optional
+        the name of the node
+    domain : iterable of str, optional
+        list of domains the parameter is valid over, defaults to empty list
 
     *Extends:* :class:`Symbol`
     """
 
-    def __init__(self, entries, name=None):
+    def __init__(self, entries, name=None, domain=[]):
         if name is None:
             name = str(entries)
-        super().__init__(name)
+        super().__init__(name, domain=domain)
         self._entries = entries
 
     @property
diff --git a/pybamm/expression_tree/concatenations.py b/pybamm/expression_tree/concatenations.py
index 30aecaa6..1d54cc27 100644
--- a/pybamm/expression_tree/concatenations.py
+++ b/pybamm/expression_tree/concatenations.py
@@ -6,6 +6,7 @@ from __future__ import print_function, unicode_literals
 import pybamm
 
 import numpy as np
+import numbers
 
 
 class Concatenation(pybamm.Symbol):
@@ -23,7 +24,6 @@ class Concatenation(pybamm.Symbol):
     def __init__(self, *children, name=None):
         if name is None:
             name = "concatenation"
-
         domain = self.get_children_domains(children)
         super().__init__(name, children, domain=domain)
 
@@ -73,6 +73,116 @@ class NumpyConcatenation(Concatenation):
 
         super().__init__(*children, name="numpy concatenation")
 
-    def evaluate(self, t, y):
+    def evaluate(self, t=None, y=None):
         """ See :meth:`pybamm.Symbol.evaluate()`. """
         return np.concatenate([child.evaluate(t, y) for child in self.children])
+
+
+class DomainConcatenation(Concatenation):
+    """A node in the expression tree representing a concatenation of symbols.
+
+    It is assumed that each child has a domain, and the final concatenated vector will
+    respect the sizes and ordering of domains established in pybamm.KNOWN_DOMAINS
+
+    **Extends**: :class:`pybamm.Concatenation`
+
+    Parameters
+    ----------
+
+    children : iterable of :class:`pybamm.Symbol`
+        The symbols to concatenate
+
+    mesh : :class:`pybamm.BaseMesh` (or subclass)
+        The underlying mesh for discretisation, used to obtain the number of mesh points
+        in each domain.
+
+    """
+
+    def __init__(self, children, mesh):
+        # Convert any constant symbols in children to a Vector of the right size for
+        # concatenation
+
+        children = list(children)
+
+        for i, child in enumerate(children):
+            if child.is_constant():
+                children[i] = self.process_node_for_concatenate(child, mesh)
+
+        # Allow the base class to sort the domains into the correct order
+        super().__init__(*children, name="numpy concatenation")
+
+        # deal with "whole cell" special case.
+        # need to split the "whole cell" domain up when we calculate slices, then
+        # recombine again afterwards (see below)
+        if self.domain == ["whole cell"]:
+            self.domain = ["negative electrode", "separator", "positive electrode"]
+
+        # create dict of domain => slice of final vector
+        self._slices = self.create_slices(self, mesh)
+
+        # store size of final vector
+        self._size = self._slices[self.domain[-1]].stop
+
+        # deal with "whole cell" special case
+        if self.domain == ["negative electrode", "separator", "positive electrode"]:
+            self.domain = ["whole cell"]
+
+        # create disc of domain => slice for each child
+        self._children_slices = []
+        for child in self.children:
+            self._children_slices.append(self.create_slices(child, mesh))
+
+    def create_slices(self, node, mesh):
+        slices = {}
+        start = 0
+        end = 0
+        for dom in node.domain:
+            end += mesh[dom].npts
+            slices[dom] = slice(start, end)
+            start = end
+        return slices
+
+    def process_node_for_concatenate(self, node, mesh):
+        """
+        the node is assumed to be constant in time. this function replaces it with a
+        single Vector node with the correct length vector (according to its domain)
+
+        Parameters
+        ----------
+        node: derived from :class:`Symbol`
+            the sub-expression to process (node.is_constant() is true)
+
+        """
+
+        # node must be constant
+        value = node.evaluate()
+
+        # correct size of vector should be number of points in the domains
+        subvector_size = sum([mesh[dom].npts for dom in node.domain])
+
+        # check if its a scalar, if so convert to vector
+        if isinstance(value, numbers.Number):
+            value = np.full(subvector_size, value)
+
+        # check it is the right size
+        if value.size != subvector_size:
+            raise ValueError(
+                "Error: expression evaluated to a vector of incorrect length"
+            )
+
+        # convert to a Vector node
+        return pybamm.Vector(value, domain=node.domain)
+
+    def evaluate(self, t=None, y=None):
+        """ See :meth:`pybamm.Symbol.evaluate()`. """
+
+        # preallocate vector
+        vector = np.empty(self._size)
+
+        # loop through domains of children writing subvectors to final vector
+        for child, slices in zip(self.children, self._children_slices):
+            child_vector = child.evaluate(t, y)
+            for dom in child.domain:
+                vector[self._slices[dom]] = child_vector[slices[dom]]
+
+        return vector
diff --git a/pybamm/expression_tree/symbol.py b/pybamm/expression_tree/symbol.py
index fa0a3178..5e855da2 100644
--- a/pybamm/expression_tree/symbol.py
+++ b/pybamm/expression_tree/symbol.py
@@ -284,3 +284,21 @@ class Symbol(anytree.NodeMixin):
                 self, type(self)
             )
         )
+
+    def is_constant(self):
+        """returns true if evaluating the expression is not dependent on `t` or `y`
+
+        See Also
+        --------
+        evaluate : evaluate the expression
+
+        """
+
+        # if any of the nodes are instances of any of these types, then the whole
+        # expression depends on either t or y
+        search_types = (pybamm.Variable, pybamm.StateVector, pybamm.IndependentVariable)
+
+        # do the search, return true if no relevent nodes are found
+        return all([
+            not isinstance(n, search_types) for n in self.pre_order()
+        ])
diff --git a/pybamm/expression_tree/vector.py b/pybamm/expression_tree/vector.py
index 99132fba..b2abb123 100644
--- a/pybamm/expression_tree/vector.py
+++ b/pybamm/expression_tree/vector.py
@@ -18,31 +18,36 @@ class Vector(pybamm.Array):
         the array associated with the node
     name : str, optional
         the name of the node
+    domain : iterable of str, optional
+        list of domains the parameter is valid over, defaults to empty list
 
     """
 
-    def __init__(self, entries, name=None):
-        super().__init__(entries, name=name)
+    def __init__(self, entries, name=None, domain=[]):
+        super().__init__(entries, name=name, domain=domain)
 
 
 class StateVector(pybamm.Symbol):
     """
     node in the expression tree that holds a slice to read from an external vector type
 
-    Arguments:
+    Parameters
+    ----------
 
-    ``y_slice``
+    y_slice: slice
         the slice of an external y to read
-    ``name``
+    name: str, optional
         the name of the node
+    domain : iterable of str, optional
+        list of domains the parameter is valid over, defaults to empty list
 
     *Extends:* :class:`Array`
     """
 
-    def __init__(self, y_slice, name=None):
+    def __init__(self, y_slice, name=None, domain=[]):
         if name is None:
             name = str(y_slice)
-        super().__init__(name=name)
+        super().__init__(name=name, domain=domain)
         self._y_slice = y_slice
 
     @property


Test patch:
diff --git a/tests/test_discretisations/test_base_discretisations.py b/tests/test_discretisations/test_base_discretisations.py
index 7ff51bea..2a13c144 100644
--- a/tests/test_discretisations/test_base_discretisations.py
+++ b/tests/test_discretisations/test_base_discretisations.py
@@ -285,42 +285,6 @@ class TestDiscretise(unittest.TestCase):
         )
         np.testing.assert_array_equal(S0 * T0, model.variables["ST"].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, ["negative electrode", "whole cell"])
-        expected_vector = np.concatenate(
-            [
-                5 * np.ones_like(mesh["negative electrode"].nodes),
-                5 * np.ones_like(mesh["whole cell"].nodes),
-            ]
-        )
-        np.testing.assert_allclose(a_vec.evaluate(), expected_vector)
-
-        a = pybamm.Scalar(5, domain=["negative electrode", "whole cell"])
-        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)
@@ -355,11 +319,6 @@ class TestDiscretise(unittest.TestCase):
         )
         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, y_slices, {})
-
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/test_expression_tree/test_concatenations.py b/tests/test_expression_tree/test_concatenations.py
index a815b4dd..79fcab95 100644
--- a/tests/test_expression_tree/test_concatenations.py
+++ b/tests/test_expression_tree/test_concatenations.py
@@ -6,6 +6,15 @@ import numpy as np
 import unittest
 
 
+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[:30])
+        self["separator"] = self.submeshclass(self["whole cell"].nodes[30:40])
+        self["positive electrode"] = self.submeshclass(self["whole cell"].nodes[40:])
+
+
 class TestConcatenations(unittest.TestCase):
     def test_base_concatenation(self):
         a = pybamm.Symbol("a")
@@ -80,6 +89,51 @@ class TestConcatenations(unittest.TestCase):
             conc.evaluate(None, y), np.concatenate([y, np.array([16]), np.array([3])])
         )
 
+    def test_numpy_domain_concatenation(self):
+        mesh = MeshForTesting()
+        a_dom = ["negative electrode"]
+        b_dom = ["positive electrode"]
+        a = pybamm.Scalar(2, domain=a_dom)
+        b = pybamm.Vector(np.ones_like(mesh[b_dom[0]].nodes), domain=b_dom)
+
+        # concatenate them the "wrong" way round to check they get reordered correctly
+        conc = pybamm.DomainConcatenation([b, a], mesh)
+        np.testing.assert_array_equal(
+            conc.evaluate(),
+            np.concatenate([
+                np.full(mesh[a_dom[0]].npts, 2),
+                np.full(mesh[b_dom[0]].npts, 1)
+            ])
+        )
+
+        # vector child of wrong size will throw
+        b = pybamm.Vector(np.full(mesh[b_dom[0]].npts - 5, 1), domain=b_dom)
+        with self.assertRaises(ValueError):
+            conc = pybamm.DomainConcatenation([b, a], mesh)
+
+        # check the reordering in case a child vector has to be split up
+        a_dom = ["separator"]
+        b_dom = ["negative electrode", "positive electrode"]
+        a = pybamm.Scalar(2, domain=a_dom)
+        b = pybamm.Vector(
+            np.concatenate([np.full(mesh[b_dom[0]].npts, 1),
+                            np.full(mesh[b_dom[1]].npts, 3)]),
+            domain=b_dom
+        )
+
+        conc = pybamm.DomainConcatenation([a, b], mesh)
+        np.testing.assert_array_equal(
+            conc.evaluate(),
+            np.concatenate([
+                np.full(mesh[b_dom[0]].npts, 1),
+                np.full(mesh[a_dom[0]].npts, 2),
+                np.full(mesh[b_dom[1]].npts, 3),
+            ])
+        )
+
+        # check special case: final domain is still ["whole cell"]
+        self.assertEqual(conc.domain, ["whole cell"])
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/test_expression_tree/test_symbol.py b/tests/test_expression_tree/test_symbol.py
index 381abed3..79ecd094 100644
--- a/tests/test_expression_tree/test_symbol.py
+++ b/tests/test_expression_tree/test_symbol.py
@@ -6,6 +6,7 @@ from __future__ import print_function, unicode_literals
 import pybamm
 
 import unittest
+import numpy as np
 
 
 class TestSymbol(unittest.TestCase):
@@ -113,6 +114,25 @@ class TestSymbol(unittest.TestCase):
         with self.assertRaises(NotImplementedError):
             a.evaluate()
 
+    def test_symbol_is_constant(self):
+        a = pybamm.Variable("a")
+        self.assertFalse(a.is_constant())
+
+        a = pybamm.Parameter("a")
+        self.assertTrue(a.is_constant())
+
+        a = pybamm.Scalar(1) * pybamm.Variable("a")
+        self.assertFalse(a.is_constant())
+
+        a = pybamm.Scalar(1) * pybamm.Parameter("a")
+        self.assertTrue(a.is_constant())
+
+        a = pybamm.Scalar(1) * pybamm.StateVector(slice(10))
+        self.assertFalse(a.is_constant())
+
+        a = pybamm.Scalar(1) * pybamm.Vector(np.zeros(10))
+        self.assertTrue(a.is_constant())
+
     def test_symbol_repr(self):
         """
         test that __repr___ returns the string