apply_ufunc with dictionary or set arguments

[hassec]

I'm not sure I fully understand apply_ufunc's behavior.
Take the following simple example:

import xarray as xr
import numpy as np


def test(x, y, z):
    print(f"x: {x} | y: {y} | z: {z}")
    return x + y


x = xr.DataArray(np.linspace(0, 2, 3), dims="x")
y = xr.DataArray(np.linspace(2, 4, 3), dims="y")

# this works
z = [1, 2, 3]
# using either one of the below crashes
# z = {1, 2, 3}
# z = {"a": 1, "b": 2, "c": 3}

print(xr.apply_ufunc(test, x, y, z, input_core_dims=((), (), ("z")), vectorize=True))
# prints if using a list
# x: 0.0 | y: 2.0 | z: [1 2 3]
# x: 0.0 | y: 3.0 | z: [1 2 3]
# ...

How can I use apply_ufunc in a way where I can vectorize over these first two arguments, but pass z along unchanged to the test function.

What would happen if I then pass e.g.
z = xr.DataArray([{'a': i} for i in range(3)], dims="zs") ?

I guess what I'm asking for is for z to either be a "scalar" object, or if it is a DataArray then be interpreted as vector input.

[dcherian]

To pass arbitrary things, pass them as kwargs.

The variables in *args (here x, y, z) are interpreted as array-like. So dicts won't really work or sets, which are unordered, won't work well.

xr.apply_ufunc(test, x, y, z, input_core_dims=((), (), ("z")), vectorize=True))

Instead try

xr.apply_ufunc(test, x, y, input_core_dims=((), ()), kwargs={"z": z}, vectorize=True))

[hassec]

Thanks a lot for the reply 😊

That works, but then wouldn't allow me to still vectorize over z right?
I'm calling apply_ufunc in a bit of code where I can either get z being a dict (which should be interpreted as scalar) or a DataArray of dicts to vectorize over together with x,y.

[kmuehlbauer]

@hassec You could call different apply_ufunc-wrappers from your application depending on your input. It might be more easy for you to implement that switching logic in your code, since you know your data layout best.