code: export_file = ''  # use '/tmp/mpm3d.ply' for exporting result to disk
  
  import numpy as np
  
  import taichi as ti
  
  ti.init(arch=ti.gpu)
  
  #dim, n_grid, steps, dt = 2, 128, 20, 2e-4
  #dim, n_grid, steps, dt = 2, 256, 32, 1e-4
  dim, n_grid, steps, dt = 3, 32, 25, 4e-4
  #dim, n_grid, steps, dt = 3, 64, 25, 2e-4
  #dim, n_grid, steps, dt = 3, 128, 25, 8e-5
  
  n_particles = n_grid**dim // 2**(dim - 1)
  dx = 1 / n_grid
  
  p_rho = 1
  p_vol = (dx * 0.5)**2
  p_mass = p_vol * p_rho
  gravity = 9.8
  bound = 3
  E = 400
  
  x = ti.Vector.field(dim, float, n_particles)
  v = ti.Vector.field(dim, float, n_particles)
  C = ti.Matrix.field(dim, dim, float, n_particles)
  J = ti.field(float, n_particles)
  
  grid_v = ti.Vector.field(dim, float, (n_grid, ) * dim)
  grid_m = ti.field(float, (n_grid, ) * dim)
  
  neighbour = (3, ) * dim
  
  
  @ti.kernel
  def substep():
      for I in ti.grouped(grid_m):
          grid_v[I] = ti.zero(grid_v[I])
          grid_m[I] = 0
      ti.block_dim(n_grid)
      for p in x:
          Xp = x[p] / dx
          base = int(Xp - 0.5)
          fx = Xp - base
          w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
          stress = -dt * 4 * E * p_vol * (J[p] - 1) / dx**2
          affine = ti.Matrix.identity(float, dim) * stress + p_mass * C[p]
          for offset in ti.static(ti.grouped(ti.ndrange(*neighbour))):
              dpos = (offset - fx) * dx
              weight = 1.0
              for i in ti.static(range(dim)):
                  weight *= w[offset[i]][i]
              grid_v[base + offset] += weight * (p_mass * v[p] + affine @ dpos)
              grid_m[base + offset] += weight * p_mass
      for I in ti.grouped(grid_m):
          if grid_m[I] > 0:
              grid_v[I] /= grid_m[I]
          grid_v[I][1] -= dt * gravity
          cond = (I < bound) & (grid_v[I] < 0) | \
                 (I > n_grid - bound) & (grid_v[I] > 0)
          grid_v[I] = 0 if cond else grid_v[I]
      ti.block_dim(n_grid)
      for p in x:
          Xp = x[p] / dx
          base = int(Xp - 0.5)
          fx = Xp - base
          w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
          new_v = ti.zero(v[p])
          new_C = ti.zero(C[p])
          for offset in ti.static(ti.grouped(ti.ndrange(*neighbour))):
              dpos = (offset - fx) * dx
              weight = 1.0
              for i in ti.static(range(dim)):
                  weight *= w[offset[i]][i]
              g_v = grid_v[base + offset]
              new_v += weight * g_v
              new_C += 4 * weight * g_v.outer_product(dpos) / dx**2
          v[p] = new_v
          x[p] += dt * v[p]
          J[p] *= 1 + dt * new_C.trace()
          C[p] = new_C
  
  
  @ti.kernel
  def init():
      for i in range(n_particles):
          x[i] = ti.Vector([ti.random() for i in range(dim)]) * 0.4 + 0.15
          J[i] = 1
  
  
  def T(a):
      if dim == 2:
          return a
  
      phi, theta = np.radians(28), np.radians(32)
  
      a = a - 0.5
      x, y, z = a[:, 0], a[:, 1], a[:, 2]
      c, s = np.cos(phi), np.sin(phi)
      C, S = np.cos(theta), np.sin(theta)
      x, z = x * c + z * s, z * c - x * s
      u, v = x, y * C + z * S
      return np.array([u, v]).swapaxes(0, 1) + 0.5
  
  
  init()
  gui = ti.GUI('MPM3D', background_color=0x112F41)
  while gui.running and not gui.get_event(gui.ESCAPE):
      for s in range(steps):
          substep()
      pos = x.to_numpy()
      if export_file:
          writer = ti.PLYWriter(num_vertices=n_particles)
          writer.add_vertex_pos(pos[:, 0], pos[:, 1], pos[:, 2])
          writer.export_frame(gui.frame, export_file)
      gui.circles(T(pos), radius=1.5, color=0x66ccff)
      gui.show()
  
apis: numpy.array numpy.sin numpy.radians numpy.cos