Writing Reusable Index Functions with xarray apply_ufunc

Write the index in NumPy once, then wrap it so it runs on labelled and chunked data without change:

import numpy as np
import xarray as xr

def _nd(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    den = a + b
    return np.divide(a - b, den, out=np.full(den.shape, np.nan, "float32"), where=den != 0)

def normalised_difference(a: xr.DataArray, b: xr.DataArray) -> xr.DataArray:
    return xr.apply_ufunc(_nd, a, b, dask="parallelized", output_dtypes=["float32"])

The NumPy core stays testable and fast; the wrapper takes care of coordinates, alignment and laziness. This page belongs to band math operations with xarray in Core Raster Fundamentals & STAC Mapping.

The Split Between Core and Wrapper

The core sees arrays; the wrapper sees DataArrays apply_ufunc takes labelled, chunked DataArrays, aligns them by coordinate, and calls the NumPy core once per chunk with plain arrays. The core returns plain arrays, and apply_ufunc reassembles them into a labelled, lazy result with the original coordinates. The core never needs to know about labels or chunking. Labels and chunks outside, plain arrays inside inputs DataArrays with x, y, time Dask chunks CRS in attributes apply_ufunc aligns by coordinate calls the core per chunk _nd(a, b) plain NumPy, no labels output same coordinates still lazy declared dtype Test the core with plain arrays; trust the wrapper to do the plumbing.

The benefit is testability as much as reuse. The core function takes and returns NumPy arrays, so it can be tested exhaustively with tiny synthetic inputs — including the overflow and zero-denominator cases — without constructing any DataArrays at all.

Environment & Setup

Package Version pin Used for
xarray >=2024.1 apply_ufunc and labelled alignment
numpy >=1.23 The core implementations
dask[array] >=2024.1 Chunk-wise lazy execution
numba >=0.59 Optional: compiling heavier per-pixel cores 
pip install "xarray>=2024.1" "numpy>=1.23" "dask[array]>=2024.1" "numba>=0.59"

Complete Working Example

from typing import Callable

import numpy as np
import xarray as xr


# ---- pure NumPy cores -------------------------------------------------------
def _nd(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    a, b = a.astype("float32", copy=False), b.astype("float32", copy=False)
    den = a + b
    return np.divide(a - b, den, out=np.full(den.shape, np.nan, "float32"), where=den != 0)


def _savi(nir: np.ndarray, red: np.ndarray, L: float = 0.5) -> np.ndarray:
    nir, red = nir.astype("float32", copy=False), red.astype("float32", copy=False)
    den = nir + red + L
    return np.divide((1 + L) * (nir - red), den,
                     out=np.full(den.shape, np.nan, "float32"), where=den != 0)


# ---- a generic wrapper ------------------------------------------------------
def as_index(core: Callable, name: str, long_name: str) -> Callable[..., xr.DataArray]:
    def wrapped(*bands: xr.DataArray, **kwargs) -> xr.DataArray:
        out = xr.apply_ufunc(
            core, *bands,
            kwargs=kwargs,
            dask="parallelized",
            output_dtypes=["float32"],
            keep_attrs="drop",               # band attributes do not describe an index
            join="exact",                    # refuse to silently align mismatched grids
        )
        out.name = name
        out.attrs.update(long_name=long_name, units="1", valid_range=[-1.0, 1.0])
        if bands[0].rio.crs is not None:
            out = out.rio.write_crs(bands[0].rio.crs)
        return out
    wrapped.__name__ = name
    wrapped.__doc__ = f"{long_name}, computed lazily per chunk."
    return wrapped


ndvi = as_index(_nd, "ndvi", "normalised difference vegetation index")
savi = as_index(_savi, "savi", "soil adjusted vegetation index")

join="exact" is the argument most people omit and later need. By default apply_ufunc performs an inner join on coordinates, so two bands whose grids differ by a fraction of a pixel are silently cropped to their intersection. exact raises instead, forcing the grids to be aligned deliberately — which is exactly the discipline described in aligning two rasters with reproject_match.

Core Dimensions for Temporal Functions

Pointwise versus along-an-axis functions A pointwise index needs no core dimensions: every pixel is independent and apply_ufunc broadcasts over all axes. A temporal function such as a seasonal amplitude needs the whole time axis at each pixel, so time is declared as an input core dimension; apply_ufunc moves it to the last axis and the function reduces over it, returning one value per pixel. pointwise: no core dims along time: core dim "time" each pixel computed alone (y, x) in → (y, x) out reduce whole series at each pixel (time, y, x) in → (y, x) out Core dims must be a single chunk, or the function sees only part of each series. 
import numpy as np
import xarray as xr


def _amplitude(series: np.ndarray) -> np.ndarray:
    """Seasonal amplitude over the LAST axis, ignoring NaN."""
    return (np.nanpercentile(series, 90, axis=-1)
            - np.nanpercentile(series, 10, axis=-1)).astype("float32")


def seasonal_amplitude(ndvi_series: xr.DataArray) -> xr.DataArray:
    series = ndvi_series.chunk({"time": -1})           # the core dim must be one chunk
    return xr.apply_ufunc(
        _amplitude, series,
        input_core_dims=[["time"]],                   # moved to the last axis
        dask="parallelized", output_dtypes=["float32"],
    )

The rechunk to a single time chunk is not optional: apply_ufunc passes each chunk to the core independently, so a time axis split across chunks would hand the function fragments of each series and the percentiles would be computed on the wrong samples. The same constraint shapes the chunking advice in building temporal feature stacks from image time series.

When Not to Use apply_ufunc

For formulas that xarray arithmetic expresses directly — every normalised difference, every ratio, most linear combinations — plain operators are clearer, equally lazy and equally fast. (nir - red) / (nir + red) on Dask-backed DataArrays builds the same graph apply_ufunc would.

The wrapper earns its place in three situations. When the core is written in something that does not understand xarray, such as a Numba-compiled per-pixel loop or a function from a compiled library. When the function works along an axis and needs explicit core dimensions. And when the same logic must run on plain NumPy arrays elsewhere in the codebase — a model inference script, say — and duplicating it in two styles would let them drift apart. Outside those cases, the simpler expression is the better one.

Verification

One core, three input types, one answer The same small synthetic input is passed as a NumPy array to the core, as an eager DataArray to the wrapper, and as a Dask-backed DataArray to the wrapper. All three results must be identical, including NaN positions, which confirms the wrapper neither changes the maths nor loses the mask. Same numbers from every path NumPy core _nd(a, b) eager DataArray ndvi(A, B) lazy DataArray ndvi(A.chunk(), B.chunk()) All three equal, NaN positions included. 
import numpy as np
import xarray as xr

a = np.array([[400, 0, 60000]], dtype="uint16")
b = np.array([[3000, 0, 62000]], dtype="uint16")
A = xr.DataArray(a, dims=("y", "x"))
B = xr.DataArray(b, dims=("y", "x"))

core = _nd(b, a)
eager = ndvi(B, A).values
lazy = ndvi(B.chunk({"x": 1}), A.chunk({"x": 1})).compute().values

for other in (eager, lazy):
    assert np.array_equal(np.isnan(core), np.isnan(other))
    assert np.allclose(core, other, equal_nan=True)

The synthetic input deliberately includes a zero denominator and values near the uint16 maximum, so the test covers the two numerical hazards alongside the plumbing. The overflow case in particular is described in avoiding integer overflow in index calculations.

Common Errors

ValueError: dimension time on 0th function argument consists of multiple chunks

A core dimension is split across chunks. Rechunk it to a single chunk before calling.

The result is smaller than the inputs

The default inner join cropped mismatched grids. Use join="exact" and align explicitly.

The result is computed immediately

dask="parallelized" was omitted, so xarray loaded the inputs. Add it along with output_dtypes.

The output lost its CRS

apply_ufunc does not carry rioxarray’s spatial reference. Write it back onto the result, as the wrapper does, taking it from the first input band.

Frequently Asked Questions

Q: Why use apply_ufunc if xarray arithmetic already works? For simple formulas plain arithmetic is better. apply_ufunc earns its place when the core function is written in NumPy, Numba or a compiled library that does not understand xarray, or when it needs control over dimensions, dtypes and vectorisation that operator overloading does not give.

Q: What does dask=‘parallelized’ do? It tells xarray to apply the function independently to each Dask chunk, producing a lazy result. The function itself only ever sees plain NumPy arrays, so it needs no Dask awareness at all.

Q: Why must output_dtypes be given? Because a lazy result needs a dtype before anything is computed. Without it xarray would have to run the function on sample data to discover one; with it, the graph is built instantly and the dtype is guaranteed.

Q: Can the core be a Numba function? Yes, and it is a good fit for per-pixel logic too complex for vectorised NumPy, such as an iterative fit per series. Compile it with numba.njit or guvectorize, keep its signature on plain arrays, and wrap it exactly as shown.