Prune data tree for Isomorphic operations (#10097)

* prune data tree function with tests

* update api.rst with prune

* refactor prune to filter_like and use self.filter in filter_like

* Added filter_like to whats-new

Author: kobebryant432

doc/api.rst

@@ -697,6 +697,7 @@ For manipulating, traversing, navigating, or mapping over the tree structure.
    DataTree.pipe
    DataTree.match
    DataTree.filter
+   DataTree.filter_like
 
 Pathlib-like Interface
 ----------------------

doc/whats-new.rst

@@ -21,6 +21,8 @@ v2025.02.0 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
+- Added :py:meth:`DataTree.filter_like` to conveniently restructure a DataTree like another DataTree (:issue:`10096`, :pull:`10097`).
+  By `Kobe Vandelanotte <https://github.com/kobebryant432>`_.
 - Added :py:meth:`Coordinates.from_xindex` as convenience for creating a new :py:class:`Coordinates` object
   directly from an existing Xarray index object if the latter supports it (:pull:`10000`)
   By `Benoit Bovy <https://github.com/benbovy>`_.

xarray/core/datatree.py

@@ -1394,6 +1394,54 @@ def filter(self: DataTree, filterfunc: Callable[[DataTree], bool]) -> DataTree:
         }
         return DataTree.from_dict(filtered_nodes, name=self.name)
 
+    def filter_like(self, other: DataTree) -> DataTree:
+        """
+        Filter a datatree like another datatree.
+
+        Returns a new tree containing only the nodes in the original tree which are also present in the other tree.
+
+        Parameters
+        ----------
+        other : DataTree
+            The tree to filter this tree by.
+
+        Returns
+        -------
+        DataTree
+
+        See Also
+        --------
+        filter
+        isomorphic
+
+        Examples
+        --------
+
+        >>> dt = DataTree.from_dict(
+        ...     {
+        ...         "/a/A": None,
+        ...         "/a/B": None,
+        ...         "/b/A": None,
+        ...         "/b/B": None,
+        ...     }
+        ... )
+        >>> other = DataTree.from_dict(
+        ...     {
+        ...         "/a/A": None,
+        ...         "/b/A": None,
+        ...     }
+        ... )
+        >>> dt.filter_like(other)
+        <xarray.DataTree>
+        Group: /
+        ├── Group: /a
+        │   └── Group: /a/A
+        └── Group: /b
+            └── Group: /b/A
+        """
+        other_keys = {key for key, _ in other.subtree_with_keys}
+        return self.filter(lambda node: node.relative_to(self) in other_keys)
+
     def match(self, pattern: str) -> DataTree:
         """
         Return nodes with paths matching pattern.

xarray/tests/test_datatree.py

@@ -1587,6 +1587,35 @@ def test_assign(self) -> None:
         result = dt.assign({"foo": xr.DataArray(0), "a": DataTree()})
         assert_equal(result, expected)
 
+    def test_filter_like(self) -> None:
+        flower_tree = DataTree.from_dict(
+            {"root": None, "trunk": None, "leaves": None, "flowers": None}
+        )
+        fruit_tree = DataTree.from_dict(
+            {"root": None, "trunk": None, "leaves": None, "fruit": None}
+        )
+        barren_tree = DataTree.from_dict({"root": None, "trunk": None})
+
+        # test filter_like tree
+        filtered_tree = flower_tree.filter_like(barren_tree)
+
+        assert filtered_tree.equals(barren_tree)
+        assert "flowers" not in filtered_tree.children
+
+        # test symetrical pruning results in isomorphic trees
+        assert flower_tree.filter_like(fruit_tree).isomorphic(
+            fruit_tree.filter_like(flower_tree)
+        )
+
+        # test "deep" pruning
+        dt = DataTree.from_dict(
+            {"/a/A": None, "/a/B": None, "/b/A": None, "/b/B": None}
+        )
+        other = DataTree.from_dict({"/a/A": None, "/b/A": None})
+
+        filtered = dt.filter_like(other)
+        assert filtered.equals(other)
+
 
 class TestPipe:
     def test_noop(self, create_test_datatree: Callable[[], DataTree]) -> None: