nrt and dask

[loicdtx]

@kenoz and @floriandeboissieu,

We've been discussing the integration of Dask with nrt, here's some simple experiments to illustrate the "problem" and look for alternatives

Currently, Dask does not fully integrate with nrt. When we pass a Dask-based DataArray to the fit() method, it immediately triggers the computation of the array, negating Dask's lazy evaluation. While the computation can be parallelized at a lower level using numba, the necessity to fully evaluate the array before starting the fit() computation create a bottleneck. The example below illustrates this behavior:

from nrt import data
from nrt.monitor.iqr import IQR

romania = data.romania_20m(chunks={'time': 1, 'y': 10, 'x': 10})
model = IQR()
model.fit(romania.B8A)

The following snippet provides a simplified illustration of what happens internally in nrt. The main issue is with the input array flattening via fancy indexing (y_flat = y[:, mask]), which Dask does not support:

import numpy as np
import dask.array as da
import numba

@numba.njit()
def nanlstsq(X, y):
    """Return the least-squares solution to a linear matrix equation.

    Analogous to ``numpy.linalg.lstsq`` for dependent variables containing ``NaN`` values.
    """
    beta = np.zeros((X.shape[1], y.shape[1]), dtype=np.float64)
    for idx in numba.prange(y.shape[1]):
        # Subset y and X where y is not NaN
        isna = np.isnan(y[:, idx])
        X_sub = X[~isna]
        y_sub = y[~isna, idx]
        beta[:, idx] = np.linalg.solve(X_sub.T @ X_sub, X_sub.T @ y_sub)
    return beta

# Input array (time,y,x)
# y = np.random.rand(200,100,100).astype(np.float32)
# dask array with typical chunking you'll get when loading data using e.g. odc.stac
y = da.random.random((200,100,100), chunks=(1,50,50)).astype(np.float32)
shape = y.shape

# Design matrix (intercept and trend for the example, assuming regular time-series)
X = np.column_stack((np.ones(shape[0]), np.arange(shape[0]))).astype(np.float32)
beta_shape = (X.shape[1], shape[1], shape[2])

# Boolean mask (e.g., forest/non-forest)
mask = np.random.randint(0, 2, size=(shape[1], shape[2])).astype(bool)

# Initialize output array for regression coefficients
beta = np.zeros(beta_shape, dtype=np.float32)

# Attempt to flatten the array using fancy indexing (unsupported by Dask)
y_flat = y[:, mask]

# Compute regression coefficients
beta_flat = nanlstsq(X, y_flat)
beta[:, mask] = beta_flat

One potential solution is to wrap the computation in a function and apply it using map_blocks. This approach works well for this simplified example but not sure it will be as easy to port to nrt:

def process(X, y, mask):
    y = np.asarray(y)
    shape = y.shape
    beta_shape = (X.shape[1], shape[1], shape[2])
    beta = np.zeros(beta_shape, dtype=np.float32)
    y_flat = y[:, mask]
    beta_flat = nanlstsq(X, y_flat)
    beta[:, mask] = beta_flat
    return beta

beta = da.map_blocks(
    process,
    X=X,
    y=y.rechunk({0: -1}),  # Rechunk along the time dimension
    mask=mask,
    chunks=(X.shape[1], *y.chunks[1:]),
    dtype=np.float32
)

print(beta)
print(beta.compute())

Any thoughts?

[kenoz]

It seems that Fordead already uses the .map_blocks function in order to deal with dask arrays -> here (line 76).