pygame spiral markov chain pixels

spiralout.py

import pygame
import numpy as np

SCREEN_X = 256
SCREEN_Y = 256

GRID_SIZE = SCREEN_X //6

CELL_W = SCREEN_X // GRID_SIZE
CELL_H = SCREEN_Y // GRID_SIZE

CALLWIDTH = min(CELL_H // 4,CELL_W // 4)  #int(GRID_SIZE ** 0.5)

ANIM_TICK_EVENT = pygame.USEREVENT + 0

SELF_CHANCE = 0.9

def gen_1d_probs(i,n,self_chance=0.5):
    for j in range(n):
        if i == j:
            yield 0 + self_chance
        else:
            yield (1 - self_chance) / (n - 1)

def spiral_indices_outward(n):
    X = Y = n
    x = y = 0
    dx = 0
    dy = -1
    inds = []
    for _ in range(max(X, Y)**2):
        if (-X/2 < x <= X/2) and (-Y/2 < y <= Y/2):
            inds.append((x, y))
        if x == y or (x < 0 and x == -y) or (x > 0 and x == 1-y):
            dx, dy = -dy, dx
        x, y = x+dx, y+dy
    return [(x+n//2, y+n//2) for x,y in inds]

def draw_spiral(indices, canvas, pallet=[(237,106,90), (126,132,107)]):
    lastcol = 0
    plen = len(pallet)
    for i, (x, y) in enumerate(indices):
        colndx = np.random.choice(plen, p=[*gen_1d_probs(lastcol,plen,self_chance=SELF_CHANCE)])
        colour = pallet[colndx]
        lastcol = colndx
        pygame.draw.rect(canvas, colour, (x * CELL_W, y * CELL_H, CELL_W, CELL_H), 0)
        yield



def main():
    pygame.init()
    pygame.display.set_caption("Wave")
    screen = pygame.display.set_mode((SCREEN_X, SCREEN_Y))
    clock = pygame.time.Clock()
    canvas = pygame.Surface((SCREEN_X, SCREEN_Y))
    
    pallet = [(237,106,90), (126,132,107), (61,49,74), (244,241,187)]#, (255,255,255)]

    pygame.time.set_timer(ANIM_TICK_EVENT, 10)
    indices = spiral_indices_outward(GRID_SIZE)
    gen = (x for x in []) # generator that fires StopIteration immidiately
    while True:
        clock.tick(60)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return
            if event.type == ANIM_TICK_EVENT:
                    try:
                        gen.__next__()
                    except StopIteration:
                        canvas.fill((0, 0, 0))
                        gen = draw_spiral(indices, canvas, pallet=pallet)
                        gen.__next__()

        screen.fill((0, 0, 0))
        screen.blit(canvas, (0, 0))
        pygame.display.update()


if __name__ == "__main__":
    main()