Repository: pybamm-team/PyBaMM

Issue:
Store variable slices
Store variable slices in the model after discretising, for easy inspection

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4895abe06..2eb59a98d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## Features
 
+-   Store variable slices in model for inspection ([#925](https://github.com/pybamm-team/PyBaMM/pull/925))
+-   Made t_plus (optionally) a function of electrolyte concentration, and added (1 + dlnf/dlnc) to models ([#921](https://github.com/pybamm-team/PyBaMM/pull/921))
 -   Added `DummySolver` for empty models ([#915](https://github.com/pybamm-team/PyBaMM/pull/915))
 -   Added functionality to broadcast to edges ([#891](https://github.com/pybamm-team/PyBaMM/pull/891))
 -   Reformatted and cleaned up `QuickPlot` ([#886](https://github.com/pybamm-team/PyBaMM/pull/886))
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index 431c812e7..9a84cffd3 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -137,6 +137,8 @@ class Discretisation(object):
         # Set the y split for variables
         pybamm.logger.info("Set variable slices for {}".format(model.name))
         self.set_variable_slices(variables)
+        # Keep a record of y_slices in the model
+        model.y_slices = self.y_slices_explicit
 
         # now add extrapolated external variables to the boundary conditions
         # if required by the spatial method
@@ -214,6 +216,7 @@ class Discretisation(object):
         """
         # Set up y_slices
         y_slices = defaultdict(list)
+        y_slices_explicit = defaultdict(list)
         start = 0
         end = 0
         # Iterate through unpacked variables, adding appropriate slices to y_slices
@@ -233,14 +236,18 @@ class Discretisation(object):
                             submesh = domain_mesh[i]
                             end += submesh.npts_for_broadcast_to_nodes
                         y_slices[child.id].append(slice(start, end))
+                        y_slices_explicit[child].append(slice(start, end))
                         start = end
             else:
                 end += self._get_variable_size(variable)
                 y_slices[variable.id].append(slice(start, end))
+                y_slices_explicit[variable].append(slice(start, end))
                 start = end
 
         # Convert y_slices back to normal dictionary
         self.y_slices = dict(y_slices)
+        # Also keep a record of what the y_slices are, to be stored in the model
+        self.y_slices_explicit = dict(y_slices_explicit)
 
         # reset discretised_symbols
         self._discretised_symbols = {}


Test patch:
diff --git a/tests/unit/test_discretisations/test_discretisation.py b/tests/unit/test_discretisations/test_discretisation.py
index ce4cf0838..48a09b98a 100644
--- a/tests/unit/test_discretisations/test_discretisation.py
+++ b/tests/unit/test_discretisations/test_discretisation.py
@@ -132,6 +132,8 @@ class TestDiscretise(unittest.TestCase):
 
         self.assertEqual(disc.y_slices[a.id][0], slice(0, 10, None))
 
+        self.assertEqual(model.y_slices[a][0], slice(0, 10, None))
+
         b_test = np.ones((10, 1))
         np.testing.assert_array_equal(
             model.variables["b"].evaluate(inputs={"b": b_test}), b_test