如何在xarray.DataArray.coarsen.reduce()中应用自定义众数函数？

错误原因 #

xarray的coarsen.reduce()方法在调用自定义聚合函数时，会自动传入axis参数（用于指定需要聚合的维度），但你定义的find_mode函数没有声明接收这个参数，因此触发TypeError: find_mode() got an unexpected keyword argument 'axis'。

解决方案 #

你可以通过两种方式修复这个问题：

方案1：修改函数，支持接收axis参数

调整find_mode函数，让它接受axis参数，并在处理时将窗口数组沿指定轴展平（因为coarsen的窗口是多维的，需要聚合所有元素的众数）：

import numpy as np
import xarray as xr

def find_mode(window, axis=None):
    # 沿指定轴展平数组，若未指定axis则直接展平为一维
    if axis is not None:
        window = window.reshape(-1)
    uniq, counts = np.unique(window, return_counts=True)
    return uniq[np.argmax(counts)]

# 测试代码
ar = np.array([[2, 0],[3, 0]])
da = xr.DataArray(ar, dims=['x', 'y'], coords=[[0, 1], [0, 1]])

# 调用coarsen.reduce
result = da.coarsen(x=2, y=2).reduce(find_mode)
print(result)
# 输出：<xarray.DataArray ()>
# array(0)

方案2：用包装函数忽略axis参数

如果不想修改原函数，可以用functools.partial或者lambda函数包裹原函数，忽略xarray传递的axis参数：

用functools.partial的示例：

from functools import partial

# 原函数不变
def find_mode(window):
    uniq, counts = np.unique(window, return_counts=True)
    return uniq[np.argmax(counts)]

# 调用时用partial包装
result = da.coarsen(x=2, y=2).reduce(partial(find_mode))

用lambda函数的示例：

result = da.coarsen(x=2, y=2).reduce(lambda window, **kwargs: find_mode(window))

两种方案都能正确返回众数0，和NumPy数组的计算结果一致。

内容的提问来源于stack exchange，提问作者hlee