change clustering results to residue ids

Author: yusuferentunc

domhmm/analysis/base.py

@@ -140,6 +140,7 @@ def __init__(
         self.membrane = universe_or_atomgroup.select_atoms(membrane_select)
         self.membrane_unique_resids = np.unique(self.membrane.resids)
         self.resids_index_map = {resid: idx for idx, resid in enumerate(self.membrane_unique_resids)}
+        self.index_resid_map = {idx: resid for idx, resid in enumerate(self.membrane_unique_resids)}
         self.heads = heads
         self.tails = tails
         self.sterol_heads = sterol_heads

domhmm/analysis/domhmm.py

@@ -1396,7 +1396,11 @@ def result_clustering(self):
                                                         w_ii_f=self.results["Getis_Ord"][j][f"w_ii_{j}"][i],
                                                         core_lipids=core_lipids)
                 frame_number = self.start + i * self.step
-                self.results["Clustering"][str(j)][frame_number] = list(clusters.values())
+                leaflet_index_resid_map = self.get_leaflet_step_index_to_resid(j, i)
+                cluster_result = []
+                for cluster in clusters.values():
+                    cluster_result.append([leaflet_index_resid_map[leaflet_index] for leaflet_index in cluster])
+                self.results["Clustering"][str(j)][frame_number] = cluster_result
 
     @staticmethod
     def assign_core_lipids(weight_matrix_f, g_star_i_f, order_states_f, w_ii_f, z_score):
@@ -1560,7 +1564,7 @@ def get_leaflet_step_order(self, leaflet, step):
         #Init two empty lists for ...
         temp = [] #... order states prediction
         idxs = [] #... indices of lipids
-        
+
         #Iterate over lipids
         for res, data in self.results.train_data_per_type.items():
 
@@ -1583,11 +1587,11 @@ def get_leaflet_step_order(self, leaflet, step):
         return order_states
 
 
-    
+
 
     def get_leaflet_step_order_index(self, leaflet, step):
         """
-        Receive residue's indexes and positions for a specific leaflet at any frame of the trajecytory.
+        Receive residue's indexes and positions for a specific leaflet at any frame of the trajectory.
 
         Parameters
         ----------
@@ -1625,3 +1629,15 @@ def get_leaflet_step_order_index(self, leaflet, step):
                     f"resname {res} and name {self.sterol_heads[res]}")).positions
 
         return indexes, positions
+
+    def get_leaflet_step_index_to_resid(self, leaflet, step):
+        leaflet_assignment_step = self.leaflet_assignment[:, step]
+        leaflet_assignment_mask = leaflet_assignment_step == leaflet
+        result_map = {}
+        for resname in self.unique_resnames:
+            sys_index = np.where(self.membrane.residues.resnames == resname)[0]
+            sys_index =sys_index[leaflet_assignment_mask[sys_index]]
+            leaflet_index = np.where(self.membrane.residues[leaflet_assignment_mask].resnames == resname)[0]
+            for i in range(0, len(leaflet_index)):
+                result_map[leaflet_index[i]] = self.index_resid_map[sys_index[i]]
+        return result_map