python - 避免 3d 矩阵运算的嵌套 for 循环的一般选项

Question

在我对嵌套循环进行研究之后，我想提出这个一般性问题:我必须在具有例如 shape 的 3d 矩阵(多波段图像文件)上运行各种函数。 (2500, 100, 200) 或 (bands, rows, cols)。

我处理函数的一般(缓慢)过程如下所示:

import numpy as np
import time

x = np.random.random_sample([2500, 100, 200])

#  1. & 2. derivative example

def derivate(matrix):
    d1 = np.empty(shape=np.shape(matrix))
    d2 = np.empty(shape=np.shape(matrix))
    for row in range(matrix.shape[1]):
        for col in range(matrix.shape[2]):
            d1[:, row, col] = np.gradient(matrix[:, row, col])
            d2[:, row, col] = np.gradient(d1[:, row, col])
    return d1, d2


t1 = time.perf_counter()
y, z = derivate(x)
t2 = time.perf_counter()
print(t2 - t1)

>>> 4.38 seconds

由于我的项目环境中已经存在许多这种风格的函数，所以我正在寻找最简单的解决方案来大幅加快这些嵌套 for 循环函数的速度。

我已经看到这里的许多问题都是针对嵌套 for 循环的，但我没有发现任何通用且可转移的解决方案。如果能提供许多可用的方法，我将不胜感激，干杯!

Answer 1

你可以像这样实现它:

import numpy as np

def derivate_vec(matrix):
    d1 = np.gradient(matrix, axis=0)
    d2 = np.gradient(d1, axis=0)
    return d1, d2

快速比较:

import numpy as np
import time

def derivate(matrix):
    d1 = np.empty(shape=np.shape(matrix))
    d2 = np.empty(shape=np.shape(matrix))
    for row in range(matrix.shape[1]):
        for col in range(matrix.shape[2]):
            d1[:, row, col] = np.gradient(matrix[:, row, col])
            d2[:, row, col] = np.gradient(d1[:, row, col])
    return d1, d2

def derivate_vec(matrix):
    d1 = np.gradient(matrix, axis=0)
    d2 = np.gradient(d1, axis=0)
    return d1, d2

# Input data
np.random.seed(0)
x = np.random.random_sample([2500, 100, 200])

# Check results
d11, d12 = derivate(x)
d21, d22 = derivate_vec(x)
print(np.allclose(d11, d21))
# True
print(np.allclose(d12, d22))
# True

# Measure time
t1 = time.perf_counter()
derivate(x)
t2 = time.perf_counter()
derivate_vec(x)
t3 = time.perf_counter()
print(t2 - t1)
# 3.6777221
print(t3 - t2)
# 0.9414466999999997