warp-lang 1.1.0__py3-none-macosx_10_13_universal2.whl → 1.2.1__py3-none-macosx_10_13_universal2.whl

warp/examples/benchmarks/benchmark_api.py

@@ -32,8 +32,6 @@ def inc_kernel(a: wp.array(dtype=float)):
     a[tid] = a[tid] + 1.0
 
 
-wp.init()
-
 # configure devices
 for target_device in wp.get_cuda_devices():
     try:

warp/examples/benchmarks/benchmark_cloth_warp.py

@@ -7,7 +7,6 @@
 
 import warp as wp
 
-wp.init()
 wp.build.clear_kernel_cache()
 
 

warp/examples/benchmarks/benchmark_launches.py

@@ -128,8 +128,6 @@ def ksz(s: Sz):
     tid = wp.tid()  # noqa: F841
 
 
-wp.init()
-
 wp.build.clear_kernel_cache()
 
 devices = wp.get_devices()

warp/examples/core/example_dem.py

@@ -19,8 +19,6 @@ import numpy as np
 import warp as wp
 import warp.render
 
-wp.init()
-
 
 @wp.func
 def contact_force(n: wp.vec3, v: wp.vec3, c: float, k_n: float, k_d: float, k_f: float, k_mu: float):

warp/examples/core/example_fluid.py

@@ -18,8 +18,6 @@ import math
 import warp as wp
 import warp.render
 
-wp.init()
-
 grid_width = wp.constant(256)
 grid_height = wp.constant(128)
 

warp/examples/core/example_graph_capture.py

@@ -16,8 +16,6 @@ import numpy as np
 
 import warp as wp
 
-wp.init()
-
 
 @wp.kernel
 def fbm(

warp/examples/core/example_marching_cubes.py

@@ -17,8 +17,6 @@
 import warp as wp
 import warp.render
 
-wp.init()
-
 
 @wp.func
 def sdf_create_box(pos: wp.vec3, size: wp.vec3):

warp/examples/core/example_mesh.py

@@ -24,8 +24,6 @@ import warp as wp
 import warp.examples
 import warp.render
 
-wp.init()
-
 
 @wp.kernel
 def deform(positions: wp.array(dtype=wp.vec3), t: float):

warp/examples/core/example_mesh_intersect.py

@@ -21,8 +21,6 @@ import warp as wp
 import warp.examples
 import warp.render
 
-wp.init()
-
 
 @wp.func
 def cw_min(a: wp.vec3, b: wp.vec3):

warp/examples/core/example_nvdb.py

@@ -23,8 +23,6 @@ import warp as wp
 import warp.examples
 import warp.render
 
-wp.init()
-
 
 @wp.func
 def volume_grad(volume: wp.uint64, p: wp.vec3):

warp/examples/core/example_raycast.py

@@ -21,8 +21,6 @@ from pxr import Usd, UsdGeom
 import warp as wp
 import warp.examples
 
-wp.init()
-
 
 @wp.kernel
 def draw(mesh: wp.uint64, cam_pos: wp.vec3, width: int, height: int, pixels: wp.array(dtype=wp.vec3)):

warp/examples/core/example_raymarch.py

@@ -16,8 +16,6 @@
 
 import warp as wp
 
-wp.init()
-
 
 @wp.func
 def sdf_sphere(p: wp.vec3, r: float):

warp/examples/core/example_render_opengl.py

@@ -18,8 +18,6 @@ import numpy as np
 import warp as wp
 import warp.render
 
-wp.init()
-
 
 class Example:
     def __init__(self, num_tiles=4, custom_tile_arrangement=False):

warp/examples/core/example_sph.py

@@ -25,8 +25,6 @@ import numpy as np
 import warp as wp
 import warp.render
 
-wp.init()
-
 
 @wp.func
 def square(x: float):

warp/examples/core/example_torch.py

@@ -23,9 +23,6 @@ import torch
 
 import warp as wp
 
-wp.init()
-
-
 pvec2 = wp.types.vector(length=2, dtype=wp.float32)
 
 

warp/examples/core/example_wave.py

@@ -18,8 +18,6 @@ import math
 import warp as wp
 import warp.render
 
-wp.init()
-
 
 @wp.func
 def sample(f: wp.array(dtype=float), x: int, y: int, width: int, height: int):

warp/examples/fem/example_apic_fluid.py

@@ -8,7 +8,8 @@
 ###########################################################################
 # Example APIC Fluid Simulation
 #
-# Shows how to implement a apic fluid simulation.
+# Shows how to implement a minimalist APIC fluid simulation using a NanoVDB
+# grid and the PicQuadrature class.
 ###########################################################################
 
 import numpy as np
@@ -16,16 +17,22 @@ import numpy as np
 import warp as wp
 import warp.fem as fem
 import warp.sim.render
-from warp.fem import Domain, Field, Sample, at_node, div, grad, integrand, lookup, normal
+from warp.fem import Domain, Field, Sample, at_node, div, grad, integrand
 from warp.sim import Model, State
-from warp.sparse import BsrMatrix, bsr_copy, bsr_mm, bsr_mv, bsr_transposed, bsr_zeros
+from warp.sparse import BsrMatrix, bsr_mm, bsr_mv, bsr_transposed
 
 try:
     from .bsr_utils import bsr_cg
 except ImportError:
     from bsr_utils import bsr_cg
 
-wp.init()
+
+@wp.func
+def collision_sdf(x: wp.vec3):
+    # Arbitrary sdf representing collision geometry
+    # Here an inverted half-ball of radius 10
+    x[1] = wp.min(x[1], 0.0)
+    return 10.0 - wp.length(x), -wp.normalize(x)
 
 
 @integrand
@@ -48,6 +55,13 @@ def integrate_velocity(
     vel_apic = velocities[s.qp_index] + velocity_gradients[s.qp_index] * node_offset
 
     vel_adv = vel_apic + dt * gravity
+
+    # if inside collider, remove normal velocity
+    sdf, sdf_gradient = collision_sdf(domain(s))
+    if sdf <= 0:
+        v_n = wp.dot(vel_adv, sdf_gradient)
+        vel_adv -= wp.max(v_n, 0.0) * sdf_gradient
+
     return wp.dot(u(s), vel_adv)
 
 
@@ -63,31 +77,33 @@ def update_particles(
     vel_grad: wp.array(dtype=wp.mat33),
 ):
     """Read particle velocity from grid and advect positions"""
-    vel[s.qp_index] = grid_vel(s)
+    p_vel = grid_vel(s)
     vel_grad[s.qp_index] = grad(grid_vel, s)
 
-    pos_adv = pos_prev[s.qp_index] + dt * vel[s.qp_index]
+    pos_adv = pos_prev[s.qp_index] + dt * p_vel
 
-    # Project onto domain
-    pos_proj = domain(lookup(domain, pos_adv))
-    pos[s.qp_index] = pos_proj
+    pos[s.qp_index] = pos_adv
+    vel[s.qp_index] = p_vel
 
 
 @integrand
 def velocity_boundary_projector_form(s: Sample, domain: Domain, u: Field, v: Field):
     """Projector for velocity-Dirichlet boundary conditions"""
 
-    n = normal(domain, s)
-    if n[1] > 0.0:
-        # Neuman  on top
+    x = domain(s)
+    sdf, sdf_normal = collision_sdf(x)
+
+    if sdf > 0.0:
+        # Neuman
         return 0.0
 
-    # Free-slip on other sides
-    return wp.dot(u(s), n) * wp.dot(v(s), n)
+    # Free-slip on boundary
+    return wp.dot(u(s), sdf_normal) * wp.dot(v(s), sdf_normal)
 
 
 @integrand
-def divergence_form(s: Sample, u: Field, psi: Field):
+def divergence_form(s: Sample, domain: Domain, u: Field, psi: Field):
+    # Divergence bilinear form
     return div(u, s) * psi(s)
 
 
@@ -95,10 +111,7 @@ def divergence_form(s: Sample, u: Field, psi: Field):
 def invert_volume_kernel(values: wp.array(dtype=float)):
     i = wp.tid()
     m = values[i]
-    if m <= 1.0e-8:
-        values[i] = 0.0
-    else:
-        values[i] = 1.0 / m
+    values[i] = wp.select(m == 0.0, 1.0 / m, 0.0)
 
 
 @wp.kernel
@@ -117,6 +130,8 @@ def scale_transposed_divergence_mat(
     tr_divergence_mat_values: wp.array(dtype=wp.mat(shape=(3, 1), dtype=float)),
     inv_fraction_int: wp.array(dtype=float),
 ):
+    # In-place scaling of gradient operator rows wiht inverse mass
+
     u_i = wp.tid()
     block_beg = tr_divergence_mat_offsets[u_i]
     block_end = tr_divergence_mat_offsets[u_i + 1]
@@ -125,6 +140,15 @@ def scale_transposed_divergence_mat(
         tr_divergence_mat_values[b] = tr_divergence_mat_values[b] * inv_fraction_int[u_i]
 
 
+@wp.kernel
+def compute_particle_ijk(positions: wp.array(dtype=wp.vec3), voxel_size: float, ijks: wp.array(dtype=wp.vec3i)):
+    # Index-space coordinates of grid cell containing each particle
+
+    p = wp.tid()
+    pos = positions[p] / voxel_size
+    ijks[p] = wp.vec3i(int(wp.floor(pos[0])), int(wp.floor(pos[1])), int(wp.floor(pos[2])))
+
+
 def solve_incompressibility(divergence_mat: BsrMatrix, inv_volume, pressure, velocity, quiet: bool = False):
     """Solve for divergence-free velocity delta:
 
@@ -144,43 +168,45 @@ def solve_incompressibility(divergence_mat: BsrMatrix, inv_volume, pressure, vel
         ],
     )
 
-    # For simplicity, assemble schur complement and solve with CG
+    # For simplicity, assemble Schur complement and solve with CG
     schur = bsr_mm(divergence_mat, transposed_divergence_mat)
 
     rhs = wp.zeros_like(pressure)
     bsr_mv(A=divergence_mat, x=velocity, y=rhs, alpha=-1.0, beta=0.0)
-    bsr_cg(schur, b=rhs, x=pressure, quiet=quiet)
+    bsr_cg(schur, b=rhs, x=pressure, quiet=quiet, tol=1.0e-6)
 
     # Apply pressure to velocity
     bsr_mv(A=transposed_divergence_mat, x=pressure, y=velocity, alpha=1.0, beta=1.0)
 
 
 class Example:
-    def __init__(self, quiet=False, stage_path="example_apic_fluid.usd", res=(32, 64, 16)):
+    def __init__(self, quiet=False, stage_path="example_apic_fluid.usd", voxel_size=1.0):
         fps = 60
         self.frame_dt = 1.0 / fps
         self.current_frame = 0
 
         self.sim_substeps = 1
         self.sim_dt = self.frame_dt / self.sim_substeps
+        self.voxel_size = voxel_size
 
         self._quiet = quiet
 
-        # grid dimensions and particle emission
-        grid_res = np.array(res, dtype=int)
-        particle_fill_frac = np.array([0.5, 0.5, 1.0])
-        grid_lo = wp.vec3(0.0)
-        grid_hi = wp.vec3(50, 100, 25)
+        # particle emission
+        particle_grid_lo = wp.vec3(-5)
+        particle_grid_hi = wp.vec3(5)
 
-        grid_cell_size = np.array(grid_hi - grid_lo) / grid_res
+        grid_cell_size = voxel_size
         grid_cell_volume = np.prod(grid_cell_size)
 
-        PARTICLES_PER_CELL_DIM = 3
+        PARTICLES_PER_CELL_DIM = 2
         self.radius = float(np.max(grid_cell_size) / (2 * PARTICLES_PER_CELL_DIM))
 
-        particle_grid_res = np.array(particle_fill_frac * grid_res * PARTICLES_PER_CELL_DIM, dtype=int)
+        particle_grid_res = (
+            np.array((particle_grid_hi - particle_grid_lo) / voxel_size, dtype=int) * PARTICLES_PER_CELL_DIM
+        )
         particle_grid_offset = wp.vec3(self.radius, self.radius, self.radius)
 
+        # Initialize warp.sim model, spawn particles
         np.random.seed(0)
         builder = wp.sim.ModelBuilder()
         builder.add_particle_grid(
@@ -190,60 +216,17 @@ class Example:
             cell_x=self.radius * 2.0,
             cell_y=self.radius * 2.0,
             cell_z=self.radius * 2.0,
-            pos=wp.vec3(0.0, 0.0, 0.0) + particle_grid_offset,
+            pos=particle_grid_lo + particle_grid_offset,
             rot=wp.quat_identity(),
             vel=wp.vec3(0.0, 0.0, 0.0),
             mass=grid_cell_volume / PARTICLES_PER_CELL_DIM**3,
             jitter=self.radius * 1.0,
-            radius_mean=self.radius,
-        )
-
-        self.grid = fem.Grid3D(wp.vec3i(grid_res), grid_lo, grid_hi)
-
-        # Function spaces
-        self.velocity_space = fem.make_polynomial_space(self.grid, dtype=wp.vec3, degree=1)
-        self.fraction_space = fem.make_polynomial_space(self.grid, dtype=float, degree=1)
-        self.strain_space = fem.make_polynomial_space(
-            self.grid,
-            dtype=float,
-            degree=0,
         )
-
-        self.pressure_field = self.strain_space.make_field()
-        self.velocity_field = self.velocity_space.make_field()
-
-        # Test and trial functions
-        self.domain = fem.Cells(self.grid)
-        self.velocity_test = fem.make_test(self.velocity_space, domain=self.domain)
-        self.velocity_trial = fem.make_trial(self.velocity_space, domain=self.domain)
-        self.fraction_test = fem.make_test(self.fraction_space, domain=self.domain)
-        self.strain_test = fem.make_test(self.strain_space, domain=self.domain)
-        self.strain_trial = fem.make_trial(self.strain_space, domain=self.domain)
-
-        # Enforcing the Dirichlet boundary condition the hard way;
-        # build projector for velocity left- and right-hand-sides
-        boundary = fem.BoundarySides(self.grid)
-        u_bd_test = fem.make_test(space=self.velocity_space, domain=boundary)
-        u_bd_trial = fem.make_trial(space=self.velocity_space, domain=boundary)
-        u_bd_projector = fem.integrate(
-            velocity_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True, output_dtype=float
-        )
-
-        fem.normalize_dirichlet_projector(u_bd_projector)
-        self.vel_bd_projector = u_bd_projector
+        self.model: Model = builder.finalize()
 
         # Storage for temporary variables
         self.temporary_store = fem.TemporaryStore()
 
-        self._divergence_matrix = bsr_zeros(
-            self.strain_space.node_count(),
-            self.velocity_space.node_count(),
-            block_type=wp.mat(shape=(1, 3), dtype=float),
-        )
-
-        # Warp.sim model
-        self.model: Model = builder.finalize()
-
         if not self._quiet:
             print("Particle count:", self.model.particle_count)
 
@@ -265,79 +248,122 @@ class Example:
         fem.set_default_temporary_store(self.temporary_store)
 
         self.current_frame = self.current_frame + 1
+
+        particle_ijk = wp.empty(self.state_0.particle_count, dtype=wp.vec3i)
+
         with wp.ScopedTimer(f"simulate frame {self.current_frame}", active=True):
             for _s in range(self.sim_substeps):
-                # Bin particles to grid cells
-                pic = fem.PicQuadrature(
-                    domain=fem.Cells(self.grid), positions=self.state_0.particle_q, measures=self.model.particle_mass
+                # Compute the voxel coordinates for each particle.
+                # `Volume.allocate_by_voxels` accepts world positions, but allocates
+                # the voxels with the closest origin rather than the enclosing ones
+                # (i.e, it "round"s the positions, while here we eant to "floor" it)
+                wp.launch(
+                    compute_particle_ijk,
+                    dim=particle_ijk.shape,
+                    inputs=[self.state_0.particle_q, self.voxel_size, particle_ijk],
+                )
+
+                # Allocate the voxels and create the warp.fem geometry
+                volume = wp.Volume.allocate_by_voxels(
+                    voxel_points=particle_ijk,
+                    voxel_size=self.voxel_size,
+                )
+                grid = fem.Nanogrid(volume)
+
+                # Define function spaces: linear (Q1) for velocity and volume fraction,
+                # piecewise-constant for pressure
+                linear_basis_space = fem.make_polynomial_basis_space(grid, degree=1)
+                velocity_space = fem.make_collocated_function_space(linear_basis_space, dtype=wp.vec3)
+                fraction_space = fem.make_collocated_function_space(linear_basis_space, dtype=float)
+                strain_space = fem.make_polynomial_space(
+                    grid,
+                    dtype=float,
+                    degree=0,
+                    discontinuous=True,
                 )
 
-                # Borrow some temporary arrays for storing integration results
-                inv_volume_temporary = fem.borrow_temporary(
-                    self.temporary_store, shape=(self.fraction_space.node_count()), dtype=float
+                pressure_field = strain_space.make_field()
+                velocity_field = velocity_space.make_field()
+
+                # Define test and trial functions and integrating linear and bilinear forms
+                domain = fem.Cells(grid)
+                velocity_test = fem.make_test(velocity_space, domain=domain)
+                velocity_trial = fem.make_trial(velocity_space, domain=domain)
+                fraction_test = fem.make_test(fraction_space, domain=domain)
+                strain_test = fem.make_test(strain_space, domain=domain)
+
+                # Build projector for Dirichlet boundary conditions
+                vel_projector = fem.integrate(
+                    velocity_boundary_projector_form,
+                    fields={"u": velocity_trial, "v": velocity_test},
+                    nodal=True,
+                    output_dtype=float,
                 )
-                velocity_int_temporary = fem.borrow_temporary(
-                    self.temporary_store, shape=(self.velocity_space.node_count()), dtype=wp.vec3
+                fem.normalize_dirichlet_projector(vel_projector)
+
+                # Bin particles to grid cells
+                pic = fem.PicQuadrature(
+                    domain=domain, positions=self.state_0.particle_q, measures=self.model.particle_mass
                 )
-                inv_volume = inv_volume_temporary.array
-                velocity_int = velocity_int_temporary.array
 
-                # Inverse volume fraction
-                fem.integrate(
+                # Compute inverse particle volume for each grid node
+                inv_volume = fem.integrate(
                     integrate_fraction,
                     quadrature=pic,
-                    fields={"phi": self.fraction_test},
-                    accumulate_dtype=float,
-                    output=inv_volume,
+                    fields={"phi": fraction_test},
+                    output_dtype=float,
                 )
                 wp.launch(kernel=invert_volume_kernel, dim=inv_volume.shape, inputs=[inv_volume])
 
                 # Velocity right-hand side
-                fem.integrate(
+                velocity_int = fem.integrate(
                     integrate_velocity,
                     quadrature=pic,
-                    fields={"u": self.velocity_test},
+                    fields={"u": velocity_test},
                     values={
                         "velocities": self.state_0.particle_qd,
                         "velocity_gradients": self.state_0.particle_qd_grad,
                         "dt": self.sim_dt,
                         "gravity": self.model.gravity,
                     },
-                    accumulate_dtype=float,
-                    output=velocity_int,
+                    output_dtype=wp.vec3,
                 )
 
                 # Compute constraint-free velocity
                 wp.launch(
                     kernel=scalar_vector_multiply,
                     dim=inv_volume.shape[0],
-                    inputs=[inv_volume, velocity_int, self.velocity_field.dof_values],
+                    inputs=[inv_volume, velocity_int, velocity_field.dof_values],
                 )
 
                 # Apply velocity boundary conditions:
-                # velocity -= vel_bd_projector * velocity
-                wp.copy(src=self.velocity_field.dof_values, dest=velocity_int)
-                bsr_mv(A=self.vel_bd_projector, x=velocity_int, y=self.velocity_field.dof_values, alpha=-1.0, beta=1.0)
+                # velocity -= vel_projector * velocity
+                bsr_mv(
+                    A=vel_projector,
+                    x=velocity_field.dof_values,
+                    y=velocity_field.dof_values,
+                    alpha=-1.0,
+                    beta=1.0,
+                )
 
-                # Divergence matrix
-                fem.integrate(
+                # Assemble divergence operator matrix
+                divergence_matrix = fem.integrate(
                     divergence_form,
                     quadrature=pic,
-                    fields={"u": self.velocity_trial, "psi": self.strain_test},
-                    accumulate_dtype=float,
-                    output=self._divergence_matrix,
+                    fields={"u": velocity_trial, "psi": strain_test},
+                    output_dtype=float,
                 )
 
                 # Project matrix to enforce boundary conditions
-                divergence_mat_tmp = bsr_copy(self._divergence_matrix)
-                bsr_mm(alpha=-1.0, x=divergence_mat_tmp, y=self.vel_bd_projector, z=self._divergence_matrix, beta=1.0)
+                # divergence_matrix -= divergence_matrix * vel_projector
+                bsr_mm(alpha=-1.0, x=divergence_matrix, y=vel_projector, z=divergence_matrix, beta=1.0)
 
                 # Solve unilateral incompressibility
                 solve_incompressibility(
-                    self._divergence_matrix,
+                    divergence_matrix,
                     inv_volume,
-                    self.pressure_field.dof_values,
-                    self.velocity_field.dof_values,
+                    pressure_field.dof_values,
+                    velocity_field.dof_values,
                     quiet=self._quiet,
                 )
 
@@ -352,7 +378,7 @@ class Example:
                         "vel_grad": self.state_1.particle_qd_grad,
                         "dt": self.sim_dt,
                     },
-                    fields={"grid_vel": self.velocity_field},
+                    fields={"grid_vel": velocity_field},
                 )
 
                 # swap states
@@ -385,19 +411,18 @@ if __name__ == "__main__":
         default="example_apic_fluid.usd",
         help="Path to the output USD file.",
     )
-    parser.add_argument("--num_frames", type=int, default=1000, help="Total number of frames.")
+    parser.add_argument("--num_frames", type=int, default=250, help="Total number of frames.")
     parser.add_argument("--quiet", action="store_true")
     parser.add_argument(
-        "--res",
-        type=lambda s: [int(item) for item in s.split(",")],
-        default="32,64,16",
-        help="Delimited list specifying resolution in x, y, and z.",
+        "--voxel_size",
+        type=float,
+        default=0.25,
     )
 
     args = parser.parse_known_args()[0]
 
     with wp.ScopedDevice(args.device):
-        example = Example(quiet=args.quiet, stage_path=args.stage_path, res=args.res)
+        example = Example(quiet=args.quiet, stage_path=args.stage_path, voxel_size=args.voxel_size)
 
         for _ in range(args.num_frames):
             example.step()