递归地在四叉树中找到包含给定点的单元格

Question

作为一个更大程序的一部分，我有这个函数，它有 4 个参数：

r

、

z

、

c

和

grid

。

r

和

z

代表点的坐标。

c

是我们首先检查该点是否在此单元格中的父单元格。如果该点在单元格中并且该单元格没有子节点，则返回单元格和

found

的布尔值，否则，我们对其所有子节点（及其子节点等）重复该过程。

grid

是一个列表，其中包含四叉树中的所有单元格。每个单元格

c

是一个结构化的 NumPy 数组。

这是代码

def locate_photon_cell_by_tree(r: float, z: float, c, grid):
    """
    Given r and z and start from c, find out a cell containing (r,z).
    """
    NMAX = 1e6

    found = False
    cout = c

    for j in range(int(NMAX)):
        if ((cout['xmin'][0]**2 <= r <= cout['xmax'][0]**2) and 
            (cout['ymin'][0]    <= z <= cout['ymax'][0])):
            if (cout['children'][0][0] == 0):
                found = True
                return cout, found
            
            flag = True
            non_zero_indices = np.nonzero(cout['children'][0])[0]
            child_indices = cout['children'][0][non_zero_indices]

            for index in child_indices:
                if ((grid[index]['xmin'][0]**2 <= r <= grid[index]['xmax'][0]**2) and 
                    (grid[index]['ymin'][0]    <= z <= grid[index]['ymax'][0])):
                    cout = grid[index]
                    flag = False
                    break
            
            if (flag):
                # cout = None
                return cout, found
        else:
            if cout['parent'][0] is not None:
                cout = grid[cout['parent'][0]]
            else:
                # cout = None
                return cout, found
    
    # cout = None
    return cout, found

代码有效，但我正在寻找优化。我在非 python 模式下尝试了 Numba，但每次在第一次迭代中我都会收到以下错误（循环无限期地卡住），这就是我在做

KeyboardInterrupt

:

时得到的结果

SystemError: CPUDispatcher(<function locate_photon_cell_by_tree at 0x7fa2cd332b00>) returned NULL without setting an exception

这很奇怪，因为该函数不会在任何地方返回 None/Null（我已经对这些部分进行了评论）并且它在没有 JIT 的情况下也能完美运行。

我正在寻找如何在没有 JIT 的情况下优化它，或者如何使其与 JIT 兼容。任何帮助将不胜感激。