如何在xarray最内层等长维度上应用np.linalg.det()计算行列式？

其实apply_ufunc()就是正确的方法，你可能是参数设置有误。下面是具体实现步骤：

假设你的DataArray是da，最内层的两个维度是row和col（长度相等），可以参考以下代码：

import numpy as np
import xarray as xr

# 示例：创建包含外层维度的方阵DataArray
da = xr.DataArray(
    np.random.rand(3, 4, 2, 2),  # 外层维度为(3,4)，内层(2,2)构成方阵
    dims=['batch1', 'batch2', 'row', 'col']
)

# 调用apply_ufunc计算行列式
det_da = xr.apply_ufunc(
    np.linalg.det,
    da,
    input_core_dims=[['row', 'col']],  # 指定构成方阵的核心维度
    vectorize=True,  # 自动对所有外层维度做批量处理
    output_dtypes=[float]
)

关键参数说明： #

• input_core_dims=[['row', 'col']]：明确告知xarray，这两个维度是要传入np.linalg.det的核心维度，剩余的维度会被当作批量维度自动遍历。
• vectorize=True：保证函数可以批量作用于每个外层维度对应的方阵上。
• output_dtypes=[float]：指定输出结果的数据类型，由于行列式是标量，这里设置为float即可。

执行后，det_da的维度会是原DataArray去掉内层两个核心维度后的维度（比如示例中的('batch1', 'batch2')），每个元素对应原数据中对应位置方阵的行列式值。

如果你的内层维度不是最后两个，只需修改input_core_dims中的维度名称，无需调整数据结构。

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