Mandelbrot Set分型动画 使用Taichi编程

import taichi as ti

ti.init(arch=ti.gpu)

n = 320
pixels = ti.var(dt=ti.f32, shape=(n * 2, n))

@ti.func
def Complex_sqr(z):
    return ti.Vector([z[0]**2 - z[1]**2, z[1] * z[0] * 2])


@ti.func
def Complex_quadratic(z,c):
    temp = Complex_sqr(z)
    return ti.Vector([temp[0] + c[0], temp[1] + c[1]])

@ti.func
def test(c,maxi):
    temp = c
    flag = 1
    for i in range(maxi):
        if temp.norm()>2 :
            flag = 0
            break
        else:
            temp = Complex_quadratic(temp,c)

    return flag



@ti.kernel
def draw(a:ti.f32):
    for i,j in pixels:
        if test(ti.Vector([i/(n+a*10)+0.3,j/(n+a*10)-0.0005*a]),10000):
            pixels[i,j] = 0
        else:
            pixels[i,j] = 1


gui = ti.GUI("MandelbrotS Set", res=(n * 2, n))
#result_dir = "./results"
#video_manger = ti.VideoManager(output_dir=result_dir, framerate=24, automatic_build=False)

for a in range(100000):
    draw(a)
    gui.set_image(pixels)
    gui.show()

    #video_manger.write_frame(pixels.to_numpy())
    #print(f'\rFrame {a+1}/100000 is recorded', end='')


#print()
#print('Exporting .mp4 and .gif videos...')
#video_manger.make_video(gif=True, mp4=True)
#print(f'MP4 video is saved to {video_manger.get_output_filename(".mp4")}')
#print(f'GIF video is saved to {video_manger.get_output_filename(".gif")}')