Function group_subtrees

xarray/core/treenode.py:803–866

Iterate over subtrees grouped by relative paths in breadth-first order. `group_subtrees` allows for applying operations over all nodes of a collection of DataTree objects with nodes matched by their relative paths. Example usage:: outputs = {} for path, (node_a, node_b

(
    *trees: AnyNamedNode,
)

def group_subtrees(
    *trees: AnyNamedNode,
) -> Iterator[tuple[str, tuple[AnyNamedNode, ...]]]:
    """Iterate over subtrees grouped by relative paths in breadth-first order.

    `group_subtrees` allows for applying operations over all nodes of a
    collection of DataTree objects with nodes matched by their relative paths.

    Example usage::

        outputs = {}
        for path, (node_a, node_b) in group_subtrees(tree_a, tree_b):
            outputs[path] = f(node_a, node_b)
        tree_out = DataTree.from_dict(outputs)

    Parameters
    ----------
    *trees : Tree
        Trees to iterate over.

    Yields
    ------
    A tuple of the relative path and corresponding nodes for each subtree in the
    inputs.

    Raises
    ------
    TreeIsomorphismError
        If trees are not isomorphic, i.e., they have different structures.

    See Also
    --------
    DataTree.subtree
    DataTree.subtree_with_keys
    """
    if not trees:
        raise TypeError("must pass at least one tree object")

    # https://en.wikipedia.org/wiki/Breadth-first_search#Pseudocode
    queue = collections.deque([(NodePath(), trees)])

    while queue:
        path, active_nodes = queue.popleft()

        # yield before raising an error, in case the caller chooses to exit
        # iteration early
        yield str(path), active_nodes

        first_node = active_nodes[0]
        if any(
            sibling.children.keys() != first_node.children.keys()
            for sibling in active_nodes[1:]
        ):
            path_str = "root node" if not path.parts else f"node {str(path)!r}"
            child_summary = " vs ".join(
                str(list(node.children)) for node in active_nodes
            )
            raise TreeIsomorphismError(