raster_tools.geo_map_blocks#

raster_tools.geo_map_blocks(func, *rasters, dtype=None, null_value=None, meta=None, out_bands=None, return_mask=False, **kwargs)[source]#

Apply func block-wise across one or more aligned rasters, handing it georeferenced xarray.DataArray blocks.

Same shape and contract as map_blocks(), but each raster’s data block is wrapped in a georeferenced xr.DataArray (with band / y / x coords from the block’s geobox, the raster’s CRS, and the raster’s null value attached as nodata) via to_dataarray() before being passed to func. Useful when the per-block function wants to operate in xarray-land (rio accessors, xr.where, etc.) without having to rebuild coordinates itself.

Available both as a module function and as a Raster method: r1.geo_map_blocks(func, r2, ...) is equivalent to geo_map_blocks(func, r1, r2, ...) – in the method form the calling raster is the first input. References to the “first input” below mean r1 in either spelling.

Per-block contract

Per-block kwargs (opt-in): input_masks, input_null_values, block_info, block_id, out_null_value, geo_block_info. Name any of these in func’s signature to receive them per chunk; see below.

By default the output Raster’s mask is rebuilt from the output data and the resolved output null value (out_data == null_value, or np.isnan(out_data) for NaN nulls) – write the sentinel only at cells you want masked. Cells your func happens to leave equal to the sentinel will appear masked even if you didn’t intend them to; cells you wanted masked but didn’t write the sentinel to will not. This is true regardless of the input rasters’ masks: input masks do not carry through unchanged. To set the output mask explicitly instead – decoupling nullness from the data values – pass return_mask=True and have func return a (data, mask) pair; see return_mask below.

Always passed positionally:

func(*data_dataarrays, **kwargs)

where data_dataarrays is a tuple of N georeferenced xarray.DataArray blocks, one per input raster, in input order.

The user can also opt in to receive per-block extras by including named parameters in func’s signature. Detection mirrors dask’s block_info= mechanism (via dask.utils.has_keyword()). Recognized names:

input_masks – tuple of N xr.DataArray (bool) per-block mask arrays, parallel to the data DataArrays. Same name as map_blocks(), but elements are xr.DataArray here (vs. np.ndarray there).
input_null_values – tuple of N scalars, each input’s null_value (None if unset).
block_info – dask’s standard per-block info dict.
block_id – dask’s standard per-block index tuple: the block’s chunk-location (equivalently block_info[None]["chunk-location"]). None during the meta inference call.
out_null_value – scalar; the resolved output null value the wrapper will use to derive the output mask. Write this sentinel at cells you want masked. See “Output null value resolution” in Notes below.
geo_block_info – the per-chunk GeoBlockInfo (the geo-aware analog of dask’s block_info). None during the meta inference call.

A function whose only kwargs absorber is **kwargs does NOT trigger any of these injections – name the kwargs you want.

Reserved kwargs

input_masks, input_null_values, block_info, block_id, out_null_value, and geo_block_info are reserved. Passing any of them via geo_map_blocks’s own **kwargs raises ValueError.

Parameters

func (callable) – Per-block function. See “Per-block contract” above. May return either an xarray.DataArray (its .data is extracted, preserving backend) or any array-like that dask can ingest (NumPy ndarray, cupy ndarray, sparse array, etc.) with the same shape as a single data block. For non-numpy backends, also pass meta=.
*rasters (Raster or str) – The input rasters. In the method form the calling raster is the first input and *rasters holds any additional ones; in the function form at least one is required. Path strings are accepted. All inputs must be on the same grid (CRS, affine, shape) within the established sub-pixel tolerance; mismatched inputs raise ValueError. Use r2.reproject(r1.geobox) to align inputs first if needed. Inputs are then auto-rechunked to the first input’s chunk structure (reproject does not adopt the target’s chunking), so the output stays on the first input’s grid and chunking.
dtype (dtype-like, optional) – Output dtype. When None (default), dask infers the dtype by calling func on tiny meta samples.
null_value (scalar, optional) –
Output null value.
- None (default): if there is exactly one input raster and the output dtype matches its dtype, the output inherits that input’s null value. Otherwise, a dtype-appropriate default from raster_tools.masking.get_default_null_value().
- scalar: used as-is.
meta (array-like, optional) – Empty array with the desired output array type. Forwarded to dask.array.map_blocks(). When provided, dask uses this as the output meta and skips the 0-shape sample call it would otherwise make to derive one – useful when func cannot tolerate 0-shape DataArray inputs. The wrapper extracts .data from any returned xarray.DataArray, so meta describes the wrapper’s array output (whatever backend the user’s func produces). When None (default), dask derives a NumPy meta by calling func on 0-shape inputs.
out_bands (int, optional) – Number of bands in the output. None (default) is shape-preserving. A positive integer lets func change the band count (the y/x grid is unchanged); func must return exactly out_bands bands per block (a mismatch raises ValueError). See map_blocks() for the full semantics: the input band axis is collapsed to a single chunk (so func sees all bands of a spatial tile, on possibly re-sized y/x tiles), the output is restored to the input’s original y/x chunking with band coord np.arange(out_bands) + 1, and passing dtype= / meta= is recommended. Setting out_bands to the input band count gives func an all-bands view of each tile (vs the default per-band blocks) without changing the count; see map_blocks().
return_mask (bool, optional) –
If True, func must return a (data, mask) pair instead of a single array/DataArray. The returned mask – not a sentinel comparison – defines the output Raster’s null cells. mask is a boolean array the same shape as data; masked cells are set to the resolved output null value (burned in). Use this to decouple which cells are null from what value they hold, avoiding the sentinel-collision pitfalls described above. The default is False (sentinel-derived mask). Notes:
- Each element of the pair may be an xarray.DataArray or a bare array; the wrapper uses each element’s .data (any coords / CRS are discarded, same as for the data block).
- The two arrays are carried through dask packed into a single NumPy structured-dtype block, then split apart again – an internal detail; func just returns the plain pair.
- Passing dtype= (or meta=) describing the data dtype is recommended: it lets dask skip its 0-shape probe. Without a hint the func must tolerate that probe.
- Requires NumPy-backed blocks. Composes with out_bands, input_masks, and geo_block_info. out_null_value injection is unnecessary (though harmless) when return_mask=True.
**kwargs – Extra keyword arguments forwarded per-block to func. The reserved names listed above are not allowed here.

Returns

A new lazy Raster on the first input’s grid (with out_bands bands when that argument is set).

Return type

Raster

Notes

Warning

func must not alter the grid. The output Raster lands on the first input’s grid (its CRS / affine / x / y), regardless of any coords / CRS / nodata your func sets on a returned DataArray – the wrapper extracts .data and discards everything else. Operations that change the grid (reproject, clip, coarsen, sel, manual coord assignment, etc.) will silently produce a Raster whose values were computed on a different grid than the one it claims, with no error at construction time. Use Raster.reproject() / Raster.clip() / etc. before or after this call instead.

The output mask is all-False if no null value is set (see the per-block contract above for how the mask is built when one is). To write the output mask directly, pass return_mask=True and return a (data, mask) pair from func (see return_mask).

Dask invokes func once on 0-shape DataArrays (with zero-filled placeholder coords, not real geocoordinates) to derive the output array meta – this happens whether or not dtype= is provided. Passing dtype= only skips the additional sample call dask would otherwise make to infer the output dtype; it does not skip the 0-shape meta call. (Exception: with return_mask=True, dtype= is folded into a structured meta= and so does skip the 0-shape meta call – see return_mask.) During the meta call geo_block_info is None if the func opts in. Most NumPy / xarray ops handle 0-shape inputs fine. If dask raises dtype inference failed in map_blocks. Please specify the dtype explicitly using the dtype kwarg, that’s the sample call – dtype= per dask’s hint is usually enough; pass meta= instead if your func also can’t tolerate 0-shape inputs (dask silently swallows that crash, but your downstream output meta will be wrong).

Output null value resolution

When func opts in to out_null_value, the wrapper resolves the scalar per-chunk without an extra dtype-inference pass:

With meta= or dtype= set, from that hint (dask leaves block_info[None]["dtype"] as None whenever meta= is set, which is why the hint is consulted first).
Otherwise, from block_info[None]["dtype"].

During dask’s meta inference call (where block_info is None), out_null_value is a typed zero of the first input’s dtype so funcs like np.where(m, out_null_value, d) infer the same dtype as their input rather than collapsing to object. If your func’s output dtype depends on the specific out_null_value scalar, pass meta= to skip dask’s 0-shape meta call entirely; dtype= skips only the additional sample call, not the meta call.

Examples

Use the per-chunk geobox to clip a vector layer and rasterize into the chunk. The output dtype is uint8, not the input’s float32 – pass meta= so dask uses the right dtype and skips the 0-shape meta call entirely (no geo_block_info is None guard needed):

>>> def clip_and_rasterize(xda, *, geo_block_info, gdf):
...     bbox = geo_block_info.bbox
...     ...                                          
...     return burned_uint8
>>> out = r.geo_map_blocks(                          
...     clip_and_rasterize,
...     meta=np.array((), dtype=np.uint8),
...     gdf=lines_gdf,
... )

For dtype-preserving funcs (output dtype equals input dtype), you can skip meta= and guard on the meta call instead:

>>> def scale(xda, *, geo_block_info=None):          
...     if geo_block_info is None:
...         return xda                # safe: same dtype, same shape
...     return xda * 2