diff --git a/spinn_pdp2/input_vertex.py b/spinn_pdp2/input_vertex.py
index 1ad1184..1503e23 100644
--- a/spinn_pdp2/input_vertex.py
+++ b/spinn_pdp2/input_vertex.py
@@ -294,15 +294,13 @@ def generate_machine_data_specification(
         spec.switch_write_focus (MLPRegions.ROUTING.value)
 
         # write link keys: fwd
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.fwd_link)
-        assert key is not None
         spec.write_value(key, data_type=DataType.UINT32)
 
         # write link keys: bkp
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.bkp_link)
-        assert key is not None
         spec.write_value (key, data_type = DataType.UINT32)
 
         # write link keys: bps (padding)
diff --git a/spinn_pdp2/sum_vertex.py b/spinn_pdp2/sum_vertex.py
index 9158283..6d72b5b 100644
--- a/spinn_pdp2/sum_vertex.py
+++ b/spinn_pdp2/sum_vertex.py
@@ -319,15 +319,13 @@ def generate_machine_data_specification(
         spec.switch_write_focus (MLPRegions.ROUTING.value)
 
         # write link keys: fwd
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.fwd_link)
-        assert key is not None
         spec.write_value(key, data_type=DataType.UINT32)
 
         # write link keys: bkp
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.bkp_link)
-        assert key is not None
         spec.write_value(key, data_type=DataType.UINT32)
 
         # write link keys: bps (padding)
@@ -337,15 +335,13 @@ def generate_machine_data_specification(
         spec.write_value (0, data_type = DataType.UINT32)
 
         # write link keys: lds
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.lds_link)
-        assert key is not None
         spec.write_value(key, data_type=DataType.UINT32)
 
         # write link keys: fsg
-        key = routing_info.get_first_key_from_pre_vertex (
+        key = routing_info.get_key_from (
             self, self.fsg_link)
-        assert key is not None
         spec.write_value(key, data_type=DataType.UINT32)
 
         # Reserve and write the stage configuration region
diff --git a/spinn_pdp2/threshold_vertex.py b/spinn_pdp2/threshold_vertex.py
index 2932203..b01c0e1 100644
--- a/spinn_pdp2/threshold_vertex.py
+++ b/spinn_pdp2/threshold_vertex.py
@@ -483,24 +483,21 @@ def generate_machine_data_specification(
         spec.switch_write_focus (MLPRegions.ROUTING.value)
 
         # write link keys: fwd
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.fwd_link)
-        assert key is not None
         spec.write_value(key, data_type=DataType.UINT32)
 
         # write link keys: bkp
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.bkp_link)
-        assert key is not None
         spec.write_value (key, data_type = DataType.UINT32)
 
         # write link keys: bps (padding)
         spec.write_value (0, data_type = DataType.UINT32)
 
         # write link keys: stp
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.stp_link)
-        assert key is not None
         spec.write_value(key, data_type=DataType.UINT32)
 
         # write link keys: lds (padding)
diff --git a/spinn_pdp2/weight_vertex.py b/spinn_pdp2/weight_vertex.py
index f7598dd..f3bb8eb 100644
--- a/spinn_pdp2/weight_vertex.py
+++ b/spinn_pdp2/weight_vertex.py
@@ -353,15 +353,13 @@ def generate_machine_data_specification(
         spec.switch_write_focus (MLPRegions.ROUTING.value)
 
         # write link keys: fwd
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.fwd_link)
-        assert key is not None
         spec.write_value(key, data_type=DataType.UINT32)
 
         # write link keys: bkp
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.bkp_link)
-        assert key is not None
         spec.write_value(key, data_type=DataType.UINT32)
 
         # write link keys: bps (padding)
@@ -371,15 +369,13 @@ def generate_machine_data_specification(
         spec.write_value (0, data_type = DataType.UINT32)
 
         # write link keys: lds
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.lds_link)
-        assert key is not None
         spec.write_value(key, data_type=DataType.UINT32)
 
         # write link keys: fsg
-        key = routing_info.get_first_key_from_pre_vertex(
+        key = routing_info.get_key_from(
             self, self.fsg_link)
-        assert key is not None
         spec.write_value(key, data_type=DataType.UINT32)
 
         # Reserve and write the stage configuration region