Pygame + taichi fire graphic effect

colors.py

def hsl_to_rgb(h: float, s: float, l: float):
    r, g, b = 0.0, 0.0, 0.0
    h = h / 256.0
    s = s / 256.0
    l = l / 256.0

    if s == 0:
        r = g = b = l
    else:
        if l < 0.5:
            temp2 = l * (1 + s)
        else:
            temp2 = (l + s) - (l * s)

        temp1 = 2 * l - temp2
        tempr = h + 1.0 / 3.0
        if tempr > 1.0:
            tempr -= 1
        tempg = h
        tempb = h-1.0 / 3.0
        if tempb < 0.0:
            tempb += 1

        #     //red
        if tempr < 1.0 / 6.0:
            r = temp1 + (temp2 - temp1) * 6.0 * tempr
        elif tempr < 0.5:
            r = temp2
        elif tempr < 2.0 / 3.0:
            r = temp1 + (temp2 - temp1) * ((2.0 / 3.0) - tempr) * 6.0
        else:
            r = temp1

#      //green
        if tempg < 1.0 / 6.0:
            g = temp1 + (temp2 - temp1) * 6.0 * tempg
        elif tempg < 0.5:
            g = temp2
        elif tempg < 2.0 / 3.0:
            g = temp1 + (temp2 - temp1) * ((2.0 / 3.0) - tempg) * 6.0
        else:
            g = temp1

#     //blue
        if tempb < 1.0 / 6.0:
            b = temp1 + (temp2 - temp1) * 6.0 * tempb
        elif tempb < 0.5:
            b = temp2
        elif tempb < 2.0 / 3.0:
            b = temp1 + (temp2 - temp1) * ((2.0 / 3.0) - tempb) * 6.0
        else:
            b = temp1
    return (int(r * 255.0), int(g * 255.0), int(b * 255.0))

main.py

import numpy as np
import taichi as ti
import taichi.math as tm
from random import random
import pygame as pg

from colors import hsl_to_rgb

ti.init(arch=ti.gpu)

pg.init()
clock = pg.time.Clock()


N = 600
SCREEN_SIZE = (N//2, N)
SCREEN_SIZE_2X = (N, 2*N)

palette = ti.Vector.field(3, dtype=float, shape=(1, 255))

for i in range(256):
    r, g, b = hsl_to_rgb(i / 3, 255, min(255, i * 2))
    palette[0, i][0] = r
    palette[0, i][1] = g
    palette[0, i][2] = b


pixels = ti.Vector.field(3, dtype=float, shape=SCREEN_SIZE)
fire = ti.field(dtype=ti.f32, shape=SCREEN_SIZE)
next_fire = ti.field(dtype=ti.f32, shape=SCREEN_SIZE)
screen = pg.display.set_mode(SCREEN_SIZE_2X, pg.HWACCEL)


@ti.kernel
def update_fire():
    w, h = SCREEN_SIZE
    # Fills the image with random gray
    for x, y in fire:
        if y < h-1:
            fire[x, y] = ((fire[(x - 1 + w) % w, (y + 1) % h] + fire[(x) %
                                                                     w, (y + 1) % h] + fire[(x + 1) % w, (y + 1) % h] + fire[(x) % w, (y + 2) % h]) * 32) / 129


@ti.kernel
def randomize_bottom():
    for x in range(SCREEN_SIZE[0]):
        # print(abs(32768 + ti.randn()) % 256, end=" ")
        fire[x, SCREEN_SIZE[1] - 1] = abs(32768 + ti.randn()) % 256


@ti.kernel
def update_pixels():
    for i, j in pixels:
        pixels[i, j] = palette[0, int(fire[i, j])]
    # for x, y in fire:
    #     fire[x, y] = next_fire[x, y]
    # fire = next_fire.copy()


def update():
    global running
    randomize_bottom()

    update_fire()
    update_pixels()
    # print(next_fire)
    # running = False


running = True
fps = clock.get_fps()

while running:
    for event in pg.event.get():
        if event.type == pg.QUIT:
            running = False
    screen.fill((200, 200, 200))
    update()
    s = pg.surfarray.make_surface(pixels.to_numpy())
    s = pg.transform.scale2x(s)
    screen.blit(s, (0, 0))

    pg.display.flip()
    fps = 0.9 * fps + 0.1 * clock.get_fps()
    pg.display.set_caption("Fire Demo [FPS:" + str(fps) + "]")
    clock.tick(60)


pg.quit()