[contacts._md_compute_contacts] use new _residue_sidechain_membership…

… to gracefully compute residue memberships for the sidechain schemes

mdciao/contacts/_md_compute_contacts.py

@@ -28,10 +28,11 @@
 #
 #    The modifications consist in:
 #    1. including the indices of the closest atom-pairs in the returned values.
-#    2. adapting 'sidechain' and 'sidechain-heavy' schemes for the cases of hydrogen-less glycines.
+#    2. adapting 'sidechain' and 'sidechain-heavy' schemes for the cases of hydrogen-less glycines
+#       and non-protein residues.
 #
 #    The modified lines are in the documentation for the returned value `atom_pairs`
-#    and the 'sidechain' and 'sidechain-heavy', and elsewhere marked with the comment
+#    and the 'sidechain' and 'sidechain-heavy' arguments, and elsewhere marked with the comment
 #    '#mdciao' in the file.
 #
 #    The modifications were applied on mdtraj release v1.10.0rc1
@@ -74,7 +75,7 @@
 from mdtraj.core import element
 from mdtraj.utils import ensure_type
 
-
+from mdciao.utils.residue_and_atom import _residue_sidechain_membership #mdciao
 def compute_contacts(
     traj,
     contacts="all",
@@ -105,11 +106,11 @@ def compute_contacts(
                 any two non-hydrogen atoms in the residues
             'sidechain' : distance is the closest distance between any
                 two atoms in residue sidechains. For glycine,
-                'sidechain' try to use sidechain hydrogens first
+                'sidechain' tries to use sidechain hydrogens first
                  and if none are present fall back to any atom of the glycine.
             'sidechain-heavy' : distance is the closest distance between
                 any two non-hydrogen atoms in residue sidechains
-                For glycine, 'sidechain-heavy' try to fall back
+                For glycine, 'sidechain-heavy' tries to fall back
                 first to sidechain hydrogens and if none are present
                 to any atom of the glycine.
     ignore_nonprotein : bool
@@ -268,17 +269,7 @@ def compute_contacts(
                     "is present, the distances involving glycine residues will be " #mdciao
                     "computed using the all glycine atoms."  # mdciao
                 ) #mdciao
-            residue_membership = [ #mdciao
-                [ #mdciao
-                    atom.index for atom in residue.atoms #mdciao
-                    if atom.is_sidechain #mdciao
-                ]
-                if not residue.name=="GLY" #mdciao
-                else [[atom.index for atom in residue.atoms if atom.is_sidechain] #mdciao
-                      if len([atom.index for atom in residue.atoms if atom.is_sidechain]) > 0 #mdciao
-                      else [atom.index for atom in residue.atoms]][0] #mdciao
-                for residue in traj.topology.residues #mdciao
-            ] #mdciao
+            residue_membership = [_residue_sidechain_membership(scheme,residue) for residue in traj.topology.residues]
         elif scheme == "sidechain-heavy":
             # then remove the hydrogens from the above list
             if "GLY" in [residue.name for residue in traj.topology.residues]:
@@ -290,21 +281,7 @@ def compute_contacts(
                     "computed using the sidechain hydrogen instead. If no sidechain hydrogen is present," #mdciao
                     " all glycine atoms will be used." #mdciao
                 )
-
-            residue_membership = [
-                (
-                    [
-                        atom.index
-                        for atom in residue.atoms
-                        if atom.is_sidechain and not (atom.element == element.hydrogen)
-                    ]
-                    if not residue.name == "GLY"
-                    else [[atom.index for atom in residue.atoms if atom.is_sidechain]
-                          if len([atom.index for atom in residue.atoms if atom.is_sidechain]) > 0 #mdciao
-                          else [atom.index for atom in residue.atoms]][0] #mdciao
-                )
-                for residue in traj.topology.residues
-            ]
+            residue_membership = [_residue_sidechain_membership(scheme, residue) for residue in traj.topology.residues] #mdciao
 
         residue_lens = [len(ainds) for ainds in residue_membership]