Cellular Automata

import numpy as np
from scipy.ndimage.filters import convolve
from skimage.morphology import disk
from matplotlib import pyplot as plt

from dataclasses import dataclass

@dataclass
class Parameters:
    birth: list
    survive: list
    neighborhood_width: int = 3
    neighborhood_height: int = 3
    kernel_shape: str = "square"
    
    def get_kernel(self):
        if self.kernel_shape == "square":
            # neighbors are in a grid of height*width around a cell, not including the cell
            self.kernel = np.ones((self.neighborhood_height, self.neighborhood_width))
            self.kernel[self.neighborhood_height//2, self.neighborhood_width//2] = 0
        elif self.kernel_shape == "disk":
            # neighbors are in a circular disk with radius height/2 (=width/2) around a cell, not including the cell
            assert self.neighborhood_height == self.neighborhood_width
            self.kernel = disk(self.neighborhood_height//2)
        else:
            self.kernel = np.ones(self.neighborhood_height, self.neighborhood_width)

class CellularAutomata:
    def __init__(self, init_world: np.ndarray, parameters: Parameters, cmap: str="gray"):
        self.height, self.width = init_world.shape
        self.world = init_world
        self.memory = np.array(self.world, dtype=np.float64)
        self.cmap = cmap
        self.set_parameters(parameters)

    def set_parameters(self, parameters):
        self.parameters = parameters
        self.parameters.get_kernel()
        
    def procreate(self):
        """
        calculates the next round, params same as iterate
        """
        new_world = np.array(self.world, dtype=np.float64)
        # finds the sum of the values of the neighbors
        neighbors = convolve(new_world, self.parameters.kernel, mode='constant')
        # checks which dead cells have enough live neighbors to come back alive in the next round
        birth_filter = np.logical_or.reduce([(b0 - 1 < neighbors) & (neighbors < b1 + 1) for b0, b1 in self.parameters.birth]) & (new_world < 0.5)
        # checks which live cells have enough live neighbors to survive to the next round
        survive_filter = np.logical_or.reduce([(s0 - 1 < neighbors) & (neighbors < s1 + 1) for s0, s1 in self.parameters.survive]) & (new_world > 0.5)
        new_world[...] = 0
        new_world[birth_filter | survive_filter] = 1
        self.world = new_world
        self.memory += self.world

    def iterate(self, n_iter: int, plot: bool=False, plot_memory: bool=False):
        """
        Runs the world and plots.
        n_iter: number of iterations
        plot: if True, plots each iteration
        plot_memory: 
            if True, plots matrix sum of all previous iterations + current iteration 
            instead of just the current iteration
        """
        old_world = np.array(self.world)
        for _ in range(n_iter):
            if not np.any(self.world):
                break
            self.procreate()
            # stop if the world has stabilized
            if np.array_equal(self.world, old_world):
                break
            old_world = np.array(self.world)
            if plot:
                self.plot(memory=plot_memory)
            

    def plot(self, memory: bool=False):
        """
        Plot the world. If memory is True, then plot the matrix sum of all previous iterations
        """
        if memory:
            plt.imshow(self.memory, cmap=self.cmap)
        else:
            plt.imshow(self.world, cmap=self.cmap)
        plt.axis('off')
        plt.show()

def create_random_world(height: int, width: int, life_probability: float=0.5):
    return np.random.choice(np.array([0, 1], dtype=np.float64), 
                            (height, width), 
                            p=[life_probability, 1.-life_probability])