How to apply a function in parallel?

[monocongo]

I want to apply a function to a 3-D dataset in parallel. The function takes a 1-D numpy array of values as input and returns a 1-D array of values of the same length. The function's input data is assumed to be a time series of values for a lat/lon slice from a lat/lon/time cube. Conceptually the intention is to compute a lat/lon/time cube of the same size/shape as the input by applying the function to each time series slice (pseudocode):

# input.dims == [lats, lons, times]
result = input.copy() # to get an equivalent size result array, we'll overwrite the values below
for lat in range(input.shape[0]):
    for lon in range(input.shape[1]):
        result[lat, lon] = my_function(input[lat, lon])

The data originates as a NetCDF with the variable having dimensions (time, lat, lon).

What I've tried so far is to read the NetCDF and then re-orient the data with a transpose to make the time dimension inner-most since this should make the retrieval of each slice more efficient. Then I have used apply_ufunc, like so:

from climate_indices.compute import Periodicity
from climate_indices.indices import Distribution, spi
import dask.array as da
import numpy as np
import xarray as xr

# extract the data from a NetCDF file
ds = xr.open_dataset('/home/james/tmp/nclimgrid_prcp.nc', engine='h5netcdf').chunk({'lat': 20, 'lon': 50, 'time': -1})

# the variable of interest is named 'prcp'
precipitation_da = ds['prcp']

# transpose the array to make time the inner-most dimension
precipitation_da = precipitation_da.transpose('lat', 'lon', 'time')

# define the parameters for the spi function
scale = 3
distribution = Distribution.gamma
data_start_year = 1895
calibration_year_initial = 1920
calibration_year_final = 2000
periodicity = Periodicity.monthly
fitting_params = None

# define the custom function to be applied using apply_ufunc
def apply_spi(
    values: np.ndarray,
) -> np.ndarray:
    return spi(
        values=values,
        scale=scale,
        distribution=distribution,
        data_start_year=data_start_year,
        calibration_year_initial=calibration_year_initial,
        calibration_year_final=calibration_year_final,
        periodicity=periodicity,
        fitting_params=fitting_params,
    )

# use xarray's apply_ufunc to apply the function across the time dimension
spi_da = xr.apply_ufunc(
    apply_spi,
    precipitation_da,
    input_core_dims=[['time']],  
    output_core_dims=[['time']],
    vectorize=True,
    dask='parallelized',
    output_dtypes=[np.float64], 
)

When I run the code above it takes hours to complete on a dataset that's from a NetCDF only 1.1GB so not very big. When it runs I don't see my machine's CPU using more than a single core so it seems that this doesn't run in parallel using multiprocessing. I feel like I must be doing something wrong.

Maybe I've not started a required Dask service or somewhere have a configuration setting that's preventing multi-core execution, etc.? Perhaps I need to somehow "vectorize" the function?

I have developed a way to utilize multiprocessing and shared memory arrays to get a full parallelization of this process, but it's a Rube Goldberg machine that I'd love to replace if I can get the split-apply-combine approach with xarray to work.

Please advise if there's anything I'm doing above that's wrong or any other optimizations I can try to use for this. Thanks in advance for any comments or suggestions.

[dcherian]

By default dask will use the threaded scheduler, this won't help you with netCDF. First, I don't understand where you are computing values in the example but using spi_da.compute(scheduler="processes") is what you want.

Re: vectorization; I recommend starting here: https://tutorial.xarray.dev/advanced/apply_ufunc/automatic-vectorizing-numpy.html and reading the rest of the series.