用 Taichi 进行物理模拟

上面的动图很好地模拟了一块布料落到一个球体上。 动图中的布料建模使用了弹簧质点系统，其中包含 10,000 多个质点和大约 100,000个弹簧。 模拟如此大规模的物理系统并实时渲染绝不是一项容易的任务。

Taichi 让物理模拟程序变得更易读和直观，同时仍然达到与 C++ 或 CUDA 相当的性能。 只需拥有基本 Python 编程技能，就可以使用 Taichi 用更少的代码编写高性能并行程序，从而关注较高层次的算法本身，把诸如性能优化的任务交由 Taichi 处理。

本篇文档将展示如何编写一个模拟和渲染布料掉落球面的完整 Python 程序。 在你继续阅读之前，可以先猜猜这个程序包含多少行代码。

开始

要在 Python 程序中使用 Taichi，需要导入 Taichi 到你的命名空间并初始化 Taichi：

1. 导入 Taichi：

import taichi as ti

2. 初始化 Taichi：

# Choose any of the following backend when initializing Taichi
# - ti.cpu
# - ti.gpu
# - ti.cuda
# - ti.vulkan
# - ti.metal
# - ti.opengl
ti.init(arch=ti.cpu)

我们在此选择 ti.cpu ，尽管在 GPU 后端运行 Taichi 会更快。 这主要是因为我们需要确保你可以运行我们的源代码，而无需对平台进行任何编辑或额外配置。 请注意：

• 如果你选择 GPU 后端，例如 ti. cuda，请确保你的系统上已经安装了该后端；否则 Taichi 会报错。
• 用于 3D 渲染的 GGUI 系统目前仅支持 CUDA 和 Vulkan 后端以及 x86 处理器。 如果你选择其他后端，请考虑更换切换源代码中提供的 GGUI 系统。

建模

本节内容包括：

• 概括和简化布料模拟中涉及的模型。
• 用 Taichi 提供的 数据容器 表示掉落的布料和球。

模型简化

以下内容概括和简化了布料模拟中涉及的模型。

布料：一个弹簧质点系统

在这个程序中，布料的模型是一个弹簧质点系统。 更具体地说，我们用一个 n×n 的质点网格表示布料，网格中相邻的点通过弹簧连接。 下图来自 Matthew Fisher，说明了这个结构，网格中的红色顶点代表质点，白色边缘代表弹簧。

这个弹簧质点系统可以用两个数组表示：

• 一个 n×n 的数组表示质点位置，
• 一个 n×n 的数组表示质点速度。

此外，该系统的位置和移动受到以下四个因素的影响：

• 重力
• 弹簧的内力
• 阻尼
• 与中间的球体的碰撞

我们的模拟开始时间是 t = 0，以一个小常量 dt 为间隔推进时间。 每个时间步 dt 结束时，程序执行以下操作：

• 估计上述四个因素对弹簧质点系统的影响，
• 相应地更新每个质点的位置和速度。

表示球体

通常可以利用球体的中心和半径来表示球体。

数据结构

在这个程序中，我们用 Taichi 提供的 数据容器 来表示正在掉落的布料和球体。

布料的数据结构

初始化 Taichi 后，你可以声明表示布料的数据结构。

1. 声明两个数组 x 和 v，用于存储质点的位置和速度。 Taichi 将这样的数组命名为 field。

n = 128
# x is an n x n field consisting of 3D floating-point vectors
# representing the mass points' positions
x = ti.Vector.field(3, dtype=float, shape=(n, n))
# v is an n x n field consisting of 3D floating-point vectors
# representing the mass points' velocities
v = ti.Vector.field(3, dtype=float, shape=(n, n))

2. 初始化已定义的 field x 和 v：

# The n x n grid is normalized
# The distance between two x- or z-axis adjacent points
# is 1.0 / n
quad_size = 1.0 / n

# The @ti.kernel decorator instructs Taichi to
# automatically parallelize all top-level for loops
# inside initialize_mass_points()
@ti.kernel
def initialize_mass_points():
    # A random offset to apply to each mass point
    random_offset = ti.Vector([ti.random() - 0.5, ti.random() - 0.5]) * 0.1

    # Field x stores the mass points' positions
    for i, j in x:
        # The piece of cloth is 0.6 (y-axis) above the original point
        #
        # By taking away 0.5 from each mass point's [x,z] coordinate value
        # you move the cloth right above the original point
        x[i, j] = [
            i * quad_size - 0.5 + random_offset[0], 0.6,
            j * quad_size - 0.5 + random_offset[1]
        ]
        # The initial velocity of each mass point is set to 0
        v[i, j] = [0, 0, 0]

球体的数据结构

此处，球体中心是个一维 field，其唯一元素是一个三维浮点数向量。

ball_radius = 0.3
# Use a 1D field for storing the position of the ball center
# The only element in the field is a 3-dimentional floating-point vector
ball_center = ti.Vector.field(3, dtype=float, shape=(1, ))
# Place the ball center at the original point
ball_center[0] = [0, 0, 0]

模拟

每个时间步 dt 结束时，程序模拟上述四个因素对弹簧质点系统的影响。 在本例中，定义了一个 kernel 函数 substep()。

# substep() works out the *accumulative* effects
# of gravity, internal force, damping, and collision
# on the mass-spring system
@ti.kernel
def substep():

重力

# Gravity is a force applied in the negative direction of the y axis,
# and so is set to [0, -9.8, 0]
gravity = ti.Vector([0, -9.8, 0])

# The @ti.kernel decorator instructs Taichi to
# automatically parallelize all top-level for loops
# inside substep()
@ti.kernel
def substep():
    # The for loop iterates over all elements of the field v
    for i in ti.grouped(x):
        v[i] += gravity * dt

note

for i in ti.grouped(x) 是 Taichi 的一个重要功能。 这意味着这个 for 循环将 x 作为一个一维数组自动遍历其中所有元素，而不论其形状如何，省去你手写多层 for 循环的麻烦。

使用 for i in ti.grouped(x) 或 for i in ti.grouped(v) 都可以，因为 field x 和 field v 的形状相同。

弹簧的内力

如下图所示，我们做出以下假设：

• 一个给定的点最多可能受到 12 个邻点的影响，其他质点施加的影响不计。
• 这些邻点通过弹簧对给定点施加内力。

note

此处的内力宽泛地表示弹簧弹性形变导致的内力以及两点间相对运动导致的阻尼。

以下代码的作用包括：

• 遍历弹簧质点网格，
• 累积邻点施加在每个质点上的内力，
• 将内力转变为速度。

quad_size = 1.0 / n
# Elastic coefficient of springs
spring_Y = 3e4
# Damping coefficient caused by
# the relative movement of the two mass points
# The assumption here is:
# A mass point can have at most 12 'influential` points
dashpot_damping = 1e4

# The cloth is modeled as a mass-spring grid. Assume that:
# a mass point, whose relative index is [0, 0],
# can be affected by at most 12 surrounding points
#
# spring_offsets is such a list storing
# the relative indices of these 'influential' points
spring_offsets = []
for i in range(-1, 2):
    for j in range(-1, 2):
        if (i, j) != (0, 0):
            spring_offsets.append(ti.Vector([i, j]))

@ti.kernel
def substep():

    # Traverses the field x as a 1D array
    #
    # The `i` here refers to the *absolute* index
    # of an element in the field x
    #
    # Note that `i` is a 2-dimentional vector here
    for i in ti.grouped(x):
        # Initial force exerted to a specific mass point
        force = ti.Vector([0.0, 0.0, 0.0])
        # Traverse the surrounding mass points
        for spring_offset in ti.static(spring_offsets):
            # j is the *absolute* index of an 'influential' point
            # Note that j is a 2-dimensional vector here
            j = i + spring_offset
            # If the 'influential` point is in the n x n grid,
            # then work out the internal force that it exerts
            # on the current mass point
            if 0 <= j[0] < n and 0 <= j[1] < n:
                # The relative displacement of the two points
                # The internal force is related to it
                x_ij = x[i] - x[j]
                # The relative movement of the two points
                v_ij = v[i] - v[j]
                # d is a normalized vector (its norm is 1)
                d = x_ij.normalized()
                current_dist = x_ij.norm()
                original_dist = quad_size * float(i - j).norm()
                # Internal force of the spring
                force += -spring_Y * d * (current_dist / original_dist - 1)
                # Continues to apply the damping force
                # from the relative movement
                # of the two points
                force += -v_ij.dot(d) * d * dashpot_damping * quad_size

        # Continues to add the velocity caused by the internal forces
        # to the current velocity
        v[i] += force * dt

阻尼

在现实世界中，当弹簧振动时，弹簧中储存的能量随着振动消失而扩散到周围环境。 为了取得这种效果，我们在每个时间步都略微降低网格中每个点的速度：

# Damping coefficient of springs
drag_damping = 1

@ti.kernel
def substep():

    # Traverse the elements in field v
    for i in ti.grouped(x):
        v[i] *= ti.exp(-drag_damping * dt)

与球体的碰撞

我们假定与球体的碰撞是一个弹性碰撞：当一个质点与球碰撞时， 其在碰撞点的法向量上的速度分量发生变化。

需注意，每个质点的位置根据每个时间步结束时的速度更新。

# Damping coefficient of springs
drag_damping = 1

@ti.kernel
def substep():

    # Traverse the elements in field v
    for i in ti.grouped(x):

        offset_to_center = x[i] - ball_center[0]
        if offset_to_center.norm() <= ball_radius:
            # Velocity projection
            normal = offset_to_center.normalized()
            v[i] -= min(v[i].dot(normal), 0) * normal
        # After working out the accumulative v[i],
        # work out the positions of each mass point
        x[i] += dt * v[i]

就是这样！ 以上就是实现一个弹簧质点网格系统的并行模拟程序所需的全部代码。

渲染

我们使用 Taichi 的 GPU GUI 系统（也称为 GGUI）进行 3D 渲染。 GGUI 支持渲染两种类型的 3D 对象：三角网格和粒子。 在本例中，我们可以将布料作为三角网格、将球体作为粒子渲染。

GGUI 用两个 Taichi field 来表示一个三角网格：vertices 和 indices。 vertices 是一个一维 field，其中每个元素都是一个代表顶点位置的三维向量，可能多个三角形共用一个顶点。 在我们的应用中，每个质点都是一个三角形的顶点，所以我们可以把代码直接从 x 拷贝到 vertices中：

vertices = ti.Vector.field(3, dtype=float, shape=n * n)

@ti.kernel
def update_vertices():
    for i, j in ti.ndrange(n, n):
        vertices[i * n + j] = x[i, j]

需注意，每帧都要调用 update_vertices，因为顶点位置在模拟中持续更新。

布料由一个 n × n 的质点网格表示，也可把该网格看作一个由小正方形组成的 n-1 × n-1 的网格。 每个正方形将作为两个三角形渲染。 因此，共有 (n - 1) * (n - 1) * 2 个三角形。 每个三角形都由 vertices field 中的三个整数表示，记录了 vertices field 中三角形的顶点索引。 以下代码片段描述了这一结构：

@ti.kernel
def initialize_mesh_indices():
    for i, j in ti.ndrange(n - 1, n - 1):
        quad_id = (i * (n - 1)) + j
        # First triangle of the square
        indices[quad_id * 6 + 0] = i * n + j
        indices[quad_id * 6 + 1] = (i + 1) * n + j
        indices[quad_id * 6 + 2] = i * n + (j + 1)
        # Second triangle of the square
        indices[quad_id * 6 + 3] = (i + 1) * n + j + 1
        indices[quad_id * 6 + 4] = i * n + (j + 1)
        indices[quad_id * 6 + 5] = (i + 1) * n + j

    for i, j in ti.ndrange(n, n):
        if (i // 4 + j // 4) % 2 == 0:
            colors[i * n + j] = (0.22, 0.72, 0.52)
        else:
            colors[i * n + j] = (1, 0.334, 0.52)

需注意，不同于 update_vertices()，initialize_mesh_indices() 只需被调用一次。 这是因为三角顶点的索引实际上没有变化——只有位置发生了改变。

至于球体的渲染，先前定义的 ball_center 和 ball_radius 变量已经足够。

源代码

import taichi as ti
ti.init(arch=ti.vulkan)  # Alternatively, ti.init(arch=ti.cpu)

n = 128
quad_size = 1.0 / n
dt = 4e-2 / n
substeps = int(1 / 60 // dt)

gravity = ti.Vector([0, -9.8, 0])
spring_Y = 3e4
dashpot_damping = 1e4
drag_damping = 1

ball_radius = 0.3
ball_center = ti.Vector.field(3, dtype=float, shape=(1, ))
ball_center[0] = [0, 0, 0]

x = ti.Vector.field(3, dtype=float, shape=(n, n))
v = ti.Vector.field(3, dtype=float, shape=(n, n))

num_triangles = (n - 1) * (n - 1) * 2
indices = ti.field(int, shape=num_triangles * 3)
vertices = ti.Vector.field(3, dtype=float, shape=n * n)
colors = ti.Vector.field(3, dtype=float, shape=n * n)

bending_springs = False

@ti.kernel
def initialize_mass_points():
    random_offset = ti.Vector([ti.random() - 0.5, ti.random() - 0.5]) * 0.1

    for i, j in x:
        x[i, j] = [
            i * quad_size - 0.5 + random_offset[0], 0.6,
            j * quad_size - 0.5 + random_offset[1]
        ]
        v[i, j] = [0, 0, 0]


@ti.kernel
def initialize_mesh_indices():
    for i, j in ti.ndrange(n - 1, n - 1):
        quad_id = (i * (n - 1)) + j
        # 1st triangle of the square
        indices[quad_id * 6 + 0] = i * n + j
        indices[quad_id * 6 + 1] = (i + 1) * n + j
        indices[quad_id * 6 + 2] = i * n + (j + 1)
        # 2nd triangle of the square
        indices[quad_id * 6 + 3] = (i + 1) * n + j + 1
        indices[quad_id * 6 + 4] = i * n + (j + 1)
        indices[quad_id * 6 + 5] = (i + 1) * n + j

    for i, j in ti.ndrange(n, n):
        if (i // 4 + j // 4) % 2 == 0:
            colors[i * n + j] = (0.22, 0.72, 0.52)
        else:
            colors[i * n + j] = (1, 0.334, 0.52)

initialize_mesh_indices()

spring_offsets = []
if bending_springs:
    for i in range(-1, 2):
        for j in range(-1, 2):
            if (i, j) != (0, 0):
                spring_offsets.append(ti.Vector([i, j]))

else:
    for i in range(-2, 3):
        for j in range(-2, 3):
            if (i, j) != (0, 0) and abs(i) + abs(j) <= 2:
                spring_offsets.append(ti.Vector([i, j]))

@ti.kernel
def substep():
    for i in ti.grouped(x):
        v[i] += gravity * dt

    for i in ti.grouped(x):
        force = ti.Vector([0.0, 0.0, 0.0])
        for spring_offset in ti.static(spring_offsets):
            j = i + spring_offset
            if 0 <= j[0] < n and 0 <= j[1] < n:
                x_ij = x[i] - x[j]
                v_ij = v[i] - v[j]
                d = x_ij.normalized()
                current_dist = x_ij.norm()
                original_dist = quad_size * float(i - j).norm()
                # Spring force
                force += -spring_Y * d * (current_dist / original_dist - 1)
                # Dashpot damping
                force += -v_ij.dot(d) * d * dashpot_damping * quad_size

        v[i] += force * dt

    for i in ti.grouped(x):
        v[i] *= ti.exp(-drag_damping * dt)
        offset_to_center = x[i] - ball_center[0]
        if offset_to_center.norm() <= ball_radius:
            # Velocity projection
            normal = offset_to_center.normalized()
            v[i] -= min(v[i].dot(normal), 0) * normal
        x[i] += dt * v[i]

@ti.kernel
def update_vertices():
    for i, j in ti.ndrange(n, n):
        vertices[i * n + j] = x[i, j]

window = ti.ui.Window("Taichi Cloth Simulation on GGUI", (1024, 1024),
                      vsync=True)
canvas = window.get_canvas()
canvas.set_background_color((1, 1, 1))
scene = ti.ui.Scene()
camera = ti.ui.Camera()

current_t = 0.0
initialize_mass_points()

while window.running:
    if current_t > 1.5:
        # Reset
        initialize_mass_points()
        current_t = 0

    for i in range(substeps):
        substep()
        current_t += dt
    update_vertices()

    camera.position(0.0, 0.0, 3)
    camera.lookat(0.0, 0.0, 0)
    scene.set_camera(camera)

    scene.point_light(pos=(0, 1, 2), color=(1, 1, 1))
    scene.ambient_light((0.5, 0.5, 0.5))
    scene.mesh(vertices,
               indices=indices,
               per_vertex_color=colors,
               two_sided=True)

    # Draw a smaller ball to avoid visual penetration
    scene.particles(ball_center, radius=ball_radius * 0.95, color=(0.5, 0.42, 0.8))
    canvas.scene(scene)
    window.show()

代码行数：145。