""" AI Voxel — Minimal GPU Experiment Kit ===================================== A self-contained PyTorch implementation of the core XYFLOW pipeline: 1. Fourier-basis vector field F(h) = Σ_k a_k sin(ω_k·h) + b_k cos(ω_k·h) 2. RK4 ODE integration dh/dt = F(h) 3. Implicit boundary S(p) via small MLP 4. Transverse flux 𝔽(p) = ∇S · F(p) ← "missing information" Run on GPU: python ai_voxel_experiment.py Run on CPU: set CUDA_VISIBLE_DEVICES="" (or set device='cpu' below) Requires: torch (pip install torch) Optional: torchdiffeq (if you want to swap RK4 for adaptive solvers) """ import math import numpy as np import torch import torch.nn as nn # ────────────────────────────────────────────────────────────────────────── # Configuration — tweak these for small experiments # ────────────────────────────────────────────────────────────────────────── DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") D = 16 # latent dimension (keep small for fast experiments) K = 32 # number of Fourier modes DHIDDEN = 32 # boundary MLP hidden width STEPS = 8 # RK4 steps T_TOTAL = 0.5 # total integration time BATCH = 4096 # query points per batch SEED = 42 torch.manual_seed(SEED) # ════════════════════════════════════════════════════════════════════════ # 1. FOURIER VECTOR FIELD F(h) = Σ_k a_k sin(ω_k·h) + b_k cos(ω_k·h) # ════════════════════════════════════════════════════════════════════════ class FourierField(nn.Module): """ Distilled vector field in Fourier basis. All parameters are learnable so you can train the field to match a target dynamics, or just use random init for visual experiments. """ def __init__(self, dim: int = D, num_modes: int = K): super().__init__() self.D = dim self.K = num_modes # frequencies ω [K, D] self.omega = nn.Parameter(torch.randn(num_modes, dim) * 0.5) # amplitudes a [K, D] self.a = nn.Parameter(torch.randn(num_modes, dim) * 0.01) # amplitudes b [K, D] self.b = nn.Parameter(torch.randn(num_modes, dim) * 0.01) def forward(self, h: torch.Tensor) -> torch.Tensor: """ Evaluate F(h). h: [..., D] → F: [..., D] """ # projections ω_k · h → [..., K] proj = h @ self.omega.t() # [..., K] sin_p = torch.sin(proj) # [..., K] cos_p = torch.cos(proj) # [..., K] # F = sin_p @ a + cos_p @ b → [..., D] return sin_p @ self.a + cos_p @ self.b def num_params(self) -> int: return sum(p.numel() for p in self.parameters()) # ════════════════════════════════════════════════════════════════════════ # 2. RK4 ODE INTEGRATOR integrate dh/dt = F(h) # ════════════════════════════════════════════════════════════════════════ def rk4_integrate(field: nn.Module, h0: torch.Tensor, t_total: float = T_TOTAL, steps: int = STEPS) -> torch.Tensor: """ 4th-order Runge-Kutta integration of dh/dt = F(h). h0: [B, D] initial conditions returns: [B, D] final state h(T) """ dt = t_total / steps h = h0 for _ in range(steps): k1 = field(h) k2 = field(h + 0.5 * dt * k1) k3 = field(h + 0.5 * dt * k2) k4 = field(h + dt * k3) h = h + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) return h # ════════════════════════════════════════════════════════════════════════ # 3. IMPLICIT BOUNDARY S(p) = 0 defines the decision surface # ════════════════════════════════════════════════════════════════════════ class ImplicitBoundary(nn.Module): """ Small MLP that learns a signed-distance-like function S(p). S(p) > 0 → inside attractor basin A S(p) < 0 → inside attractor basin B S(p) = 0 → on the boundary (ambiguous without flux info) """ def __init__(self, dim: int = D, hidden: int = DHIDDEN): super().__init__() self.net = nn.Sequential( nn.Linear(dim, hidden), nn.Tanh(), nn.Linear(hidden, 1), ) def forward(self, p: torch.Tensor) -> torch.Tensor: """p: [..., D] → S: [..., 1]""" return self.net(p) # ════════════════════════════════════════════════════════════════════════ # 4. BOUNDARY FLUX 𝔽(p) = ∇S · F(p) ← THE MISSING INFORMATION # ════════════════════════════════════════════════════════════════════════ def compute_flux(boundary: nn.Module, field: FourierField, p: torch.Tensor, create_graph: bool = True): """ Compute first-order and (optionally) second-order transverse flux. Returns: flux1: ∇S · F — tells which way the trajectory crosses S=0 flux2: ∇(∇S·F) · F — super-resolution: how sharply S bends grad_S: ∇S — boundary normal f_val: F(p) — field velocity """ p = p.detach().requires_grad_(True) # S(p) and ∇S s_val = boundary(p) grad_S = torch.autograd.grad(s_val, p, create_graph=create_graph)[0] # F(p) f_val = field(p) # First-order flux flux1 = (grad_S * f_val).sum(dim=-1, keepdim=True) # Second-order flux flux2 = None if create_graph: grad_flux = torch.autograd.grad( flux1.sum(), p, create_graph=create_graph )[0] flux2 = (grad_flux * f_val).sum(dim=-1, keepdim=True) return flux1, flux2, grad_S, f_val # ════════════════════════════════════════════════════════════════════════ # 5. FULL AI VOXEL MODULE (field + boundary + integrator + flux) # ════════════════════════════════════════════════════════════════════════ class AIVoxel(nn.Module): """ A complete AI Voxel: field + boundary + RK4 + flux. forward(p) → {density, color, flux, edge_sharpness} query(p) → full diagnostics dict """ def __init__(self, dim: int = D, num_modes: int = K, hidden: int = DHIDDEN, steps: int = STEPS, t_total: float = T_TOTAL): super().__init__() self.field = FourierField(dim, num_modes) self.boundary = ImplicitBoundary(dim, hidden) self.steps = steps self.t_total = t_total self.D = dim def forward(self, p: torch.Tensor) -> dict: """Batch inference: p [B, D] → outputs dict.""" # 1. Integrate the ODE h_final = rk4_integrate(self.field, p, self.t_total, self.steps) # 2. Evaluate boundary S = self.boundary(h_final) # [B, 1] # 3. Compute flux (the missing information) with torch.enable_grad(): flux1, flux2, grad_S, f_val = compute_flux( self.boundary, self.field, h_final, create_graph=False ) # 4. Generate outputs density = torch.sigmoid(S * 10.0) # [B, 1] color = 0.5 + 0.5 * torch.tanh(h_final[:, :3]) # [B, 3] edge_sharpness = flux1.abs() / (1.0 + S.abs()) # [B, 1] return { "density": density.squeeze(-1), "color": color, "boundary_val": S.squeeze(-1), "flux": flux1.squeeze(-1), "edge_sharpness": edge_sharpness.squeeze(-1), "h_final": h_final, } def query(self, p: torch.Tensor) -> dict: """Single-point diagnostic query with full flux info.""" with torch.enable_grad(): h_final = rk4_integrate(self.field, p, self.t_total, self.steps) flux1, flux2, grad_S, f_val = compute_flux( self.boundary, self.field, h_final, create_graph=True ) S = self.boundary(h_final) return { "input": p.detach(), "h_final": h_final.detach(), "S": S.detach().item(), "flux1": flux1.detach().item(), "flux2": flux2.detach().item() if flux2 is not None else None, "grad_S": grad_S.detach(), "field_velocity": f_val.detach(), } # -- export coefficients for GLSL shader (Stage 6) ---------------- def export_coefficients(self) -> dict: """Export flat numpy arrays ready for SSBO upload.""" params = list(self.boundary.net.parameters()) return { "omega": self.field.omega.data.cpu().numpy(), # [K, D] "a": self.field.a.data.cpu().numpy(), # [K, D] "b": self.field.b.data.cpu().numpy(), # [K, D] "K": self.field.K, "D": self.field.D, "bW1": params[0].data.cpu().numpy(), # [DHIDDEN, D] "bB1": params[1].data.cpu().numpy(), # [DHIDDEN] "bW2": params[2].data.squeeze(0).cpu().numpy(), # [DHIDDEN] "bB2": params[3].data.cpu().numpy(), # [1] "Dhidden": self.boundary.net[0].out_features, "steps": self.steps, "T": self.t_total, } # ════════════════════════════════════════════════════════════════════════ # 6. TRAINING LOOP — distill a *target* dynamics into the Fourier field # ════════════════════════════════════════════════════════════════════════ def train_field_to_target( voxel: AIVoxel, target_fn, # callable: h → dh/dt num_samples: int = 8192, epochs: int = 2000, lr: float = 2e-3, dist_scale: float = 0.5, print_every: int = 200, ): """ Fit the Fourier field to match an arbitrary target vector field (e.g. a residual block from a HF model, or any dh/dt = f(h)). Uses simple MSE on the instantaneous derivative F(h) ≈ target_fn(h). """ opt = torch.optim.Adam(voxel.field.parameters(), lr=lr) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs) # Pre-sample random query points h_samples = (torch.randn(num_samples, voxel.D, device=DEVICE) * dist_scale) for epoch in range(epochs): # Random mini-batch idx = torch.randint(0, num_samples, (min(512, num_samples),), device=DEVICE) h_batch = h_samples[idx] f_pred = voxel.field(h_batch) # [B, D] with torch.no_grad(): f_target = target_fn(h_batch) # [B, D] loss = torch.nn.functional.mse_loss(f_pred, f_target) opt.zero_grad() loss.backward() opt.step() scheduler.step() if epoch % print_every == 0: rel_err = (loss.sqrt() / (f_target.abs().mean() + 1e-8)).item() print(f" epoch {epoch:5d} loss={loss.item():.6e} " f"rel_err≈{rel_err:.4f} lr={scheduler.get_last_lr()[0]:.2e}") print(f" Final loss: {loss.item():.6e}") print(f" Field params: {voxel.field.num_params():,} " f"(≈{voxel.field.num_params() * 4 / 1024:.1f} KB)") # ════════════════════════════════════════════════════════════════════════ # 7. BOUNDARY TRAINING — physics-informed loss (eikonal + flux) # ════════════════════════════════════════════════════════════════════════ def train_boundary( voxel: AIVoxel, p_data: torch.Tensor, # [N, D] sample points labels: torch.Tensor, # [N, 1] +1 or -1 epochs: int = 1000, lr: float = 1e-3, lambda_eik: float = 0.1, lambda_flux: float = 0.5, print_every: int = 200, ): """ Train S(p) with three losses: 1. Margin classification (off-boundary points) 2. Eikonal regularization |∇S| = 1 (signed distance property) 3. Flux consistency ∇S·F aligns with label direction """ opt = torch.optim.Adam(voxel.boundary.parameters(), lr=lr) for epoch in range(epochs): idx = torch.randint(0, len(p_data), (min(256, len(p_data)),), device=DEVICE) p_batch = p_data[idx].detach().requires_grad_(True) y_batch = labels[idx] # S(p) s_val = voxel.boundary(p_batch) # Classification loss (margin ranking) cls_loss = torch.nn.functional.margin_ranking_loss( s_val.squeeze(), y_batch.squeeze(), -torch.ones_like(y_batch.squeeze()), margin=0.5, ) # ∇S grad_S = torch.autograd.grad(s_val.sum(), p_batch, create_graph=True)[0] # Eikonal: |∇S| ≈ 1 eik_loss = ((grad_S.norm(dim=-1) - 1.0) ** 2).mean() # Flux consistency: sign(∇S·F) should match label f_val = voxel.field(p_batch).detach() flux = (grad_S * f_val).sum(dim=-1) flux_loss = torch.relu(-flux * y_batch.squeeze()).mean() total = cls_loss + lambda_eik * eik_loss + lambda_flux * flux_loss opt.zero_grad() total.backward() opt.step() if epoch % print_every == 0: print(f" epoch {epoch:5d} total={total.item():.4f} " f"cls={cls_loss.item():.4f} eik={eik_loss.item():.4f} " f"flux={flux_loss.item():.4f}") print(f" Boundary training complete.") # ════════════════════════════════════════════════════════════════════════ # 8. DEMO / EXPERIMENT ENTRY POINT # ════════════════════════════════════════════════════════════════════════ def demo_target_field(h: torch.Tensor) -> torch.Tensor: """ A simple target dynamics for distillation experiments: a damped oscillator in the first 2 dims, decay elsewhere. Replace this with any function (e.g. a HF model's residual block). """ dh = torch.zeros_like(h) # damped pendulum: dh0/dt = h1, dh1/dt = -h0 - 0.1*h1 dh[:, 0] = h[:, 1] dh[:, 1] = -h[:, 0] - 0.1 * h[:, 1] # decay remaining dims dh[:, 2:] = -0.05 * h[:, 2:] return dh def main(): print("=" * 64) print(" AI VOXEL — Small GPU Experiment Kit") print(f" Device: {DEVICE} D={D} K={K} steps={STEPS} T={T_TOTAL}") print("=" * 64) # --- Build the AI Voxel --- voxel = AIVoxel(dim=D, num_modes=K, hidden=DHIDDEN, steps=STEPS, t_total=T_TOTAL).to(DEVICE) total_params = sum(p.numel() for p in voxel.parameters()) print(f"\nVoxel parameters: {total_params:,} " f"(≈{total_params * 4 / 1024:.1f} KB)") # --- Stage 3: Distill field to target dynamics --- print("\n[Stage 3] Distilling Fourier field → target dynamics ...") train_field_to_target(voxel, demo_target_field, num_samples=8192, epochs=2000, lr=2e-3) # --- Stage 4: Train boundary --- print("\n[Stage 4] Training implicit boundary S(p) ...") # Synthetic labeled data: label = sign of h[0] after integration p_data = torch.randn(2048, D, device=DEVICE) * 0.5 with torch.no_grad(): h_final = rk4_integrate(voxel.field, p_data, T_TOTAL, STEPS) labels = torch.sign(h_final[:, 0:1]) labels[labels == 0] = 1.0 train_boundary(voxel, p_data, labels, epochs=1000, lr=1e-3) # --- Inference benchmark --- print("\n[Inference] Batch query benchmark ...") q = torch.randn(BATCH, D, device=DEVICE) * 0.5 if DEVICE.type == "cuda": torch.cuda.synchronize() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() out = voxel(q) if DEVICE.type == "cuda": end.record() torch.cuda.synchronize() ms = start.elapsed_time(end) print(f" Batch size: {BATCH}") print(f" GPU time: {ms:.3f} ms ({BATCH / (ms / 1e3):,.0f} queries/sec)") else: print(f" Batch size: {BATCH} (CPU mode)") print(f" Output keys: {list(out.keys())}") print(f" density range: [{out['density'].min():.3f}, {out['density'].max():.3f}]") print(f" flux range: [{out['flux'].min():.4f}, {out['flux'].max():.4f}]") # --- Single-point diagnostic query --- print("\n[Diagnostic] Single-point flux query ...") p_test = torch.zeros(1, D, device=DEVICE) p_test[0, 0] = 0.3 p_test[0, 1] = 0.1 diag = voxel.query(p_test) print(f" S(p) = {diag['S']:.6f}") print(f" flux1 (∇S·F) = {diag['flux1']:.6f} ← missing information") print(f" flux2 (2nd) = {diag['flux2']:.6f} ← super-resolution") print(f" |∇S| = {diag['grad_S'].norm():.6f} (target: 1.0)") print(f" |F(p)| = {diag['field_velocity'].norm():.6f}") # --- Export for GLSL shader --- print("\n[Export] Coefficients for GLSL shader ...") coeffs = voxel.export_coefficients() field_bytes = 3 * coeffs["K"] * coeffs["D"] * 4 # omega + a + b, float32 print(f" Field coefficients: {coeffs['K']} modes × {coeffs['D']} dims") print(f" Field size: {field_bytes / 1024:.1f} KB") print(f" Ready for SSBO upload to ai_voxel_shader.glsl") # --- Save coefficients for the GLSL renderer --- save_path = "ai_voxel_coeffs.npz" np.savez(save_path, **coeffs) print(f"\n Saved coefficients → {save_path}") print(f" Run: python ai_voxel_render.py to render on GPU") print("\n" + "=" * 64) print(" Done. Tweak D, K, STEPS, T_TOTAL at the top to experiment.") print("=" * 64) if __name__ == "__main__": main()