fix: support new datatree and numpy 2.0 (#169)

tests/models/test_complex_eof_rotator.py

@@ -47,8 +47,8 @@ def test_fit(ceof_model):
     assert hasattr(
         ceof_rotator, "data"
     ), 'The attribute "data" should be populated after fitting.'
-    assert type(ceof_rotator.model) == ComplexEOF
-    assert type(ceof_rotator.data) == DataContainer
+    assert isinstance(ceof_rotator.model, ComplexEOF)
+    assert isinstance(ceof_rotator.data, DataContainer)
 
 
 @pytest.mark.parametrize(

tests/models/test_decomposer.py

@@ -176,4 +176,4 @@ def test_random_state(
     U2 = decomposer.U_.data
 
     # Check that the results are the same
-    assert np.alltrue(U1 == U2)
+    assert np.all(U1 == U2)

tests/models/test_eof_rotator.py

@@ -49,8 +49,8 @@ def test_fit(eof_model):
     assert hasattr(
         eof_rotator, "data"
     ), 'The attribute "data" should be populated after fitting.'
-    assert type(eof_rotator.model) == EOF
-    assert type(eof_rotator.data) == DataContainer
+    assert isinstance(eof_rotator.model, EOF)
+    assert isinstance(eof_rotator.data, DataContainer)
 
 
 @pytest.mark.parametrize(

xeofs/data_container/data_container.py

@@ -3,7 +3,11 @@
 
 import dask
 from dask.diagnostics.progress import ProgressBar
-from datatree import DataTree
+
+try:
+    from xarray.core.datatree import DataTree
+except ImportError:
+    from datatree import DataTree
 
 from ..utils.data_types import DataArray
 

xeofs/models/_base_cross_model.py

@@ -5,9 +5,13 @@
 
 import dask
 import xarray as xr
-from datatree import DataTree
 from dask.diagnostics.progress import ProgressBar
 
+try:
+    from xarray.core.datatree import DataTree
+except ImportError:
+    from datatree import DataTree
+
 from .eof import EOF
 from ..preprocessing.preprocessor import Preprocessor
 from ..data_container import DataContainer

xeofs/models/_base_model.py

@@ -14,9 +14,13 @@
 
 import dask
 import xarray as xr
-from datatree import DataTree
 from dask.diagnostics.progress import ProgressBar
 
+try:
+    from xarray.core.datatree import DataTree
+except ImportError:
+    from datatree import DataTree
+
 from ..preprocessing.preprocessor import Preprocessor
 from ..data_container import DataContainer
 from ..utils.data_types import DataObject, Data, DataArray

xeofs/preprocessing/preprocessor.py

@@ -2,7 +2,12 @@
 from typing_extensions import Self
 
 import numpy as np
-from datatree import DataTree
+
+try:
+    from xarray.core.datatree import DataTree
+except ImportError:
+    from datatree import DataTree
+
 
 from .list_processor import GenericListTransformer
 from .dimension_renamer import DimensionRenamer
@@ -48,7 +53,7 @@ def extract_new_dim_names(X: List[DimensionRenamer]) -> Tuple[Dims, DimsList]:
     for x in X:
         new_sample_dims.append(x.sample_dims_after)
         new_feature_dims.append(x.feature_dims_after)
-    new_sample_dims: Dims = tuple(np.unique(np.asarray(new_sample_dims)))
+    new_sample_dims: Dims = tuple(np.unique(np.asarray(new_sample_dims)).tolist())
     return new_sample_dims, new_feature_dims
 
 

xeofs/preprocessing/transformer.py

@@ -5,9 +5,13 @@
 
 import pandas as pd
 import xarray as xr
-from datatree import DataTree
 from sklearn.base import BaseEstimator, TransformerMixin
 
+try:
+    from xarray.core.datatree import DataTree
+except ImportError:
+    from datatree import DataTree
+
 from ..utils.data_types import Dims, DataArray, DataSet, Data
 
 

xeofs/utils/io.py

@@ -3,7 +3,12 @@
 
 import numpy as np
 import xarray as xr
-from datatree import DataTree, open_datatree
+
+try:
+    from xarray.core.datatree import DataTree
+    from xarray.backends.api import open_datatree
+except ImportError:
+    from datatree import DataTree, open_datatree
 
 
 def write_model_tree(
@@ -20,9 +25,11 @@ def write_model_tree(
         raise ValueError(f"Unknown engine {engine}")
 
 
-def open_model_tree(path: str, engine: str = "zarr", chunks={}, **kwargs) -> DataTree:
+def open_model_tree(path: str, engine: str = "zarr", **kwargs) -> DataTree:
     """Open a DataTree from a file."""
-    dt = open_datatree(path, engine=engine, chunks=chunks, **kwargs)
+    if engine == "zarr" and "chunks" not in kwargs:
+        kwargs["chunks"] = {}
+    dt = open_datatree(path, engine=engine, **kwargs)
     if engine in ["netcdf4", "h5netcdf"]:
         dt = _desanitize_attrs_nc(dt)
     return dt