warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/sim/particles.py

@@ -1,121 +0,0 @@
-# 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
-
-
-@wp.func
-def particle_force(n: wp.vec3, v: wp.vec3, c: float, k_n: float, k_d: float, k_f: float, k_mu: float):
-    # compute normal and tangential friction force for a single contact
-    vn = wp.dot(n, v)
-    jn = c * k_n
-    jd = min(vn, 0.0) * k_d
-
-    # contact force
-    fn = jn + jd
-
-    # friction force
-    vt = v - n * vn
-    vs = wp.length(vt)
-
-    if vs > 0.0:
-        vt = vt / vs
-
-    # Coulomb condition
-    ft = wp.min(vs * k_f, k_mu * wp.abs(fn))
-
-    # total force
-    return -n * fn - vt * ft
-
-
-@wp.kernel
-def eval_particle_forces_kernel(
-    grid: wp.uint64,
-    particle_x: wp.array(dtype=wp.vec3),
-    particle_v: wp.array(dtype=wp.vec3),
-    particle_radius: wp.array(dtype=float),
-    particle_flags: wp.array(dtype=wp.uint32),
-    k_contact: float,
-    k_damp: float,
-    k_friction: float,
-    k_mu: float,
-    k_cohesion: float,
-    max_radius: float,
-    # outputs
-    particle_f: wp.array(dtype=wp.vec3),
-):
-    tid = wp.tid()
-
-    # order threads by cell
-    i = wp.hash_grid_point_id(grid, tid)
-    if i == -1:
-        # hash grid has not been built yet
-        return
-    if (particle_flags[i] & PARTICLE_FLAG_ACTIVE) == 0:
-        return
-
-    x = particle_x[i]
-    v = particle_v[i]
-    radius = particle_radius[i]
-
-    f = wp.vec3()
-
-    # particle contact
-    query = wp.hash_grid_query(grid, x, radius + max_radius + k_cohesion)
-    index = int(0)
-
-    count = int(0)
-
-    while wp.hash_grid_query_next(query, index):
-        if (particle_flags[index] & PARTICLE_FLAG_ACTIVE) != 0 and index != i:
-            # compute distance to point
-            n = x - particle_x[index]
-            d = wp.length(n)
-            err = d - radius - particle_radius[index]
-
-            count += 1
-
-            if err <= k_cohesion:
-                n = n / d
-                vrel = v - particle_v[index]
-
-                f = f + particle_force(n, vrel, err, k_contact, k_damp, k_friction, k_mu)
-
-    particle_f[i] = f
-
-
-def eval_particle_forces(model, state, forces):
-    if model.particle_count > 1 and model.particle_max_radius > 0.0:
-        wp.launch(
-            kernel=eval_particle_forces_kernel,
-            dim=model.particle_count,
-            inputs=[
-                model.particle_grid.id,
-                state.particle_q,
-                state.particle_qd,
-                model.particle_radius,
-                model.particle_flags,
-                model.particle_ke,
-                model.particle_kd,
-                model.particle_kf,
-                model.particle_mu,
-                model.particle_cohesion,
-                model.particle_max_radius,
-            ],
-            outputs=[forces],
-            device=model.device,
-        )

warp/sim/render.py

@@ -1,431 +0,0 @@
-# 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.
-
-from collections import defaultdict
-
-import numpy as np
-
-import warp as wp
-import warp.render
-import warp.sim
-from warp.render.utils import solidify_mesh, tab10_color_map
-
-# TODO allow NaNs in Warp kernels
-NAN = wp.constant(-1.0e8)
-
-
-@wp.kernel
-def compute_contact_points(
-    body_q: wp.array(dtype=wp.transform),
-    shape_body: wp.array(dtype=int),
-    contact_count: wp.array(dtype=int),
-    contact_shape0: wp.array(dtype=int),
-    contact_shape1: wp.array(dtype=int),
-    contact_point0: wp.array(dtype=wp.vec3),
-    contact_point1: wp.array(dtype=wp.vec3),
-    # outputs
-    contact_pos0: wp.array(dtype=wp.vec3),
-    contact_pos1: wp.array(dtype=wp.vec3),
-):
-    tid = wp.tid()
-    count = contact_count[0]
-    if tid >= count:
-        contact_pos0[tid] = wp.vec3(NAN, NAN, NAN)
-        contact_pos1[tid] = wp.vec3(NAN, NAN, NAN)
-        return
-    shape_a = contact_shape0[tid]
-    shape_b = contact_shape1[tid]
-    if shape_a == shape_b:
-        contact_pos0[tid] = wp.vec3(NAN, NAN, NAN)
-        contact_pos1[tid] = wp.vec3(NAN, NAN, NAN)
-        return
-
-    body_a = shape_body[shape_a]
-    body_b = shape_body[shape_b]
-    X_wb_a = wp.transform_identity()
-    X_wb_b = wp.transform_identity()
-    if body_a >= 0:
-        X_wb_a = body_q[body_a]
-    if body_b >= 0:
-        X_wb_b = body_q[body_b]
-
-    contact_pos0[tid] = wp.transform_point(X_wb_a, contact_point0[tid])
-    contact_pos1[tid] = wp.transform_point(X_wb_b, contact_point1[tid])
-
-
-def CreateSimRenderer(renderer):
-    class SimRenderer(renderer):
-        use_unique_colors = True
-
-        def __init__(
-            self,
-            model: warp.sim.Model,
-            path,
-            scaling=1.0,
-            fps=60,
-            up_axis=None,
-            show_rigid_contact_points=False,
-            contact_points_radius=1e-3,
-            show_joints=False,
-            **render_kwargs,
-        ):
-            # create USD stage
-            if up_axis is None:
-                up_axis = "XYZ"[model.up_axis]
-            super().__init__(path, scaling=scaling, fps=fps, up_axis=up_axis, **render_kwargs)
-            self.scaling = scaling
-            self.cam_axis = "XYZ".index(up_axis.upper())
-            self.show_rigid_contact_points = show_rigid_contact_points
-            self.show_joints = show_joints
-            self.contact_points_radius = contact_points_radius
-            self.populate(model)
-
-        def populate(self, model: warp.sim.Model):
-            self.skip_rendering = False
-
-            self.model = model
-            self.num_envs = model.num_envs
-            self.body_names = []
-
-            if self.show_rigid_contact_points and model.rigid_contact_max:
-                self.contact_points0 = wp.array(
-                    np.zeros((model.rigid_contact_max, 3)), dtype=wp.vec3, device=model.device
-                )
-                self.contact_points1 = wp.array(
-                    np.zeros((model.rigid_contact_max, 3)), dtype=wp.vec3, device=model.device
-                )
-
-                self.contact_points0_colors = [(1.0, 0.5, 0.0)] * model.rigid_contact_max
-                self.contact_points1_colors = [(0.0, 0.5, 1.0)] * model.rigid_contact_max
-
-            self.body_env = []  # mapping from body index to its environment index
-            env_id = 0
-            self.bodies_per_env = model.body_count // self.num_envs
-            # create rigid body nodes
-            for b in range(model.body_count):
-                body_name = f"body_{b}_{self.model.body_name[b].replace(' ', '_')}"
-                self.body_names.append(body_name)
-                self.register_body(body_name)
-                if b > 0 and b % self.bodies_per_env == 0:
-                    env_id += 1
-                self.body_env.append(env_id)
-
-            # create rigid shape children
-            if self.model.shape_count:
-                # mapping from hash of geometry to shape ID
-                self.geo_shape = {}
-
-                self.instance_count = 0
-
-                self.body_name = {}  # mapping from body name to its body ID
-                self.body_shapes = defaultdict(list)  # mapping from body index to its shape IDs
-
-                shape_body = model.shape_body.numpy()
-                shape_geo_src = model.shape_geo_src
-                shape_geo_type = model.shape_geo.type.numpy()
-                shape_geo_scale = model.shape_geo.scale.numpy()
-                shape_geo_thickness = model.shape_geo.thickness.numpy()
-                shape_geo_is_solid = model.shape_geo.is_solid.numpy()
-                shape_transform = model.shape_transform.numpy()
-                shape_visible = model.shape_visible.numpy()
-
-                p = np.zeros(3, dtype=np.float32)
-                q = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32)
-                scale = np.ones(3)
-                color = (1.0, 1.0, 1.0)
-                # loop over shapes excluding the ground plane
-                for s in range(model.shape_count - 1):
-                    geo_type = shape_geo_type[s]
-                    geo_scale = [float(v) for v in shape_geo_scale[s]]
-                    geo_thickness = float(shape_geo_thickness[s])
-                    geo_is_solid = bool(shape_geo_is_solid[s])
-                    geo_src = shape_geo_src[s]
-                    name = f"shape_{s}"
-
-                    # shape transform in body frame
-                    body = int(shape_body[s])
-                    if body >= 0 and body < len(self.body_names):
-                        body = self.body_names[body]
-                    else:
-                        body = None
-
-                    if self.use_unique_colors and body is not None:
-                        color = self._get_new_color()
-
-                    # shape transform in body frame
-                    X_bs = wp.transform_expand(shape_transform[s])
-                    # check whether we can instance an already created shape with the same geometry
-                    geo_hash = hash((int(geo_type), geo_src, *geo_scale, geo_thickness, geo_is_solid))
-                    if geo_hash in self.geo_shape:
-                        shape = self.geo_shape[geo_hash]
-                    else:
-                        if geo_type == warp.sim.GEO_PLANE:
-                            if s == model.shape_count - 1 and not model.ground:
-                                continue  # hide ground plane
-
-                            # plane mesh
-                            width = geo_scale[0] if geo_scale[0] > 0.0 else 100.0
-                            length = geo_scale[1] if geo_scale[1] > 0.0 else 100.0
-
-                            shape = self.render_plane(
-                                name, p, q, width, length, color, parent_body=body, is_template=True
-                            )
-
-                        elif geo_type == warp.sim.GEO_SPHERE:
-                            shape = self.render_sphere(
-                                name, p, q, geo_scale[0], parent_body=body, is_template=True, color=color
-                            )
-
-                        elif geo_type == warp.sim.GEO_CAPSULE:
-                            shape = self.render_capsule(
-                                name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True, color=color
-                            )
-
-                        elif geo_type == warp.sim.GEO_CYLINDER:
-                            shape = self.render_cylinder(
-                                name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True, color=color
-                            )
-
-                        elif geo_type == warp.sim.GEO_CONE:
-                            shape = self.render_cone(
-                                name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True, color=color
-                            )
-
-                        elif geo_type == warp.sim.GEO_BOX:
-                            shape = self.render_box(
-                                name, p, q, geo_scale, parent_body=body, is_template=True, color=color
-                            )
-
-                        elif geo_type == warp.sim.GEO_MESH:
-                            if not geo_is_solid:
-                                faces, vertices = solidify_mesh(geo_src.indices, geo_src.vertices, geo_thickness)
-                            else:
-                                faces, vertices = geo_src.indices, geo_src.vertices
-
-                            shape = self.render_mesh(
-                                name,
-                                vertices,
-                                faces,
-                                pos=p,
-                                rot=q,
-                                scale=geo_scale,
-                                colors=color,
-                                parent_body=body,
-                                is_template=True,
-                            )
-
-                        elif geo_type == warp.sim.GEO_SDF:
-                            continue
-
-                        self.geo_shape[geo_hash] = shape
-
-                    if shape_visible[s]:
-                        # TODO support dynamic visibility
-                        self.add_shape_instance(
-                            name, shape, body, X_bs.p, X_bs.q, scale, custom_index=s, visible=shape_visible[s]
-                        )
-                    self.instance_count += 1
-
-                if self.show_joints and model.joint_count:
-                    joint_type = model.joint_type.numpy()
-                    joint_axis = model.joint_axis.numpy()
-                    joint_axis_start = model.joint_axis_start.numpy()
-                    joint_axis_dim = model.joint_axis_dim.numpy()
-                    joint_parent = model.joint_parent.numpy()
-                    joint_child = model.joint_child.numpy()
-                    joint_tf = model.joint_X_p.numpy()
-                    shape_collision_radius = model.shape_collision_radius.numpy()
-                    y_axis = wp.vec3(0.0, 1.0, 0.0)
-                    color = (1.0, 0.0, 1.0)
-
-                    shape = self.render_arrow(
-                        "joint_arrow",
-                        None,
-                        None,
-                        base_radius=0.01,
-                        base_height=0.4,
-                        cap_radius=0.02,
-                        cap_height=0.1,
-                        parent_body=None,
-                        is_template=True,
-                        color=color,
-                    )
-                    for i, t in enumerate(joint_type):
-                        if t not in {
-                            warp.sim.JOINT_REVOLUTE,
-                            # warp.sim.JOINT_PRISMATIC,
-                            warp.sim.JOINT_UNIVERSAL,
-                            warp.sim.JOINT_COMPOUND,
-                            warp.sim.JOINT_D6,
-                        }:
-                            continue
-                        tf = joint_tf[i]
-                        body = int(joint_parent[i])
-                        if body >= 0 and body < len(self.body_names):
-                            body = self.body_names[body]
-                        else:
-                            body = None
-                        # if body == -1:
-                        #     continue
-                        num_linear_axes = int(joint_axis_dim[i][0])
-                        num_angular_axes = int(joint_axis_dim[i][1])
-
-                        # find a good scale for the arrow based on the average radius
-                        # of the shapes attached to the joint child body
-                        scale = np.ones(3)
-                        child = int(joint_child[i])
-                        if child >= 0:
-                            radii = []
-                            for s in model.body_shapes[child]:
-                                radii.append(shape_collision_radius[s])
-                            if len(radii) > 0:
-                                scale *= np.mean(radii) * 2.0
-
-                        for a in range(num_linear_axes, num_linear_axes + num_angular_axes):
-                            index = joint_axis_start[i] + a
-                            axis = joint_axis[index]
-                            if np.linalg.norm(axis) < 1e-6:
-                                continue
-                            p = wp.vec3(tf[:3])
-                            q = wp.quat(tf[3:])
-                            # compute rotation between axis and y
-                            axis = axis / np.linalg.norm(axis)
-                            q = q * wp.quat_between_vectors(wp.vec3(axis), y_axis)
-                            name = f"joint_{i}_{a}"
-                            self.add_shape_instance(name, shape, body, p, q, scale, color1=color, color2=color)
-                            self.instance_count += 1
-
-            if model.ground:
-                self.render_ground(plane=model.ground_plane_params)
-
-            if hasattr(self, "complete_setup"):
-                self.complete_setup()
-
-        def _get_new_color(self):
-            return tab10_color_map(self.instance_count)
-
-        def render(self, state: warp.sim.State):
-            """
-            Updates the renderer with the given simulation state.
-
-            Args:
-                state (warp.sim.State): The simulation state to render.
-            """
-            if self.skip_rendering:
-                return
-
-            if self.model.particle_count:
-                particle_q = state.particle_q.numpy()
-
-                # render particles
-                self.render_points(
-                    "particles", particle_q, radius=self.model.particle_radius.numpy(), colors=(0.8, 0.3, 0.2)
-                )
-
-                # render tris
-                if self.model.tri_count:
-                    self.render_mesh(
-                        "surface",
-                        particle_q,
-                        self.model.tri_indices.numpy().flatten(),
-                        colors=(0.75, 0.25, 0.0),
-                    )
-
-                # render springs
-                if self.model.spring_count:
-                    self.render_line_list(
-                        "springs", particle_q, self.model.spring_indices.numpy().flatten(), (0.25, 0.5, 0.25), 0.02
-                    )
-
-            # render muscles
-            if self.model.muscle_count:
-                body_q = state.body_q.numpy()
-
-                muscle_start = self.model.muscle_start.numpy()
-                muscle_links = self.model.muscle_bodies.numpy()
-                muscle_points = self.model.muscle_points.numpy()
-                muscle_activation = self.model.muscle_activation.numpy()
-
-                # for s in self.skeletons:
-
-                #     # for mesh, link in s.mesh_map.items():
-
-                #     #     if link != -1:
-                #     #         X_sc = wp.transform_expand(self.state.body_X_sc[link].tolist())
-
-                #     #         #self.renderer.add_mesh(mesh, "../assets/snu/OBJ/" + mesh + ".usd", X_sc, 1.0, self.render_time)
-                #     #         self.renderer.add_mesh(mesh, "../assets/snu/OBJ/" + mesh + ".usd", X_sc, 1.0, self.render_time)
-
-                for m in range(self.model.muscle_count):
-                    start = int(muscle_start[m])
-                    end = int(muscle_start[m + 1])
-
-                    points = []
-
-                    for w in range(start, end):
-                        link = muscle_links[w]
-                        point = muscle_points[w]
-
-                        X_sc = wp.transform_expand(body_q[link][0])
-
-                        points.append(wp.transform_point(X_sc, point).tolist())
-
-                    self.render_line_strip(
-                        name=f"muscle_{m}", vertices=points, radius=0.0075, color=(muscle_activation[m], 0.2, 0.5)
-                    )
-
-            # update bodies
-            if self.model.body_count:
-                self.update_body_transforms(state.body_q)
-
-                if self.show_rigid_contact_points and self.model.rigid_contact_max:
-                    wp.launch(
-                        kernel=compute_contact_points,
-                        dim=self.model.rigid_contact_max,
-                        inputs=[
-                            state.body_q,
-                            self.model.shape_body,
-                            self.model.rigid_contact_count,
-                            self.model.rigid_contact_shape0,
-                            self.model.rigid_contact_shape1,
-                            self.model.rigid_contact_point0,
-                            self.model.rigid_contact_point1,
-                        ],
-                        outputs=[
-                            self.contact_points0,
-                            self.contact_points1,
-                        ],
-                        device=self.model.device,
-                    )
-
-                    self.render_points(
-                        "contact_points0",
-                        self.contact_points0.numpy(),
-                        radius=self.contact_points_radius * self.scaling,
-                        colors=self.contact_points0_colors,
-                    )
-                    self.render_points(
-                        "contact_points1",
-                        self.contact_points1.numpy(),
-                        radius=self.contact_points_radius * self.scaling,
-                        colors=self.contact_points1_colors,
-                    )
-
-    return SimRenderer
-
-
-SimRendererUsd = CreateSimRenderer(wp.render.UsdRenderer)
-SimRendererOpenGL = CreateSimRenderer(wp.render.OpenGLRenderer)
-SimRenderer = SimRendererUsd