warp-lang 1.0.0b5__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

examples/example_sim_granular.py

@@ -44,9 +44,9 @@ class Example:
             cell_x=self.radius * 2.0,
             cell_y=self.radius * 2.0,
             cell_z=self.radius * 2.0,
-            pos=(0.0, 1.0, 0.0),
+            pos=wp.vec3(0.0, 1.0, 0.0),
             rot=wp.quat_identity(),
-            vel=(5.0, 0.0, 0.0),
+            vel=wp.vec3(5.0, 0.0, 0.0),
             mass=0.1,
             jitter=self.radius * 0.1,
         )

examples/example_sim_granular_collision_sdf.py

@@ -48,9 +48,9 @@ class Example:
             cell_x=self.radius * 2.0,
             cell_y=self.radius * 2.0,
             cell_z=self.radius * 2.0,
-            pos=(0.0, 20.0, 0.0),
+            pos=wp.vec3(0.0, 20.0, 0.0),
             rot=wp.quat_identity(),
-            vel=(2.0, 0.0, 0.0),
+            vel=wp.vec3(2.0, 0.0, 0.0),
             mass=0.1,
             jitter=self.radius * 0.1,
         )
@@ -67,9 +67,9 @@ class Example:
             mu=0.5,
             sdf=rock_sdf,
             body=-1,
-            pos=(0.0, 0.0, 0.0),
-            rot=wp.quat_from_axis_angle((1.0, 0.0, 0.0), -0.5 * math.pi),
-            scale=(0.01, 0.01, 0.01),
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -0.5 * math.pi),
+            scale=wp.vec3(0.01, 0.01, 0.01),
         )
 
         mins = np.array([-3.0, -3.0, -3.0])
@@ -89,7 +89,7 @@ class Example:
         sphere_vdb = wp.Volume.load_from_numpy(sphere_sdf_np, mins, voxel_size, rad + 3.0 * voxel_size)
         sphere_sdf = wp.sim.SDF(sphere_vdb)
 
-        self.sphere_pos = (3.0, 15.0, 0.0)
+        self.sphere_pos = wp.vec3(3.0, 15.0, 0.0)
         self.sphere_scale = 1.0
         self.sphere_radius = rad
         builder.add_shape_sdf(
@@ -100,7 +100,7 @@ class Example:
             sdf=sphere_sdf,
             body=-1,
             pos=self.sphere_pos,
-            scale=(self.sphere_scale, self.sphere_scale, self.sphere_scale),
+            scale=wp.vec3(self.sphere_scale, self.sphere_scale, self.sphere_scale),
         )
 
         self.model = builder.finalize()
@@ -136,21 +136,21 @@ class Example:
 
             self.renderer.begin_frame(time)
 
-            # Note the extra wp.quat_from_axis_angle((1.0, 0.0, 0.0), math.pi) is because .usd is oriented differently from .nvdb
+            # Note the extra wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi) is because .usd is oriented differently from .nvdb
             self.renderer.render_ref(
                 name="collision",
                 path=os.path.join(os.path.dirname(__file__), "assets/rocks.usd"),
-                pos=(0.0, 0.0, 0.0),
-                rot=wp.quat_from_axis_angle((1.0, 0.0, 0.0), -0.5 * math.pi)
-                * wp.quat_from_axis_angle((1.0, 0.0, 0.0), math.pi),
-                scale=(0.01, 0.01, 0.01),
+                pos=wp.vec3(0.0, 0.0, 0.0),
+                rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -0.5 * math.pi)
+                * wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi),
+                scale=wp.vec3(0.01, 0.01, 0.01),
             )
 
             self.renderer.render_sphere(
                 name="sphere",
                 pos=self.sphere_pos,
                 radius=self.sphere_scale * self.sphere_radius,
-                rot=(0.0, 0.0, 0.0, 1.0),
+                rot=wp.quat(0.0, 0.0, 0.0, 1.0),
             )
 
             self.renderer.render(self.state_0)

examples/example_sim_neo_hookean.py

@@ -80,9 +80,9 @@ class Example:
         center = cell_size * cell_dim * 0.5
 
         builder.add_soft_grid(
-            pos=(-center, 0.0, -center),
+            pos=wp.vec3(-center, 0.0, -center),
             rot=wp.quat_identity(),
-            vel=(0.0, 0.0, 0.0),
+            vel=wp.vec3(0.0, 0.0, 0.0),
             dim_x=cell_dim,
             dim_y=cell_dim,
             dim_z=cell_dim,
@@ -117,7 +117,7 @@ class Example:
         with wp.ScopedTimer("simulate"):
             xform = wp.transform(
                 (0.0, self.lift_speed * self.sim_time, 0.0),
-                wp.quat_from_axis_angle((0.0, 1.0, 0.0), self.rot_speed * self.sim_time),
+                wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), self.rot_speed * self.sim_time),
             )
             wp.launch(
                 kernel=twist_points,

examples/example_sim_particle_chain.py

@@ -38,13 +38,13 @@ class Example:
         builder = wp.sim.ModelBuilder()
 
         # anchor
-        builder.add_particle((0.0, 1.0, 0.0), (0.0, 0.0, 0.0), 0.0)
+        builder.add_particle(wp.vec3(0.0, 1.0, 0.0), wp.vec3(0.0, 0.0, 0.0), 0.0)
 
         # chain
         for i in range(1, 10):
             radius = math.sqrt(i) * 0.2
             mass = math.pi * radius * radius * radius
-            builder.add_particle((i, 1.0, 0.0), (0.0, 0.0, 0.0), mass, radius=radius)
+            builder.add_particle(wp.vec3(i, 1.0, 0.0), wp.vec3(0.0, 0.0, 0.0), mass, radius=radius)
             builder.add_spring(i - 1, i, 1.0e6, 0.0, 0)
 
         self.model = builder.finalize()

examples/example_sim_quadruped.py

@@ -70,12 +70,13 @@ def compute_env_offsets(num_envs, env_offset=(5.0, 0.0, 5.0), up_axis="Y"):
 
 class Example:
     def __init__(self, stage=None, num_envs=1, enable_rendering=True, print_timers=True):
+        self.device = wp.get_device()
         self.num_envs = num_envs
         articulation_builder = wp.sim.ModelBuilder()
         wp.sim.parse_urdf(
             os.path.join(os.path.dirname(__file__), "assets/quadruped.urdf"),
             articulation_builder,
-            xform=wp.transform([0.0, 0.7, 0.0], wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.5)),
+            xform=wp.transform([0.0, 0.7, 0.0], wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)),
             floating=True,
             density=1000,
             armature=0.01,
@@ -135,13 +136,15 @@ class Example:
 
         if self.use_graph:
             # create update graph
-            wp.capture_begin()
-            self.update()
-            self.graph = wp.capture_end()
+            wp.capture_begin(self.device)
+            try:
+                self.update()
+            finally:
+                self.graph = wp.capture_end(self.device)
 
     def update(self):
         with wp.ScopedTimer("simulate", active=True, print=self.print_timers):
-            if self.use_graph is False or self.graph is None:
+            if not self.use_graph or self.graph is None:
                 for _ in range(self.sim_substeps):
                     self.state_0.clear_forces()
                     wp.sim.collide(self.model, self.state_0)

examples/example_sim_rigid_chain.py

@@ -27,6 +27,7 @@ wp.init()
 
 class Example:
     def __init__(self, stage):
+        self.device = wp.get_device()
         self.chain_length = 8
         self.chain_width = 1.0
         self.chain_types = [
@@ -65,7 +66,7 @@ class Example:
 
                 # create shape
                 builder.add_shape_box(
-                    pos=(self.chain_width * 0.5, 0.0, 0.0),
+                    pos=wp.vec3(self.chain_width * 0.5, 0.0, 0.0),
                     hx=self.chain_width * 0.5,
                     hy=0.1,
                     hz=0.1,
@@ -146,13 +147,15 @@ class Example:
 
         if self.use_graph:
             # create update graph
-            wp.capture_begin()
-            self.update()
-            self.graph = wp.capture_end()
+            wp.capture_begin(self.device)
+            try:
+                self.update()
+            finally:
+                self.graph = wp.capture_end(self.device)
 
     def update(self):
         with wp.ScopedTimer("simulate", active=True):
-            if self.use_graph is False or self.graph is None:
+            if not self.use_graph or self.graph is None:
                 for _ in range(self.sim_substeps):
                     self.state_0.clear_forces()
                     self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)

examples/example_sim_rigid_contact.py

@@ -28,6 +28,7 @@ wp.init()
 
 class Example:
     def __init__(self, stage):
+        self.device = wp.get_device()
         builder = wp.sim.ModelBuilder()
 
         self.sim_time = 0.0
@@ -50,7 +51,7 @@ class Example:
             b = builder.add_body(origin=wp.transform((i, 1.0, 0.0), wp.quat_identity()))
 
             builder.add_shape_box(
-                pos=(0.0, 0.0, 0.0),
+                pos=wp.vec3(0.0, 0.0, 0.0),
                 hx=0.5 * self.scale,
                 hy=0.2 * self.scale,
                 hz=0.2 * self.scale,
@@ -65,7 +66,7 @@ class Example:
             b = builder.add_body(origin=wp.transform((i, 1.0, 2.0), wp.quat_identity()))
 
             builder.add_shape_sphere(
-                pos=(0.0, 0.0, 0.0), radius=0.25 * self.scale, body=b, ke=self.ke, kd=self.kd, kf=self.kf
+                pos=wp.vec3(0.0, 0.0, 0.0), radius=0.25 * self.scale, body=b, ke=self.ke, kd=self.kd, kf=self.kf
             )
 
         # capsules
@@ -73,7 +74,7 @@ class Example:
             b = builder.add_body(origin=wp.transform((i, 1.0, 6.0), wp.quat_identity()))
 
             builder.add_shape_capsule(
-                pos=(0.0, 0.0, 0.0),
+                pos=wp.vec3(0.0, 0.0, 0.0),
                 radius=0.25 * self.scale,
                 half_height=self.scale * 0.5,
                 up_axis=0,
@@ -93,15 +94,15 @@ class Example:
             b = builder.add_body(
                 origin=wp.transform(
                     (i * 0.5 * self.scale, 1.0 + i * 1.7 * self.scale, 4.0 + i * 0.5 * self.scale),
-                    wp.quat_from_axis_angle((0.0, 1.0, 0.0), math.pi * 0.1 * i),
+                    wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), math.pi * 0.1 * i),
                 )
             )
 
             builder.add_shape_mesh(
                 body=b,
                 mesh=bunny,
-                pos=(0.0, 0.0, 0.0),
-                scale=(self.scale, self.scale, self.scale),
+                pos=wp.vec3(0.0, 0.0, 0.0),
+                scale=wp.vec3(self.scale, self.scale, self.scale),
                 ke=self.ke,
                 kd=self.kd,
                 kf=self.kf,
@@ -126,9 +127,11 @@ class Example:
 
         if self.use_graph:
             # create update graph
-            wp.capture_begin()
-            self.update()
-            self.graph = wp.capture_end()
+            wp.capture_begin(self.device)
+            try:
+                self.update()
+            finally:
+                self.graph = wp.capture_end(self.device)
 
     def load_mesh(self, filename, path):
         asset_stage = Usd.Stage.Open(filename)
@@ -141,7 +144,7 @@ class Example:
 
     def update(self):
         with wp.ScopedTimer("simulate", active=True):
-            if self.use_graph is False or self.graph is None:
+            if not self.use_graph or self.graph is None:
                 for _ in range(self.sim_substeps):
                     self.state_0.clear_forces()
                     wp.sim.collide(self.model, self.state_0)

examples/example_sim_rigid_fem.py

@@ -40,9 +40,9 @@ class Example:
         builder.default_particle_radius = 0.01
 
         builder.add_soft_grid(
-            pos=(0.0, 0.0, 0.0),
+            pos=wp.vec3(0.0, 0.0, 0.0),
             rot=wp.quat_identity(),
-            vel=(0.0, 0.0, 0.0),
+            vel=wp.vec3(0.0, 0.0, 0.0),
             dim_x=20,
             dim_y=10,
             dim_z=10,

examples/example_sim_rigid_gyroscopic.py

@@ -36,7 +36,7 @@ class Example:
 
         # axis shape
         builder.add_shape_box(
-            pos=(0.3 * self.scale, 0.0, 0.0),
+            pos=wp.vec3(0.3 * self.scale, 0.0, 0.0),
             hx=0.25 * self.scale,
             hy=0.1 * self.scale,
             hz=0.1 * self.scale,
@@ -46,7 +46,7 @@ class Example:
 
         # tip shape
         builder.add_shape_box(
-            pos=(0.0, 0.0, 0.0), hx=0.05 * self.scale, hy=0.2 * self.scale, hz=1.0 * self.scale, density=100.0, body=b
+            pos=wp.vec3(0.0, 0.0, 0.0), hx=0.05 * self.scale, hy=0.2 * self.scale, hz=1.0 * self.scale, density=100.0, body=b
         )
 
         # initial spin

examples/example_sim_rigid_kinematics.py

@@ -59,7 +59,7 @@ class Example:
             # base
             builder.joint_q[coord_start : coord_start + 3] = [i * 2.0, 0.70, 0.0]
             builder.joint_q[coord_start + 3 : coord_start + 7] = wp.quat_from_axis_angle(
-                (1.0, 0.0, 0.0), -math.pi * 0.5
+                wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5
             )
 
             # joints

examples/example_sim_trajopt.py

@@ -16,7 +16,6 @@
 
 import os
 
-import matplotlib.pyplot as plt
 import numpy as np
 
 import warp as wp
@@ -72,7 +71,7 @@ class Example:
         builder = wp.sim.ModelBuilder(gravity=0.0)
         builder.add_articulation()
         b = builder.add_body(origin=wp.transform())
-        builder.add_shape_box(pos=(0.0, 0.0, 0.0), hx=0.5, hy=0.5, hz=0.5, density=100.0, body=b)
+        builder.add_shape_box(pos=wp.vec3(0.0, 0.0, 0.0), hx=0.5, hy=0.5, hz=0.5, density=100.0, body=b)
 
         # compute reference trajectory
         rad = np.linspace(0.0, np.pi * 2, self.episode_frames)
@@ -194,6 +193,8 @@ class Example:
 
 
 if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+
     stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_trajopt.usd")
     example = Example(stage_path, device=wp.get_preferred_device(), verbose=True)
 

examples/fem/example_apic_fluid.py

@@ -64,8 +64,6 @@ def update_particles(
     pos_proj = domain(lookup(domain, pos_adv))
     pos[s.qp_index] = pos_proj
 
-    return 0.0
-
 
 @integrand
 def velocity_boundary_projector_form(s: Sample, domain: Domain, u: Field, v: Field):
@@ -171,10 +169,10 @@ class Example:
         grid_cell_volume = np.prod(grid_cell_size)
 
         PARTICLES_PER_CELL_DIM = 3
-        self.radius = np.max(grid_cell_size) / (2 * PARTICLES_PER_CELL_DIM)
+        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_offset = self.radius * np.ones(3)
+        particle_grid_offset = wp.vec3(self.radius, self.radius, self.radius)
 
         np.random.seed(0)
         builder = wp.sim.ModelBuilder()
@@ -185,9 +183,9 @@ class Example:
             cell_x=self.radius * 2.0,
             cell_y=self.radius * 2.0,
             cell_z=self.radius * 2.0,
-            pos=(0.0, 0.0, 0.0) + particle_grid_offset,
+            pos=wp.vec3(0.0, 0.0, 0.0) + particle_grid_offset,
             rot=wp.quat_identity(),
-            vel=(0.0, 0.0, 0.0),
+            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,
@@ -331,7 +329,7 @@ class Example:
                 )
 
                 # (A)PIC advection
-                fem.integrate(
+                fem.interpolate(
                     update_particles,
                     quadrature=pic,
                     values={

examples/fem/example_diffusion_mgpu.py

@@ -6,19 +6,18 @@ from typing import Tuple
 
 import warp as wp
 import warp.fem as fem
-
 from warp.sparse import bsr_axpy, bsr_mv
 from warp.utils import array_cast
 
 # Import example utilities
 # Make sure that works both when imported as module and run as standalone file
 try:
-    from .example_diffusion import linear_form, diffusion_form
     from .bsr_utils import bsr_cg
+    from .example_diffusion import diffusion_form, linear_form
     from .plot_utils import Plot
 except ImportError:
-    from example_diffusion import linear_form, diffusion_form
     from bsr_utils import bsr_cg
+    from example_diffusion import diffusion_form, linear_form
     from plot_utils import Plot
 
 
@@ -100,14 +99,18 @@ class Example:
         self._quiet = quiet
 
         self._geo = fem.Grid2D(res=wp.vec2i(25))
-        self._scalar_space = fem.make_polynomial_space(self._geo, degree=3)
-        self._scalar_field = self._scalar_space.make_field()
+
+        self._main_device = wp.get_device("cuda")
+
+        with wp.ScopedDevice(self._main_device):
+            self._scalar_space = fem.make_polynomial_space(self._geo, degree=3)
+            self._scalar_field = self._scalar_space.make_field()
 
         self.renderer = Plot(stage)
 
     def update(self):
         devices = wp.get_cuda_devices()
-        main_device = devices[0]
+        main_device = self._main_device
 
         rhs_vecs = []
         res_vecs = []
@@ -141,16 +144,15 @@ class Example:
         A.tmp_buf = tmp
         A.rank_data = (matrices, rhs_vecs, res_vecs, indices)
 
-        with wp.ScopedDevice(main_device):
-            bsr_cg(
-                A,
-                x=global_res,
-                b=glob_rhs,
-                use_diag_precond=False,
-                quiet=self._quiet,
-                mv_routine=mv_routine,
-                device=main_device,
-            )
+        bsr_cg(
+            A,
+            x=global_res,
+            b=glob_rhs,
+            use_diag_precond=False,
+            quiet=self._quiet,
+            mv_routine=mv_routine,
+            device=main_device,
+        )
 
         array_cast(in_array=global_res, out_array=self._scalar_field.dof_values)
 

warp/__init__.py

@@ -26,7 +26,9 @@ from warp.types import spatial_matrix, spatial_matrixh, spatial_matrixf, spatial
 
 # geometry types
 from warp.types import Bvh, Mesh, HashGrid, Volume, MarchingCubes
-from warp.types import bvh_query_t, mesh_query_aabb_t, hash_grid_query_t
+from warp.types import bvh_query_t, hash_grid_query_t, mesh_query_aabb_t, mesh_query_point_t, mesh_query_ray_t
+
+
 
 # device-wide gemms
 from warp.types import matmul, adj_matmul, batched_matmul, adj_batched_matmul, from_ptr
@@ -57,7 +59,6 @@ from warp.context import (
 )
 from warp.context import set_module_options, get_module_options, get_module
 from warp.context import capture_begin, capture_end, capture_launch
-from warp.context import print_builtins, export_builtins, export_stubs
 from warp.context import Kernel, Function, Launch
 from warp.context import Stream, get_stream, set_stream, synchronize_stream
 from warp.context import Event, record_event, wait_event, wait_stream

warp/build_dll.py

@@ -8,11 +8,24 @@
 import sys
 import os
 import subprocess
+import platform
 
 import warp.config
 from warp.utils import ScopedTimer
 
 
+# return a canonical machine architecture string
+# - "x86_64" for x86-64, aka. AMD64, aka. x64
+# - "aarch64" for AArch64, aka. ARM64
+def machine_architecture() -> str:
+    machine = platform.machine()
+    if machine == "x86_64" or machine == "AMD64":
+        return "x86_64"
+    if machine == "aarch64" or machine == "arm64":
+        return "aarch64"
+    raise RuntimeError(f"Unrecognized machine architecture {machine}")
+
+
 def run_cmd(cmd, capture=False):
     if warp.config.verbose:
         print(cmd)
@@ -182,12 +195,15 @@ def build_dll_for_arch(dll_path, cpp_paths, cu_path, libs, mode, arch, verify_fp
                 "-gencode=arch=compute_75,code=sm_75",  # Turing
                 "-gencode=arch=compute_80,code=sm_80",  # Ampere
                 "-gencode=arch=compute_86,code=sm_86",
-                # SASS for supported mobile architectures (e.g. Tegra/Jetson)
-                # "-gencode=arch=compute_53,code=sm_53",
-                # "-gencode=arch=compute_62,code=sm_62",
-                # "-gencode=arch=compute_72,code=sm_72",
-                # "-gencode=arch=compute_87,code=sm_87",
             ]
+            if arch == "aarch64" and sys.platform == "linux":
+                gencode_opts += [
+                    # SASS for supported mobile architectures (e.g. Tegra/Jetson)
+                    "-gencode=arch=compute_53,code=sm_53",  # X1
+                    "-gencode=arch=compute_62,code=sm_62",  # X2
+                    "-gencode=arch=compute_72,code=sm_72",  # Xavier
+                    "-gencode=arch=compute_87,code=sm_87",  # Orin
+                ]
 
             # support for Ada and Hopper is available with CUDA Toolkit 11.8+
             if ctk_version >= (11, 8):
@@ -354,11 +370,15 @@ def build_dll(dll_path, cpp_paths, cu_path, libs=[], mode="release", verify_fp=F
     if sys.platform == "darwin":
         # create a universal binary by combining x86-64 and AArch64 builds
         build_dll_for_arch(dll_path + "-x86_64", cpp_paths, cu_path, libs, mode, "x86_64", verify_fp, fast_math, quick)
-        build_dll_for_arch(dll_path + "-arm64", cpp_paths, cu_path, libs, mode, "arm64", verify_fp, fast_math, quick)
+        build_dll_for_arch(
+            dll_path + "-aarch64", cpp_paths, cu_path, libs, mode, "aarch64", verify_fp, fast_math, quick
+        )
 
-        run_cmd(f"lipo -create -output {dll_path} {dll_path}-x86_64 {dll_path}-arm64")
+        run_cmd(f"lipo -create -output {dll_path} {dll_path}-x86_64 {dll_path}-aarch64")
         os.remove(f"{dll_path}-x86_64")
-        os.remove(f"{dll_path}-arm64")
+        os.remove(f"{dll_path}-aarch64")
 
     else:
-        build_dll_for_arch(dll_path, cpp_paths, cu_path, libs, mode, "x86_64", verify_fp, fast_math, quick)
+        build_dll_for_arch(
+            dll_path, cpp_paths, cu_path, libs, mode, machine_architecture(), verify_fp, fast_math, quick
+        )