feat: normalized option in inverse_transform (#124)

tests/models/test_eof.py

@@ -510,13 +510,11 @@ def test_save_load(dim, mock_data_array, tmp_path, engine):
 
     # Test that the recreated model can be used to transform new data
     assert np.allclose(
-        original.scores(), loaded.transform(mock_data_array), rtol=1e-3, atol=1e-3
+        original.transform(mock_data_array), loaded.transform(mock_data_array)
     )
 
     # The loaded model should also be able to inverse_transform new data
     assert np.allclose(
         original.inverse_transform(original.scores()),
         loaded.inverse_transform(loaded.scores()),
-        rtol=1e-3,
-        atol=1e-3,
     )

tests/models/test_eof_rotator.py

@@ -236,12 +236,11 @@ def test_save_load(dim, mock_data_array, tmp_path, engine):
 
     # Test that the recreated model can be used to transform new data
     assert np.allclose(
-        original.scores(), loaded.transform(mock_data_array), rtol=1e-3, atol=1e-3
+        original.transform(mock_data_array), loaded.transform(mock_data_array)
     )
 
     # The loaded model should also be able to inverse_transform new data
     assert np.allclose(
         original.inverse_transform(original.scores()),
         loaded.inverse_transform(loaded.scores()),
-        rtol=1e-2,
     )

tests/models/test_mca.py

@@ -423,16 +423,12 @@ def test_save_load(dim, mock_data_array, tmp_path, engine):
 
     # Test that the recreated model can be used to transform new data
     assert np.allclose(
-        original.scores(),
+        original.transform(mock_data_array, mock_data_array),
         loaded.transform(mock_data_array, mock_data_array),
-        rtol=1e-3,
-        atol=1e-3,
     )
 
     # The loaded model should also be able to inverse_transform new data
     assert np.allclose(
         original.inverse_transform(*original.scores()),
         loaded.inverse_transform(*loaded.scores()),
-        rtol=1e-3,
-        atol=1e-3,
     )

tests/models/test_mca_rotator.py

@@ -295,16 +295,12 @@ def test_save_load(dim, mock_data_array, tmp_path, engine):
 
     # Test that the recreated model can be used to transform new data
     assert np.allclose(
-        original.scores(),
-        loaded.transform(data1=mock_data_array, data2=mock_data_array),
-        rtol=1e-3,
-        atol=1e-3,
+        original.transform(mock_data_array, mock_data_array),
+        loaded.transform(mock_data_array, mock_data_array),
     )
 
     # The loaded model should also be able to inverse_transform new data
     assert np.allclose(
         original.inverse_transform(*original.scores()),
         loaded.inverse_transform(*loaded.scores()),
-        rtol=1e-3,
-        atol=1e-3,
     )

tests/models/test_orthogonality.py

@@ -455,7 +455,8 @@ def test_crmca_scores(dim, use_coslat, power, squared_loadings, mock_data_array)
         (("lon", "lat"), False),
     ],
 )
-def test_eof_transform(dim, use_coslat, mock_data_array):
+@pytest.mark.parametrize("normalized", [True, False])
+def test_eof_transform(dim, use_coslat, mock_data_array, normalized):
     """Transforming the original data results in the model scores"""
     model = EOF(
         n_modes=5,
@@ -464,8 +465,8 @@ def test_eof_transform(dim, use_coslat, mock_data_array):
         random_state=5,
     )
     model.fit(mock_data_array, dim=dim)
-    scores = model.scores()
-    pseudo_scores = model.transform(mock_data_array)
+    scores = model.scores(normalized=normalized)
+    pseudo_scores = model.transform(mock_data_array, normalized=normalized)
     assert np.allclose(
         scores, pseudo_scores, atol=1e-4
     ), "Transformed data does not match the scores"
@@ -480,13 +481,14 @@ def test_eof_transform(dim, use_coslat, mock_data_array):
         (("lon", "lat"), False),
     ],
 )
-def test_ceof_transform(dim, use_coslat, mock_data_array):
+@pytest.mark.parametrize("normalized", [True, False])
+def test_ceof_transform(dim, use_coslat, mock_data_array, normalized):
     """Not implemented yet"""
     model = ComplexEOF(n_modes=5, standardize=True, use_coslat=use_coslat)
     model.fit(mock_data_array, dim=dim)
-    scores = model.scores()
+    scores = model.scores(normalized=normalized)
     with pytest.raises(NotImplementedError):
-        pseudo_scores = model.transform(mock_data_array)
+        pseudo_scores = model.transform(mock_data_array, normalized=normalized)
 
 
 # Rotated EOF
@@ -501,14 +503,15 @@ def test_ceof_transform(dim, use_coslat, mock_data_array):
         (("lon", "lat"), False, 2),
     ],
 )
-def test_reof_transform(dim, use_coslat, power, mock_data_array):
+@pytest.mark.parametrize("normalized", [True, False])
+def test_reof_transform(dim, use_coslat, power, mock_data_array, normalized):
     """Transforming the original data results in the model scores"""
     model = EOF(n_modes=5, standardize=True, use_coslat=use_coslat, random_state=5)
     model.fit(mock_data_array, dim=dim)
     rot = EOFRotator(n_modes=5, power=power)
     rot.fit(model)
-    scores = rot.scores()
-    pseudo_scores = rot.transform(mock_data_array)
+    scores = rot.scores(normalized=normalized)
+    pseudo_scores = rot.transform(mock_data_array, normalized=normalized)
     np.testing.assert_allclose(
         scores,
         pseudo_scores,
@@ -529,15 +532,16 @@ def test_reof_transform(dim, use_coslat, power, mock_data_array):
         (("lon", "lat"), False, 2),
     ],
 )
-def test_creof_transform(dim, use_coslat, power, mock_data_array):
+@pytest.mark.parametrize("normalized", [True, False])
+def test_creof_transform(dim, use_coslat, power, mock_data_array, normalized):
     """not implemented yet"""
     model = ComplexEOF(n_modes=5, standardize=True, use_coslat=use_coslat)
     model.fit(mock_data_array, dim=dim)
     rot = ComplexEOFRotator(n_modes=5, power=power)
     rot.fit(model)
-    scores = rot.scores()
+    scores = rot.scores(normalized=normalized)
     with pytest.raises(NotImplementedError):
-        pseudo_scores = rot.transform(mock_data_array)
+        pseudo_scores = rot.transform(mock_data_array, normalized=normalized)
 
 
 # MCA
@@ -687,13 +691,14 @@ def r2_score(x, y, dim=None):
         (("lon", "lat"), False),
     ],
 )
-def test_eof_inverse_transform(dim, use_coslat, mock_data_array):
+@pytest.mark.parametrize("normalized", [True, False])
+def test_eof_inverse_transform(dim, use_coslat, mock_data_array, normalized):
     """Inverse transform produces an approximate reconstruction of the original data"""
     data = mock_data_array
     model = EOF(n_modes=19, standardize=True, use_coslat=use_coslat)
     model.fit(data, dim=dim)
-    scores = model.data["scores"]
-    data_rec = model.inverse_transform(scores)
+    scores = model.scores(normalized=normalized)
+    data_rec = model.inverse_transform(scores, normalized=normalized)
     r2 = r2_score(data, data_rec, dim=dim)
     r2 = r2.mean()
     # Choose a threshold of 0.95; a bit arbitrary
@@ -709,13 +714,14 @@ def test_eof_inverse_transform(dim, use_coslat, mock_data_array):
         (("lon", "lat"), False),
     ],
 )
-def test_ceof_inverse_transform(dim, use_coslat, mock_data_array):
+@pytest.mark.parametrize("normalized", [True, False])
+def test_ceof_inverse_transform(dim, use_coslat, mock_data_array, normalized):
     """Inverse transform produces an approximate reconstruction of the original data"""
     data = mock_data_array
     model = ComplexEOF(n_modes=19, standardize=True, use_coslat=use_coslat)
     model.fit(data, dim=dim)
-    scores = model.data["scores"]
-    data_rec = model.inverse_transform(scores).real
+    scores = model.scores(normalized=normalized)
+    data_rec = model.inverse_transform(scores, normalized=normalized).real
     r2 = r2_score(data, data_rec, dim=dim)
     r2 = r2.mean()
     # Choose a threshold of 0.95; a bit arbitrary
@@ -734,15 +740,16 @@ def test_ceof_inverse_transform(dim, use_coslat, mock_data_array):
         (("lon", "lat"), False, 2),
     ],
 )
-def test_reof_inverse_transform(dim, use_coslat, power, mock_data_array):
+@pytest.mark.parametrize("normalized", [True, False])
+def test_reof_inverse_transform(dim, use_coslat, power, mock_data_array, normalized):
     """Inverse transform produces an approximate reconstruction of the original data"""
     data = mock_data_array
     model = EOF(n_modes=19, standardize=True, use_coslat=use_coslat)
     model.fit(data, dim=dim)
     rot = EOFRotator(n_modes=19, power=power)
     rot.fit(model)
-    scores = rot.data["scores"]
-    data_rec = rot.inverse_transform(scores).real
+    scores = rot.scores(normalized=normalized)
+    data_rec = rot.inverse_transform(scores, normalized=normalized).real
     r2 = r2_score(data, data_rec, dim=dim)
     r2 = r2.mean()
     # Choose a threshold of 0.95; a bit arbitrary
@@ -763,15 +770,16 @@ def test_reof_inverse_transform(dim, use_coslat, power, mock_data_array):
         (("lon", "lat"), False, 2),
     ],
 )
-def test_creof_inverse_transform(dim, use_coslat, power, mock_data_array):
+@pytest.mark.parametrize("normalized", [True, False])
+def test_creof_inverse_transform(dim, use_coslat, power, mock_data_array, normalized):
     """Inverse transform produces an approximate reconstruction of the original data"""
     data = mock_data_array
     model = ComplexEOF(n_modes=19, standardize=True, use_coslat=use_coslat)
     model.fit(data, dim=dim)
     rot = ComplexEOFRotator(n_modes=19, power=power)
     rot.fit(model)
-    scores = rot.data["scores"]
-    data_rec = rot.inverse_transform(scores).real
+    scores = rot.scores(normalized=normalized)
+    data_rec = rot.inverse_transform(scores, normalized=normalized).real
     r2 = r2_score(data, data_rec, dim=dim)
     r2 = r2.mean()
     # Choose a threshold of 0.95; a bit arbitrary

xeofs/models/_base_model.py

@@ -282,7 +282,9 @@ def fit_transform(
         """
         return self.fit(data, dim, weights).transform(data, **kwargs)
 
-    def inverse_transform(self, scores: DataObject) -> DataObject:
+    def inverse_transform(
+        self, scores: DataObject, normalized: bool = True
+    ) -> DataObject:
         """Reconstruct the original data from transformed data.
 
         Parameters
@@ -291,13 +293,17 @@ def inverse_transform(self, scores: DataObject) -> DataObject:
             Transformed data to be reconstructed. This could be a subset
             of the `scores` data of a fitted model, or unseen data. Must
             have a 'mode' dimension.
+        normalized: bool, default=True
+            Whether the scores data have been normalized by the L2 norm.
 
         Returns
         -------
         data: DataArray | Dataset | List[DataArray]
             Reconstructed data.
 
         """
+        if normalized:
+            scores = scores * self.data["norms"]
         data_reconstructed = self._inverse_transform_algorithm(scores)
         return self.preprocessor.inverse_transform_data(data_reconstructed)
 

xeofs/models/mca_rotator.py

@@ -7,7 +7,7 @@
 from .mca import MCA, ComplexMCA
 from ..preprocessing.preprocessor import Preprocessor
 from ..utils.rotation import promax
-from ..utils.data_types import DataArray
+from ..utils.data_types import DataArray, DataObject
 from ..utils.xarray_utils import argsort_dask, get_deterministic_sign_multiplier
 from ..data_container import DataContainer
 from .._version import __version__
@@ -319,7 +319,9 @@ def _sort_by_variance(self):
                     )
         self.sorted = True
 
-    def transform(self, **kwargs) -> DataArray | List[DataArray]:
+    def transform(
+        self, data1: DataObject | None = None, data2: DataObject | None = None
+    ) -> DataArray | List[DataArray]:
         """Project new "unseen" data onto the rotated singular vectors.
 
         Parameters
@@ -336,7 +338,7 @@ def transform(self, **kwargs) -> DataArray | List[DataArray]:
 
         """
         # raise error if no data is provided
-        if not kwargs:
+        if data1 is None and data2 is None:
             raise ValueError("No data provided. Please provide data1 and/or data2.")
 
         n_modes = self._params["n_modes"]
@@ -348,8 +350,7 @@ def transform(self, **kwargs) -> DataArray | List[DataArray]:
 
         results = []
 
-        if "data1" in kwargs.keys():
-            data1 = kwargs["data1"]
+        if data1 is not None:
             # Select the (non-rotated) singular vectors of the first dataset
             comps1 = self.model.data["components1"].sel(mode=slice(1, n_modes))
             norm1 = self.model.data["norm1"].sel(mode=slice(1, n_modes))
@@ -375,8 +376,7 @@ def transform(self, **kwargs) -> DataArray | List[DataArray]:
 
             results.append(projections1)
 
-        if "data2" in kwargs.keys():
-            data2 = kwargs["data2"]
+        if data2 is not None:
             # Select the (non-rotated) singular vectors of the second dataset
             comps2 = self.model.data["components2"].sel(mode=slice(1, n_modes))
             norm2 = self.model.data["norm2"].sel(mode=slice(1, n_modes))