How to make it work when function in apply_ufunc returns different-sized output?

[jajcayn]

Hey there,

Suppose I have input data as space x time, and for each spatial point, I'd like to compute a decomposition of any kind. This means that the input to the function is 1D array (time), and my function returns decomposition as (mode, time). Now, the tricky part is that the number of modes depends on the input time series, which can be between 4-6. No matter how I set xr.apply_ufunc I am getting ValueError: could not broadcast input array from the shape (5,500) into shape (6,500).

Here is the dummy code:

import numpy as np
import xarray as xr

# random data - space x time
data = xr.DataArray(
    np.random.rand(4, 500),
    dims=["space", "time"],
    coords={"time": np.linspace(0, 2, 500)},
)

# decomposition here - the number of "components" depends on data and can be between 4 and 6
# i.e. array (500,) -> (4, 500) or (5, 500) or (6, 500)
# mimicked here with random integer
def func(a):
    assert a.ndim == 1, a.ndim
    random_ = np.random.randint(low=4, high=7)
    out = []
    for r in range(random_):
        out.append(a + r)
    return np.array(out)

# doesn't work
out = xr.apply_ufunc(
    func,
    data,
    input_core_dims=[["time"]],
    output_core_dims=[["frequency", "time"]],
    vectorize=True,
)

out

Is there any way to make this work? The ideal solution would be to use the highest number of modes in the frequency dimension and fill non-existent values with nan or something else.

Thanks!

[keewis]

Your function needs to return a array of constant length for apply_ufunc to be able to concatenate it.

If I understand your example correctly, you can simply create a constant-length array and fill that as needed:

def func(a):
    assert a.ndim == 1, a.ndim
    random_ = np.random.randint(low=4, high=7)
    out = np.full((7,), fill_value=np.nan)
    out[:random_] = a + np.arange(random_)

    return out