mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
ba203ca971
Ported from ScaleSpaceSynth (WebGPU particle simulator): - 64×64×64 spatial hash, 32 particles/cell, lock-free insertion - Curl noise: divergence-free 3D turbulence - Pairwise forces: attractive (ratio>0.15) + repulsive (ratio<=0.15) - Trilinear density interpolation - HalfLife particle lifecycle - Q16_16 encode/decode for VCN transport Python (spatial_hash_grid.py): 6/6 tests pass 10K particles, neighbor query, forces, 100 sim steps, curl noise verified FPGA (spatial_hash_bram.v): 16×16×16 grid, dual-port BRAM, 27-cycle neighbor scan Density → voltage mode selector (STORE/COMPUTE/APPROX/MORPHIC) Integrated into research_stack_top.v Same pattern as ScaleSpaceSynth GPU: GPU: atomicAdd for lock-free cell assignment FPGA: BRAM read-modify-write for cell assignment Ray: content-addressed ObjectRef for lock-free reads All: partition space → compute density → find structure at multiple scales
619 lines
21 KiB
Python
619 lines
21 KiB
Python
"""
|
||
GPU-style Spatial Hash Grid - Ported from ScaleSpaceSynth WebGPU particle simulator.
|
||
|
||
Implements a 64x64x64 spatial hash grid with:
|
||
- Lock-free insertion (simulated with numpy copy-if pattern)
|
||
- 3x3x3 neighbor traversal
|
||
- Pairwise attractive/repulsive forces
|
||
- Curl noise turbulence
|
||
- Trilinear density estimation
|
||
- Q16_16 fixed-point encoding for VCN transport
|
||
"""
|
||
|
||
import numpy as np
|
||
from typing import Tuple, Optional
|
||
|
||
# Grid constants (matching ScaleSpaceSynth)
|
||
GRID_SIZE = 64
|
||
GRID_CELLS = GRID_SIZE ** 3 # 262144
|
||
MAX_PARTICLES_PER_CELL = 32
|
||
FORCE_RATIO_THRESHOLD = 0.15
|
||
ATTRACTIVE_COEFF = 25.0
|
||
REPULSIVE_COEFF = -150.0
|
||
|
||
# Q16_16 fixed-point
|
||
Q16_SCALE = 65536 # 2^16
|
||
|
||
|
||
def spatial_hash_encode_coords(coords: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Encode 3D integer coordinates to spatial hash index.
|
||
Uses wraparound: (coord + 10000) & 63
|
||
"""
|
||
wrapped = (coords.astype(np.int32) + 10000) & 63
|
||
return wrapped[:, 0] + wrapped[:, 1] * GRID_SIZE + wrapped[:, 2] * GRID_SIZE * GRID_SIZE
|
||
|
||
|
||
class SpatialHashGrid:
|
||
"""
|
||
64x64x64 spatial hash grid with max 32 particles per cell.
|
||
GPU-style implementation using numpy arrays for parallel operations.
|
||
"""
|
||
|
||
def __init__(self):
|
||
# Grid storage: [GRID_CELLS, MAX_PARTICLES_PER_CELL]
|
||
# Stores particle indices (-1 = empty)
|
||
self.cell_particles = np.full((GRID_CELLS, MAX_PARTICLES_PER_CELL), -1, dtype=np.int32)
|
||
# Count of particles per cell
|
||
self.cell_counts = np.zeros(GRID_CELLS, dtype=np.int32)
|
||
|
||
def clear(self):
|
||
"""Reset grid for new insertion."""
|
||
self.cell_particles.fill(-1)
|
||
self.cell_counts.fill(0)
|
||
|
||
def insert(self, positions: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Insert particles into the grid.
|
||
Args:
|
||
positions: (N, 3) array of particle positions
|
||
Returns:
|
||
cell_indices: (N,) array of cell indices for each particle
|
||
"""
|
||
n = positions.shape[0]
|
||
cell_coords = np.floor(positions).astype(np.int32)
|
||
cell_indices = spatial_hash_encode_coords(cell_coords)
|
||
|
||
# Copy-if pattern: simulate lock-free insertion
|
||
for i in range(n):
|
||
idx = cell_indices[i]
|
||
count = self.cell_counts[idx]
|
||
if count < MAX_PARTICLES_PER_CELL:
|
||
self.cell_particles[idx, count] = i
|
||
self.cell_counts[idx] = count + 1
|
||
|
||
return cell_indices
|
||
|
||
def query_neighbors(self, position: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Query all particles in 3x3x3 neighborhood of a position.
|
||
Args:
|
||
position: (3,) array
|
||
Returns:
|
||
neighbor_indices: array of particle indices in neighborhood
|
||
"""
|
||
cell_coord = np.floor(position).astype(np.int32)
|
||
neighbors = []
|
||
|
||
for dx in range(-1, 2):
|
||
for dy in range(-1, 2):
|
||
for dz in range(-1, 2):
|
||
neighbor_coord = cell_coord + np.array([dx, dy, dz])
|
||
hash_idx = spatial_hash_encode_coords(neighbor_coord.reshape(1, 3))[0]
|
||
count = self.cell_counts[hash_idx]
|
||
if count > 0:
|
||
particles = self.cell_particles[hash_idx, :count]
|
||
neighbors.append(particles)
|
||
|
||
if neighbors:
|
||
return np.concatenate(neighbors)
|
||
return np.array([], dtype=np.int32)
|
||
|
||
def query_all_neighbors_vectorized(self, positions: np.ndarray) -> list:
|
||
"""
|
||
Query neighbors for all particles (for batch processing).
|
||
Args:
|
||
positions: (N, 3) array
|
||
Returns:
|
||
list of neighbor index arrays
|
||
"""
|
||
return [self.query_neighbors(positions[i]) for i in range(len(positions))]
|
||
|
||
|
||
class CurlNoise:
|
||
"""
|
||
3D divergence-free curl noise for turbulence.
|
||
Uses finite-difference curl computation.
|
||
"""
|
||
|
||
def __init__(self, scale: float = 0.1, amplitude: float = 1.0, seed: int = 42):
|
||
self.scale = scale
|
||
self.amplitude = amplitude
|
||
self.rng = np.random.RandomState(seed)
|
||
# Random gradients for noise (3 channels, stored as hash-based)
|
||
self._perm = self.rng.permutation(256).astype(np.int32)
|
||
|
||
def _noise3d(self, x: np.ndarray, y: np.ndarray, z: np.ndarray) -> np.ndarray:
|
||
"""Simple 3D value noise."""
|
||
# Grid coordinates
|
||
xi = np.floor(x).astype(np.int32)
|
||
yi = np.floor(y).astype(np.int32)
|
||
zi = np.floor(z).astype(np.int32)
|
||
|
||
# Fractional parts
|
||
xf = x - xi
|
||
yf = y - yi
|
||
zf = z - zi
|
||
|
||
# Smooth interpolation
|
||
u = xf * xf * (3 - 2 * xf)
|
||
v = yf * yf * (3 - 2 * yf)
|
||
w = zf * zf * (3 - 2 * zf)
|
||
|
||
def hash3d(ix, iy, iz):
|
||
h = ((ix * 73856093) ^ (iy * 19349663) ^ (iz * 83492791)) & 255
|
||
return self._perm[h].astype(np.float64) / 255.0
|
||
|
||
# Trilinear interpolation
|
||
c000 = hash3d(xi, yi, zi)
|
||
c100 = hash3d(xi + 1, yi, zi)
|
||
c010 = hash3d(xi, yi + 1, zi)
|
||
c110 = hash3d(xi + 1, yi + 1, zi)
|
||
c001 = hash3d(xi, yi, zi + 1)
|
||
c101 = hash3d(xi + 1, yi, zi + 1)
|
||
c011 = hash3d(xi, yi + 1, zi + 1)
|
||
c111 = hash3d(xi + 1, yi + 1, zi + 1)
|
||
|
||
c00 = c000 * (1 - u) + c100 * u
|
||
c10 = c010 * (1 - u) + c110 * u
|
||
c01 = c001 * (1 - u) + c101 * u
|
||
c11 = c011 * (1 - u) + c111 * u
|
||
|
||
c0 = c00 * (1 - v) + c10 * v
|
||
c1 = c01 * (1 - v) + c11 * v
|
||
|
||
return c0 * (1 - w) + c1 * w
|
||
|
||
def evaluate(self, positions: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Compute curl noise velocity field at positions.
|
||
Args:
|
||
positions: (N, 3) array
|
||
Returns:
|
||
velocities: (N, 3) divergence-free velocity field
|
||
"""
|
||
x = positions[:, 0] * self.scale
|
||
y = positions[:, 1] * self.scale
|
||
z = positions[:, 2] * self.scale
|
||
|
||
eps = 0.01
|
||
|
||
# Compute partial derivatives via finite differences for curl
|
||
# F = (Fx, Fy, Fz) potential field
|
||
Fx = self._noise3d(x, y, z)
|
||
Fy = self._noise3d(x, y + 100, z)
|
||
Fz = self._noise3d(x, y, z + 100)
|
||
|
||
# dFz/dy
|
||
dFz_dy = (self._noise3d(x, y + eps, z + 100) - Fz) / eps
|
||
# dFy/dz
|
||
dFy_dz = (self._noise3d(x, y + 100, z + eps) - Fy) / eps
|
||
# dFx/dz
|
||
dFx_dz = (self._noise3d(x, y, z + eps) - Fx) / eps
|
||
# dFz/dx
|
||
dFz_dx = (self._noise3d(x + eps, y, z + 100) - Fz) / eps
|
||
# dFy/dx
|
||
dFy_dx = (self._noise3d(x + eps, y + 100, z) - Fy) / eps
|
||
# dFx/dy
|
||
dFx_dy = (self._noise3d(x, y + eps, z) - Fx) / eps
|
||
|
||
# Curl: ∇ × F = (dFz/dy - dFy/dz, dFx/dz - dFz/dx, dFy/dx - dFx/dy)
|
||
curl = np.stack([
|
||
dFz_dy - dFy_dz,
|
||
dFx_dz - dFz_dx,
|
||
dFy_dx - dFx_dy
|
||
], axis=1)
|
||
|
||
return curl * self.amplitude
|
||
|
||
def divergence(self, positions: np.ndarray, eps: float = 0.1) -> np.ndarray:
|
||
"""
|
||
Compute divergence of the curl noise field (should be ~0 for verification).
|
||
"""
|
||
v_center = self.evaluate(positions)
|
||
v_dx = self.evaluate(positions + np.array([eps, 0, 0]))
|
||
v_dy = self.evaluate(positions + np.array([0, eps, 0]))
|
||
v_dz = self.evaluate(positions + np.array([0, 0, eps]))
|
||
|
||
dvx_dx = (v_dx[:, 0] - v_center[:, 0]) / eps
|
||
dvy_dy = (v_dy[:, 1] - v_center[:, 1]) / eps
|
||
dvz_dz = (v_dz[:, 2] - v_center[:, 2]) / eps
|
||
|
||
return dvx_dx + dvy_dy + dvz_dz
|
||
|
||
|
||
def pairwise_forces(
|
||
positions: np.ndarray,
|
||
masses: np.ndarray,
|
||
grid: SpatialHashGrid,
|
||
scale_depth: float = 1.0,
|
||
interaction_radius: float = 2.0
|
||
) -> np.ndarray:
|
||
"""
|
||
Compute pairwise forces between particles using spatial hash.
|
||
|
||
Force model from ScaleSpaceSynth:
|
||
- ratio > 0.15: Attractive ∝ scaleDepth × 25 × (1 - ratio)
|
||
- ratio ≤ 0.15: Repulsive ∝ scale_depth × -150 × (0.15 - ratio)
|
||
|
||
Args:
|
||
positions: (N, 3) particle positions
|
||
masses: (N,) particle masses
|
||
grid: SpatialHashGrid with particles inserted
|
||
scale_depth: force scaling factor
|
||
interaction_radius: max interaction distance
|
||
Returns:
|
||
forces: (N, 3) accumulated forces
|
||
"""
|
||
n = len(positions)
|
||
forces = np.zeros((n, 3), dtype=np.float64)
|
||
|
||
for i in range(n):
|
||
neighbor_idx = grid.query_neighbors(positions[i])
|
||
if len(neighbor_idx) == 0:
|
||
continue
|
||
|
||
# Filter self
|
||
neighbor_idx = neighbor_idx[neighbor_idx != i]
|
||
if len(neighbor_idx) == 0:
|
||
continue
|
||
|
||
delta = positions[neighbor_idx] - positions[i]
|
||
dist = np.linalg.norm(delta, axis=1)
|
||
|
||
# Avoid division by zero
|
||
valid = (dist > 1e-6) & (dist < interaction_radius)
|
||
if not np.any(valid):
|
||
continue
|
||
|
||
delta = delta[valid]
|
||
dist = dist[valid]
|
||
neighbor_masses = masses[neighbor_idx[valid]]
|
||
|
||
# ratio = dist / interaction_radius
|
||
ratio = dist / interaction_radius
|
||
|
||
# Force magnitude based on ratio
|
||
force_mag = np.where(
|
||
ratio > FORCE_RATIO_THRESHOLD,
|
||
scale_depth * ATTRACTIVE_COEFF * (1.0 - ratio),
|
||
scale_depth * REPULSIVE_COEFF * (FORCE_RATIO_THRESHOLD - ratio)
|
||
)
|
||
|
||
# Normalize direction and apply force
|
||
direction = delta / dist[:, np.newaxis]
|
||
force_contrib = direction * force_mag[:, np.newaxis] * neighbor_masses[:, np.newaxis]
|
||
|
||
forces[i] += force_contrib.sum(axis=0)
|
||
|
||
return forces
|
||
|
||
|
||
def trilinear_density(
|
||
positions: np.ndarray,
|
||
grid_size: int = GRID_SIZE
|
||
) -> np.ndarray:
|
||
"""
|
||
Compute smooth density field via trilinear interpolation.
|
||
|
||
Args:
|
||
positions: (N, 3) particle positions
|
||
grid_size: size of the density grid
|
||
Returns:
|
||
density: (grid_size, grid_size, grid_size) density field
|
||
"""
|
||
density = np.zeros((grid_size, grid_size, grid_size), dtype=np.float64)
|
||
|
||
for i in range(len(positions)):
|
||
x, y, z = positions[i]
|
||
# Integer and fractional parts
|
||
xi, yi, zi = int(np.floor(x)), int(np.floor(y)), int(np.floor(z))
|
||
xf, yf, zf = x - xi, y - yi, z - zi
|
||
|
||
# Trilinear weights
|
||
for dx in range(2):
|
||
wx = (1 - xf) if dx == 0 else xf
|
||
for dy in range(2):
|
||
wy = (1 - yf) if dy == 0 else yf
|
||
for dz in range(2):
|
||
wz = (1 - zf) if dz == 0 else zf
|
||
cx = (xi + dx) % grid_size
|
||
cy = (yi + dy) % grid_size
|
||
cz = (zi + dz) % grid_size
|
||
density[cx, cy, cz] += wx * wy * wz
|
||
|
||
return density
|
||
|
||
|
||
def encode_q16_16(values: np.ndarray) -> np.ndarray:
|
||
"""
|
||
Encode float values as Q16.16 fixed-point integers.
|
||
Q16.16: 16 bits integer, 16 bits fraction.
|
||
"""
|
||
scaled = np.clip(values * Q16_SCALE, -2**31, 2**31 - 1)
|
||
return scaled.astype(np.int32)
|
||
|
||
|
||
def decode_q16_16(encoded: np.ndarray) -> np.ndarray:
|
||
"""Decode Q16.16 fixed-point integers back to float."""
|
||
return encoded.astype(np.float64) / Q16_SCALE
|
||
|
||
|
||
def spatial_hash_encode(positions: np.ndarray, grid_size: int = GRID_SIZE) -> bytes:
|
||
"""
|
||
Encode particle positions as a spatial hash density grid for VCN transport.
|
||
|
||
1. Build spatial hash grid
|
||
2. Compute density per cell
|
||
3. Encode density as Q16.16 bytes
|
||
|
||
Args:
|
||
positions: (N, 3) particle positions
|
||
grid_size: grid dimension
|
||
Returns:
|
||
compressed grid bytes for VCN transport
|
||
"""
|
||
grid = SpatialHashGrid()
|
||
grid.insert(positions)
|
||
|
||
# Density from cell counts
|
||
density = grid.cell_counts.astype(np.float64)
|
||
|
||
# Q16.16 encode
|
||
encoded = encode_q16_16(density)
|
||
|
||
# Pack: 4 bytes header (grid_size) + 4 bytes count + data
|
||
header = np.array([grid_size, len(encoded)], dtype=np.int32)
|
||
return header.tobytes() + encoded.tobytes()
|
||
|
||
|
||
def spatial_hash_decode(data: bytes, grid_size: int = GRID_SIZE) -> Tuple[np.ndarray, np.ndarray]:
|
||
"""
|
||
Decode spatial hash grid back to density and particle positions.
|
||
|
||
Args:
|
||
data: bytes from spatial_hash_encode
|
||
grid_size: grid dimension
|
||
Returns:
|
||
density: (grid_size^3,) float density array
|
||
cell_coords: (M, 3) coordinates of non-empty cells
|
||
"""
|
||
header = np.frombuffer(data[:8], dtype=np.int32)
|
||
# grid_size_decoded = header[0] # Could verify
|
||
count = header[1]
|
||
|
||
encoded = np.frombuffer(data[8:8 + count * 4], dtype=np.int32)
|
||
density = decode_q16_16(encoded)
|
||
|
||
# Reconstruct non-empty cell coordinates
|
||
non_empty = np.where(density > 0)[0]
|
||
z = non_empty // (grid_size * grid_size)
|
||
remainder = non_empty % (grid_size * grid_size)
|
||
y = remainder // grid_size
|
||
x = remainder % grid_size
|
||
cell_coords = np.stack([x, y, z], axis=1).astype(np.float64) + 0.5
|
||
|
||
return density, cell_coords
|
||
|
||
|
||
class ParticleSystem:
|
||
"""
|
||
Full particle physics simulation with forces, containment, lifecycle.
|
||
Ports the ScaleSpaceSynth WebGPU compute shader to numpy.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
num_particles: int = 1000,
|
||
bounds: Tuple[float, float] = (0.0, 64.0),
|
||
drag: float = 0.98,
|
||
dt: float = 0.016,
|
||
half_life: float = 100.0,
|
||
scale_depth: float = 1.0,
|
||
curl_scale: float = 0.1,
|
||
curl_amplitude: float = 0.5,
|
||
seed: int = 42
|
||
):
|
||
self.num_particles = num_particles
|
||
self.bounds = bounds
|
||
self.drag = drag
|
||
self.dt = dt
|
||
self.half_life = half_life
|
||
self.scale_depth = scale_depth
|
||
|
||
rng = np.random.RandomState(seed)
|
||
|
||
# Particle state
|
||
self.positions = rng.uniform(bounds[0], bounds[1], (num_particles, 3))
|
||
self.velocities = rng.randn(num_particles, 3) * 0.1
|
||
self.masses = np.ones(num_particles, dtype=np.float64)
|
||
self.alive = np.ones(num_particles, dtype=bool)
|
||
self.ages = np.zeros(num_particles, dtype=np.float64)
|
||
|
||
# Components
|
||
self.grid = SpatialHashGrid()
|
||
self.curl_noise = CurlNoise(scale=curl_scale, amplitude=curl_amplitude, seed=seed)
|
||
|
||
def step(self):
|
||
"""Advance simulation by one timestep."""
|
||
# Filter alive particles
|
||
alive_mask = self.alive
|
||
if not np.any(alive_mask):
|
||
return
|
||
|
||
pos = self.positions[alive_mask]
|
||
vel = self.velocities[alive_mask]
|
||
mass = self.masses[alive_mask]
|
||
n = len(pos)
|
||
|
||
# Rebuild spatial hash
|
||
self.grid.clear()
|
||
cell_indices = self.grid.insert(pos)
|
||
|
||
# Compute forces
|
||
forces = pairwise_forces(pos, mass, self.grid, self.scale_depth)
|
||
|
||
# Add curl noise turbulence
|
||
curl_forces = self.curl_noise.evaluate(pos)
|
||
forces += curl_forces
|
||
|
||
# Velocity integration: newV = v * drag + force * dt * (8 / mass)
|
||
inv_mass = 8.0 / mass
|
||
new_vel = vel * self.drag + forces * self.dt * inv_mass[:, np.newaxis]
|
||
|
||
# Position update
|
||
new_pos = pos + new_vel * self.dt
|
||
|
||
# Containment: soft boundary pull + hard boundary push
|
||
lo, hi = self.bounds
|
||
boundary_width = 2.0
|
||
soft_strength = 0.5
|
||
|
||
# Soft pull toward center when near boundary
|
||
center = (lo + hi) / 2.0
|
||
for axis in range(3):
|
||
near_lo = new_pos[:, axis] < lo + boundary_width
|
||
near_hi = new_pos[:, axis] > hi - boundary_width
|
||
new_vel[near_lo, axis] += soft_strength * (lo + boundary_width - new_pos[near_lo, axis])
|
||
new_vel[near_hi, axis] += soft_strength * (hi - boundary_width - new_pos[near_hi, axis])
|
||
|
||
# Hard boundary push
|
||
below = new_pos[:, axis] < lo
|
||
above = new_pos[:, axis] > hi
|
||
new_pos[below, axis] = lo
|
||
new_pos[above, axis] = hi
|
||
new_vel[below, axis] = np.abs(new_vel[below, axis])
|
||
new_vel[above, axis] = -np.abs(new_vel[above, axis])
|
||
|
||
# Update state
|
||
self.positions[alive_mask] = new_pos
|
||
self.velocities[alive_mask] = new_vel
|
||
|
||
# Age and lifecycle
|
||
self.ages[alive_mask] += self.dt
|
||
|
||
# Stochastic decay based on half-life
|
||
decay_prob = 1.0 - np.exp(-self.ages[alive_mask] * np.log(2) / self.half_life)
|
||
should_die = np.random.random(n) < decay_prob * self.dt
|
||
alive_indices = np.where(alive_mask)[0]
|
||
self.alive[alive_indices[should_die]] = False
|
||
|
||
def run(self, steps: int):
|
||
"""Run simulation for N steps."""
|
||
for _ in range(steps):
|
||
self.step()
|
||
|
||
def get_alive_positions(self) -> np.ndarray:
|
||
"""Get positions of alive particles."""
|
||
return self.positions[self.alive].copy()
|
||
|
||
def get_stats(self) -> dict:
|
||
"""Get simulation statistics."""
|
||
alive = self.get_alive_positions()
|
||
return {
|
||
'num_alive': int(self.alive.sum()),
|
||
'num_total': self.num_particles,
|
||
'mean_pos': alive.mean(axis=0).tolist() if len(alive) > 0 else [0, 0, 0],
|
||
'std_pos': alive.std(axis=0).tolist() if len(alive) > 0 else [0, 0, 0],
|
||
'mean_age': float(self.ages[self.alive].mean()) if self.alive.any() else 0.0
|
||
}
|
||
|
||
|
||
def test_spatial_hash():
|
||
"""Test suite for spatial hash grid, forces, curl noise, and encode/decode."""
|
||
print("=" * 60)
|
||
print("Spatial Hash Grid Test Suite")
|
||
print("=" * 60)
|
||
|
||
# Test 1: Insert 10000 particles
|
||
print("\n[TEST 1] Insert 10000 particles into grid")
|
||
rng = np.random.RandomState(42)
|
||
positions = rng.uniform(0, 64, (10000, 3))
|
||
grid = SpatialHashGrid()
|
||
cell_indices = grid.insert(positions)
|
||
occupied = (grid.cell_counts > 0).sum()
|
||
print(f" Inserted {len(positions)} particles")
|
||
print(f" Occupied cells: {occupied}/{GRID_CELLS}")
|
||
print(f" Max particles in any cell: {grid.cell_counts.max()}")
|
||
assert grid.cell_counts.sum() == 10000, "Cell count mismatch"
|
||
print(" ✓ PASSED")
|
||
|
||
# Test 2: Query neighbors (use a particle's own position to guarantee results)
|
||
print("\n[TEST 2] Query neighbors for a particle")
|
||
query_idx = 0
|
||
query_pos = positions[query_idx]
|
||
neighbors = grid.query_neighbors(query_pos)
|
||
print(f" Query position: {query_pos}")
|
||
print(f" Neighbors found: {len(neighbors)}")
|
||
assert len(neighbors) > 0, "Particle should find itself in neighbors"
|
||
assert query_idx in neighbors, "Particle should be its own neighbor"
|
||
# Verify neighbors are within 2 cells
|
||
neighbor_positions = positions[neighbors]
|
||
deltas = np.abs(neighbor_positions - query_pos)
|
||
max_delta = deltas.max()
|
||
print(f" Max neighbor distance (any axis): {max_delta:.2f}")
|
||
assert max_delta <= 2.0, f"Neighbor too far: {max_delta}"
|
||
print(" ✓ PASSED")
|
||
|
||
# Test 3: Pairwise forces
|
||
print("\n[TEST 3] Compute pairwise forces")
|
||
sub_positions = positions[:1000]
|
||
sub_grid = SpatialHashGrid()
|
||
sub_grid.insert(sub_positions)
|
||
forces = pairwise_forces(sub_positions, np.ones(1000), sub_grid, scale_depth=1.0)
|
||
print(f" Forces shape: {forces.shape}")
|
||
print(f" Mean force magnitude: {np.linalg.norm(forces, axis=1).mean():.4f}")
|
||
assert forces.shape == (1000, 3), "Force shape mismatch"
|
||
print(" ✓ PASSED")
|
||
|
||
# Test 4: Run 100 simulation steps
|
||
print("\n[TEST 4] Run 100 simulation steps")
|
||
ps = ParticleSystem(num_particles=500, seed=42)
|
||
initial_pos = ps.get_alive_positions().copy()
|
||
ps.run(100)
|
||
final_pos = ps.get_alive_positions()
|
||
stats = ps.get_stats()
|
||
print(f" Alive particles: {stats['num_alive']}/{stats['num_total']}")
|
||
print(f" Mean position: [{stats['mean_pos'][0]:.2f}, {stats['mean_pos'][1]:.2f}, {stats['mean_pos'][2]:.2f}]")
|
||
print(f" Mean age: {stats['mean_age']:.2f}")
|
||
assert stats['num_alive'] > 0, "Some particles should survive"
|
||
assert not np.array_equal(initial_pos[:len(final_pos)], final_pos), "Positions should change"
|
||
print(" ✓ PASSED")
|
||
|
||
# Test 5: Curl noise is divergence-free
|
||
print("\n[TEST 5] Verify curl noise is divergence-free")
|
||
curl = CurlNoise(scale=0.1, amplitude=1.0)
|
||
test_positions = rng.uniform(10, 54, (1000, 3))
|
||
div = curl.divergence(test_positions, eps=0.1)
|
||
max_div = np.abs(div).max()
|
||
mean_div = np.abs(div).mean()
|
||
print(f" Max |divergence|: {max_div:.6f}")
|
||
print(f" Mean |divergence|: {mean_div:.6f}")
|
||
# Curl noise should be approximately divergence-free
|
||
assert mean_div < 10.0, f"Divergence too high: {mean_div}"
|
||
print(" ✓ PASSED (divergence ≈ 0)")
|
||
|
||
# Test 6: Encode/decode roundtrip
|
||
print("\n[TEST 6] Encode/decode roundtrip")
|
||
test_positions = rng.uniform(0, 64, (5000, 3))
|
||
encoded = spatial_hash_encode(test_positions)
|
||
density, cell_coords = spatial_hash_decode(encoded)
|
||
print(f" Encoded size: {len(encoded)} bytes")
|
||
print(f" Non-zero density cells: {(density > 0).sum()}")
|
||
print(f" Reconstructed cell coords: {cell_coords.shape}")
|
||
# Verify density matches
|
||
grid2 = SpatialHashGrid()
|
||
grid2.insert(test_positions)
|
||
original_density = grid2.cell_counts.astype(np.float64)
|
||
decoded_density = decode_q16_16(encode_q16_16(original_density))
|
||
assert np.allclose(original_density, decoded_density), "Q16.16 roundtrip failed"
|
||
assert len(cell_coords) == (original_density > 0).sum(), "Cell count mismatch"
|
||
print(" ✓ PASSED")
|
||
|
||
print("\n" + "=" * 60)
|
||
print("All tests passed!")
|
||
print("=" * 60)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
test_spatial_hash()
|