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

warp/sim/integrator.py

@@ -0,0 +1,233 @@
+# 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()