-
Notifications
You must be signed in to change notification settings - Fork 0
/
mass_spring_explicit.py
117 lines (91 loc) · 3.54 KB
/
mass_spring_explicit.py
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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import taichi as ti
ti.init(debug=True)
max_num_particles = 256
dt = 1e-3
num_particles = ti.var(ti.i32, shape=())
spring_stiffness = ti.var(ti.f32, shape=())
paused = ti.var(ti.i32, shape=())
damping = ti.var(ti.f32, shape=())
particle_mass = 1
bottom_y = 0.05
x = ti.Vector(2, dt=ti.f32, shape=max_num_particles)
v = ti.Vector(2, dt=ti.f32, shape=max_num_particles)
A = ti.Matrix(2, 2, dt=ti.f32, shape=(max_num_particles, max_num_particles))
b = ti.Vector(2, dt=ti.f32, shape=max_num_particles)
# rest_length[i, j] = 0 means i and j are not connected
rest_length = ti.var(ti.f32, shape=(max_num_particles, max_num_particles))
connection_radius = 0.15
gravity = [0, -9.8]
@ti.kernel
def substep():
# Compute force and new velocity
n = num_particles[None]
for i in range(n):
v[i] *= ti.exp(-dt * damping[None]) # damping
total_force = ti.Vector(gravity) * particle_mass
for j in range(n):
if rest_length[i, j] != 0:
x_ij = x[i] - x[j]
total_force += -spring_stiffness[None] * (x_ij.norm() - rest_length[i, j]) * x_ij.normalized()
v[i] += dt * total_force / particle_mass
# Collide with ground
for i in range(n):
if x[i].y < bottom_y:
x[i].y = bottom_y
v[i].y = 0
# Compute new position
for i in range(num_particles[None]):
x[i] += v[i] * dt
@ti.kernel
def new_particle(pos_x: ti.f32, pos_y: ti.f32): # Taichi doesn't support using Matrices as kernel arguments yet
new_particle_id = num_particles[None]
x[new_particle_id] = [pos_x, pos_y]
v[new_particle_id] = [0, 0]
num_particles[None] += 1
# Connect with existing particles
for i in range(new_particle_id):
dist = (x[new_particle_id] - x[i]).norm()
if dist < connection_radius:
rest_length[i, new_particle_id] = 0.1
rest_length[new_particle_id, i] = 0.1
gui = ti.GUI('Mass Spring System', res=(512, 512), background_color=0xdddddd)
spring_stiffness[None] = 10000
damping[None] = 20
new_particle(0.3, 0.3)
new_particle(0.3, 0.4)
new_particle(0.4, 0.4)
while True:
for e in gui.get_events(ti.GUI.PRESS):
if e.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]:
exit()
elif e.key == gui.SPACE:
paused[None] = not paused[None]
elif e.key == ti.GUI.LMB:
new_particle(e.pos[0], e.pos[1])
elif e.key == 'c':
num_particles[None] = 0
rest_length.fill(0)
elif e.key == 's':
if gui.is_pressed('Shift'):
spring_stiffness[None] /= 1.1
else:
spring_stiffness[None] *= 1.1
elif e.key == 'd':
if gui.is_pressed('Shift'):
damping[None] /= 1.1
else:
damping[None] *= 1.1
if not paused[None]:
for step in range(10):
substep()
X = x.to_numpy()
gui.circles(X[:num_particles[None]], color=0xffaa77, radius=5)
gui.line(begin=(0.0, bottom_y), end=(1.0, bottom_y), color=0x0, radius=1)
for i in range(num_particles[None]):
for j in range(i + 1, num_particles[None]):
if rest_length[i, j] != 0:
gui.line(begin=X[i], end=X[j], radius=2, color=0x445566)
gui.text(content=f'C: clear all; Space: pause', pos=(0, 0.95), color=0x0)
gui.text(content=f'S: Spring stiffness {spring_stiffness[None]:.1f}', pos=(0, 0.9), color=0x0)
gui.text(content=f'D: damping {damping[None]:.2f}', pos=(0, 0.85), color=0x0)
gui.show()