[JTH] add different row dimesnsion recon in pca

Author: tausiaj

bluemath_tk/core/models.py

@@ -34,7 +34,7 @@ class BlueMathModel(ABC):
     @abstractmethod
     def __init__(self) -> None:
         self._logger: logging.Logger = None
-        self._exclude_attributes: List[str] = ["_logger"]
+        self._exclude_attributes: List[str] = []
 
     def __getstate__(self):
         """Exclude certain attributes from being pickled."""

bluemath_tk/datamining/pca.py

@@ -482,13 +482,7 @@ def _reshape_data(self, X: np.ndarray, destandarize: bool = True) -> xr.Dataset:
             )
         }
 
-        return xr.Dataset(
-            X_reshaped_vars_dict,
-            coords={
-                self.pca_dim_for_rows: self.data[self.pca_dim_for_rows],
-                **{coord: self.data[coord] for coord in self.coords_to_stack},
-            },
-        )
+        return X_reshaped_vars_dict
 
     @validate_data_pca
     def fit(
@@ -671,7 +665,15 @@ def inverse_transform(self, PCs: Union[np.ndarray, xr.Dataset]) -> xr.Dataset:
 
         self.logger.info("Inverse transforming data using PCA model")
         X_transformed = self.pca.inverse_transform(X=X)
-        data_transformed = self._reshape_data(X=X_transformed, destandarize=True)
+        data_reshaped_vars_dict = self._reshape_data(X=X_transformed, destandarize=True)
+        # Create xarray Dataset with the transformed data
+        data_transformed = xr.Dataset(
+            data_reshaped_vars_dict,
+            coords={
+                self.pca_dim_for_rows: PCs[self.pca_dim_for_rows].values,
+                **{coord: self.data[coord] for coord in self.coords_to_stack},
+            },
+        )
         # Transpose dimensions based on the original data
         data_transformed = data_transformed.transpose(*self.data.dims)
 

tests/datamining/test_pca.py

@@ -67,10 +67,10 @@ def test_inverse_transform(self):
             coords_to_stack=["coord1", "coord2"],
             pca_dim_for_rows="coord3",
         )
-        reconstructed_ds = self.pca.inverse_transform(PCs=pcs)
+        reconstructed_ds = self.pca.inverse_transform(PCs=pcs.isel(coord3=slice(0, 5)))
         self.assertAlmostEqual(
-            self.ds.isel(coord1=5, coord2=5, coord3=5),
-            reconstructed_ds.isel(coord1=5, coord2=5, coord3=5),
+            self.ds.isel(coord1=5, coord2=5, coord3=1),
+            reconstructed_ds.isel(coord1=5, coord2=5, coord3=1),
         )
 
     def test_incremental_fit(self):