In xarray, how can I apply a non-vectorize, non-universal function to a DataArrray, such that I can map each element in the values to a new one ? The custom function shall take a scalar value and return a scalar value:
import numpy as np
import xarray as xr
def custom_function(x):
# imagine some non-vectorized, non-numpy-ufunc stuff
return type(x) # dummy example function
data = xr.DataArray([1, 2, 3, 4], dims='x')
# this doesn't work, custom_function actually gets send the whole array [1, 2, 3, 4]
# xr.apply_ufunc(custom_function, data)
# I'd expect something like this, where this is basically a loop on all elementss
# xr.apply(custom_function, data)
# in pandas, I would just use the .apply or .map method of Series
import pandas as pd
s = data.to_series()
s.apply(custom_function)
s.map(custom_function)
Using apply_ufunc as suggested, I get the expected result with vectorize=True
:
data = xr.DataArray([1, 2, 3, 4], dims='x')
# xr.apply_ufunc(custom_function, data) # fails, equivalent to vectorize=False
xr.apply_ufunc(custom_function, data, vectorize=True) # working
<xarray.DataArray (x: 4)>
array([<class 'int'>, <class 'int'>, <class 'int'>, <class 'int'>],
dtype=object)
Dimensions without coordinates: x
Explanation: using vectorize=True
wraps the custom_function with numpy.vectorize
, basically turning it into a for-loop, which is what I needed.