Repository: tinosulzer/PyBaMM

Issue:
Make Spatial Variables class
Class should replace the current `Space` class and allow input of `x`, `y`, `z` or `r` for different spatial variables. Will the need to update 
- Geometry classes to take `SpatialVariable` instead of `IndependentVariable`
- Discretisation of `Space` to discretise the specific variable given (return Vector on the appropriate mesh for `x` and `r`, raise `NotImplementedError` for `y` and `z` to be returned to later when implementing 2D/3D)

Discussion:

Patch:
diff --git a/docs/source/expression_tree/independent_variable.rst b/docs/source/expression_tree/independent_variable.rst
index 731abc9d..bc621d07 100644
--- a/docs/source/expression_tree/independent_variable.rst
+++ b/docs/source/expression_tree/independent_variable.rst
@@ -4,6 +4,11 @@ Independent Variable
 .. autoclass:: pybamm.IndependentVariable
   :members:
 
+.. autoclass:: pybamm.Time
+  :members:
+
+.. autoclass:: pybamm.SpatialVariable
+  :members:
+
 .. autodata:: pybamm.t
   :annotation: = the independent variable time
-
diff --git a/pybamm/__init__.py b/pybamm/__init__.py
index 81627d99..c6133f12 100644
--- a/pybamm/__init__.py
+++ b/pybamm/__init__.py
@@ -86,7 +86,11 @@ from .expression_tree.unary_operators import (
 from .expression_tree.function_parameter import FunctionParameter
 from .expression_tree.scalar import Scalar
 from .expression_tree.variable import Variable
-from .expression_tree.independent_variable import IndependentVariable, Time, Space
+from .expression_tree.independent_variable import (
+    IndependentVariable,
+    Time,
+    SpatialVariable,
+)
 from .expression_tree.independent_variable import t
 from .expression_tree.vector import Vector, StateVector
 
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index 0d6637d2..eb807b48 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -237,7 +237,7 @@ class Discretisation(object):
         elif isinstance(symbol, pybamm.Variable):
             return pybamm.StateVector(self._y_slices[symbol.id])
 
-        elif isinstance(symbol, pybamm.Space):
+        elif isinstance(symbol, pybamm.SpatialVariable):
             return self._spatial_methods[symbol.domain[0]].spatial_variable(symbol)
 
         elif isinstance(symbol, pybamm.Concatenation):
diff --git a/pybamm/expression_tree/independent_variable.py b/pybamm/expression_tree/independent_variable.py
index d7b968bb..ca500770 100644
--- a/pybamm/expression_tree/independent_variable.py
+++ b/pybamm/expression_tree/independent_variable.py
@@ -41,23 +41,33 @@ class Time(IndependentVariable):
         return t
 
 
-class Space(IndependentVariable):
-    """A node in the expression tree representing time
+class SpatialVariable(IndependentVariable):
+    """A node in the expression tree representing a spatial variable
 
     Parameters
     ----------
     name : str
-        name of the node
+        name of the node ("x", "y", "z" or "r")
     domain : iterable of str
         list of domains that this variable is valid over
 
     *Extends:* :class:`Symbol`
     """
 
-    def __init__(self, domain):
-        if domain is []:
+    def __init__(self, name, domain):
+        if name not in ["x", "y", "z", "r"]:
+            raise ValueError(
+                "name must be 'x', 'y', 'z' or 'r' but is '{}'".format(name)
+            )
+        if domain == []:
             raise ValueError("domain must be provided")
-        name = "space ({})".format(domain)
+        if name == "r" and domain not in [["negative particle"], ["positive particle"]]:
+            raise pybamm.DomainError("domain must be particle if name is 'r'")
+        elif name in ["x", "y", "z"] and any(["particle" in dom for dom in domain]):
+            raise pybamm.DomainError(
+                "domain cannot be particle if name is '{}'".format(name)
+            )
+
         super().__init__(name, domain=domain)
 
 
diff --git a/pybamm/geometry/geometry.py b/pybamm/geometry/geometry.py
index 019dbc4f..2179e760 100644
--- a/pybamm/geometry/geometry.py
+++ b/pybamm/geometry/geometry.py
@@ -38,7 +38,7 @@ class Geometry1DMacro(Geometry):
     def __init__(self, custom_geometry={}):
         super().__init__()
         whole_cell = ["negative electrode", "separator", "positive electrode"]
-        x = pybamm.IndependentVariable("x", whole_cell)
+        x = pybamm.SpatialVariable("x", whole_cell)
         ln = pybamm.standard_parameters.ln
         ls = pybamm.standard_parameters.ls
 
@@ -54,13 +54,14 @@ class Geometry1DMicro(Geometry):
     def __init__(self, custom_geometry={}):
         super().__init__()
 
-        r = pybamm.IndependentVariable("r", ["negative particle", "positive particle"])
+        rn = pybamm.SpatialVariable("r", ["negative particle"])
+        rp = pybamm.SpatialVariable("r", ["positive particle"])
 
         self["negative particle"] = {
-            r: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)}
+            rn: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)}
         }
         self["positive particle"] = {
-            r: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)}
+            rp: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)}
         }
         # update with custom geometry if non empty
         self.update(custom_geometry)
@@ -81,8 +82,8 @@ class Geometry3DMacro(Geometry1DMacro):
         super().__init__()
 
         whole_cell = ["negative electrode", "separator", "positive electrode"]
-        y = pybamm.IndependentVariable("y", whole_cell)
-        z = pybamm.IndependentVariable("z", whole_cell)
+        y = pybamm.SpatialVariable("y", whole_cell)
+        z = pybamm.SpatialVariable("z", whole_cell)
 
         y_lim = {"min": pybamm.Scalar(0), "max": pybamm.standard_parameters.ly}
 
diff --git a/pybamm/spatial_methods/finite_volume.py b/pybamm/spatial_methods/finite_volume.py
index 86cc3f68..18968ad5 100644
--- a/pybamm/spatial_methods/finite_volume.py
+++ b/pybamm/spatial_methods/finite_volume.py
@@ -23,7 +23,7 @@ class FiniteVolume(pybamm.SpatialMethod):
     """
 
     def __init__(self, mesh):
-        self._mesh = mesh
+        super().__init__(mesh)
 
     def spatial_variable(self, symbol):
         """
@@ -40,10 +40,12 @@ class FiniteVolume(pybamm.SpatialMethod):
         :class:`pybamm.Vector`
             Contains the discretised spatial variable
         """
-
         # for finite volume we use the cell centres
-        symbol_mesh = self._mesh.combine_submeshes(*symbol.domain)
-        return pybamm.Vector(symbol_mesh.nodes)
+        if symbol.name in ["x", "r"]:
+            symbol_mesh = self.mesh.combine_submeshes(*symbol.domain)
+            return pybamm.Vector(symbol_mesh.nodes)
+        else:
+            raise NotImplementedError("3D meshes not yet implemented")
 
     def broadcast(self, symbol, domain):
         """
@@ -54,7 +56,7 @@ class FiniteVolume(pybamm.SpatialMethod):
         """
 
         # for finite volume we send variables to cells and so use number_of_cells
-        number_of_cells = {dom: submesh.npts for dom, submesh in self._mesh.items()}
+        number_of_cells = {dom: submesh.npts for dom, submesh in self.mesh.items()}
         broadcasted_symbol = pybamm.NumpyBroadcast(symbol, domain, number_of_cells)
 
         # if the broadcasted symbol evaluates to a constant value, replace the
@@ -108,7 +110,7 @@ class FiniteVolume(pybamm.SpatialMethod):
             The (sparse) finite volume gradient matrix for the domain
         """
         # Create appropriate submesh by combining submeshes in domain
-        submesh = self._mesh.combine_submeshes(*domain)
+        submesh = self.mesh.combine_submeshes(*domain)
 
         # Create matrix using submesh
         n = submesh.npts
@@ -146,7 +148,7 @@ class FiniteVolume(pybamm.SpatialMethod):
             # implement spherical operator
             divergence_matrix = self.divergence_matrix(domain)
 
-            submesh = self._mesh[domain[0]]
+            submesh = self.mesh[domain[0]]
             r = pybamm.Vector(submesh.nodes)
             r_edges = pybamm.Vector(submesh.edges)
 
@@ -175,7 +177,7 @@ class FiniteVolume(pybamm.SpatialMethod):
             The (sparse) finite volume divergence matrix for the domain
         """
         # Create appropriate submesh by combining submeshes in domain
-        submesh = self._mesh.combine_submeshes(*domain)
+        submesh = self.mesh.combine_submeshes(*domain)
 
         # Create matrix using submesh
         n = submesh.npts + 1
diff --git a/pybamm/spatial_methods/spatial_method.py b/pybamm/spatial_methods/spatial_method.py
index b92104b7..18927420 100644
--- a/pybamm/spatial_methods/spatial_method.py
+++ b/pybamm/spatial_methods/spatial_method.py
@@ -1,8 +1,6 @@
 #
 # A general spatial method class
 #
-from __future__ import absolute_import, division
-from __future__ import print_function, unicode_literals
 
 
 class SpatialMethod:
@@ -19,7 +17,11 @@ class SpatialMethod:
     """
 
     def __init__(self, mesh):
-        self.mesh = mesh
+        self._mesh = mesh
+
+    @property
+    def mesh(self):
+        return self._mesh
 
     def spatial_variable(self, symbol):
         """


Test patch:
diff --git a/tests/shared.py b/tests/shared.py
index 3c7fb6cd..4b1f44e8 100644
--- a/tests/shared.py
+++ b/tests/shared.py
@@ -14,7 +14,6 @@ class SpatialMethodForTesting(pybamm.SpatialMethod):
         super().__init__(mesh)
 
     def spatial_variable(self, symbol):
-        # for finite volume we use the cell centres
         symbol_mesh = self._mesh.combine_submeshes(*symbol.domain)
         return pybamm.Vector(symbol_mesh.nodes)
 
diff --git a/tests/test_discretisations/test_base_discretisations.py b/tests/test_discretisations/test_base_discretisations.py
index 9bd1e0ff..c3006a23 100644
--- a/tests/test_discretisations/test_base_discretisations.py
+++ b/tests/test_discretisations/test_base_discretisations.py
@@ -505,22 +505,19 @@ class TestDiscretise(unittest.TestCase):
         )
         np.testing.assert_allclose(eqn_disc.evaluate(), expected_vector)
 
-    def test_discretise_space(self):
-        # variables
-        x1 = pybamm.Space(["negative electrode"])
-        x2 = pybamm.Space(["negative electrode", "separator"])
-        x3 = 3 * pybamm.Space(["negative electrode"])
-
+    def test_discretise_spatial_variable(self):
         # create discretisation
         disc = get_discretisation_for_testing()
 
         # space
+        x1 = pybamm.SpatialVariable("x", ["negative electrode"])
         x1_disc = disc.process_symbol(x1)
         self.assertIsInstance(x1_disc, pybamm.Vector)
         np.testing.assert_array_equal(
             x1_disc.evaluate(), disc.mesh["negative electrode"].nodes
         )
 
+        x2 = pybamm.SpatialVariable("x", ["negative electrode", "separator"])
         x2_disc = disc.process_symbol(x2)
         self.assertIsInstance(x2_disc, pybamm.Vector)
         np.testing.assert_array_equal(
@@ -528,10 +525,11 @@ class TestDiscretise(unittest.TestCase):
             disc.mesh.combine_submeshes("negative electrode", "separator").nodes,
         )
 
-        x3_disc = disc.process_symbol(x3)
-        self.assertIsInstance(x3_disc.children[1], pybamm.Vector)
+        r = 3 * pybamm.SpatialVariable("r", ["negative particle"])
+        r_disc = disc.process_symbol(r)
+        self.assertIsInstance(r_disc.children[1], pybamm.Vector)
         np.testing.assert_array_equal(
-            x3_disc.evaluate(), 3 * disc.mesh["negative electrode"].nodes
+            r_disc.evaluate(), 3 * disc.mesh["negative particle"].nodes
         )
 
 
diff --git a/tests/test_expression_tree/test_independent_variable.py b/tests/test_expression_tree/test_independent_variable.py
index df0eacd5..82afb695 100644
--- a/tests/test_expression_tree/test_independent_variable.py
+++ b/tests/test_expression_tree/test_independent_variable.py
@@ -33,12 +33,27 @@ class TestIndependentVariable(unittest.TestCase):
         with self.assertRaises(ValueError):
             t.evaluate(None)
 
-    def test_space(self):
-        x = pybamm.Space(["negative electrode"])
-        self.assertEqual(x.name, "space (['negative electrode'])")
+    def test_spatial_variable(self):
+        x = pybamm.SpatialVariable("x", ["negative electrode"])
+        self.assertEqual(x.name, "x")
+        y = pybamm.SpatialVariable("y", ["separator"])
+        self.assertEqual(y.name, "y")
+        z = pybamm.SpatialVariable("z", ["positive electrode"])
+        self.assertEqual(z.name, "z")
+        r = pybamm.SpatialVariable("r", ["negative particle"])
+        self.assertEqual(r.name, "r")
         with self.assertRaises(NotImplementedError):
             x.evaluate()
 
+        with self.assertRaisesRegex(ValueError, "name must be"):
+            pybamm.SpatialVariable("not a variable", ["negative electrode"])
+        with self.assertRaisesRegex(ValueError, "domain must be"):
+            pybamm.SpatialVariable("x", [])
+        with self.assertRaises(pybamm.DomainError):
+            pybamm.SpatialVariable("r", ["negative electrode"])
+        with self.assertRaises(pybamm.DomainError):
+            pybamm.SpatialVariable("x", ["negative particle"])
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/test_geometry/test_geometry.py b/tests/test_geometry/test_geometry.py
index 3cea6591..cbdcd133 100644
--- a/tests/test_geometry/test_geometry.py
+++ b/tests/test_geometry/test_geometry.py
@@ -9,7 +9,7 @@ class TestGeometry1DMacro(unittest.TestCase):
     def test_add_custom_geometry(self):
         geometry = pybamm.Geometry1DMacro()
         whole_cell = ["negative electrode", "separator", "positive electrode"]
-        x = pybamm.IndependentVariable("x", whole_cell)
+        x = pybamm.SpatialVariable("x", whole_cell)
         custom_geometry = {
             "negative electrode": {
                 x: {"min": pybamm.Scalar(1), "max": pybamm.Scalar(2)}
@@ -24,7 +24,7 @@ class TestGeometry1DMacro(unittest.TestCase):
         geometry = pybamm.Geometry1DMacro()
         for spatial_vars in geometry.values():
             all(
-                self.assertIsInstance(spatial_var, pybamm.IndependentVariable)
+                self.assertIsInstance(spatial_var, pybamm.SpatialVariable)
                 for spatial_var in spatial_vars.keys()
             )
 
@@ -33,7 +33,7 @@ class TestGeometry1DMicro(unittest.TestCase):
     def test_add_custom_geometry(self):
         geometry = pybamm.Geometry1DMicro()
         whole_cell = ["negative electrode", "separator", "positive electrode"]
-        x = pybamm.IndependentVariable("x", whole_cell)
+        x = pybamm.SpatialVariable("x", whole_cell)
         custom_geometry = {
             "negative electrode": {
                 x: {"min": pybamm.Scalar(1), "max": pybamm.Scalar(2)}
@@ -49,7 +49,7 @@ class TestGeometry1DMicro(unittest.TestCase):
         geometry = pybamm.Geometry1DMicro()
         for spatial_vars in geometry.values():
             all(
-                self.assertIsInstance(spatial_var, pybamm.IndependentVariable)
+                self.assertIsInstance(spatial_var, pybamm.SpatialVariable)
                 for spatial_var in spatial_vars.keys()
             )
 
@@ -58,7 +58,7 @@ class TestGeometry3DMacro(unittest.TestCase):
     def test_add_custom_geometry(self):
         geometry = pybamm.Geometry3DMacro()
         whole_cell = ["negative electrode", "separator", "positive electrode"]
-        x = pybamm.IndependentVariable("x", whole_cell)
+        x = pybamm.SpatialVariable("x", whole_cell)
         custom_geometry = {
             "negative electrode": {
                 x: {"min": pybamm.Scalar(1), "max": pybamm.Scalar(2)}
@@ -74,7 +74,7 @@ class TestGeometry3DMacro(unittest.TestCase):
         geometry = pybamm.Geometry3DMacro()
         for spatial_vars in geometry.values():
             all(
-                self.assertIsInstance(spatial_var, pybamm.IndependentVariable)
+                self.assertIsInstance(spatial_var, pybamm.SpatialVariable)
                 for spatial_var in spatial_vars.keys()
             )
 
@@ -97,7 +97,7 @@ class TestGeometry(unittest.TestCase):
 
         # update with custom geometry
         whole_cell = ["negative electrode", "separator", "positive electrode"]
-        x = pybamm.IndependentVariable("x", whole_cell)
+        x = pybamm.SpatialVariable("x", whole_cell)
         custom_geometry = {
             "negative electrode": {
                 x: {"min": pybamm.Scalar(1), "max": pybamm.Scalar(2)}
diff --git a/tests/test_spatial_methods/test_finite_volume.py b/tests/test_spatial_methods/test_finite_volume.py
index 76590990..637c1ff7 100644
--- a/tests/test_spatial_methods/test_finite_volume.py
+++ b/tests/test_spatial_methods/test_finite_volume.py
@@ -619,6 +619,43 @@ class TestFiniteVolume(unittest.TestCase):
         rates = np.log2(errs[:-1] / errs[1:])
         np.testing.assert_array_less(1.99 * np.ones_like(rates), rates)
 
+    def test_discretise_spatial_variable(self):
+        # create discretisation
+        mesh = get_mesh_for_testing()
+        spatial_methods = {
+            "macroscale": pybamm.FiniteVolume,
+            "negative particle": pybamm.FiniteVolume,
+            "positive particle": pybamm.FiniteVolume,
+        }
+        disc = pybamm.Discretisation(mesh, spatial_methods)
+
+        # space
+        x1 = pybamm.SpatialVariable("x", ["negative electrode"])
+        x1_disc = disc.process_symbol(x1)
+        self.assertIsInstance(x1_disc, pybamm.Vector)
+        np.testing.assert_array_equal(
+            x1_disc.evaluate(), disc.mesh["negative electrode"].nodes
+        )
+
+        z = pybamm.SpatialVariable("z", ["negative electrode"])
+        with self.assertRaises(NotImplementedError):
+            disc.process_symbol(z)
+
+        x2 = pybamm.SpatialVariable("x", ["negative electrode", "separator"])
+        x2_disc = disc.process_symbol(x2)
+        self.assertIsInstance(x2_disc, pybamm.Vector)
+        np.testing.assert_array_equal(
+            x2_disc.evaluate(),
+            disc.mesh.combine_submeshes("negative electrode", "separator").nodes,
+        )
+
+        r = 3 * pybamm.SpatialVariable("r", ["negative particle"])
+        r_disc = disc.process_symbol(r)
+        self.assertIsInstance(r_disc.children[1], pybamm.Vector)
+        np.testing.assert_array_equal(
+            r_disc.evaluate(), 3 * disc.mesh["negative particle"].nodes
+        )
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")