将 xarray 中最内层的两个维度作为矩阵进行操作

Question

我在 xarray 中有一个多维 DataArray，其中最里面的两个维度具有相等的长度。我想将 np.linalg.det() 应用于由二维形成的方阵。我该怎么做？

我查看了文档并花了一些时间让 apply_ufunc() 工作，但我不认为这是正确的方法。

Answer 1

我认为 apply_ufunc 确实是可行的方法。这个怎么办？

import xarray as xr
import numpy as np

# sample dataarray with 3 dimensions: n, m, x

n = 100
m = 100
x = 50

data = np.random.rand(n, m, x)

da = xr.DataArray(
    data,
    dims=["n", "m", "x"],
    coords={"n": np.arange(n), "m": np.arange(m), "x": np.arange(x)},
)
print(da)

输出：

<xarray.DataArray (n: 100, m: 100, x: 50)>
array([[[0.94590714, 0.09124397, 0.76948787, ..., 0.47377442,
         0.49205666, 0.2548719 ],
        [0.76343264, 0.41485065, 0.14439847, ..., 0.59190488,
         0.24071555, 0.00809688],
        [0.35608536, 0.44927753, 0.50033374, ..., 0.63904572,
         0.95172323, 0.24194656],
        ...,
        [0.71689346, 0.44678367, 0.92199706, ..., 0.99269701,
        [0.11836402, 0.23515399, 0.48816566, ..., 0.60960886,
...
Coordinates:
  * n        (n) int64 0 1 2 3 4 5 6 7 8 9 10 ... 90 91 92 93 94 95 96 97 98 99
  * m        (m) int64 0 1 2 3 4 5 6 7 8 9 10 ... 90 91 92 93 94 95 96 97 98 99
  * x        (x) int64 0 1 2 3 4 5 6 7 8 9 10 ... 40 41 42 43 44 45 46 47 48 49

然后：

result = xr.apply_ufunc(
    np.linalg.det,
    da,
    input_core_dims=[["n", "m"]],
    output_core_dims=[[]],
)
print(result)

输出：

<xarray.DataArray (x: 50)>
array([-3.05070222e+25,  4.43036648e+25,  1.46495819e+25,  4.70795386e+24,
       ...
       -2.37514032e+24, -2.44225924e+25, -1.29536184e+24,  6.68729607e+24,
       -5.06400130e+24,  9.12843726e+24, -2.96175281e+25, -1.79864819e+26,
       -3.26686808e+24,  2.83969992e+23])
Coordinates:
  * x        (x) int64 0 1 2 3 4 5 6 7 8 9 10 ... 40 41 42 43 44 45 46 47 48 49