使用并行处理从多维数组中集成维度

Question

我希望找到一些聪明的方法来解决我一直在努力解决的并行处理问题。基本上，我正在处理 20,160 个大小为 (72,35,25,20) 的多维数组。目前，我通过简单地在嵌套 for 循环中进行梯形积分来积分尺寸为 72 的维度。我的最终目标是获得一个大小为 (20160,35,25,20) 的输出数组。

for idx,filename in enumerate(filenames):
    #Read NetCDF Data File as 'raw_data'
    flux=raw_data['FluxHydrogen'][:]   #This is size (72,35,25,20)
    PA=raw_data['PitchAngleGrid'][:]   #This is size (72)
    for i in range(35):
        for j in range(25):
            for k in range(20):
                dir_flux=flux[:,i,j,k]
                omni_flux=np.trapz(dir_flux*np.sin(PA),PA)
                data[idx,i,j,k]=omni_flux   #This will have size (20160,35,25,20)

我相信在嵌套的 for 循环中实现较低的并行化是最有益的，但似乎无法弄清楚如何实现。我搜索了一些常见问题，但 [我发现] 没有一个问题提供足够的洞察力来说明如何实现共享内存、将多维数组传递到池和/或重塑结果数组。任何帮助或见解将不胜感激。

Answer 1

您可以使用麻波所以要大大加快这段代码的速度。 Numba 是一个 JIT 编译器，能够将基于 Numpy 的代码编译为快速本机代码（因此循环不是问题，事实上，在 Numba 中使用循环是个好主意）。

首先要做的是预计算np.sin(PA)一次从而避免重复计算。然后，可以使用 for 循环计算 dir_flux * np.sin(PA)，并将结果存储在预分配数组所以不要执行数百万昂贵的小数组分配.外循环可以使用多线程使用 prange 和 Numba 标志 parallel=True。假设输入值不特殊（如 NaN 或 Inf 或非常非常小：请参阅 subnormal numbers），可以使用标志 fastmath=True 进一步加速它。

虽然理论上这应该足以获得快速代码，但 np.trapz 的当前实现效率不高，因为它执行昂贵的分配。人们可以很容易地重写函数，这样就不会分配任何额外的数组。

下面是结果代码：

import numpy as np
import numba as nb

@nb.njit('(float64[::1], float64[::1])')
def trapz(y, x):
    s = 0.0
    for i in range(x.size-1):
        dx = x[i+1] - x[i]
        dy = y[i] + y[i+1]
        s += dx * dy
    return s * 0.5

@nb.njit('(float64[:,:,:,:], float64[:])', parallel=True)
def compute(flux, PA):
    sl, si, sj, sk = flux.shape
    assert sl == PA.size
    data = np.empty((si, sj, sk))
    flattenPA = np.ascontiguousarray(PA)
    sinPA = np.sin(flattenPA)
    for i in nb.prange(si):
        tmp = np.empty(sl)
        for j in range(sj):
            for k in range(sk):
                dir_flux = flux[:, i, j, k]
                for l in range(sl):
                    tmp[l] = dir_flux[l] * sinPA[l]
                omni_flux = trapz(tmp, flattenPA)
                data[i, j, k] = omni_flux
    return data

for idx,filename in enumerate(filenames):
    # Read NetCDF Data File as 'raw_data'
    flux=raw_data['FluxHydrogen'][:]   #This is size (72,35,25,20)
    PA=raw_data['PitchAngleGrid'][:]   #This is size (72)
    data[idx] = compute(flux, PA)

注意flux 和PA 必须是 Numpy 数组。另请注意，只要 len(PA) 相对较小且 np.std(PA) 不大，trapz 就是准确的。否则，成对求和甚至（偏执狂）Kahan summation 应该会有所帮助（注意 Numpy 使用成对求和）。实际上，随机正态数的结果是相同的。

---

进一步优化

通过使 flux 访问更连续，可以使代码更快。可以使用有效的转置来做到这一点（Numpy 的转置效率不高）。然而，这在 4D 阵列上并不简单。另一种解决方案是在 k 维度的整行上计算 trapz 操作。这使得计算非常高效并且在我的机器上几乎受内存限制。这是代码：

@nb.njit('(float64[:,:,:,:], float64[:])', fastmath=True, parallel=True)
def compute(flux, PA):
    sl, si, sj, sk = flux.shape
    assert sl == PA.size
    data = np.empty((si, sj, sk))
    sinPA = np.sin(PA)
    premultPA = PA * 0.5
    for i in nb.prange(si):
        for j in range(sj):
            dir_flux = flux[:, i, j, :]
            data[i, j, :].fill(0.0)
            for l in range(sl-1):
                dx = premultPA[l+1] - premultPA[l]
                fact1 = dx * sinPA[l]
                fact2 = dx * sinPA[l+1]
                for k in range(sk):
                    data[i, j, k] += fact1 * dir_flux[l, k] + fact2 * dir_flux[l+1, k]
    return data

请注意，预乘会使计算的精度略有下降。

---

结果

以下是我的 6 核机器（i5-9600KF 处理器）上随机数的结果（如 @DominikStańczak 使用的）：

Initial sequential solution:                       193.14 ms  (±  1.8 ms)
DominikStańczak sequential vectorized solution:      8.68 ms  (± 48.1 µs)
Numba parallel solution without fastmath:            0.48 ms  (±  6.7 µs)
Numba parallel solution without fastmath:            0.38 ms  (±  9.5 µs)
Best Numba solution (with fastmath):                 0.32 ms  (±  5.2 µs)
Optimal lower-bound execution:                       0.24 ms  (RAM bandwidth saturation)

因此，最快的 Numba 版本比 @DominikStańczak 的（顺序）版本快 27 倍，并且比最初的快 604 倍.这几乎是最优的。

Answer 2

作为第一步，让我们向量化代码本身。我现在只打算在每个文件的基础上处理这个问题，向您展示如何摆脱嵌套的 for 循环：

shape = (72, 35, 25, 20)

flux = np.random.normal(size=shape)
PA = np.random.normal(size=shape[0])

现在，为您的实施计时，稍微重写一下：

%%timeit
data = np.empty(shape[1:])
for i in range(shape[1]):
    for j in range(shape[2]):
        for k in range(shape[3]):
            dir_flux=flux[:,i,j,k]
            omni_flux=np.trapz(dir_flux*np.sin(PA),PA)
            data[i,j,k]=omni_flux
# 211 ms ± 4.86 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

我的第一个想法是将 sin 拉出 for 循环，因为不需要每次都重新计算它，但这让我达到了 10ms 的最高点。但是，如果我们不使用 for 循环，而是通过广播使用简单的 numpy 向量化，将 sin_PA 转换为 (72, 1, 1, 1) 形数组：

%%timeit
sin_PA = np.sin(PA).reshape(-1, 1, 1, 1)
data = np.trapz(flux * sin_PA, x=PA, axis=0)
# 9.03 ms ± 554 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

这是一个 20 倍的速度，没有什么可嘲笑的。我估计您的所有文件大约需要三分钟。您还可以使用 np.allclose 来验证结果是否符合浮点错误。

---

如果你之后仍然需要并行化，我会使用dask.array 事实上，如果你的数据在 netcdf4 文件中，我会使用xarray（这对多维数据有帮助）来读取它们，然后在后端的 with Dask enabled 上运行 trapz 计算。我认为这是在这种情况下实现简单多处理的最简单方法。这是一个快速草图：

import xarray
from Dask.distributed import Client
client = Client()

file_data = xarray.open_mfdataset(filenames, parallel=True)

# massage the data a little, probably

flux = file_data["FluxHydrogen"]
PA = file_data["PitchAngleGrid"]
integrand = flux * np.sin(PA)   # most element-wise numpy operations work on xarray ones or Dask based ones without a hitch
data = integrand.integrate(coord="PitchAngle")   # or some such name for the dimension you're integrating out