问题描述
通过taichi实现基于PBD的固液耦合模拟，在实现过程中出现了，一个kernel内对field赋值失败的情况，
代码的561~575实现了一个基于网格的近邻搜索，为了调试打印出了邻居大于0的粒子的编号。
运行结果如下
在561~575行，fr_neighbors_num[23]应该被赋值为2，但是循环结束后，fr_neighbors_num[23]变为了0，请问这可能是由什么原因导致的呢？
源代码如下
import math
import numpy as np
import taichi as ti
from taichi.lang.matrix import MatrixType
from utils import Cube, compute_cube_sdf, compute_cube_sdf_grident
ti.init(arch=ti.cuda)
fluid_color = 0x068587
rigid_color = 0xebaca2
screen_res = (400, 400)
screen_to_world_ratio = 10.0
boundary = (screen_res[0] / screen_to_world_ratio,
screen_res[1] / screen_to_world_ratio)
cell_size = 2.51
cell_recpr = 1.0 / cell_size
def round_up(f, s):
return (math.floor(f * cell_recpr / s) + 1) * s
grid_size = (round_up(boundary[0], 1), round_up(boundary[1], 1))
max_neighbors_num = 100
particle_radius = 3.0
particle_radius_in_world = particle_radius / screen_to_world_ratio
num_particles_x = 60
num_particles = num_particles_x * 20
h = 1.1
neighbor_radius = h * 1.05
dim = 2
epsilon = 1e-5
g = 9.8
fps = 20
time_delta = 1 / fps
# fluid
poly6_factor = 315.0 / 64.0 / math.pi
spiky_grad_factor = -45.0 / math.pi
fluid_mass = 1.0
rho0 = 1.0
lambda_eps = 100.0
corrK = 0.001
corrN = 4
corr_deltaQ_coeff = 0.3
solver_iterations = 5
cs = 1.0
gamma = 1.0
B = rho0 * cs * cs / gamma
# rigid
rigid_mass = 1.0
damping = 0.99
coef = 1.5
vec2: MatrixType = ti.types.vector(2, ti.f32)
vec3: MatrixType = ti.types.vector(3, ti.f32)
mat2x2: MatrixType = ti.types.matrix(2, 2, ti.f32)
mat3x3: MatrixType = ti.types.matrix(3, 3, ti.f32)
@ti.data_oriented
class ParticleSystem:
def __init__(self, max_particls_num: int) -> None:
n = max_particls_num
self.max_particls_num = n
self.particles_num = ti.field(ti.i32, shape=())
self.x = vec2.field(shape=(n,))
self.v = vec2.field(shape=(n,))
self.f = vec2.field(shape=(n,)) # F / m
self.phase = ti.field(ti.i32, shape=(n,))
def dump(self):
n = self.particles_num[None]
x = self.x.to_numpy()
return x[:n, :]
def reset(self):
self.particles_num[None] = 0
@ti.func
def poly6_value(self, s, h):
result = 0.0
if 0 < s and s < h:
x = (h * h - s * s) / (h * h * h)
result = poly6_factor * x * x * x
return result
@ti.func
def spiky_gradient(self, r, h):
result = ti.Vector([0.0, 0.0])
r_len = r.norm()
if 0 < r_len and r_len < h:
x = (h - r_len) / (h * h * h)
g_factor = spiky_grad_factor * x * x
result = r * g_factor / r_len
return result
class RigidSolver(ParticleSystem):
def __init__(self, max_particls_num: int, max_rigid_num: int = 10) -> None:
super().__init__(max_particls_num)
n = max_particls_num
self.max_rigid_num = max_rigid_num
# debug
self.debug = ti.field(ti.i32, shape=())
self.n_x = vec2.field(shape=(n,))
self.o_x = vec2.field(shape=(n,)) # origin x
self.d_x = vec2.field(shape=(n,))
self.count_x = ti.field(ti.i32, shape=(n,))
self.sdf = ti.field(ti.f32, shape=(n,))
self.sdf_gradient = vec2.field(shape=(n,))
# grid for search neighbors
self.ns_shape = grid_size
self.ns_grid = ti.field(ti.i32, shape=self.ns_shape + (max_neighbors_num,))
self.ns_grid_num = ti.field(ti.i32, shape=self.ns_shape)
self.neighbors_num = ti.field(ti.i32, shape=(n,))
self.neighbors = ti.field(ti.i32, shape=(n, max_neighbors_num))
# number of rigid-body
self.rigid_num = ti.field(ti.i32, shape=())
self.rigid_particls_num = ti.field(ti.i32, shape=(max_rigid_num,))
self.rigid_particls_offset = ti.field(ti.i32, shape=(max_rigid_num,))
self.rigidR = mat2x2.field(shape=(max_rigid_num,))
def reset(self):
self.debug[None] = 0
self.particles_num[None] = 0
self.rigid_num[None] = 0
@ti.kernel
def reset_f(self):
m = self.particles_num[None]
for i in range(m):
self.f[i] = ti.Vector([0., -g])
@ti.kernel
def add_particle(self,
vec_p: ti.types.ndarray(),
sdf: ti.types.ndarray(),
sdf_grident: ti.types.ndarray(),
n: int,
):
rigid_id = self.rigid_num[None]
m = self.particles_num[None]
for i in range(n):
self.o_x[m + i] = ti.Vector([vec_p[i, 0], vec_p[i, 1]])
self.x[m + i] = ti.Vector([vec_p[i, 0], vec_p[i, 1]])
self.n_x[m + i] = ti.Vector([vec_p[i, 0], vec_p[i, 1]])
self.v[m + i] = ti.Vector([0., 0.])
self.f[m + i] = ti.Vector([0., -g])
self.sdf[m + i] = sdf[i]
self.sdf_gradient[m + i] = ti.Vector([sdf_grident[i, 0], sdf_grident[i, 1]])
self.phase[m + i] = rigid_id
self.rigid_particls_offset[rigid_id] = m
self.rigid_particls_num[rigid_id] = n
self.particles_num[None] += n
self.rigid_num[None] += 1
def add_cube(self, center: np.ndarray, W: float, H: float):
center = center / screen_to_world_ratio
W /= screen_to_world_ratio
H /= screen_to_world_ratio
r = particle_radius_in_world
nW = int((W - 2 * r) / (2 * r) + 0.5)
nH = int((H - 2 * r) / (2 * r) + 0.5)
pos_x = np.linspace(center[0] - W / 2 + r, center[0] + W / 2 - r, nW)
pos_y = np.linspace(center[1] - H / 2 + r, center[1] + H / 2 - r, nH)
pos_x, pos_y = np.meshgrid(pos_x, pos_y)
pos = np.stack([pos_x.ravel(), pos_y.ravel()], axis=-1)
n = pos.shape[0]
sdf = np.zeros((n,))
sdf_grident = np.zeros((n, 2))
cube = Cube(center, W, H)
for i in range(n):
sdf[i] = compute_cube_sdf(cube, pos[i])
sdf_grident[i] = compute_cube_sdf_grident(cube, pos[i])
if self.particles_num[None] + n <= self.max_particls_num and self.rigid_num[None] < self.max_rigid_num:
self.add_particle(pos, sdf, sdf_grident, n)
@ti.func
def to_grid_index(self, x): # Vector is Matrix with m = 1
I = (x * cell_recpr).cast(ti.i32)
return I
@ti.func
def grid_index_is_valid(self, I):
return (0 <= I[0] < self.ns_grid.shape[0]) and (0 <= I[1] < self.ns_grid.shape[1])
@ti.kernel
def find_neighbors(self):
for I in ti.grouped(self.ns_grid_num):
self.ns_grid_num[I] = 0
for p_i in range(self.particles_num[None]):
I = self.to_grid_index(self.n_x[p_i])
j = ti.atomic_add(self.ns_grid_num[I], 1)
self.ns_grid[I, j] = p_i
for p_i in range(self.particles_num[None]):
I = self.to_grid_index(self.n_x[p_i])
j = 0
for dI in ti.grouped(ti.ndrange((-1, 2), (-1, 2))):
nI = I + dI
if self.grid_index_is_valid(nI):
for J in range(self.ns_grid_num[nI]):
p_j = self.ns_grid[nI, J]
dist = (self.n_x[p_i] - self.n_x[p_j]).norm()
if p_i != p_j and dist < neighbor_radius and j < max_neighbors_num:
self.neighbors[p_i, j] = p_j
j += 1
if j >= max_neighbors_num:
break
self.neighbors_num[p_i] = j
@ti.kernel
def predict_position(self, dt: ti.f32):
for i in range(self.particles_num[None]):
vel = self.v[i] + dt * self.f[i]
pos = self.x[i] + dt * vel
self.n_x[i] = pos
@ti.kernel
def reset_dx(self):
for p_i in range(self.particles_num[None]):
self.d_x[p_i] = ti.Vector([0., 0.])
self.count_x[p_i] = 0
@ti.kernel
def shape_matching(self, r_id: ti.i32):
offset = self.rigid_particls_offset[r_id]
nx_cm = ti.Vector([0., 0.]) # mass center
ox_cm = ti.Vector([0., 0.])
ti.loop_config(serialize=True)
for i in range(self.rigid_particls_num[r_id]):
n = ti.cast(i, ti.f32)
nx_cm = n / (n + 1) * nx_cm + 1.0 / (n + 1) * self.n_x[offset + i]
ox_cm = n / (n + 1) * ox_cm + 1.0 / (n + 1) * self.o_x[offset + i]
Apq = ti.Matrix([[0., 0.], [0., 0.]])
for i in range(self.rigid_particls_num[r_id]):
p = self.n_x[offset + i] - nx_cm
q = self.o_x[offset + i] - ox_cm
Apq += rigid_mass * (p @ q.transpose())
R, _ = ti.polar_decompose(Apq)
self.rigidR[r_id] = R
for i in range(self.rigid_particls_num[r_id]):
self.d_x[offset + i] += (R @ (self.o_x[offset + i] - ox_cm)) + nx_cm - self.n_x[offset + i]
self.count_x[offset + i] += 1
@ti.kernel
def collision_solving(self):
for p_i in range(self.particles_num[None]):
I = self.to_grid_index(self.n_x[p_i])
for dI in ti.grouped(ti.ndrange((-1, 2), (-1, 2))):
nI = I + dI
if self.grid_index_is_valid(nI):
for p_j_index in range(self.ns_grid_num[nI]):
p_j = self.ns_grid[nI, p_j_index]
dist = (self.n_x[p_i] - self.n_x[p_j]).norm()
r = particle_radius_in_world
if self.phase[p_i] != self.phase[p_j] and dist < 2 * r:
sdf_grad_i = self.sdf_gradient[p_i]
sdf_grad_j = self.sdf_gradient[p_j]
sdf_i = ti.abs(self.sdf[p_i])
sdf_j = ti.abs(self.sdf[p_j])
d = min(sdf_i, sdf_j)
self.count_x[p_i] += 1
if sdf_i < sdf_j:
r_id = self.phase[p_i]
n_ij = self.rigidR[r_id] @ sdf_grad_i
self.d_x[p_i] += -0.5 * d * n_ij
else:
r_id = self.phase[p_j]
n_ji = self.rigidR[r_id] @ sdf_grad_j
self.d_x[p_i] += 0.5 * d * n_ji
@ti.kernel
def apply_dx(self):
for p_i in range(self.particles_num[None]):
if self.count_x[p_i] > 0:
self.n_x[p_i] += self.d_x[p_i] / ti.cast(self.count_x[p_i], ti.f32) * ti.min(coef, 1.0 * self.count_x[p_i])
@ti.kernel
def apply_dx_prev(self):
for p_i in range(self.particles_num[None]):
if self.count_x[p_i] > 0:
self.x[p_i] += self.d_x[p_i] / ti.cast(self.count_x[p_i], ti.f32) * ti.min(coef, 1.0 * self.count_x[p_i])
@ti.kernel
def boundary_solving(self):
bmin = ti.Vector([particle_radius_in_world, particle_radius_in_world])
bmax = ti.Vector([boundary[0], boundary[1]]) - particle_radius_in_world
for j in range(self.particles_num[None]):
p = self.n_x[j]
for i in ti.static(range(dim)):
# Use randomness to prevent particles from sticking into each other after clamping
if p[i] <= bmin[i]:
p[i] = bmin[i] # bmin[i] + epsilon * ti.random()
elif bmax[i] <= p[i]:
p[i] = bmax[i] # bmax[i] - epsilon * ti.random()
if (p - self.n_x[j]).norm() > 0:
self.d_x[j] += p - self.n_x[j]
self.count_x[j] += 1
def substep(self):
self.reset_dx()
for r_id in range(self.rigid_num[None]):
self.shape_matching(r_id)
self.apply_dx()
self.reset_dx()
self.collision_solving()
self.apply_dx()
self.reset_dx()
self.boundary_solving()
self.apply_dx()
@ti.kernel
def update(self, dt: ti.f32):
for p_i in range(self.particles_num[None]):
self.v[p_i] = (self.n_x[p_i] - self.x[p_i]) / dt * damping
if (self.v[p_i]).norm() <= 0.05 * dt:
self.v[p_i] = ti.Vector([0., 0.])
for i in range(self.particles_num[None]):
self.x[i] = self.n_x[i]
def step(self):
dt = time_delta
self.predict_position(dt)
self.find_neighbors()
for _ in range(solver_iterations):
self.substep()
self.update(dt)
@ti.data_oriented
class FluidSolver(ParticleSystem):
def __init__(self, max_particls_num: int) -> None:
super().__init__(max_particls_num)
n = max_particls_num
self.n_x = vec2.field(shape=(n,))
# params for pbf
self.d_x = vec2.field(shape=(n,))
self.lam = ti.field(ti.f32, shape=(n,))
# grid for search neighbors
self.ns_shape = grid_size
self.ns_grid = ti.field(ti.i32, shape=self.ns_shape + (max_neighbors_num,))
self.ns_grid_num = ti.field(ti.i32, shape=self.ns_shape)
self.neighbors_num = ti.field(ti.i32, shape=(n,))
self.neighbors = ti.field(ti.i32, shape=(n, max_neighbors_num))
@ti.kernel
def reset_f(self):
m = self.particles_num[None]
for i in range(m):
self.f[i] = ti.Vector([0., -g])
@ti.kernel
def add_particle(self,
vec_p: ti.types.ndarray(),
n: int,
):
m = self.particles_num[None]
for i in range(n):
self.x[m + i] = ti.Vector([vec_p[i, 0], vec_p[i, 1]])
self.n_x[m + i] = ti.Vector([vec_p[i, 0], vec_p[i, 1]])
self.v[m + i] = ti.Vector([0., 0.])
self.f[m + i] = ti.Vector([0., -g])
self.phase[m + i] = 0
self.particles_num[None] += n
def add_cube(self, center: np.ndarray, W: float, H: float):
center = center / screen_to_world_ratio
W /= screen_to_world_ratio
H /= screen_to_world_ratio
r = particle_radius_in_world
nW = int((W - 2 * r) / (2 * r) + 0.5)
nH = int((H - 2 * r) / (2 * r) + 0.5)
pos_x = np.linspace(center[0] - W / 2 + r, center[0] + W / 2 - r, nW)
pos_y = np.linspace(center[1] - H / 2 + r, center[1] + H / 2 - r, nH)
pos_x, pos_y = np.meshgrid(pos_x, pos_y)
pos = np.stack([pos_x.ravel(), pos_y.ravel()], axis=-1)
n = pos.shape[0]
if self.particles_num[None] + n <= self.max_particls_num:
self.add_particle(pos, n)
@ti.func
def boundary_solving(self, p):
bmin = particle_radius_in_world
bmax = ti.Vector([boundary[0], boundary[1]]) - particle_radius_in_world
for i in ti.static(range(dim)):
# Use randomness to prevent particles from sticking into each other after clamping
if p[i] <= bmin:
p[i] = bmin + epsilon * ti.random()
elif bmax[i] <= p[i]:
p[i] = bmax[i] - epsilon * ti.random()
return p
@ti.kernel
def predict_position(self, dt: ti.f32):
for i in range(self.particles_num[None]):
vel = self.v[i] + dt * self.f[i]
pos = self.x[i] + dt * vel
self.n_x[i] = self.boundary_solving(pos)
@ti.func
def to_grid_index(self, x): # Vector is Matrix with m = 1
I = (x * cell_recpr).cast(ti.i32)
return I
@ti.func
def grid_index_is_valid(self, I):
return (0 <= I[0] < self.ns_grid.shape[0]) and (0 <= I[1] < self.ns_grid.shape[1])
@ti.kernel
def find_neighbors(self):
for I in ti.grouped(self.ns_grid_num):
self.ns_grid_num[I] = 0
for p_i in range(self.particles_num[None]):
I = self.to_grid_index(self.n_x[p_i])
j = ti.atomic_add(self.ns_grid_num[I], 1)
self.ns_grid[I, j] = p_i
for p_i in range(self.particles_num[None]):
I = self.to_grid_index(self.n_x[p_i])
j = 0
for dI in ti.grouped(ti.ndrange((-1, 2), (-1, 2))):
nI = I + dI
if self.grid_index_is_valid(nI):
for J in range(self.ns_grid_num[nI]):
p_j = self.ns_grid[nI, J]
dist = (self.n_x[p_i] - self.n_x[p_j]).norm()
if p_i != p_j and dist < neighbor_radius and j < max_neighbors_num:
self.neighbors[p_i, j] = p_j
j += 1
if j >= max_neighbors_num:
break
self.neighbors_num[p_i] = j
@ti.func
def compute_scorr(self, pos_ji):
# Eq (13)
x = self.poly6_value(pos_ji.norm(), h) / self.poly6_value(corr_deltaQ_coeff * h, h)
# pow(x, 4)
x = x * x
x = x * x
return (-corrK) * x
@ti.kernel
def substep(self):
# compute lambdas
# Eq (8) ~ (11)
for p_i in range(self.particles_num[None]):
pos_i = self.n_x[p_i]
grad_i = ti.Vector([0.0, 0.0])
sum_gradient_sqr = 0.0
density_constraint = 0.0
for j in range(self.neighbors_num[p_i]):
p_j = self.neighbors[p_i, j]
pos_ji = pos_i - self.n_x[p_j]
grad_j = self.spiky_gradient(pos_ji, h)
grad_i += grad_j
sum_gradient_sqr += grad_j.norm() ** 2
# Eq(2)
density_constraint += self.poly6_value(pos_ji.norm(), h)
# Eq(1)
density_constraint = (fluid_mass * density_constraint / rho0) - 1.0
sum_gradient_sqr += grad_i.norm() ** 2
self.lam[p_i] = (-density_constraint) / (sum_gradient_sqr + lambda_eps)
# compute position deltas
# Eq(12), (14)
for p_i in range(self.particles_num[None]):
pos_i = self.n_x[p_i]
lambda_i = self.lam[p_i]
pos_delta_i = ti.Vector([0.0, 0.0])
for j in range(self.neighbors_num[p_i]):
p_j = self.neighbors[p_i, j]
lambda_j = self.lam[p_j]
pos_ji = pos_i - self.n_x[p_j]
scorr_ij = self.compute_scorr(pos_ji)
pos_delta_i += (lambda_i + lambda_j + scorr_ij) * \
self.spiky_gradient(pos_ji, h)
pos_delta_i /= rho0
self.d_x[p_i] = pos_delta_i
# apply position deltas
for i in range(self.particles_num[None]):
self.n_x[i] += self.d_x[i]
@ti.kernel
def update(self, dt: ti.f32):
for i in range(self.particles_num[None]):
self.n_x[i] = self.boundary_solving(self.n_x[i])
for p_i in range(self.particles_num[None]):
self.v[p_i] = (self.n_x[p_i] - self.x[p_i]) / dt
for i in range(self.particles_num[None]):
self.x[i] = self.n_x[i]
def step(self):
dt = time_delta
self.predict_position(dt)
self.find_neighbors()
for _ in range(solver_iterations):
self.substep()
self.update(dt)
@ti.data_oriented
class CouplingSolver:
def __init__(self, max_particls_num: int, max_rigid_num: int) -> None:
self.fluid_solver = FluidSolver(max_particls_num)
self.rigid_solver = RigidSolver(max_particls_num, max_rigid_num)
n = max_neighbors_num
self.ns_shape = grid_size
self.fr_ns_grid = ti.field(ti.i32, shape=self.ns_shape + (max_neighbors_num,))
self.fr_ns_grid_num = ti.field(ti.i32, shape=self.ns_shape)
self.fr_neighbors_num = ti.field(ti.i32, shape=(n,))
self.fr_neighbors = ti.field(ti.i32, shape=(n, max_neighbors_num))
self.fluid_p = ti.field(ti.f32, shape=(n,))
self.fluid_rho = ti.field(ti.f32, shape=(n,))
def reset(self):
self.fluid_solver.reset()
self.rigid_solver.reset()
def add_rigid_cube(self, center: np.ndarray, W: float, H: float):
self.rigid_solver.add_cube(center, W, H)
def add_fluid_cube(self, center: np.ndarray, W: float, H: float):
self.fluid_solver.add_cube(center, W, H)
@ti.kernel
def compute_ext_force(self):
print("BEGIN-----------------------------------------------------------")
n = self.fluid_solver.particles_num[None]
m = self.rigid_solver.particles_num[None]
ti.loop_config(serialize=True)
for I in ti.grouped(self.fr_ns_grid_num):
self.fr_ns_grid_num[I] = 0
ti.loop_config(serialize=True)
for b_i in range(m):
I = self.fluid_solver.to_grid_index(self.rigid_solver.n_x[b_i])
j = ti.atomic_add(self.fr_ns_grid_num[I], 1)
self.fr_ns_grid[I, j] = b_i
ti.loop_config(serialize=True)
for p_i in range(n):
I = self.fluid_solver.to_grid_index(self.fluid_solver.n_x[p_i])
j = 0
for dI in ti.grouped(ti.ndrange((-1, 2), (-1, 2))):
nI = I + dI
if self.fluid_solver.grid_index_is_valid(nI):
for J in range(self.fr_ns_grid_num[nI]):
b_j = self.fr_ns_grid[nI, J]
dist = (self.fluid_solver.n_x[p_i] - self.rigid_solver.n_x[b_j]).norm()
if dist < neighbor_radius and j < max_neighbors_num:
self.fr_neighbors[p_i, j] = b_j
j += 1
if j >= max_neighbors_num:
break
self.fr_neighbors_num[p_i] = j
if j > 0:
print("p_i:", p_i, "num:", self.fr_neighbors_num[p_i])
print("-----------------------------------------------------------")
print("p_i:", 23, "num:", self.fr_neighbors_num[23])
print("-----------------------------------------------------------")
for p_i in range(n):
pos_i = self.fluid_solver.n_x[p_i]
rho_i = 0.0
for j in range(self.fluid_solver.neighbors_num[p_i]):
p_j = self.fluid_solver.neighbors[p_i, j]
pos_j = self.fluid_solver.n_x[p_j]
pos_ij = pos_i - pos_j
rho_i += fluid_mass * self.fluid_solver.poly6_value(pos_ij.norm(), h)
for j in range(self.fr_neighbors_num[p_i]):
b_j = self.fr_neighbors[p_i, j]
pos_j = self.rigid_solver.n_x[b_j]
pos_ij = pos_i - pos_j
rho_i += rigid_mass * self.rigid_solver.poly6_value(pos_ij.norm(), h)
# print("p_i:", p_i, "b_j:", b_j, " norm:", pos_ij.norm(), "pos_i:", pos_i, "pos_j:", pos_j)
self.fluid_rho[p_i] = rho_i
for p_i in range(n):
rho_i = self.fluid_rho[p_i]
self.fluid_p[p_i] = (ti.pow(rho_i / rho0, gamma) - 1)
for p_i in range(n):
pos_i = self.fluid_solver.n_x[p_i]
pre_i = self.fluid_p[p_i] # pressure
rho_i = self.fluid_rho[p_i]
for j in range(self.fr_neighbors_num[p_i]):
b_j = self.fr_neighbors[p_i, j]
pos_j = self.rigid_solver.n_x[b_j]
pos_ij = pos_i - pos_j
Fp_ij = fluid_mass * rigid_mass * (pre_i / ti.pow(rho_i, 2)) * \
self.fluid_solver.spiky_gradient(pos_ij, h)
self.fluid_solver.f[p_i] += -Fp_ij / fluid_mass
self.rigid_solver.f[b_j] += Fp_ij / rigid_mass
# for p_i in range(n):
# if self.neighbors_num[p_i] > 0:
# print("p_i:", p_i, "neighbors_num:", self.neighbors_num[p_i], "p:", self.fluid_p[p_i], "rho:", self.fluid_rho[p_i])
def step(self):
self.fluid_solver.reset_f()
self.rigid_solver.reset_f()
self.fluid_solver.find_neighbors()
self.rigid_solver.find_neighbors()
self.compute_ext_force()
dt = time_delta
self.fluid_solver.predict_position(dt)
self.fluid_solver.find_neighbors()
self.rigid_solver.predict_position(dt)
self.rigid_solver.find_neighbors()
for _ in range(solver_iterations):
self.fluid_solver.substep()
self.rigid_solver.substep()
self.fluid_solver.update(dt)
self.rigid_solver.update(dt)
def dump(self):
return self.fluid_solver.dump(), self.rigid_solver.dump()
def main():
pbd_solver = CouplingSolver(10000, 10)
pbd_solver.add_fluid_cube(np.array([100, 200]), 80, 300)
pbd_solver.add_rigid_cube(np.array([300, 200]), 30, 30)
gui = ti.GUI("PBD", screen_res)
while gui.running:
for e in gui.get_events(ti.GUI.PRESS):
if e.key == 'r':
pbd_solver.reset()
if e.key == 'g':
pbd_solver.step()
x0, x1 = pbd_solver.dump()
for j in range(dim):
x0[:, j] *= screen_to_world_ratio / screen_res[j]
x1[:, j] *= screen_to_world_ratio / screen_res[j]
gui.circles(x0, radius=3.0, color=fluid_color)
gui.circles(x1, radius=3.0, color=rigid_color)
gui.show()
if __name__ == '__main__':
main()