taichi_gpu_julia_set_fractal_renderer_with_animation.py

1import taichi as ti
2import taichi.math as tm
3
4ti.init(arch=ti.gpu)
5
6n = 320
7pixels = ti.Vector.field(3, dtype=float, shape=(n * 2, n))
8
9@ti.func
10def complex_sqr(z):  # complex square of a 2D vector
11    return tm.vec2(z[0] * z[0] - z[1] * z[1], 2 * z[0] * z[1])
12
13@ti.kernel
14def paint(t: float):
15    for i, j in pixels:  # Parallelized over all pixels
16        c = tm.vec2(-0.8, tm.cos(t) * 0.2)
17        z = tm.vec2(i / n - 1, j / n - 0.5) * 2
18        iterations = 0
19        while z.norm() < 20 and iterations < 50:
20            z = complex_sqr(z) + c
21            iterations += 1
22        pixels[i, j] = tm.vec3(1 - iterations * 0.02, 1 - iterations * 0.05, 1)
23
24gui = ti.GUI("Julia Set", res=(n * 2, n))
25
26for i in range(1000000):
27    paint(i * 0.03)
28    gui.set_image(pixels)
29    gui.show()