[cmdupuis3]

Closes #1846

Overview

This PR adds map_blocks methods for UxDataset and UxDataArray.

ux.map_blocks (also UxDataArray.map_blocks / UxDataset.map_blocks) wraps xarray.map_blocks so each block reaches func with a grid attached and the result comes back with the original grid:

• grid dims unsplit (e.g. chunked along time): every block gets obj.uxgrid
• n_face split: each block gets uxgrid.isel(n_face=...) for its faces, located via face positions passed as a hidden map_blocks argument; node and edge data in args is gathered onto the block's sub-grid
• n_node or n_edge split: blocks keep the whole grid; grid variables passed in args are rechunked to match and arrive sliced to the block
• splitting more than one grid dim, or mixing grid dims, raises DimensionError

The docstring notes that grid-dim splits are only correct for element-local functions (and, with n_face split, faces computed from their own nodes and edges), not for neighbour-based ones like gradient.

Grid.chunks reports the grid's chunk sizes. Grid.chunk gains inplace: inplace=False returns a new chunked Grid; leaving it unset keeps the in-place behaviour with a FutureWarning about the coming default.

Face-area computation now handles zero-face grids (dtype probe instead of arr[0], np.min with initial), which xarray's template inference needs.

Expected Usage

import numpy as np
import uxarray as ux

uxds = ux.open_dataset(
    "test/meshfiles/ugrid/quad-hexagon/grid.nc",
    "test/meshfiles/ugrid/quad-hexagon/multi_dim_data.nc",
)
uxda = uxds["multi_dim_data"]  # dims (time, lev, n_face)


# 1. Chunked along time: every block carries the full grid.
def area_weighted(block):
    return block * block.uxgrid.face_areas

result = uxda.chunk(time=2).map_blocks(area_weighted)
result.compute()

# 2. Chunked along n_face: each block carries the sub-grid of its own faces,
#    so per-face grid quantities line up with the block.
result = uxda.chunk(n_face=2).map_blocks(area_weighted)

# 3. Extra arguments: grid variables are rechunked to match and split
#    alongside the block.
def mask_tropics(block, face_lat):
    return block.where(np.abs(face_lat) < 23.5)

result = uxda.chunk(n_face=2).map_blocks(mask_tropics, args=[uxda.uxgrid.face_lat])

PR Checklist

General

• An issue is created and linked
• Added appropriate labels (if your uxarray repo permissions allow it)
• Filled out Overview and Expected Usage (if applicable) sections

Testing & Benchmarking

• There is adequate test coverage of changes from this PR (add new tests if needed)
• If this PR could affect performance, ran ASV benchmarks and confirmed they show expected behavior (add a new benchmark if necessary)

Documentation and Examples

• Docstrings updated with any function changes, and included in all new functions
• User (public) functions added to docs/api.rst; internal (private) function names start with an underscore (_)

AI Disclosure

AI Usage: Claude Opus 5.5

• I have tested and take responsibility for all AI-generated content in my PR.