taichi_life.py

from kiwisolver import strength
import taichi as ti
import numpy as np

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

ti.init(arch=ti.gpu)

N = 5000
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)
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))

gravity.from_numpy(np.random.uniform(-1, 1, (N_COLOR, N_COLOR)).astype(np.float32))
radius.from_numpy(np.random.uniform(40, 100, (N_COLOR, N_COLOR)).astype(np.float32))
#radius.from_numpy(np.full((N_COLOR, N_COLOR), 80, dtype=np.float32))

@ti.kernel
def initialize():
    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]


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

    initialize()
    print(color.to_numpy().max())
    while gui.running:
        update()
        gui.circles(pos.to_numpy() / window_size, palette=COLORS, palette_indices=color, radius=1.5)
        gui.show()


if __name__ == '__main__':
    main()        
    # import IPython
    # c = color.to_numpy()
    # p = pos.to_numpy()
    # v = vel.to_numpy()
    # IPython.embed()