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

In the geosciensces while porting code from Fortran to python I see variations of these nested for loops(sometimes double nested and sometimes triple nested) that I would like to vectorize(shown here as an minimum reproducible example)

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()

What have I tried ?

I have gotten as far as elminating the istart and getting j and using outer addition to get istart j. But how do I use the index k to get a 2d tChunked array in a single line is where I am stuck.

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

CodePudding user response：

numpy will output a 2D array if the index is 2D. So you simply do this.

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

CodePudding user response：

If you have to deal with a lot of nested loops, maybe the solution is to use numba as it can result in better performance than native numpy. Specially for non-python functions like the one you showed.

As easy as:

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]