Added radius argument to list-neighbours

q2_makarsa/_neighbours.py

@@ -1,6 +1,6 @@
 # Everybody needs good neighbours.
 
-from networkx import Graph
+from networkx import Graph, shortest_path_length
 from pandas import DataFrame
 from qiime2 import Metadata
 
@@ -10,6 +10,7 @@
 def list_neighbours(
         network: Graph,
         feature_id: str,
+        radius: int = 1,
         metadata: Metadata = None) -> Metadata:
     """List the neighbours of a feature in a network.
 
@@ -39,24 +40,27 @@ def list_neighbours(
         raise ValueError(f"Feature {feature_id} not found in network")
     feature_node = feature_node[0]
 
-    # Find the neighbours of the feature node
-    neighbours = list(network.neighbors(feature_node))
+    # Find the nodes within the given radius
+    nodes_within_radius = [
+        (node, path_length) for node, path_length in
+        shortest_path_length(network, source=feature_node).items()
+        if 0 < path_length <= radius]
 
     # Create a table of the neighbours, where the index is
     # the feature ID, and the columns are the node attributes.
     result = DataFrame(
-        [dict(network.nodes[x]) for x in neighbours],
-        index=neighbours)
+        [dict(network.nodes[node],
+              **{f"Weight to {feature_id}":
+                 network.edges[feature_node, node]["weight"]
+                 if path_length == 1 else float("nan"),
+                 f"Path Length to {feature_id}": path_length})
+         for node, path_length in nodes_within_radius],
+        index=list(zip(*nodes_within_radius))[0])
     result.index = result["Feature"]
     result = result.drop(columns=["Feature"])
     result.index.name = "Feature ID"
     for column in result.columns:
         if column.startswith("Taxon Level "):
             result = result.drop(columns=[column])
 
-    # Add a column to result for the edge weight between
-    # the feature and each neighbour.
-    result[f"Weight to {feature_id}"] = [
-        network.edges[feature_node, x]["weight"] for x in neighbours]
-
     return Metadata(result)

q2_makarsa/plugin_setup.py

@@ -312,14 +312,15 @@ def _6(data: q2.Metadata) -> NodeMapFormat:
 plugin.methods.register_function(
     function=list_neighbours,
     inputs={"network": Network},
-    parameters={"feature_id": Str, "metadata": Metadata},
+    parameters={"feature_id": Str, "radius": Int, "metadata": Metadata},
     outputs=[("neighbours", ImmutableMetadata)],
     input_descriptions={
         "network": "The network to search."
     },
     parameter_descriptions={
         "metadata": "Any additional metadata to include in the output.",
-        "feature_id": "The feature for which to list neighbours."
+        "feature_id": "The feature for which to list neighbours.",
+        "radius": "Include neighbourhood up to radius away."
     },
     output_descriptions={
         "neighbours": "The neighbours of the feature, with metadata."

q2_makarsa/tests/data/neighbourhood.tsv

@@ -0,0 +1,30 @@
+Feature ID	Community	Degree Centrality	Betweenness Centrality	Closeness Centrality	Eigenvector Centrality	Assortativity	Weight to rbpyb	Path Length to rbpyb
+i0tPB	7.0	0.05084745762711865	0.021784809774289607	0.2226502311248074	0.08378280613317068	9.666666666666666	-0.0968280795126354	1.0
+X4K3Hh	7.0	0.1016949152542373	0.0009935710111046172	0.22520942918371326	0.2037895497454346	10.166666666666666	-0.0201288984160626	1.0
+gc64H	8.0	0.1016949152542373	0.0024454278107112707	0.2389417114510128	0.20882079490729386	11.0	-0.157175005292387	1.0
+X871v	7.0	0.1016949152542373	0.0006582060059558598	0.2226502311248074	0.21100645843176735	10.5	0.128295555295925	1.0
+TYAZm	7.0	0.0847457627118644	0.00015054106228857834	0.2177024482109228	0.18547909239044177	11.2	0.152439276133693	1.0
+CTvJP	2.0	0.2033898305084746	0.045156119936949866	0.2721280602636535	0.2943277003055724	7.333333333333333	-0.112638376699115	1.0
+xBo9q	7.0	0.0847457627118644	0.037405026300409115	0.22520942918371326	0.1558295655013325	10.0	0.070345334195137	1.0
+X6nItH	7.0	0.13559322033898305	0.00630526691309859	0.2480154473288994	0.26041170386839774	9.75	-0.250909414101931	1.0
+BDqLV	8.0	0.2711864406779661	0.06906212496510566	0.2968669748330765	0.3971926455568289	7.0	-0.0918249884094724	1.0
+X0aHk0	8.0	0.1864406779661017	0.02345366497908871	0.2647732478240953	0.31349862803890255	8.818181818181818	-0.0688174382569234	1.0
+N0gX0	7.0	0.2542372881355932	0.06299333824173099	0.2924361244624336	0.3680983379998797	7.0	-0.28013675186314	1.0
+aCVgY	0.0	0.05084745762711865	0.02333829233303225	0.19399228058399062	0.015272030389552476	2.6666666666666665		2.0
+aITW7	8.0	0.11864406779661017	0.035783230990712	0.2578055307760928	0.14998934038322886	7.0		2.0
+X7H3D	2.0	0.0847457627118644	0.021065375302663437	0.21296978629329402	0.11113268306470583	7.4		2.0
+pVaq5	2.0	0.03389830508474576	0.0	0.1902254401843015	0.04832065099196568	8.5		2.0
+kga4g	2.0	0.05084745762711865	0.0	0.21531011361519836	0.1262797563701204	14.333333333333334		2.0
+NmKte	2.0	0.1016949152542373	0.021517167162988314	0.23050847457627122	0.1794335804081999	10.0		2.0
+FCfLi	8.0	0.03389830508474576	0.0	0.20199196225755722	0.08241173609156337	14.0		2.0
+BfyaV	8.0	0.13559322033898305	0.07710666825219718	0.2721280602636535	0.18640609227390426	7.75		2.0
+xIkpV	2.0	0.05084745762711865	0.0	0.21531011361519836	0.1262797563701204	14.333333333333334		2.0
+Xnu7S	4.0	0.03389830508474576	0.01928696668614845	0.16604424016087332	0.01883843182118977	3.0		2.0
+FPTtD	8.0	0.03389830508474576	0.0	0.20199196225755722	0.05523627286802827	9.0		2.0
+bPSEb	8.0	0.05084745762711865	0.0003010821245771567	0.2226502311248074	0.08395814147270868	10.0		2.0
+X8Hlao	0.0	0.01694915254237288	0.0	0.1462180622312168	0.0018200842038891662	3.0		3.0
+psRbW	0.0	0.06779661016949153	0.09389714123787764	0.23050847457627122	0.04254418034508322	5.5		3.0
+pOqLH	2.0	0.01694915254237288	0.0	0.1567457627118644	0.013244209531342244	5.0		3.0
+CKzvN	2.0	0.01694915254237288	0.0	0.16604424016087332	0.02138396384166743	6.0		3.0
+X3KD3L	8.0	0.01694915254237288	0.0	0.18660209846650527	0.02221503001224044	8.0		3.0
+X0XOLj	4.0	0.01694915254237288	0.0	0.12975642608598048	0.0022450434620669172	2.0		3.0

q2_makarsa/tests/data/neighbours.tsv

@@ -1,3 +1,3 @@
-Feature ID	Community	Degree Centrality	Betweenness Centrality	Closeness Centrality	Eigenvector Centrality	Assortativity	Weight to PfvSC
-xXbIa	9.0	0.03389830508474576	0.0011689070718877848	0.038135593220338986	1.3011000855318294e-09	1.5	-0.0691309772699902
-ZS7qE	9.0	0.01694915254237288	0.0	0.025423728813559324	8.041240756241885e-10	2.0	-0.00061793520262609
+Feature ID	Community	Degree Centrality	Betweenness Centrality	Closeness Centrality	Eigenvector Centrality	Assortativity	Weight to PfvSC	Path Length to PfvSC
+xXbIa	9.0	0.03389830508474576	0.0011689070718877848	0.038135593220338986	1.3011000855318294e-09	1.5	-0.0691309772699902	1.0
+ZS7qE	9.0	0.01694915254237288	0.0	0.025423728813559324	8.041240756241885e-10	2.0	-0.00061793520262609	1.0

q2_makarsa/tests/test_neighbours.py

@@ -19,5 +19,20 @@ def test_list_neighbours(self):
             sep='\t', header=0, index_col=0)
         metadata.index.name = "feature id"
         metadata = q2.Metadata(metadata)
-        observed = list_neighbours(network, "PfvSC", metadata).to_dataframe()
+        observed = list_neighbours(
+            network, "PfvSC", metadata=metadata).to_dataframe()
+        pd.testing.assert_frame_equal(observed, expected)
+
+    def test_list_neighbours_with_radius(self):
+        expected = pd.read_csv(
+            self.get_data_path("neighbourhood.tsv"),
+            sep='\t', header=0, index_col=0)
+        network = read_graphml(self.get_data_path("network.graphml"))
+        metadata = pd.read_csv(
+            self.get_data_path("louvain.tsv"),
+            sep='\t', header=0, index_col=0)
+        metadata.index.name = "feature id"
+        metadata = q2.Metadata(metadata)
+        observed = list_neighbours(
+            network, "rbpyb", radius=3, metadata=metadata).to_dataframe()
         pd.testing.assert_frame_equal(observed, expected)