Repository: MESMER-group/mesmer
Pull request #627: drop support for xarray-datatree
Author: login: mathause
Created: 2025-04-03T14:48:00Z | Merged: 2025-04-04T10:11:51Z

<!-- Feel free to remove check-list items aren't relevant to your change -->

 - [x] Closes #606
 - [x] Closed #106
 - [x] Tests added
 - [x] Fully documented, including `CHANGELOG.rst`

Drop support for `xarray-datatree` and require `xarray >= 2025.02`. WIP for now. I just saw something I want to address first in a separate PR.

Linked issue #606: porting `DataTree` to xarray (from datatree)
We wrote the `DataTree` layer with the retired datatree package. Porting this to the new `DataTree` interface in xarray is not-trivial because they added several breaking changes. I have started this endeavor but found several issues right away. I think the only at the moment is pydata/xarray#10013
Linked issue comments: user: login: veni-vidi-vici-dormivi
created_at: 2025-01-31T13:43:20Z
body: Thank you! Yes I tried it too and that was also where I got stuck. 
Other issues are:
- datatree cannot be built from `xr.DataArray` anymore, only from `xr.Dataset` (xarray issue #9539)
- and `map_over_subtree` does not handle empty nodes well (xarray issue #9693)
author_association: COLLABORATOR

Diff:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 2d6db15f7..19892a5c8 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,7 +6,7 @@ v0.11.0 - unreleased
 
 New Features
 ^^^^^^^^^^^^
-- Implemented new data structure using xarray-datatree, see `Data structure using xarray-datatree`_.
+- Implemented new data structure using ``xr.DataTree``, see `Data structure using DataTree`_.
 - Integrated MESMER-M into the code base, see `Integration of MESMER-M`_.
 - Added number of observations to the output of the AR process (`#395 <https://github.com/MESMER-group/mesmer/pull/395>`_).
   By `Victoria Bauer`_.
@@ -24,8 +24,9 @@ Breaking changes
 - The supported versions of some dependencies were changed
   (`#399 <https://github.com/MESMER-group/mesmer/pull/399>`_,
   `#405 <https://github.com/MESMER-group/mesmer/pull/405>`_,
-  `#503 <https://github.com/MESMER-group/mesmer/pull/503>`_, and
-  `#621 <https://github.com/MESMER-group/mesmer/pull/621>`_):
+  `#503 <https://github.com/MESMER-group/mesmer/pull/503>`_,
+  `#621 <https://github.com/MESMER-group/mesmer/pull/621>`_, and
+  `#627 <https://github.com/MESMER-group/mesmer/pull/627>`_):
 
   ================= ============= =========
   Package           Old           New
@@ -46,7 +47,7 @@ Breaking changes
   **scipy**         not specified 1.12
   **shapely**       not specified 2.0
   **statsmodels**   not specified 0.14
-  **xarray**        2023.04       2024.2
+  **xarray**        2023.04       2025.03
   ================= ============= =========
 
 Deprecations
@@ -78,25 +79,31 @@ Internal Changes
 - Use ruff instead of isort and flake8 to lint the code base (`#490 <https://github.com/MESMER-group/mesmer/pull/490>`_).
   By `Mathias Hauser`_.
 
-Data structure using xarray-datatree
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Data structure using DataTree
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-This release implements using :py:class:`DataTree` from `xarray-datatree` to handle multiple scenarios.
+This release uses :py:class:`xr.DataTree` as data structure to handle multiple scenarios.
+This was originally done with the prototype `xarray-datatree` package. After the port of
+``DataTree`` to xarray stabilized, mesmer briefly supported both ``DataTree`` versions
+(the one in xarray-datatree and in xarray) before dropping support for `xarray-datatree`.
 
-- Add `filefisher` as dependency to handle file paths of several scenarios (\
-  `#586 <https://github.com/MESMER-group/mesmer/pull/586>`_,
-  `#592 <https://github.com/MESMER-group/mesmer/pull/592>`_, and
-  `#629 <https://github.com/MESMER-group/mesmer/pull/629>`_).
 - Enable passing a :py:class:`DataTree` to the auto regression functions (`#570 <https://github.com/MESMER-group/mesmer/pull/570>`_).
-- Add weighting function for several scenarios (`#567 <https://github.com/MESMER-group/mesmer/pull/567>`_).
 - Enable passing :py:class:`DataTree` and :py:class:`xr.Dataset` to :py:class:`LinearRegression` (`#566 <https://github.com/MESMER-group/mesmer/pull/566>`_).
-- Add upper pin to `xarray` version to support `xarray-datatree`(`#559 <https://github.com/MESMER-group/mesmer/pull/559>`_).
+- Add weighting function for several scenarios (`#567 <https://github.com/MESMER-group/mesmer/pull/567>`_).
+- Add function to compute anomalies over several scenarios stored in a :py:class:`DataTree` (`#625 <https://github.com/MESMER-group/mesmer/pull/625>`_).
 - Add utility functions for :py:class:`DataTree` (`#556 <https://github.com/MESMER-group/mesmer/pull/556>`_).
-- Add `xarray-datatree` as dependency (`#554 <https://github.com/MESMER-group/mesmer/pull/554>`_)
 - Add calibration integration tests for multiple scenarios and change parameter files to netcdfs with new naming structure (`#537 <https://github.com/MESMER-group/mesmer/pull/537>`_)
 - Add new integration tests for drawing realisations (`#599 <https://github.com/MESMER-group/mesmer/pull/599>`_)
-- Port the functionality to xarray's :py:class:`DataTree` implementation (`#607 <https://github.com/MESMER-group/mesmer/pull/607>`_).
-- Add function to compute anomalies over several scenarios stored in a DataTree (`#625 <https://github.com/MESMER-group/mesmer/pull/625>`_).
+- PRs related to xarray and xarray-datatree:
+
+  - Add `xarray-datatree` as dependency (`#554 <https://github.com/MESMER-group/mesmer/pull/554>`_)
+  - Add upper pin to `xarray` version to support `xarray-datatree` (`#559 <https://github.com/MESMER-group/mesmer/pull/559>`_).
+  - Port the functionality to xarray's :py:class:`DataTree` implementation (`#607 <https://github.com/MESMER-group/mesmer/pull/607>`_).
+  - Drop support for `xarray-datatree`  (`#627 <https://github.com/MESMER-group/mesmer/pull/627>`_).
+- Add `filefisher` as dependency to handle file paths of several scenarios (\
+  `#586 <https://github.com/MESMER-group/mesmer/pull/586>`_,
+  `#592 <https://github.com/MESMER-group/mesmer/pull/592>`_, and
+  `#629 <https://github.com/MESMER-group/mesmer/pull/629>`_).
 
 By `Victoria Bauer`_ and `Mathias Hauser`_.
 
diff --git a/ci/min_deps_check.py b/ci/min_deps_check.py
index 0d319e65c..870b6b704 100644
--- a/ci/min_deps_check.py
+++ b/ci/min_deps_check.py
@@ -34,7 +34,8 @@
 
 POLICY_MONTHS = {"python": 30, "numpy": 18}
 POLICY_MONTHS_DEFAULT = 12
-POLICY_OVERRIDE: dict[str, tuple[int, int]] = {}
+# xarray v2025.03 is needed for xr.DataTree support
+POLICY_OVERRIDE: dict[str, tuple[int, int]] = {"xarray": (2025, 3)}
 errors = []
 
 
diff --git a/ci/requirements/min-all-deps.yml b/ci/requirements/min-all-deps.yml
index 201860b9a..03de0d652 100644
--- a/ci/requirements/min-all-deps.yml
+++ b/ci/requirements/min-all-deps.yml
@@ -22,8 +22,7 @@ dependencies:
  - scipy=1.12
  - shapely=2.0
  - statsmodels=0.14
- - xarray=2024.2
- - xarray-datatree=0.0.13
+ - xarray=2025.03
 # for testing
  - pytest
  - pytest-cov
diff --git a/docs/source/installation.rst b/docs/source/installation.rst
index 6a6a2862b..37064c0c4 100644
--- a/docs/source/installation.rst
+++ b/docs/source/installation.rst
@@ -19,7 +19,6 @@ Required dependencies
 - `scipy <https://scipy.org/>`__
 - `statsmodels <https://www.statsmodels.org/stable/index.html>`__
 - `xarray <http://xarray.pydata.org/>`__
-- `xarray-datatree <https://xarray-datatree.readthedocs.io/en/latest/index.html>`__
 
 Please note that we only support datatree version 0.0.13 and will switch to the xarray internal datatree module in the future.
 
diff --git a/environment.yml b/environment.yml
index 5facbd15a..f37bb6855 100644
--- a/environment.yml
+++ b/environment.yml
@@ -22,8 +22,7 @@ dependencies:
  - scikit-learn
  - scipy
  - statsmodels>=0.14
- - xarray>=2024.02, <2024.10 # upper pin for xarray-datatree
- - xarray-datatree==0.0.13
+ - xarray>=2025.03
 # for testing
  - black!=23
  - pytest
diff --git a/mesmer/core/_datatreecompat.py b/mesmer/core/_datatreecompat.py
index ba8b78731..5f9724a8c 100644
--- a/mesmer/core/_datatreecompat.py
+++ b/mesmer/core/_datatreecompat.py
@@ -3,44 +3,59 @@
 import xarray as xr
 from packaging.version import Version
 
-if Version(xr.__version__) < Version("2024.10"):
+if Version(xr.__version__) < Version("2025.03"):
+    raise ImportError(
+        f"xarray version {xr.__version__} not supported - please upgrade to v2025.03 ("
+        "or later)"
+    )
 
-    from datatree import DataTree, map_over_subtree, open_datatree
 
-    def map_over_datasets(func, *args, kwargs=None):
-        "compatibility layer for older xarray versions"
+def skip_empty_nodes(func):
+    @functools.wraps(func)
+    def _func(ds, *args, **kwargs):
+        if not ds:
+            return ds
+        return func(ds, *args, **kwargs)
 
-        if kwargs is None:
-            kwargs = {}
+    return _func
 
-        return map_over_subtree(func)(*args, **kwargs)
 
-elif Version(xr.__version__) > Version("2025.01"):
+def map_over_datasets(func, *args, kwargs=None):
+    """
+    Applies a function to every dataset in one or more DataTree objects with
+    the same structure (ie.., that are isomorphic), returning new trees which
+    store the results.
 
-    def skip_empty_nodes(func):
-        @functools.wraps(func)
-        def _func(ds, *args, **kwargs):
-            if not ds:
-                return ds
-            return func(ds, *args, **kwargs)
+    adapted version of xr.map_over_datasets which skips empty nodes
 
-        return _func
+    Parameters
+    ----------
+    func : callable
+        Function to apply to datasets with signature:
 
-    from xarray import DataTree, open_datatree
-    from xarray import map_over_datasets as _map_over_datasets
+        `func(*args: Dataset, **kwargs) -> Union[Dataset, tuple[Dataset, ...]]`.
 
-    def map_over_datasets(func, *args, kwargs=None):
+        (i.e. func must accept at least one Dataset and return at least one Dataset.)
+    *args : tuple, optional
+        Positional arguments passed on to `func`. Any DataTree arguments will be
+        converted to Dataset objects via `.dataset`.
+    kwargs : dict, optional
+        Optional keyword arguments passed directly to ``func``.
 
-        return _map_over_datasets(skip_empty_nodes(func), *args, kwargs=kwargs)
 
-else:
-    raise ImportError(
-        f"xarray version {xr.__version__} not supported - please upgrade to v2025.02 ("
-        "or later) or downgrade to v2024.09"
-    )
+    See Also
+    --------
+    xr.map_over_datasets
+
+    Notes
+    -----
+    For the discussion in xarray see https://github.com/pydata/xarray/issues/9693
+
+    """
+
+    return xr.map_over_datasets(skip_empty_nodes(func), *args, kwargs=kwargs)
+
 
 __all__ = [
-    "DataTree",
     "map_over_datasets",
-    "open_datatree",
 ]
diff --git a/mesmer/core/anomaly.py b/mesmer/core/anomaly.py
index adfe7f211..4aae62b42 100644
--- a/mesmer/core/anomaly.py
+++ b/mesmer/core/anomaly.py
@@ -1,11 +1,17 @@
 import operator
 
-from mesmer.core._datatreecompat import DataTree, map_over_datasets
+import xarray as xr
+
+from mesmer.core._datatreecompat import map_over_datasets
 
 
 def calc_anomaly(
-    dt: DataTree, reference_period: slice, *, time_dim="time", ref_scenario="historical"
-) -> DataTree:
+    dt: xr.DataTree,
+    reference_period: slice,
+    *,
+    time_dim="time",
+    ref_scenario="historical",
+) -> xr.DataTree:
     """subtract mean over the reference period
 
     Parameters
@@ -50,26 +56,22 @@ def calc_anomaly(
     # anomalies = dt - ref.ds
     anomalies = map_over_datasets(operator.sub, dt, ref.ds)
 
-    map_over_datasets(_assert_same_coords, dt, anomalies, ref_scenario)
+    _assert_same_coords(dt, anomalies, ref_scenario)
 
     return anomalies
 
 
 def _assert_same_coords(ref, anom, ref_scenario):
 
-    # TODO: use xr.group_subtrees?
-
-    # import xarray as xr
-
-    # for path, (l, r) in xr.group_subtrees(dt, dt_anom):
-    if not ref.coords.equals(anom.coords):
+    for path, (ref_scen, anom_scen) in xr.group_subtrees(ref, anom):
+        if not ref_scen.coords.equals(anom_scen.coords):
 
-        msg = (
-            f"Subtracting the reference changed the coordinates. "
-            f"Most likely because the ref_scenario ({ref_scenario}) is missing "
-            "some ensemble members.\n"
-        )
+            msg = (
+                f"Subtracting the reference changed the coordinates for '{path}'. "
+                f"Most likely because the ref_scenario ({ref_scenario}) is missing "
+                "some ensemble members.\n"
+            )
 
-        raise ValueError(msg)
+            raise ValueError(msg)
 
     return ref
diff --git a/mesmer/core/datatree.py b/mesmer/core/datatree.py
index e89666837..0fd41e6d5 100644
--- a/mesmer/core/datatree.py
+++ b/mesmer/core/datatree.py
@@ -2,10 +2,12 @@
 
 import xarray as xr
 
-from mesmer.core._datatreecompat import DataTree, map_over_datasets
+from mesmer.core._datatreecompat import map_over_datasets
 
 
-def _extract_single_dataarray_from_dt(dt: DataTree, name: str = "node") -> xr.DataArray:
+def _extract_single_dataarray_from_dt(
+    dt: xr.DataTree, name: str = "node"
+) -> xr.DataArray:
     """
     Extract a single DataArray from a DataTree node, holding a ``Dataset`` with one ``DataArray``.
     """
@@ -26,7 +28,7 @@ def _extract_single_dataarray_from_dt(dt: DataTree, name: str = "node") -> xr.Da
 
 
 def collapse_datatree_into_dataset(
-    dt: DataTree, dim: str, **concat_kwargs
+    dt: xr.DataTree, dim: str, **concat_kwargs
 ) -> xr.Dataset:
     """
     Take a ``DataTree`` and collapse **all its subtrees** into a single ``xr.Dataset`` along dim.
@@ -65,35 +67,35 @@ def collapse_datatree_into_dataset(
 
 @overload
 def stack_datatrees_for_linear_regression(
-    predictors: DataTree,
-    target: DataTree,
+    predictors: xr.DataTree,
+    target: xr.DataTree,
     weights: None = None,
     *,
     stacking_dims: list[str],
     collapse_dim: str = "scenario",
     stacked_dim: str = "sample",
-) -> tuple[DataTree, xr.Dataset, None]: ...
+) -> tuple[xr.DataTree, xr.Dataset, None]: ...
 @overload
 def stack_datatrees_for_linear_regression(
-    predictors: DataTree,
-    target: DataTree,
-    weights: DataTree,
+    predictors: xr.DataTree,
+    target: xr.DataTree,
+    weights: xr.DataTree,
     *,
     stacking_dims: list[str],
     collapse_dim: str = "scenario",
     stacked_dim: str = "sample",
-) -> tuple[DataTree, xr.Dataset, xr.Dataset]: ...
+) -> tuple[xr.DataTree, xr.Dataset, xr.Dataset]: ...
 
 
 def stack_datatrees_for_linear_regression(
-    predictors: DataTree,
-    target: DataTree,
-    weights: DataTree | None = None,
+    predictors: xr.DataTree,
+    target: xr.DataTree,
+    weights: xr.DataTree | None = None,
     *,
     stacking_dims: list[str],
     collapse_dim: str = "scenario",
     stacked_dim: str = "sample",
-) -> tuple[DataTree, xr.Dataset, xr.Dataset | None]:
+) -> tuple[xr.DataTree, xr.Dataset, xr.Dataset | None]:
     """
     prepares data for Linear Regression:
     1. Broadcasts predictors to target
@@ -155,7 +157,7 @@ def stack_datatrees_for_linear_regression(
     exclude_dim = set(target.leaves[0].ds.dims) - set(stacking_dims)
 
     # prepare predictors
-    predictors_stacked = DataTree()
+    predictors_stacked = xr.DataTree()
     for key, pred in predictors.items():
         # 1) broadcast to target
         # TODO: use DataTree method again, once available
@@ -167,7 +169,7 @@ def stack_datatrees_for_linear_regression(
         # 2) collapse into DataSets
         predictor_ds = collapse_datatree_into_dataset(pred_broadcast, dim=collapse_dim)
         # 3) stack
-        predictors_stacked[key] = DataTree(
+        predictors_stacked[key] = xr.DataTree(
             predictor_ds.stack(stack_dim, create_index=False)
         )
     # TODO: use DataTree method again, once available
diff --git a/mesmer/core/weighted.py b/mesmer/core/weighted.py
index 3a7a909bb..36cd0074f 100644
--- a/mesmer/core/weighted.py
+++ b/mesmer/core/weighted.py
@@ -3,7 +3,7 @@
 import numpy as np
 import xarray as xr
 
-from mesmer.core._datatreecompat import DataTree, map_over_datasets
+from mesmer.core._datatreecompat import map_over_datasets
 
 
 def _weighted_if_dim(obj, weights, dims):
@@ -111,8 +111,8 @@ def global_mean(data, weights=None, x_dim="lon", y_dim="lat"):
 
 
 def equal_scenario_weights_from_datatree(
-    dt: DataTree, ens_dim: str = "member", time_dim: str = "time"
-) -> DataTree:
+    dt: xr.DataTree, ens_dim: str = "member", time_dim: str = "time"
+) -> xr.DataTree:
     """
     Create a DataTree isomorphic to ``dt``, holding the weights for each scenario to weight the ensemble members of each
     scenario such that each scenario contributes equally to some fitting procedure.
diff --git a/mesmer/stats/_auto_regression.py b/mesmer/stats/_auto_regression.py
index f4c12cf24..76c58e963 100644
--- a/mesmer/stats/_auto_regression.py
+++ b/mesmer/stats/_auto_regression.py
@@ -7,7 +7,7 @@
 import scipy
 import xarray as xr
 
-from mesmer.core._datatreecompat import DataTree, map_over_datasets
+from mesmer.core._datatreecompat import map_over_datasets
 from mesmer.core.datatree import (
     collapse_datatree_into_dataset,
 )
@@ -18,10 +18,10 @@
 )
 
 
-def _scen_ens_inputs_to_dt(objs: Sequence) -> DataTree:
+def _scen_ens_inputs_to_dt(objs: Sequence) -> xr.DataTree:
     """Helper function to convert a sequence of objects to a DataTree"""
 
-    if isinstance(objs[0], DataTree):
+    if isinstance(objs[0], xr.DataTree):
         if len(objs) != 1:
             raise ValueError("Only one DataTree can be passed.")
         dt = objs[0]
@@ -31,7 +31,7 @@ def _scen_ens_inputs_to_dt(objs: Sequence) -> DataTree:
         # see https://github.com/pydata/xarray/issues/9486
         # with just da_dict = {f"da_{i}": da for i, da in enumerate(objs)}
         ds_dict = {f"scen_{i}": da._to_temp_dataset() for i, da in enumerate(objs)}
-        dt = DataTree.from_dict(ds_dict)
+        dt = xr.DataTree.from_dict(ds_dict)
 
     elif isinstance(objs[0], dict):
         if len(objs) != 1:
@@ -40,7 +40,7 @@ def _scen_ens_inputs_to_dt(objs: Sequence) -> DataTree:
         # TODO: in the future might be able to use DataTree.from_array_dict(da_dict)
         # see https://github.com/pydata/xarray/issues/9486
         ds_dict = {f"{key}": da._to_temp_dataset() for key, da in da_dict.items()}
-        dt = DataTree.from_dict(ds_dict)
+        dt = xr.DataTree.from_dict(ds_dict)
 
     else:
         raise ValueError(
@@ -67,7 +67,7 @@ def inner(ds: xr.Dataset, **kwargs) -> xr.Dataset:
 
 
 def select_ar_order_scen_ens(
-    *objs: xr.DataArray | dict[str, xr.DataArray] | DataTree,
+    *objs: xr.DataArray | dict[str, xr.DataArray] | xr.DataTree,
     dim: str,
     ens_dim: str | None,
     maxlag: int,
@@ -106,7 +106,7 @@ def select_ar_order_scen_ens(
 
 
 def _select_ar_order_scen_ens_dt(
-    dt: DataTree,
+    dt: xr.DataTree,
     dim: str,
     ens_dim: str | None,
     maxlag: int,
@@ -172,7 +172,7 @@ def _ens_quantile(ds, ens_dim):
 
 
 def fit_auto_regression_scen_ens(
-    *objs: xr.DataArray | dict[str, xr.DataArray] | DataTree,
+    *objs: xr.DataArray | dict[str, xr.DataArray] | xr.DataTree,
     dim: str,
     ens_dim: str | None,
     lags: int | xr.DataArray,
@@ -212,7 +212,7 @@ def fit_auto_regression_scen_ens(
 
 
 def _fit_auto_regression_scen_ens_dt(
-    dt: DataTree, dim: str, ens_dim: str | None, lags: int | xr.DataArray
+    dt: xr.DataTree, dim: str, ens_dim: str | None, lags: int | xr.DataArray
 ) -> xr.Dataset:
     """
     fit an auto regression and potentially calculate the mean over ensemble members
@@ -424,7 +424,7 @@ def draw_auto_regression_uncorrelated(
 
     """
 
-    if isinstance(seed, DataTree):
+    if isinstance(seed, xr.DataTree):
         return map_over_datasets(
             _draw_auto_regression_uncorrelated,
             seed,
@@ -561,7 +561,7 @@ def draw_auto_regression_correlated(
 
     """
 
-    if isinstance(seed, DataTree):
+    if isinstance(seed, xr.DataTree):
 
         return map_over_datasets(
             _draw_auto_regression_correlated,
@@ -1034,7 +1034,7 @@ def draw_auto_regression_monthly(
 
     """
 
-    if isinstance(seed, DataTree):
+    if isinstance(seed, xr.DataTree):
 
         return map_over_datasets(
             _draw_auto_regression_monthly,
diff --git a/mesmer/stats/_linear_regression.py b/mesmer/stats/_linear_regression.py
index 0858019dc..d783e4642 100644
--- a/mesmer/stats/_linear_regression.py
+++ b/mesmer/stats/_linear_regression.py
@@ -1,7 +1,7 @@
 import numpy as np
 import xarray as xr
 
-from mesmer.core._datatreecompat import DataTree, map_over_datasets
+from mesmer.core._datatreecompat import map_over_datasets
 from mesmer.core.datatree import (
     _extract_single_dataarray_from_dt,
     collapse_datatree_into_dataset,
@@ -21,7 +21,7 @@ def __init__(self):
 
     def fit(
         self,
-        predictors: dict[str, xr.DataArray] | DataTree | xr.Dataset,
+        predictors: dict[str, xr.DataArray] | xr.DataTree | xr.Dataset,
         target: xr.DataArray,
         dim: str,
         weights: xr.DataArray | None = None,
@@ -60,7 +60,7 @@ def fit(
 
     def predict(
         self,
-        predictors: dict[str, xr.DataArray] | DataTree | xr.Dataset,
+        predictors: dict[str, xr.DataArray] | xr.DataTree | xr.Dataset,
         exclude: str | set[str] | None = None,
     ) -> xr.DataArray:
         """
@@ -111,7 +111,7 @@ def predict(
 
         # if predictors is a DataTree, rename all data variables to "pred" to avoid conflicts
         # not necessaey if predictors is empty DataTree or only data is in root, i.e. depth == 0
-        if isinstance(predictors, DataTree) and not predictors.depth == 0:
+        if isinstance(predictors, xr.DataTree) and not predictors.depth == 0:
             predictors = map_over_datasets(
                 lambda ds: ds.rename({var: "pred" for var in ds.data_vars}), predictors
             )
@@ -121,21 +121,21 @@ def predict(
             # TODO: fix once .transpose() is possible for DataTree
             signal = predictors[key] * params[key]
 
-            if isinstance(signal, DataTree):
+            if isinstance(signal, xr.DataTree):
                 signal = map_over_datasets(xr.Dataset.transpose, signal)
             else:
                 signal = signal.transpose()
 
             prediction = signal + prediction
 
-        if isinstance(prediction, DataTree):
+        if isinstance(prediction, xr.DataTree):
             prediction = _extract_single_dataarray_from_dt(prediction)
 
         return prediction.rename("prediction")
 
     def residuals(
         self,
-        predictors: dict[str, xr.DataArray] | DataTree | xr.Dataset,
+        predictors: dict[str, xr.DataArray] | xr.DataTree | xr.Dataset,
         target: xr.DataArray,
     ) -> xr.DataArray:
         """
@@ -222,7 +222,7 @@ def to_netcdf(self, filename: str, **kwargs):
 
 
 def _fit_linear_regression_xr(
-    predictors: dict[str, xr.DataArray] | DataTree | xr.Dataset,
+    predictors: dict[str, xr.DataArray] | xr.DataTree | xr.Dataset,
     target: xr.DataArray,
     dim: str,
     weights: xr.DataArray | None = None,
@@ -254,10 +254,10 @@ def _fit_linear_regression_xr(
         individual DataArray.
     """
     # if DataTree only has data in root, extract Dataset
-    if isinstance(predictors, DataTree) and predictors.depth == 0:
+    if isinstance(predictors, xr.DataTree) and predictors.depth == 0:
         predictors = predictors.to_dataset()
 
-    if not isinstance(predictors, dict | DataTree | xr.Dataset):
+    if not isinstance(predictors, dict | xr.DataTree | xr.Dataset):
         raise TypeError(
             f"predictors should be a dict, DataTree or xr.Dataset, got {type(predictors)}."
         )
@@ -271,7 +271,7 @@ def _fit_linear_regression_xr(
         raise ValueError("dim cannot currently be 'predictor'.")
 
     for key, pred in predictors.items():
-        if isinstance(pred, DataTree):
+        if isinstance(pred, xr.DataTree):
             pred = _extract_single_dataarray_from_dt(pred, name=f"predictor: {key}")
 
         _check_dataarray_form(pred, ndim=1, required_dims=dim, name=f"predictor: {key}")
@@ -286,9 +286,9 @@ def _fit_linear_regression_xr(
         predictors_concat = predictors_concat.assign_coords(
             {"predictor": list(predictors.keys())}
         )
-    elif isinstance(predictors, DataTree):
+    elif isinstance(predictors, xr.DataTree):
         # rename all data variables to "pred" to avoid conflicts when concatenating
-        def _rename_vars(ds) -> DataTree:
+        def _rename_vars(ds) -> xr.DataTree:
             (var,) = ds.data_vars
             return ds.rename({var: "pred"})
 
diff --git a/setup.cfg b/setup.cfg
index eea12d39e..4d9a6100e 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -46,8 +46,7 @@ install_requires =
     scikit-learn >=1.4 # only for the tests
     scipy >=1.12
     statsmodels >=0.14
-    xarray >=2024.2, <2024.10 # upper pin for xarray-datatree
-    xarray-datatree ==0.0.13
+    xarray >=2025.3 # for xr.DataTree
 
 [options.extras_require]
 complete =
diff --git a/tests/integration/test_calibrate_mesmer_newcodepath.py b/tests/integration/test_calibrate_mesmer_newcodepath.py
index 9733c6e26..069c2e0f7 100644
--- a/tests/integration/test_calibrate_mesmer_newcodepath.py
+++ b/tests/integration/test_calibrate_mesmer_newcodepath.py
@@ -6,7 +6,7 @@
 
 import mesmer
 import mesmer.core.datatree
-from mesmer.core._datatreecompat import DataTree, map_over_datasets
+from mesmer.core._datatreecompat import map_over_datasets
 
 
 @pytest.mark.filterwarnings("ignore:No local minimum found")
@@ -133,7 +133,7 @@ def test_calibrate_mesmer(
     scenarios_incl_hist.append("historical")
 
     # load data for each scenario
-    dt = DataTree()
+    dt = xr.DataTree()
     for scen in scenarios_incl_hist:
         files = fc_all.search(scenario=scen)
 
@@ -150,7 +150,7 @@ def test_calibrate_mesmer(
         # create a Dataset that holds each member along the member dimension
         scen_data = xr.concat(members, dim="member")
         # put the scenario dataset into the DataTree
-        dt[scen] = DataTree(scen_data)
+        dt[scen] = xr.DataTree(scen_data)
 
     # load additional data
     if use_hfds:
@@ -163,7 +163,7 @@ def test_calibrate_mesmer(
             member=unique_scen_members,
         )
 
-        dt_hfds = DataTree()
+        dt_hfds = xr.DataTree()
         for scen in scenarios_incl_hist:
             files = fc_hfds.search(scenario=scen)
 
@@ -177,7 +177,7 @@ def test_calibrate_mesmer(
                 members.append(ds)
 
             scen_data = xr.concat(members, dim="member")
-            dt_hfds[scen] = DataTree(scen_data)
+            dt_hfds[scen] = xr.DataTree(scen_data)
     else:
         dt_hfds = None
 
@@ -255,7 +255,7 @@ def mask_and_stack(ds, threshold_land):
     # broadcast so all datasets have all the dimensions
     # gridcell can be excluded because it will be mapped in the Linear Regression
     target = tas_stacked
-    predictors = DataTree.from_dict(
+    predictors = xr.DataTree.from_dict(
         {"tas": tas_globmean_smoothed, "tas_resids": tas_resid_novolc}
     )
     if use_tas2:
@@ -292,9 +292,9 @@ def mask_and_stack(ds, threshold_land):
         sample=("time", "member", "scenario")
     ).unstack("sample")
 
-    dt_resids = DataTree()
+    dt_resids = xr.DataTree()
     for scenario in tas_un_stacked_residuals.scenario.values:
-        dt_resids[scenario] = DataTree(
+        dt_resids[scenario] = xr.DataTree(
             tas_un_stacked_residuals.sel(scenario=scenario)
             .dropna("member", how="all")
             .dropna("time")
diff --git a/tests/integration/test_draw_realisations_mesmer_newcodepath.py b/tests/integration/test_draw_realisations_mesmer_newcodepath.py
index 6c3c71853..bede213c6 100644
--- a/tests/integration/test_draw_realisations_mesmer_newcodepath.py
+++ b/tests/integration/test_draw_realisations_mesmer_newcodepath.py
@@ -5,7 +5,7 @@
 from filefisher import FileFinder
 
 import mesmer
-from mesmer.core._datatreecompat import DataTree, map_over_datasets, open_datatree
+from mesmer.core._datatreecompat import map_over_datasets
 
 
 def create_forcing_data(test_data_root_dir, scenarios, use_hfds, use_tas2):
@@ -62,7 +62,7 @@ def load_hist(meta, fc_hist):
         return xr.open_dataset(fN, use_cftime=True)
 
     def load_hist_scen_continuous(fc_hist, fc_scens):
-        dt = DataTree()
+        dt = xr.DataTree()
         for scen in fc_scens.df.scenario.unique():
             files = fc_scens.search(scenario=scen)
 
@@ -99,7 +99,7 @@ def load_hist_scen_continuous(fc_hist, fc_scens):
             # create a Dataset that holds each member along the member dimension
             scen_data = xr.concat(members, dim="member")
             # put the scenario dataset into the DataTree
-            dt[scen] = DataTree(scen_data)
+            dt[scen] = xr.DataTree(scen_data)
         return dt
 
     tas = load_hist_scen_continuous(fc_hist, fc_scens)
@@ -265,14 +265,14 @@ def test_make_realisations(
     # we need a maximum of 4 seeds if there are max 2 scenarios (1 for global and 1 for local)
     seed_list = [981, 314, 272, 42]
 
-    seed_global_variability = DataTree()
-    seed_local_variability = DataTree()
+    seed_global_variability = xr.DataTree()
+    seed_local_variability = xr.DataTree()
 
     for scen in scenarios:
-        seed_global_variability[scen] = DataTree(
+        seed_global_variability[scen] = xr.DataTree(
             xr.Dataset({"seed": xr.DataArray(seed_list.pop())})
         )
-        seed_local_variability[scen] = DataTree(
+        seed_local_variability[scen] = xr.DataTree(
             xr.Dataset({"seed": xr.DataArray(seed_list.pop())})
         )
 
@@ -311,9 +311,9 @@ def test_make_realisations(
     # 3.) compute the local forced response
     lr = mesmer.stats.LinearRegression().from_netcdf(local_forced_file)
 
-    predictors = DataTree()
+    predictors = xr.DataTree()
     for scen in scenarios:
-        predictors[scen] = DataTree.from_dict(
+        predictors[scen] = xr.DataTree.from_dict(
             {"tas": tas_forcing[scen], "tas_resids": global_variability[scen]}
         )
 
@@ -323,11 +323,11 @@ def test_make_realisations(
             predictors[scen]["hfds"] = hfds[scen]
 
     # uses ``exclude`` to split the linear response
-    local_forced_response = DataTree()
-    local_variability_from_global_var = DataTree()
+    local_forced_response = xr.DataTree()
+    local_variability_from_global_var = xr.DataTree()
 
     for scen in predictors.children:
-        local_forced_response[scen] = DataTree(
+        local_forced_response[scen] = xr.DataTree(
             lr.predict(predictors[scen], exclude={"tas_resids"})
             .rename("tas")
             .to_dataset()
@@ -340,7 +340,7 @@ def test_make_realisations(
         if use_hfds:
             exclude.add("hfds")
 
-        local_variability_from_global_var[scen] = DataTree(
+        local_variability_from_global_var[scen] = xr.DataTree(
             lr.predict(predictors[scen], exclude=exclude).rename("tas").to_dataset()
         )
 
@@ -371,7 +371,7 @@ def test_make_realisations(
         result.to_netcdf(expected_output_file)
         pytest.skip("Updated expected output file.")
     else:
-        expected = open_datatree(expected_output_file, use_cftime=True)
+        expected = xr.open_datatree(expected_output_file, use_cftime=True)
         for scen in scenarios:
             exp_scen = expected[scen].to_dataset()
             res_scen = result[scen].to_dataset()
diff --git a/tests/integration/test_draw_realisations_mesmer_x.py b/tests/integration/test_draw_realisations_mesmer_x.py
index e9a61e5b8..e2e291015 100644
--- a/tests/integration/test_draw_realisations_mesmer_x.py
+++ b/tests/integration/test_draw_realisations_mesmer_x.py
@@ -3,7 +3,7 @@
 import pytest
 import xarray as xr
 
-# from mesmer.core._datatreecompat import Datatree, map_over_datasets
+# from mesmer.core._datatreecompat import map_over_datasets
 import mesmer
 import mesmer.mesmer_x
 
diff --git a/tests/unit/test_anomaly.py b/tests/unit/test_anomaly.py
index 35d3b1827..d09a066a8 100644
--- a/tests/unit/test_anomaly.py
+++ b/tests/unit/test_anomaly.py
@@ -5,7 +5,7 @@
 import xarray as xr
 
 import mesmer
-from mesmer.core._datatreecompat import DataTree, map_over_datasets
+from mesmer.core._datatreecompat import map_over_datasets
 
 
 @pytest.fixture
@@ -37,7 +37,7 @@ def example_tas():
     )
     proj = xr.Dataset(data_vars={"tas": dta})
 
-    tas = DataTree.from_dict({"historical": hist, "proj": proj})
+    tas = xr.DataTree.from_dict({"historical": hist, "proj": proj})
 
     return tas
 
@@ -65,7 +65,7 @@ def test_calc_anomaly_errors(example_tas):
     # example_tas["proj1"] = example_tas["proj"].dataset.assign_coords(ens=["0", "3"])
 
     ds = example_tas["proj"].to_dataset().assign_coords(ens=["0", "3"])
-    example_tas["proj1"] = DataTree(ds)
+    example_tas["proj1"] = xr.DataTree(ds)
 
     with pytest.raises(
         ValueError, match="Subtracting the reference changed the coordinates."
diff --git a/tests/unit/test_auto_regression.py b/tests/unit/test_auto_regression.py
index d01da7228..e51140fd1 100644
--- a/tests/unit/test_auto_regression.py
+++ b/tests/unit/test_auto_regression.py
@@ -6,7 +6,6 @@
 import xarray as xr
 
 import mesmer
-from mesmer.core._datatreecompat import DataTree
 from mesmer.core.utils import LinAlgWarning, _check_dataarray_form, _check_dataset_form
 from mesmer.stats import _auto_regression
 from mesmer.testing import trend_data_1D, trend_data_2D, trend_data_3D
@@ -175,7 +174,7 @@ def test_draw_auto_regression_uncorrelated(
 
 
 def test_draw_auto_regression_uncorrelated_dt(ar_params_1D):
-    seeds = DataTree.from_dict(
+    seeds = xr.DataTree.from_dict(
         {
             "scen1": xr.DataArray(np.array([25]), name="seed").to_dataset(),
             "scen2": xr.DataArray(np.array([42]), name="seed").to_dataset(),
@@ -331,7 +330,7 @@ def test_draw_auto_regression_correlated(
 
 
 def test_draw_auto_regression_correlated_dt(ar_params_2D, covariance):
-    seeds = DataTree.from_dict(
+    seeds = xr.DataTree.from_dict(
         {
             "scen1": xr.DataArray(np.array([25]), name="seed").to_dataset(),
             "scen2": xr.DataArray(np.array([42]), name="seed").to_dataset(),
@@ -887,7 +886,7 @@ def test_draw_auto_regression_monthly(seed):
 
 def test_draw_auto_regression_monthly_dt():
 
-    seeds = DataTree.from_dict(
+    seeds = xr.DataTree.from_dict(
         {
             "scen1": xr.DataArray(np.array([25]), name="seed").to_dataset(),
             "scen2": xr.DataArray(np.array([42]), name="seed").to_dataset(),
diff --git a/tests/unit/test_auto_regression_scen_ens.py b/tests/unit/test_auto_regression_scen_ens.py
index f2c31d66b..64a42435a 100644
--- a/tests/unit/test_auto_regression_scen_ens.py
+++ b/tests/unit/test_auto_regression_scen_ens.py
@@ -7,7 +7,6 @@
 
 import mesmer
 import mesmer.stats._auto_regression
-from mesmer.core._datatreecompat import DataTree
 
 
 def generate_ar_samples(ar, std=1, n_timesteps=100, n_ens=4):
@@ -35,7 +34,7 @@ def _prepare_data(*dataarrays, data_format):
         data = {
             f"scen{i}": xr.Dataset({"tas": da}) for i, da in enumerate(dataarrays, 1)
         }
-        return DataTree.from_dict(data)  # type: ignore[attr-defined] not typed in datatree
+        return xr.DataTree.from_dict(data)  # type: ignore[attr-defined] not typed in datatree
     else:
         raise ValueError(f"Unknown format_type: {data_format}")
 
@@ -187,7 +186,7 @@ def test_fit_auto_regression_scen_ens_errors():
 
     # data tree
     data = {"scen1": xr.Dataset({"tas": da})}
-    dt1 = DataTree.from_dict(data)  # type: ignore[attr-defined] not typed in datatree
+    dt1 = xr.DataTree.from_dict(data)  # type: ignore[attr-defined] not typed in datatree
     dt2 = dt1.copy()
 
     with pytest.raises(ValueError, match="Only one DataTree can be passed."):
@@ -196,7 +195,7 @@ def test_fit_auto_regression_scen_ens_errors():
         )
 
     data = {"scen1": xr.Dataset({"tas": da, "tas2": da})}
-    dt = DataTree.from_dict(data)  # type: ignore[attr-defined]
+    dt = xr.DataTree.from_dict(data)  # type: ignore[attr-defined]
 
     with pytest.raises(ValueError, match="Dataset must have only one data variable."):
         mesmer.stats.fit_auto_regression_scen_ens(dt, dim="time", ens_dim="ens", lags=3)
diff --git a/tests/unit/test_datatree.py b/tests/unit/test_datatree.py
index 8d3842d38..eaa48b12f 100644
--- a/tests/unit/test_datatree.py
+++ b/tests/unit/test_datatree.py
@@ -3,7 +3,7 @@
 import xarray as xr
 
 import mesmer
-from mesmer.core._datatreecompat import DataTree, map_over_datasets
+from mesmer.core._datatreecompat import map_over_datasets
 from mesmer.core.utils import _check_dataarray_form
 from mesmer.testing import trend_data_1D, trend_data_2D
 
@@ -19,7 +19,7 @@ def test_collapse_datatree_into_dataset():
     dim = xr.Variable("member", np.arange(2))
     leaf2 = xr.concat([ds1, ds2], dim=dim)
 
-    dt = DataTree.from_dict({"scen1": leaf1, "scen2": leaf2})
+    dt = xr.DataTree.from_dict({"scen1": leaf1, "scen2": leaf2})
 
     collapse_dim = "scenario"
     res = mesmer.datatree.collapse_datatree_into_dataset(dt, dim=collapse_dim)
@@ -32,7 +32,7 @@ def test_collapse_datatree_into_dataset():
 
     # error if data set has no coords along dim (bc then it is not concatenable if lengths differ)
     leaf_missing_coords = leaf1.drop_vars("member")
-    dt = DataTree.from_dict({"scen1": leaf_missing_coords, "scen2": leaf2})
+    dt = xr.DataTree.from_dict({"scen1": leaf_missing_coords, "scen2": leaf2})
     with pytest.raises(
         ValueError, match="cannot reindex or align along dimension 'member'"
     ):
@@ -41,7 +41,7 @@ def test_collapse_datatree_into_dataset():
     # Dimension along which to concatenate already exists
     leaf1_scen = leaf1.assign_coords({"scenario": "scen1"}).expand_dims(collapse_dim)
     leaf2_scen = leaf2.assign_coords({"scenario": "scen2"}).expand_dims(collapse_dim)
-    dt = DataTree.from_dict({"scen1": leaf1_scen, "scen2": leaf2_scen})
+    dt = xr.DataTree.from_dict({"scen1": leaf1_scen, "scen2": leaf2_scen})
 
     res = mesmer.datatree.collapse_datatree_into_dataset(dt, dim=collapse_dim)
     assert isinstance(res, xr.Dataset)
@@ -50,7 +50,7 @@ def test_collapse_datatree_into_dataset():
     xr.testing.assert_equal(scen1.drop_vars("scenario"), leaf1)
 
     # only one leaf works
-    dt = DataTree.from_dict({"scen1": leaf1})
+    dt = xr.DataTree.from_dict({"scen1": leaf1})
     res = mesmer.datatree.collapse_datatree_into_dataset(dt, dim=collapse_dim)
 
     assert isinstance(res, xr.Dataset)
@@ -61,7 +61,7 @@ def test_collapse_datatree_into_dataset():
     xr.testing.assert_equal(scen1.drop_vars(collapse_dim), leaf1)
 
     # test data in root works
-    dt = DataTree(leaf1, name="scen1")
+    dt = xr.DataTree(leaf1, name="scen1")
     res = mesmer.datatree.collapse_datatree_into_dataset(dt, dim=collapse_dim)
 
     assert isinstance(res, xr.Dataset)
@@ -72,10 +72,10 @@ def test_collapse_datatree_into_dataset():
     xr.testing.assert_equal(scen1.drop_vars(collapse_dim), leaf1)
 
     # nested DataTree works
-    dt = DataTree()
-    dt["scen1/sub_scen1"] = DataTree(leaf1)
-    dt["scen1/sub_scen2"] = DataTree(leaf2)
-    dt["scen2"] = DataTree(leaf2)
+    dt = xr.DataTree()
+    dt["scen1/sub_scen1"] = xr.DataTree(leaf1)
+    dt["scen1/sub_scen2"] = xr.DataTree(leaf2)
+    dt["scen2"] = xr.DataTree(leaf2)
 
     res = mesmer.datatree.collapse_datatree_into_dataset(dt, dim=collapse_dim)
     assert isinstance(res, xr.Dataset)
@@ -89,7 +89,7 @@ def test_collapse_datatree_into_dataset():
     leaf3 = xr.merge(
         [ds1.assign_coords({"member": 1}), ds.assign_coords({"member": 1})]
     ).expand_dims("member")
-    dt = DataTree.from_dict({"scen1": leaf1, "scen2": leaf2, "scen3": leaf3})
+    dt = xr.DataTree.from_dict({"scen1": leaf1, "scen2": leaf2, "scen3": leaf3})
 
     res = mesmer.datatree.collapse_datatree_into_dataset(dt, dim=collapse_dim)
     assert isinstance(res, xr.Dataset)
@@ -103,7 +103,7 @@ def test_collapse_datatree_into_dataset():
     ds_with_different_time = ds1.shift(time=1)
 
     badleaf = ds_with_different_time.assign_coords({"member": 0}).expand_dims("member")
-    dt = DataTree.from_dict({"scen1": leaf1, "scen2": badleaf})
+    dt = xr.DataTree.from_dict({"scen1": leaf1, "scen2": badleaf})
 
     res = mesmer.datatree.collapse_datatree_into_dataset(dt, dim=collapse_dim)
 
@@ -117,13 +117,13 @@ def test_collapse_datatree_into_dataset():
     da2 = mesmer.testing.trend_data_2D(n_timesteps=n_ts, n_lat=n_lat, n_lon=n_lon)
     ds2 = xr.Dataset({"tas": da2})
 
-    dt = DataTree.from_dict({"mem1": ds1, "mem2": ds2})
+    dt = xr.DataTree.from_dict({"mem1": ds1, "mem2": ds2})
     res = mesmer.datatree.collapse_datatree_into_dataset(dt, dim="members")
 
     # empty nodes are removed before concatenating
     # NOTE: implicitly this is already there in the other tests, since the root node is always empty
     # but it is nice to have it explicitly too
-    dt = DataTree.from_dict({"scen1": leaf1, "scen2": DataTree()})
+    dt = xr.DataTree.from_dict({"scen1": leaf1, "scen2": xr.DataTree()})
     res = mesmer.datatree.collapse_datatree_into_dataset(dt, dim=collapse_dim)
     expected = leaf1.expand_dims(collapse_dim).assign_coords(
         {collapse_dim: np.array(["scen1"])}
@@ -133,19 +133,19 @@ def test_collapse_datatree_into_dataset():
 
 def test_extract_single_dataarray_from_dt():
     da = trend_data_1D(n_timesteps=30).rename("tas")
-    dt = DataTree.from_dict({"/": xr.Dataset({"tas": da})})
+    dt = xr.DataTree.from_dict({"/": xr.Dataset({"tas": da})})
 
     res = mesmer.datatree._extract_single_dataarray_from_dt(dt)
     xr.testing.assert_equal(res, da)
 
-    dt = DataTree(xr.Dataset({"tas": da, "tas2": da}))
+    dt = xr.DataTree(xr.Dataset({"tas": da, "tas2": da}))
     with pytest.raises(
         ValueError,
         match="Node must only contain one data variable, node has tas2 and tas.",
     ):
         mesmer.datatree._extract_single_dataarray_from_dt(dt)
 
-    dt = DataTree.from_dict(
+    dt = xr.DataTree.from_dict(
         {"scen1": xr.Dataset({"tas": da, "tas2": da}), "scen2": xr.Dataset({"tas": da})}
     )
 
@@ -165,7 +165,7 @@ def test_extract_single_dataarray_from_dt():
 
     # passing empty Dataree
     with pytest.raises(ValueError, match="node has no data."):
-        mesmer.datatree._extract_single_dataarray_from_dt(DataTree())
+        mesmer.datatree._extract_single_dataarray_from_dt(xr.DataTree())
 
 
 def test_stack_linear_regression_datatrees():
@@ -191,16 +191,16 @@ def test_stack_linear_regression_datatrees():
         {member_dim: np.arange(2)}
     )
 
-    target = DataTree.from_dict({"scen1": leaf1, "scen2": leaf2})
+    target = xr.DataTree.from_dict({"scen1": leaf1, "scen2": leaf2})
 
     d1D_1 = xr.Dataset({"tas": trend_data_1D(n_timesteps=n_ts)})
     d1D_2 = d1D_1 * 2
     d1D_3 = d1D_1 * 3
     d1D_4 = d1D_1 * 4
-    predictors = DataTree.from_dict(
+    predictors = xr.DataTree.from_dict(
         {
-            "pred1": DataTree.from_dict({"scen1": d1D_1, "scen2": d1D_2}),
-            "pred2": DataTree.from_dict({"scen1": d1D_3, "scen2": d1D_4}),
+            "pred1": xr.DataTree.from_dict({"scen1": d1D_1, "scen2": d1D_2}),
+            "pred2": xr.DataTree.from_dict({"scen1": d1D_3, "scen2": d1D_4}),
         }
     )
 
diff --git a/tests/unit/test_linear_regression.py b/tests/unit/test_linear_regression.py
index 627857a87..e0a1adf86 100644
--- a/tests/unit/test_linear_regression.py
+++ b/tests/unit/test_linear_regression.py
@@ -7,11 +7,10 @@
 
 import mesmer
 import mesmer.stats._linear_regression
-from mesmer.core._datatreecompat import DataTree
 from mesmer.testing import trend_data_1D, trend_data_2D
 
 
-def convert_to(dct: dict, data_type: str) -> dict | DataTree | xr.Dataset:
+def convert_to(dct: dict, data_type: str) -> dict | xr.DataTree | xr.Dataset:
     if data_type == "dict":
         return dct
     elif data_type == "DataTree":
@@ -21,7 +20,7 @@ def convert_to(dct: dict, data_type: str) -> dict | DataTree | xr.Dataset:
             for key, value in dct.items()
         }
 
-        return DataTree.from_dict(dct)
+        return xr.DataTree.from_dict(dct)
     elif data_type == "xr_dataset":
         return xr.Dataset(dct)
     else:
@@ -643,7 +642,7 @@ def test_linear_regression_datatree_data_in_root(
     pred0 = trend_data_1D(slope=1, scale=0).rename("bar")
     pred1 = trend_data_1D(slope=1, scale=0).rename("foo")
     ds = xr.Dataset({"pred0": pred0, "pred1": pred1})
-    dt = DataTree(ds, name="root")
+    dt = xr.DataTree(ds, name="root")
     tgt = trend_data_1D_or_2D(as_2D=as_2D, slope=slope, scale=0, intercept=intercept)
 
     result = lr_method_or_function(dt, tgt, "time")
diff --git a/tests/unit/test_weighted.py b/tests/unit/test_weighted.py
index b7001f126..d8aee5606 100644
--- a/tests/unit/test_weighted.py
+++ b/tests/unit/test_weighted.py
@@ -3,7 +3,6 @@
 import xarray as xr
 
 import mesmer
-from mesmer.core._datatreecompat import DataTree
 
 
 def data_lon_lat(as_dataset, x_dim="lon", y_dim="lat"):
@@ -174,7 +173,7 @@ def test_global_mean_weights_passed(as_dataset):
 
 
 def test_equal_scenario_weights_from_datatree():
-    dt = DataTree()
+    dt = xr.DataTree()
 
     n_members_ssp119 = 3
     n_members_ssp585 = 2
@@ -197,19 +196,19 @@ def test_equal_scenario_weights_from_datatree():
         }
     )
     ssp585 = ssp585.assign_coords(member=np.arange(n_members_ssp585))
-    dt = DataTree()
-    dt["ssp119"] = DataTree(ssp119)
-    dt["ssp585"] = DataTree(ssp585)
+    dt = xr.DataTree()
+    dt["ssp119"] = xr.DataTree(ssp119)
+    dt["ssp585"] = xr.DataTree(ssp585)
 
     result1 = mesmer.weighted.equal_scenario_weights_from_datatree(dt)
-    expected = DataTree.from_dict(
+    expected = xr.DataTree.from_dict(
         {
-            "ssp119": DataTree(
+            "ssp119": xr.DataTree(
                 xr.full_like(ssp119, fill_value=1 / n_members_ssp119).rename(
                     {"tas": "weights"}
                 )
             ),
-            "ssp585": DataTree(
+            "ssp585": xr.DataTree(
                 xr.full_like(ssp585, fill_value=1 / n_members_ssp585).rename(
                     {"tas": "weights"}
                 )
@@ -220,10 +219,10 @@ def test_equal_scenario_weights_from_datatree():
     # TODO: replace with datatree testing funcs when switching to xarray internal DataTree
     assert result1.equals(expected)
 
-    dt["ssp119"] = DataTree(
+    dt["ssp119"] = xr.DataTree(
         dt.ssp119.ds.expand_dims(gridcell=np.arange(n_gridcells), axis=1)
     )
-    dt["ssp585"] = DataTree(
+    dt["ssp585"] = xr.DataTree(
         dt.ssp585.ds.expand_dims(gridcell=np.arange(n_gridcells), axis=1)
     )
 
@@ -236,26 +235,26 @@ def test_equal_scenario_weights_from_datatree():
 
 def test_create_equal_scenario_weights_from_datatree_checks():
 
-    dt = DataTree()
+    dt = xr.DataTree()
     ssp119 = xr.Dataset(
         {"tas": xr.DataArray(np.ones((20, 2)), dims=("time", "member"))}
     )
     ssp585 = xr.Dataset(
         {"tas": xr.DataArray(np.ones((20, 3)), dims=("time", "member"))}
     )
-    dt = DataTree()
-    dt["ssp119"] = DataTree(ssp119)
-    dt["ssp585"] = DataTree(ssp585)
+    dt = xr.DataTree()
+    dt["ssp119"] = xr.DataTree(ssp119)
+    dt["ssp585"] = xr.DataTree(ssp585)
 
     # too deep
     dt_too_deep = dt.copy()
-    dt_too_deep["ssp585/1"] = DataTree(ssp585)
+    dt_too_deep["ssp585/1"] = xr.DataTree(ssp585)
     with pytest.raises(ValueError, match="DataTree must have a depth of 1, not 2."):
         mesmer.weighted.equal_scenario_weights_from_datatree(dt_too_deep)
 
     # missing member dimension
     dt_no_member = dt.copy()
-    dt_no_member["ssp119"] = DataTree(dt_no_member.ssp119.ds.sel(member=1))
+    dt_no_member["ssp119"] = xr.DataTree(dt_no_member.ssp119.ds.sel(member=1))
     with pytest.raises(
         ValueError, match="Member dimension 'member' not found in dataset."
     ):
@@ -263,13 +262,13 @@ def test_create_equal_scenario_weights_from_datatree_checks():
 
     # missing time dimension
     dt_no_time = dt.copy()
-    dt_no_time["ssp119"] = DataTree(dt_no_time.ssp119.ds.sel(time=1))
+    dt_no_time["ssp119"] = xr.DataTree(dt_no_time.ssp119.ds.sel(time=1))
     with pytest.raises(ValueError, match="Time dimension 'time' not found in dataset."):
         mesmer.weighted.equal_scenario_weights_from_datatree(dt_no_time)
 
     # multiple data variables
     dt_multiple_vars = dt.copy()
-    dt_multiple_vars["ssp119"] = DataTree(
+    dt_multiple_vars["ssp119"] = xr.DataTree(
         xr.Dataset(
             {
                 "tas": xr.DataArray(np.ones((20, 2)), dims=("time", "member")),


Reviews:
user: login: mathause
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-04T06:53:16Z
user: login: veni-vidi-vici-dormivi
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-04T09:14:54Z

Review comments:
diff_hunk: @@ -22,8 +22,8 @@ dependencies:
  - scipy=1.12
  - shapely=2.0
  - statsmodels=0.14
- - xarray=2024.2
- - xarray-datatree=0.0.13
+ - xarray=2025.03
+
path: ci/requirements/min-all-deps.yml
user: login: mathause
body: ```suggestion
```
created_at: 2025-04-04T06:45:21Z
author_association: MEMBER
diff_hunk: @@ -36,11 +23,9 @@ def map_over_datasets(func, *args, kwargs=None):
 else:
     raise ImportError(
         f"xarray version {xr.__version__} not supported - please upgrade to v2025.02 ("
path: mesmer/core/_datatreecompat.py
user: login: mathause
body: ```suggestion
        f"xarray version {xr.__version__} not supported - please upgrade to v2025.03 ("
```
created_at: 2025-04-04T06:45:59Z
author_association: MEMBER
diff_hunk: @@ -3,19 +3,7 @@
 import xarray as xr
 from packaging.version import Version
 
-if Version(xr.__version__) < Version("2024.10"):
-
-    from datatree import DataTree, map_over_subtree, open_datatree
-
-    def map_over_datasets(func, *args, kwargs=None):
-        "compatibility layer for older xarray versions"
-
-        if kwargs is None:
-            kwargs = {}
-
-        return map_over_subtree(func)(*args, **kwargs)
-
-elif Version(xr.__version__) > Version("2025.01"):
+if Version(xr.__version__) > Version("2025.01"):
path: mesmer/core/_datatreecompat.py
user: login: mathause
body: ```suggestion
if Version(xr.__version__) >= Version("2025.03"):
```
created_at: 2025-04-04T06:46:10Z
author_association: MEMBER
diff_hunk: @@ -46,8 +46,7 @@ install_requires =
     scikit-learn >=1.4 # only for the tests
     scipy >=1.12
     statsmodels >=0.14
-    xarray >=2024.2, <2024.10 # upper pin for xarray-datatree
-    xarray-datatree ==0.0.13
+    xarray >=2025.2 # for xr.DataTree
path: setup.cfg
user: login: mathause
body: ```suggestion
    xarray >=2025.3 # for xr.DataTree
```
created_at: 2025-04-04T06:51:55Z
author_association: MEMBER
diff_hunk: @@ -46,7 +47,7 @@ Breaking changes
   **scipy**         not specified 1.12
   **shapely**       not specified 2.0
   **statsmodels**   not specified 0.14
-  **xarray**        2023.04       2024.2
+  **xarray**        2023.04       2025.3
path: CHANGELOG.rst
user: login: veni-vidi-vici-dormivi
body: ```suggestion
  **xarray**        2023.04       2025.03
```
created_at: 2025-04-04T09:06:01Z
author_association: COLLABORATOR

PR comments:
user: login: codecov[bot]
created_at: 2025-04-03T14:52:50Z
body: ## [Codecov](https://app.codecov.io/gh/MESMER-group/mesmer/pull/627?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=MESMER-group) Report
Attention: Patch coverage is `97.50000%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 80.75%. Comparing base [(`82a5e25`)](https://app.codecov.io/gh/MESMER-group/mesmer/commit/82a5e25e9c8255a2e4d7a97945d113af6fbb01f4?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=MESMER-group) to head [(`1883187`)](https://app.codecov.io/gh/MESMER-group/mesmer/commit/18831877cad84641fd8c5e976aa901e54945334d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=MESMER-group).
> Report is 77 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/MESMER-group/mesmer/pull/627?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=MESMER-group) | Patch % | Lines |
|---|---|---|
| [mesmer/core/\_datatreecompat.py](https://app.codecov.io/gh/MESMER-group/mesmer/pull/627?src=pr&el=tree&filepath=mesmer%2Fcore%2F_datatreecompat.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=MESMER-group#diff-bWVzbWVyL2NvcmUvX2RhdGF0cmVlY29tcGF0LnB5) | 90.90% | [1 Missing :warning: ](https://app.codecov.io/gh/MESMER-group/mesmer/pull/627?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=MESMER-group) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #627      +/-   ##
==========================================
+ Coverage   80.40%   80.75%   +0.34%     
==========================================
  Files          51       51              
  Lines        3139     3133       -6     
==========================================
+ Hits         2524     2530       +6     
+ Misses        615      603      -12     
```

| [Flag](https://app.codecov.io/gh/MESMER-group/mesmer/pull/627/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=MESMER-group) | Coverage Δ | |
|---|---|---|
| [unittests](https://app.codecov.io/gh/MESMER-group/mesmer/pull/627/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=MESMER-group) | `80.75% <97.50%> (+0.34%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=MESMER-group#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/MESMER-group/mesmer/pull/627?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=MESMER-group).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=MESMER-group).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: mathause
created_at: 2025-04-04T09:00:14Z
body: @veni-vidi-vici-dormivi: would you mind reviewing the [changelog](https://mesmer-emulator--627.org.readthedocs.build/en/627/changelog.html#data-structure-using-datatree)? The rest of the PR is actually quite straight-forward. I will have another look at `map_over_datasets` but that probably needs to stay as is for now.

@yquilcaille & @l-pierini this PR drops support for `xarray-datatree` you will need to re-create the environment (or at least upgrade xarray) after this PR is merged

author_association: MEMBER
user: login: veni-vidi-vici-dormivi
created_at: 2025-04-04T09:15:15Z
body: Very cool, thanks!
author_association: COLLABORATOR