Trouble applying custom function over 3D dataset

[rmnvncnt]

I'm working on the following 3D dataset :

<xarray.Dataset>
Dimensions:  (index: 10000, model: 34, time: 7670)
Coordinates:
  * index    (index) int64 31882 58759 18178 81683 ... 8776 121253 59596 34761
    lat      (index) float64 dask.array<chunksize=(1,), meta=np.ndarray>
    lon      (index) float64 dask.array<chunksize=(1,), meta=np.ndarray>
  * model    (model) <U16 'ACCESS-CM2' 'ACCESS-ESM1-5' ... 'UKESM1-0-LL'
  * time     (time) datetime64[ns] 1990-01-01 1990-01-02 ... 2010-12-31
Data variables:
    sfcWind  (model, time, index) float32 dask.array<chunksize=(1, 7670, 1), meta=np.ndarray>

My goal is to apply the following function over each index * model combination :

def fit(da, time_index):
    # expects two numpy arrays of size (7670,)
    data = pd.Series(
        data=da
        index=time_index,
    )

    threshold = data.dropna().sort_values().iloc[-2]

    model = pyextremes.EVA(data=data)
    model.get_extremes("POT", threshold=threshold, r="24h")
    model.fit_model()

    summary = model.get_summary(
        return_period=[10, 100],
        alpha=0.95,
        n_samples=10000,
    )

    return summary.loc[10,"return value"]

My understanding is that xarray.apply_ufunc() is a good candidate for this task. I followed the documentation and went through the tutorial to come up with this :

def process_continent():
    rechunked_store = "../data/asia_rechunked/"
    ds = xr.open_zarr(rechunked_store, consolidated=True)

    result = xr.apply_ufunc(
        fit,
        ds,
        ds.time.values,
        input_core_dims=[['time'], []],
        dask="allowed",
        vectorize=True,
    )

What I expect is a 2D dataset with index * model as dimensions containing the result of the function as variables. But I'm unable to find the right combination of input_core_dims, output_core_dims and exclude_dims that would make it work. I invariably get the following error : ValueError: Length of values (34) does not match length of index (7670) or ValueError: Length of values (10000) does not match length of index (7670).

Am I missing something here?

Edit

I could make it work by loading the dataset into memory before passing it to apply_ufunc() (using ds.compute()). I don't know why dimensions are understood properly this way. This solution is not really desirable however, as the dataset is quite large and would not fit on all systems.

[mathause]

Would you need to use dask="parallelized"?

(

• I assume your real-world function actually uses summary result?
• It will be faster if you determine your threshold using numpy (could matter in a loop)

  import numpy as np
  import pandas as pd
  
  arr = np.array([1, 3, 2, np.nan])
  ser = pd.Series([1, 3, 2, np.nan], index=[0, 1, 2, 3])
  
  %timeit np.sort(arr[~np.isnan(arr)])[-2]
  %timeit ser.dropna().sort_values().iloc[-2]

)

[rmnvncnt]

Would you need to use dask="parallelized"?

I did try, but maybe I did not push it far enough. I'm going to try it again.

I assume your real-world function actually uses summary result?

You're right. I edited the question accordingly.

It will be faster if you determine your threshold using numpy (could matter in a loop)

Unfortunately, EVA model accepts Series as input only.

[mathause]

(but threshold is a scalar [not that it matters for your actual question])