使用numpy数组和递归函数优化Python

使用numpy数组和递归函数优化Python,python,numpy,recursion,numba,Python,Numpy,Recursion,Numba,我正在创建一个递归函数,该函数允许我为numpy数组中的任何gridcell查找最低的下坡邻居。这个函数将被调用很多次,所以我需要优化这个函数。 我使用numba来优化函数调用,但是我有点执着于在numba中指定函数定义和定义它 在我写的简单例子中,我可以使用@jit(nopython=True)直接优化装饰 但我认为函数定义可以再优化一点(完整测试程序见下文) 我主要关心的是将numpy数组传递给函数 有人对进一步优化有什么好主意吗 非常感谢 威廉 import os import sys i

我正在创建一个递归函数,该函数允许我为numpy数组中的任何gridcell查找最低的下坡邻居。这个函数将被调用很多次,所以我需要优化这个函数。 我使用numba来优化函数调用,但是我有点执着于在numba中指定函数定义和定义它

在我写的简单例子中,我可以使用@jit(nopython=True)直接优化装饰 但我认为函数定义可以再优化一点(完整测试程序见下文)

我主要关心的是将numpy数组传递给函数

有人对进一步优化有什么好主意吗

非常感谢 威廉

import os
import sys
import numpy as np
from time import time
from numba import jit

@jit("intc(intc,intc,intc)",nopython=True)
def rowcol2index(mapcols,row,col):
    return row * mapcols+col

@jit("intc(intc,intc)",nopython=True)
def index2row(mapcols, index):
    return index // mapcols 

@jit("intc(intc,intc)",nopython=True)
def index2col(mapcols, index):
    return index % mapcols
    

#Create a numpy array with random values
np.random.seed(1)
mydemarray = np.random.randint(1,100, size=(4000,500))
maprows=mydemarray.shape[0]
mapcols=mydemarray.shape[1]

demarrayflattened = mydemarray.flatten()

#@jit("intc(intc[:],intc,intc,intc)",nopython=True)
@jit(nopython=True)
def findlowestcandidate2(elevationmapflattened,maprows, mapcols,centralmapindex):
        myheight = elevationmapflattened[centralmapindex]
        myrow=index2row(mapcols, centralmapindex)
        mycol=index2col(mapcols, centralmapindex)
        if myrow==0 or myrow==maprows-1 or mycol==0 or mycol==mapcols-1:
            return -1

        lowestslope = 0;
        lowestslopeindex = centralmapindex
        neighbormapindex=rowcol2index(mapcols,myrow+1,mycol)
        neighborhead=elevationmapflattened[neighbormapindex]
        myslope = myheight-neighborhead
        if myslope > lowestslope:
            lowestslope=myslope
            lowestslopeindex = neighbormapindex
    
        neighbormapindex=rowcol2index(mapcols,myrow-1,mycol)
        neighborhead=elevationmapflattened[neighbormapindex]
        myslope = myheight-neighborhead
        if myslope > lowestslope:
            lowestslope=myslope
            lowestslopeindex = neighbormapindex
    
        neighbormapindex=rowcol2index(mapcols,myrow,mycol+1)
        neighborhead=elevationmapflattened[neighbormapindex]
        myslope = myheight-neighborhead
        if myslope > lowestslope:
            lowestslope=myslope
            lowestslopeindex = neighbormapindex
    
        neighbormapindex=rowcol2index(mapcols,myrow,mycol-1)
        neighborhead=elevationmapflattened[neighbormapindex]
        myslope = myheight-neighborhead
        if myslope > lowestslope:
            lowestslope=myslope
            lowestslopeindex = neighbormapindex
        
        if lowestslope > 0:
            return findlowestcandidate2(elevationmapflattened,maprows, mapcols,lowestslopeindex)
        else:
            return lowestslopeindex

starttime= time()
for row in range(maprows):
    for col in range(mapcols):
        elev_index=rowcol2index(mapcols,row,col)
        draintoindex = findlowestcandidate2(demarrayflattened,maprows, mapcols,elev_index)

endtime=time()
print("timer",endtime-starttime)


print("draintoindex",draintoindex)