Conway's Game of Life (3D)

import numpy as np
import pyvista as pv
from matplotlib import cm
from matplotlib.colors import LinearSegmentedColormap, to_rgb

def initialize_grid(size, prob=0.2):
    """Initialize a 2D grid with random live cells."""
    return np.random.choice([0, 1], size=size, p=[1-prob, prob])

def count_neighbors(grid, x, y):
    """Count neighbors for the cell at (x, y) in the grid."""
    offsets = [-1, 0, 1]
    neighbors = 0
    for dx in offsets:
        for dy in offsets:
            if dx == 0 and dy == 0:
                continue
            nx, ny = x + dx, y + dy
            if 0 <= nx < grid.shape[0] and 0 <= ny < grid.shape[1]:
                neighbors += grid[nx, ny]
    return neighbors

def evolve(grid):
    """Evolve the 2D grid to the next generation."""
    new_grid = grid.copy()
    for x in range(grid.shape[0]):
        for y in range(grid.shape[1]):
            neighbors = count_neighbors(grid, x, y)
            if grid[x, y] == 1 and (neighbors < 2 or neighbors > 3):
                new_grid[x, y] = 0
            elif grid[x, y] == 0 and neighbors == 3:
                new_grid[x, y] = 1
    return new_grid

def create_3d_grid(history):
    """Combine all generations into a single 3D grid."""
    max_generations = len(history)
    size_x, size_y = history[0].shape
    grid_3d = np.zeros((max_generations, size_x, size_y), dtype=int)
    for z, layer in enumerate(history):
        grid_3d[z, :, :] = layer
    return grid_3d

def get_colors(history, color_mode, start_color=None, end_color=None):
    """Generate colors for the cubes based on the chosen mode."""
    num_generations = len(history)

    if color_mode == "gradient":
        # Use a predefined colormap
        cmap = cm.get_cmap("viridis", num_generations)
        return [cmap(i / (num_generations - 1))[:3] for i in range(num_generations)]

    elif color_mode == "custom_gradient":
        # Create a custom gradient using start and end colors
        if not start_color or not end_color:
            raise ValueError("Start and end colors must be specified for custom gradient.")
        cmap = LinearSegmentedColormap.from_list("custom", [start_color, end_color], num_generations)
        return [cmap(i / (num_generations - 1))[:3] for i in range(num_generations)]

    elif color_mode == "fixed":
        # Use a set of fixed contrasting colors
        fixed_colors = [
            (1, 0, 0),  # Red
            (0, 1, 0),  # Green
            (0, 0, 1),  # Blue
            (1, 1, 0),  # Yellow
            (1, 0, 1),  # Magenta
            (0, 1, 1),  # Cyan
        ]
        return [fixed_colors[i % len(fixed_colors)] for i in range(num_generations)]

    else:
        raise ValueError("Invalid color mode. Choose 'gradient', 'custom_gradient', or 'fixed'.")

def visualize_pyvista_with_colors(history, color_mode="gradient", start_color=None, end_color=None):
    """Visualize the 3D grid with PyVista using cubes and flexible colors."""
    grid_3d = create_3d_grid(history)
    points = []
    colors = get_colors(history, color_mode, start_color, end_color)

    for z, layer in enumerate(grid_3d):
        for x in range(layer.shape[0]):
            for y in range(layer.shape[1]):
                if grid_3d[z, x, y] == 1:
                    points.append((x, y, z))

    points = np.array(points)

    # Create PyVista mesh with cubes and dynamic colors
    plotter = pv.Plotter()
    for i, point in enumerate(points):
        gen = int(point[2])  # The generation determines the color
        cube = pv.Cube(center=point, x_length=1, y_length=1, z_length=1)
        plotter.add_mesh(cube, color=colors[gen], show_edges=True)
    plotter.show()

def main():
    size = (20, 20)  # Define the grid size
    generations = 10

    # Set this to "random", "custom", or "list"
    initial_input_mode = "list"

    if initial_input_mode == "random":
        # Generate a random initial grid
        initial_grid = initialize_grid(size)
    elif initial_input_mode == "custom":
        # Define a custom initial grid (modify as needed)
        initial_grid = np.zeros(size, dtype=int)
        initial_grid[10, 9:12] = 1  # Example: Horizontal line at the center
    elif initial_input_mode == "list":
        # User provides a Python list for the initial grid
        user_input_list = [
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 1, 1, 1, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        ]
        # Convert list to numpy array
        initial_grid = np.array(user_input_list, dtype=int)
        size = initial_grid.shape  # Update size based on user input
    else:
        raise ValueError("Invalid input mode. Choose 'random', 'custom', or 'list'.")

    print("Initial 2D grid (0 = dead, 1 = alive):")
    print(initial_grid)

    history = [initial_grid]
    current_grid = initial_grid

    for _ in range(generations):
        next_grid = evolve(current_grid)
        history.append(next_grid)
        current_grid = next_grid

    # Choose the desired color mode
    color_mode = "custom_gradient"  # "gradient", "custom_gradient", or "fixed"
    start_color = (1, 0, 0)  # RGB (for custom_gradient)
    end_color = (0, 0, 1)    # RGB (for custom_gradient)

    visualize_pyvista_with_colors(history, color_mode=color_mode, start_color=start_color, end_color=end_color)

if __name__ == "__main__":
    main()