fix: handle indexing changes in new xarray versions (#159)

xarray-contrib · Mar 30, 2024 · 7d07d30 · 7d07d30
1 parent a3cb204
commit 7d07d30
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 14 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -14,7 +14,7 @@ python = ">=3.10,<3.12"
 poetry = "^1.6"
 numpy = "^1.24"
 pandas = "^2"
-xarray = "<=2023.12"
+xarray = ">=2023"
 scikit-learn = "^1.0.2"
 pooch = "^1.6.0"
 tqdm = "^4.64.0"

diff --git a/xeofs/models/_base_cross_model.py b/xeofs/models/_base_cross_model.py
@@ -248,6 +248,12 @@ def inverse_transform(
             Reconstructed data of right field.
 
         """
+        # Handle scalar mode in xr.dot
+        if "mode" not in scores1.dims:
+            scores1 = scores1.expand_dims("mode")
+        if "mode" not in scores2.dims:
+            scores2 = scores2.expand_dims("mode")
+
         data1, data2 = self._inverse_transform_algorithm(scores1, scores2)
 
         # Unstack and rescale the data

diff --git a/xeofs/models/_base_model.py b/xeofs/models/_base_model.py
@@ -305,6 +305,11 @@ def inverse_transform(
         if normalized:
             norms = self.data["norms"].sel(mode=scores.mode)
             scores = scores * norms
+
+        # Handle scalar mode in xr.dot
+        if "mode" not in scores.dims:
+            scores = scores.expand_dims("mode")
+
         data_reconstructed = self._inverse_transform_algorithm(scores)
 
         # Reconstructing the data using a single mode introduces a

diff --git a/xeofs/preprocessing/transformer.py b/xeofs/preprocessing/transformer.py
@@ -96,14 +96,19 @@ def _serialize_data(self, key: str, data: Data) -> DataSet:
             ds = data
         else:
             # Convert DataArray to Dataset
+            coords = {}
+            data_vars = {}
             if data.name in data.coords:
                 # Convert a coord-like DataArray to Dataset and note multiindexes
                 if isinstance(data.to_index(), pd.MultiIndex):
                     multiindexes[data.name] = [n for n in data.to_index().names]
-            # Make sure the DataArray has some name so we can create a string mapping
-            elif data.name is None:
-                data.name = key
-            ds = xr.Dataset(data_vars={data.name: data})
+                coords[data.name] = data
+            else:
+                # Make sure the DataArray has some name so we can create a string mapping
+                if data.name is None:
+                    data.name = key
+                data_vars[data.name] = data
+            ds = xr.Dataset(data_vars=data_vars, coords=coords)
             name_map = data.name
 
         # Drop multiindexes and record for later