xr.apply_ufunc returns unexpected dimensions

[ameliafdezrodriguez]

Describe the bug
Im trying to compute SPI from Python script instead of console with the input file example. From process_climate_indices works fine but not from python. It gives an error using the xr.apply_ufunc function

To Reproduce
Very simple code:

from climate_indices import compute, indices
import xarray as xr
import numpy as np

infile = '/example_climate_indices-master/input/nclimgrid_lowres_prcp.nc'

ds = xr.open_dataset(infile)

da_precip_lo = ds['prcp']

preferred_dims = ("lat", "lon", "time")
da_precip_lo = da_precip_lo.transpose(*preferred_dims)

data_arrays = {
    "low": da_precip_lo,
}
for label, da in data_arrays.items():
    if da["lat"][0] > da["lat"][1]:
        print(f"The {label}-resolution DataArray's lats are descending -- flipping")
        da["lat"] = np.flip(da["lat"])
    if da["lon"][0] > da["lon"][1]:
        print(f"The {label}-resolution DataArray's lons are descending -- flipping")
        da["lon"] = np.flip(da["lon"])



def apply_spi_gamma_monthly(
    data_array: xr.DataArray,
    months: int,
    data_start_year: int = 1895,
    calibration_year_initial: int = 1900,
    calibration_year_final: int = 2000,
) -> xr.DataArray:
    # stack the lat and lon dimensions into a new dimension named point, so at each lat/lon
    # we'll have a time series for the geospatial point, and group by these points
    da_precip_groupby = data_array.stack(point=("lat", "lon")).groupby("point")

    spi_args = {
        "scale": months,
        "distribution": indices.Distribution.gamma,
        "data_start_year": data_start_year,
        "calibration_year_initial": calibration_year_initial,
        "calibration_year_final": calibration_year_final,
        "periodicity": compute.Periodicity.monthly,
    }

    # apply the SPI function to the data array
    da_spi = xr.apply_ufunc(
        indices.spi,
        da_precip_groupby,
        kwargs=spi_args,
    )

    # unstack the array back into original dimensions
    da_spi = da_spi.unstack("point")

    return da_spi


da_spi_lo = apply_spi_gamma_monthly(da_precip_lo, 3)`

Expected behavior
Values of SPI

Error

837 data = as_compatible_data(data)
    838 if data.ndim != len(dims):
--> [839](file:///E:/InfoSequia/.venv/Lib/site-packages/xarray/computation/apply_ufunc.py:839)     raise ValueError(
    840         "applied function returned data with an unexpected "
    841         f"number of dimensions. Received {data.ndim} dimension(s) but "
    842         f"expected {len(dims)} dimensions with names {dims!r}, from:\n\n"
    843         f"{short_array_repr(data)}"
    844     )
    846 var = Variable(dims, data, fastpath=True)
    847 for dim, new_size in var.sizes.items():

ValueError: applied function returned data with an unexpected number of dimensions. Received 1 dimension(s) but expected 2 dimensions with names ('time', 'point'), from:

array([nan, nan, nan, ..., nan, nan, nan], dtype=float32)

[monocongo]

I have tried several times over the years to get parallelism with xarray to work as you're attempting here and I never was successful at it, hence I wrote a convoluted Rube Goldberg machine in the main processing script to use shared memory objects with multiprocessing, and it sped things up tremendously but it's complex/fragile code, I'm afraid to touch it now. I know that doesn't help solve this issue for you, but hopefully some additional context may be insightful.

One idea I've had is to pull data into a polars data frame object and perform the parallelization there (i.e. leveraging what polars has already built in). If you try that or any other approaches that work please report back, we'd definitely like to leverage your discoveries.

Finally thanks for using the package, I'm glad it's still useful to others.

[ameliafdezrodriguez]

Thank you for your quick answer!

I was surprised recently by this error beacuse I used to run the same code earlier and it works. Just checked that python version was 3.9.9 and xarray version 0.20.1 back then. So apparently xarray ufunc function has changed it's internal behaviour since then.

What would be your recommendation then? just iterate by lat lon as in this example?