Repository: xarray-contrib/xeofs

Issue:
Support for complex input data
**Is your feature request related to a problem? Please describe.**
Although complex data is accepted as input for `xeofs` models, not all models and methods (e.g. `inverse_transform`) work properly in this case.

**Describe the solution you'd like**
Complex input data can be decomposed into complex scores and components. Reconstructing data via `inverse_transform` results in the original complex data.

**Describe alternatives you've considered**
-NA-

**Additional context**
See use case in #149 .

**To Do**
- [x] `inverse_transform` does not enforce real output
- [x] restrict output to be real only for Complex EOF analysis
- [x] create data type checks for the other models as well
- [x] update documentation

Discussion:

Solution patch:
diff --git a/docs/overview_2_features.rst b/docs/overview_2_features.rst
index 5d17422..9c62684 100644
--- a/docs/overview_2_features.rst
+++ b/docs/overview_2_features.rst
@@ -141,6 +141,13 @@ the input type. For instance, executing PCA on a singular ``xr.DataArray`` will
   
   A mixed list containing both ``xr.DataArray`` and ``xr.Dataset`` objects is not currently supported.
 
+.. note::
+
+    Some methods like PCA/EOF analysis can also handle complex-valued input data. However, due to a limitation_ in the 
+    underlying ``dask`` library, complex-valued data is not supported when using ``dask``.
+
+.. _limitation: https://github.com/dask/dask/issues/7639
+
 Handling Missing Values
 -----------------------
 
diff --git a/xeofs/models/_base_cross_model.py b/xeofs/models/_base_cross_model.py
index 4989cdc..7ecb886 100644
--- a/xeofs/models/_base_cross_model.py
+++ b/xeofs/models/_base_cross_model.py
@@ -216,6 +216,38 @@ class _BaseCrossModel(ABC):
 
         return self._transform_algorithm(data1, data2)
 
+    def inverse_transform(
+        self, scores1: DataArray, scores2: DataArray
+    ) -> Tuple[DataObject, DataObject]:
+        """Reconstruct the original data from transformed data.
+
+        Parameters
+        ----------
+        scores1: DataObject
+            Transformed left field 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.
+        scores2: DataObject
+            Transformed right field 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.
+
+        Returns
+        -------
+        Xrec1: DataArray | Dataset | List[DataArray]
+            Reconstructed data of left field.
+        Xrec2: DataArray | Dataset | List[DataArray]
+            Reconstructed data of right field.
+
+        """
+        data1, data2 = self._inverse_transform_algorithm(scores1, scores2)
+
+        # Unstack and rescale the data
+        data1 = self.preprocessor1.inverse_transform_data(data1)
+        data2 = self.preprocessor2.inverse_transform_data(data2)
+
+        return data1, data2
+
     @abstractmethod
     def _fit_algorithm(self, data1: DataArray, data2: DataArray) -> Self:
         """
@@ -247,9 +279,32 @@ class _BaseCrossModel(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    def inverse_transform(
-        self, scores1: DataObject, scores2: DataObject
-    ) -> Tuple[DataObject, DataObject]:
+    def _inverse_transform_algorithm(
+        self, scores1: DataArray, scores2: DataArray
+    ) -> Tuple[DataArray, DataArray]:
+        """
+        Reconstruct the original data from transformed data. This method needs to be implemented in the respective
+        subclass.
+
+        Parameters
+        ----------
+        scores1: DataArray
+            Transformed left field 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.
+        scores2: DataArray
+            Transformed right field 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.
+
+        Returns
+        -------
+        Xrec1: DataArray
+            Reconstructed data of left field.
+        Xrec2: DataArray
+            Reconstructed data of right field.
+
+        """
         raise NotImplementedError
 
     def components(self) -> Tuple[DataObject, DataObject]:
diff --git a/xeofs/models/cca.py b/xeofs/models/cca.py
index 1c859d9..d50ec43 100644
--- a/xeofs/models/cca.py
+++ b/xeofs/models/cca.py
@@ -24,6 +24,7 @@ from xeofs.models import EOF
 from .._version import __version__
 from ..preprocessing.preprocessor import Preprocessor
 from ..utils.data_types import DataArray, DataList, DataObject
+from ..utils.sanity_checks import assert_not_complex
 
 
 def _check_parameter_number(parameter_name: str, parameter, n_views: int):
@@ -305,6 +306,9 @@ class CCA(CCABaseModel):
         self.eps = eps
 
     def _fit_algorithm(self, views: List[DataArray]) -> Self:
+        # Check input data
+        [assert_not_complex(view) for view in views]
+
         self.c = _process_parameter("c", self.c, 0, self.n_views_)
         eigvals, eigvecs = self._solve_gevp(views)
         self.eigvals = eigvals
diff --git a/xeofs/models/eof.py b/xeofs/models/eof.py
index 257b9a2..ba1cfe4 100644
--- a/xeofs/models/eof.py
+++ b/xeofs/models/eof.py
@@ -8,6 +8,7 @@ from .decomposer import Decomposer
 from ..utils.data_types import DataObject, DataArray, Dims
 from ..utils.hilbert_transform import hilbert_transform
 from ..utils.xarray_utils import total_variance as compute_total_variance
+from ..utils.sanity_checks import assert_not_complex
 
 
 class EOF(_BaseModel):
@@ -150,9 +151,6 @@ class EOF(_BaseModel):
         reconstructed_data = xr.dot(comps.conj(), scores, dims="mode")
         reconstructed_data.name = "reconstructed_data"
 
-        # Enforce real output
-        reconstructed_data = reconstructed_data.real
-
         return reconstructed_data
 
     def components(self) -> DataObject:
@@ -343,6 +341,8 @@ class ComplexEOF(EOF):
         self._params.update({"padding": padding, "decay_factor": decay_factor})
 
     def _fit_algorithm(self, data: DataArray) -> Self:
+        assert_not_complex(data)
+
         sample_name = self.sample_name
         feature_name = self.feature_name
 
@@ -387,6 +387,11 @@ class ComplexEOF(EOF):
     def _transform_algorithm(self, data: DataArray) -> DataArray:
         raise NotImplementedError("Complex EOF does not support transform method.")
 
+    def _inverse_transform_algorithm(self, scores: DataArray) -> DataArray:
+        Xrec = super()._inverse_transform_algorithm(scores)
+        # Enforce real output
+        return Xrec.real
+
     def components_amplitude(self) -> DataObject:
         """Return the amplitude of the (EOF) components.
 
diff --git a/xeofs/models/gwpca.py b/xeofs/models/gwpca.py
index 9f01407..4e123a9 100644
--- a/xeofs/models/gwpca.py
+++ b/xeofs/models/gwpca.py
@@ -6,7 +6,7 @@ from sklearn.utils.extmath import randomized_svd
 from xeofs.utils.data_types import DataArray
 from xeofs.utils.data_types import Data
 from ._base_model import _BaseModel
-from ..utils.sanity_checks import validate_input_type
+from ..utils.sanity_checks import assert_not_complex
 from ..utils.xarray_utils import convert_to_dim_type
 from ..utils.constants import (
     VALID_CARTESIAN_X_NAMES,
@@ -136,6 +136,9 @@ class GWPCA(_BaseModel):
         # Hide numba imports here to greatly speed up module import time
         from ..utils.numba_utils import _local_pcas
 
+        # Check input type
+        assert_not_complex(X)
+
         # Convert Dask arrays
         if not isinstance(X.data, np.ndarray):
             print(
diff --git a/xeofs/models/mca.py b/xeofs/models/mca.py
index bd155b9..3bb5922 100644
--- a/xeofs/models/mca.py
+++ b/xeofs/models/mca.py
@@ -12,6 +12,7 @@ from ..utils.statistics import pearson_correlation
 from ..utils.hilbert_transform import hilbert_transform
 from ..utils.dimension_renamer import DimensionRenamer
 from ..utils.xarray_utils import argsort_dask
+from ..utils.sanity_checks import assert_not_complex
 
 
 class MCA(_BaseCrossModel):
@@ -258,25 +259,27 @@ class MCA(_BaseCrossModel):
 
         return results
 
-    def inverse_transform(self, scores1: DataObject, scores2: DataObject):
+    def _inverse_transform_algorithm(
+        self, scores1: DataArray, scores2: DataArray
+    ) -> Tuple[DataArray, DataArray]:
         """Reconstruct the original data from transformed data.
 
         Parameters
         ----------
-        scores1: DataObject
+        scores1: DataArray
             Transformed left field 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.
-        scores2: DataObject
+        scores2: DataArray
             Transformed right field 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.
 
         Returns
         -------
-        Xrec1: DataArray | Dataset | List[DataArray]
+        Xrec1: DataArray
             Reconstructed data of left field.
-        Xrec2: DataArray | Dataset | List[DataArray]
+        Xrec2: DataArray
             Reconstructed data of right field.
 
         """
@@ -292,14 +295,6 @@ class MCA(_BaseCrossModel):
         data1 = xr.dot(scores1, comps1.conj() * norm1, dims="mode")
         data2 = xr.dot(scores2, comps2.conj() * norm2, dims="mode")
 
-        # Enforce real output
-        data1 = data1.real
-        data2 = data2.real
-
-        # Unstack and rescale the data
-        data1 = self.preprocessor1.inverse_transform_data(data1)
-        data2 = self.preprocessor2.inverse_transform_data(data2)
-
         return data1, data2
 
     def squared_covariance(self):
@@ -679,6 +674,9 @@ class ComplexMCA(MCA):
         self._params.update({"padding": padding, "decay_factor": decay_factor})
 
     def _fit_algorithm(self, data1: DataArray, data2: DataArray) -> Self:
+        assert_not_complex(data1)
+        assert_not_complex(data2)
+
         sample_name = self.sample_name
         feature_name = self.feature_name
 
@@ -908,6 +906,12 @@ class ComplexMCA(MCA):
     def transform(self, data1: DataObject, data2: DataObject):
         raise NotImplementedError("Complex MCA does not support transform method.")
 
+    def _inverse_transform_algorithm(self, scores1: DataArray, scores2: DataArray):
+        data1, data2 = super()._inverse_transform_algorithm(scores1, scores2)
+
+        # Enforce real output
+        return data1.real, data2.real
+
     def homogeneous_patterns(self, correction=None, alpha=0.05):
         raise NotImplementedError(
             "Complex MCA does not support homogeneous_patterns method."
diff --git a/xeofs/models/opa.py b/xeofs/models/opa.py
index f26924f..08a4b53 100644
--- a/xeofs/models/opa.py
+++ b/xeofs/models/opa.py
@@ -8,6 +8,7 @@ from ._base_model import _BaseModel
 from .eof import EOF
 from .decomposer import Decomposer
 from ..utils.data_types import DataObject, DataArray
+from ..utils.sanity_checks import assert_not_complex
 
 
 class OPA(_BaseModel):
@@ -127,6 +128,8 @@ class OPA(_BaseModel):
         )
 
     def _fit_algorithm(self, data: DataArray) -> Self:
+        assert_not_complex(data)
+
         sample_name = self.sample_name
         feature_name = self.feature_name
 
diff --git a/xeofs/utils/sanity_checks.py b/xeofs/utils/sanity_checks.py
index ddc5bb9..975622e 100644
--- a/xeofs/utils/sanity_checks.py
+++ b/xeofs/utils/sanity_checks.py
@@ -93,3 +93,10 @@ def validate_input_type(X) -> None:
             raise TypeError(err_msg)
     else:
         raise TypeError(err_msg)
+
+
+def assert_not_complex(da: xr.DataArray) -> None:
+    if da.dtype == "complex":
+        raise TypeError(
+            f"Invalid input type. This method does not support complex data types."
+        )


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 7927083..c2d6623 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -284,6 +284,31 @@ def mock_dask_data_array(mock_data_array):
     return mock_data_array.chunk({"lon": 2, "lat": 2, "time": -1})
 
 
+@pytest.fixture
+def mock_complex_data_array():
+    def f1(x, t):
+        return xr.DataArray(
+            1.0 / np.cosh(x[np.newaxis, :] + 3) * np.exp(2.3j * t[:, np.newaxis]),
+            coords=[("time", t), ("x", x)],
+        )
+
+    def f2(x, t):
+        return xr.DataArray(
+            2.0
+            / np.cosh(x[np.newaxis, :])
+            * np.tanh(x)
+            * np.exp(2.8j * t[:, np.newaxis]),
+            coords=[("time", t), ("x", x)],
+        )
+
+    x = np.linspace(-5, 5, 128)
+    t = np.linspace(0, 4 * np.pi, 256)
+
+    X1 = f1(x, t)
+    X2 = f2(x, t)
+    return X1 + X2
+
+
 # =============================================================================
 # Intermediate data
 # =============================================================================
diff --git a/tests/models/test_eof.py b/tests/models/test_eof.py
index c9053f4..b1d7e27 100644
--- a/tests/models/test_eof.py
+++ b/tests/models/test_eof.py
@@ -574,3 +574,19 @@ def test_serialize_deserialize_dataset(dim, mock_dataset):
     assert np.allclose(
         model.transform(mock_dataset), rebuilt_model.transform(mock_dataset)
     )
+
+
+def test_complex_input_inverse_transform(mock_complex_data_array):
+    """Test that the EOF model can handle complex input data."""
+
+    model = EOF(n_modes=128)
+    model.fit(mock_complex_data_array, dim="time")
+    scores = model.scores()
+
+    # Test the inverse_transform method
+    X_rec = model.inverse_transform(scores)
+    assert isinstance(X_rec, xr.DataArray)
+    # Check that the reconstructed data has the same dimensions as the original data
+    assert set(X_rec.dims) == set(mock_complex_data_array.dims)
+    # Reconstructed data should be close to the original data
+    xr.testing.assert_allclose(mock_complex_data_array, X_rec)