Change node numbering to atomic data

LBC-LNBio · Sep 21, 2021 · bd6d6f5 · bd6d6f5
1 parent ce5439b
commit bd6d6f5
Show file tree

Hide file tree

Showing 2 changed files with 121 additions and 41 deletions.
diff --git a/SERD/__init__.py b/SERD/__init__.py
@@ -824,22 +824,24 @@ def r2g(
     if selection != "all":
         atomic = _get_atom_selection(atomic, selection=selection)
 
-    # Keep solvent exposed residues
-    # https://stackoverflow.com/questions/51352527/check-for-identical-rows-in-different-numpy-arrays
-    atomic = atomic[(atomic[:, 0:3][:, None] == residues).all(-1).any(-1)]
-
     # Calculate distance
     distance = _calculate_distance(atomic[:, 4:7])
 
     if selection == "all":
         distance = _all_atoms_to_residues(atomic, distance)
+        atomic = _get_atom_selection(atomic, selection=selection)
 
     # Calculate adjacency matrix
     if intraresidual:
         adjacency = (distance < cutoff).astype(int)
     else:
         adjacency = numpy.logical_and(distance > 0.0, distance < cutoff).astype(int)
 
+    # Keep solvent exposed residues
+    # https://stackoverflow.com/questions/51352527/check-for-identical-rows-in-different-numpy-arrays
+    adjacency[:, ~(atomic[:, 0:3][:, None] == residues).all(-1).any(-1)] = 0
+    adjacency[~(atomic[:, 0:3][:, None] == residues).all(-1).any(-1), :] = 0
+
     # Create networkx.Graph
     G = networkx.Graph()
     G.add_edges_from(numpy.argwhere(adjacency))

diff --git a/notebooks/1FMO.ipynb b/notebooks/1FMO.ipynb