Finding cells in a 3D grid with free neighbors

Question

I have a program, where the get_free function below has the most major computational cost by far, and I was wondering if there was a faster way to do it. Perhaps an all-numpy solution would be faster.

The task is to find the indices of the third dimension in a three-dimensional boolean array where the cell (as specified by the two first dimensions) and its neighbors (as given by the neighbors2 function) are 0. That's what get_free does.

Don't worry too much about understanding the neighbors2 function, it doesn't do much heavy lifting. The indices can be returned in two formats, as specified by the docstring, which can be helpful for some numpy-functions.

Any feedback, regarding code quality in general or for optimizing performance is appreciated.

import numpy as np
import timeit

rows, cols = 7, 7
depth = 70
state = np.random.choice([0, 1], size=(rows, cols, depth)).astype(bool)


def neighbors2(row, col, separate=False):
    """
    If 'separate' is True, return ([r1, r2, ...], [c1, c2, ...])
    if not, return [(r1, c1), (r2, c2), ...]

    Returns a list of indices of neighbors (in an hexogonal grid)
    within a radius of 2, not including self.
    """
    if separate:
        rs = []
        cs = []
    else:
        idxs = []

    r_low = max(0, row-2)
    r_hi = min(rows-1, row+2)
    c_low = max(0, col-2)
    c_hi = min(cols-1, col+2)
    if col % 2 == 0:
        cross1 = row-2
        cross2 = row+2
    else:
        cross1 = row+2
        cross2 = row-2
    for r in range(r_low, r_hi+1):
        for c in range(c_low, c_hi+1):
            if not ((r, c) == (row, col) or
                    (r, c) == (cross1, col-2) or
                    (r, c) == (cross1, col-1) or
                    (r, c) == (cross1, col+1) or
                    (r, c) == (cross1, col+2) or
                    (r, c) == (cross2, col-2) or
                    (r, c) == (cross2, col+2)):
                if separate:
                    rs.append(r)
                    cs.append(c)
                else:
                    idxs.append((r, c))
    if separate:
        return (rs, cs)
    else:
        return idxs


def get_free(cell):
    """
    Return the indices of a a cell that are 0 and
    where all its neighbors are 0 for the same depth
    """
    candidates = np.where(state[cell] == 0)[0]
    neighs = neighbors2(*cell, False)
    free = []
    # Exclude elements that have non-zero value in neighboring cells
    for candidate in candidates:
        non_zero = False
        for neigh in neighs:
            if state[neigh][candidate]:
                non_zero = True
                break
        if not non_zero:
            free.append(candidate)


print(timeit.timeit("get_free((4, 4))", number=100000,
      setup="from __main__ import get_free"))

Answer 1

It seems like your neighbors2 is a pure function. I think you could speed things up by memoizing the results:

def neighbors2(row, col, separate=False):
    global _neighbors2_memoized

    args = (row, col, separate)

    if not _neighbors2_memoized:
        _neighbors2_memoized = { } # Start with empty dict
    elif row_col_sep in _neighbors2_memoized:
        return _neighbors2_memoized[args]

    # rest of your code

    if separate:
        result = (rs, cs)
    else:
        result = idxs

    _neighbors2_memoized[args] = result
    return result

Of course, a little searching will yield hundreds of memoization decorators and the like for python.

Update

Also, I think what you are doing is looking for "layers" or "levels" at which the given cell has all its neighbors empty.

If so, consider recoding your check to use intersect1d(). The neighbors are going to be constant - it's the same neighbors at each level.

So make a similar candidates array for each neighbor, and treat it like a set: "this is the set of levels where neighbor (ij) is zero".

Then, compute the intersection of that set with the candidates set (which is just the set of levels where cell is zero).

This is basically a large and: cell is zero AND neighbor(ij) is zero AND ...

You can execute len(neighbors) numpy operations instead of looping in python.

def get_free_setwise(cell):
    """
    Return the indices of a a cell that are 0 and
    where all its neighbors are 0 for the same depth
    """
    neighbors = neighbors2(*cell, False)

    candidates = np.where(state[cell] == 0)[0]

    for n in neighbors:
        candidates = np.intersect1d(candidates, np.where(state[n] == 0)[0])
        if len(candidates) == 0:
            break

    return candidates


print(timeit.timeit("get_free((4, 4))", number=10000,
      setup="from __main__ import get_free"))
print(timeit.timeit("get_free_setwise((4, 4))", number=10000,
      setup="from __main__ import get_free_setwise"))

Output of this was (note: number=10,000):

0.8349439720041119
6.3053770850092405

About seven times slower. However, when I raised depth to 700, the times were almost equal, and when I raised depth to 7000, the times were

3.6607714939891594
2.716483567011892

So the numpy-centric version does better as depth increases (FWIW).

Update 2

Using bitwise ops I was able to get better performance:

# Neighborhood pre-computation for get_free5. This is a one-time thing - the
# neighborhoods will never change as long as the rows,cols stay the same.
# NOTE however that these arrays are TRANSPOSED, because that's what free5
# does.

neighborhood = []
for r in range(rows):
    neighborhood.append([])
    for c in range(cols):
        hood = np.zeros((rows, cols), dtype=np.bool_)
        hood[r,c] = 1
        for nb in neighbors2(r, c):
            hood[nb] = 1
        hood = hood.T # Transpose!
        # Now hood is all 1's for neighbors. Do this same computation
        # in the same fashion that the state-updater does.
        bytestr = np.packbits(hood).tobytes() + b'\0'
        uint = np.frombuffer(bytestr, dtype=np.uint64)
        neighborhood[r].append(uint[0])

neighborhood = np.array(neighborhood, dtype=np.uint64)
# Neighborhood[r,c] is a bitmask with all neighbors + (r,c) set to 1
# NOTE: +(r,c) - the neighborhood explicitly includes the center point,
# which is not true for the neighbors function. I am using this to test for
# the center ALSO being zero at the same time - just one pass through the
# state array.

def get_free5(cell, *, verbose=False):
    """
    Return the indices of a a cell that are 0 and
    where all its neighbors are 0 for the same depth
    """
    result = list(np.nonzero(state_asbits & neighborhood[cell] == 0)[0])

    # This if statement is used in my self-test code. Delete all this code
    # for production use.
    if verbose:
        orig = get_free(cell)
        missing = [n for n in orig if n not in result]
        extra = [n for n in result if n not in orig]
        print("Missing from result, found in orig:")
        for n in missing:
            print("orig: (transposed)", n)
            print(etats[n])
            print("bits: {:64b}".format(state_asbits[n]))
            print("mask: {:64b}".format(neighborhood[cell]))

        print("Found in result, not found in orig:")
        for n in extra:
            print("orig: (transposed)", n)
            print(etats[n])
            print("bits: {:64b}".format(state_asbits[n]))
            print("mask: {:64b}".format(neighborhood[cell]))

    return result

# Setup for get_free5(). This code has to be run whenever the state[]
# array changes or is computed. It's an alternate encoding of the state array
# so if you don't need state[] for anything except this check, you can get
# rid of it. Otherwise, you'll have to call this before you start calling
# get_free5()

# Transpose, to get all the values at the same depth in a contiguous area.
etats = state.T
bitplane = []
for d in range(state.shape[-1]):
    #bytestr = np.packbits(etats[d]).tobytes()
    #if len(bytestr) % 8 != 0:
    #    bytestr += b'\00' * (8 - len(bytestr) % 8)

    # Hard-coding for 7x7 - append one extra byte
    bytestr = np.packbits(etats[d]).tobytes() + b'\0'
    uint = np.frombuffer(bytestr, dtype=np.uint64)
    bitplane.append(uint[0])

# state_asbits[n] is an array of 64-bit numbers, holding bits for each of the
# 49 cells at depth=n. If the value in state[r,c,n] is 1, then the bit at
# state_asbits[n].get_bit(r,c) == 1. These values can be trivially checked
# using bitwise AND operations, with the mask values in the neighbors array.
state_asbits = np.array(bitplane)

This code runs faster than the original version, but that doesn't include the computations I have to do for the state_asbits array. I think that's pretty cheap, but I don't know how often you are updating it.

$ python test.py
original code 8.314752619000501
free5: bitwise 1.1757051760068862

Answer 2

One improvement, with bitwise or:

 def get_free2(cell):
    """
    Return the indices of a a cell that are 0 and
    where all its neighbors are 0 for the same depth
    """
    neighs = neighbors2(*cell, False)
    f = np.bitwise_or(state[cell], state[neighs[0]])
    for n in neighs[1:]:
        f = np.bitwise_or(f, state[n])
    free = np.where(f == 0)[0]
    return free

The results, 200000 iterations:

get_free (original) 8.616401553998003
get_free2 (this post) 6.426983985002153
get_free5 (@Hastings, including encode) 37.249181182996836