Realtime Mandelbrot fractal with tachi and pygame

import taichi as ti
import taichi.math as tm

import pygame as pg


ti.init(arch=ti.gpu, default_fp=ti.f64, fast_math=False)

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


n = 1200
pixels = ti.Vector.field(3, dtype=float, shape=(n // 2, n))


@ti.func
def complex_sqr(z):  # complex square of a 2D vector
    return tm.vec2(z[0] * z[0] - z[1] * z[1], 2 * z[0] * z[1])


@ti.kernel
def paint(max_iter: int, offset_x: float, offset_y: float, thezoom: float):
    zoom = 1.0 / (thezoom + 0.0001)
    for i, j in pixels:  # Parallelized over all pixels
        c = tm.vec2(
            zoom * 4.0 * (i / n) - (2.0 * zoom) + offset_x / 100.0,
            zoom * 4 * (j / n) - (zoom * 2.0) + offset_y / 100.0,
        )
        z = tm.vec2(0.0, 0.0)
        iterations = 0
        while z.norm() < 2 and iterations < max_iter:
            z = complex_sqr(z) + c
            iterations += 1
        pixels[i, j][0] = 255.0 * (1 - (iterations / float(max_iter)))
        pixels[i, j][1] = 55.0 * (iterations / float(max_iter))
        pixels[i, j][2] = 128.0 * (1 - (iterations / float(max_iter)))


# gui = ti.GUI("Mandelbrot Set", res=(n * 2, n))


def get_output():
    paint(iterations, center_x, center_y, zoom)

    arr = pixels.to_numpy()
    return pg.surfarray.make_surface(arr)


# while True:
#     gui.set_image(pixels)
#     gui.show()

screen = pg.display.set_mode((n // 2, n), pg.HWACCEL)

running = True


iterations = 100
center_x = 0.0
center_y = 0.0
zoom = 1.0

output = get_output()

fps = clock.get_fps()

while running:
    for event in pg.event.get():
        if event.type == pg.QUIT:
            running = False
        if event.type == pg.KEYDOWN and event.key == pg.K_DOWN:
            iterations = max(iterations - 1, 0)
            output = get_output()
            print(iterations)
        elif event.type == pg.KEYDOWN and event.key == pg.K_UP:
            iterations = min(iterations + 1, 280)
            output = get_output()
            print(iterations)
        elif event.type == pg.KEYDOWN and event.key == pg.K_r:
            zoom = 100.0
            output = get_output()
            print(zoom)

    pressed = pg.key.get_pressed()
    if pressed[pg.K_a]:
        center_x -= 1 * zoom
        output = get_output()
        print(f"({center_x},{center_y})")
    elif pressed[pg.K_d]:
        center_x += 1 / zoom
        output = get_output()
        print(f"({center_x},{center_y})")
    elif pressed[pg.K_w]:
        center_y -= 1 / zoom
        output = get_output()
        print(f"({center_x},{center_y})")
    elif pressed[pg.K_s]:
        center_y += 1 / zoom
        output = get_output()
        print(f"({center_x},{center_y})")
    elif pressed[pg.K_PAGEUP]:
        zoom += 0.1
        output = get_output()
        print(zoom)
    elif pressed[pg.K_PAGEDOWN]:
        zoom -= 0.1
        zoom = max(zoom, 0.000001)
        output = get_output()
        print(zoom)

    screen.fill((200, 200, 200))
    screen.blit(output, (0, 0))

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


pg.quit()