Region Selection API

Access-side helpers for HEALPix stores — required for high-level stores written with implicit coordinates (save_pyramid(write_coords=False)), and convenient for any dense store. See the shared tutorial Accessing High-Resolution Regional Data.

select_bbox(ds, *, lon, lat, level=None, query_delta=_DEFAULT_QUERY_DELTA, load=True)

Extract all cells covering a geographic bounding box.

The box is rasterised at the coarse level level - query_delta; every coarse parent expands to one contiguous fine-level range of 4**query_delta cells. The result therefore covers the box (cells straddling the edge are included) with read granularity set by query_delta — align it with the store's chunk exponent for reads with zero waste.

Parameters:

Name	Type	Description	Default
ds	Dataset	HEALPix dataset (see select_cells).	required
lon	tuple[float, float]	(west, east) longitude bounds in degrees.	required
lat	tuple[float, float]	(south, north) latitude bounds in degrees.	required
level	int | None	HEALPix level override.	None
query_delta	int	Level offset for the coverage query. Larger values mean fewer, larger contiguous reads; smaller values follow the box outline more tightly.	_DEFAULT_QUERY_DELTA
load	bool	Load the selection into memory (default).	True

Returns:

Type	Description
Dataset	Compact subset covering the box.

Source code in grid_doctor/select.py

def select_bbox(
    ds: xr.Dataset,
    *,
    lon: tuple[float, float],
    lat: tuple[float, float],
    level: int | None = None,
    query_delta: int = _DEFAULT_QUERY_DELTA,
    load: bool = True,
) -> xr.Dataset:
    """Extract all cells covering a geographic bounding box.

    The box is rasterised at the coarse level ``level - query_delta``;
    every coarse parent expands to one contiguous fine-level range of
    ``4**query_delta`` cells.  The result therefore *covers* the box
    (cells straddling the edge are included) with read granularity set
    by *query_delta* — align it with the store's chunk exponent for
    reads with zero waste.

    Parameters
    ----------
    ds:
        HEALPix dataset (see [`select_cells`][grid_doctor.select_cells]).
    lon:
        ``(west, east)`` longitude bounds in degrees.
    lat:
        ``(south, north)`` latitude bounds in degrees.
    level:
        HEALPix level override.
    query_delta:
        Level offset for the coverage query.  Larger values mean fewer,
        larger contiguous reads; smaller values follow the box outline
        more tightly.
    load:
        Load the selection into memory (default).

    Returns
    -------
    xarray.Dataset
        Compact subset covering the box.
    """
    resolved_level = _dataset_level(ds, level)
    query_level = max(0, resolved_level - query_delta)
    module, kwargs = _require_healpix_geo_module(nest=True)
    coverage = np.asarray(
        module.zone_coverage(
            (lon[0], lat[0], lon[1], lat[1]),
            query_level,
            flat=True,
            **kwargs,
        )
    )
    parents = coverage[0].astype(np.int64)
    ranges = _parents_to_ranges(parents, resolved_level - query_level)
    logger.info(
        "Bounding box resolves to %d parent cells at level %d "
        "(%d contiguous ranges at level %d).",
        parents.size,
        query_level,
        len(ranges),
        resolved_level,
    )
    cells = np.concatenate([np.arange(a, b, dtype=np.int64) for a, b in ranges])
    return select_cells(ds, cells, level=resolved_level, load=load)

select_cone(ds, *, lon, lat, radius, level=None, query_delta=_DEFAULT_QUERY_DELTA, load=True)

Extract all cells within radius degrees of a centre point.

Parameters:

Name	Type	Description	Default
ds	Dataset	HEALPix dataset (see select_cells).	required
lon	float	Centre coordinates in degrees.	required
lat	float	Centre coordinates in degrees.	required
radius	float	Angular radius in degrees.	required
level	int | None	HEALPix level override.	None
query_delta	int	Level offset for the coverage query.	_DEFAULT_QUERY_DELTA
load	bool	Load the selection into memory (default).	True

Returns:

Type	Description
Dataset	Compact subset covering the cone.

Source code in grid_doctor/select.py

def select_cone(
    ds: xr.Dataset,
    *,
    lon: float,
    lat: float,
    radius: float,
    level: int | None = None,
    query_delta: int = _DEFAULT_QUERY_DELTA,
    load: bool = True,
) -> xr.Dataset:
    """Extract all cells within *radius* degrees of a centre point.

    Parameters
    ----------
    ds:
        HEALPix dataset (see [`select_cells`][grid_doctor.select_cells]).
    lon, lat:
        Centre coordinates in degrees.
    radius:
        Angular radius in degrees.
    level:
        HEALPix level override.
    query_delta:
        Level offset for the coverage query.
    load:
        Load the selection into memory (default).

    Returns
    -------
    xarray.Dataset
        Compact subset covering the cone.
    """
    resolved_level = _dataset_level(ds, level)
    query_level = max(0, resolved_level - query_delta)
    module, kwargs = _require_healpix_geo_module(nest=True)
    coverage = np.asarray(
        module.cone_coverage((lon, lat), radius, query_level, flat=True, **kwargs)
    )
    parents = coverage[0].astype(np.int64)
    ranges = _parents_to_ranges(parents, resolved_level - query_level)
    cells = np.concatenate([np.arange(a, b, dtype=np.int64) for a, b in ranges])
    return select_cells(ds, cells, level=resolved_level, load=load)

select_cells(ds, cells, *, level=None, load=True)

Extract the given HEALPix cells from a dataset.

Works on dense datasets (positional index equals cell ID, with or without materialised coordinates) and on compact subsets (cell coordinate holds actual IDs). Contiguous ID runs are read as contiguous slices, so a spatially compact selection touches only the chunks it needs.

Parameters:

Name	Type	Description	Default
ds	Dataset	HEALPix dataset, typically opened with xr.open_zarr(store, chunks=None).	required
cells	Int64Array | list[int]	Global HEALPix cell IDs to extract (any order; duplicates are dropped).	required
level	int | None	HEALPix level override when the healpix_level attribute is missing.	None
load	bool	Load the selected data into memory (default). The selection is small by construction — that is the point of selecting.	True

Returns:

Type	Description
Dataset	Compact subset: cell coordinate holds the requested IDs, cell-centre latitude/longitude are attached, and grid_doctor_sparse = 1 is set.

Source code in grid_doctor/select.py

def select_cells(
    ds: xr.Dataset,
    cells: Int64Array | list[int],
    *,
    level: int | None = None,
    load: bool = True,
) -> xr.Dataset:
    """Extract the given HEALPix cells from a dataset.

    Works on dense datasets (positional index equals cell ID, with or
    without materialised coordinates) and on compact subsets (``cell``
    coordinate holds actual IDs).  Contiguous ID runs are read as
    contiguous slices, so a spatially compact selection touches only
    the chunks it needs.

    Parameters
    ----------
    ds:
        HEALPix dataset, typically opened with
        ``xr.open_zarr(store, chunks=None)``.
    cells:
        Global HEALPix cell IDs to extract (any order; duplicates are
        dropped).
    level:
        HEALPix level override when the ``healpix_level`` attribute is
        missing.
    load:
        Load the selected data into memory (default).  The selection is
        small by construction — that is the point of selecting.

    Returns
    -------
    xarray.Dataset
        Compact subset: ``cell`` coordinate holds the requested IDs,
        cell-centre ``latitude``/``longitude`` are attached, and
        ``grid_doctor_sparse = 1`` is set.
    """
    _require_nested(ds)
    resolved_level = _dataset_level(ds, level)
    wanted = np.unique(np.asarray(cells, dtype=np.int64))
    if wanted.size == 0:
        raise ValueError("No cells requested.")
    npix = 12 * 4**resolved_level
    if int(wanted[0]) < 0 or int(wanted[-1]) >= npix:
        raise ValueError(
            f"Cell IDs must be within [0, {npix}) for level {resolved_level}."
        )

    if "cell" in ds.coords:
        # Compact dataset (or legacy dense store with coordinates):
        # positions are found through the coordinate values.
        coord = np.asarray(ds["cell"].values, dtype=np.int64)
        pos = np.searchsorted(coord, wanted)
        pos = np.clip(pos, 0, coord.size - 1)
        present = coord[pos] == wanted
        if not present.all():
            missing = wanted[~present]
            raise KeyError(
                f"{missing.size} requested cells are not present in the "
                f"dataset (first missing: {int(missing[0])})."
            )
        pieces = [ds.isel(cell=slice(a, b)) for a, b in _contiguous_runs(pos)]
    else:
        # Dense store without materialised coordinates: positional
        # index *is* the cell ID.
        pieces = [ds.isel(cell=slice(a, b)) for a, b in _contiguous_runs(wanted)]

    subset = (
        pieces[0]
        if len(pieces) == 1
        else xr.concat(pieces, dim="cell", data_vars="minimal", coords="minimal")
    )
    if load:
        subset = subset.load()
    return attach_cell_coords(subset, wanted, level=resolved_level, attrs=ds.attrs)

attach_cell_coords(ds, cells, *, level, attrs=None)

Attach computed cell coordinates to a compact subset.

HEALPix coordinates are a pure function of the cell index; stores written with write_coords=False carry none, and this function reconstructs them for exactly the cells at hand.

Parameters:

Name	Type	Description	Default
ds	Dataset	Subset whose cell dimension corresponds to cells.	required
cells	Int64Array	Global HEALPix cell IDs, one per position along cell.	required
level	int	HEALPix level of the IDs.	required
attrs	Any	Optional attribute mapping to merge (e.g. the source store's attributes).	None

Returns:

Type	Description
Dataset	Subset with cell, latitude, longitude, and crs coordinates, grid_mapping tags, and grid_doctor_sparse set.

Source code in grid_doctor/select.py

def attach_cell_coords(
    ds: xr.Dataset,
    cells: Int64Array,
    *,
    level: int,
    attrs: Any = None,
) -> xr.Dataset:
    """Attach computed cell coordinates to a compact subset.

    HEALPix coordinates are a pure function of the cell index; stores
    written with ``write_coords=False`` carry none, and this function
    reconstructs them for exactly the cells at hand.

    Parameters
    ----------
    ds:
        Subset whose ``cell`` dimension corresponds to *cells*.
    cells:
        Global HEALPix cell IDs, one per position along ``cell``.
    level:
        HEALPix level of the IDs.
    attrs:
        Optional attribute mapping to merge (e.g. the source store's
        attributes).

    Returns
    -------
    xarray.Dataset
        Subset with ``cell``, ``latitude``, ``longitude``, and ``crs``
        coordinates, ``grid_mapping`` tags, and ``grid_doctor_sparse``
        set.
    """
    cells = np.asarray(cells, dtype=np.int64)
    if ds.sizes.get("cell") != cells.size:
        raise ValueError(
            f"Dataset has {ds.sizes.get('cell')} cells but {cells.size} "
            "IDs were provided."
        )
    module, kwargs = _require_healpix_geo_module(nest=True)
    lon_deg, lat_deg = module.healpix_to_lonlat(cells, level, **kwargs)

    result = ds.assign_coords(
        cell=cells,
        latitude=("cell", np.asarray(lat_deg, dtype=np.float64)),
        longitude=(
            "cell",
            _canonical_lon(np.asarray(lon_deg, dtype=np.float64)),
        ),
        crs=_make_crs_variable(level=level, nside=2**level, order="nested"),
    )
    for name in result.data_vars:
        if "cell" in result[name].dims:
            result[name].attrs["grid_mapping"] = "crs"
    if attrs:
        merged = dict(attrs)
        merged.update(result.attrs)
        result.attrs = merged
    result.attrs["healpix_level"] = level
    result.attrs["healpix_nside"] = 2**level
    result.attrs["healpix_order"] = "nested"
    result.attrs["grid_doctor_sparse"] = 1
    return result