Schelling Segregation Model

def make_grid(size=GRID_SIZE, empty_fraction=EMPTY_FRACTION, seed=SEED):
    """Return a random grid of RED and BLUE agents with some EMPTY cells.

    Each cell contains 0 (EMPTY), 1 (RED), or 2 (BLUE).  The two agent
    types are equal in number; the remaining cells are empty.
    """
    rng = np.random.default_rng(seed)
    n_cells = size * size
    n_empty = int(n_cells * empty_fraction)
    n_agents = n_cells - n_empty
    n_red = n_agents // 2
    n_blue = n_agents - n_red
    flat = np.array([EMPTY] * n_empty + [RED] * n_red + [BLUE] * n_blue)
    rng.shuffle(flat)
    return flat.reshape(size, size)

def same_neighbor_fraction(grid, row, col):
    """Return the fraction of occupied neighbours of the same type.

    Uses the Moore neighbourhood (up to 8 surrounding cells).
    Returns 1.0 when no neighbours are occupied so isolated agents
    are always considered satisfied.
    """
    size = grid.shape[0]
    agent_type = grid[row, col]
    n_same = 0
    n_occupied = 0
    for dr in (-1, 0, 1):
        for dc in (-1, 0, 1):
            if dr == 0 and dc == 0:
                continue
            r, c = row + dr, col + dc
            if 0 <= r < size and 0 <= c < size and grid[r, c] != EMPTY:
                n_occupied += 1
                if grid[r, c] == agent_type:
                    n_same += 1
    if n_occupied == 0:
        return 1.0
    return n_same / n_occupied

def step(grid, threshold=THRESHOLD, rng=None):
    """Return a new grid after moving all dissatisfied agents.

    Each dissatisfied agent is paired with a randomly chosen empty cell
    and moved there.  Agents that cannot be matched (too few vacancies)
    stay in place.
    """
    if rng is None:
        rng = np.random.default_rng(SEED)
    size = grid.shape[0]
    dissatisfied = [
        (r, c)
        for r in range(size)
        for c in range(size)
        if grid[r, c] != EMPTY and same_neighbor_fraction(grid, r, c) < threshold
    ]
    empties = [(r, c) for r in range(size) for c in range(size) if grid[r, c] == EMPTY]
    rng.shuffle(dissatisfied)
    rng.shuffle(empties)
    new_grid = grid.copy()
    for i, (r, c) in enumerate(dissatisfied):
        if i < len(empties):
            er, ec = empties[i]
            new_grid[er, ec] = grid[r, c]
            new_grid[r, c] = EMPTY
    return new_grid

def satisfaction_rate(grid, threshold=THRESHOLD):
    """Return the fraction of agents that are satisfied.

    A higher satisfaction rate indicates that agents have clustered with
    like neighbours; it approaches 1.0 as stable clusters form.
    """
    size = grid.shape[0]
    n_agents = 0
    n_satisfied = 0
    for r in range(size):
        for c in range(size):
            if grid[r, c] != EMPTY:
                n_agents += 1
                if same_neighbor_fraction(grid, r, c) >= threshold:
                    n_satisfied += 1
    if n_agents == 0:
        return 1.0
    return n_satisfied / n_agents

def run(grid, n_steps=N_STEPS, threshold=THRESHOLD, seed=SEED):
    """Run the simulation and return a list of grid snapshots.

    Element 0 is the initial grid; element i is the grid after i steps.
    """
    rng = np.random.default_rng(seed)
    snapshots = [grid.copy()]
    for _ in range(n_steps):
        grid = step(grid, threshold, rng)
        snapshots.append(grid.copy())
    return snapshots

def plot_snapshots(snapshots, steps_to_show, filename):
    """Save a figure showing the grid state at the given step indices.

    White cells are empty, red cells are type RED, blue cells are type BLUE.
    """
    n = len(steps_to_show)
    cmap = mcolors.ListedColormap(["white", "tomato", "steelblue"])
    fig, axes = plt.subplots(1, n, figsize=(3 * n, 3))
    for ax, idx in zip(axes, steps_to_show):
        ax.imshow(snapshots[idx], cmap=cmap, vmin=0, vmax=2, interpolation="nearest")
        ax.set_title(f"Step {idx}")
        ax.axis("off")
    fig.tight_layout()
    fig.savefig(filename, dpi=150)
    plt.close(fig)

import numpy as np
import pytest
from schelling import (
    EMPTY,
    RED,
    BLUE,
    GRID_SIZE,
    SEED,
    make_grid,
    same_neighbor_fraction,
    step,
    satisfaction_rate,
    run,
)


def test_grid_shape():
    # Grid must be a square with side GRID_SIZE.
    grid = make_grid()
    assert grid.shape == (GRID_SIZE, GRID_SIZE)


def test_grid_cell_values():
    # Every cell must be EMPTY, RED, or BLUE.
    grid = make_grid()
    assert set(np.unique(grid)).issubset({EMPTY, RED, BLUE})


def test_grid_agent_counts_balanced():
    # The two agent types are built as equal halves of the agent pool,
    # so their counts differ by at most 1 (rounding when n_agents is odd).
    grid = make_grid()
    assert abs(np.sum(grid == RED) - np.sum(grid == BLUE)) <= 1


def test_neighbor_fraction_isolated_agent():
    # An agent with no occupied neighbours has same_neighbor_fraction 1.0
    # so it is treated as satisfied and never moves.
    grid = np.zeros((5, 5), dtype=int)
    grid[2, 2] = RED
    assert same_neighbor_fraction(grid, 2, 2) == pytest.approx(1.0)


def test_neighbor_fraction_all_same():
    # An agent surrounded entirely by the same type has fraction 1.0.
    grid = np.full((3, 3), RED, dtype=int)
    assert same_neighbor_fraction(grid, 1, 1) == pytest.approx(1.0)


def test_neighbor_fraction_all_different():
    # An agent surrounded entirely by the opposite type has fraction 0.0.
    grid = np.full((3, 3), BLUE, dtype=int)
    grid[1, 1] = RED
    assert same_neighbor_fraction(grid, 1, 1) == pytest.approx(0.0)


def test_stable_grid_unchanged():
    # A perfectly segregated 4x4 grid (RED left, BLUE right) has every
    # agent meeting the 0.3 threshold, so no movement should occur.
    # Border agents (column 1 or 2) have at least 5 same-type neighbours
    # out of 8, giving a fraction of 0.625 >= THRESHOLD.
    grid = np.zeros((4, 4), dtype=int)
    grid[:, :2] = RED
    grid[:, 2:] = BLUE
    rng = np.random.default_rng(SEED)
    assert np.array_equal(step(grid, rng=rng), grid)


def test_satisfaction_increases():
    # Running 20 steps from the random initial grid must raise the
    # satisfaction rate; agents self-organise into clusters.
    grid = make_grid()
    initial = satisfaction_rate(grid)
    snapshots = run(grid)
    assert satisfaction_rate(snapshots[-1]) > initial