Search code examples
python-xarray

XArray apply function to the coordinates

How do I create a data array by applying a function over the values of the coordinates of a XArray Dataset?

Ex.

ds = xr.Dataset(
    coords={
        'x': [1.5, 4, 2],
        'y': [2, -3, -1],
    },
)
# some function like *apply that takes the coordinates as a list?
r = ds.apply*(lambda coords: math.sqrt(coords.x**2 + coords.x**2))

Should output

r = [ [ [2.5, ..],
        [.., 5, ..], .. ]
Solution

  • You can do this using xr.apply_ufunc

    ds = xr.Dataset(
    coords={
        'x': [1.5, 4, 2],
        'y': [2, -3, -1],
    },
)

r = xr.apply_ufunc(lambda x,y: math.sqrt(x**2 + y**2), ds.x, ds.y, vectorize=True)

    Outputs:

    array([[2.5       , 3.35410197, 1.80277564],
       [4.47213595, 5.        , 4.12310563],
       [2.82842712, 3.60555128, 2.23606798]])