warp-lang 0.11.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/examples/sim/example_rigid_soft_contact.py

@@ -0,0 +1,124 @@
+# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+###########################################################################
+# Example Sim Rigid FEM
+#
+# Shows how to set up a rigid sphere colliding with an FEM beam
+# using wp.sim.ModelBuilder().
+#
+###########################################################################
+
+import os
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+wp.init()
+
+
+class Example:
+    def __init__(self, stage):
+        self.sim_width = 8
+        self.sim_height = 8
+
+        self.sim_fps = 60.0
+        self.frame_dt = 1.0 / self.sim_fps
+        self.sim_substeps = 32
+        self.sim_duration = 5.0
+        self.sim_frames = int(self.sim_duration * self.sim_fps)
+        self.sim_dt = (1.0 / self.sim_fps) / self.sim_substeps
+        self.sim_time = 0.0
+        self.sim_iterations = 1
+        self.sim_relaxation = 1.0
+        self.profiler = {}
+
+        builder = wp.sim.ModelBuilder()
+        builder.default_particle_radius = 0.01
+
+        builder.add_soft_grid(
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            rot=wp.quat_identity(),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            dim_x=20,
+            dim_y=10,
+            dim_z=10,
+            cell_x=0.1,
+            cell_y=0.1,
+            cell_z=0.1,
+            density=100.0,
+            k_mu=50000.0,
+            k_lambda=20000.0,
+            k_damp=0.0,
+        )
+
+        b = builder.add_body(origin=wp.transform((0.5, 2.5, 0.5), wp.quat_identity()))
+        builder.add_shape_sphere(body=b, radius=0.75, density=100.0)
+
+        self.model = builder.finalize()
+        self.model.ground = True
+        self.model.soft_contact_ke = 1.0e3
+        self.model.soft_contact_kd = 0.0
+        self.model.soft_contact_kf = 1.0e3
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+
+        self.state_0 = self.model.state()
+        self.state_1 = self.model.state()
+
+        self.renderer = None
+        if stage:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=1.0)
+
+        self.use_graph = wp.get_device().is_cuda
+        if self.use_graph:
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
+
+    def simulate(self):
+        for s in range(self.sim_substeps):
+            wp.sim.collide(self.model, self.state_0)
+
+            self.state_0.clear_forces()
+            self.state_1.clear_forces()
+
+            self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
+
+            # swap states
+            (self.state_0, self.state_1) = (self.state_1, self.state_0)
+
+    def step(self):
+        with wp.ScopedTimer("step", dict=self.profiler):
+            if self.use_graph:
+                wp.capture_launch(self.graph)
+            else:
+                self.simulate()
+            self.sim_time += self.frame_dt
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render", active=True):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render(self.state_0)
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "example_rigid_soft_contact.usd")
+
+    example = Example(stage_path)
+
+    for i in range(example.sim_frames):
+        example.step()
+        example.render()
+
+    if example.renderer:
+        example.renderer.save()

warp/examples/sim/example_soft_body.py

@@ -0,0 +1,178 @@
+# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+###########################################################################
+# Example Sim Neo-Hookean
+#
+# Shows a simulation of an Neo-Hookean FEM beam being twisted through a
+# 180 degree rotation.
+#
+###########################################################################
+
+import math
+import os
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+wp.init()
+
+
+@wp.kernel
+def twist_points(
+    rest: wp.array(dtype=wp.vec3), points: wp.array(dtype=wp.vec3), mass: wp.array(dtype=float), xform: wp.transform
+):
+    tid = wp.tid()
+
+    r = rest[tid]
+    p = points[tid]
+    m = mass[tid]
+
+    # twist the top layer of particles in the beam
+    if m == 0 and p[1] != 0.0:
+        points[tid] = wp.transform_point(xform, r)
+
+
+@wp.kernel
+def compute_volume(points: wp.array(dtype=wp.vec3), indices: wp.array2d(dtype=int), volume: wp.array(dtype=float)):
+    tid = wp.tid()
+
+    i = indices[tid, 0]
+    j = indices[tid, 1]
+    k = indices[tid, 2]
+    l = indices[tid, 3]
+
+    x0 = points[i]
+    x1 = points[j]
+    x2 = points[k]
+    x3 = points[l]
+
+    x10 = x1 - x0
+    x20 = x2 - x0
+    x30 = x3 - x0
+
+    v = wp.dot(x10, wp.cross(x20, x30)) / 6.0
+
+    wp.atomic_add(volume, 0, v)
+
+
+class Example:
+    def __init__(self, stage):
+        sim_fps = 60.0
+        self.sim_substeps = 64
+        sim_duration = 5.0
+        self.sim_frames = int(sim_duration * sim_fps)
+        self.frame_dt = 1.0 / sim_fps
+        self.sim_dt = (1.0 / sim_fps) / self.sim_substeps
+        self.sim_time = 0.0
+        self.lift_speed = 2.5 / sim_duration * 2.0  # from Smith et al.
+        self.rot_speed = math.pi / sim_duration
+
+        builder = wp.sim.ModelBuilder()
+
+        cell_dim = 15
+        cell_size = 2.0 / cell_dim
+
+        center = cell_size * cell_dim * 0.5
+
+        builder.add_soft_grid(
+            pos=wp.vec3(-center, 0.0, -center),
+            rot=wp.quat_identity(),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            dim_x=cell_dim,
+            dim_y=cell_dim,
+            dim_z=cell_dim,
+            cell_x=cell_size,
+            cell_y=cell_size,
+            cell_z=cell_size,
+            density=100.0,
+            fix_bottom=True,
+            fix_top=True,
+            k_mu=1000.0,
+            k_lambda=5000.0,
+            k_damp=0.0,
+        )
+
+        self.model = builder.finalize()
+        self.model.ground = False
+        self.model.gravity[1] = 0.0
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+
+        self.rest = self.model.state()
+        self.rest_vol = (cell_size * cell_dim) ** 3
+
+        self.state_0 = self.model.state()
+        self.state_1 = self.model.state()
+
+        self.volume = wp.zeros(1, dtype=wp.float32)
+
+        self.renderer = None
+        if stage:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=20.0)
+
+        self.use_graph = wp.get_device().is_cuda
+        if self.use_graph:
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
+
+    def simulate(self):
+        for _ in range(self.sim_substeps):
+            self.state_0.clear_forces()
+            self.state_1.clear_forces()
+
+            self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
+
+            # swap states
+            (self.state_0, self.state_1) = (self.state_1, self.state_0)
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            xform = wp.transform(
+                (0.0, self.lift_speed * self.sim_time, 0.0),
+                wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), self.rot_speed * self.sim_time),
+            )
+            wp.launch(
+                kernel=twist_points,
+                dim=len(self.state_0.particle_q),
+                inputs=[self.rest.particle_q, self.state_0.particle_q, self.model.particle_mass, xform],
+            )
+            if self.use_graph:
+                wp.capture_launch(self.graph)
+            else:
+                self.simulate()
+            self.volume.zero_()
+            wp.launch(
+                kernel=compute_volume,
+                dim=self.model.tet_count,
+                inputs=[self.state_0.particle_q, self.model.tet_indices, self.volume],
+            )
+        self.sim_time += self.frame_dt
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render"):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render(self.state_0)
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "example_soft_body.usd")
+
+    example = Example(stage_path)
+
+    for i in range(example.sim_frames):
+        example.step()
+        example.render()
+
+    if example.renderer:
+        example.renderer.save()

warp/fabric.py

@@ -108,6 +108,8 @@ class fabricarray(noncontiguous_array_base[T]):
     def __init__(self, data=None, attrib=None, dtype=Any, ndim=None):
         super().__init__(ARRAY_TYPE_FABRIC)
 
+        self.deleter = None
+
         if data is not None:
             from .context import runtime
 
@@ -174,24 +176,29 @@ class fabricarray(noncontiguous_array_base[T]):
 
             num_buckets = len(pointers)
             size = 0
-
             buckets = (fabricbucket_t * num_buckets)()
-            for i in range(num_buckets):
-                buckets[i].index_start = size
-                buckets[i].index_end = size + counts[i]
-                buckets[i].ptr = pointers[i]
-                if array_lengths:
-                    buckets[i].lengths = array_lengths[i]
-                size += counts[i]
-
-            if self.device.is_cuda:
-                # copy bucket info to device
-                with warp.ScopedStream(self.device.null_stream):
+
+            if num_buckets > 0:
+                for i in range(num_buckets):
+                    buckets[i].index_start = size
+                    buckets[i].index_end = size + counts[i]
+                    buckets[i].ptr = pointers[i]
+                    if array_lengths:
+                        buckets[i].lengths = array_lengths[i]
+                    size += counts[i]
+
+                if self.device.is_cuda:
+                    # copy bucket info to device
                     buckets_size = ctypes.sizeof(buckets)
-                    buckets_ptr = self.device.allocator.alloc(buckets_size)
-                    runtime.core.memcpy_h2d(self.device.context, buckets_ptr, ctypes.addressof(buckets), buckets_size)
+                    allocator = self.device.get_allocator()
+                    buckets_ptr = allocator.alloc(buckets_size)
+                    cuda_stream = self.device.stream.cuda_stream
+                    runtime.core.memcpy_h2d(self.device.context, buckets_ptr, ctypes.addressof(buckets), buckets_size, cuda_stream)
+                    self.deleter = allocator.deleter
+                else:
+                    buckets_ptr = ctypes.addressof(buckets)
             else:
-                buckets_ptr = ctypes.addressof(buckets)
+                buckets_ptr = None
 
             self.buckets = buckets
             self.size = size
@@ -211,11 +218,13 @@ class fabricarray(noncontiguous_array_base[T]):
             self.ctype = fabricarray_t()
 
     def __del__(self):
-        # release the GPU copy of bucket info
-        if self.buckets is not None and self.device.is_cuda:
-            buckets_size = ctypes.sizeof(self.buckets)
-            with self.device.context_guard:
-                self.device.allocator.free(self.ctype.buckets, buckets_size)
+        # release the bucket info if needed
+        if self.deleter is None:
+            return
+
+        buckets_size = ctypes.sizeof(self.buckets)
+        with self.device.context_guard:
+            self.deleter(self.ctype.buckets, buckets_size)
 
     def __ctype__(self):
         return self.ctype

warp/fem/cache.py

@@ -223,7 +223,6 @@ class Temporary:
                 pinned=array.pinned,
                 capacity=array.capacity,
                 copy=False,
-                owner=False,
                 grad=None if array.grad is None else _view_reshaped_truncated(array.grad),
             )
 

warp/fem/dirichlet.py

@@ -23,7 +23,6 @@ def normalize_dirichlet_projector(projector_matrix: BsrMatrix, fixed_value: Opti
             data=None,
             ptr=projector_values.ptr,
             capacity=projector_values.capacity,
-            owner=False,
             device=projector_values.device,
             dtype=wp.mat(shape=projector_matrix.block_shape, dtype=projector_matrix.scalar_type),
             shape=projector_values.shape[0],
@@ -47,7 +46,6 @@ def normalize_dirichlet_projector(projector_matrix: BsrMatrix, fixed_value: Opti
                 data=None,
                 ptr=fixed_value.ptr,
                 capacity=fixed_value.capacity,
-                owner=False,
                 device=fixed_value.device,
                 dtype=wp.vec(length=projector_matrix.block_shape[0], dtype=projector_matrix.scalar_type),
                 shape=fixed_value.shape[0],

warp/fem/integrate.py

@@ -969,7 +969,6 @@ def _launch_integrate_kernel(
                 data=None,
                 ptr=array.ptr,
                 capacity=array.capacity,
-                owner=False,
                 device=array.device,
                 shape=(test.space_partition.node_count(), test.space.VALUE_DOF_COUNT),
                 dtype=wp.types.type_scalar_type(output_dtype),

warp/jax.py

@@ -6,6 +6,7 @@
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
 import warp
+from warp.context import type_str
 
 
 def device_to_jax(wp_device):
@@ -34,6 +35,50 @@ def device_from_jax(jax_device):
         raise RuntimeError(f"Unknown or unsupported Jax device platform '{jax_device.platform}'")
 
 
+def dtype_to_jax(wp_dtype):
+    import jax.numpy as jp
+
+    warp_to_jax_dict = {
+        warp.float16: jp.float16,
+        warp.float32: jp.float32,
+        warp.float64: jp.float64,
+        warp.int8: jp.int8,
+        warp.int16: jp.int16,
+        warp.int32: jp.int32,
+        warp.int64: jp.int64,
+        warp.uint8: jp.uint8,
+        warp.uint16: jp.uint16,
+        warp.uint32: jp.uint32,
+        warp.uint64: jp.uint64,
+    }
+    jax_dtype = warp_to_jax_dict.get(wp_dtype)
+    if jax_dtype is None:
+        raise TypeError(f"Invalid or unsupported data type: {type_str(wp_dtype)}")
+    return jax_dtype
+
+
+def dtype_from_jax(jax_dtype):
+    import jax.numpy as jp
+
+    jax_to_warp_dict = {
+        jp.float16: warp.float16,
+        jp.float32: warp.float32,
+        jp.float64: warp.float64,
+        jp.int8: warp.int8,
+        jp.int16: warp.int16,
+        jp.int32: warp.int32,
+        jp.int64: warp.int64,
+        jp.uint8: warp.uint8,
+        jp.uint16: warp.uint16,
+        jp.uint32: warp.uint32,
+        jp.uint64: warp.uint64,
+    }
+    wp_dtype = jax_to_warp_dict.get(jax_dtype)
+    if wp_dtype is None:
+        raise TypeError(f"Invalid or unsupported data type: {jax_dtype}")
+    return wp_dtype
+
+
 def to_jax(wp_array):
     import jax.dlpack