在 xarray 或 Numpy 中计算卷积类型的积分

如何解决在 xarray 或 Numpy 中计算卷积类型的积分

我有以下积分要计算：

其中 H(.) 是 Heaviside function。

现在，我知道如何手动迭代数组并以这种方式计算 I，但我显然希望从 xarray 或至少 numpy 中使用它来利用加速。

>

这里一个巨大的复杂问题是我的数据集很大，我必须及时分块才能处理它们。但我在到达那里之前就失败了。

例如，我试图从一个仅依赖于下面 MWE 中的 x、y 和 z 的数组开始，但是我仍然收到错误调用 get_z_star() 是因为显然 xr.apply_ufunc 传递的是一个 numpy 数组而不是 DataArray，这正是我想要的：

import numpy as np
import xarray as xr

x = np.linspace(0,1,10)
b_tot = xr.DataArray(np.random.randn(10,10,10),dims=('x','y','z','t'),coords=dict(x=x,y=x,z=x,t=x))
def get_H_func_int(b,b_prime):
    """ b needs to be a float
        b_prime needs to be a DataArray but can't depend on time
    """
    H_func = np.heaviside(b_prime-b,1/2)
    return H_func.integrate(('x','z'))
    
def get_z_star(b,b_prime,**kwargs):
    """ b here can be an dataArray but can't depend on time
        b_prime needs to be a DataArray but can't depend on time
    """
    I = xr.apply_ufunc(get_H_func_int,b,**kwargs)
    return I
    
b = b_tot.min()
b_prime = b_tot.isel(t=0)
print(get_H_func_int(b,b_prime))
print(get_z_star(b_prime,b_prime))

有什么想法吗？