Repository: pybamm-team/PyBaMM
Pull request #653 (2019-10-09 14:38:08)

#652 cache evaluation array
# Description

Reuse evaluation array when making a copy of a state vector
Fixes #653

## Type of change

- [x] New feature (non-breaking change which adds functionality)

# Key checklist:

- [x] No style issues: `$ flake8`
- [ ] All tests pass: `$ python run-tests.py --unit`
- [x] The documentation builds: `$ cd docs` and then `$ make clean; make html`

You can run all three at once, using `$ python run-tests.py --quick`.

## Further checks: 

- [x] Code is commented, particularly in hard-to-understand areas
- [x] Tests added that prove fix is effective or that feature works
- [x] Any dependent changes have been merged and published in downstream modules



Patch:
diff --git a/pybamm/expression_tree/state_vector.py b/pybamm/expression_tree/state_vector.py
index 1023764a6..dbdd7c425 100644
--- a/pybamm/expression_tree/state_vector.py
+++ b/pybamm/expression_tree/state_vector.py
@@ -22,11 +22,21 @@ class StateVector(pybamm.Symbol):
         list of domains the parameter is valid over, defaults to empty list
     auxiliary_domains : dict of str, optional
         dictionary of auxiliary domains
+    evaluation_array : list, optional
+        List of boolean arrays representing slices. Default is None, in which case the
+        evaluation_array is computed from y_slices.
 
     *Extends:* :class:`Array`
     """
 
-    def __init__(self, *y_slices, name=None, domain=None, auxiliary_domains=None):
+    def __init__(
+        self,
+        *y_slices,
+        name=None,
+        domain=None,
+        auxiliary_domains=None,
+        evaluation_array=None,
+    ):
         for y_slice in y_slices:
             if not isinstance(y_slice, slice):
                 raise TypeError("all y_slices must be slice objects")
@@ -52,7 +62,7 @@ class StateVector(pybamm.Symbol):
         self._y_slices = y_slices
         self._first_point = y_slices[0].start
         self._last_point = y_slices[-1].stop
-        self.set_evaluation_array(y_slices)
+        self.set_evaluation_array(y_slices, evaluation_array)
         super().__init__(name=name, domain=domain, auxiliary_domains=auxiliary_domains)
 
     @property
@@ -76,12 +86,15 @@ class StateVector(pybamm.Symbol):
     def size(self):
         return self.evaluation_array.count(True)
 
-    def set_evaluation_array(self, y_slices):
+    def set_evaluation_array(self, y_slices, evaluation_array):
         "Set evaluation array using slices"
-        array = np.zeros(y_slices[-1].stop)
-        for y_slice in y_slices:
-            array[y_slice] = True
-        self._evaluation_array = [bool(x) for x in array]
+        if evaluation_array is not None and pybamm.settings.debug_mode is False:
+            self._evaluation_array = evaluation_array
+        else:
+            array = np.zeros(y_slices[-1].stop)
+            for y_slice in y_slices:
+                array[y_slice] = True
+            self._evaluation_array = [bool(x) for x in array]
 
     def set_id(self):
         """ See :meth:`pybamm.Symbol.set_id()` """
@@ -156,7 +169,8 @@ class StateVector(pybamm.Symbol):
             *self.y_slices,
             name=self.name,
             domain=self.domain,
-            auxiliary_domains=self.auxiliary_domains
+            auxiliary_domains=self.auxiliary_domains,
+            evaluation_array=self.evaluation_array,
         )
 
     def evaluate_for_shape(self):


Test patch:
diff --git a/tests/unit/test_expression_tree/test_state_vector.py b/tests/unit/test_expression_tree/test_state_vector.py
index 78d8368de..f6809565d 100644
--- a/tests/unit/test_expression_tree/test_state_vector.py
+++ b/tests/unit/test_expression_tree/test_state_vector.py
@@ -45,6 +45,18 @@ class TestStateVector(unittest.TestCase):
         )
         self.assertEqual(sv.name, "y[0:10,20:30,...,60:70]")
 
+    def test_pass_evaluation_array(self):
+        # Turn off debug mode for this test
+        original_debug_mode = pybamm.settings.debug_mode
+        pybamm.settings.debug_mode = False
+        # Test that evaluation array gets passed down (doesn't have to be the correct
+        # array for this test)
+        array = np.array([1, 2, 3, 4, 5])
+        sv = pybamm.StateVector(slice(0, 10), evaluation_array=array)
+        np.testing.assert_array_equal(sv.evaluation_array, array)
+        # Turn debug mode back to what is was before
+        pybamm.settings.debug_mode = original_debug_mode
+
     def test_failure(self):
         with self.assertRaisesRegex(TypeError, "all y_slices must be slice objects"):
             pybamm.StateVector(slice(0, 10), 1)