CODEWITHSUNDEEP
pythonnumbajitquadtree
peakcipher
peakcipher

Reputation: 51

Finding the cell containing a given point in quadtree recursively

As part of a larger program, I have this function which takes 4 arguments: r, z, c, and grid. r and z represents the coordinates of the point. c is the parent cell for which we first make the check that if the point is in this cell. If the point is in the cell and if this cell has no children, we return the cell and a boolean of found, otherwise, we repeat the process for all its children (and their children and so on). grid is a list which contains all the cells in the quadtree. Each cell c is a structured NumPy array.

Here is the code

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

The code works, but I am looking for optimization. I tried Numba in no-python mode but I am getting the following error every time in the first iteration (the loop gets stuck indefinitely), This is what I get when do KeyboardInterrupt:

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

This is weird because the function is not returning None/Null anywhere (I have commented on those parts) and it works perfectly without JIT.

I am looking for ways how can I optimize it without JIT, or how can I make it JIT-compatible. Any help will be appreciated.

Upvotes: 1

Views: 39

Answers (0)

Related Questions