mass_spring_base.py

# Authored by Yuanming Hu, Taichi Graphics.
import taichi as ti

ti.init(arch=ti.gpu)

max_num_particles = 1024
substeps = 10
spring_Y = 1000  # Young's modulus
drag_damping = 1
dashpot_damping = 100
particle_mass = 1.0
dt = 1e-3
connection_radius = 0.15


num_particles = ti.field(dtype=ti.i32, shape=())
# Boolean state
paused = ti.field(dtype=ti.i32, shape=())
# Container holding particle positions
x = ti.Vector.field(2, dtype=ti.f32, shape=max_num_particles)
# Container holding particle velocities
v = ti.Vector.field(2, dtype=ti.f32, shape=max_num_particles)
# Container holding particle forces
f = ti.Vector.field(2, dtype=ti.f32, shape=max_num_particles)
# Boolean state per particle
fixed = ti.field(dtype=ti.i32, shape=max_num_particles)

# rest_length[i, j] == 1: i and j are connected
# rest_length[i, j] == 0: i and j are NOT connected
rest_length = ti.field(dtype=ti.f32,
                       shape=(max_num_particles, max_num_particles))

@ti.kernel
def new_particle():
    # TODO: what happens if you add a new fixed or free falling particle?
    # Hints:
    #  1. Save new particle i's position to the corresponding slot in `x`.
    #  2. Compute the distance between i and any old particle j
    #  3. Set rest_length[i, j] & rest_length[j, i] to 0.1 if their distance
    #     is smaller than `connection_radius`.
    pass

@ti.kernel
def substep():
    # TODO: implement behavior of particles in a substep
    # Hints:
    #   Using `d` as unit vector of distance between particle i and particle j.
    #   1. Gravity for free falling particles)
    #      f = [0, -9.8] * particle_mass
    #   2. Spring force between connected particles
    #      f += spring_Y * delta_length * d
    #   3. Dashpot damping
    #      f += -dashpot_damping * delta_v_ij * d
    #   4. Drag damping
    #      v *= exp(-dt * drag_damping)
    #   5. force and velocity
    #      v += dt * f / particle_mass
    #   6. velocity and position
    #      x += v * dt
    #   7. Set window border
    #      Reset position to border if a particle falls out of the window.
    pass

def main():
    gui = ti.GUI('Explicit Mass Spring System',
                 res=(512, 512),
                 background_color=0xDDDDDD)

    while True:
        for e in gui.get_events(ti.GUI.PRESS):
            if e.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]:
                print('Exiting')
                exit()
            elif e.key == gui.SPACE:
                paused[None] = not paused[None]
                print('Paused: ', paused[None])
            elif e.key == ti.GUI.LMB:
                is_fixed = gui.is_pressed(ti.GUI.SHIFT)
                state_str = 'fixed' if is_fixed else 'free falling'
                print(f'Adding new {state_str} particle ({e.pos[0]}, {e.pos[1]})!')
                new_particle()
            elif e.key == 'c':
                print('Clear everything')
                num_particles[None] = 0
                rest_length.fill(0)

        if not paused[None]:
            for step in range(substeps):
                substep()

        X = x.to_numpy()
        n = num_particles[None]

        # Draw the springs
        for i in range(n):
            for j in range(i + 1, n):
                if rest_length[i, j] != 0:
                    gui.line(begin=X[i], end=X[j], radius=2, color=0x444444)

        # Draw the particles
        for i in range(n):
            c = 0xFF0000 if fixed[i] else 0x111111
            gui.circle(pos=X[i], color=c, radius=5)


        gui.show()

if __name__ == '__main__':
    main()