warp-lang 1.0.2__py3-none-manylinux2014_x86_64.whl → 1.2.0__py3-none-manylinux2014_x86_64.whl

warp/sim/integrator.py

@@ -1,233 +1,234 @@
-# 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.
-
-import warp as wp
-from .model import Model, State, Control, PARTICLE_FLAG_ACTIVE
-
-
-@wp.kernel
-def integrate_particles(
-    x: wp.array(dtype=wp.vec3),
-    v: wp.array(dtype=wp.vec3),
-    f: wp.array(dtype=wp.vec3),
-    w: wp.array(dtype=float),
-    particle_flags: wp.array(dtype=wp.uint32),
-    gravity: wp.vec3,
-    dt: float,
-    v_max: float,
-    x_new: wp.array(dtype=wp.vec3),
-    v_new: wp.array(dtype=wp.vec3),
-):
-    tid = wp.tid()
-    if (particle_flags[tid] & PARTICLE_FLAG_ACTIVE) == 0:
-        return
-
-    x0 = x[tid]
-    v0 = v[tid]
-    f0 = f[tid]
-
-    inv_mass = w[tid]
-
-    # simple semi-implicit Euler. v1 = v0 + a dt, x1 = x0 + v1 dt
-    v1 = v0 + (f0 * inv_mass + gravity * wp.step(-inv_mass)) * dt
-    # enforce velocity limit to prevent instability
-    v1_mag = wp.length(v1)
-    if v1_mag > v_max:
-        v1 *= v_max / v1_mag
-    x1 = x0 + v1 * dt
-
-    x_new[tid] = x1
-    v_new[tid] = v1
-
-
-@wp.func
-def integrate_rigid_body(
-    q: wp.transform,
-    qd: wp.spatial_vector,
-    f: wp.spatial_vector,
-    com: wp.vec3,
-    inertia: wp.mat33,
-    inv_mass: float,
-    inv_inertia: wp.mat33,
-    gravity: wp.vec3,
-    angular_damping: float,
-    dt: float,
-):
-    # unpack transform
-    x0 = wp.transform_get_translation(q)
-    r0 = wp.transform_get_rotation(q)
-
-    # unpack spatial twist
-    w0 = wp.spatial_top(qd)
-    v0 = wp.spatial_bottom(qd)
-
-    # unpack spatial wrench
-    t0 = wp.spatial_top(f)
-    f0 = wp.spatial_bottom(f)
-
-    x_com = x0 + wp.quat_rotate(r0, com)
-
-    # linear part
-    v1 = v0 + (f0 * inv_mass + gravity * wp.nonzero(inv_mass)) * dt
-    x1 = x_com + v1 * dt
-
-    # angular part (compute in body frame)
-    wb = wp.quat_rotate_inv(r0, w0)
-    tb = wp.quat_rotate_inv(r0, t0) - wp.cross(wb, inertia * wb)  # coriolis forces
-
-    w1 = wp.quat_rotate(r0, wb + inv_inertia * tb * dt)
-    r1 = wp.normalize(r0 + wp.quat(w1, 0.0) * r0 * 0.5 * dt)
-
-    # angular damping
-    w1 *= 1.0 - angular_damping * dt
-
-    q_new = wp.transform(x1 - wp.quat_rotate(r1, com), r1)
-    qd_new = wp.spatial_vector(w1, v1)
-
-    return q_new, qd_new
-
-
-# semi-implicit Euler integration
-@wp.kernel
-def integrate_bodies(
-    body_q: wp.array(dtype=wp.transform),
-    body_qd: wp.array(dtype=wp.spatial_vector),
-    body_f: wp.array(dtype=wp.spatial_vector),
-    body_com: wp.array(dtype=wp.vec3),
-    m: wp.array(dtype=float),
-    I: wp.array(dtype=wp.mat33),
-    inv_m: wp.array(dtype=float),
-    inv_I: wp.array(dtype=wp.mat33),
-    gravity: wp.vec3,
-    angular_damping: float,
-    dt: float,
-    # outputs
-    body_q_new: wp.array(dtype=wp.transform),
-    body_qd_new: wp.array(dtype=wp.spatial_vector),
-):
-    tid = wp.tid()
-
-    # positions
-    q = body_q[tid]
-    qd = body_qd[tid]
-    f = body_f[tid]
-
-    # masses
-    inv_mass = inv_m[tid]  # 1 / mass
-
-    inertia = I[tid]
-    inv_inertia = inv_I[tid]  # inverse of 3x3 inertia matrix
-
-    com = body_com[tid]
-
-    q_new, qd_new = integrate_rigid_body(
-        q,
-        qd,
-        f,
-        com,
-        inertia,
-        inv_mass,
-        inv_inertia,
-        gravity,
-        angular_damping,
-        dt,
-    )
-
-    body_q_new[tid] = q_new
-    body_qd_new[tid] = qd_new
-
-
-class Integrator:
-    """
-    Generic base class for integrators. Provides methods to integrate rigid bodies and particles.
-    """
-
-    def integrate_bodies(
-        self,
-        model: Model,
-        state_in: State,
-        state_out: State,
-        dt: float,
-        angular_damping: float = 0.0,
-    ):
-        """
-        Integrate the rigid bodies of the model.
-
-        Args:
-            model (Model): The model to integrate.
-            state_in (State): The input state.
-            state_out (State): The output state.
-            dt (float): The time step (typically in seconds).
-            angular_damping (float, optional): The angular damping factor. Defaults to 0.0.
-        """
-        if model.body_count:
-            wp.launch(
-                kernel=integrate_bodies,
-                dim=model.body_count,
-                inputs=[
-                    state_in.body_q,
-                    state_in.body_qd,
-                    state_in.body_f,
-                    model.body_com,
-                    model.body_mass,
-                    model.body_inertia,
-                    model.body_inv_mass,
-                    model.body_inv_inertia,
-                    model.gravity,
-                    angular_damping,
-                    dt,
-                ],
-                outputs=[state_out.body_q, state_out.body_qd],
-                device=model.device,
-            )
-
-    def integrate_particles(
-        self,
-        model: Model,
-        state_in: State,
-        state_out: State,
-        dt: float,
-    ):
-        """
-        Integrate the particles of the model.
-
-        Args:
-            model (Model): The model to integrate.
-            state_in (State): The input state.
-            state_out (State): The output state.
-            dt (float): The time step (typically in seconds).
-        """
-        if model.particle_count:
-            wp.launch(
-                kernel=integrate_particles,
-                dim=model.particle_count,
-                inputs=[
-                    state_in.particle_q,
-                    state_in.particle_qd,
-                    state_in.particle_f,
-                    model.particle_inv_mass,
-                    model.particle_flags,
-                    model.gravity,
-                    dt,
-                    model.particle_max_velocity,
-                ],
-                outputs=[state_out.particle_q, state_out.particle_qd],
-                device=model.device,
-            )
-
-    def simulate(self, model: Model, state_in: State, state_out: State, dt: float, control: Control = None):
-        """
-        Simulate the model for a given time step using the given control input.
-
-        Args:
-            model (Model): The model to simulate.
-            state_in (State): The input state.
-            state_out (State): The output state.
-            dt (float): The time step (typically in seconds).
-            control (Control): The control input. Defaults to `None` which means the control values from the :class:`Model` are used.
-        """
-        raise NotImplementedError()
+# 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.
+
+import warp as wp
+
+from .model import PARTICLE_FLAG_ACTIVE, Control, Model, State
+
+
+@wp.kernel
+def integrate_particles(
+    x: wp.array(dtype=wp.vec3),
+    v: wp.array(dtype=wp.vec3),
+    f: wp.array(dtype=wp.vec3),
+    w: wp.array(dtype=float),
+    particle_flags: wp.array(dtype=wp.uint32),
+    gravity: wp.vec3,
+    dt: float,
+    v_max: float,
+    x_new: wp.array(dtype=wp.vec3),
+    v_new: wp.array(dtype=wp.vec3),
+):
+    tid = wp.tid()
+    if (particle_flags[tid] & PARTICLE_FLAG_ACTIVE) == 0:
+        return
+
+    x0 = x[tid]
+    v0 = v[tid]
+    f0 = f[tid]
+
+    inv_mass = w[tid]
+
+    # simple semi-implicit Euler. v1 = v0 + a dt, x1 = x0 + v1 dt
+    v1 = v0 + (f0 * inv_mass + gravity * wp.step(-inv_mass)) * dt
+    # enforce velocity limit to prevent instability
+    v1_mag = wp.length(v1)
+    if v1_mag > v_max:
+        v1 *= v_max / v1_mag
+    x1 = x0 + v1 * dt
+
+    x_new[tid] = x1
+    v_new[tid] = v1
+
+
+@wp.func
+def integrate_rigid_body(
+    q: wp.transform,
+    qd: wp.spatial_vector,
+    f: wp.spatial_vector,
+    com: wp.vec3,
+    inertia: wp.mat33,
+    inv_mass: float,
+    inv_inertia: wp.mat33,
+    gravity: wp.vec3,
+    angular_damping: float,
+    dt: float,
+):
+    # unpack transform
+    x0 = wp.transform_get_translation(q)
+    r0 = wp.transform_get_rotation(q)
+
+    # unpack spatial twist
+    w0 = wp.spatial_top(qd)
+    v0 = wp.spatial_bottom(qd)
+
+    # unpack spatial wrench
+    t0 = wp.spatial_top(f)
+    f0 = wp.spatial_bottom(f)
+
+    x_com = x0 + wp.quat_rotate(r0, com)
+
+    # linear part
+    v1 = v0 + (f0 * inv_mass + gravity * wp.nonzero(inv_mass)) * dt
+    x1 = x_com + v1 * dt
+
+    # angular part (compute in body frame)
+    wb = wp.quat_rotate_inv(r0, w0)
+    tb = wp.quat_rotate_inv(r0, t0) - wp.cross(wb, inertia * wb)  # coriolis forces
+
+    w1 = wp.quat_rotate(r0, wb + inv_inertia * tb * dt)
+    r1 = wp.normalize(r0 + wp.quat(w1, 0.0) * r0 * 0.5 * dt)
+
+    # angular damping
+    w1 *= 1.0 - angular_damping * dt
+
+    q_new = wp.transform(x1 - wp.quat_rotate(r1, com), r1)
+    qd_new = wp.spatial_vector(w1, v1)
+
+    return q_new, qd_new
+
+
+# semi-implicit Euler integration
+@wp.kernel
+def integrate_bodies(
+    body_q: wp.array(dtype=wp.transform),
+    body_qd: wp.array(dtype=wp.spatial_vector),
+    body_f: wp.array(dtype=wp.spatial_vector),
+    body_com: wp.array(dtype=wp.vec3),
+    m: wp.array(dtype=float),
+    I: wp.array(dtype=wp.mat33),
+    inv_m: wp.array(dtype=float),
+    inv_I: wp.array(dtype=wp.mat33),
+    gravity: wp.vec3,
+    angular_damping: float,
+    dt: float,
+    # outputs
+    body_q_new: wp.array(dtype=wp.transform),
+    body_qd_new: wp.array(dtype=wp.spatial_vector),
+):
+    tid = wp.tid()
+
+    # positions
+    q = body_q[tid]
+    qd = body_qd[tid]
+    f = body_f[tid]
+
+    # masses
+    inv_mass = inv_m[tid]  # 1 / mass
+
+    inertia = I[tid]
+    inv_inertia = inv_I[tid]  # inverse of 3x3 inertia matrix
+
+    com = body_com[tid]
+
+    q_new, qd_new = integrate_rigid_body(
+        q,
+        qd,
+        f,
+        com,
+        inertia,
+        inv_mass,
+        inv_inertia,
+        gravity,
+        angular_damping,
+        dt,
+    )
+
+    body_q_new[tid] = q_new
+    body_qd_new[tid] = qd_new
+
+
+class Integrator:
+    """
+    Generic base class for integrators. Provides methods to integrate rigid bodies and particles.
+    """
+
+    def integrate_bodies(
+        self,
+        model: Model,
+        state_in: State,
+        state_out: State,
+        dt: float,
+        angular_damping: float = 0.0,
+    ):
+        """
+        Integrate the rigid bodies of the model.
+
+        Args:
+            model (Model): The model to integrate.
+            state_in (State): The input state.
+            state_out (State): The output state.
+            dt (float): The time step (typically in seconds).
+            angular_damping (float, optional): The angular damping factor. Defaults to 0.0.
+        """
+        if model.body_count:
+            wp.launch(
+                kernel=integrate_bodies,
+                dim=model.body_count,
+                inputs=[
+                    state_in.body_q,
+                    state_in.body_qd,
+                    state_in.body_f,
+                    model.body_com,
+                    model.body_mass,
+                    model.body_inertia,
+                    model.body_inv_mass,
+                    model.body_inv_inertia,
+                    model.gravity,
+                    angular_damping,
+                    dt,
+                ],
+                outputs=[state_out.body_q, state_out.body_qd],
+                device=model.device,
+            )
+
+    def integrate_particles(
+        self,
+        model: Model,
+        state_in: State,
+        state_out: State,
+        dt: float,
+    ):
+        """
+        Integrate the particles of the model.
+
+        Args:
+            model (Model): The model to integrate.
+            state_in (State): The input state.
+            state_out (State): The output state.
+            dt (float): The time step (typically in seconds).
+        """
+        if model.particle_count:
+            wp.launch(
+                kernel=integrate_particles,
+                dim=model.particle_count,
+                inputs=[
+                    state_in.particle_q,
+                    state_in.particle_qd,
+                    state_in.particle_f,
+                    model.particle_inv_mass,
+                    model.particle_flags,
+                    model.gravity,
+                    dt,
+                    model.particle_max_velocity,
+                ],
+                outputs=[state_out.particle_q, state_out.particle_qd],
+                device=model.device,
+            )
+
+    def simulate(self, model: Model, state_in: State, state_out: State, dt: float, control: Control = None):
+        """
+        Simulate the model for a given time step using the given control input.
+
+        Args:
+            model (Model): The model to simulate.
+            state_in (State): The input state.
+            state_out (State): The output state.
+            dt (float): The time step (typically in seconds).
+            control (Control): The control input. Defaults to `None` which means the control values from the :class:`Model` are used.
+        """
+        raise NotImplementedError()