Optimizing a node search method given a 2-D spatial point

Question

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 (xmin, xmax, ymin, 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 None.

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 xmin and 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).

Answer 1

flags

I appreciate the clarity of this, thank you.

                found = True
                return cout_index, found

That is, I know exactly what the return tuple means. It's much clearer than return cout_index, True.

Consider creating a namedtuple so you can concisely return ... , found=True). That way things are very clear to the caller, as well, for example during a debugging session.

Here is the only other place we assign that variable:

    found = False

There are three returns. Maybe we don't need that variable, and could have three literal False values in return statements?

This is not a terrific name:

            flag = True

It has a related meaning. Maybe it wants to claim the "found" name? Or maybe we could come up with a more informative identifier.

The whole for i loop is really begging for Extract Helper. Then each boolean would operate within its own scope, reducing the coupling which prevented me from immediately grasping Author's Intent behind these loops.

---

sentinel

Please define a symbolic name for the -999999 sentinel you use in a few places.

Consider defining another name for -1, which is used with cout_parent.

---

extra parens

Extra ( ) parentheses are Good in those 4-line conditionals. Thank you for avoiding \ backwhack line continuations in that way.

But this isn't C. Avoid saying things like this:

            if (cout['nChildren'] == 0):

            if (flag):

A simple if flag: suffices, and is definitely preferrable.

I know, I know, "it's a lotta rules!" But hey, there's no need to remember all of them. That's what the machine is for. Just run "$ black *.py" and you're done, it's all fixed up yet still means the same. The whole idea is to prevent trivial spelling nits like this from even surfacing during code review.

---

range for bbox

First, you're storing {xmin, xmax} values but you might find it more convenient to store the square of those values, so you square just once. Or compute r_sqrt just once on the way in, similar to z_abs.

Second, you're storing a pair of numbers which are only used in comparisons. You might find it more convenient to store cout['x_range'] = range(xmin**2, xmax**2) so that later you can just ask if r in cout['x_range']. (I know, there's a closed vs half-open interval detail, but we can finesse that.)

It would be helpful for the source code to cite a reference, perhaps an URL, so we understand where the (r, z) terminology came from. Initially I was mentally reading that as (radius, theta), until I saw it doesn't apply here.

---

caching, discretizing

You didn't describe what distributions we see over r and z values. Let's imagine that both range over the unit interval, and that some values are more common than others, perhaps due to a normal distribution.

I'm going to assume

1. (r, z) completely specifies a grid location since c_index is just a starting hint,
2. there's a limited number of cell mirror locations, maybe a thousand, and
3. values are high-resolution, say six or nine decimal places.

Given that, you might want to decorate a helper with @lru_cache. The idea would be to discretize values to something coarser but "good enough", say three or four places, and call the decorated helper with those values. Fix r and z for the moment. Then a call with (r + 2ϵ, z) might return index 7. And a subsequent (r + ϵ, z) call would enjoy a cache hit, also returning 7, without nested loops.

You know the spacing between mirror cell locations much better than I do, and better than that LRU utility. Feel free to implement your own simple caching approach, which would be free to not cache (r, z) values that are "too close to the border", forcing careful evaluation of which cell to return.

---

datastructure, Voronoi

There are plenty of Voronoi diagram implementations available on pypi. You might use Voronoi cell boundaries to map (r, z) to nearest cell.

Or brute force it. I hear that memory is cheap. Allocate a numpy array, impose a sufficiently fine discrete grid on your space, and exhaustively calculate "nearest mirror cell" for every point. This can be cheap to do: iterate over each mirror, fill in "I am nearest!" for some limited radius of grid points, then go to next mirror and repeat, overwriting in the case where that mirror turns out to be nearest.

---

datastructure, spatial index

Why does this even need to be your problem? PostGIS and similar products have already solved it. Use postgres queries to rapidly identify containing shapes. Note that a 2-D index is very different from having a pair of indexes, one on the r column and another on the z column.

Consider annotating those shapes with a size or area attribute, so you can rapidly identify the smallest containing shape.

Answer 2

Faster lookups

To speed up this code, I suggest you use an existing spatial index rather than reinventing the wheel (even if you didn’t invent this particular wheel). Look up, for example, R-trees or R*-trees.

Code review

Variable names

I’m aware that you translated someone else’s Fortran code, and likely took over variable names as-is. But I think the choice of names doesn’t aid reading the code.

Your members xmin and xmax are compared to the square root of r, and ymin and ymax are compared to the absolute value of z. We don’t know the meaning of r and z, but setting x = sqrt(r) and y = abs(z) at the top of the function would certainly make things less jarring.

You have cout for the current cell, whose index is cout_index. You also have child_cell (the one good name here), and cout_parent, which is not a cell but an index. I would suggest current_cell and current_index, and parent_index.

Code repetition

The comparison

        if (cout['xmin']**2 <= r and
            cout['xmax']**2 >= r and 
            cout['ymin']    <= abs_z and
            cout['ymax']    >= abs_z):

is repeated twice. This should be in a function. Don’t repeat yourself.

Code logic

It is hard to follow what is going on. As I read it, you find the leaf node containing a particular set of coordinates, and you search starting at a given node. If the node doesn’t contain the point, you go up until you find a node that does, then you go back down through its children. If the root node doesn’t contain the point, or if a node contains a point but none of its children do, then your search fails.

And you do all of this in no more than NMAX = 1000000 steps. Given this is a tree, making 1 million steps would mean the tree has more nodes in it than could possibly fit in memory (unless it is really poorly balanced with most nodes having only one child). Remember that the depth of the tree is O(log(N)). So this limit is not really necessary or useful.

I think the way to salvage the readability of this logic is to separate out for example the portion that finds the child node containing the coordinates into a separate function. You could also separate the loop where you go up from the one where you go down. There is no need to put these into the same loop.

In fact, this separation would speed up the code by not checking cout for inclusion after the first step down into a child, because you already verified the inclusion when you made that step.

while not included and parent exists:
    current = parent

while True:
    if leaf:
        return success
    current = find_child(current)
    if not current (ie failed to find a child):
        return failure

Having the find_child() here avoids the need for the flag variable.

Return value

You return the tuple (cout_index, found). But the first is always -999999 if the second is false, and a valid index if true. Why do you need two values? Just return an index, if it’s negative, you didn’t succeed.

Loop over range

This code:

            for i in range(cout['nChildren']):
                child_cell = grid[cout['children'][i]]
                if ...:
                    cout_index = cout['children'][i]

loops over indices, but only uses those indices to get the elements of cout['children']. It is better to directly loop over cout['children']. The loop would look like this:

            for child_index in cout['children']:
                child_cell = grid[child_index]
                if ...:
                    cout_index = child_index

I am assuming here that cout['children'] has always exactly cout['nChildren'] elements, and doesn't have a bunch of unused elements at the end. If so, cout['nChildren'] is redundant and can be left out. If not, then you need to change how cout['children'] is defined, because it makes the code simpler.