boids.py

Part of Deep Sea Aquarium.

 1"""Vectorised numpy boids flocking: separation, alignment, cohesion.
 2
 3Every rule is a masked reduction over the pairwise distance matrix, so a
 4school costs one pass of numpy work per frame rather than a Python loop
 5per fish.
 6"""
 7
 8import numpy as np
 9
10
11def compute_boids(
12    positions: np.ndarray,
13    velocities: np.ndarray,
14    separation_radius: float = 2.5,
15    alignment_radius: float = 5.0,
16    cohesion_radius: float = 7.0,
17    separation_weight: float = 3.0,
18    alignment_weight: float = 0.8,
19    cohesion_weight: float = 0.5,
20    bounds: float = 12.0,
21    bounds_weight: float = 1.0,
22) -> np.ndarray:
23    """Compute boids acceleration for all agents.
24
25    Args:
26        positions: (N, 3) array of world positions.
27        velocities: (N, 3) array of current velocities.
28
29    Returns:
30        (N, 3) acceleration array to add to velocities.
31    """
32    n = len(positions)
33    if n < 2:
34        return np.zeros_like(positions)
35
36    # Pairwise displacement and distance
37    diff = positions[:, None, :] - positions[None, :, :]  # (N, N, 3)
38    dist = np.linalg.norm(diff, axis=2)  # (N, N)
39    np.fill_diagonal(dist, 1e10)  # Ignore self
40
41    accel = np.zeros_like(positions)
42
43    # Separation: steer away from neighbours too close. diff[i, j] already points
44    # FROM neighbour j TO self, weighted by inverse square distance.
45    falloff = np.where(dist < separation_radius, 1.0 / (dist * dist + 0.01), 0.0)  # (N, N)
46    accel += (falloff[:, :, None] * diff).sum(axis=1) * separation_weight
47
48    # Alignment: match the mean velocity of nearby neighbours
49    align_mask = (dist < alignment_radius).astype(positions.dtype)  # (N, N)
50    align_count = align_mask.sum(axis=1, keepdims=True)
51    avg_vel = (align_mask @ velocities) / np.maximum(align_count, 1.0)
52    accel += np.where(align_count > 0, avg_vel - velocities, 0.0) * alignment_weight
53
54    # Cohesion: steer toward the centre of the nearby flock
55    coh_mask = (dist < cohesion_radius).astype(positions.dtype)
56    coh_count = coh_mask.sum(axis=1, keepdims=True)
57    centre = (coh_mask @ positions) / np.maximum(coh_count, 1.0)
58    accel += np.where(coh_count > 0, centre - positions, 0.0) * cohesion_weight
59
60    # Boundary avoidance: soft steering near aquarium walls
61    for axis in range(3):
62        over = positions[:, axis] > bounds
63        under = positions[:, axis] < -bounds
64        accel[over, axis] -= bounds_weight * (positions[over, axis] - bounds)
65        accel[under, axis] -= bounds_weight * (positions[under, axis] + bounds)
66
67    # Vertical bias: keep fish roughly horizontal (-3 to 3 Y range)
68    y_high = positions[:, 1] > 3.0
69    y_low = positions[:, 1] < -2.0
70    accel[y_high, 1] -= 1.0
71    accel[y_low, 1] += 1.0
72
73    return accel