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 tChunked 数组是我卡住的地方。

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 更好的性能。专门用于 non-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]

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

Python 3 vectorizing nested for loop where inner loop depends on parameter

for-loop

numpy

vectorization

python-3.8

array-broadcasting