Added sweep and prune collision detection

docs/_toc.yml

@@ -27,6 +27,7 @@ parts:
   - file: examples/lagrangian/lagrangian_intro
     sections:
     - file: examples/lagrangian/basic_lagrangian_box
+    - file: examples/lagrangian/sweep_and_prune
   - file: examples/equilibria/equilibria_intro
     sections:
     - file: examples/equilibria/activity_part1

particula/lagrangian/boundary.py

@@ -11,7 +11,8 @@ def wrapped_cube(
     This function modifies positions that exceed the cubic domain side,
     wrapping them around to the opposite side of the domain. It handles both
     positive and negative overflows. The center of the cube is assumed to be
-    at zero.
+    at zero. If a particle is way outside the cube, it is wrapped around to
+    the opposite side of the cube.
 
     Parameters:
     - position (torch.Tensor): A tensor representing positions that might
@@ -40,4 +41,17 @@ def wrapped_cube(
         position < -half_cube_side,
         position + cube_side,
         position)
+
+    # Wrap around for very large positive overflow
+    position = torch.where(
+        position > half_cube_side,
+        -half_cube_side,
+        position)
+
+    # Wrap around for very large negative overflow
+    position = torch.where(
+        position < -half_cube_side,
+        half_cube_side,
+        position)
+
     return position

particula/lagrangian/collisions.py

@@ -67,8 +67,10 @@ def find_collisions(
 
 
 def coalescence(
+        position: torch.Tensor,
         velocity: torch.Tensor,
         mass: torch.Tensor,
+        radius: torch.Tensor,
         collision_indices_pairs: torch.Tensor,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """
@@ -80,9 +82,12 @@ def coalescence(
     according to the conservation of mass and momentum.
 
     Parameters:
+    position (torch.Tensor): A 2D tensor of shape [n_dimensions, n_particles]
+        representing the positions of particles.
     velocity (torch.Tensor): A 2D tensor of shape [n_dimensions, n_particles]
         representing the velocities of particles.
     mass (torch.Tensor): A 1D tensor containing the mass of each particle.
+    radius (torch.Tensor): A 1D tensor containing the radius of each particle.
     collision_indices_pairs (torch.Tensor): A 2D tensor containing pairs of
         indices representing colliding particles.
     remove_duplicates_func (function): A function to remove duplicate entries
@@ -102,6 +107,17 @@ def coalescence(
     sorted_pairs, _ = torch.sort(collision_indices_pairs, dim=1)
     unique_left_indices = particle_pairs.remove_duplicates(sorted_pairs, 0)
     unique_indices = particle_pairs.remove_duplicates(unique_left_indices, 1)
+    # fast return if no collisions
+    if unique_indices.shape[0] == 0:
+        return velocity, mass
+    unique_indices = particle_pairs.validate_pair_distance(
+        collision_indices_pairs=unique_indices,
+        position=position,
+        radius=radius,
+    )
+    # fast return if no collisions
+    if unique_indices.shape[0] == 0:
+        return velocity, mass
 
     # Update velocities based on conservation of momentum
     total_mass = mass[unique_indices[:, 0]] + mass[unique_indices[:, 1]]

particula/lagrangian/particle_pairs.py

@@ -1,5 +1,6 @@
 """Lagrangian particle pairwise distances and pairwise operations."""
 
+from typing import Tuple
 import torch
 
 
@@ -73,3 +74,221 @@ def calculate_pairwise_distance(position: torch.Tensor) -> torch.Tensor:
     detla_position = position.unsqueeze(2) - position.unsqueeze(1)
     # Compute pairwise Euclidean distances
     return torch.sqrt(torch.sum(detla_position**2, dim=0))
+
+
+def validate_pair_distance(
+    collision_indices_pairs: torch.Tensor,
+    position: torch.Tensor,
+    radius: torch.Tensor
+) -> torch.Tensor:
+    """
+    Validates if the Euclidean distances between pairs of points are smaller
+    than the sum of their radii.
+
+    Args:
+        collision_indices_pairs (torch.Tensor): A tensor containing pairs of
+            indices of potentially colliding particles.
+        position (torch.Tensor): A 2D tensor of particle positions, where each
+            column represents a particle, and each row represents an axis.
+        radius (torch.Tensor): A 1D tensor representing the radius of each
+            particle.
+
+    Returns:
+        torch.Tensor: A tensor containing the indices of the pairs of
+            particles that are actually colliding.
+    """
+    # Fast return if there are no particle pairs
+    if collision_indices_pairs.numel() == 0:
+        return torch.tensor([], dtype=torch.int64)
+
+    # Calculate 3D distance for each pair of particles
+    delta_position = position[:, collision_indices_pairs[:, 0]] \
+        - position[:, collision_indices_pairs[:, 1]]
+    # Euclidean distance between particles
+    distance = torch.sqrt(torch.sum(delta_position**2, axis=0))
+    # radius sum of both particles
+    distance_threshold = radius[collision_indices_pairs[:, 0]] \
+        + radius[collision_indices_pairs[:, 1]]
+
+    # Return the pairs of particles where the distance is less than the sum of
+    # their radii
+    return collision_indices_pairs[distance < distance_threshold]
+
+
+def single_axis_sweep_and_prune(
+    position_axis: torch.Tensor,
+    radius: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Sweep and prune algorithm for collision detection along a single axis.
+    This function identifies pairs of particles that are close enough to
+    potentially collide.
+
+    Args:
+        position_axis (torch.Tensor): The position of particles along a single
+            axis.
+        radius (torch.Tensor): The radius of particles.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Two tensors containing the indices
+        of potentially colliding particles.
+    """
+
+    # Fast return if there are no particles
+    if position_axis.shape[0] == 0:
+        return torch.tensor([], dtype=torch.int64), torch.tensor(
+            [], dtype=torch.int64)
+
+    # Apply sweep and prune to find pairs of particles that are close enough
+    # to collide
+    sweep = torch.sort(position_axis)
+    sweep_diff = torch.diff(sweep.values)
+    radius_sum = radius[sweep.indices[:-1]] + radius[sweep.indices[1:]]
+
+    # Select indices of particles that are close enough to collide
+    prune_bool = sweep_diff < radius_sum
+    left_overlap_indices = sweep.indices[torch.cat(
+        [prune_bool, torch.tensor([False], dtype=torch.bool)])]
+    right_overlap_indices = sweep.indices[torch.cat(
+        [torch.tensor([False], dtype=torch.bool), prune_bool])]
+
+    return left_overlap_indices, right_overlap_indices
+
+
+# pylint: disable=too-many-locals
+def full_sweep_and_prune(
+        position: torch.Tensor,
+        radius: torch.Tensor
+) -> torch.Tensor:
+    """
+    Sweep and prune algorithm for collision detection along all three axes
+    (x, y, z). This function identifies pairs of particles that are close
+    enough to potentially collide in 3D space.
+
+    Args:
+        position (torch.Tensor): The 2D tensor of particle positions,
+            where each row represents an axis (x, y, z).
+        radius (torch.Tensor): The radius of particles.
+
+    Returns:
+        torch.Tensor: A tensor containing pairs of indices of potentially
+            colliding particles.
+    """
+    # select only particles with positive radius
+    valid_radius = radius > 0
+    valid_radius_indices = torch.arange(radius.shape[0])[valid_radius]
+    # sweep x axis
+    left_x_overlap_shifted, right_x_overlap_shifted = \
+        single_axis_sweep_and_prune(
+            position_axis=position[0, valid_radius_indices],
+            radius=radius[valid_radius_indices]
+        )
+    # fast return if there are no particles overlapping in x
+    if left_x_overlap_shifted.shape[0] == 0:
+        return torch.tensor([])
+    # unshift from relative valid radius to position index
+    left_x_overlap_indices = valid_radius_indices[left_x_overlap_shifted]
+    right_x_overlap_indices = valid_radius_indices[right_x_overlap_shifted]
+
+    # cobine left and right indices for next step
+    all_overlaps_x = torch.cat(
+        [left_x_overlap_indices, right_x_overlap_indices])
+    # select unique indices
+    indices_x_unique = torch.unique(all_overlaps_x)
+
+    # sweep y axis
+    left_y_overlap_shifted, right_y_overlap_shifted = \
+        single_axis_sweep_and_prune(
+            position_axis=position[1][indices_x_unique],
+            radius=radius[indices_x_unique]
+        )
+    # fast return if there are no particles overlapping in y
+    if left_y_overlap_shifted.shape[0] == 0:
+        return torch.tensor([])
+    # unshift from x relative index to position index
+    left_y_overlap_indices = indices_x_unique[left_y_overlap_shifted]
+    right_y_overlap_indices = indices_x_unique[right_y_overlap_shifted]
+
+    # combine left and right indices for next step
+    all_overlaps_y = torch.cat(
+        [left_y_overlap_indices, right_y_overlap_indices])
+    # select unique indices
+    indices_y_unique = torch.unique(all_overlaps_y)
+
+    # sweep z axis
+    left_z_overlap_shifted, right_z_overlap_shifted = \
+        single_axis_sweep_and_prune(
+            position_axis=position[2][indices_y_unique],
+            radius=radius[indices_y_unique]
+        )
+    # fast return if there are no particles overlapping in z
+    if left_z_overlap_shifted.shape[0] == 0:
+        return torch.tensor([])
+    # unshift from y relative index to position index
+    left_z_overlap_indices = indices_y_unique[left_z_overlap_shifted]
+    right_z_overlap_indices = indices_y_unique[right_z_overlap_shifted]
+
+    return torch.cat(
+        [
+            left_z_overlap_indices.unsqueeze(1),
+            right_z_overlap_indices.unsqueeze(1),
+        ],
+        dim=1,
+    )
+
+
+def full_sweep_and_prune_simplified(
+    position: torch.Tensor,
+    radius: torch.Tensor
+) -> torch.Tensor:
+    """
+    A simplified version of the full sweep and prune algorithm for collision
+    written above, it is not working yet. there is an error in the update of
+    the indices in the y and z axis.
+
+    Sweep and prune algorithm for collision detection along all three axes
+    (x, y, z). This function identifies pairs of particles that are close
+    enough to potentially collide in 3D space.
+
+    Args:
+        position (torch.Tensor): The 2D tensor of particle positions,
+            where each row represents an axis (x, y, z).
+        radius (torch.Tensor): The radius of particles.
+
+    Returns:
+        torch.Tensor: A tensor containing pairs of indices of potentially
+            colliding particles.
+    """
+    if radius.shape[0] == 0:
+        return torch.tensor([])
+
+    valid_indices = torch.arange(radius.shape[0])[radius > 0]
+    unique_indices = valid_indices
+
+    for axis in range(position.shape[0]):
+        if unique_indices.shape[0] == 0:
+            break
+
+        left_overlap_indices, right_overlap_indices = \
+            single_axis_sweep_and_prune(
+                position_axis=position[axis, unique_indices],
+                radius=radius[unique_indices]
+            )
+
+        if left_overlap_indices.shape[0] == 0:
+            return torch.tensor([])
+
+        # Combine indices to form collision pairs, may still have duplicates
+        all_overlaps = torch.cat(
+            [unique_indices[left_overlap_indices],
+             unique_indices[right_overlap_indices]])
+        if axis < position.shape[0]:  # not last axis
+            unique_indices = torch.unique(all_overlaps)  # remove duplicates
+
+    return torch.cat(
+        [
+            unique_indices[left_overlap_indices].unsqueeze(1),
+            unique_indices[right_overlap_indices].unsqueeze(1),
+        ],
+        dim=1,
+    )

particula/lagrangian/tests/collisions_test.py

@@ -23,13 +23,22 @@ def test_coalescence():
                              [7.0, 8.0, 9.0],
                              [0.0, 0.0, 0.0]]).T
     mass = torch.tensor([1.0, 2.0, 3.0, 1.0])
+    radius = torch.tensor([1.0, 2.0, 0.5, 0.5])
+    position = torch.tensor([[1.0, 1.0, 2.0],
+                             [1.0, 1.0, 3.0],
+                             [7.0, 8.0, 9.0],
+                             [0.0, 0.0, 0.0]]).T
     collision_indices_pairs = torch.tensor([[0, 1]], dtype=torch.int32)
     expected_velocity = torch.tensor([[1.0, 2.0, 5.0],
                                       [0.0, 0.0, 0.0],
                                       [7.0, 8.0, 9.0],
                                       [0.0, 0.0, 0.0]]).T
     expected_mass = torch.tensor([3.0, 0.0, 3.0, 1.0])
     result_velocity, result_mass = collisions.coalescence(
-        velocity, mass, collision_indices_pairs)
+        position=position,
+        velocity=velocity,
+        mass=mass,
+        radius=radius,
+        collision_indices_pairs=collision_indices_pairs)
     assert torch.allclose(result_mass, expected_mass, atol=1e-4)
     assert torch.allclose(result_velocity, expected_velocity, atol=1e-4)

particula/lagrangian/tests/particle_pairs_test.py

@@ -21,3 +21,46 @@ def test_calculate_pairwise_distance():
                                     [2.2361, 1.0, 0.0]])
     result = particle_pairs.calculate_pairwise_distance(position)
     assert torch.allclose(result, expected_output, atol=1e-4)
+
+
+def test_single_axis_sweep_and_prune():
+    """Test single axis sweep and prune algorithm."""
+    # Test case with some overlapping particles
+    position_axis = torch.tensor([1.0, 2.0, 4.0, 5.0])
+    radius = torch.tensor([0.5, 1.5, 0.2, 0.2])
+    left_indices, right_indices = particle_pairs.single_axis_sweep_and_prune(
+        position_axis, radius)
+
+    assert torch.equal(left_indices, torch.tensor([0]))
+    assert torch.equal(right_indices, torch.tensor([1]))
+
+    # Test case with no particles
+    position_axis = torch.tensor([])
+    radius = torch.tensor([])
+    left_indices, right_indices = particle_pairs.single_axis_sweep_and_prune(
+        position_axis, radius)
+
+    assert torch.equal(left_indices, torch.tensor([]))
+    assert torch.equal(right_indices, torch.tensor([]))
+
+
+def test_validate_pair_distance():
+    """Test validating pair distances."""
+    # Mock data
+    collision_indices_pairs = torch.tensor([[0, 1], [1, 2]])
+    position = torch.tensor([[0.0, 1.0, 1.0],
+                             [0.0, 0.0, 0.0],
+                             [5.0, 5.0, 0.0]])
+    radius = torch.tensor([1.5, 1.5, 0.5])
+
+    # Expected output: Only the first pair should collide
+    expected_output = torch.tensor([[0, 1]])
+
+    # Run the function
+    actual_output = particle_pairs.validate_pair_distance(
+        collision_indices_pairs=collision_indices_pairs,
+        position=position,
+        radius=radius)
+
+    # Assert the result
+    assert torch.equal(actual_output, expected_output)