How to apply function which based on DataArrays of Dataset?

[zxdawn]

Hi, I have three DataArrays inside one Dataset and want to apply a function based on them. Here's an example.

top, middle, and bottom DataArrays are three layers for calculating the integration of Gaussian distribution.
I have to integrate the x and y dims to calculate the ratio of integrations. Because the calculation is independent for each x/y grid, is it possible to use something like apply_func to solve it?

import numpy as np
from scipy import integrate
import xarray as xr

# three 2D arrays
top = np.arange(100., 104.).reshape(2,2)
middle = np.arange(500., 504.).reshape(2,2)
bottom = np.arange(1000., 1004.).reshape(2,2)

ds = xr.Dataset(
    {"top": (["x", "y"], top),
     "middle": (["x", "y"], middle),
     "bottom":  (["x", "y"], bottom)
     })

# the Gaussian function
def gaussian(x):
    return np.exp(-np.power(x - 600, 2)/(2*np.power(60, 2)))

# create empty array
ds['ratio'] = ds['top'].copy()

# calculate the ratio of integration of Gaussian between two layers
for x in ds['x']:
    for y in ds['y']:
        ds['ratio'].loc[dict(x=x, y=y)] = integrate.quad(gaussian, ds['top'].sel(x=x, y=y), ds['middle'].sel(x=x, y=y))[0] \
                / integrate.quad(gaussian, ds['top'].sel(x=x, y=y), ds['bottom'].sel(x=x, y=y))[0]

[dcherian]

IIUC you want something like

xr.apply_ufunc(
    integrate_at_single_x_y,
	ds.top,
	ds.middle,
	vectorize=True, # does the for loop for you
)

Neither x or y are a core dimension here, because the lowest-dimensional input your function requires is a scalar.

[zxdawn]

Thanks! Here's the example:

def integrate_at_single_x_y(top, low):
    return integrate.quad(lambda x: np.exp(-np.power(x - 600, 2)/(2*np.power(60, 2))),
                          top, low)[0]

integ_middle = xr.apply_ufunc(integrate_at_single_x_y, ds.top, ds.middle, vectorize=True)
integ_bottom = xr.apply_ufunc(integrate_at_single_x_y, ds.top, ds.bottom, vectorize=True)
ratio = integ_middle/integ_bottom