Integrating numba with xarray/dask: Guidance on parallelizing operations over large datasets

[nicrie]

Hey there!

I recently began using numba to accelerate my computations. The initial experience was seamless and promising. However, I've hit a roadblock when attempting to combine dask and numba. While I understand that these two tools operate on "orthogonal" principles, I'm struggling to integrate them effectively. Below is a minimal working example (MWE) illustrating my current approach:

Minimal Working Example (MWE)

Consider a matrix A with dimensions (samples, features) and another matrix b, which can be thought of as a weighting vector, with dimensions (samples, 2). I aim to apply a transformation to each sample (row). Since each sample can be processed independently, the task should be fully parallelizable. Furthermore, the required transformation can be done using "numba -friendly" arithmetic operations.

import numpy as np
import xarray as xr
import numba 

A = np.random.rand(20, 5)  # 20 samples, 5 features
b = np.random.rand(20, 2)  # 20 samples, (fixed) extra dimension

@numba.njit(fastmath=True, parallel=True)
def foo_nb(A, b, n_out: int = 3):
    n_samples, n_features = A.shape

    res1 = np.empty((n_samples, n_features, n_out))
    res2 = np.empty((n_samples, n_out))
    res3 = np.empty((n_samples,))
    for i in range(n_samples):
        # here numba arithmetics happen
        # note that actually X will have a dimensionalty of (n_samples_reduced, n_features)
        # with n_samples_reduced ~ 20% n_samples
        X = A * np.sum(b[i] ** 2)
        # ... 

        U, s, VT = np.linalg.svd(X)
        res1[i] = VT[:n_out].T
        res2[i] = s[:n_out]
        res3[i] = np.sum(s)

    return res1, res2, res3

foo_nb(A, b, n_out=3)

This works flawlessly!

However, when trying to adapt this code to work with DataArrays, I'm uncertain about the steps. While I'm inclined to utilize xr.apply_ufunc, my initial attempts have been unsuccessful:

A = xr.DataArray(A, dims=["sample", "feature"])
b = xr.DataArray(b, dims=["sample", "extra_dim"])

# Attempt to parallelize over samples, so -> core dimensions
xr.apply_ufunc(
    foo_nb,
    A,
    b,
    input_core_dims=[["sample"], ["sample"]],
    output_core_dims=[["sample"], ["sample"], ["sample"]],
    # dask="parallelized",
)

Questions

1. How can I adapt the above example to work with both numba and xarray/dask?
2. Is there an optimized or more efficient solution to my current approach? Would leveraging numba CUDA/GPUs offer any advantage for such tasks?

For context, in some real-world scenarios, I anticipate handling datasets ranging from thousands to hundreds of thousands of samples. Additionally, the X matrix, which undergoes decomposition, will likely be significantly smaller, approximately 20% of the sample size of A.

[dcherian]

Look here: https://tutorial.xarray.dev/advanced/apply_ufunc/apply_ufunc.html and let us know how it goes. If you see opportunities to improve that material, PRs are very welcome!

[nicrie]

Thanks Deepak! The tutorial is incredibly well-structured and effectively communicates the core concepts. It helped me recognize that my current approach with @njit(parallel=True) might not be suitable, so I'll likely need to adapt my function to use @guvectorize. I'm in the process of figuring out how to do that, and I'll update once I've made progress.