Example of parellelization of xr.apply_ufunc as an alternative to groupby('y','x').apply(func)

xr_apply_ufunc_parallel.py

def interpolate_(dataarray, intdates):
    """
    Interpolate time series in a data array for the specified dates. Tries to
    parallelize using dask.
    
    Args:
        dataarray (xarray.DataArray): data array (time,y,x) with time series
        intdates (np.ndarray(np.datetime64)): array of dates to interpolate
        
    Returns:
        result (xarray.DataArray): data array with interpolated values stacked
                                    in new dimension itime (interpolated-time)
    """
    # Apply ufunc-- inputs xr.DataArray and dates for interpolation
    # returns data array with interpolated values for int_dates
    result = xr.apply_ufunc(ufunc_cubic_spline, dataarray,
                            input_core_dims=[['time']],
                            output_core_dims=[['itime']],
                            kwargs={'axis': -1,
                                    'orig_times': dataarray.time.values,
                                    'new_times': intdates},
                            dask='parallelized',
                            output_dtypes=[np.float32],
                            output_sizes={'itime':intdates.shape[0]})
    result['itime'] = ('itime', intdates)
    return result

def ufunc_cubic_spline(array, axis, orig_times, new_times):
    """
    Ufunc to fit a cubic spline on time series (z=f(t)) and interpolate
    values for the specified dates. y,x represent Lat, Lon.
    
    Args:
        array (np.ndarray): 3-D array with shape (y,x,time) having the response
                                values (y)
        axis (int): axis of the time dimension to apply the function
        orig_times: dates of the original data
        new_times: dates to interpolate
        
    Returns:
        interpolated (np.ndarray): 3-D array with interpolated values, stacked
                                    by interpolated dates in the third axis.
    """
    # Fit cubic spline and interpolate dates
    interpolated = np.apply_along_axis(int_cubic_spline,
                                       axis,
                                       array,
                                       orig_times=orig_times,
                                       new_times=new_times)
    
    return interpolated

def int_cubic_spline(y, orig_times, new_times):
    """
    Cubic spline fitting and interpolation removing NaN values.
    
    Args:
        y (np.ndarray): 1D array with response to fit spline
        orig_times (np.ndarray): 1D array with predictor to fit spline
        new_times (np.ndarray): 1D array with predictor values to interpolate
    
    Returns:
        interpolated (np 1d-array): interpolated values
    """
    # Filter NaNs in response
    nans = np.isnan(y)
    
    # Try to fit cubic spline with filtered y values
    try:
        spl = interpolate.CubicSpline(orig_times.astype('d')[~nans], y[~nans])
        
        interpolated = spl(new_times.astype('d'))
    
    except ValueError:
        warnings.warn('CubicSpline could not be fitted for one or more pixels')
        ## When spline cannot be fitted(not enought data), return NaN
        interpolated = np.empty(new_times.shape[0])
        interpolated[:] = np.nan
        
    return interpolated