Move code here for vertices too

pacman/model/graphs/application/application_vertex.py

@@ -289,3 +289,57 @@ def add_incoming_edge(self, edge, partition):
         :type partition:
             ~pacman.model.graphs.application.ApplicationEdgePartition
         """
+
+    def get_key_ordered_indices(self, indices=None):
+        """
+        Get indices of the vertex in the order that atoms appear when the
+        vertex is split into cores as determined by max_atoms_per_core.  When
+        a multi-dimensional vertex is split into cores, the atoms on each
+        vertex is not linear but rather a hyper-rectangle of the atoms, thus
+        the order of the atoms in the vertex as a whole is not the same as the
+        order of the vertex when scanning over the cores.
+
+        :param indices:
+            Optional subset of indices to convert.  If not provided all indices
+            will be converted.
+        :type indices: numpy.ndarray or None.
+        :rtype: numpy.ndarray
+        """
+        if indices is None:
+            indices = numpy.arange(self.n_atoms)
+        atoms_shape = self.atoms_shape
+        n_dims = len(atoms_shape)
+        if n_dims == 1:
+            return indices
+        atoms_per_core = self.get_max_atoms_per_dimension_per_core()
+        remainders = numpy.array(indices)
+        cum_per_core = 1
+        cum_cores_per_dim = 1
+        core_index = numpy.zeros(len(indices))
+        atom_index = numpy.zeros(len(indices))
+        for n in range(n_dims):
+            # Work out the global index in this dimension
+            global_index_d = remainders % atoms_shape[n]
+
+            # Work out the core index in this dimension
+            core_index_d = global_index_d // atoms_per_core[n]
+
+            # Update the core index by multiplying the current value by the
+            # last dimension size and then add the current dimension value
+            core_index += core_index_d * cum_cores_per_dim
+
+            # Work out the atom index on the core in this dimension
+            atom_index_d = global_index_d - (atoms_per_core[n] * core_index_d)
+
+            # Update the atom index by multiplying the current value by the
+            # last dimension size and then add the current dimension value
+            atom_index += atom_index_d * cum_per_core
+
+            # Update the remainders for next time based on the values in
+            # this dimension, and save the sizes for this dimension
+            remainders = remainders // atoms_shape[n]
+            cum_per_core *= atoms_per_core[n]
+            cum_cores_per_dim *= atoms_shape[n] / atoms_per_core[n]
+
+        return ((core_index * self.get_max_atoms_per_core()) +
+                atom_index).astype(numpy.uint32)

pacman/model/graphs/common/mdslice.py

@@ -154,3 +154,28 @@ def from_string(cls, as_str):
         return MDSlice(
             lo_atom, lo_atom + size - 1, tuple(shape), tuple(start),
             atoms_shape)
+
+    @overrides(Slice.get_relative_indices)
+    def get_relative_indices(self, app_vertex_indices):
+        n_dims = len(self._atoms_shape)
+        remainders = app_vertex_indices
+        last_per_core = 0
+        rel_index = numpy.zeros(len(app_vertex_indices))
+        for n in n_dims:
+            # Work out the index in this dimension
+            global_index_d = remainders % self._atoms_shape[n]
+
+            # Work out the index in this dimension relative to the core start
+            rel_index_d = global_index_d - self._starts[n]
+
+            # Update the total relative index using the position in this
+            # dimension
+            rel_index = (rel_index * last_per_core) + rel_index_d
+
+            # Prepare for next round of the loop by removing what we used
+            # of the global index and remembering the sizes in this
+            # dimension
+            remainders = remainders // self._atoms_shape[n]
+            last_per_core = self._shape[n]
+
+        return rel_index

pacman/model/graphs/common/slice.py

@@ -202,3 +202,16 @@ def from_string(cls, as_str):
         lo_atom = int(parts[0])
         hi_atom = int(parts[1])
         return Slice(lo_atom, hi_atom)
+
+    def get_relative_indices(self, app_vertex_indices):
+        """
+        Convert from global indices to slice-level indices.
+
+        Note that no checking is done on the given indices; they should be
+        within this slice!
+
+        :param numpy.ndarray app_vertex_indices:
+            The global application vertex indices to convert
+        :return The local core-level indices relative to this slice
+        """
+        return app_vertex_indices - self._lo_atom

unittests/model_tests/application_graph_tests/test_application_vertex.py

@@ -19,6 +19,7 @@
     PacmanConfigurationException, PacmanInvalidParameterException)
 from pacman.model.graphs.common import Slice, ChipAndCore
 from pacman.model.graphs.machine import SimpleMachineVertex
+from pacman.model.graphs.application import ApplicationVertex
 from pacman.model.partitioner_splitters import SplitterFixedLegacy
 from pacman_test_objects import SimpleTestVertex
 
@@ -31,6 +32,22 @@ def reset_called(self):
         self.reset_seen = True
 
 
+class SimpleMDVertex(ApplicationVertex):
+
+    def __init__(self, max_atoms_per_core, atoms_shape):
+        super(SimpleMDVertex, self).__init__(
+            max_atoms_per_core=max_atoms_per_core)
+        self.__atoms_shape = atoms_shape
+
+    @property
+    def n_atoms(self):
+        return numpy.prod(self.__atoms_shape)
+
+    @property
+    def atoms_shape(self):
+        return self.__atoms_shape
+
+
 class TestApplicationGraphModel(unittest.TestCase):
     """
     tests which test the application graph object
@@ -123,3 +140,23 @@ def test_set_label(self):
         vert.addedToGraph()
         with self.assertRaises(PacmanConfigurationException):
             vert.set_label("test 2")
+
+    def test_get_key_ordered_indices(self):
+        vtx = SimpleMDVertex((3, 3), (6, 6))
+        # From the interdimensional compatibility documentation diagram, we can
+        # check we get what we expect
+        first_half_indices = vtx.get_key_ordered_indices(numpy.arange(18))
+        assert numpy.array_equal(
+            first_half_indices,
+            [0, 1, 2, 9, 10, 11, 3, 4, 5, 12, 13, 14, 6, 7, 8, 15, 16, 17])
+
+        # Much more complex - suffice to check that each index exists exactly
+        # once
+        size = (6, 9, 2)
+        n_atoms = numpy.prod(size)
+        vtx_2 = SimpleMDVertex((2, 3, 2), size)
+        row_indices = vtx_2.get_key_ordered_indices(numpy.arange(n_atoms))
+        print(row_indices)
+        test_array = numpy.zeros(n_atoms)
+        test_array[row_indices] += 1
+        assert all(numpy.equal(test_array, 1))