Homework2: TaichiMD - 实时交互式粒子模拟引擎

肝了好几天终于把想要做的部分都弄好了,来交个overdue的作业

简介

TaichiMD 是基于 Taichi 编程语言,以分子动力学为核心的实时交互式粒子模拟引擎。其特性和目标包括:

  • 通过 Taichi 编程语言搭建计算机图形学到和分子模拟/计算化学的桥梁
  • 通过GPU加速实现美观,实时,可交互的粒子模拟
  • 为新的模拟算法的快速实现和可微分模拟提供平台

最终效果

512个丙烷分子的分子动力学模拟,使用我们组的TraPPE-UA力场加上简谐形式的化学键伸缩(其实我一直没加每个粒子的质量,所以不是完全准确的)
propane
(gif被压得完全变色了,原图看这儿)

Taichi 的“99行代码冰雪奇缘”经典示例 MPM99,使用32768个MPM粒子和64 * 64 * 64网格

代码示例

现在版本的TaichiMD在作业1的基础上增加的大量的功能,比如多组分混合物的模拟,化学键弯曲的bending potential,以及一系列基于网格的方法,包括Neighbor List加速模拟,PIC和APIC。最后,MPM99的三种材料模拟被压缩到了48行:

import taichi as ti
from taichimd import Simulation, MDRenderer, DIM
from taichimd.grid import APIC, QuadraticKernel
from taichimd.integrator import ForwardEulerIntegrator as FE
MDRenderer.radius = 0.2

class MPMGrid(APIC):
    @ti.func
    def affine(self, i):
        self.system.F[i] = (ti.Matrix.identity(ti.f32, DIM) + self.system.dt * self.system.C[i]) @ self.system.F[i]
        h = 0.3 if self.system.type[i] == 1 else ti.exp(10 * (1.0 - self.system.Jp[i]))
        la = self.lambda_0 * h
        U, sig, V = ti.svd(self.system.F[i])
        J = 1.0
        for d in ti.static(range(DIM)):
            new_sig = min(max(sig[d, d], 1 - 2.5e-2), 1 + 4.5e-3) if self.system.type[i] == 2 else sig[d, d]
            self.system.Jp[i] *= sig[d, d] / new_sig
            sig[d, d] = new_sig
            J *= new_sig
        if self.system.type[i] == 0:  # Reset deformation gradient to avoid numerical instability
            self.system.F[i] = ti.Matrix.identity(ti.f32, DIM)
            self.system.F[i][0, 0] = J
        elif self.system.type[i] == 2:
            self.system.F[i] = U @ sig @ V.T() # Reconstruct elastic deformation gradient after plasticity
        stress = ti.Matrix.identity(ti.f32, DIM) * la * J * (J - 1)
        if self.system.type[i] != 0:
            stress += 2 * self.mu_0 * h * (self.system.F[i] - U @ V.T()) @ self.system.F[i].T()
        stress = (-self.system.dt * self.vol * 4 * self.inv_dx * self.inv_dx) * stress
        return stress + self.mass * self.system.C[i]

ti.init(arch=ti.cuda, device_memory_GB=3)
dt = 2e-4
n_particles, n_grid = 32768, 64
p_vol, p_rho = (0.5 / n_grid)**2, 1
mpmgrid = MPMGrid(QuadraticKernel(), n_grid, mass=p_vol * p_rho, gravity=50)
mpmgrid.vol = p_vol
E, nu = 0.1e4, 0.2 # Young's modulus and Poisson's ratio
mpmgrid.mu_0, mpmgrid.lambda_0 = E / (2 * (1 + nu)), E * nu / ((1+nu) * (1 - 2 * nu)) # Lame parameters
sim = Simulation(n_particles, integrator=FE(dt, False), grid=mpmgrid)
sim.add_attr("F", dims=(DIM, DIM), dtype=ti.f32) # affine velocity field, deformation gradient
sim.add_attr("Jp", dims=(), dtype=ti.f32) # plastic deformation
sim.gui.set_colors([[6/255, 133/255, 135/255], [237/255, 85/255, 89/255], [238/255, 238/255, 240/255]]) # mpm99 colors
sim.init_random(center=(0.4, 0.15, 0.7), length=0.2, start=0, end=n_particles//3, inittype=0)
sim.init_random(center=(0.4, 0.45, 0.7), length=0.2, start=n_particles//3, end=2*n_particles//3, inittype=1)
sim.init_random(center=(0.5, 0.75, 0.7), length=0.2, start=2*n_particles//3, end=n_particles, inittype=2)
sim.build()
sim.velocity.fill(0); sim.F.fill(((1,0,0),(0,1,0),(0,0,1))); sim.C.fill(0); sim.Jp.fill(1)
sim.run(irender=10, pause=True)
  1. 前五行都是import,还有设定粒子渲染的半径比例。之后我们定义了一个MPMGrid的类,MLS-MPM和APIC关系很近,所以我们就直接继承自带的APIC网格类。MLS-MPM和APIC的区别是每个粒子的affine matrix如何计算,而P2G和G2P的其他部分可以完全一样,所以就只重写affine()这个方法,抄一遍mpm99例子的代码即可。
  2. 定义完MPMGrid之后就进行初始化和设定mpm99对应的各种参数。然后我们新建一个Simulation对象作为我们的模拟体系,这个体系使用不含力的显示积分和我们刚才写的MPM网格。
  3. 模拟体系看到我们使用了APIC网格的子类,会给我们自动新建每个粒子的位置、速度和affine matrix场(C[i])以及每个粒子的类型。然而为了实现mpm99我们还需要给每个粒子加入F和Jp这两个属性(sim.add_attr(...))
  4. 然后我们设定每种粒子的颜色和初始化的位置,再调用sim.build()这个方法,会触发taichi的materialize过程,并place上模拟体系中的各种粒子和网格变量。
  5. 在开始模拟之前我们把各种物理量设为初始值,然后用sim.run()开始模拟。pause=True意味着模拟不会自动运行,按下空格键才会开始运行(并且编译运行时需要的各种kernel)

图形渲染

既然要做基于Taichi的交互式模拟引擎,那我们肯定需要一个完全用Taichi实现的实时渲染器。既然要实时,那么taichi例子给的各种path tracing肯定是不行的,这就得自己造轮子。好在@archibate于斌的Taichi-Three开了个好头,就可以站在巨人的肩膀上少造一些轮子。

TaichiMD是一个粒子模拟引擎,只要能渲染球形粒子(和圆柱体的化学键)就可以了。那么除了传统的三角形栅格化方法,我们就可以做一些更“直接”的解决方案,否则每个粒子是一个几百面的球,几万个粒子渲染的开销是巨大的。所以TaichiMD使用的是完全implicit的几何进行实时渲染,换句话说上面那两个GIF里面,其实一个三角形都没有。这样最直接的好处就是随粒子数量的scalability非常强,3万多个粒子的栅格化渲染在我的2080S上可以做到40–50fps,然而坏处就是因为taichi只能并行遍历每个粒子的for loop,假如1个粒子占了大半个屏幕就反而会卡了。

Taichi-Three自带只有最简单的phong lighting model,不够好看,所以我又自己写了一个Physically-based的Cook-Torrance BRDF。然而到这个时候我们还是只有direct lighting,依然不够好看(本科帮实验室里画图的,强迫症)。想要模拟indirect lighting最简单的办法是搞一个SSAO,但是粒子模拟的场景只有球体,肯定有更优雅的解法。于是我就照着GPU Gems 2里的一章最简单近似下推了个球形体系全局光(GI)的解,用Taichi实现了居然真能管用!下面请欣赏乞丐版RTX ON对比图(模拟界面内按tab打开全局光):
[全局光照]


[全局光照]

[全局光照]


[全局光照]

可以看到效果是很明显的,第一张图的红色大果冻的diffuse反光会在地面上显示,粒子的小球与小球之间的occlusion也能显示。虽然不一定比SSAO快(需要对每个fragment迭代一遍每个粒子),但这个方法有一点好,因为它是closed-form的解,所以直接渲染到frame buffer里,不需要后处理就几乎没有任何噪声。至于详细的推导过程有时间了可以写一篇文章发到我的博客上。

总结

GAMES201终于迎来了大结局。在此期间我学到了很多东西,TaichiMD也从homework 0的200行小程序到现在集粒子、网格、微观模拟、宏观模拟、实时渲染为一体的逐步成熟的项目,说这门课它是我本科和研究生加起来最有用的几个都不为过(本科学的化学所以90%课都已经没用了hhhh)。

最后这么长的帖子能你一直读到这也不容易,祝学业、工作、科研顺利!

15 Likes