raster_tools.map_blocks#

raster_tools.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.

Thin wrapper over dask.array.map_blocks(). Each call to func receives one block from each input raster, in input order. The output Raster adopts its CRS, affine, and x/y coords from the first input; its mask is derived from the output data by default, or set explicitly by func via return_mask (see Notes).

Available both as a module function and as a Raster method: r1.map_blocks(func, r2, ...) is equivalent to 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. 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 and avoiding both failure modes above – pass return_mask=True and have func return a (data, mask) pair; see return_mask below.

Always passed positionally:

func(*input_data, **kwargs)

where input_data is a tuple of N NumPy 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 own block_info= / block_id= mechanism and uses inspect.signature() (via dask.utils.has_keyword()). Recognized names:

input_masks – tuple of N np.ndarray (bool) per-block mask arrays, parallel to input_data.
input_null_values – tuple of N scalars, each input raster’s null_value (None if unset).
block_info – dask’s standard per-block info dict (see dask.array.map_blocks()).
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.

A function whose only kwargs absorber is **kwargs (e.g. def f(*args, **kwargs):) does NOT trigger any of these injections – inspect.signature only sees explicit parameter names. Name the kwargs you want.

Reserved kwargs

input_masks, input_null_values, block_info, block_id, and out_null_value are reserved. Passing any of them via map_blocks’s own **kwargs raises ValueError – otherwise the wrapper’s injection would silently clobber the caller’s value.

Parameters

func (callable) – Per-block function. See “Per-block contract” above for the full signature rules. Must return an array-like that dask can ingest (NumPy ndarray, cupy ndarray, sparse array, etc.) with the same shape as a single input data block. For non-numpy backends, also pass meta= so dask’s output meta is correct.
*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. Only the 3D shape is validated; CRS and affine are not checked. The caller is responsible for aligning inputs – typically via r2.reproject(r1.geobox). Inputs are auto-rechunked to the first input’s chunk structure, so same-shape inputs with differing chunking are handled. For a geo-aware variant that strictly requires matching grids, see geo_map_blocks().
dtype (dtype-like, optional) – Output dtype. When None (default), dask infers the dtype by calling func on tiny meta samples (matches NumPy promotion for the typical elementwise case, and reflects any in-func cast such as .astype).
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 (preserves the sentinel for identity-like single-input ops). Otherwise, the value is a dtype-appropriate default from raster_tools.masking.get_default_null_value() against the resolved output dtype – always representable, never overflows when the dtype changes.
- scalar: used as-is.
To force inheriting the first input’s null value across other cases, pass it explicitly: null_value=r1.null_value.
meta (array-like, optional) – Empty array with the desired output array type (e.g. np.array((), dtype=np.float32), or a CuPy / sparse equivalent). 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 input. 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: the output has the same band count as the input and func returns same-shape blocks. A positive integer lets func change the band count (the y/x grid is unchanged), e.g. out_bands=3 to emit three per-pixel statistics from a 1-band input, or out_bands=1 to collapse a multi-band stack. func must return exactly out_bands bands per block (a mismatch raises ValueError). Consequences when set:
- The input band axis is collapsed to a single chunk so func sees every band of a spatial tile at once, and the y/x tiles are transparently re-sized so a block holds roughly one dask array.chunk-size worth of data – so func may see different spatial tile sizes than the input’s native chunking.
- The output Raster is restored to the input’s original y/x chunking, with per-band band chunks and band coord np.arange(out_bands) + 1.
- Passing dtype= or meta= is recommended, since dask’s 0-shape meta call still runs and a band-reshaping func may not produce the right dtype/shape at 0-shape.
Setting out_bands to the input’s band count is a valid way to give func an all-bands view of each spatial tile: the default per-band blocking calls func once per band, whereas the band collapse above hands it every band of a tile at once. The count is unchanged, but the exact-out_bands return guard and the y/x re-tile still apply.
return_mask (bool, optional) –
If True, func must return a (data, mask) pair instead of a single array. 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:
- 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 entirely. Without a hint the func must tolerate that probe (same caveat as out_bands).
- Requires NumPy-backed blocks (the structured-dtype carrier is a NumPy concept); cupy / sparse outputs are not supported with return_mask. Composes with out_bands (the returned mask must also have out_bands bands).
- 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

Cross-input CRS or affine mismatches are not validated; the caller is responsible.

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 inputs to derive the output array meta – this happens whether or not dtype= is provided. Pass meta= to skip the call entirely; dtype= only skips the additional sample call dask would otherwise make to infer the output dtype, not 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.) Most NumPy 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

Plain elementwise on one input (no special injection):

>>> def double(d):
...     return d * 2
>>> doubled = r.map_blocks(double)               

Mask-aware multi-input – write the null sentinel where either input was masked:

>>> def add_skip_nulls(a, b, *, input_masks, **kwargs):
...     ma, mb = input_masks
...     out = a + b
...     out[ma | mb] = -9999
...     return out
>>> summed = map_blocks(                          
...     add_skip_nulls, r1, r2, null_value=-9999,
... )

Use dask’s per-block info inside func:

>>> def by_chunk_id(d, *, block_info, **kwargs):
...     bi = block_info[None]["chunk-location"]
...     ...                                       