""" Flow field system — faithful port from Android's FlowFieldSystem.java. Features: - 40³ velocity grid with trilinear interpolation - Particle lifecycle (lattice init, pre-integration, lifespan, wrapping) - Turbo colormap (exact polynomial from Android) - Ball/probe following flow field - Pre-computed animation frames for dashboard Constants matching Android: GRID = 40, LIFESPAN = 57, PRE_AGE = 34, DT = 0.02 """ from __future__ import annotations from typing import Optional, Tuple import numpy as np # ═══════════════════════════════════════════════════ # TURBO COLORMAP (exact polynomial from FlowFieldSystem.java) # ═══════════════════════════════════════════════════ def turbo_colormap(t: np.ndarray, out: Optional[np.ndarray] = None) -> np.ndarray: """Apply Turbo colormap to speed values. Exact polynomial coefficients from FlowFieldSystem.java. Works with scalar or array input. Args: t: Speed values normalized to [0, 1]. Can be scalar or (N,) array. out: Optional pre-allocated (N, 3) output buffer to avoid allocation. Returns: (N, 3) array of RGB values in [0, 1]. """ t = np.clip(np.atleast_1d(t).astype(np.float64), 0.0, 1.0) r = (34.61 + t * (1172.33 + t * (-10793.56 + t * (33300.12 + t * (-38345.17 + 14829.80 * t))))) / 255.0 g = (23.31 + t * (557.33 + t * (1225.33 + t * (-3574.96 + t * 2199.29)))) / 255.0 b = (27.20 + t * (3211.10 + t * (-15327.97 + t * (34592.87 + t * (-30538.66 + 9347.97 * t))))) / 255.0 if out is not None and out.shape == (len(t), 3): np.clip(r, 0, 1, out=out[:, 0]) np.clip(g, 0, 1, out=out[:, 1]) np.clip(b, 0, 1, out=out[:, 2]) return out rgb = np.stack([np.clip(r, 0, 1), np.clip(g, 0, 1), np.clip(b, 0, 1)], axis=-1) return rgb # ═══════════════════════════════════════════════════ # DIVERGING COLORMAP (matching FlowFieldSystem.java) # ═══════════════════════════════════════════════════ def diverging_colormap(t: np.ndarray) -> np.ndarray: """Red ↔ Blue ↔ Green diverging colormap. Args: t: Values in [-1, 1]. Negative → red, zero → blue, positive → green. Returns: (N, 3) array of RGB values in [0, 1]. """ t = np.clip(np.atleast_1d(t).astype(np.float64), -1.0, 1.0) RED = np.array([1.0, 0.1, 0.1]) BLUE = np.array([0.18, 0.28, 1.0]) GREEN = np.array([0.0, 0.78, 0.0]) rgb = np.zeros((len(t), 3)) neg_mask = t < 0 pos_mask = ~neg_mask # Negative: blue → red if np.any(neg_mask): a = -t[neg_mask] rgb[neg_mask] = BLUE[None, :] * (1 - a[:, None]) + RED[None, :] * a[:, None] # Positive: blue → green if np.any(pos_mask): a = t[pos_mask] rgb[pos_mask] = BLUE[None, :] * (1 - a[:, None]) + GREEN[None, :] * a[:, None] return rgb # ═══════════════════════════════════════════════════ # SCORE COLORMAP (red → yellow → green for 0 → 0.5 → 1) # ═══════════════════════════════════════════════════ def score_colormap(t: np.ndarray) -> np.ndarray: """Map score values [0, 1] to red → yellow → green gradient. 0.0 = red (bad/failure), 0.5 = yellow (neutral), 1.0 = green (success). Args: t: Score values in [0, 1]. Can be scalar or (N,) array. Returns: (N, 3) array of RGB values in [0, 1]. """ t = np.clip(np.atleast_1d(t).astype(np.float64), 0.0, 1.0) rgb = np.zeros((len(t), 3)) # Red channel: 1.0 at t=0, 1.0 at t=0.5, 0.0 at t=1.0 rgb[:, 0] = np.where(t <= 0.5, 1.0, 2.0 * (1.0 - t)) # Green channel: 0.0 at t=0, 1.0 at t=0.5, 1.0 at t=1.0 rgb[:, 1] = np.where(t <= 0.5, 2.0 * t, 1.0) # Blue stays near 0 for vivid colors rgb[:, 2] = 0.05 return np.clip(rgb, 0.0, 1.0) # ═══════════════════════════════════════════════════ # FLOW FIELD SYSTEM (matching FlowFieldSystem.java) # ═══════════════════════════════════════════════════ class FlowFieldSystem: """Particle flow field system. Faithful port from Android's FlowFieldSystem.java. Uses a 40³ velocity grid with trilinear interpolation, particle lifecycle with wrapping, and Turbo speed coloring. Args: velocity_grid: (G, G, G, 3) velocity field array. axis_min: (3,) minimum bounds of the domain. axis_max: (3,) maximum bounds of the domain. particle_grid: Particle grid resolution (default 20 → 8000 particles). Use 40 for full Android fidelity (64,000 particles). """ # Constants matching Android VELOCITY_GRID_SIZE = 40 LIFESPAN = 57 # particle lifetime in frames PRE_AGE = 34 # pre-integration steps DT = 0.02 # world-units per step def __init__( self, velocity_grid: np.ndarray, axis_min: np.ndarray, axis_max: np.ndarray, particle_grid: int = 20, path_points: Optional[np.ndarray] = None, confidence_grid: Optional[np.ndarray] = None, ): self.velocity_grid = velocity_grid.astype(np.float32) self.grid_size = velocity_grid.shape[0] # typically 40 self.axis_min = np.asarray(axis_min, dtype=np.float32) self.axis_max = np.asarray(axis_max, dtype=np.float32) self.span = self.axis_max - self.axis_min self.particle_grid = particle_grid self.particle_count = particle_grid ** 3 self.speed_multiplier = 1.0 # MDN confidence grid (None for RBF models) self._confidence_grid = confidence_grid # (G, G, G) values in [0, 1] self.confidence_strength = 0.0 # 0 = ignore confidence, 1 = full effect # Build path occupancy mask for blob-based spawning self._path_mask = None self._blob_opacity = None self._score_grid = None self._mask_res = 0 self.blob_enabled = True # can be toggled off for full-cube spawning if path_points is not None and len(path_points) >= 2: self._build_path_mask(path_points) # Particle state + pre-allocated buffers (before init/pre-integration) N = self.particle_count self.pos = np.zeros((N, 3), dtype=np.float32) self.age = np.zeros(N, dtype=np.int32) self._buf_velocities = np.zeros((N, 3), dtype=np.float32) self._buf_speeds = np.zeros(N, dtype=np.float32) self._buf_colors = np.zeros((N, 3), dtype=np.float32) self._buf_alphas = np.ones(N, dtype=np.float32) self._buf_pos_out = np.zeros((N, 3), dtype=np.float32) self._buf_normalized = np.zeros((N, 3), dtype=np.float32) # Initialize on lattice (filtered by path mask if available) self._init_lattice() # Save blob-constrained positions as respawn points BEFORE # pre-integration so particles always respawn inside the blob self.orig_pos = self.pos.copy() # Pre-integrate to displace from regular lattice (matching Android) self._dry_integrate(self.PRE_AGE) # After pre-integration, snap any particles that drifted outside # the blob back to their blob-constrained origin if self._blob_opacity is not None and self.blob_enabled: for i in range(self.particle_count): if not self._point_in_blob_world(*self.pos[i]): self.pos[i] = self.orig_pos[i] # Randomize ages to stagger respawns (matching Android initParticles) self.age = np.random.randint(0, self.LIFESPAN, N, dtype=np.int32) # Ball state self.ball_pos = np.array([ (axis_min[0] + axis_max[0]) / 2, (axis_min[1] + axis_max[1]) / 2, (axis_min[2] + axis_max[2]) / 2, ], dtype=np.float32) self.ball_trail: list = [] self.ball_flowing = False # ── Blob tuning constants ────────────────────────────────── # BLOB_RESOLUTION: occupancy grid resolution (higher = finer blob shape) BLOB_RESOLUTION = 42 # BLOB_RADIUS: minimum influence radius around each path point, as # fraction of the per-axis span. The actual radius is computed # adaptively from data density (median NN distance × 1.5) and # floored at this value. Increase for a looser blob. BLOB_RADIUS = 0.03 # BLOB_MARGIN: additional cells beyond the core blob where opacity # fades from 1.0 to 0.0 (soft boundary falloff). BLOB_MARGIN_CELLS = 3 def _build_path_mask(self, path_points: np.ndarray): """Build a 3D float opacity grid from path sample points. The goal is to cut out the empty void where there is absolutely no training data, while keeping everything that is even remotely near any data or path as one continuous cloud. The blob is NOT meant to be a tight segmentation — it should be generous. Steps: 1. Compute adaptive radius from data density (median NN × 1.5, floored at BLOB_RADIUS). 2. Mark core cells around each path point. 3. If the core has multiple disconnected components, dilate until they merge into one connected blob. 4. Add soft opacity margin at the boundary. Args: path_points: (M, 3) densely sampled points along semantic paths. """ from scipy.ndimage import distance_transform_edt, label, binary_dilation res = self.BLOB_RESOLUTION self._mask_res = res pp = np.asarray(path_points, dtype=np.float32) # Normalize path points to [0, 1] per axis norm = np.zeros_like(pp) for a in range(3): if self.span[a] > 0: norm[:, a] = (pp[:, a] - self.axis_min[a]) / self.span[a] # Adaptive radius: median nearest-neighbor distance × 1.5 if len(norm) > 1: sample_idx = np.random.default_rng(0).choice( len(norm), min(500, len(norm)), replace=False ) sample = norm[sample_idx] from scipy.spatial import cKDTree tree = cKDTree(sample) dists, _ = tree.query(sample, k=2) nn_dists = dists[:, 1] adaptive_r = float(np.median(nn_dists)) * 1.5 radius_frac = max(adaptive_r, self.BLOB_RADIUS) else: radius_frac = self.BLOB_RADIUS # Core boolean mask core = np.zeros((res, res, res), dtype=bool) r_cells = int(np.ceil(radius_frac * (res - 1))) for p in range(len(norm)): ci = int(round(norm[p, 0] * (res - 1))) cj = int(round(norm[p, 1] * (res - 1))) ck = int(round(norm[p, 2] * (res - 1))) i0, i1 = max(0, ci - r_cells), min(res, ci + r_cells + 1) j0, j1 = max(0, cj - r_cells), min(res, cj + r_cells + 1) k0, k1 = max(0, ck - r_cells), min(res, ck + r_cells + 1) core[i0:i1, j0:j1, k0:k1] = True # Ensure one connected component — dilate until all islands merge labeled, n_components = label(core) max_dilations = res // 2 # safety cap dilations = 0 while n_components > 1 and dilations < max_dilations: core = binary_dilation(core) labeled, n_components = label(core) dilations += 1 # Build opacity grid: core=1.0, margin=falloff, outside=0.0 margin = self.BLOB_MARGIN_CELLS opacity = np.zeros((res, res, res), dtype=np.float32) opacity[core] = 1.0 if margin > 0: dist = distance_transform_edt(~core).astype(np.float32) margin_mask = (dist > 0) & (dist <= margin) opacity[margin_mask] = 1.0 - dist[margin_mask] / (margin + 1) self._path_mask = core self._blob_opacity = opacity self._path_mask_norm = norm def _sample_blob_opacity(self, x, y, z) -> float: """Sample the blob opacity at a world-space point.""" if self._blob_opacity is None: return 1.0 r = self._mask_res nx = (x - self.axis_min[0]) / self.span[0] if self.span[0] > 0 else 0.5 ny = (y - self.axis_min[1]) / self.span[1] if self.span[1] > 0 else 0.5 nz = (z - self.axis_min[2]) / self.span[2] if self.span[2] > 0 else 0.5 ix = max(0, min(r - 1, int(round(nx * (r - 1))))) iy = max(0, min(r - 1, int(round(ny * (r - 1))))) iz = max(0, min(r - 1, int(round(nz * (r - 1))))) return float(self._blob_opacity[ix, iy, iz]) def _point_in_blob_world(self, x, y, z): """Check if a world-space point is inside the path blob (core or margin).""" if self._blob_opacity is None: return True return self._sample_blob_opacity(x, y, z) > 0.0 def _check_outside_blob_batch(self, positions: np.ndarray) -> np.ndarray: """Vectorized check: returns boolean mask of particles outside the blob.""" r = self._mask_res norm = np.zeros_like(positions) for a in range(3): if self.span[a] > 0: norm[:, a] = (positions[:, a] - self.axis_min[a]) / self.span[a] # Convert to grid indices ix = np.clip(np.round(norm[:, 0] * (r - 1)).astype(np.int32), 0, r - 1) iy = np.clip(np.round(norm[:, 1] * (r - 1)).astype(np.int32), 0, r - 1) iz = np.clip(np.round(norm[:, 2] * (r - 1)).astype(np.int32), 0, r - 1) return self._blob_opacity[ix, iy, iz] <= 0.0 def _init_lattice(self): """Initialize particles, constrained to path blob if available. Particles spawned in the margin zone get a baked blob_alpha < 1.0. """ margin = 0.05 g = self.particle_grid self.blob_alpha = np.ones(self.particle_count, dtype=np.float32) if self._blob_opacity is None: # Original bounding-box lattice idx = 0 for i in range(g): for j in range(g): for k in range(g): fx = i / (g - 1) if g > 1 else 0.5 fy = j / (g - 1) if g > 1 else 0.5 fz = k / (g - 1) if g > 1 else 0.5 self.pos[idx, 0] = self.axis_min[0] + (margin + fx * (1 - 2 * margin)) * self.span[0] self.pos[idx, 1] = self.axis_min[1] + (margin + fy * (1 - 2 * margin)) * self.span[1] self.pos[idx, 2] = self.axis_min[2] + (margin + fz * (1 - 2 * margin)) * self.span[2] self.age[idx] = 0 idx += 1 else: # Blob-constrained: uniform random sampling inside core + margin idx = 0 rng = np.random.default_rng(42) batch = max(self.particle_count * 4, 10000) while idx < self.particle_count: candidates = np.empty((batch, 3), dtype=np.float32) for a in range(3): lo = self.axis_min[a] + margin * self.span[a] hi = self.axis_max[a] - margin * self.span[a] candidates[:, a] = rng.uniform(lo, hi, batch).astype(np.float32) for c in range(len(candidates)): if idx >= self.particle_count: break x, y, z = candidates[c] opa = self._sample_blob_opacity(x, y, z) if opa > 0.0: self.pos[idx] = candidates[c] self.blob_alpha[idx] = opa self.age[idx] = 0 idx += 1 def _dry_integrate(self, steps: int): """Pre-integrate without recording, to break lattice regularity. Matches Android's dryIntegrate(): advects and wraps only, does NOT increment age (ages are randomized after). """ for _ in range(steps): velocities = self.sample_velocity_batch(self.pos) self.pos += velocities * self.DT self._wrap_all() def sample_velocity(self, x: float, y: float, z: float) -> np.ndarray: """Trilinear interpolation on the velocity grid. Ported from FlowFieldSystem.java VelocityField.sample(). """ G = self.grid_size # Normalize to grid coordinates [0, G-1] nx = (x - self.axis_min[0]) / self.span[0] * (G - 1) if self.span[0] > 0 else 0 ny = (y - self.axis_min[1]) / self.span[1] * (G - 1) if self.span[1] > 0 else 0 nz = (z - self.axis_min[2]) / self.span[2] * (G - 1) if self.span[2] > 0 else 0 ix = int(np.floor(nx)) iy = int(np.floor(ny)) iz = int(np.floor(nz)) tx = nx - ix ty = ny - iy tz = nz - iz ix = np.clip(ix, 0, G - 2) iy = np.clip(iy, 0, G - 2) iz = np.clip(iz, 0, G - 2) result = np.zeros(3, dtype=np.float32) for dx in range(2): for dy in range(2): for dz in range(2): w = ((1 - tx) if dx == 0 else tx) * \ ((1 - ty) if dy == 0 else ty) * \ ((1 - tz) if dz == 0 else tz) result += w * self.velocity_grid[ix + dx, iy + dy, iz + dz] return result def sample_velocity_batch(self, positions: np.ndarray, out: Optional[np.ndarray] = None) -> np.ndarray: """Trilinear interpolation for all particles using scipy's C backend. Args: positions: (N, 3) array of particle positions. out: Optional pre-allocated (N, 3) output buffer. Returns: (N, 3) array of velocities. """ from scipy.ndimage import map_coordinates G = self.grid_size N = len(positions) if out is not None and out.shape == (N, 3): result = out else: result = np.zeros((N, 3), dtype=np.float32) # Normalize to grid coordinates (reuse buffer if possible) if N == self.particle_count: normalized = self._buf_normalized else: normalized = np.zeros_like(positions) for a in range(3): if self.span[a] > 0: normalized[:, a] = (positions[:, a] - self.axis_min[a]) / self.span[a] * (G - 1) # scipy's map_coordinates expects (ndim, N_points) coordinate array coords = normalized.T # (3, N) — no copy, just transpose view # Interpolate each velocity component via C-implemented trilinear for c in range(3): result[:, c] = map_coordinates( self.velocity_grid[:, :, :, c], coords, order=1, mode='nearest', ) return result def _wrap_all(self): """Apply periodic wrapping to all particles.""" for a in range(3): if self.span[a] <= 0: continue dist = self.pos[:, a] - self.axis_min[a] dist = dist % self.span[a] self.pos[:, a] = self.axis_min[a] + dist def _respawn(self, indices: np.ndarray): """Respawn particles at their pre-integrated origin positions. Matches Android: age[i] = 0; pos[i] = origPos[i]. """ self.pos[indices] = self.orig_pos[indices] self.age[indices] = 0 def step(self, dt_scale: float = 1.0) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Advance all particles by one frame. Uses pre-allocated buffers to avoid per-frame memory allocation. Args: dt_scale: Multiplier for the time step (1.0 = normal 60 Hz frame). Use >1 to compensate for lower frame rates. Returns: (positions, colors, alphas, speeds) — all (N, ...) arrays. positions is a copy; colors/alphas/speeds are internal buffers valid until the next step() call. """ vel = self._buf_velocities speeds = self._buf_speeds colors = self._buf_colors alphas = self._buf_alphas # Compute velocities into pre-allocated buffer self.sample_velocity_batch(self.pos, out=vel) # Compute speeds without allocating new array np.einsum('ij,ij->i', vel, vel, out=speeds) np.sqrt(speeds, out=speeds) # Advect (dt_scale compensates for variable frame rate) self.pos += vel * (self.DT * self.speed_multiplier * dt_scale) self._wrap_all() # Kill particles that escaped the blob (vectorized check) if self._blob_opacity is not None and self.blob_enabled: escaped = self._check_outside_blob_batch(self.pos) if np.any(escaped): self.age[escaped] = self.LIFESPAN # force respawn # Age and respawn self.age += 1 expired = self.age >= self.LIFESPAN if np.any(expired): self._respawn(np.where(expired)[0]) # Speed-based colors (Turbo colormap) into pre-allocated buffer max_speed = speeds.max() if speeds.max() > 0 else 1.0 turbo_colormap(speeds / max_speed, out=colors) # Alpha based on age (fade in/out) into pre-allocated buffer age_frac = self.age.astype(np.float32) / self.LIFESPAN alphas[:] = 1.0 fade_in = age_frac < 0.1 alphas[fade_in] = age_frac[fade_in] / 0.1 fade_out = age_frac > 0.8 alphas[fade_out] = (1.0 - age_frac[fade_out]) / 0.2 np.clip(alphas, 0, 1, out=alphas) alphas *= self.blob_alpha # Copy positions (renderer needs stable reference) np.copyto(self._buf_pos_out, self.pos) return self._buf_pos_out, colors, alphas, speeds def is_outside_blob(self, x, y, z) -> bool: """Check if a world-space point is outside the path blob.""" return self._path_mask is not None and not self._point_in_blob_world(x, y, z) def get_blob_surface_points(self) -> Optional[np.ndarray]: """Return world-space points on the blob surface for debug viz. Samples all mask-True cells and returns their centres. """ if self._path_mask is None: return None r = self._mask_res indices = np.argwhere(self._path_mask) # (N, 3) of (i,j,k) if len(indices) == 0: return None # Convert grid indices to world coords pts = np.zeros((len(indices), 3), dtype=np.float32) for a in range(3): pts[:, a] = self.axis_min[a] + (indices[:, a] / (r - 1)) * self.span[a] return pts def advance_ball(self, dt_scale: float = 1.0) -> np.ndarray: """Advance the ball probe by one step along the flow field. Uses 0.8 * DT for slower movement (matching Android's dtSlow). Decelerates heavily when outside the path blob. Args: dt_scale: Multiplier for the time step (1.0 = normal 60 Hz frame). Returns: New ball position (3,) array. """ v = self.sample_velocity(*self.ball_pos) dt_slow = self.DT * 0.8 * self.speed_multiplier * dt_scale # Slow down to 10% speed when outside blob if self.is_outside_blob(*self.ball_pos): dt_slow *= 0.1 self.ball_pos = self.ball_pos + v * dt_slow # Clamp to domain self.ball_pos = np.clip(self.ball_pos, self.axis_min, self.axis_max) # Update trail (max 100 points, matching Android MAX_TRAIL) self.ball_trail.append(self.ball_pos.copy()) if len(self.ball_trail) > 100: self.ball_trail.pop(0) return self.ball_pos.copy() def set_ball_position(self, x: float, y: float, z: float): """Set ball position manually (from slider interaction).""" self.ball_pos = np.array([x, y, z], dtype=np.float32) def set_particle_grid(self, new_grid: int): """Reinitialize particles with a new grid resolution.""" self.particle_grid = new_grid self.particle_count = new_grid ** 3 N = self.particle_count self.pos = np.zeros((N, 3), dtype=np.float32) self.age = np.zeros(N, dtype=np.int32) # Reallocate step() buffers for new particle count self._buf_velocities = np.zeros((N, 3), dtype=np.float32) self._buf_speeds = np.zeros(N, dtype=np.float32) self._buf_colors = np.zeros((N, 3), dtype=np.float32) self._buf_alphas = np.ones(N, dtype=np.float32) self._buf_pos_out = np.zeros((N, 3), dtype=np.float32) self._buf_normalized = np.zeros((N, 3), dtype=np.float32) self._init_lattice() self.orig_pos = self.pos.copy() # save BEFORE pre-integration self._dry_integrate(self.PRE_AGE) if self._blob_opacity is not None and self.blob_enabled: for i in range(self.particle_count): if not self._point_in_blob_world(*self.pos[i]): self.pos[i] = self.orig_pos[i] self.age = np.random.randint(0, self.LIFESPAN, self.particle_count, dtype=np.int32) def set_path_points(self, path_points: np.ndarray): """Update the path mask and reinitialize particles.""" if path_points is not None and len(path_points) >= 2: self._build_path_mask(path_points) else: self._path_mask = None self._blob_opacity = None # Reinitialize with current grid size self.pos = np.zeros((self.particle_count, 3), dtype=np.float32) self.age = np.zeros(self.particle_count, dtype=np.int32) self._init_lattice() self.orig_pos = self.pos.copy() # save BEFORE pre-integration self._dry_integrate(self.PRE_AGE) if self._blob_opacity is not None and self.blob_enabled: for i in range(self.particle_count): if not self._point_in_blob_world(*self.pos[i]): self.pos[i] = self.orig_pos[i] self.age = np.random.randint(0, self.LIFESPAN, self.particle_count, dtype=np.int32) def start_ball_flow(self): """Start ball following the flow field.""" self.ball_flowing = True self.ball_trail.clear() self.ball_trail.append(self.ball_pos.copy()) def stop_ball_flow(self): """Stop ball following the flow field.""" self.ball_flowing = False def build_score_grid(self, data_points: np.ndarray, data_scores: np.ndarray): """Precompute a 3D score grid for fast particle-score lookup. Each grid cell gets the score of the nearest data point. Uses the same grid resolution as the velocity field. Args: data_points: (M, 3) data point positions. data_scores: (M,) score values per data point. """ G = self.grid_size # Build grid cell centers in world coords grid_coords = np.zeros((G, G, G, 3), dtype=np.float32) for a in range(3): linspace = np.linspace(self.axis_min[a], self.axis_max[a], G) if a == 0: grid_coords[:, :, :, 0] = linspace[:, None, None] elif a == 1: grid_coords[:, :, :, 1] = linspace[None, :, None] else: grid_coords[:, :, :, 2] = linspace[None, None, :] # Flatten to (G³, 3) and find nearest data point for each cell flat = grid_coords.reshape(-1, 3) # Chunked to avoid huge memory: process 1000 cells at a time score_flat = np.zeros(len(flat), dtype=np.float32) chunk = 2000 for i in range(0, len(flat), chunk): batch = flat[i:i + chunk] dists = np.linalg.norm(batch[:, None, :] - data_points[None, :, :], axis=2) nearest = np.argmin(dists, axis=1) score_flat[i:i + chunk] = data_scores[nearest] self._score_grid = score_flat.reshape(G, G, G) def _sample_scalar_grid(self, grid: np.ndarray, positions: np.ndarray) -> np.ndarray: """Trilinear interpolation on a scalar (G,G,G) grid using scipy C backend.""" from scipy.ndimage import map_coordinates normalized = np.zeros_like(positions) for a in range(3): if self.span[a] > 0: normalized[:, a] = ((positions[:, a] - self.axis_min[a]) / self.span[a] * (grid.shape[0] - 1)) return map_coordinates(grid, normalized.T, order=1, mode='nearest' ).astype(np.float32) def sample_score_batch(self, positions: np.ndarray) -> np.ndarray: """Sample precomputed score grid at particle positions. Returns (N,) interpolated score values, or 0.5 if no score grid. """ if self._score_grid is None: return np.full(len(positions), 0.5, dtype=np.float32) return self._sample_scalar_grid(self._score_grid, positions) def sample_confidence_batch(self, positions: np.ndarray) -> np.ndarray: """Sample confidence grid at particle positions. Returns (N,) values in [0, 1] where 1 = high confidence. Returns all-ones if no confidence grid. """ if self._confidence_grid is None: return np.ones(len(positions), dtype=np.float32) return self._sample_scalar_grid(self._confidence_grid, positions) # ═══════════════════════════════════════════════════ # Attractor detection # ═══════════════════════════════════════════════════ # Bump this version when the attractor detection algorithm changes # to auto-invalidate cached results. _ATTRACTOR_CACHE_VERSION = 18 @staticmethod def _sensitivity_params(s: float) -> dict: """Map sensitivity (0.0–1.0) to attractor detection parameters. At s=0.7 (default) the values match the hardcoded defaults. Lower sensitivity shrinks large basins first (via max_basin_frac and basin_thresh_pct), then removes weak attractors (peak_floor_pct). """ s = max(0.0, min(1.0, s)) # Interpolation helper: s=0 → val_lo, s=0.7 → val_mid, s=1.0 → val_hi def _lerp(val_lo, val_mid, val_hi): if s <= 0.7: t = s / 0.7 return val_lo + t * (val_mid - val_lo) else: t = (s - 0.7) / 0.3 return val_mid + t * (val_hi - val_mid) # max_basin_frac uses a power curve for steeper drop at low # sensitivity — large basins shrink first, small ones are spared. t_basin = s / 0.7 if s <= 0.7 else 1.0 + (s - 0.7) / 0.3 if s <= 0.7: # Quadratic: drops fast at first, then flattens near default max_bf = 0.002 + (0.02 - 0.002) * (s / 0.7) ** 1.5 else: max_bf = 0.02 + (0.05 - 0.02) * ((s - 0.7) / 0.3) return { 'peak_floor_pct': _lerp(0.18, 0.05, 0.01), 'basin_thresh_pct': _lerp(0.70, 0.35, 0.12), 'max_basin_frac': max_bf, 'speed_gate_mult': _lerp(1.2, 2.0, 4.0), } def _attractor_fingerprint(self, sensitivity: float = 0.7) -> str: """Build a fingerprint from the velocity grid + sensitivity.""" import hashlib h = hashlib.sha256() h.update(f"v{self._ATTRACTOR_CACHE_VERSION}".encode()) h.update(self.velocity_grid.tobytes()[:4096]) # first 4KB is enough h.update(f"|G={self.grid_size}|s={sensitivity:.3f}|".encode()) return h.hexdigest()[:16] def save_attractors(self, path: str, attractors: list, sensitivity: float = 0.7): """Cache attractor results to disk as .npz for instant reload. Args: path: Base path (without extension). Creates {path}_attractors.npz. """ import os os.makedirs(os.path.dirname(path) if os.path.dirname(path) else '.', exist_ok=True) npz_path = path + "_attractors.npz" arrays = {} meta_list = [] for i, att in enumerate(attractors): arrays[f"basin_{i}"] = att['basin_mask'] arrays[f"position_{i}"] = att['position'] meta_list.append({ 'strength': att['strength'], 'divergence': att['divergence'], 'basin_size': att['basin_size'], 'basin_fraction': att['basin_fraction'], 'mean_score': att['mean_score'], }) import json as _json meta_obj = { 'attractors': meta_list, 'fingerprint': self._attractor_fingerprint(sensitivity), } arrays['_meta'] = np.array([_json.dumps(meta_obj)], dtype=object) np.savez_compressed(npz_path, **arrays) print(f"[ATTRACTORS] Cached {len(attractors)} attractor(s) to {npz_path}") def load_attractors(self, path: str, sensitivity: float = 0.7) -> Optional[list]: """Load cached attractor results from disk. Validates the fingerprint against the current velocity grid so stale caches are automatically discarded when data changes. Args: path: Base path (without extension). Looks for {path}_attractors.npz. Returns: List of attractor dicts, or None if cache missing/stale. """ npz_path = path + "_attractors.npz" import os if not os.path.exists(npz_path): return None try: data = np.load(npz_path, allow_pickle=True) import json as _json meta_obj = _json.loads(str(data['_meta'][0])) # Support both old format (list) and new format (dict with fingerprint) if isinstance(meta_obj, list): # Old cache without fingerprint — discard print("[ATTRACTORS] Old cache format without fingerprint — recomputing") return None meta_list = meta_obj['attractors'] cached_fp = meta_obj.get('fingerprint', '') # Validate fingerprint current_fp = self._attractor_fingerprint(sensitivity) if cached_fp != current_fp: print(f"[ATTRACTORS] Cache fingerprint mismatch — recomputing") return None attractors = [] for i, meta in enumerate(meta_list): attractors.append({ 'basin_mask': data[f'basin_{i}'], 'position': data[f'position_{i}'], 'strength': meta['strength'], 'divergence': meta['divergence'], 'basin_size': meta['basin_size'], 'basin_fraction': meta['basin_fraction'], 'mean_score': meta['mean_score'], }) print(f"[ATTRACTORS] Loaded {len(attractors)} cached attractor(s) from {npz_path}") return attractors except Exception as e: print(f"[ATTRACTORS] Cache load failed: {e}") return None def find_attractors(self, score_grid: Optional[np.ndarray] = None, min_confidence: float = 0.3, sim_steps: int = 800, sample_interval: int = 8, cache_path: Optional[str] = None, sensitivity: float = 0.7) -> list: """Detect flow attractors using residence density + local peak finding. Three-stage approach: 1. Simulate particles with renderer-compatible rules (wrapping, blob escape) but LONGER lifespan to reduce spawn bias. Accumulate occupancy = where particles spend time. 2. Find LOCAL maxima in occupancy (not global threshold) so weaker attractors aren't erased by dominant ones. 3. Validate each peak with negative divergence (field structure) so we're not just finding traffic jams. basin_mask = high-occupancy region around each peak (what gets drawn). Args: score_grid: Optional (G,G,G) score values for basin scoring. min_confidence: Minimum confidence to consider a cell. cache_path: If provided, try loading from / saving to this path. sim_steps: Total simulation steps. sample_interval: Sample occupancy every N steps. sensitivity: 0.0–1.0 controls attractor count & basin size. Default 0.6 matches hardcoded defaults. Lower = fewer, smaller basins (large ones shrink first). Higher = more, larger basins. """ # ── Derive parameters from sensitivity ── sp = self._sensitivity_params(sensitivity) # ── Try loading from cache first ── if cache_path is not None: cached = self.load_attractors(cache_path, sensitivity) if cached is not None: return cached from scipy import ndimage from scipy.ndimage import maximum_filter, label G = self.grid_size vg = self.velocity_grid # ── Blob mask resampled to velocity grid ── # Use the CORE mask (not the soft-margin opacity) so attractor # basins don't expand into the fade-out margin zone. blob_valid = np.ones((G, G, G), dtype=bool) if self._path_mask is not None: res = self._path_mask.shape[0] bi = np.round(np.arange(G) / (G - 1) * (res - 1)).astype(int) bi = np.clip(bi, 0, res - 1) blob_valid = self._path_mask[np.ix_(bi, bi, bi)] conf = self._confidence_grid if self._confidence_grid is not None \ else np.ones((G, G, G), dtype=np.float32) valid = (conf >= min_confidence) & blob_valid if np.sum(valid) < 4: return [] # ═══════════════════════════════════════════════════ # STAGE 1: Simulate particles, measure occupancy # Uses renderer rules but LONGER lifespan (5× normal) # to let particles reach distant attractors. # ═══════════════════════════════════════════════════ LONG_LIFESPAN = self.LIFESPAN * 5 # 285 frames vs 57 N_test = min(20000, max(5000, int(np.sum(valid)) * 8)) rng = np.random.default_rng(42) # Seed inside blob positions = np.zeros((N_test, 3), dtype=np.float32) idx_p = 0 batch = 2000 while idx_p < N_test: candidates = np.zeros((batch, 3), dtype=np.float32) for a in range(3): candidates[:, a] = rng.uniform( self.axis_min[a], self.axis_max[a], batch ).astype(np.float32) for c in range(len(candidates)): if idx_p >= N_test: break if self._point_in_blob_world(*candidates[c]): positions[idx_p] = candidates[c] idx_p += 1 # Pre-integrate for _ in range(self.PRE_AGE): vel = self.sample_velocity_batch(positions) positions += vel * self.DT for a in range(3): if self.span[a] > 0: positions[:, a] = self.axis_min[a] + \ (positions[:, a] - self.axis_min[a]) % self.span[a] orig_pos = positions.copy() ages = rng.integers(0, LONG_LIFESPAN, N_test, dtype=np.int32) # Run simulation, accumulate occupancy occupancy = np.zeros((G, G, G), dtype=np.float32) n_samples = 0 warmup = 150 # let particles spread out before counting for step_i in range(sim_steps): vel = self.sample_velocity_batch(positions) positions += vel * self.DT # Wrap at boundaries (same as MDN) for a in range(3): if self.span[a] > 0: positions[:, a] = self.axis_min[a] + \ (positions[:, a] - self.axis_min[a]) % self.span[a] # Blob escape → force respawn (same as MDN) if self._blob_opacity is not None: escaped = self._check_outside_blob_batch(positions) if np.any(escaped): ages[escaped] = LONG_LIFESPAN # Age and respawn with LONGER lifespan ages += 1 expired = ages >= LONG_LIFESPAN if np.any(expired): exp_idx = np.where(expired)[0] positions[exp_idx] = orig_pos[exp_idx] ages[exp_idx] = 0 # Sample occupancy after warmup if step_i >= warmup and step_i % sample_interval == 0: gi = np.zeros((N_test, 3), dtype=int) for a in range(3): gi[:, a] = np.clip( ((positions[:, a] - self.axis_min[a]) / self.span[a] * (G - 1)).astype(int), 0, G - 1) np.add.at(occupancy, (gi[:, 0], gi[:, 1], gi[:, 2]), 1.0) n_samples += 1 if n_samples == 0: return [] occupancy /= n_samples occupancy = ndimage.gaussian_filter(occupancy, sigma=0.8) occupancy[~valid] = 0.0 # Speed grid speed = np.linalg.norm(vg, axis=3) speed_ref = max(float(np.percentile(speed[valid], 90)), 1e-8) speed_factor = 1.0 / (1.0 + speed / speed_ref) # Divergence (for validation + info) dvx_dx = np.gradient(vg[:, :, :, 0], axis=0) dvy_dy = np.gradient(vg[:, :, :, 1], axis=1) dvz_dz = np.gradient(vg[:, :, :, 2], axis=2) divergence = dvx_dx + dvy_dy + dvz_dz # Combined residence score residence_score = occupancy * speed_factor residence_score[~valid] = 0.0 rs_max = float(np.max(residence_score)) if rs_max < 1e-10: return [] # ═══════════════════════════════════════════════════ # STAGE 2: Find LOCAL maxima (not global threshold) # ═══════════════════════════════════════════════════ neighborhood = 5 local_max = maximum_filter(residence_score, size=neighborhood) is_peak = (residence_score == local_max) & (residence_score > 0) # Absolute floor: peak_floor_pct of max OR 80th percentile (whichever lower) abs_floor = min( rs_max * sp['peak_floor_pct'], float(np.percentile(residence_score[valid], 80)) ) is_peak = is_peak & (residence_score > abs_floor) peak_coords = np.argwhere(is_peak) if len(peak_coords) == 0: return [] peak_scores = np.array([residence_score[tuple(p)] for p in peak_coords]) order = np.argsort(-peak_scores) peak_coords = peak_coords[order] peak_scores = peak_scores[order] # Reject peaks where the local flow is not actually converging. # A real attractor must have negative divergence in its neighborhood # (not just at the exact peak cell, which can be noisy). # Smooth divergence and require neighborhood mean < 0. div_smooth = ndimage.gaussian_filter(divergence, sigma=1.5) real_peaks = [] for pc in peak_coords: pi, pj, pk = int(pc[0]), int(pc[1]), int(pc[2]) # 3×3×3 neighborhood mean divergence sl = tuple(slice(max(0, c - 1), min(G, c + 2)) for c in (pi, pj, pk)) neigh_div = float(np.mean(div_smooth[sl])) if neigh_div < 0: real_peaks.append(pc) peak_coords = np.array(real_peaks) if real_peaks else np.empty((0, 3), dtype=int) peak_scores = np.array([residence_score[tuple(p)] for p in peak_coords]) \ if len(peak_coords) > 0 else np.array([]) if len(peak_coords) == 0: return [] # Merge peaks within 3 cells merge_dist = 3.0 kept = [] for pc in peak_coords: too_close = False for kc in kept: if np.linalg.norm(pc.astype(float) - kc.astype(float)) < merge_dist: too_close = True break if not too_close: kept.append(pc) peak_coords = np.array(kept) if kept else np.empty((0, 3), dtype=int) if len(peak_coords) == 0: return [] # ═══════════════════════════════════════════════════ # STAGE 3: Build basins using divergence-modulated expansion # # Key insight: "parking lots" (real sinks) have negative # divergence, while "rivers" (slow flow) have ~zero divergence. # Instead of separate hacks (absolute floor, divergence penalty), # we build a SINGLE basin expansion score that naturally gives # generous basins to real sinks and chokes off rivers: # # convergence = clamp(-div / div_ref, 0, 1) # basin_score = residence_score * (0.1 + 0.9 * convergence) # # Near a sink: convergence≈1 → basin_score ≈ residence_score # In a river: convergence≈0 → basin_score ≈ 0.1 * residence # Near a source: convergence=0 → basin_score ≈ 0.1 * residence # # Then flood-fill each peak's basin in basin_score space. # Rivers self-limit because their basin_score is 10× lower. # ═══════════════════════════════════════════════════ # Normalized divergence: scale so typical negative div → ~1.0 div_ref = max(float(np.percentile(np.abs(divergence[valid]), 90)), 1e-8) convergence = np.clip(-divergence / div_ref, 0.0, 1.0) # 1=sink, 0=source/neutral # Basin expansion score: occupancy weighted by convergence basin_score = residence_score * (0.1 + 0.9 * convergence) basin_score[~valid] = 0.0 bs_max = float(np.max(basin_score)) if bs_max < 1e-10: bs_max = rs_max # fallback # Absolute floor: median of valid basin_scores (or 8% of max). # This prevents shallow hills from claiming huge territories. # The local percentage handles sharp peaks; the floor handles flat ones. valid_bs = basin_score[valid] abs_basin_floor = max( float(np.median(valid_bs)), bs_max * 0.08, ) attractors_out = [] n_valid = int(np.sum(valid)) claimed = np.zeros((G, G, G), dtype=bool) for pc in peak_coords: pi, pj, pk = int(pc[0]), int(pc[1]), int(pc[2]) peak_val = float(residence_score[pi, pj, pk]) peak_bs = float(basin_score[pi, pj, pk]) peak_div = float(divergence[pi, pj, pk]) # ── Per-attractor speed gate ── # Basin mesh = where particles SETTLE. Only cells with speed # close to the peak's own (slow) speed belong in the basin. # Cells much faster are approach corridors, not settling zones. # Gate = 2× peak neighborhood speed, floored so we don't # choke basins where the entire field is uniformly slow. sl = tuple(slice(max(0, c - 1), min(G, c + 2)) for c in (pi, pj, pk)) peak_speed = float(np.mean(speed[sl])) # 3×3×3 neighborhood avg speed_gate = max(peak_speed * sp['speed_gate_mult'], speed_ref * 0.08) basin_slow = speed <= speed_gate # Basin threshold: HIGHER of local percentage and absolute floor. # - Sharp peaks: local 20% is high → reasonable basin size # - Shallow hills: local 20% is tiny → abs floor kicks in → tight basin local_thresh = max(peak_bs * sp['basin_thresh_pct'], abs_basin_floor) # Adaptive tightening: if basin exceeds 2% of valid cells, # raise threshold until it fits. This prevents flat, uniformly # convergent flow fields from producing basins that swallow # half the grid. max_basin_cells = max(int(n_valid * sp['max_basin_frac']), 20) for _tighten in range(8): basin_candidates = ( (basin_score >= local_thresh) & valid & (~claimed) & basin_slow ) basin_labeled, _ = label(basin_candidates) peak_label = basin_labeled[pi, pj, pk] if peak_label == 0: break basin_mask = basin_labeled == peak_label basin_size = int(np.sum(basin_mask)) if basin_size <= max_basin_cells: break # Tighten: raise threshold toward peak value local_thresh = local_thresh + (peak_bs - local_thresh) * 0.3 if peak_label == 0: continue basin_size = int(np.sum(basin_mask)) if basin_size < 5: continue claimed |= basin_mask position = np.array([ self.axis_min[0] + pi / (G - 1) * self.span[0], self.axis_min[1] + pj / (G - 1) * self.span[1], self.axis_min[2] + pk / (G - 1) * self.span[2], ], dtype=np.float32) basin_fraction = basin_size / n_valid if n_valid > 0 else 0.0 mean_score = None if score_grid is not None and score_grid.shape == (G, G, G): bs = score_grid[basin_mask] if len(bs) > 0: mean_score = float(np.mean(bs)) # Strength from basin_score (already incorporates divergence) attractors_out.append({ 'position': position, 'strength': float(peak_bs / bs_max) if bs_max > 0 else 0.0, 'divergence': peak_div, 'basin_mask': basin_mask, 'basin_size': basin_size, 'basin_fraction': basin_fraction, 'mean_score': mean_score, }) # ═══════════════════════════════════════════════════ # STAGE 4: Vortex attractor detection (supplementary) # # Occupancy-based detection (stages 1-3) misses vortex centers # because particles ORBIT around the center rather than sitting # on it. RBF flow fields often create vortices that are the # true global attractors (probes get trapped there forever). # # Detection: find grid cells where the flow is nearly zero but # the curl (vorticity) is high and divergence is negative # (spiral-in). Then verify with a short probe simulation that # a particle actually gets trapped. # ═══════════════════════════════════════════════════ # Curl magnitude dvz_dy = np.gradient(vg[:, :, :, 2], axis=1) dvy_dz = np.gradient(vg[:, :, :, 1], axis=2) dvx_dz = np.gradient(vg[:, :, :, 0], axis=2) dvz_dx = np.gradient(vg[:, :, :, 2], axis=0) dvy_dx = np.gradient(vg[:, :, :, 1], axis=0) dvx_dy = np.gradient(vg[:, :, :, 0], axis=1) curl_mag = np.sqrt( (dvz_dy - dvy_dz) ** 2 + (dvx_dz - dvz_dx) ** 2 + (dvy_dx - dvx_dy) ** 2 ) # Vortex score: high curl, low speed, negative divergence speed_30 = max(float(np.percentile(speed[valid], 30)), 1e-8) slow_mask = valid & (speed < speed_30) & (divergence < 0) vortex_score = np.zeros_like(speed) vortex_score[slow_mask] = ( curl_mag[slow_mask] / (speed[slow_mask] + 1e-6) * (-divergence[slow_mask]) ) if float(np.max(vortex_score)) > 0: # Find vortex peaks vortex_local_max = maximum_filter(vortex_score, size=5) is_vortex = (vortex_score == vortex_local_max) & (vortex_score > 0) vortex_peaks_arr = np.argwhere(is_vortex) if len(vortex_peaks_arr) > 0: vp_scores = np.array([vortex_score[tuple(p)] for p in vortex_peaks_arr]) vp_order = np.argsort(-vp_scores) # Only consider strong vortex candidates (top 10%, at least score > 1.0) vp_thresh = max(float(np.percentile(vp_scores, 90)), 1.0) n_vortex_added = 0 for vpi in vp_order: if vp_scores[vpi] < vp_thresh: break if n_vortex_added >= 3: # max 3 vortex attractors break vpc = vortex_peaks_arr[vpi] vi, vj, vk = int(vpc[0]), int(vpc[1]), int(vpc[2]) # Skip if too close to an existing attractor (within 4 cells) vpos_world = np.array([ self.axis_min[0] + vi / (G - 1) * self.span[0], self.axis_min[1] + vj / (G - 1) * self.span[1], self.axis_min[2] + vk / (G - 1) * self.span[2], ], dtype=np.float32) too_close = False for existing in attractors_out: ep = existing['position'] # Convert to grid distance edist = np.linalg.norm( (vpos_world - ep) / self.span * (G - 1) ) if edist < 4: too_close = True break if too_close: continue # Verify: simulate a probe from nearby and check it gets trapped. # "Trapped" means the probe stays confined in a small region — # either reaching near-zero speed OR orbiting tightly (vortex). # We check confinement by comparing the bounding box of the # last 100 steps to the grid cell size. cell_size = float(np.mean(self.span / (G - 1))) confinement_radius = cell_size * 3 # trapped if last 100 steps within 3 cells def _probe_is_trapped(start_pos, n_steps=1000): """Simulate probe and check if it gets confined.""" p = start_pos.copy() history = [] for step in range(n_steps): pv = self.sample_velocity(p[0], p[1], p[2]) ps = float(np.linalg.norm(pv)) if ps < 1e-7: return True, p # dead stop p = p + pv * self.DT p = np.clip(p, self.axis_min, self.axis_max) if step >= n_steps - 200: history.append(p.copy()) if len(history) < 50: return False, p hist = np.array(history) extent = hist.max(axis=0) - hist.min(axis=0) max_extent = float(np.max(extent)) return max_extent < confinement_radius, p probe_start = vpos_world.copy() + self.span / (G - 1) * 3 probe_start = np.clip(probe_start, self.axis_min, self.axis_max) trapped, final_probe = _probe_is_trapped(probe_start) if not trapped: # Try from opposite side probe_start2 = vpos_world.copy() - self.span / (G - 1) * 3 probe_start2 = np.clip(probe_start2, self.axis_min, self.axis_max) trapped, final_probe = _probe_is_trapped(probe_start2) if not trapped: continue # Use the center of the final orbit as the attractor position # (run another short sim to get the centroid) centroid_pts = [] p = final_probe.copy() for _ in range(200): pv = self.sample_velocity(p[0], p[1], p[2]) p = p + pv * self.DT p = np.clip(p, self.axis_min, self.axis_max) centroid_pts.append(p.copy()) final_pos = np.mean(centroid_pts, axis=0).astype(np.float32) # Build basin using the SAME Stage 3 flood-fill logic # so vortex basins have organic shapes, not cubes. fi = int(np.clip( (final_pos[0] - self.axis_min[0]) / self.span[0] * (G - 1) + 0.5, 0, G - 1)) fj = int(np.clip( (final_pos[1] - self.axis_min[1]) / self.span[1] * (G - 1) + 0.5, 0, G - 1)) fk = int(np.clip( (final_pos[2] - self.axis_min[2]) / self.span[2] * (G - 1) + 0.5, 0, G - 1)) # Speed gate: same as Stage 3 sl = tuple(slice(max(0, c - 1), min(G, c + 2)) for c in (fi, fj, fk)) peak_spd = float(np.mean(speed[sl])) spd_gate = max(peak_spd * sp['speed_gate_mult'], speed_ref * 0.08) basin_slow = speed <= spd_gate peak_bs_v = float(basin_score[fi, fj, fk]) local_thresh_v = max(peak_bs_v * sp['basin_thresh_pct'], abs_basin_floor) max_basin_cells = max(int(n_valid * sp['max_basin_frac']), 20) basin_mask = np.zeros((G, G, G), dtype=bool) for _tighten in range(8): basin_candidates = ( (basin_score >= local_thresh_v) & valid & (~claimed) & basin_slow ) basin_labeled_v, _ = label(basin_candidates) peak_label_v = basin_labeled_v[fi, fj, fk] if peak_label_v == 0: break basin_mask = basin_labeled_v == peak_label_v bsz = int(np.sum(basin_mask)) if bsz <= max_basin_cells: break local_thresh_v = local_thresh_v + ( peak_bs_v - local_thresh_v) * 0.3 basin_size = int(np.sum(basin_mask)) if basin_size < 3: # Fallback: very small basin from nearby convergent cells basin_mask = np.zeros((G, G, G), dtype=bool) for di in range(-1, 2): for dj in range(-1, 2): for dk in range(-1, 2): ni, nj, nk = fi + di, fj + dj, fk + dk if 0 <= ni < G and 0 <= nj < G and 0 <= nk < G: if (valid[ni, nj, nk] and not claimed[ni, nj, nk] and divergence[ni, nj, nk] < 0): basin_mask[ni, nj, nk] = True basin_size = int(np.sum(basin_mask)) if basin_size < 3: continue claimed |= basin_mask basin_fraction = basin_size / n_valid if n_valid > 0 else 0.0 mean_score = None if score_grid is not None and score_grid.shape == (G, G, G): bs_vals = score_grid[basin_mask] if len(bs_vals) > 0: mean_score = float(np.mean(bs_vals)) # Strength: verified vortex attractors get a floor of 0.5 occ_at_vortex = float(residence_score[fi, fj, fk]) vortex_strength = max( occ_at_vortex / rs_max if rs_max > 0 else 0.0, 0.5 ) attractors_out.append({ 'position': final_pos, 'strength': vortex_strength, 'divergence': float(divergence[fi, fj, fk]), 'basin_mask': basin_mask, 'basin_size': basin_size, 'basin_fraction': basin_fraction, 'mean_score': mean_score, }) n_vortex_added += 1 if n_vortex_added > 0: print(f"[ATTRACTORS] Added {n_vortex_added} vortex attractor(s)") # ── Merge encapsulated attractors ── # When one attractor's bounding box is contained inside another's, # they look like nested shells. Merge by absorbing the weaker # one's basin into the stronger one. merged = True while merged: merged = False for i in range(len(attractors_out)): if attractors_out[i] is None: continue bi = attractors_out[i]['basin_mask'] ci = np.argwhere(bi) if len(ci) == 0: continue mni, mxi = ci.min(0), ci.max(0) for j in range(len(attractors_out)): if i == j or attractors_out[j] is None: continue bj = attractors_out[j]['basin_mask'] cj = np.argwhere(bj) if len(cj) == 0: continue mnj, mxj = cj.min(0), cj.max(0) # Check if j's bbox is inside i's bbox (or vice versa) j_inside_i = np.all(mnj >= mni) and np.all(mxj <= mxi) i_inside_j = np.all(mni >= mnj) and np.all(mxi <= mxj) if j_inside_i or i_inside_j: # Keep the stronger one, absorb the other's basin si = attractors_out[i]['strength'] sj = attractors_out[j]['strength'] keep, drop = (i, j) if si >= sj else (j, i) attractors_out[keep]['basin_mask'] = ( attractors_out[keep]['basin_mask'] | attractors_out[drop]['basin_mask'] ) attractors_out[keep]['basin_size'] = int( np.sum(attractors_out[keep]['basin_mask']) ) attractors_out[keep]['basin_fraction'] = ( attractors_out[keep]['basin_size'] / n_valid if n_valid > 0 else 0.0 ) attractors_out[drop] = None merged = True break if merged: break attractors_out = [a for a in attractors_out if a is not None] # ── Probe escape validation ── # Drop any attractor where a probe starting at its position escapes # beyond a threshold. Real attractors hold probes nearby; false # positives from occupancy noise let them fly away. cell_diag = float(np.linalg.norm(self.span / (G - 1))) escape_thresh = cell_diag * 5 # must stay within 5 cell diagonals validated = [] for att in attractors_out: p = att['position'].copy().astype(np.float32) start = p.copy() for _ in range(500): pv = self.sample_velocity(p[0], p[1], p[2]) if np.linalg.norm(pv) < 1e-8: break p = p + pv * self.DT p = np.clip(p, self.axis_min, self.axis_max) drift = float(np.linalg.norm(p - start)) if drift < escape_thresh: validated.append(att) attractors_out = validated # ── Probe-based convergence merge ── # Simulate a short probe from each attractor. If two probes end # up at essentially the same point they are duplicate detections # of the same convergence structure. Keep it short (300 steps) so # we only merge genuinely redundant peaks without traversing across # distinct basins in smooth (MDN) fields. merge_radius = cell_diag * 3 for att in attractors_out: p = att['position'].copy().astype(np.float32) for _ in range(300): pv = self.sample_velocity(p[0], p[1], p[2]) if np.linalg.norm(pv) < 1e-8: break p = p + pv * self.DT p = np.clip(p, self.axis_min, self.axis_max) att['_converged_to'] = p.copy() merged = True while merged: merged = False for i in range(len(attractors_out)): if attractors_out[i] is None: continue ci = attractors_out[i]['_converged_to'] for j in range(i + 1, len(attractors_out)): if attractors_out[j] is None: continue cj = attractors_out[j]['_converged_to'] dist = float(np.linalg.norm(ci - cj)) if dist < merge_radius: si = attractors_out[i]['strength'] sj = attractors_out[j]['strength'] keep, drop = (i, j) if si >= sj else (j, i) attractors_out[keep]['basin_mask'] = ( attractors_out[keep]['basin_mask'] | attractors_out[drop]['basin_mask'] ) attractors_out[keep]['basin_size'] = int( np.sum(attractors_out[keep]['basin_mask']) ) attractors_out[keep]['basin_fraction'] = ( attractors_out[keep]['basin_size'] / n_valid if n_valid > 0 else 0.0 ) attractors_out[drop] = None merged = True break if merged: break attractors_out = [a for a in attractors_out if a is not None] # Clean up temp field for att in attractors_out: att.pop('_converged_to', None) # ── Connectivity cleanup ── # After merges, a basin_mask may have disconnected components. # Keep only the connected component containing the attractor center. for att in attractors_out: bm = att['basin_mask'] labeled_bm, n_comp = label(bm) if n_comp <= 1: continue pos = att['position'] ci = int(np.clip( (pos[0] - self.axis_min[0]) / self.span[0] * (G - 1) + 0.5, 0, G - 1)) cj = int(np.clip( (pos[1] - self.axis_min[1]) / self.span[1] * (G - 1) + 0.5, 0, G - 1)) ck = int(np.clip( (pos[2] - self.axis_min[2]) / self.span[2] * (G - 1) + 0.5, 0, G - 1)) center_comp = labeled_bm[ci, cj, ck] if center_comp > 0: att['basin_mask'] = labeled_bm == center_comp else: # Center not inside any component — keep largest comp_sizes = [(c, int(np.sum(labeled_bm == c))) for c in range(1, n_comp + 1)] largest = max(comp_sizes, key=lambda x: x[1])[0] att['basin_mask'] = labeled_bm == largest att['basin_size'] = int(np.sum(att['basin_mask'])) att['basin_fraction'] = (att['basin_size'] / n_valid if n_valid > 0 else 0.0) # ── Convergence-point patching ── # The basin is built around the occupancy peak, but probes may # converge to a slightly different point. Patch the basin to # include a small neighborhood of convergent cells around the # actual convergence point so the probe rests inside the mesh. for att in attractors_out: p = att['position'].copy().astype(np.float32) for _ in range(500): pv = self.sample_velocity(p[0], p[1], p[2]) if np.linalg.norm(pv) < 1e-8: break p = p + pv * self.DT p = np.clip(p, self.axis_min, self.axis_max) # Grid index of convergence point ei = int(np.clip( (p[0] - self.axis_min[0]) / self.span[0] * (G - 1) + 0.5, 0, G - 1)) ej = int(np.clip( (p[1] - self.axis_min[1]) / self.span[1] * (G - 1) + 0.5, 0, G - 1)) ek = int(np.clip( (p[2] - self.axis_min[2]) / self.span[2] * (G - 1) + 0.5, 0, G - 1)) if not att['basin_mask'][ei, ej, ek]: # Convergence point outside basin — add a small patch # of convergent cells (negative divergence) around it. for di in range(-1, 2): for dj in range(-1, 2): for dk in range(-1, 2): ni = ei + di nj = ej + dj nk = ek + dk if 0 <= ni < G and 0 <= nj < G and 0 <= nk < G: if valid[ni, nj, nk] and divergence[ni, nj, nk] < 0: att['basin_mask'][ni, nj, nk] = True att['basin_size'] = int(np.sum(att['basin_mask'])) att['basin_fraction'] = (att['basin_size'] / n_valid if n_valid > 0 else 0.0) attractors_out.sort(key=lambda a: a['strength'], reverse=True) # Cap at 8, but always keep at least the strongest one attractors_out = attractors_out[:8] # Guarantee at least 1 attractor: if all were filtered, take the # strongest peak and give it a minimal basin if not attractors_out and len(peak_coords) > 0: pc = peak_coords[0] pi, pj, pk = int(pc[0]), int(pc[1]), int(pc[2]) basin_mask = np.zeros((G, G, G), dtype=bool) # 3×3×3 cube around peak for di in range(-1, 2): for dj in range(-1, 2): for dk in range(-1, 2): ni, nj, nk = pi+di, pj+dj, pk+dk if 0 <= ni < G and 0 <= nj < G and 0 <= nk < G: if valid[ni, nj, nk]: basin_mask[ni, nj, nk] = True position = np.array([ self.axis_min[0] + pi / (G - 1) * self.span[0], self.axis_min[1] + pj / (G - 1) * self.span[1], self.axis_min[2] + pk / (G - 1) * self.span[2], ], dtype=np.float32) attractors_out.append({ 'position': position, 'strength': 1.0, 'divergence': float(divergence[pi, pj, pk]), 'basin_mask': basin_mask, 'basin_size': int(np.sum(basin_mask)), 'basin_fraction': int(np.sum(basin_mask)) / n_valid if n_valid > 0 else 0.0, 'mean_score': None, }) print(f"[ATTRACTORS] {N_test} particles, {sim_steps} steps " f"(lifespan={LONG_LIFESPAN}), {n_samples} samples") print(f"[ATTRACTORS] Occupancy max={float(np.max(occupancy)):.1f}, " f"div_ref={div_ref:.4f}, abs_floor={abs_basin_floor:.3f}, " f"speed_ref={speed_ref:.4f}, " f"peaks={len(peak_coords)}, output={len(attractors_out)}") # ── Save to cache ── if cache_path is not None and attractors_out: try: self.save_attractors(cache_path, attractors_out, sensitivity) except Exception as e: print(f"[ATTRACTORS] Cache save failed: {e}") return attractors_out def precompute_frames(self, n_frames: int = 200) -> list: """Pre-compute N animation frames for efficient playback. Returns: List of (positions, colors, alphas) tuples. """ frames = [] for _ in range(n_frames): pos, colors, alphas, speeds = self.step() frames.append({ "positions": pos, "colors": colors, "alphas": alphas, "speeds": speeds, }) return frames def build_flow_figure( result, flow_system: FlowFieldSystem, frame_data: dict, show_data_points: bool = True, ) -> dict: """Build Plotly trace data for one flow frame. Returns a dict of trace data that can be added to a figure. """ from tracescope.visualization.scatter3d import CLUSTER_COLORS positions = frame_data["positions"] colors = frame_data["colors"] alphas = frame_data["alphas"] # Convert colors to Plotly format plotly_colors = [ f"rgba({int(c[0]*255)},{int(c[1]*255)},{int(c[2]*255)},{a:.2f})" for c, a in zip(colors, alphas) ] traces = [] # Flow particles traces.append(dict( type="scatter3d", x=positions[:, 0].tolist(), y=positions[:, 1].tolist(), z=positions[:, 2].tolist(), mode="markers", marker=dict( size=3, color=plotly_colors, opacity=1.0, # Per-particle alpha is in the color ), name="Flow particles", showlegend=False, hoverinfo="skip", )) return traces