# Homework 0 rope simulator

``````# rope_simulator.py

import taichi as ti

ti.init(arch=ti.gpu)

n = 320
image = ti.var(dt=ti.f32, shape=(n * 2, n))
# 10个质点
# 各质点位置
positions = ti.Vector(2, dt=ti.f32, shape=10)
# 各质点受力
forces = ti.Vector(2, dt=ti.f32, shape=10)
# 各质点速度
velocities = ti.Vector(2, dt=ti.f32, shape=10)
# 重力加速度
gravity = ti.Vector([0, -1])
# 原长
rest_length = ti.var(dt=ti.f32, shape=())
# 弹簧系数
k = ti.var(dt=ti.f32, shape=())
# 质点质量
mass = ti.var(dt=ti.f32, shape=())
# 步长
delta_t = ti.var(dt=ti.f32, shape=())

def my_abs(x):
if x < 0:
return -x
else:
return x

@ti.func
def draw_line(a, b):
x0 = a[0]
x1 = b[0]
y0 = a[1]
y1 = b[1]
steep = False
if my_abs(x0 - x1) < my_abs(y0 - y1):
temp = x0
x0 = y0
y0 = temp
temp = x1
x1 = y1
y1 = temp
steep = True
if x0 > x1:
temp = x0
x0 = x1
x1 = temp
temp = y0
y0 = y1
y1 = temp
for x in range(x0, x1):
t = (x - x0) / (x1 - x0)
y = y0 + (y1 - y0) * t
if steep:
image[ti.cast(y, ti.i32), ti.cast(x, ti.i32)] = 0.3
else:
image[ti.cast(x, ti.i32), ti.cast(y, ti.i32)] = 0.3

@ti.kernel
def init():
rest_length[None] = 20
k[None] = 100
mass[None] = 1
delta_t[None] = 1 / 16
for i in range(10):
positions[i] = [220 + i * rest_length[None], 280]

@ti.kernel
def update():
for i, j in image:
image[i, j] = 0
for i in range(9):
b_a = positions[i + 1] - positions[i]
fb_a = -k[None] * b_a * (b_a.norm() - rest_length[None]) / b_a.norm()
forces[i] -= fb_a
forces[i + 1] += fb_a
for i in range(1, 10):
forces[i] += mass[None] * gravity
accel = forces[i] / mass[None]
cur_velocity = velocities[i] + accel * delta_t[None]

forces[i] -= 0.005 * cur_velocity
accel = forces[i] / mass[None]

velocities[i] += accel * delta_t[None]
positions[i] += velocities[i] * delta_t[None]
for force in forces:
forces[force] = [0, 0]
for i in range(9):
draw_line(positions[i], positions[i + 1])
for i in range(10):
image[ti.cast(positions[i][0], ti.i32), ti.cast(positions[i][1], ti.i32)] = 1

gui = ti.GUI("Rope Simulator", (n * 2, n))

init()
for i in range(100000000):
update()
gui.set_image(image)
gui.show()
``````
4 Likes

Taichi GUI可以直接用来画line, circle等功能哈。

``````# rope_simulator.py

import taichi as ti
import numpy as np

ti.init(arch=ti.gpu)

w = 640
h = 480
# 10个质点
# 各质点位置
positions = ti.Vector(2, dt=ti.f32, shape=10)
draw_positions = ti.Vector(2, dt=ti.f32, shape=10)
# 各质点受力
forces = ti.Vector(2, dt=ti.f32, shape=10)
# 各质点速度
velocities = ti.Vector(2, dt=ti.f32, shape=10)
# 重力加速度
gravity = ti.Vector([0, -1])
# 原长
rest_length = ti.var(dt=ti.f32, shape=())
# 弹簧系数
k = ti.var(dt=ti.f32, shape=())
# 质点质量
mass = ti.var(dt=ti.f32, shape=())
# 步长
delta_t = ti.var(dt=ti.f32, shape=())

def my_abs(x):
if x < 0:
return -x
else:
return x

@ti.kernel
def init():
rest_length[None] = 30
k[None] = 100
mass[None] = 1
delta_t[None] = 1 / 10
x = 0.5 * w
y = 0.8 * h
for i in range(10):
positions[i] = [x + i * rest_length[None], y]

@ti.kernel
def update():
for i in range(9):
b_a = positions[i + 1] - positions[i]
fb_a = -k[None] * b_a * (b_a.norm() - rest_length[None]) / b_a.norm()
forces[i] -= fb_a
forces[i + 1] += fb_a
for i in range(10):
if i:
forces[i] += mass[None] * gravity
acceleration = forces[i] / mass[None]
cur_velocity = velocities[i] + acceleration * delta_t[None]

forces[i] -= 0.005 * cur_velocity
acceleration = forces[i] / mass[None]

velocities[i] += acceleration * delta_t[None]
positions[i] += velocities[i] * delta_t[None]
forces[i] = [0, 0]
draw_positions[i] = positions[i]
for i in range(10):
draw_positions[i][0] /= w
draw_positions[i][1] /= h

gui = ti.GUI("Rope Simulator", (w, h), 0xffffff)

init()
for i in range(100000000):
for e in gui.get_events(ti.GUI.PRESS):
if e.key in [ti.GUI.ESCAPE]:
exit()
update()
for i in range(10):
if i:
gui.line(begin=[draw_positions[i - 1][0], draw_positions[i - 1][1]],
end=[draw_positions[i][0], draw_positions[i][1]],