Conway's game of life in python (numpy)

game-of-life.py

import numpy as np
from numpy.lib.shape_base import expand_dims

def add_walls(board) -> np.ndarray:
    '''add 2 rows & columns around the board
    @param board: numpy 2d array
    @return: a new board
    '''
    h, w = board.shape # get heigh & width
    # now sandwich it!
    s = np.vstack((np.zeros(w, dtype=int), board, np.zeros(w, dtype=int)))
    s = np.hstack((np.expand_dims(np.zeros(h+2, dtype=int), axis=1), s, expand_dims(np.zeros(h+2, dtype=int), axis=1)))
    return s

def neighbors(y_x, board) -> np.ndarray:
    '''get all neighboring cells (max 8)
    @param x_y: tuple of coord in board; x for horizontal axis, y for vertical axis, starting from top-left
    @param board: numpy 2d array
    @return: a flattened list of neighbor cells
    '''
    h, w = board.shape
    board_copy = board.copy() # create copy to not affect the board

    y, x = y_x
    nb: np.ndarray = board_copy[max(y-1, 0): min(y+2, h), max(x-1, 0):min(x+2, w)]

    board_copy[y][x] = -1 # filter out current cell
    f = nb.flatten()

    return f[f >= 0]

def trim(board) -> np.ndarray:
    '''remove zero lines from a numpy 2d array'''
    nz = np.nonzero(board)
    return board[nz[0].min():nz[0].max()+1, nz[1].min():nz[1].max()+1]

def get_generation(cells, generations = 1) -> list:
    board = np.array(cells)
    b_padded = board

    for _ in range(generations):
        # since this is unlimited version, we expand the board to include 'out-of-bound' cells
        b_padded = add_walls(b_padded)
        b_padded_copy = b_padded.copy()
        
        h, w = b_padded.shape
        for i in range(h):
            for j in range(w):
                nb = neighbors((i, j), b_padded)
                living_nb_count = np.count_nonzero(nb)
                if b_padded[i][j] == 1: # live cell
                    if living_nb_count < 2 or living_nb_count > 3:
                        b_padded_copy[i][j] = 0 # die
                else:
                    if living_nb_count == 3:
                        b_padded_copy[i][j] = 1 # become live

        b_padded = trim(b_padded_copy).copy()

    return b_padded.tolist()