Применить функцию к набору данных Xarray

Question

Я пытаюсь применить функцию к набору данных Xarray, используя маску dataset.where, чтобы решить, где ее применить. Я не уверен, как это сделать.

Некоторый контекст: набор данных имеет две переменные (A и B), которые представляют собой два перекрывающихся растровых изображения (одинаковый размер и координаты). Я хочу запустить функцию (определенную мной), которая выполняет вычисление для каждой ячейки сетки растрового изображения с использованием двух переменных. Это не u_fun c, это просто функция, которая принимает значения двух перекрывающихся ячеек сетки, выполняет некоторые вычисления и возвращает третье значение, которое я хочу сохранить в третьем растровом изображении (C).

Как мне это сделать?

Я зашел так далеко, но это не работает, потому что ds.A передает функции весь массив, а не только значение ds.A в этой конкретной ячейке сетки :

def my_func(x, y):
    ..do things
    return result

ds = xr.Dataset({
    "A": xarray.open_rasterio("A.tif"), 
    "B": xarray.open_rasterio("B.tif"),
    "C": xarray.open_rasterio("C.tif"),
})

ds.C.where(ds.A > 0, my_func(ds.A, ds.B))