diff --git a/demo/demo_3d.py b/demo/demo_3d.py index 50b7954..3e89b9a 100644 --- a/demo/demo_3d.py +++ b/demo/demo_3d.py @@ -25,6 +25,12 @@ mpm.set_gravity((0, -50, 0)) +mpm.add_wind_field( + lower_bound=[0, 0, 0], + upper_bound=[3, 10, 10], + force=[200, 0, 0] +) + for frame in range(1500): mpm.step(4e-3) colors = np.array([0x068587, 0xED553B, 0xEEEEF0, 0xFFFF00], diff --git a/engine/mpm_solver.py b/engine/mpm_solver.py index b898c37..3098ec4 100644 --- a/engine/mpm_solver.py +++ b/engine/mpm_solver.py @@ -318,6 +318,26 @@ def set_gravity(self, g): assert len(g) == self.dim self.gravity[None] = g + def add_wind_field(self, lower_bound, upper_bound, force): + lower_bound = list(lower_bound) + upper_bound = list(upper_bound) + force = list(force) + assert len(lower_bound) == self.dim + assert len(upper_bound) == self.dim + assert len(force) == self.dim + + @ti.kernel + def apply_wind(t: ti.f32, dt: ti.f32, grid_v: ti.template()): + for I in ti.grouped(grid_v): + grid_pos = I * self.dx + inside = True + for d in ti.static(range(self.dim)): + inside = inside and (lower_bound[d] <= grid_pos[d] < upper_bound[d]) + if inside: + grid_v[I] += dt * ti.Vector(force) + + self.grid_postprocess.append(apply_wind) + @ti.func def sand_projection(self, sigma, p): sigma_out = ti.Matrix.zero(ti.f32, self.dim, self.dim)