taichi + py5 failure (taichi needs the main thread)

taichi_py5.py

import taichi as ti
import taichi.math as tm
import py5

#ti.init(arch=ti.gpu)
ti.init()

@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(t: float):
    for i, j in pixels:  # Parallelized over all pixels
        c = tm.vec2(-0.8, tm.cos(t) * 0.2)
        z = tm.vec2(i / n - 1, j / n - 0.5) * 2
        iterations = 0
        while z.norm() < 20 and iterations < 50:
            z = complex_sqr(z) + c
            iterations += 1
        pixels[i, j] = 1 - iterations * 0.02

#gui = ti.GUI("Julia Set", res=(n * 2, n))
# i = 0
# while gui.running:
#     paint(i * 0.03)
#     gui.set_image(pixels)
#     gui.show()
#     i += 1

def setup():
    global pixels, n
    py5.size(640, 320)
    n = 320
    pixels = ti.field(dtype=float, shape=(n * 2, n))

def draw():
    #paint(py5.frame_count * 0.03)
    py5.set_np_pixels(pixels.to_numpy())
    
py5.run_sketch(block=False)

Hi! Yes, I'm on Linux... I'm not sure it is incompatible, my mental model of the whole thing is sketchy... my first intuitions (and attempts) are that it might require the main thread in a way incompatible to py5? but again, you are the expert :D

I tried using taichi a few times in the past and didn't get very far.

Usually the "main thread issue" is for MacOS. It could be taichi uses the GPU in a special way that requires it to be the context thread, much like how OpenGL works. You aren't using P2D or P3D here though. It could be confused though because calling taichi method sfrom inside draw like that means the executing thread originates with Java, not Python. If taichi is finicky about what threads it will run in, it might not be happy about that.