I have a tree-like structure called
grid. I have designed it as a structured numpy array. Each element of
grid is a tree-node. Each node itself is a structured numpy array, with fields that describe its bounding box (
ymax) in 2-D space. Each node has an ID which is basically the index of that node in
grid. Each node has a field
parent which is the ID of its parent. Each node has a field called
children which is a numpy array of integers containing the IDs of its children.
nChildren obviously denotes number of children that node has (this tree is not a strict binary/quadtree). Root node has parent
ID = -1 and
-99999 is just a flag for when I want to return an integer instead of
Given below is a function, whose arguments are
(r, z), a spatial-point in 2-D space,
c_index which is the node we start with, and of course, the whole
grid object. Task is to find the smallest node given
(r, z) and a starting
c_index that contains the point (speaking of that, if anybody has an idea why we use squares of the
xmax when checking if the point is in the cell, please tell me).
I have done profiling of the function. Without using Numba-JIT with
nopython=True, the function takes around ~75 seconds for around ~180K calls. With Numba-JIT with
nopython=True, it takes around ~14 seconds for around the same number of calls. That is good and all, but I desire a bit more performance as you can tell it is called an obscenely large number of times. The problem is that these results are for a test run with a small number of parameters than I will be actually using. When the codebase will be actually deployed, this function will be called probably around a million times, so times add up.
Here is the function:
def locate_photon_cell_mirror(r, z, c_index, grid): NMAX = 1000000 found = False cout_index = c_index abs_z = np.abs(z) for j in range(NMAX): cout = grid[cout_index] if (cout['xmin']**2 <= r and cout['xmax']**2 >= r and cout['ymin'] <= abs_z and cout['ymax'] >= abs_z): if (cout['nChildren'] == 0): found = True return cout_index, found flag = True for i in range(cout['nChildren']): child_cell = grid[cout['children'][i]] if (child_cell['xmin']**2 <= r and child_cell['xmax']**2 >= r and child_cell['ymin'] <= abs_z and child_cell['ymax'] >= abs_z): cout_index = cout['children'][i] flag = False break if (flag): cout_index = -999999 return cout_index, found else: cout_parent = cout['parent'] if cout_parent != -1: cout_index = cout_parent else: cout_index = -999999 return cout_index, found cout_index = -999999 return cout_index, found
As Numba-JIT is not enough for me, I'm looking for a faster algorithm to achieve this. If I understand it correctly, the problem is basically to do this: given a point in a plane and a rectangle, what is the smallest rectangle that contains the point? (as all child nodes will be part of the parent node, as is the case in 2-D space-partitioning trees).