taichi example里的stable_fluid.py，只改了压力求解器。但改完之后爆炸了，不知道什么原因，希望大佬指教
爆炸效果gif：
共轭梯度求解器部分的代码如下：
import taichi as ti
ti.init(arch=ti.cpu)
@ti.data_oriented
class CG:
def __init__(self, velocity_divs: ti.template()):
self.velocity_divs = velocity_divs
self.x = ti.field(dtype=ti.f32, shape=velocity_divs.shape)
self.r = ti.field(dtype=ti.f32, shape=velocity_divs.shape)
self.d = ti.field(dtype=ti.f32, shape=velocity_divs.shape)
self.Ad = ti.field(dtype=ti.f32, shape=velocity_divs.shape)
self.alpha = ti.field(dtype=ti.f32, shape=())
self.sum = ti.field(dtype=ti.f32, shape=())
self.beta = ti.field(dtype=ti.f32, shape=())
self.delta_old = ti.field(dtype=ti.f32, shape=())
self.delta_new = ti.field(dtype=ti.f32, shape=())
@ti.func
def sample(self, c, u, v):
I = ti.Vector([int(u), int(v)])
I = max(0, min(512 - 1, I)) # 网格大小512*512
return c[I]
@ti.func
def sum_neighbor(self, c, I):
sum = self.sample(c, I.x - 1, I.y) + self.sample(c, I.x + 1, I.y) + self.sample(c, I.x, I.y - 1) + self.sample(
c, I.x, I.y + 1)
return sum
@ti.kernel
def compute_Ad(self):
for I in ti.grouped(self.velocity_divs):
self.Ad[I] = -2 * 2 * self.d[I] + self.sum_neighbor(self.d, I)
@ti.kernel
def compute_rTr(self) -> ti.f32:
self.sum[None] = 0.0
for I in ti.grouped(self.r):
self.sum[None] += self.r[I] * self.r[I]
return self.sum[None]
@ti.kernel
def update_delta(self):
self.delta_old[None] = self.delta_new[None]
@ti.kernel
def update_alpha(self, eps: ti.template()):
self.sum[None] = 0.0
for I in ti.grouped(self.d):
self.sum[None] += self.d[I] * self.Ad[I]
self.alpha[None] = self.delta_new[None] / ti.max(self.sum[None], eps)
@ti.kernel
def update_beta(self, eps: ti.template()):
self.beta[None] = self.delta_new[None] / ti.max(self.delta_old[None], eps)
@ti.kernel
def update_x(self):
for I in ti.grouped(self.x):
self.x[I] += (self.alpha[None] * self.d[I])
@ti.kernel
def update_r(self):
for I in ti.grouped(self.r):
self.r[I] -= (self.alpha[None] * self.Ad[I])
@ti.kernel
def update_d(self):
for I in ti.grouped(self.d):
self.d[I] = self.r[I] + (self.beta[None] * self.d[I])
@ti.kernel
def init(self):
for I in ti.grouped(self.x):
self.x[I] = 0.0
self.r[I] = self.velocity_divs[I]
self.d[I] = self.r[I]
def solve(self, maxiters=100):
self.init()
self.delta_new[None] = self.compute_rTr()
for i in ti.static(range(maxiters)):
self.compute_Ad() # q=Ad
self.update_alpha(1e-12) # alpha=rTr/dTAd
self.update_x() # x=x+alpha*d
self.update_r() # r=r-alpha*Ad
self.update_delta() # delta_old=delta_new
self.delta_new[None] = self.compute_rTr() # delta_new=rTr
if self.delta_new[None] < 1e-12:
# print('迭代次数：', i, '误差：', self.delta_new)
break
self.update_beta(1e-12) # beta=delta_new/delta_old
self.update_d() # d=r+beta*d
@ti.kernel
def update_result(self, pf_cur: ti.template()):
for I in ti.grouped(self.x):
pf_cur[I] = self.x[I]
def solve_pressure_CG():
cg_solver.solve(maxiters=cg_iters)
cg_solver.update_result(pressures_pair.cur)
cg_solver = CG(velocity_divs)
