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

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

• If you don’t know why you’re using the square of xmax, why did you write the code this way? I assume you wrote the code? Jul 19 at 1:03
• Not really, I adapted this code from a Fortran code base. I was confused about squares as well, but that's what they had used. @CrisLuengo Jul 19 at 1:05
• Which data structure are you actually implementing? There may be an existing optimized Python library for it; if it fits your requirements, it could save you a lot of programming time and computation time. It looks like it could be an R-tree — have you looked at rtree.readthedocs.io/en/latest? For KD-trees, you have docs.scipy.org/doc/scipy/reference/generated/…. Both are backed by C code.
– Pkkm
Jul 19 at 20:36
• Also, you haven't shown us the code that creates this data. What type is cout?
– Pkkm
Jul 19 at 20:44
• @Pkkm The thing is, I did not come up with this data structure design, I merely adapted it from a Fortran code base. I have looked into R-trees and k-d trees, but that will require redesigning the Fortran code at the core level, which I don't have the time for, unfortunately. Complexity of that code base is the primary reason. Jul 19 at 22:33

# 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.

• Sorry about the bad coding practices, I am aware of them and will correct them promptly. I tried storing squares of r co-ordinates, but I saw no performance improvement whatsoever. I would like to implement caching, but at least @lru_cache is not Numba-compatible, so I will need to spend some time figuring that out. Voronoi is an good suggestion, but does I see it does not take advantage of the tree structure? I mean the parent-child relationship. Will also look into PostGIS, thank for all the suggestions. I will try implementing them and return soon. Hopefully something works. Jul 19 at 2:07
• I tried you methods but after some profiling I have come to the conclusion that I don't think I will be able to get this thing any faster. With Numba-JIT, one call of it takes around ~10^-5 seconds. The problem is that this function is just called too many times, adding the times up. Jul 19 at 14:27
• Well, maybe you're calling it "too many times". Consider batching up those (r, z) locations and asking the question about a batch of a thousand, rather than asking a thousand individual questions. Which would be best addressed by merge sort. Or send those thousand locations to PostGIS and let it do the merge sort against its 2-D index.
– J_H
Jul 19 at 14:43
• The larger program which calls this function multiple times, is sequential. It is basically a photon random walk algorithm in 2-D grid. So this function which I have shared is basically called everytime a photon walks. (r, z) are the photon's location. So I am not sure how to make batches of (r, z). Also because of the sequential nature of program, I cannot parallelize it. Each photon random walk of course depends upon it's previous position. Further, I simulate 4 ~million photons, and each photon's wavelength is calculated using previous photon's wavelength. Jul 19 at 22:50

# 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.

• cout['children'] unfortunately does not have the same size as specified in cout['nChildren']. The reason for this, is again, tied to the existing Fortran code base and redesigning things around it would be very difficult. And yes, that comparison with r and z was originally inside a function. I can do that. I will get back to you soon after I try you suggestion in the Code Logic section Jul 19 at 22:39
• Also as I said, faster lookups than they currently are might just not be possible (with Numba-JIT). This function is at the bottom of a larger program, and it is the one which is called the maximum number of times. The original Fortran code takes around 40 minutes for a average run (of the larger program). My Python equivalent with Numba-JIT takes around 60-70 minutes. The goal was to match the Fortran speed or best it. Jul 19 at 22:42
• @matts Fortran is hard to beat for speed. Numba is good, but is not going to beat Fortran, especially with the overhead of Python object storage. Jul 19 at 22:49
• I guess so :(. Well, this project will stay in Python, But do you think Julia could do any better? Jul 19 at 22:51
• I don’t know, I have little experience with Julia. Questions I see on Stack Overflow about slow Julia code seem to indicate that it’s only fast if you program in just the right way. Not sure how difficult it is to learn to do that. Jul 19 at 22:52