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

参考文档为:https://www.cnblogs.com/shushen/p/5473264.html
直接按照上面的公式实现的,但是直接就爆炸了

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

connection_radius = 0.15  # 两个质点连接的最大距离

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()
        if dist < connection_radius:
            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()
    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.text(content=f'Number of particles {num_particles[None]:.0f}', pos=(0, 0.80), color=0x0)
    gui.show()