def layout(self):
self.system.add_field("grid_n_particles", dims=(), dtype=ti.i32)
self.cell_snode = self.system.grid_snode.dense(ti.indices(DIM),
self.max_cell)
self.system.add_layout("grid_particles",
self.cell_snode,
dims=(),
dtype=ti.i32)
n = self.system.n_particles
self.neighbor_snode = ti.root.bitmasked(ti.ij, (n, n))
self.system.add_layout("neighbors",
self.neighbor_snode,
dims=(),
dtype=ti.i32)
def layout(self):
self.cell_snode = self.system.grid_snode.dense(ti.indices(DIM),
self.max_cell)
#self.cell_snode = self.get_snode().dynamic(ti.indices(DIM), self.max_cell)
#self.neighbor_snode = ti.root.dense(ti.i, self.system.n_particles).dynamic(ti.j, self.max_neighbors)
self.system.add_field("grid_n_particles", dims=(), dtype=ti.i32)
self.system.add_layout("grid_particles",
self.cell_snode,
dims=(),
dtype=ti.i32)
if self.rcut > 0:
self.system.add_attr("n_neighbors", dims=(), dtype=ti.i32)
self.neighbor_snode = self.system.particle_snode.dense(
ti.j, self.max_neighbors)
self.system.add_layout("neighbors",
self.neighbor_snode,
dims=(),
dtype=ti.i32)
def __init__(self, res):
dim = len(res)
self.dx = 1 / res[0]
self.inv_dx = 1.0 / self.dx
self.pid = ti.field(ti.i32)
self.x = ti.Vector.field(dim, dtype=ti.f32)
self.grid_m = ti.field(dtype=ti.f32)
indices = ti.ij
self.grid = ti.root.pointer(indices, 32)
block = self.grid.pointer(indices, 16)
voxel = block.dense(indices, 8)
voxel.place(self.grid_m)
block.dynamic(ti.indices(dim), 1024 * 1024,
chunk_size=4096).place(self.pid)
ti.root.dynamic(ti.i, 2**25, 2**20).place(self.x)
self.substeps = 0
for i in range(10000):
self.x[i] = [random.random() * 0.5, random.random() * 0.5]
# N-body gravity simulation in 300 lines of Taichi, tree method, no multipole, O(N log N)
# Author: archibate <*****@*****.**>, all left reserved
import taichi as ti
import taichi_glsl as tl
ti.init()
if not hasattr(ti, 'jkl'):
ti.jkl = ti.indices(1, 2, 3)
kUseTree = True
#kDisplay = 'tree mouse pixels cmap save_result'
kDisplay = 'pixels'
kResolution = 512
kShapeFactor = 1
kMaxParticles = 8192
kMaxDepth = kMaxParticles * 1
kMaxNodes = kMaxParticles * 4
kDim = 2
dt = 0.00005
LEAF = -1
TREE = -2
particle_mass = ti.var(ti.f32)
particle_pos = ti.Vector(kDim, ti.f32)
particle_vel = ti.Vector(kDim, ti.f32)
particle_table = ti.root.dense(ti.i, kMaxParticles)
particle_table.place(particle_pos).place(particle_vel).place(particle_mass)
particle_table_len = ti.var(ti.i32, ())
if kUseTree:
def bls_test_template(dim,
N,
bs,
stencil,
block_dim=None,
scatter=False,
benchmark=0,
dense=False):
x, y, y2 = ti.field(ti.i32), ti.field(ti.i32), ti.field(ti.i32)
index = ti.indices(*range(dim))
mismatch = ti.field(ti.i32, shape=())
if not isinstance(bs, (tuple, list)):
bs = [bs for _ in range(dim)]
grid_size = [N // bs[i] for i in range(dim)]
if dense:
create_block = lambda: ti.root.dense(index, grid_size)
else:
create_block = lambda: ti.root.pointer(index, grid_size)
if scatter:
block = create_block()
block.dense(index, bs).place(x)
block.dense(index, bs).place(y)
block.dense(index, bs).place(y2)
else:
create_block().dense(index, bs).place(x)
create_block().dense(index, bs).place(y)
create_block().dense(index, bs).place(y2)
ndrange = ((bs[i], N - bs[i]) for i in range(dim))
if block_dim is None:
block_dim = 1
for i in range(dim):
block_dim *= bs[i]
@ti.kernel
def populate():
for I in ti.grouped(ti.ndrange(*ndrange)):
s = 0
for i in ti.static(range(dim)):
s += I[i]**(i + 1)
x[I] = s
@ti.kernel
def apply(use_bls: ti.template(), y: ti.template()):
if ti.static(use_bls and not scatter):
ti.cache_shared(x)
if ti.static(use_bls and scatter):
ti.cache_shared(y)
ti.block_dim(block_dim)
for I in ti.grouped(x):
if ti.static(scatter):
for offset in ti.static(stencil):
y[I + ti.Vector(offset)] += x[I]
else:
# gather
s = 0
for offset in ti.static(stencil):
s = s + x[I + ti.Vector(offset)]
y[I] = s
populate()
if benchmark:
for i in range(benchmark):
x.snode.parent().deactivate_all()
if not scatter:
populate()
y.snode.parent().deactivate_all()
y2.snode.parent().deactivate_all()
apply(False, y2)
apply(True, y)
else:
# Simply test
apply(False, y2)
apply(True, y)
@ti.kernel
def check():
for I in ti.grouped(y2):
if y[I] != y2[I]:
print('check failed', I, y[I], y2[I])
mismatch[None] = 1
check()
ti.kernel_profiler_print()
assert mismatch[None] == 0
trash_base_parent = ti.field(ti.i32)
trash_base_geo_center = ti.Vector.field(kDim, ti.f32)
trash_base_geo_size = ti.field(ti.f32)
trash_table = ti.root.dense(ti.i, kMaxDepth)
trash_table.place(trash_particle_id)
trash_table.place(trash_base_parent, trash_base_geo_size)
trash_table.place(trash_base_geo_center)
trash_table_len = ti.field(ti.i32, ())
node_mass = ti.field(ti.f32)
node_weighted_pos = ti.Vector.field(kDim, ti.f32)
node_particle_id = ti.field(ti.i32)
node_children = ti.field(ti.i32)
node_table = ti.root.dense(ti.i, kMaxNodes)
node_table.place(node_mass, node_particle_id, node_weighted_pos)
node_table.dense(ti.indices(*list(range(1, 1 + kDim))),
2).place(node_children)
node_table_len = ti.field(ti.i32, ())
if 'mouse' in kDisplay:
display_image = ti.Vector.field(3, ti.f32, (kResolution, kResolution))
elif len(kDisplay):
display_image = ti.field(ti.f32, (kResolution, kResolution))
@ti.func
def alloc_node():
ret = ti.atomic_add(node_table_len[None], 1)
assert ret < kMaxNodes
node_mass[ret] = 0
node_weighted_pos[ret] = particle_pos[0] * 0
def __init__(
self,
res,
size=1,
max_num_particles=2**27,
# Max 128 MB particles
padding=3,
unbounded=False,
dt_scale=1,
E_scale=1,
voxelizer_super_sample=2):
self.dim = len(res)
assert self.dim in (
2, 3), "MPM solver supports only 2D and 3D simulations."
self.res = res
self.n_particles = ti.field(ti.i32, shape=())
self.dx = size / res[0]
self.inv_dx = 1.0 / self.dx
self.default_dt = 2e-2 * self.dx / size * dt_scale
self.p_vol = self.dx**self.dim
self.p_rho = 1000
self.p_mass = self.p_vol * self.p_rho
self.max_num_particles = max_num_particles
self.gravity = ti.Vector.field(self.dim, dtype=ti.f32, shape=())
self.source_bound = ti.Vector.field(self.dim, dtype=ti.f32, shape=2)
self.source_velocity = ti.Vector.field(self.dim,
dtype=ti.f32,
shape=())
self.pid = ti.field(ti.i32)
# position
self.x = ti.Vector.field(self.dim, dtype=ti.f32)
# velocity
self.v = ti.Vector.field(self.dim, dtype=ti.f32)
# affine velocity field
self.C = ti.Matrix.field(self.dim, self.dim, dtype=ti.f32)
# deformation gradient
self.F = ti.Matrix.field(self.dim, self.dim, dtype=ti.f32)
# material id
self.material = ti.field(dtype=ti.i32)
self.color = ti.field(dtype=ti.i32)
# plastic deformation volume ratio
self.Jp = ti.field(dtype=ti.f32)
if self.dim == 2:
indices = ti.ij
else:
indices = ti.ijk
offset = tuple(-self.grid_size // 2 for _ in range(self.dim))
self.offset = offset
# grid node momentum/velocity
self.grid_v = ti.Vector.field(self.dim, dtype=ti.f32)
# grid node mass
self.grid_m = ti.field(dtype=ti.f32)
grid_block_size = 128
self.grid = ti.root.pointer(indices, self.grid_size // grid_block_size)
if self.dim == 2:
self.leaf_block_size = 16
else:
self.leaf_block_size = 8
block = self.grid.pointer(indices,
grid_block_size // self.leaf_block_size)
def block_component(c):
block.dense(indices, self.leaf_block_size).place(c, offset=offset)
block_component(self.grid_m)
for v in self.grid_v.entries:
block_component(v)
block.dynamic(ti.indices(self.dim),
1024 * 1024,
chunk_size=self.leaf_block_size**self.dim * 8).place(
self.pid, offset=offset + (0, ))
self.padding = padding
# Young's modulus and Poisson's ratio
self.E, self.nu = 1e6 * size * E_scale, 0.2
# Lame parameters
self.mu_0, self.lambda_0 = self.E / (
2 * (1 + self.nu)), self.E * self.nu / ((1 + self.nu) *
(1 - 2 * self.nu))
# Sand parameters
friction_angle = math.radians(45)
sin_phi = math.sin(friction_angle)
self.alpha = math.sqrt(2 / 3) * 2 * sin_phi / (3 - sin_phi)
self.particle = ti.root.dynamic(ti.i, max_num_particles, 2**20)
self.particle.place(self.x, self.v, self.C, self.F, self.material,
self.color, self.Jp)
self.total_substeps = 0
self.unbounded = unbounded
if self.dim == 2:
self.voxelizer = None
self.set_gravity((0, -9.8))
else:
if USE_IN_BLENDER:
from .voxelizer import Voxelizer
else:
from engine.voxelizer import Voxelizer
self.voxelizer = Voxelizer(res=self.res,
dx=self.dx,
padding=self.padding,
super_sample=voxelizer_super_sample)
self.set_gravity((0, -9.8, 0))
self.voxelizer_super_sample = voxelizer_super_sample
self.grid_postprocess = []
self.add_bounding_box(self.unbounded)
self.writers = []
def get_snode(self):
return ti.root.dense(ti.indices(*range(DIM)), (self.gridsize, ) * DIM)
def __init__(
self,
res,
quant=False,
size=1,
max_num_particles=2 ** 30,
# Max 1 G particles
padding=3,
unbounded=False,
dt_scale=1,
E_scale=1,
voxelizer_super_sample=2,
use_g2p2g=False, # ref: A massively parallel and scalable multi-GPU material point method
use_bls=True,
g2p2g_allowed_cfl=0.9, # 0.0 for no CFL limit
water_density=1.0,
support_plasticity=True):
self.dim = len(res)
self.quant = quant
self.use_g2p2g = use_g2p2g
self.use_bls = use_bls
self.g2p2g_allowed_cfl = g2p2g_allowed_cfl
self.water_density = water_density
assert self.dim in (
2, 3), "MPM solver supports only 2D and 3D simulations."
self.t = 0.0
self.res = res
self.n_particles = ti.field(ti.i32, shape=())
self.dx = size / res[0]
self.inv_dx = 1.0 / self.dx
self.default_dt = 2e-2 * self.dx / size * dt_scale
self.p_vol = self.dx ** self.dim
self.p_rho = 1000
self.p_mass = self.p_vol * self.p_rho
self.max_num_particles = max_num_particles
self.gravity = ti.Vector.field(self.dim, dtype=ti.f32, shape=())
self.source_bound = ti.Vector.field(self.dim, dtype=ti.f32, shape=2)
self.source_velocity = ti.Vector.field(self.dim,
dtype=ti.f32,
shape=())
self.input_grid = 0
self.all_time_max_velocity = 0.0
self.support_plasticity = support_plasticity
self.F_bound = 4.0
# affine velocity field
if not self.use_g2p2g:
self.C = ti.Matrix.field(self.dim, self.dim, dtype=ti.f32)
# deformation gradient
if quant:
ci21 = ti.type_factory.custom_int(21, True)
cft = ti.type_factory.custom_float(significand_type=ci21,
scale=1 / (2 ** 19))
self.x = ti.Vector.field(self.dim, dtype=cft)
cu6 = ti.type_factory.custom_int(7, False)
ci19 = ti.type_factory.custom_int(19, True)
cft = ti.type_factory.custom_float(significand_type=ci19,
exponent_type=cu6)
self.v = ti.Vector.field(self.dim, dtype=cft)
ci16 = ti.type_factory.custom_int(16, True)
cft = ti.type_factory.custom_float(significand_type=ci16,
scale=(self.F_bound + 0.1) /
(2 ** 15))
self.F = ti.Matrix.field(self.dim, self.dim, dtype=cft)
else:
self.v = ti.Vector.field(self.dim, dtype=ti.f32)
self.x = ti.Vector.field(self.dim, dtype=ti.f32)
self.F = ti.Matrix.field(self.dim, self.dim, dtype=ti.f32)
self.last_time_final_particles = ti.field(dtype=ti.i32, shape=())
# material id
if quant and self.dim == 3:
self.material = ti.field(dtype=ti.quant.int(16, False))
else:
self.material = ti.field(dtype=ti.i32)
# particle color
self.color = ti.field(dtype=ti.i32)
# plastic deformation volume ratio
if self.support_plasticity:
self.Jp = ti.field(dtype=ti.f32)
if self.dim == 2:
indices = ti.ij
else:
indices = ti.ijk
offset = tuple(-self.grid_size // 2 for _ in range(self.dim))
self.offset = offset
self.num_grids = 2 if self.use_g2p2g else 1
grid_block_size = 128
if self.dim == 2:
self.leaf_block_size = 16
else:
# TODO: use 8?
self.leaf_block_size = 4
self.grid = []
self.grid_v = []
self.grid_m = []
self.pid = []
for g in range(self.num_grids):
# grid node momentum/velocity
grid_v = ti.Vector.field(self.dim, dtype=ti.f32)
grid_m = ti.field(dtype=ti.f32)
pid = ti.field(ti.i32)
self.grid_v.append(grid_v)
# grid node mass
self.grid_m.append(grid_m)
grid = ti.root.pointer(indices, self.grid_size // grid_block_size)
block = grid.pointer(indices,
grid_block_size // self.leaf_block_size)
self.grid.append(grid)
def block_component(c):
block.dense(indices, self.leaf_block_size).place(c,
offset=offset)
block_component(grid_m)
for v in grid_v.entries:
block_component(v)
self.pid.append(pid)
# TODO ? why
block_offset = tuple(o // self.leaf_block_size for o in self.offset)
block.dynamic(ti.indices(self.dim),
1024 * 1024,
chunk_size=self.leaf_block_size ** self.dim * 8).place(
pid, offset=block_offset + (0,))
self.padding = padding
# Young's modulus and Poisson's ratio
self.E, self.nu = 1e6 * size * E_scale, 0.2
# Lame parameters
self.mu_0, self.lambda_0 = self.E / (
2 * (1 + self.nu)), self.E * self.nu / ((1 + self.nu) *
(1 - 2 * self.nu))
# Sand parameters
friction_angle = math.radians(45)
sin_phi = math.sin(friction_angle)
self.alpha = math.sqrt(2 / 3) * 2 * sin_phi / (3 - sin_phi)
# An empirically optimal chunk size is 1/10 of the expected particle number
chunk_size = 2 ** 20 if self.dim == 2 else 2 ** 23
self.particle = ti.root.dynamic(ti.i, max_num_particles, chunk_size)
if self.quant:
if not self.use_g2p2g:
self.particle.place(self.C)
if self.support_plasticity:
self.particle.place(self.Jp)
self.particle.bit_struct(num_bits=64).place(self.x)
self.particle.bit_struct(num_bits=64).place(self.v,
shared_exponent=True)
if self.dim == 3:
self.particle.bit_struct(num_bits=32).place(
self.F(0, 0), self.F(0, 1))
self.particle.bit_struct(num_bits=32).place(
self.F(0, 2), self.F(1, 0))
self.particle.bit_struct(num_bits=32).place(
self.F(1, 1), self.F(1, 2))
self.particle.bit_struct(num_bits=32).place(
self.F(2, 0), self.F(2, 1))
self.particle.bit_struct(num_bits=32).place(
self.F(2, 2), self.material)
else:
assert self.dim == 2
self.particle.bit_struct(num_bits=32).place(
self.F(0, 0), self.F(0, 1))
self.particle.bit_struct(num_bits=32).place(
self.F(1, 0), self.F(1, 1))
# no quantization on particle material in 2D
self.particle.place(self.material)
self.particle.place(self.color)
else:
self.particle.place(self.x, self.v, self.F, self.material,
self.color)
if self.support_plasticity:
self.particle.place(self.Jp)
if not self.use_g2p2g:
self.particle.place(self.C)
self.total_substeps = 0
self.unbounded = unbounded
if self.dim == 2:
self.voxelizer = None
self.set_gravity((0, -9.8))
else:
if USE_IN_BLENDER:
from .voxelizer import Voxelizer
else:
from engine.voxelizer import Voxelizer
self.voxelizer = Voxelizer(res=self.res,
dx=self.dx,
padding=self.padding,
super_sample=voxelizer_super_sample)
self.set_gravity((0, -9.8, 0))
self.voxelizer_super_sample = voxelizer_super_sample
self.grid_postprocess = []
self.add_bounding_box(self.unbounded)
self.writers = []
if not self.use_g2p2g:
self.grid = self.grid[0]
self.grid_v = self.grid_v[0]
self.grid_m = self.grid_m[0]
self.pid = self.pid[0]
trash_base_parent = ti.field(ti.i32)
trash_base_geo_center = ti.Vector.field(kDim, ti.f32)
trash_base_geo_size = ti.field(ti.f32)
trash_table = ti.root.dense(ti.i, kMaxDepth)
trash_table.place(trash_particle_id)
trash_table.place(trash_base_parent, trash_base_geo_size)
trash_table.place(trash_base_geo_center)
trash_table_len = ti.field(ti.i32, ())
node_mass = ti.field(ti.f32)
node_weighted_pos = ti.Vector.field(kDim, ti.f32)
node_particle_id = ti.field(ti.i32)
node_children = ti.field(ti.i32)
node_table = ti.root.dense(ti.i, kMaxNodes)
node_table.place(node_mass, node_particle_id, node_weighted_pos)
node_table.dense(ti.indices(*list(range(1, 1 + kDim))), 2).place(node_children)
node_table_len = ti.field(ti.i32, ())
@ti.func
def alloc_node():
ret = ti.atomic_add(node_table_len[None], 1)
assert ret < kMaxNodes
node_mass[ret] = 0
node_weighted_pos[ret] = particle_pos[0] * 0
node_particle_id[ret] = LEAF
for which in ti.grouped(ti.ndrange(*([2] * kDim))):
node_children[ret, which] = LEAF
return ret
