diff --git a/.gitignore b/.gitignore
index 2865c03..c8100f6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -15,3 +15,5 @@
 **/results*.npy
 c_models/Debug/
 c_models/Release/
+/.mypy_cache/
+/.pytest_cache/
diff --git a/mcmc/mcmc_coordinator_vertex.py b/mcmc/mcmc_coordinator_vertex.py
index 6e42ba0..f788181 100644
--- a/mcmc/mcmc_coordinator_vertex.py
+++ b/mcmc/mcmc_coordinator_vertex.py
@@ -156,14 +156,14 @@ def get_data_window_size(self, placement):
 
     def get_sequence_mask(self, placement, routing_info):
         if self._is_receiver_placement(placement):
-            mask = routing_info.get_routing_info_from_pre_vertex(
+            mask = routing_info.get_info_from(
                 self, self._data_partition_name).mask
             return ~mask & 0xFFFFFFFF
         return 0
 
     def get_acknowledge_key(self, placement, routing_info):
         if self._is_receiver_placement(placement):
-            key = routing_info.get_first_key_from_pre_vertex(
+            key = routing_info.get_key_from(
                 placement.vertex, self._acknowledge_partition_name)
             return key
         return 0
@@ -224,10 +224,9 @@ def generate_data_specification(
             mcmc_placement = FecDataView.get_placement_of_vertex(vertex)
             self._mcmc_placements.append(mcmc_placement)
             if self._is_receiver_placement(mcmc_placement):
-                key = routing_info.get_first_key_from_pre_vertex(
+                key = routing_info.get_key_from(
                     vertex, self._acknowledge_partition_name)
-                if key is not None:
-                    keys.append(key)
+                keys.append(key)
         keys.sort()
 
         # Write the data size in words
@@ -239,17 +238,16 @@ def generate_data_specification(
         spec.write_value(len(keys), data_type=DataType.UINT32)
 
         # Write the key
-        vertex_routing_info = routing_info.get_routing_info_from_pre_vertex(
+        vtx_routing_info = routing_info.get_info_from(
             self, self._data_partition_name)
-        assert vertex_routing_info is not None
-        spec.write_value(vertex_routing_info.key, data_type=DataType.UINT32)
+        spec.write_value(vtx_routing_info.key, data_type=DataType.UINT32)
 
         # Write the window size
         spec.write_value(self._window_size, data_type=DataType.UINT32)
 
         # Write the sequence mask
         spec.write_value(
-            ~vertex_routing_info.mask & 0xFFFFFFFF, data_type=DataType.UINT32)
+            ~vtx_routing_info.mask & 0xFFFFFFFF, data_type=DataType.UINT32)
 
         # Write the timer
         spec.write_value(self._send_timer, data_type=DataType.UINT32)
diff --git a/mcmc/mcmc_vertex.py b/mcmc/mcmc_vertex.py
index 77e24fb..74262a1 100644
--- a/mcmc/mcmc_vertex.py
+++ b/mcmc/mcmc_vertex.py
@@ -147,7 +147,7 @@ def _get_model_state_array(self):
     def get_result_key(self, placement):
         if self._is_receiver_placement(placement):
             routing_info = FecDataView.get_routing_infos()
-            key = routing_info.get_first_key_from_pre_vertex(
+            key = routing_info.get_key_from(
                 placement.vertex, self._result_partition_name)
             return key
         return 0
@@ -163,7 +163,7 @@ def _is_receiver_placement(self, placement):
     def get_cholesky_result_key(self, placement):
         if self._is_cholesky_receiver_placement(placement):
             routing_info = FecDataView.get_routing_infos()
-            key = routing_info.get_first_key_from_pre_vertex(
+            key = routing_info.get_key_from(
                 placement.vertex, self._cholesky_result_partition_name)
             return key
         return 0
@@ -259,19 +259,17 @@ def generate_data_specification(
         spec.write_value(self._coordinator.acknowledge_timer)
 
         # Write the (first) key for sending parameter data, if needed
-        if (self._model.root_finder):
-            routing_info_rf = routing_info.get_routing_info_from_pre_vertex(
+        if self._model.root_finder:
+            rinfo_rf = routing_info.get_info_from(
                 self, self._parameter_partition_name)
-            assert routing_info_rf is not None
-            spec.write_value(routing_info_rf.key, data_type=DataType.UINT32)
+            spec.write_value(rinfo_rf.key, data_type=DataType.UINT32)
         else:
             spec.write_value(0, data_type=DataType.UINT32)
 
-        if (self._model.cholesky):
-            routing_info_ch = routing_info.get_routing_info_from_pre_vertex(
+        if self._model.cholesky:
+            rinfo_ch = routing_info.get_info_from(
                 self, self._cholesky_partition_name)
-            assert routing_info_ch is not None
-            spec.write_value(routing_info_ch.key, data_type=DataType.UINT32)
+            spec.write_value(rinfo_ch.key, data_type=DataType.UINT32)
         else:
             spec.write_value(0, data_type=DataType.UINT32)