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 = []
        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
        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:
                    idxs.append((r, c))
    if separate:
        return (rs, cs)
        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
        if not non_zero:

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

2 Answers 2


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)
        result = idxs

    _neighbors2_memoized[args] = result
    return result

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


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:

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


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


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):
    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 = 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("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("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)

# 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
  • \$\begingroup\$ I'll implement memoizing this evening, I'm sure it'll help and it's a useful technique to remember. As for the set intersection, the depth parameter is domain determined and so will not increase beyond 70; I should have mentioned it. \$\endgroup\$
    – tsorn
    Sep 26, 2017 at 20:34
  • \$\begingroup\$ Are the rows,cols also guaranteed to be <7? \$\endgroup\$
    – aghast
    Sep 26, 2017 at 21:21
  • \$\begingroup\$ (because 49 booleans will fit into an int64, which opens a whole bunch of solutions involving bitwise ops) \$\endgroup\$
    – aghast
    Sep 26, 2017 at 21:32
  • \$\begingroup\$ Sure, we can make that assumption, so long as converting to/from array-form is fast, such that the rest of the code does not have to deal with a change in data structure. \$\endgroup\$
    – tsorn
    Sep 26, 2017 at 21:58
  • \$\begingroup\$ Including the encode, not including decode, it's about 4 times slower. There's just one change in the array, followed by get_free, followed by another change then get_free again and so on. So encode/decode is definitely not negligible, however I might consider changing all the other code considering how fast it is excluding encode. \$\endgroup\$
    – tsorn
    Sep 27, 2017 at 12:52

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

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.