Source code for jaxsim.parsers.descriptions.collision

from __future__ import annotations

import abc
import dataclasses

import numpy as np
import numpy.typing as npt

from jaxsim import logging

from .link import LinkDescription




@dataclasses.dataclass
class CollidablePoint:
    """
    Represents a collidable point associated with a parent link.

    Attributes:
        parent_link: The parent link to which the collidable point is attached.
        position: The position of the collidable point relative to the parent link.
        enabled: A flag indicating whether the collidable point is enabled for collision detection.
    """

    parent_link: LinkDescription
    enabled: bool = True
    _position: tuple[float] = dataclasses.field(default=(0.0, 0.0, 0.0))

    @property
    def position(self) -> npt.NDArray:
        """
        Get the position of the collidable point.

        Returns:
            The position of the collidable point.
        """
        return np.array(self._position)

    @position.setter
    def position(self, value: npt.NDArray) -> None:
        """
        Set the position of the collidable point.

        Args:
            value: The new position of the collidable point.
        """
        self._position = tuple(value.tolist())



    def change_link(
        self, new_link: LinkDescription, new_H_old: npt.NDArray
    ) -> CollidablePoint:
        """
        Move the collidable point to a new parent link.

        Args:
            new_link: The new parent link to which the collidable point is moved.
            new_H_old: The transformation matrix from the new link's frame to the old link's frame.

        Returns:
            CollidablePoint: A new collidable point associated with the new parent link.
        """

        msg = f"Moving collidable point: {self.parent_link.name} -> {new_link.name}"
        logging.debug(msg=msg)

        return CollidablePoint(
            parent_link=new_link,
            _position=tuple(
                (new_H_old @ np.hstack([self.position, 1.0])).squeeze()[0:3].tolist()
            ),
            enabled=self.enabled,
        )


    def __str__(self) -> str:
        return (
            f"{self.__class__.__name__}("
            + f"parent_link={self.parent_link.name}"
            + f", position={self.position}"
            + f", enabled={self.enabled}"
            + ")"
        )





@dataclasses.dataclass
class CollisionShape(abc.ABC):
    """
    Abstract base class for representing collision shapes.

    Attributes:
        collidable_points: A list of collidable points associated with the collision shape.
    """

    collidable_points: tuple[CollidablePoint]

    def __str__(self):
        return (
            f"{self.__class__.__name__}("
            + "collidable_points=[\n    "
            + ",\n    ".join(str(cp) for cp in self.collidable_points)
            + "\n])"
        )





@dataclasses.dataclass
class BoxCollision(CollisionShape):
    """
    Represents a box-shaped collision shape.

    Attributes:
        center: The center of the box in the local frame of the collision shape.
    """

    center: tuple[float, float, float]





@dataclasses.dataclass
class SphereCollision(CollisionShape):
    """
    Represents a spherical collision shape.

    Attributes:
        center: The center of the sphere in the local frame of the collision shape.
    """

    center: tuple[float, float, float]





@dataclasses.dataclass
class MeshCollision(CollisionShape):
    """
    Represents a mesh-shaped collision shape.

    Attributes:
        center: The center of the mesh in the local frame of the collision shape.
    """

    center: tuple[float, float, float]