将 3D Array 过滤为一个 1D Array 的最快方法

问题描述

我正在尝试按条件过滤多个 3D 数组，并将结果用作 scipy.stats.binned_statistic_2d() 或 np.histogram2d() 的输入。

这是一个简单的例子：

import numpy as np
import xarray as xr

array = xr.DataArray(np.random.randn(20,2000,3000))

stacked = array.stack(z=('dim_0','dim_1','dim_2'))
res = stacked.where(stacked>0,drop=True)

stack 步骤的运行速度非常慢：

16.9 s ± 1 s per loop (mean ± std. dev. of 7 runs,1 loop each)

是否可以更快地获取 res 并将 dim_0、dim_1 和 dim_2 的值也保存为一维变量？

解决方法

你在这里使用 xarray 究竟是为了什么？

由于 scipy 和 numpy 都适用于裸 numpy 数组，您可以简单地将数据作为 numpy 数组，并使用 ravel 获得相同内存的一维视图。

如果您还不熟悉副本和视图，请参阅 numpy 手册：https://numpy.org/devdocs/reference/arrays.ndarray.html?highlight=view

这是闪电般的速度：

ravelled = array.data.ravel()

只要它是一个一维数组，并且你只需要布尔索引，numpy 就可以了：

result = ravelled[ravelled > 0]