Repository: pybamm-team/PyBaMM

Issue:
Overload `symbol.__getitem__` to return `Index`

Discussion:

Patch:
diff --git a/pybamm/expression_tree/symbol.py b/pybamm/expression_tree/symbol.py
index 09e4873db..16a54905b 100644
--- a/pybamm/expression_tree/symbol.py
+++ b/pybamm/expression_tree/symbol.py
@@ -303,6 +303,10 @@ class Symbol(anytree.NodeMixin):
         """return an :class:`AbsoluteValue` object"""
         return pybamm.AbsoluteValue(self)
 
+    def __getitem__(self, key):
+        """return a :class:`Index` object"""
+        return pybamm.Index(self, key)
+
     def diff(self, variable):
         """
         Differentiate a symbol with respect to a variable. Default behaviour is to
diff --git a/pybamm/spatial_methods/finite_volume.py b/pybamm/spatial_methods/finite_volume.py
index 4cfb67e82..e82fb5231 100644
--- a/pybamm/spatial_methods/finite_volume.py
+++ b/pybamm/spatial_methods/finite_volume.py
@@ -404,11 +404,11 @@ class FiniteVolume(pybamm.SpatialMethod):
             if lbc_value.evaluates_to_number():
                 lbc_i = lbc_value
             else:
-                lbc_i = pybamm.Index(lbc_value, i)
+                lbc_i = lbc_value[i]
             if rbc_value.evaluates_to_number():
                 rbc_i = rbc_value
             else:
-                rbc_i = pybamm.Index(rbc_value, i)
+                rbc_i = rbc_value[i]
 
             if lbc_type == "Dirichlet":
                 left_ghost_constant = 2 * lbc_i


Test patch:
diff --git a/tests/unit/test_expression_tree/test_unary_operators.py b/tests/unit/test_expression_tree/test_unary_operators.py
index 0cdeceb8e..deeb5914d 100644
--- a/tests/unit/test_expression_tree/test_unary_operators.py
+++ b/tests/unit/test_expression_tree/test_unary_operators.py
@@ -111,13 +111,19 @@ class TestUnaryOperators(unittest.TestCase):
     def test_index(self):
         vec = pybamm.Vector(np.array([1, 2, 3, 4, 5]))
         # with integer
-        ind = pybamm.Index(vec, 3)
+        ind = vec[3]
+        self.assertIsInstance(ind, pybamm.Index)
+        self.assertEqual(ind.slice, slice(3, 4))
         self.assertEqual(ind.evaluate(), 4)
         # with slice
-        ind = pybamm.Index(vec, slice(1, 3))
+        ind = vec[1:3]
+        self.assertIsInstance(ind, pybamm.Index)
+        self.assertEqual(ind.slice, slice(1, 3))
         np.testing.assert_array_equal(ind.evaluate(), np.array([[2], [3]]))
         # with only stop slice
-        ind = pybamm.Index(vec, slice(3))
+        ind = vec[:3]
+        self.assertIsInstance(ind, pybamm.Index)
+        self.assertEqual(ind.slice, slice(3))
         np.testing.assert_array_equal(ind.evaluate(), np.array([[1], [2], [3]]))
 
         # errors