Search code examples
pythonpandaspython-xarray

Apply function element-wise


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)

Solution

  • 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.