How can I get a core dimension in a group?

[ricardobarroslourenco]

I'm processing the following data with xr.apply_ufunc:

And when I pass:

result = xr.apply_ufunc(
    global_correlation,  # first the function
    lai_array['lai_north_masked_30_groupby_year_and_season_mean_detrended'],
    lai_array['temp_north_masked_30_groupby_year_and_season_mean_detrended'],
    input_core_dims=[['year'],['year']],
    vectorize=True,
    dask='parallelized',
    output_dtypes=['float']
)

I get the error:
ValueError: operand to apply_ufunc has required core dimensions ['year'], but some of these dimensions are absent on an input variable: ['year']

How can I refer to year, which is part of the group (together with season)? I was able to make it work in the version without traversing all coords, as well when I pass input_core_dims=[['group'],['group']] (which gives me wrong results, since I want to calculate separately for each season).

[dcherian]

You'll need to unstack("group") first which should work OK with your current chunking.

It seems like you're stacking/unstacking to work around the limitation that Xarray only lets you group by a single variable.

Using flox will be better with dask, and more natural: https://flox.readthedocs.io/en/latest/intro.html#grouping-by-multiple-variables

[ricardobarroslourenco]

Thanks for the reference and the solution! Now it works beautifully and calculates fast.