Python 3 向量化嵌套 for 循环，其中内部循环取决于参数

Question

在将代码从 Fortran 移植到 python 的地球科学中，我看到了这些嵌套 for 循环的变体（有时是双重嵌套，有时是三重嵌套），我想对其进行矢量化（此处显示为最小可重现示例）

import numpy as np
import sys
import math
def main():
    t = np.arange(0,300)
    n1=7
    tc = test(n1,t)

def test(n1,t):
    n2 = int(2*t.size/(n1+1))
    print(n2)
    tChunked = np.zeros(shape = (n1,n2))
    for i in range(0,n1):
        istart = int(i*n2/2)
        for j in range(0,n2):
            tChunked[i,j] = t[istart+j]



  return  tChunked

main()

我尝试了什么？

我已经达到了消除 istart 并获得 j 以及使用外部加法获得 istart+j 的程度。但是如何使用索引 k 在一行中获取 2d Chunked 数组是我遇到的问题。

istart = np.linspace(0,math.ceil(n1*n2/2),num=n1,endpoint=False,dtype=np.int32)
jstart = np.linspace(0,n2,num=n2,endpoint=False,dtype=np.int32)

k = jstart[:,np.newaxis]+istart

Answer 1

如果索引是二维的，numpy 将输出一个二维数组。所以你只需这样做。

def test2(n1, t):
    n2 = int(2 * t.size / (n1 + 1))
    istart = np.linspace(0, math.ceil(n1 * n2 / 2), num=n1, endpoint=False, dtype=np.int32)
    jstart = np.linspace(0, n2, num=n2, endpoint=False, dtype=np.int32)
    k = istart[:, np.newaxis] + jstart  # Note: I switched i and j.

    tChunked = t[k]  # This creates an array of the same shape as k.

    return tChunked

Answer 2

如果你必须处理很多嵌套循环，也许解决方案是使用numba，因为它可以产生比原生 numpy 更好的性能。特别适用于您展示的非 python 函数。

一样容易：

from numba import njit

@njit
def test(n1,t):
    n2 = int(2*t.size/(n1+1))
    print(n2)
    tChunked = np.zeros(shape = (n1,n2))
    for i in range(0,n1):
        istart = int(i*n2/2)
        for j in range(0,n2):
            tChunked[i,j] = t[istart+j]

Answer 3

让我爱上 Python 的实际上是 NumPy，尤其是它令人惊叹的 indexing 和 indexing routines！

在 test_extra_crispy() 中，我们可以使用 zip() 将我们的鸭子（初始条件）排成一行，然后使用偏移量进行索引以对值块进行“移植”：

i_values = np.arange(7)
istarts = (i_values * n2 / 2).astype(int)
for i, istart in zip(i_values, istarts):
    tChunked[i, :n2] = t[istart:istart+n2]

也可以看看

• How to np.roll() faster?（以及该页面上的所有链接问题
• Best way to insert values of 3D array inside of another larger array（更高维度）

我们可以看到对于

t = np.arange(10000000)
n1 = 7

“extra crispy”比原来的快很多（91 vs 4246 ms），但只比 Zaero Divide's answer 的 test2() 快一点，考虑到它比我的蛮力处理更仔细地检查，这并不重要。

如果您需要在数组中寻址一个形状更随机的卷，您可以像这样使用索引：

array = np.array([[0, 0, 1, 0, 0], [0, 1, 0, 1, 0], [1, 0, 0, 0, 1], [0, 1, 0, 1, 0], [0, 0, 1, 0, 0]])
print(array)

给

[[0 0 1 0 0]
 [0 1 0 1 0]
 [1 0 0 0 1]
 [0 1 0 1 0]
 [0 0 1 0 0]]

我们可以这样得到 1 的索引：

i, j =  np.where(array == 1)
print(i)
print(j)

如果我们想从一个归零数组开始并通过 numpy 索引插入那些 1，就这样做

array = np.zeros((5, 5), dtype=int)
array[i, j] = 1

import numpy as np
import matplotlib.pyplot as plt
import time

def test_original(n1, t):
    n2 = int(2*t.size / (n1 + 1))
    tChunked = np.zeros(shape = (n1, n2))
    for i in range(n1):
        istart = int(i * n2 / 2)
        for j in range(0, n2):
            tChunked[i, j] = t[istart + j]
    return  tChunked

t = np.arange(1000000)
n1 = 70

t_start = time.process_time()

tc_original = test_original(n1, t)

print('original process time (ms)', round(1000*(time.process_time() - t_start), 3))
# print('tc_original.shape: ', tc_original.shape)

fig, ax = plt.subplots(1, 1)
for thing in tc_original:
    ax.plot(thing)
plt.show()

def test_extra_crispy(n1, t):
    n2 = int(2*t.size / (n1 + 1))
    tChunked = np.zeros(shape = (n1, n2))
    i_values = np.arange(7)
    istarts = (i_values * n2 / 2).astype(int)
    for i, istart in zip(i_values, istarts):
        tChunked[i, :n2] = t[istart:istart+n2]
    return  tChunked


t_start = time.process_time()

tc_extra_crispy = test_extra_crispy(n1, t)

print('extra crispy process time (ms)', round(1000*(time.process_time() - t_start), 3))
# print('tc_extra_crispy.shape: ', tc_extra_crispy.shape)

print('np.all(tc_extra_crispy == tc_original): ', np.all(tc_extra_crispy == tc_original))

import math

def test2(n1, t): # https://stackoverflow.com/a/72492815/3904031
    n2 = int(2 * t.size / (n1 + 1))
    istart = np.linspace(0, math.ceil(n1 * n2 / 2), num=n1, endpoint=False, dtype=np.int32)
    jstart = np.linspace(0, n2, num=n2, endpoint=False, dtype=np.int32)
    k = istart[:, np.newaxis] + jstart  # Note: I switched i and j.

    tChunked = t[k]  # This creates an array of the same shape as k.

    return tChunked

t_start = time.process_time()

tc_test2 = test2(n1, t)

print('test2 process time (ms)', round(1000*(time.process_time() - t_start), 3))
# print('tc_test2.shape: ', tc_test2.shape)

print('np.all(tc_test2 == tc_original): ', np.all(tc_test2 == tc_original))