Function _datatree_concat

xarray/structure/concat.py:929–979

Concatenate a sequence of datatrees along a new or existing dimension

(
    objs: Iterable[DataTree],
    dim: Hashable | Variable | T_DataArray | pd.Index | Any,
    data_vars: T_DataVars | Iterable[Hashable] | CombineKwargDefault,
    coords: ConcatOptions | Iterable[Hashable] | CombineKwargDefault,
    compat: CompatOptions | CombineKwargDefault,
    positions: Iterable[Iterable[int]] | None,
    fill_value: Any,
    join: JoinOptions | CombineKwargDefault,
    combine_attrs: CombineAttrsOptions,
    create_index_for_new_dim: bool,
)

def _datatree_concat(
    objs: Iterable[DataTree],
    dim: Hashable | Variable | T_DataArray | pd.Index | Any,
    data_vars: T_DataVars | Iterable[Hashable] | CombineKwargDefault,
    coords: ConcatOptions | Iterable[Hashable] | CombineKwargDefault,
    compat: CompatOptions | CombineKwargDefault,
    positions: Iterable[Iterable[int]] | None,
    fill_value: Any,
    join: JoinOptions | CombineKwargDefault,
    combine_attrs: CombineAttrsOptions,
    create_index_for_new_dim: bool,
) -> DataTree:
    """
    Concatenate a sequence of datatrees along a new or existing dimension
    """
    from xarray.core.datatree import DataTree
    from xarray.core.treenode import TreeIsomorphismError, group_subtrees

    dim_name, _ = _calc_concat_dim_index(dim)

    objs = list(objs)
    if not all(isinstance(obj, DataTree) for obj in objs):
        raise TypeError("All objects to concatenate must be DataTree objects")

    if compat == "identical":
        if any(obj.name != objs[0].name for obj in objs[1:]):
            raise ValueError("DataTree names not identical")

    dim_in_tree = any(dim_name in node.dims for node in objs[0].subtree)

    results = {}
    try:
        for path, nodes in group_subtrees(*objs):
            datasets_to_concat = [node.to_dataset() for node in nodes]
            results[path] = _dataset_concat(
                datasets_to_concat,
                dim=dim,
                data_vars=data_vars,
                coords=coords,
                compat=compat,
                positions=positions,
                fill_value=fill_value,
                join=join,
                combine_attrs=combine_attrs,
                create_index_for_new_dim=create_index_for_new_dim,
                preexisting_dim=dim_in_tree,
            )
    except TreeIsomorphismError as e:
        raise ValueError("All trees must be isomorphic to be concatenated") from e

    return DataTree.from_dict(results, name=objs[0].name)