# Homework1 弹簧质点系统隐式时间积分的雅克比迭代，有点问题，求解答

``````import taichi as ti

ti.init(arch=ti.cpu)
# ti.init(arch=ti.cpu, debug=True)

max_num_particles = 500

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  # 最底部y

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))  # 两个质点之间的弹簧原长

gravity = [0, -9.8]  # 重力

# jacobi iteration

max_iterate = 10

A = ti.var(dt=ti.f32, shape=(2 * max_num_particles, 2 * max_num_particles))
deltaV = ti.var(dt=ti.f32, shape=2 * max_num_particles)
new_deltaV = ti.var(dt=ti.f32, shape=2 * max_num_particles)
b = ti.var(dt=ti.f32, shape=2 * max_num_particles)

dfx = ti.var(dt=ti.f32, shape=(2 * max_num_particles, 2 * max_num_particles))
dfv = ti.var(dt=ti.f32, shape=(2 * max_num_particles, 2 * max_num_particles))

# 鼠标单击添加一个新的质点
@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]  # 必须用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()
rest_length[i, new_particle_id] = dist
rest_length[new_particle_id, i] = dist

@ti.func
def init_solver():
for i in range(2 * max_num_particles):
A[i, i] = particle_mass
b[i] = 0
for j in range(2 * max_num_particles):
dfx[i, j] = 0
dfv[i, j] = 0

@ti.func
def iterate():
n = num_particles[None]  # 只迭代左上角的小矩阵
for i in range(2 * n):
r = b[i]
for j in range(2 * n):
if i != j:
r -= A[i, j] * deltaV[j]
new_deltaV[i] = r / A[i, i]

for i in range(2 * n):
deltaV[i] = new_deltaV[i]

@ti.kernel
def substep():
n = num_particles[None]
# init solver
init_solver()
# compute dfx and dfv
for i in range(n):
total_force = ti.Vector(gravity) * particle_mass
for j in range(n):
if rest_length[i, j] != 0:
x_ij = x[i] - x[j]
v_ij = v[i] - v[j]
total_force += -damping[None] * x_ij.normalized() * v_ij * x_ij.normalized()  # damping
total_force += -spring_stiffness[None] * (x_ij.norm() - rest_length[i, j]) * x_ij.normalized()  # spring
if j >= i:  # 只需要计算上三角
x_ji = x[j] - x[i]
v_ji = v[j] - v[i]
dfx[2 * i, 2 * i] = spring_stiffness[None] * ((x_ji[0] * x_ji[0] - 1) / x_ji.norm() *
(x_ji.norm() - rest_length[i, j]) -
(x_ji[0] * x_ji[0]))
dfx[2 * i + 1, 2 * i] = spring_stiffness[None] * ((x_ji[0] * x_ji[1] - 0) / x_ji.norm() *
(x_ji.norm() - rest_length[i, j]) -
(x_ji[0] * x_ji[1]))
dfx[2 * i, 2 * i + 1] = spring_stiffness[None] * ((x_ji[1] * x_ji[0] - 0) / x_ji.norm() *
(x_ji.norm() - rest_length[i, j]) -
(x_ji[1] * x_ji[0]))
dfx[2 * i + 1, 2 * i + 1] = spring_stiffness[None] * ((x_ji[1] * x_ji[1] - 1) / x_ji.norm() *
(x_ji.norm() - rest_length[i, j]) -
(x_ji[1] * x_ji[1]))
dfx[2 * j, 2 * j] = dfx[2 * i, 2 * i]
dfx[2 * j + 1, 2 * j] = dfx[2 * i + 1, 2 * i]
dfx[2 * j, 2 * j + 1] = dfx[2 * i, 2 * i + 1]
dfx[2 * j + 1, 2 * j + 1] = dfx[2 * i + 1, 2 * i + 1]

dfx[2 * i, 2 * j] = -dfx[2 * i, 2 * i]
dfx[2 * i + 1, 2 * j] = -dfx[2 * i + 1, 2 * i]
dfx[2 * i, 2 * j + 1] = -dfx[2 * i, 2 * i + 1]
dfx[2 * i + 1, 2 * j + 1] = -dfx[2 * i + 1, 2 * i + 1]

dfx[2 * j, 2 * i] = -dfx[2 * i, 2 * i]
dfx[2 * j + 1, 2 * i] = -dfx[2 * i + 1, 2 * i]
dfx[2 * j, 2 * i + 1] = -dfx[2 * i, 2 * i + 1]
dfx[2 * j + 1, 2 * i + 1] = -dfx[2 * i + 1, 2 * i + 1]

dfx[2 * i, 2 * i] += -damping[None] * ((x_ji[0] * x_ji[0] - 1) / x_ji.norm() *
((x_ji[0] * v_ji[0] + x_ji[1] * v_ji[1]) * 1 +
x_ji[0] * v_ji[0]))
dfx[2 * i + 1, 2 * i] += -damping[None] * ((x_ji[0] * x_ji[1] - 0) / x_ji.norm() *
((x_ji[0] * v_ji[0] + x_ji[1] * v_ji[1]) * 0 +
x_ji[0] * v_ji[1]))
dfx[2 * i, 2 * i + 1] += -damping[None] * ((x_ji[1] * x_ji[0] - 0) / x_ji.norm() *
((x_ji[0] * v_ji[0] + x_ji[1] * v_ji[1]) * 0 +
x_ji[1] * v_ji[0]))
dfx[2 * i + 1, 2 * i + 1] += -damping[None] * ((x_ji[1] * x_ji[1] - 1) / x_ji.norm() *
((x_ji[0] * v_ji[0] + x_ji[1] * v_ji[1]) * 1 +
x_ji[1] * v_ji[1]))

dfv[2 * i, 2 * i] = damping[None] * x_ji[0] * x_ji[0]
dfv[2 * i + 1, 2 * i] = damping[None] * x_ji[0] * x_ji[1]
dfv[2 * i, 2 * i + 1] = damping[None] * x_ji[1] * x_ji[0]
dfv[2 * i + 1, 2 * i + 1] = damping[None] * x_ji[1] * x_ji[1]

b[2 * i] = total_force[0]
b[2 * i + 1] = total_force[1]

# compute A and b
for i in range(2 * n):
for j in range(2 * n):
A[i, j] -= (dt * dfv[i, j] + dt * dt * dfx[i, j])
if j % 2 == 0:
b[i] += dt * dfx[i, j] * v[j][0]
else:
b[i] += dt * dfx[i, j] * v[j][1]

# solve
for i in range(max_iterate):
iterate()

# Compute new velocity
for i in range(n):
v[i] += deltaV[i] * dt

# Compute new position
for i in range(n):
x[i] += v[i] * dt

# Collide with ground
for i in range(n):
if x[i].y < bottom_y:
x[i].y = bottom_y
v[i].y = 0

spring_stiffness[None] = 10000
# spring_stiffness[None] = 1000000
damping[None] = 20

new_particle(0.3, 0.3)
new_particle(0.3, 0.4)
new_particle(0.4, 0.4)

gui = ti.GUI('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]:
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()