taichi 冰雪奇缘学习

一、Taichi设计目标

性能

生产力

空间稀疏计算

可微编程 --> 深度学习

元编程

二、设计决策

计算与数据结构解耦合

自动特定的编译器优化

Megakernels

实现两个尺度自动微分

嵌入进python

三、99行代码分析

https://zhuanlan.zhihu.com/p/97700605

https://www.bilibili.com/read/cv4358206?spm_id_from=333.788.b_636f6d6d656e74.55

1、基本数据类型

taichi库最常用到的数据类型是矩阵，ti.var、ti.Vector、ti.Matrix这三个函数生成的其实都是矩阵，别被名字误导

1.1、ti.Vector

功能：创建一个矩阵，矩阵的每个元素是向量

ti.Vector(每个元素是几维向量，dt=数据类型，shape=矩阵形状)
dt：data_type的意思，dt=ti.f32指32位float
shape:确定矩阵的行数和列数，如果shape是一个数字，就是单行矩阵，如果shape是一个元祖（3,4），就是3行4列矩阵
    
exp: ti.Vector(2,dt=ti.f32,shape=1000) 他是1000个2维向量组成的1行矩阵，每个数字是float32，访问最后一个元素写v[999][1]即可

1.2、ti.Matrix

功能：创建一个矩阵，矩阵的每个元素也是矩阵

ti.Matrix(每个元素行数，每个元素列数，dt=数据类型，shape=矩阵形状)

1.3、ti.var

功能：创建一个矩阵，每个元素是一个数值

ti.var(dt=数据类型，shape=矩阵形状)

1.4、数据结构分析

x=ti.Vector(2,dt=ti.f32,shape=n_particles) #position位置，每一个粒子有一个位置
v=ti.Vector(2,dt=ti.f32,shape=n_particles) #velocity速度，每一个粒子有一个速度
C=ti.Matrix(2,2,dt=ti.f32,shape=n_particles) #affine速度场，每一个粒子对应一个2X2矩阵
F=ti.Matrix(2,2,dt=ti.f32,shape=n_particles) #deformation gradient变形梯度矩阵
material=ti.var(dt=ti.i32,shape=n_particles) # material id 这个粒子里有三种材料，分别是0,1,2
Jp=ti.var(dt=ti.f32,shape=n_particles) # plastic deformation 塑形变形，不可恢复变形
grid_v=ti.Vector(2,dt=ti.f32,shape=(n_grid,n_grid)) #grid node
momemtum/velocity: 128X128矩阵,每一个元素是一个Vector2，每个格子一个总体速度
grid_m=ti.var(dt=ti.f32,shape=(n_grid,n_grid)) #grid node mass格子质量，128X128矩阵，每个元素是一个数字，每个格子一个质量

2、粒子particle，格子grid

# 99行程序开头定义
#粒子数量，网格数量=128*1
n_particles,n_grid=9000*quality**2,128*quality
#每个网格宽度dX=1/128,以及它的倒数inv_dx
dx,inv_dx=1/n_grid,float(n_grid)

所有的坐标、距离都是用0~1的小数表示

可以看出，粒子总数很多，但是格子只有128*128个，如下所示：

示意图中，格子画的很稀疏，实际上每个格子很小。红色圈是有粒子的格子，蓝色圈是没有粒子的格子，蓝色圈可以跳过计算，极大提高运算效率。

原程序的思路是：某些算法与粒子紧密相关，每帧每个粒子都要计算；而另外一些属性是总体属性，只对每个包含粒子的格子计算一次。比如重力的影响就是整体属性，要通过格子算。而且，每个格子只对周围相邻的格子有影响，影响不会超过一个格子的距离。如果没有这个假设，粒子的实时模拟就不可能做到了。

之后会看到，每一帧的计算分三段：

• ​ 粒子相关计算，粒子被归入某一个格子（particle to grid）；
• ​ 格子之内的计算，以及周边一共9个格子互相的影响
• ​ 格子的速度、速度场等属性，要再影响到每一个粒子（grid to particle）

3、taichi这个库到底做了什么

可以看出，所有粒子计算的方法，全都写在了python代码中，由此可见taichi这个库并非一个即插即用的“物理模拟引擎”，而是一种用于科学计算的基础设施。

熟悉深度学习的朋友应该对这种模式更熟悉一些，最常用的TensorFlow的底层也是在做同样的事，实际上这些库主要是解决了底层计算问题，要拿来直接用还得再做一层封装。

和深度学习的python代码相同，程序并非是在执行python脚本时进行的计算，实际的底层运算是被延后的。python脚本所做的事情都可以看成“准备工作”。例如代码中间这一行：

# 二次核函数   
w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1), 0.5 * ti.sqr(fx - 0.5)]

如果你想通过print 打印 w 的值只会无功而返，因为运行完这句话，w的值还没有开始计算呢~~这也为分析代码带来很多麻烦，因为不能通过打印log的方式观察每一步计算的结果。

4、完整代码+注释

import taichi as ti

# 计算品质，越大算得越准确，但也越慢。

quality = 1 # Use a larger value for higher-res simulations

# 粒子数量=9000，网格数量=128

n_particles, n_grid = 9000 * quality ** 2, 128 * quality

# 每个网格宽度ΔX=1/128，以及它的倒数inv_dx

dx, inv_dx = 1 / n_grid, float(n_grid)

# deltaTime，演算的时间间隔

dt = 1e-4 / quality

# 体积vol，密度 (rho就是ρ)

p_vol, p_rho = (dx * 0.5)**2, 1

# 质量 = 体积 * 密度

p_mass = p_vol * p_rho

#以下是材料系数

# E=0.1e4=1000, 泊松分布系数nu=0.2

E, nu = 0.1e4, 0.2 # Young's modulus and Poisson's ratio

# Lame系数，定义材料性质的，分别是μ和λ

mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / ((1+nu) * (1 - 2 * nu)) # Lame parameters

 

 

# 小技巧：taichi的对象类型不易查看，可以调用矩阵对象的.to_numpy()方法，把它变成numpy对象再看

x = ti.Vector(2, dt=ti.f32, shape=n_particles) # position位置, 每个粒子有一个位置

v = ti.Vector(2, dt=ti.f32, shape=n_particles) # velocity速度，每个粒子有一个速度

C = ti.Matrix(2, 2, dt=ti.f32, shape=n_particles) # affine速度场，每个粒子对应一个2x2矩阵

F = ti.Matrix(2, 2, dt=ti.f32, shape=n_particles) # deformation gradient变形梯度矩阵

 

material = ti.var(dt=ti.i32, shape=n_particles) # material id，这个例子里有3种材料，分别是0、1、2

Jp = ti.var(dt=ti.f32, shape=n_particles) # plastic deformation 塑性变形，不可恢复变形

grid_v = ti.Vector(2, dt=ti.f32, shape=(n_grid, n_grid)) # grid node momemtum/velocity 一个128x128矩阵，每个元素是一个Vector2。每个格子一个总体速度

grid_m = ti.var(dt=ti.f32, shape=(n_grid, n_grid)) # grid node mass格子质量。128x128矩阵，每个元素是一个数字，每个格子一个质量

 

ti.cfg.arch = ti.cuda # Try to run on GPU，Nvidia显卡的CUDA

 

 

# 下面的函数是每帧更新用的，先跳过这个函数看主程序初始化~~

# 留到最后看这个函数

@ti.kernel

def substep():

  # 每次都要将每个格子的总速度和总质量归0

  for i, j in ti.ndrange(n_grid, n_grid):

    grid_v[i, j] = [0, 0]

    grid_m[i, j] = 0

  # 更新每一个粒子，并让它们属于某一个格子(Particle to Grid)

  for p in range(n_particles): # Particle state update and scatter to grid (P2G)

    # 用粒子坐标x[p]，计算出它所属的格子的左下角

    base = (x[p] * inv_dx - 0.5).cast(int)

    # fx是该粒子距离格子左下角的局部坐标

    fx = x[p] * inv_dx - base.cast(float)

 

    # 二次核函数

    # Quadratic kernels  [http://mpm.graphics   Eqn. 123, with x=fx, fx-1,fx-2]

    w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1), 0.5 * ti.sqr(fx - 0.5)]

    # w是一个延迟计算的表达式，无法直接看到它的值，后面很多参数都是类似的

    # w是位置的函数。

    # w是一个权重，会影响当前格子与周围格子的质量、速度

    # 猜测：一个格子边缘的粒子，显然对旁边格子的影响更大

 

    #每帧更新F，F = somefunc(F, C, dt)

    F[p] = (ti.Matrix.identity(ti.f32, 2) + dt * C[p]) @ F[p] # deformation gradient update

 

    # h "加硬"系数（雪积累在一起变硬），是根据塑性变形参数Jp算出来的

    h = ti.exp(10 * (1.0 - Jp[p])) # Hardening coefficient: snow gets harder when compressed

    # jelly 果冻的加硬系数固定0.3

    if material[p] == 1: # jelly, make it softer

      h = 0.3

    # μ和λ的值根据h确定

    mu, la = mu_0 * h, lambda_0 * h

    # 液体的μ固定为0

    if material[p] == 0: # liquid

      mu = 0.0

    # ---------------硬核计算开始，按理说不要随意改动-----------------

    U, sig, V = ti.svd(F[p])

    J = 1.0

    for d in ti.static(range(2)):

      new_sig = sig[d, d]

      if material[p] == 2:  # Snow

        new_sig = min(max(sig[d, d], 1 - 2.5e-2), 1 + 4.5e-3)  # Plasticity 塑性

      Jp[p] *= sig[d, d] / new_sig

      sig[d, d] = new_sig

      J *= new_sig

    if material[p] == 0:  # Reset deformation gradient to avoid numerical instability

      F[p] = ti.Matrix.identity(ti.f32, 2) * ti.sqrt(J)

    elif material[p] == 2:

      F[p] = U @ sig @ V.T() # Reconstruct elastic deformation gradient after plasticity

    stress = 2 * mu * (F[p] - U @ V.T()) @ F[p].T() + ti.Matrix.identity(ti.f32, 2) * la * J * (J - 1)

    stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress

    affine = stress + p_mass * C[p]

    # ===============硬核计算结束==================

 

    # 遍历相邻8+1个格子，把粒子参数算到到格子上

    for i, j in ti.static(ti.ndrange(3, 3)): # Loop over 3x3 grid node neighborhood

      offset = ti.Vector([i, j])

      dpos = (offset.cast(float) - fx) * dx

      weight = w[i][0] * w[j][1]

      # 每个粒子的速度、质量叠加，得到当前格子与周围格子的速度、质量

      grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)

      grid_m[base + offset] += weight * p_mass

 

 

  #遍历所有的格子

  for i, j in ti.ndrange(n_grid, n_grid):

    #如果这块格子有质量才需要计算

    if grid_m[i, j] > 0: # No need for epsilon here

      # 动量转为速度。格子质量越大，格子速度变得越小

      grid_v[i, j] = (1 / grid_m[i, j]) * grid_v[i, j] # Momentum to velocity

      # 重力影响

      grid_v[i, j][1] -= dt * 50 # gravity

      # 碰到四周墙壁，格子速度强行置为0。数字3就是第几个格子算墙壁的意思，可以改大试试

      if i < 3 and grid_v[i, j][0] < 0:          grid_v[i, j][0] = 0 # Boundary conditions

      if i > n_grid - 3 and grid_v[i, j][0] > 0: grid_v[i, j][0] = 0

      if j < 3 and grid_v[i, j][1] < 0:          grid_v[i, j][1] = 0

      if j > n_grid - 3 and grid_v[i, j][1] > 0: grid_v[i, j][1] = 0

 

 

  # 最后，把格子的计算还原到粒子上去。遍历所有粒子

  for p in range(n_particles): # grid to particle (G2P)

    base = (x[p] * inv_dx - 0.5).cast(int)

    fx = x[p] * inv_dx - base.cast(float)

    w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1.0), 0.5 * ti.sqr(fx - 0.5)]

    new_v = ti.Vector.zero(ti.f32, 2)

    new_C = ti.Matrix.zero(ti.f32, 2, 2)

    # 同样，要考虑该粒子周围的几个格子

    for i, j in ti.static(ti.ndrange(3, 3)): # loop over 3x3 grid node neighborhood

      dpos = ti.Vector([i, j]).cast(float) - fx

      g_v = grid_v[base + ti.Vector([i, j])]

      weight = w[i][0] * w[j][1]

      # 新速度 += 权重 * 格子速度

      new_v += weight * g_v

      new_C += 4 * inv_dx * weight * ti.outer_product(g_v, dpos)

    # 更新粒子的速度、C

    v[p], C[p] = new_v, new_C

    # 位置 += 速度 * ΔTime

    x[p] += dt * v[p] # advection

 

# ~~~~初始化开始~~~~~

# 这里要结合运行效果分析

import random

group_size = n_particles // 3       # 分三组，每组一个正方形

for i in range(n_particles):

  # 位置都是归1化的，取值0~1

  x[i] = [random.random() * 0.2 + 0.3 + 0.10 * (i // group_size), random.random() * 0.2 + 0.05 + 0.32 * (i // group_size)]

  material[i] = i // group_size # 0: fluid流体 1: jelly果冻 2: snow雪

  v[i] = [0, 0]     # 初始速度0

  F[i] = [[1, 0], [0, 1]]       #变形梯度矩阵初始值

  Jp[i] = 1                     #塑料变形参数初始值

 

import numpy as np

# 初始化ti.GUI，显示用

gui = ti.GUI("Taichi MLS-MPM-99", res=512, background_color=0x112F41)

# 运行20000帧

for frame in range(20000):

  # s从0到18，每帧要算19次

  for s in range(int(2e-3 // dt)):

    substep()

  colors = np.array([0x068587, 0xED553B, 0xEEEEF0], dtype=np.uint32)    # 三种颜色，每组一个颜色

  gui.circles(x.to_numpy(), radius=1.5, color=colors[material.to_numpy()])    # 根据位置画小圆点，颜色3选1

  gui.show() # Change to gui.show(f'{frame:06d}.png') to write images to disk   # 每帧画一次
作者：皮皮关做游戏
https://www.bilibili.com/read/cv4358206?spm_id_from=333.788.b_636f6d6d656e74.55
出处： bilibili