Skip to content

ENH: Add core manipulation and coordinate methods to DataTree - #11587

Open
Shikhar-Kesharwani wants to merge 3 commits into
pydata:mainfrom
Shikhar-Kesharwani:enh/datatree-manipulation-methods
Open

Shikhar-Kesharwani wants to merge 3 commits into
pydata:mainfrom
Shikhar-Kesharwani:enh/datatree-manipulation-methods

Conversation

@Shikhar-Kesharwani

Copy link
Copy Markdown

What does this PR do?

Resolves #10015 (DataTree: missing methods)
Resolves #9472 (DataTree.assign_coords and similar should only set node coords)
Resolves #9336 (DataTree.drop_vars)

This PR implements core data manipulation, missing value handling, array operations, coordinate assignment, and variable/dimension dropping methods on DataTree:

1. Dimension Manipulation & Reshaping

  • DataTree.transpose(*dim, missing_dims="raise"): Transposes dimensions across all nodes in the tree. Supports custom dimension ordering, handles tree-wide dimension checks via _get_all_dims(), respects missing_dims ("raise", "warn", "ignore"), and supports trees with groups containing heterogeneous subsets of dimensions.
  • DataTree.squeeze(dim=None, drop=False, axis=None): Squeezes length-1 dimensions across the tree using self.isel(drop=drop, missing_dims="ignore", **{d: 0}).
  • DataTree.broadcast_like(other, exclude=None): Broadcasts the tree against another DataTree (by path matching) or Dataset/DataArray.

2. Missing Value Handling

  • DataTree.dropna(dim, *, how="any", thresh=None, subset=None): Computes a tree-wide non-NA mask along dim across data variables in all nodes containing dim, then indexes via self.isel({dim: mask}). This guarantees parent and child nodes remain strictly aligned along shared and inherited coordinates without misalignment or duplicate coordinate errors.
  • DataTree.fillna(value): Fills missing values with scalars, arrays, or variable-keyed dictionaries across all groups.

3. Array Operations & Filtering

  • DataTree.clip(min=None, max=None, *, keep_attrs=None): Limits values in numeric variables across all groups to [min, max].
  • DataTree.isin(test_elements): Tests element membership across the tree, returning a boolean DataTree.
  • DataTree.where(cond, other=dtypes.NA, drop=False): Filters elements across the tree according to cond. Correctly resolves relative paths (including path == ".") when both caller and condition are DataTree objects.
  • DataTree.pad(pad_width=None, mode="constant", ...): Pads arrays along specified dimensions across groups.
  • DataTree.roll(shifts=None, roll_coords=None, **shifts_kwargs): Rolls arrays along dimensions across groups.
  • DataTree.shift(shifts=None, fill_value=dtypes.NA, **shifts_kwargs): Shifts arrays along dimensions across groups.

4. Node Coordinates & Attributes (Addresses #9472, #9336)

  • DataTree.assign_coords(coords=None, **coords_kwargs): Assigns coordinates to the specific target node. As discussed in DataTree.assign_coords and similar should only set node coords #9472, this avoids attempting to redundantly map over child nodes; instead, child nodes naturally inherit indexed coordinates.
  • DataTree.assign_attrs(*args, **kwargs): Assigns attributes to the node, returning a new tree.
  • DataTree.drop_attrs(): Clears attributes on the node.
  • DataTree.drop_vars(names, *, errors="raise"): Drops variables across groups in the tree. Resolves datatree: drop_vars issue? #9336.
  • DataTree.drop_dims(drop_dims, *, errors="raise"): Drops dimensions and their associated variables across groups.

Tests

Added comprehensive test suite TestDataTreeManipulationMethods in xarray/tests/test_datatree.py covering all 16 methods, multi-node trees, varying node dimensions, error conditions, and coordinate inheritance verification (100% pass rate locally).

…#10015, pydata#9472, pydata#9336)

Implements core data manipulation, missing value handling, array operations,
coordinate assignment, and variable/dimension dropping methods on DataTree:
- Dimension manipulation: transpose, squeeze, broadcast_like
- Missing value handling: dropna, fillna
- Array operations: clip, isin, where, pad, roll, shift
- Node coordinates and attributes: assign_coords, assign_attrs, drop_attrs
- Variable and dimension dropping: drop_vars, drop_dims

All methods adhere to DataTree hierarchical structure, respect coordinate
inheritance, and keep parent and child groups aligned.

Closes pydata#10015
Closes pydata#9472
Closes pydata#9336

Signed-off-by: Shikhar Kesharwani <ayushgu02@gmail.com>
@Shikhar-Kesharwani
Shikhar-Kesharwani force-pushed the enh/datatree-manipulation-methods branch from 6ae42fe to 446bb4c Compare September 13, 2026 11:34
pre-commit-ci Bot and others added 2 commits September 13, 2026 11:35
- Replace loop with dictionary comprehension in DataTree.squeeze (PERF403)
- Use isinstance(other, DataTree) in broadcast_like to narrow type for Dataset.broadcast_like (resolves mypy type-var error)
- Place None at end of union types in map_over_datasets and to_netcdf (RUF036)
- Localize numpy and math imports in dropna to preserve module import structure

Signed-off-by: Shikhar Kesharwani <ayushgu02@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

topic-DataTree Related to the implementation of a DataTree class

Projects

None yet

Development

Successfully merging this pull request may close these issues.

DataTree: missing methods DataTree.assign_coords and similar should only set node coords datatree: drop_vars issue?

1 participant