def save_mpm_state(solver: MPMSolver, frame: int, save_dir: str):
"""
Save MPM middle state as a npz file
:param solver:
:param frame:
:param save_dir:
:return:
"""
if not solver.use_g2p2g:
raise NotImplementedError
particles = solver.particle_info() # particle information
# other meta data
phase = solver.input_grid
# save grid_v
count_activate(solver.grid_v[phase])
print(f"we have {count[None]} nodes activated")
np_grid_idx = np.ndarray((count[None], solver.dim), dtype=np.float32)
np_grid_val = np.ndarray((count[None], solver.dim), dtype=np.float32)
copy_grid(np_grid_idx, np_grid_val, solver.grid_v[phase], solver)
# save deformation gradient
np_F = np.ndarray((solver.n_particles[None], solver.dim, solver.dim),
dtype=np.float32)
copy_matrix(np_F, solver.F, solver)
particles['F'] = np_F
if mpm.support_plasticity:
np_j = np.ndarray((solver.n_particles[None], ), dtype=np.float32)
solver.copy_dynamic(np_j, solver.Jp)
particles['p_Jp'] = np_j
np.savez(save_dir,
frame=frame,
input_grid=phase,
grid_v_idx=np_grid_idx,
grid_v_val=np_grid_val,
**particles)
print(f'save {frame}th frame to {save_dir}')
return
# Try to run on GPU
ti.init(arch=ti.cuda, device_memory_GB=3.0)
gui = ti.GUI("Taichi Elements", res=512, background_color=0x112F41)
mpm = MPMSolver(res=(64, 64, 64), size=1)
triangles = np.fromfile('suzanne.npy', dtype=np.float32)
triangles = np.reshape(triangles, (len(triangles) // 9, 9)) * 0.306 + 0.501
mpm.add_mesh(triangles=triangles,
material=MPMSolver.material_elastic,
color=0xFFFF00)
mpm.set_gravity((0, -20, 0))
for frame in range(1500):
mpm.step(4e-3)
particles = mpm.particle_info()
np_x = particles['position'] / 1.0
# simple camera transform
screen_x = ((np_x[:, 0] + np_x[:, 2]) / 2**0.5) - 0.2
screen_y = (np_x[:, 1])
screen_pos = np.stack([screen_x, screen_y], axis=-1)
gui.circles(screen_pos, radius=1.1, color=particles['color'])
gui.show(f'{frame:06d}.png' if write_to_disk else None)