Fix save_clusters_config to be consistent with distmat.

hmcezar · Dec 19, 2024 · b6ca3ef · b6ca3ef
1 parent eae0650
commit b6ca3ef
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/clusttraj/distmat.py b/clusttraj/distmat.py
@@ -272,7 +272,6 @@ def compute_distmat_line(
             Pview = Pview[prr]
 
             # build the total molecule with the reordered atoms
-            # whereins = np.where(np.isin(np.arange(len(P)), reorderexcl) is True)
             whereins = np.where(np.atleast_1d(np.isin(np.arange(len(P)), exclusions)))
             Pr = np.insert(
                 Pview,

diff --git a/clusttraj/io.py b/clusttraj/io.py
@@ -844,23 +844,30 @@ def save_clusters_config(
 
             # reorder the solvent atoms separately
             if reorder:
+                # if the solute is specified, reorder just the solvent atoms in this step
+                if nsatoms:
+                    exclusions = np.unique(
+                        np.concatenate((np.arange(natoms), reorderexcl))
+                    )
+                else:
+                    exclusions = reorderexcl
+
                 # get the view without the excluded atoms
-                view = np.delete(np.arange(len(P)), reorderexcl)
+                view = np.delete(np.arange(len(P)), exclusions)
                 Pview = P[view]
                 Paview = Pa[view]
 
                 prr = reorder(Qa[view], Paview, Q[view], Pview)
                 Pview = Pview[prr]
 
                 # build the total molecule with the reordered atoms
-                # whereins = np.where(np.isin(np.arange(len(P)), reorderexcl) is True)
                 whereins = np.where(
-                    np.atleast_1d(np.isin(np.arange(len(P)), reorderexcl))
+                    np.atleast_1d(np.isin(np.arange(len(P)), exclusions))
                 )
                 Pr = np.insert(
                     Pview,
                     [x - whereins[0].tolist().index(x) for x in whereins[0]],
-                    P[reorderexcl],
+                    P[exclusions],
                     axis=0,
                 )
             else: