return array from apply_ufunc?

[gbromley]

Hello,

I am wondering how to return an array of values from apply_ufunc? For context, I have daily precipitation data in the shape of (2551, 181, 261) and I'd like to parallelize processing of this data using xarray/dask. For each grid point I would take the entire time series and calculate 1 value per year, so the output from apply_ufunc would be of shape (181, 261, 71). My function works well when the input is numpy arrays for a single point, but as soon as I try to use apply_ufunc the output is returned as (181,261,1).

This should be an embarrassingly parallel problem, but not sure how to approach it.

I can post the function code if that's helpful, but here is the apply_ufunc call:

xr.apply_ufunc(
    onset_calculation,
    precip,
    precip.jday,
    precip.time,
    input_core_dims=[["time"],["time"],["time"]],
    exclude_dims=set(["time"]),
    output_core_dims=[["onset_date"]],
    vectorize='True',
    dask = 'parallelized',
)

[dcherian]

Can you take a look at https://xarray.pydata.org/en/stable/examples/apply_ufunc_vectorize_1d.html if you haven't already? I noticed that you have a typo: vectorize=True not vectorize='True'