Source code for diffsim.collision

"""Approximate collision penalties for the explicit solver.

This inexpensive randomized vertex-pair model is useful for previews only. It has
no continuous collision detection or non-penetration guarantee. Frictionless IPC
lives in :mod:`diffsim.ipc` because it must be coupled to an implicit nonlinear solve.
"""

import torch




class SimplifiedCollisionHandler:
    """
    Fast simplified collision detection for real-time simulation

    Uses distance-based repulsion without full CCD
    """



    def __init__(
        self,
        collision_distance=0.02,
        repulsion_stiffness=1e4,
        max_checks_per_frame=1000,
    ):
        """
        Args:
            collision_distance: minimum distance threshold
            repulsion_stiffness: strength of repulsion forces
            max_checks_per_frame: maximum collision checks per frame (for performance)
        """
        self.d_min = collision_distance
        self.k_repulsion = repulsion_stiffness
        self.max_checks = max_checks_per_frame




    def compute_self_collision_forces(self, mesh, positions):
        """
        Fast self-collision detection using simple distance checks

        Args:
            mesh: TetrahedralMesh
            positions: :math:`(N, 3)` current positions

        Returns:
            forces: :math:`(N, 3)` repulsion forces
        """
        forces = torch.zeros_like(positions)

        # Vectorized collision detection for speed
        N = positions.shape[0]

        # Limit checks for performance - sample vertex pairs
        num_samples = min(self.max_checks, N * 5)

        if num_samples == 0:
            return forces

        # Generate random pairs in batch
        indices_i = torch.randint(0, N, (num_samples,), device=positions.device)
        indices_j = torch.randint(0, N, (num_samples,), device=positions.device)

        # Filter out self-pairs
        valid = indices_i != indices_j
        indices_i = indices_i[valid]
        indices_j = indices_j[valid]

        if len(indices_i) == 0:
            return forces

        # Compute distances in batch
        pos_i = positions[indices_i]
        pos_j = positions[indices_j]
        diff = pos_i - pos_j
        dist = torch.norm(diff, dim=1, keepdim=True) + 1e-12

        # Find colliding pairs
        colliding = dist.squeeze(-1) < self.d_min

        if not colliding.any():
            return forces

        # Compute repulsion forces for colliding pairs
        dist_colliding = dist[colliding]
        diff_colliding = diff[colliding]
        penetration = self.d_min - dist_colliding
        force_mag = self.k_repulsion * penetration / dist_colliding
        force = force_mag * diff_colliding

        # Accumulate forces (using scatter_add for efficiency)
        indices_i_colliding = indices_i[colliding]
        indices_j_colliding = indices_j[colliding]

        forces.index_add_(0, indices_i_colliding, force)
        forces.index_add_(0, indices_j_colliding, -force)

        return forces