bugfix to distance edge func #146

a-r-j · Mar 28, 2022 · 45b2f7c · 45b2f7c
1 parent 087296d
commit 45b2f7c
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/graphein/protein/edges/distance.py b/graphein/protein/edges/distance.py
@@ -480,9 +480,10 @@ def add_distance_threshold(
     )
     dist_mat = compute_distmat(pdb_df)
     interacting_nodes = get_interacting_atoms(threshold, distmat=dist_mat)
-    interacting_nodes = zip(interacting_nodes[0], interacting_nodes[1])
+    interacting_nodes = list(zip(interacting_nodes[0], interacting_nodes[1]))
 
-    log.info(f"Found: {len(list(interacting_nodes))} distance edges")
+    log.info(f"Found: {len(interacting_nodes)} distance edges")
+    count = 0
     for a1, a2 in interacting_nodes:
         n1 = G.graph["pdb_df"].loc[a1, "node_id"]
         n2 = G.graph["pdb_df"].loc[a2, "node_id"]
@@ -491,16 +492,20 @@ def add_distance_threshold(
         n1_position = G.graph["pdb_df"].loc[a1, "residue_number"]
         n2_position = G.graph["pdb_df"].loc[a2, "residue_number"]
 
-        condition_1 = n1_chain != n2_chain
+        condition_1 = n1_chain == n2_chain
         condition_2 = (
-            abs(n1_position - n2_position) > long_interaction_threshold
+            abs(n1_position - n2_position) < long_interaction_threshold
         )
 
-        if condition_1 or condition_2:
+        if not (condition_1 and condition_2):
+            count += 1
             if G.has_edge(n1, n2):
                 G.edges[n1, n2]["kind"].add("distance_threshold")
             else:
                 G.add_edge(n1, n2, kind={"distance_threshold"})
+    log.info(
+        f"Added {count} distance edges. ({len(list(interacting_nodes)) - count} removed by LIN)"
+    )
 
 
 def add_k_nn_edges(