Skip to content

Instantly share code, notes, and snippets.

@FantasyVR

FantasyVR/implicit_mass_spring.py

Last active October 9, 2021 06:44
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save FantasyVR/ef2896c18da649e5dc85fb20aa0808a2 to your computer and use it in GitHub Desktop.
Save FantasyVR/ef2896c18da649e5dc85fb20aa0808a2 to your computer and use it in GitHub Desktop.
Download ZIP
Add diagnal edges
Raw
implicit_mass_spring.py
# https://www.cs.cmu.edu/~baraff/papers/sig98.pdf
import taichi as ti
import numpy as np
import argparse
@ti.data_oriented
class Cloth:
def __init__(self, N):
self.N = N
self.NF = 2 * N**2 # number of faces
self.NV = (N + 1)**2 # number of vertices
self.NE = 2 * N * (N + 1) + 2 * N * N # numbser of edges
self.pos = ti.Vector.field(2, ti.f32, self.NV)
self.initPos = ti.Vector.field(2, ti.f32, self.NV)
self.vel = ti.Vector.field(2, ti.f32, self.NV)
self.force = ti.Vector.field(2, ti.f32, self.NV)
self.invMass = ti.field(ti.f32, self.NV)
self.spring = ti.Vector.field(2, ti.i32, self.NE)
self.indices = ti.field(ti.i32, 2 * self.NE)
self.Jx = ti.Matrix.field(2, 2, ti.f32,
self.NE) # Jacobian with respect to position
self.Jv = ti.Matrix.field(2, 2, ti.f32,
self.NE) # Jacobian with respect to velocity
self.rest_len = ti.field(ti.f32, self.NE)
self.ks = 1000.0 # spring stiffness
self.kd = 0.5 # damping constant
self.kf = 1.0e5 # fix point stiffness
self.gravity = ti.Vector([0.0, -2.0])
self.init_pos()
self.init_edges()
self.MassBuilder = ti.SparseMatrixBuilder(2 * self.NV,
2 * self.NV,
max_num_triplets=10000)
self.init_mass_sp(self.MassBuilder)
self.M = self.MassBuilder.build()
self.fix_vertex = [self.N, self.NV - 1]
self.Jf = ti.Matrix.field(2, 2, ti.f32, len(
self.fix_vertex)) # fix constraint hessian
@ti.kernel
def init_pos(self):
for i, j in ti.ndrange(self.N + 1, self.N + 1):
k = i * (self.N + 1) + j
self.pos[k] = ti.Vector([i, j]) / self.N * 0.5 + ti.Vector(
[0.25, 0.25])
self.initPos[k] = self.pos[k]
self.vel[k] = ti.Vector([0, 0])
self.invMass[k] = 1.0
self.invMass[self.N] = 0.0
self.invMass[self.NV - 1] = 0.0
@ti.kernel
def init_edges(self):
pos, spring, N, rest_len = ti.static(self.pos, self.spring, self.N,
self.rest_len)
for i, j in ti.ndrange(N + 1, N):
idx, idx1 = i * N + j, i * (N + 1) + j
spring[idx] = ti.Vector([idx1, idx1 + 1])
rest_len[idx] = (pos[idx1] - pos[idx1 + 1]).norm()
start = N * (N + 1)
for i, j in ti.ndrange(N, N + 1):
idx, idx1, idx2 = start + i + j * N, i * (N + 1) + j, i * (
N + 1) + j + N + 1
spring[idx] = ti.Vector([idx1, idx2])
rest_len[idx] = (pos[idx1] - pos[idx2]).norm()
start = 2 * N * (N + 1)
for i, j in ti.ndrange(N, N):
idx, idx1, idx2 = start + i * N + j, i * (N + 1) + j, (i + 1) * (
N + 1) + j + 1
spring[idx] = ti.Vector([idx1, idx2])
rest_len[idx] = (pos[idx1] - pos[idx2]).norm()
start = 2 * N * (N + 1) + N * N
for i, j in ti.ndrange(N, N):
idx, idx1, idx2 = start + i * N + j, i * (N + 1) + j + 1, (
i + 1) * (N + 1) + j
spring[idx] = ti.Vector([idx1, idx2])
rest_len[idx] = (pos[idx1] - pos[idx2]).norm()
@ti.kernel
def init_mass_sp(self, M: ti.sparse_matrix_builder()):
for i in range(self.NV):
if self.invMass[i] != 0.0:
mass = 1.0 / self.invMass[i]
M[2 * i + 0, 2 * i + 0] += mass
M[2 * i + 1, 2 * i + 1] += mass
@ti.func
def clear_force(self):
for i in self.force:
self.force[i] = ti.Vector([0.0, 0.0])
@ti.kernel
def compute_force(self):
self.clear_force()
for i in self.force:
if self.invMass[i] != 0.0:
self.force[i] += self.gravity / self.invMass[i]
for i in self.spring:
idx1, idx2 = self.spring[i][0], self.spring[i][1]
pos1, pos2 = self.pos[idx1], self.pos[idx2]
dis = pos2 - pos1
force = self.ks * (dis.norm() -
self.rest_len[i]) * dis.normalized()
self.force[idx1] += force
self.force[idx2] -= force
# fix constraint gradient
self.force[self.N] += self.kf * (self.initPos[self.N] -
self.pos[self.N])
self.force[self.NV - 1] += self.kf * (self.initPos[self.NV - 1] -
self.pos[self.NV - 1])
@ti.kernel
def compute_Jacobians(self):
for i in self.spring:
idx1, idx2 = self.spring[i][0], self.spring[i][1]
pos1, pos2 = self.pos[idx1], self.pos[idx2]
dx = pos1 - pos2
I = ti.Matrix([[1.0, 0.0], [0.0, 1.0]])
dxtdx = ti.Matrix([[dx[0] * dx[0], dx[0] * dx[1]],
[dx[1] * dx[0], dx[1] * dx[1]]])
l = dx.norm()
if l != 0.0:
l = 1.0 / l
self.Jx[i] = (I - self.rest_len[i] * l *
(I - dxtdx * l**2)) * self.ks
self.Jv[i] = self.kd * I
# fix point constraint hessian
self.Jf[0] = ti.Matrix([[self.kf, 0], [0, self.kf]])
self.Jf[1] = ti.Matrix([[self.kf, 0], [0, self.kf]])
@ti.kernel
def assemble_K(self, K: ti.sparse_matrix_builder()):
for i in self.spring:
idx1, idx2 = self.spring[i][0], self.spring[i][1]
K[2 * idx1 + 0, 2 * idx1 + 0] -= self.Jx[i][0, 0]
K[2 * idx1 + 0, 2 * idx1 + 1] -= self.Jx[i][0, 1]
K[2 * idx1 + 1, 2 * idx1 + 0] -= self.Jx[i][1, 0]
K[2 * idx1 + 1, 2 * idx1 + 1] -= self.Jx[i][1, 1]
K[2 * idx1 + 0, 2 * idx2 + 0] += self.Jx[i][0, 0]
K[2 * idx1 + 0, 2 * idx2 + 1] += self.Jx[i][0, 1]
K[2 * idx1 + 1, 2 * idx2 + 0] += self.Jx[i][1, 0]
K[2 * idx1 + 1, 2 * idx2 + 1] += self.Jx[i][1, 1]
K[2 * idx2 + 0, 2 * idx1 + 0] += self.Jx[i][0, 0]
K[2 * idx2 + 0, 2 * idx1 + 1] += self.Jx[i][0, 1]
K[2 * idx2 + 1, 2 * idx1 + 0] += self.Jx[i][1, 0]
K[2 * idx2 + 1, 2 * idx1 + 1] += self.Jx[i][1, 1]
K[2 * idx2 + 0, 2 * idx2 + 0] -= self.Jx[i][0, 0]
K[2 * idx2 + 0, 2 * idx2 + 1] -= self.Jx[i][0, 1]
K[2 * idx2 + 1, 2 * idx2 + 0] -= self.Jx[i][1, 0]
K[2 * idx2 + 1, 2 * idx2 + 1] -= self.Jx[i][1, 1]
K[2 * self.N + 0, 2 * self.N + 0] += self.Jf[0][0, 0]
K[2 * self.N + 0, 2 * self.N + 1] += self.Jf[0][0, 1]
K[2 * self.N + 1, 2 * self.N + 0] += self.Jf[0][1, 0]
K[2 * self.N + 1, 2 * self.N + 1] += self.Jf[0][1, 1]
K[2 * (self.NV - 1) + 0, 2 * (self.NV - 1) + 0] += self.Jf[1][0, 0]
K[2 * (self.NV - 1) + 0, 2 * (self.NV - 1) + 1] += self.Jf[1][0, 1]
K[2 * (self.NV - 1) + 1, 2 * (self.NV - 1) + 0] += self.Jf[1][1, 0]
K[2 * (self.NV - 1) + 1, 2 * (self.NV - 1) + 1] += self.Jf[1][1, 1]
@ti.kernel
def assemble_D(self, D: ti.sparse_matrix_builder()):
for i in self.spring:
idx1, idx2 = self.spring[i][0], self.spring[i][1]
D[2 * idx1 + 0, 2 * idx1 + 0] -= self.Jv[i][0, 0]
D[2 * idx1 + 0, 2 * idx1 + 1] -= self.Jv[i][0, 1]
D[2 * idx1 + 1, 2 * idx1 + 0] -= self.Jv[i][1, 0]
D[2 * idx1 + 1, 2 * idx1 + 1] -= self.Jv[i][1, 1]
D[2 * idx1 + 0, 2 * idx2 + 0] += self.Jv[i][0, 0]
D[2 * idx1 + 0, 2 * idx2 + 1] += self.Jv[i][0, 1]
D[2 * idx1 + 1, 2 * idx2 + 0] += self.Jv[i][1, 0]
D[2 * idx1 + 1, 2 * idx2 + 1] += self.Jv[i][1, 1]
D[2 * idx2 + 0, 2 * idx1 + 0] += self.Jv[i][0, 0]
D[2 * idx2 + 0, 2 * idx1 + 1] += self.Jv[i][0, 1]
D[2 * idx2 + 1, 2 * idx1 + 0] += self.Jv[i][1, 0]
D[2 * idx2 + 1, 2 * idx1 + 1] += self.Jv[i][1, 1]
D[2 * idx2 + 0, 2 * idx2 + 0] -= self.Jv[i][0, 0]
D[2 * idx2 + 0, 2 * idx2 + 1] -= self.Jv[i][0, 1]
D[2 * idx2 + 1, 2 * idx2 + 0] -= self.Jv[i][1, 0]
D[2 * idx2 + 1, 2 * idx2 + 1] -= self.Jv[i][1, 1]
@ti.kernel
def updatePosVel(self, h: ti.f32, dv: ti.ext_arr()):
for i in self.pos:
if self.invMass[i] != 0.0:
self.vel[i] += ti.Vector([dv[2 * i], dv[2 * i + 1]])
self.pos[i] += h * self.vel[i]
def update(self, h):
self.compute_force()
self.compute_Jacobians()
# Assemble global system
DBuilder = ti.SparseMatrixBuilder(2 * self.NV,
2 * self.NV,
max_num_triplets=10000)
self.assemble_D(DBuilder)
D = DBuilder.build()
KBuilder = ti.SparseMatrixBuilder(2 * self.NV,
2 * self.NV,
max_num_triplets=10000)
self.assemble_K(KBuilder)
K = KBuilder.build()
A = self.M - h * D - h**2 * K
vel = self.vel.to_numpy().reshape(2 * self.NV)
force = self.force.to_numpy().reshape(2 * self.NV)
b = (force + h * K @ vel) * h
# Sparse solver
solver = ti.SparseSolver(solver_type="LDLT")
solver.analyze_pattern(A)
solver.factorize(A)
# Solve the linear system
dv = solver.solve(b)
self.updatePosVel(h, dv)
def display(self, gui, radius=5, color=0xffffff):
lines = self.spring.to_numpy()
pos = self.pos.to_numpy()
edgeBegin = np.zeros(shape=(lines.shape[0], 2))
edgeEnd = np.zeros(shape=(lines.shape[0], 2))
for i in range(lines.shape[0]):
idx1, idx2 = lines[i][0], lines[i][1]
edgeBegin[i] = pos[idx1]
edgeEnd[i] = pos[idx2]
gui.lines(edgeBegin, edgeEnd, radius=2, color=0x0000ff)
gui.circles(self.pos.to_numpy(), radius, color)
@ti.kernel
def spring2indices(self):
for i in self.spring:
self.indices[2 * i + 0] = self.spring[i][0]
self.indices[2 * i + 1] = self.spring[i][1]
def displayGGUI(self, canvas, radius=0.01, color=(1.0, 1.0, 1.0)):
self.spring2indices()
canvas.lines(self.pos,
width=0.005,
indices=self.indices,
color=(0.0, 0.0, 1.0))
canvas.circles(self.pos, radius, color)
if __name__ == "__main__":
ti.init(arch=ti.cpu)
h = 0.01
cloth = Cloth(N=5)
pause = False
parser = argparse.ArgumentParser()
parser.add_argument('-g',
'--use-ggui',
action='store_true',
help='Display with GGUI')
args = parser.parse_args()
use_ggui = False
use_ggui = args.use_ggui
if not use_ggui:
gui = ti.GUI('Implicit Mass Spring System', res=(500, 500))
while gui.running:
for e in gui.get_events():
if e.key == gui.ESCAPE:
gui.running = False
elif e.key == gui.SPACE:
pause = not pause
if not pause:
cloth.update(h)
cloth.display(gui)
gui.show()
else:
window = ti.ui.Window('Implicit Mass Spring System', res=(500, 500))
while window.running:
if window.get_event(ti.ui.PRESS):
if window.event.key == ti.ui.ESCAPE:
break
if window.is_pressed(ti.ui.SPACE):
pause = not pause
if not pause:
cloth.update(h)
canvas = window.get_canvas()
cloth.displayGGUI(canvas)
window.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment