taichi_zoo_life.py

taichi_zoo_life.py

import taichi as ti

Category10_10 = ('#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', 
                 '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf')

ti.init(arch=ti.gpu)

N = 2000
N_COLOR = 10
COLORS = [int(c[1:], 16) for c in Category10_10]
window_size = 600
viscosity = 0.6
margin = 5
time_scale = 0.2

pos = ti.Vector.field(2, ti.f32, N)
pos_draw = ti.Vector.field(2, ti.f32, N)
vel = ti.Vector.field(2, ti.f32, N)
force = ti.Vector.field(2, ti.f32, N)
color = ti.field(ti.i32, N)
gravity = ti.field(ti.float32, (N_COLOR, N_COLOR))
radius = ti.field(ti.float32, (N_COLOR, N_COLOR))

@ti.kernel
def initialize():
    for i, j in gravity:
        gravity[i, j] = ti.random() * 2 - 1

    for i, j in radius:
        radius[i, j] = 40 + 60 * ti.random()

    for i in range(N):
        pos[i] = [ti.random() * window_size, ti.random() * window_size]
        vel[i] = [0, 0]
        force[i] = [0, 0]
        color[i] = ti.cast(ti.random() * N_COLOR, ti.i32)

@ti.kernel
def update():
    for i in range(N):
        force[i] = [0.0, 0.0]

    for i in range(N):
        p = pos[i]
        for j in range(N):
            if i != j:
                diff = p - pos[j]
                r = radius[color[i], color[j]]
                d = diff.norm(1e-5)
                if d > 0 and d < r:
                    g = gravity[color[i], color[j]]
                    force[i] += (g / d) * diff

        strength = 1

        if p[0] < margin:
            force[i][0] += (margin - p[0]) * strength
        elif p[0] > window_size - margin:
            force[i][0] += (window_size - margin - p[0]) * strength
        
        if p[1] < margin:
            force[i][1] += (margin - p[1]) * strength
        elif p[1] > window_size - margin:
            force[i][1] += (window_size - margin - p[1]) * strength


    for i in range(N):
        vel[i] = vel[i] * (1 - viscosity) + time_scale * force[i]
        pos[i] += vel[i]
        pos_draw[i] = pos[i] / window_size


def main():
    gui = ti.GUI('life', (window_size, window_size))

    initialize()
    while gui.running:
        gui.get_event() 
        if gui.is_pressed(ti.GUI.LMB):
            initialize()
        update()
        gui.clear(0x112F41)
        gui.circles(pos_draw.to_numpy(), palette=COLORS, palette_indices=color, radius=1.5)
        gui.show()


main()