Improvement of TMD Protein Logic

docs/source/how-to-run.rst

@@ -200,7 +200,7 @@ Probability value that is used for z-score calculation. It is a determination pe
 
 * ``tmd_protein_list``
 
-Transmembrane domain (TMD) protein list to include area per lipid calculation. TMD proteins take up space in the exoplasmic, cytoplasmic leaflets. Three backbone atoms of protein that are in close position to lipid head groups should be included in this parameter to increase the success of identification.
+Transmembrane domain (TMD) protein list to include area per lipid calculation. TMD proteins take up space in the exoplasmic, cytoplasmic leaflets. Three backbone atoms of protein that are in close position to lipid head groups should be included in this parameter to increase the success of identification. Backbone atoms can be provided by MDAnalysis atom group or a string query for MDAnalysis selection.
 
 .. code-block::
 

domhmm/analysis/base.py

@@ -259,24 +259,19 @@ def __init__(
                             "Entry for each TDM protein should be a dictionary in the format {'0': ..., '1': ...} "
                             "where 0 for lower leaflet and 1 for upper leaflet.")
                     if isinstance(query, AtomGroup):
-                        # Take center of geometry of three positions
-                        cog = np.mean(query.positions, axis=0)
-                        self.tmd_protein[leaflet].append(cog)
+                        self.tmd_protein[leaflet].append(query)
                     # Character string was provided as input, assume it contains a selection for an MDAnalysis.AtomGroup
                     elif isinstance(query, str):
                         # Try to create a MDAnalysis.AtomGroup, raise a ValueError if not selection group could be
                         # provided
                         try:
-                            cog = np.mean(self.universe.select_atoms(query).positions, axis=0)
-                            self.tmd_protein[leaflet].append(cog)
+                            self.tmd_protein[leaflet].append(self.universe.select_atoms(query))
                         except Exception as e:
                             raise ValueError("Please provide a valid MDAnalysis selection string!") from e
                     else:
                         raise ValueError(
                             "TDM Protein list should contain AtomGroup from MDAnalysis universe or a string "
                             "query for MDAnalysis selection.")
-            self.tmd_protein["0"] = np.array(self.tmd_protein["0"])
-            self.tmd_protein["1"] = np.array(self.tmd_protein["1"])
         elif tmd_protein_list is not None:
             # An unknown argument is provided for tdm_protein_list
             raise ValueError(

domhmm/analysis/domhmm.py

@@ -321,11 +321,11 @@ def _single_frame(self):
         lower_coor_xy = self.leaflet_selection[str(1)].positions
         # Check Transmembrane domain existence
         if self.tmd_protein is not None:
-            tmd_upper_coor_xy = self.tmd_protein["0"]
+            tmd_upper_coor_xy = np.array([np.mean(bb.positions, axis=0) for bb in self.tmd_protein["0"]])
             # Check if dimension of coordinates is same in both array
             if tmd_upper_coor_xy.shape[1] == upper_coor_xy.shape[1]:
                 upper_coor_xy = np.append(upper_coor_xy, tmd_upper_coor_xy, axis=0)
-            tmd_lower_coor_xy = self.tmd_protein["1"]
+            tmd_lower_coor_xy = np.array([np.mean(bb.positions, axis=0) for bb in self.tmd_protein["1"]])
             # Check if dimension of coordinates is same in both array
             if tmd_lower_coor_xy.shape[1] == lower_coor_xy.shape[1]:
                 lower_coor_xy = np.append(lower_coor_xy, tmd_lower_coor_xy, axis=0)
@@ -1403,9 +1403,16 @@ def clustering_plot(self):
                               positions[idx, 1],
                               s=100, marker="o", color=colors[j], zorder=-10)
             if self.tmd_protein is not None:
-                label_length += len(self.tmd_protein["0"])
-                for protein in self.tmd_protein["0"]:
-                    ax[k].scatter(protein[0], protein[1], s=100, marker="^", color="black", label="TMD Protein")
+                label_length += 1
+                proteins = np.array([np.mean(bb.positions, axis=0) for bb in self.tmd_protein["0"]])
+                ax[k].scatter(
+                    proteins[:, 0],
+                    proteins[:, 1],
+                    s=100,
+                    marker="^",
+                    color="black",
+                    label="TMD Protein",
+                )
             ax[k].set_xticks([])
             ax[k].set_yticks([])
             ax[k].set_aspect("equal")

domhmm/tests/test_domhmm.py

@@ -176,9 +176,15 @@ def result_parameter_check(analysis, test_type):
             assert analysis.results['HMM_Pred']['DPPC'].shape == (302, 100)
             assert analysis.results['HMM_Pred']['DIPC'].shape == (202, 100)
             assert analysis.results['HMM_Pred']['CHOL'].shape == (216, 100)
+
+            # Set the first frame for tmd protein check
+            _ = analysis.universe.trajectory[0]
             assert analysis.tmd_protein.keys() == {'0', '1'}
-            assert np.allclose(analysis.tmd_protein['0'], [[63.04445, 86.886116, 60.87222]], error_tolerance)
-            assert np.allclose(analysis.tmd_protein['1'], [[63.04445,  86.886116, 60.87222 ]], error_tolerance)
+            # Calculate center of geometry of three backbone atoms for both leaflets
+            tmd_upper_coor_xy = np.array([np.mean(bb.positions, axis=0) for bb in analysis.tmd_protein["0"]])
+            tmd_lower_coor_xy = np.array([np.mean(bb.positions, axis=0) for bb in analysis.tmd_protein["1"]])
+            assert np.allclose(tmd_upper_coor_xy, [[63.04445, 86.886116, 60.87222]], error_tolerance)
+            assert np.allclose(tmd_lower_coor_xy, [[63.04445,  86.886116, 60.87222 ]], error_tolerance)
         assert len(analysis.results['HMM_Pred']) == 3
         assert analysis.results['HMM_Pred'].keys() == {'DPPC', 'DIPC', 'CHOL'}
         assert len(analysis.results['Getis_Ord']) == 4