warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/sim/inertia.py

@@ -0,0 +1,324 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Helper functions for computing rigid body inertia properties."""
+
+import math
+from typing import List, Union
+
+import numpy as np
+
+import warp as wp
+
+
+@wp.func
+def triangle_inertia(
+    p: wp.vec3,
+    q: wp.vec3,
+    r: wp.vec3,
+    density: float,
+    com: wp.vec3,
+    # outputs
+    mass: wp.array(dtype=float, ndim=1),
+    inertia: wp.array(dtype=wp.mat33, ndim=1),
+):
+    pcom = p - com
+    qcom = q - com
+    rcom = r - com
+
+    Dm = wp.mat33(pcom[0], qcom[0], rcom[0], pcom[1], qcom[1], rcom[1], pcom[2], qcom[2], rcom[2])
+
+    volume = wp.abs(wp.determinant(Dm) / 6.0)
+
+    # accumulate mass
+    wp.atomic_add(mass, 0, 4.0 * density * volume)
+
+    alpha = wp.sqrt(5.0) / 5.0
+    mid = (com + p + q + r) / 4.0
+    off_mid = mid - com
+
+    # displacement of quadrature point from COM
+    d0 = alpha * (p - mid) + off_mid
+    d1 = alpha * (q - mid) + off_mid
+    d2 = alpha * (r - mid) + off_mid
+    d3 = alpha * (com - mid) + off_mid
+
+    # accumulate inertia
+    identity = wp.mat33(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0)
+    I = wp.dot(d0, d0) * identity - wp.outer(d0, d0)
+    I += wp.dot(d1, d1) * identity - wp.outer(d1, d1)
+    I += wp.dot(d2, d2) * identity - wp.outer(d2, d2)
+    I += wp.dot(d3, d3) * identity - wp.outer(d3, d3)
+
+    wp.atomic_add(inertia, 0, (density * volume) * I)
+
+    return volume
+
+
+@wp.kernel
+def compute_solid_mesh_inertia(
+    # inputs
+    com: wp.vec3,
+    weight: float,
+    indices: wp.array(dtype=int, ndim=1),
+    vertices: wp.array(dtype=wp.vec3, ndim=1),
+    # outputs
+    mass: wp.array(dtype=float, ndim=1),
+    inertia: wp.array(dtype=wp.mat33, ndim=1),
+    volume: wp.array(dtype=float, ndim=1),
+):
+    i = wp.tid()
+
+    p = vertices[indices[i * 3 + 0]]
+    q = vertices[indices[i * 3 + 1]]
+    r = vertices[indices[i * 3 + 2]]
+
+    vol = triangle_inertia(p, q, r, weight, com, mass, inertia)
+    wp.atomic_add(volume, 0, vol)
+
+
+@wp.kernel
+def compute_hollow_mesh_inertia(
+    # inputs
+    com: wp.vec3,
+    density: float,
+    indices: wp.array(dtype=int, ndim=1),
+    vertices: wp.array(dtype=wp.vec3, ndim=1),
+    thickness: wp.array(dtype=float, ndim=1),
+    # outputs
+    mass: wp.array(dtype=float, ndim=1),
+    inertia: wp.array(dtype=wp.mat33, ndim=1),
+    volume: wp.array(dtype=float, ndim=1),
+):
+    tid = wp.tid()
+    i = indices[tid * 3 + 0]
+    j = indices[tid * 3 + 1]
+    k = indices[tid * 3 + 2]
+
+    vi = vertices[i]
+    vj = vertices[j]
+    vk = vertices[k]
+
+    normal = -wp.normalize(wp.cross(vj - vi, vk - vi))
+    ti = normal * thickness[i]
+    tj = normal * thickness[j]
+    tk = normal * thickness[k]
+
+    # wedge vertices
+    vi0 = vi - ti
+    vi1 = vi + ti
+    vj0 = vj - tj
+    vj1 = vj + tj
+    vk0 = vk - tk
+    vk1 = vk + tk
+
+    triangle_inertia(vi0, vj0, vk0, density, com, mass, inertia)
+    triangle_inertia(vj0, vk1, vk0, density, com, mass, inertia)
+    triangle_inertia(vj0, vj1, vk1, density, com, mass, inertia)
+    triangle_inertia(vj0, vi1, vj1, density, com, mass, inertia)
+    triangle_inertia(vj0, vi0, vi1, density, com, mass, inertia)
+    triangle_inertia(vj1, vi1, vk1, density, com, mass, inertia)
+    triangle_inertia(vi1, vi0, vk0, density, com, mass, inertia)
+    triangle_inertia(vi1, vk0, vk1, density, com, mass, inertia)
+
+    # compute volume
+    a = wp.length(wp.cross(vj - vi, vk - vi)) * 0.5
+    vol = a * (thickness[i] + thickness[j] + thickness[k]) / 3.0
+    wp.atomic_add(volume, 0, vol)
+
+
+def compute_sphere_inertia(density: float, r: float) -> tuple:
+    """Helper to compute mass and inertia of a solid sphere
+
+    Args:
+        density: The sphere density
+        r: The sphere radius
+
+    Returns:
+
+        A tuple of (mass, inertia) with inertia specified around the origin
+    """
+
+    v = 4.0 / 3.0 * math.pi * r * r * r
+
+    m = density * v
+    Ia = 2.0 / 5.0 * m * r * r
+
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ia, 0.0], [0.0, 0.0, Ia]])
+
+    return (m, wp.vec3(), I)
+
+
+def compute_capsule_inertia(density: float, r: float, h: float) -> tuple:
+    """Helper to compute mass and inertia of a solid capsule extending along the y-axis
+
+    Args:
+        density: The capsule density
+        r: The capsule radius
+        h: The capsule height (full height of the interior cylinder)
+
+    Returns:
+
+        A tuple of (mass, inertia) with inertia specified around the origin
+    """
+
+    ms = density * (4.0 / 3.0) * math.pi * r * r * r
+    mc = density * math.pi * r * r * h
+
+    # total mass
+    m = ms + mc
+
+    # adapted from ODE
+    Ia = mc * (0.25 * r * r + (1.0 / 12.0) * h * h) + ms * (0.4 * r * r + 0.375 * r * h + 0.25 * h * h)
+    Ib = (mc * 0.5 + ms * 0.4) * r * r
+
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
+
+    return (m, wp.vec3(), I)
+
+
+def compute_cylinder_inertia(density: float, r: float, h: float) -> tuple:
+    """Helper to compute mass and inertia of a solid cylinder extending along the y-axis
+
+    Args:
+        density: The cylinder density
+        r: The cylinder radius
+        h: The cylinder height (extent along the y-axis)
+
+    Returns:
+
+        A tuple of (mass, inertia) with inertia specified around the origin
+    """
+
+    m = density * math.pi * r * r * h
+
+    Ia = 1 / 12 * m * (3 * r * r + h * h)
+    Ib = 1 / 2 * m * r * r
+
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
+
+    return (m, wp.vec3(), I)
+
+
+def compute_cone_inertia(density: float, r: float, h: float) -> tuple:
+    """Helper to compute mass and inertia of a solid cone extending along the y-axis
+
+    Args:
+        density: The cone density
+        r: The cone radius
+        h: The cone height (extent along the y-axis)
+
+    Returns:
+
+        A tuple of (mass, inertia) with inertia specified around the origin
+    """
+
+    m = density * math.pi * r * r * h / 3.0
+
+    Ia = 1 / 20 * m * (3 * r * r + 2 * h * h)
+    Ib = 3 / 10 * m * r * r
+
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ia]])
+
+    return (m, wp.vec3(), I)
+
+
+def compute_box_inertia(density: float, w: float, h: float, d: float) -> tuple:
+    """Helper to compute mass and inertia of a solid box
+
+    Args:
+        density: The box density
+        w: The box width along the x-axis
+        h: The box height along the y-axis
+        d: The box depth along the z-axis
+
+    Returns:
+
+        A tuple of (mass, inertia) with inertia specified around the origin
+    """
+
+    v = w * h * d
+    m = density * v
+
+    Ia = 1.0 / 12.0 * m * (h * h + d * d)
+    Ib = 1.0 / 12.0 * m * (w * w + d * d)
+    Ic = 1.0 / 12.0 * m * (w * w + h * h)
+
+    I = wp.mat33([[Ia, 0.0, 0.0], [0.0, Ib, 0.0], [0.0, 0.0, Ic]])
+
+    return (m, wp.vec3(), I)
+
+
+def compute_mesh_inertia(
+    density: float, vertices: list, indices: list, is_solid: bool = True, thickness: Union[List[float], float] = 0.001
+) -> tuple:
+    """Computes mass, center of mass, 3x3 inertia matrix, and volume for a mesh."""
+    com = wp.vec3(np.mean(vertices, 0))
+
+    indices = np.array(indices).flatten()
+    num_tris = len(indices) // 3
+
+    # compute signed inertia for each tetrahedron
+    # formed with the interior point, using an order-2
+    # quadrature: https://www.sciencedirect.com/science/article/pii/S0377042712001604#br000040
+
+    # Allocating for mass and inertia
+    I_warp = wp.zeros(1, dtype=wp.mat33)
+    mass_warp = wp.zeros(1, dtype=float)
+    vol_warp = wp.zeros(1, dtype=float)
+
+    if is_solid:
+        weight = 0.25
+        # alpha = math.sqrt(5.0) / 5.0
+        wp.launch(
+            kernel=compute_solid_mesh_inertia,
+            dim=num_tris,
+            inputs=[
+                com,
+                weight,
+                wp.array(indices, dtype=int),
+                wp.array(vertices, dtype=wp.vec3),
+            ],
+            outputs=[mass_warp, I_warp, vol_warp],
+        )
+    else:
+        weight = 0.25 * density
+        if isinstance(thickness, float):
+            thickness = [thickness] * len(vertices)
+        wp.launch(
+            kernel=compute_hollow_mesh_inertia,
+            dim=num_tris,
+            inputs=[
+                com,
+                weight,
+                wp.array(indices, dtype=int),
+                wp.array(vertices, dtype=wp.vec3),
+                wp.array(thickness, dtype=float),
+            ],
+            outputs=[mass_warp, I_warp, vol_warp],
+        )
+
+    # Extract mass and inertia and save to class attributes.
+    mass = float(mass_warp.numpy()[0] * density)
+    I = wp.mat33(*(I_warp.numpy()[0] * density))
+    volume = float(vol_warp.numpy()[0])
+    return mass, com, I, volume
+
+
+def transform_inertia(m, I, p, q):
+    R = wp.quat_to_matrix(q)
+
+    # Steiner's theorem
+    return R @ I @ wp.transpose(R) + m * (wp.dot(p, p) * wp.mat33(np.eye(3)) - wp.outer(p, p))

warp/sim/integrator.py

@@ -0,0 +1,242 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+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()