fix: only sort transform if model is already sorted (#118)

xarray-contrib · Nov 17, 2023 · dfe0001 · dfe0001
1 parent 1878406
commit dfe0001
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 9 deletions.
diff --git a/xeofs/models/eof_rotator.py b/xeofs/models/eof_rotator.py
@@ -242,9 +242,10 @@ def _transform_algorithm(self, data: DataArray) -> DataArray:
         projections = xr.dot(projections, RinvT, dims="mode_m")
         # Reorder according to variance
         # this must be done in one line: i) select modes according to their variance, ii) replace coords with modes from 1 ... n
-        projections = projections.isel(
-            mode=self.data["idx_modes_sorted"].values
-        ).assign_coords(mode=projections.mode)
+        if self.sorted:
+            projections = projections.isel(
+                mode=self.data["idx_modes_sorted"].values
+            ).assign_coords(mode=projections.mode)
 
         # Scale scores by "pseudo" norms
         projections = projections * pseudo_norms

diff --git a/xeofs/models/mca_rotator.py b/xeofs/models/mca_rotator.py
@@ -363,9 +363,10 @@ def transform(self, **kwargs) -> DataArray | List[DataArray]:
             projections1 = projections1.rename({"mode": "mode_m"})
             projections1 = xr.dot(projections1, RinvT, dims="mode_m")
             # Reorder according to variance
-            projections1 = projections1.isel(
-                mode=self.data["idx_modes_sorted"].values
-            ).assign_coords(mode=projections1.mode)
+            if self.sorted:
+                projections1 = projections1.isel(
+                    mode=self.data["idx_modes_sorted"].values
+                ).assign_coords(mode=projections1.mode)
             # Adapt the sign of the scores
             projections1 = projections1 * self.data["modes_sign"]
 
@@ -389,9 +390,10 @@ def transform(self, **kwargs) -> DataArray | List[DataArray]:
             projections2 = projections2.rename({"mode": "mode_m"})
             projections2 = xr.dot(projections2, RinvT, dims="mode_m")
             # Reorder according to variance
-            projections2 = projections2.isel(
-                mode=self.data["idx_modes_sorted"].values
-            ).assign_coords(mode=projections2.mode)
+            if self.sorted:
+                projections2 = projections2.isel(
+                    mode=self.data["idx_modes_sorted"].values
+                ).assign_coords(mode=projections2.mode)
             # Determine the sign of the scores
             projections2 = projections2 * self.data["modes_sign"]