1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import taichi as ti
import numpy as np

ti.init(arch=ti.gpu)

arr = (np.random.rand(512, 512)*4).astype(int)
arr[arr >= 3] = 0

x,y = arr.shape
sea = ti.field(dtype=ti.types.i8, shape=(x, y))
bak = ti.field(dtype=ti.types.i8, shape=(x, y))
sea.from_numpy(arr)
bak.copy_from(sea)

@ti.kernel
def forward2():
    for i, j in sea:
        ii = (i+1)%x
        if sea[i,j] == 2 and sea[ii,j] == 0:
            bak[i,j] = 0
            bak[ii,j] = 2
     
@ti.kernel
def forward1():
    for i, j in sea:
        jj = (j+1)%y
        if sea[i,j] == 1 and sea[i,jj] == 0:
            bak[i,j] = 0
            bak[i,jj] = 1

gui = ti.GUI("Sea", res=(x, y))
cnt = []
while not gui.get_event(ti.GUI.ESCAPE, ti.GUI.EXIT):
    mask = sea.to_numpy()
    img = np.zeros((*mask.shape, 3), dtype=np.uint8)
    img[mask == 0] = [1, 21, 66]
    img[mask == 1] = [3, 255, 247]
    img[mask == 2] = [247, 116, 146]
    cnt.append(img)
    gui.set_image(img)
    gui.show()
    forward1()
    sea.copy_from(bak)
    forward2()
    sea.copy_from(bak)

video_manager = ti.tools.VideoManager(output_dir="./out", framerate=60, automatic_build=False, width=2*x, height=2*y)
for img in cnt:
    video_manager.write_frame(img)
video_manager.make_video(mp4=True)