FIRE retro x80 Python numpy FTW

fire.py

import pygame
import random
import colorsys
import numpy as np
from pygame import surfarray

def color_HSL(h, s, l):
    color = colorsys.hls_to_rgb(h/255, l/255, s/255)
    return tuple( int(x*255) for x in color )

def fire_palette(color_shift=0, color_split=1):
    return [
        color_HSL((x / color_split + color_shift) % 256, 255, min(255, x*2))
        for x in range(256)
        ]

def check_exit():
    for ev in pygame.event.get():
        if ev.type == pygame.QUIT:
            raise KeyboardInterrupt()

def fire_transform_numpy(arr):
    height = arr.shape[0]
    for y in range(height-1):
        arr[y, 1:-1] = (
            arr[y+1, 1:-1]
          + arr[y+1,  :-2]
          + arr[y+1, 2:  ]
          + arr[(y+2)%height, 1:-1]
          ) * 0.2499

def optimized_demo(W, H):
    screen = pygame.display.set_mode((W, H))
    clock = pygame.time.Clock()
    
    fire = np.zeros((H, W)).astype(int)

    fire_surface = pygame.Surface((fire.shape[1], fire.shape[0]), pygame.HWPALETTE, 8)
    fire_surface.set_palette(fire_palette(0))

    while True:
        screen.fill(0)
            
        if pygame.time.get_ticks() // 1000 % 10 < 5:
            fire[-1] = np.random.randint(0, 150, (fire.shape[1],))
        
        fire_transform_numpy(fire) 

        surfarray.blit_array(fire_surface, np.transpose(fire))
        screen.blit(fire_surface, (0,0))

        check_exit()
        clock.tick(200)
        pygame.display.flip()

def main():
    try:
        optimized_demo(300, 150)
    except KeyboardInterrupt:
        pass
    finally:
        pygame.quit()

main()