diff --git a/examples/06_spectral_analysis/PhaseFormer.md b/examples/06_spectral_analysis/PhaseFormer.md new file mode 100644 index 0000000..2343b5c --- /dev/null +++ b/examples/06_spectral_analysis/PhaseFormer.md @@ -0,0 +1,741 @@ +# PhaseFormer + +A dual-stream hierarchical transformer for coherent phase accumulation +in spectral token sequences. Designed for precision parameter estimation +of chirping signals — gravitational waves, radar, sonar, spectroscopy. + +``` + ╔═══════════════════════════════╗ + ║ PhaseFormer ║ + ║ ║ + ║ "Coherence is all you need" ║ + ╚═══════════════════════════════╝ +``` + +--- + +## 1. Motivation + +The Cramér-Rao bound (CRB) on frequency estimation scales as ~1/(N·SNR), +where N is the total number of samples. Achieving this precision requires +**coherent integration** of phase across the full observation — exactly +what a matched filter does. + +A standard transformer operating on spectral tokens must learn to: +1. Identify which tokens belong to the same spectral track (frequency matching) +2. Compute phase differences between adjacent tokens (complex arithmetic) +3. Accumulate these differences over the full signal (sequential chaining) + +Step 3 is the bottleneck. Standard self-attention computes weighted sums, +not products. Phase accumulation requires chaining complex multiplications +— a fundamentally different operation. + +**PhaseFormer bakes coherent phase accumulation into the architecture.** + +--- + +## 2. Token Representation + +Each token represents one spectral peak in one time window, carrying +a **feature vector** (real-valued) and a **phase vector** (complex-valued): + +``` + Token = (f, p) + + f ∈ R^d_f feature vector: frequency, amplitude, time + p ∈ C² phase vectors: p_start, p_end (unit complex) +``` + +From the raw ToneTokenizer output `(f_start, f_end, amp, φ_start, φ_end)`: + +``` + ┌─────────────────────────────────────────────────────────┐ + │ Feature stream (real-valued, for attention routing): │ + │ │ + │ f_start, f_end raw frequency in [-1, 1] │ + │ fourier(f_start) multi-scale sin/cos encoding │ + │ fourier(f_end) multi-scale sin/cos encoding │ + │ fourier(log1p(A)) multi-scale amplitude encoding │ + │ t window time position in [-1, 1] │ + │ │ + ├─────────────────────────────────────────────────────────┤ + │ Phase stream (complex-valued, for accumulation): │ + │ │ + │ p_start = (cos φ_start, sin φ_start) │ + │ p_end = (cos φ_end, sin φ_end) │ + │ │ + └─────────────────────────────────────────────────────────┘ +``` + +The feature stream determines **which** tokens to link (Q·K attention). +The phase stream carries the **payload** that gets coherently accumulated. + +### Spectral embedding of amplitude + +Amplitude is embedded with the same multi-scale Fourier features as +frequency: `fourier(log1p(A))` → `sin/cos` at geometrically spaced scales. + +This ensures: +- **Amplitude matching via dot products**: tokens with similar brightness + score high in attention, same mechanism as frequency matching +- **Dynamic range handling**: log scale compresses the range, spectral + embedding distributes it across multiple features — very bright and + very faint signals get equal representation +- **Balanced feature space**: amplitude gets comparable dimensionality + to frequency, preventing frequency features from dominating attention + +### Higher harmonics + +Higher harmonic structure (integer multiples, or more complex +relationships like f = f₀ + m·f₁ + n·f₂ for EMRIs) is not handled +by the PhaseFormer architecture itself. Instead, harmonic detection +is delegated to an upstream peak detector that identifies harmonic +relationships and encodes them as additional token features. This +keeps the transformer focused on what it does best: frequency matching +and phase accumulation. + +--- + +## 3. Architecture Overview + +``` + Input: (B, W, K, 5) raw tokens + W = time windows, K = peaks per window + + + ┌──────────────────────────────────────────────┐ + │ SpectralTokenEmbed │ + │ raw tokens → feature vectors + phase vectors │ + └──────────────────┬───────────────────────────┘ + │ + ▼ + ┌─────────────────────┐ + │ W=50, K tokens │ ← full resolution + └─────────┬───────────┘ + │ + ┌──────────┴──────────┐ + ▼ ▼ + ┌─────────────┐ ┌──────────────────┐ + │ Intra-window │ │ Inter-window │ + │ self-attn │ │ cross-attn │ + │ (K × K) │ │ (w ↔ w±1) │ + │ │ │ │ + │ standard │ │ feature: sum │ + │ attention │ │ phase: complex │ + │ │ │ multiply │ + └──────┬──────┘ └────────┬─────────┘ + │ │ + └────────┬───────────┘ + │ + ▼ stride 2 (drop odd windows) + ┌─────────────────────┐ + │ W=25, K tokens │ + └─────────┬───────────┘ + │ + ┌──────────┴──────────┐ + ▼ ▼ + ┌─────────────┐ ┌──────────────────┐ + │ Intra-window │ │ Inter-window │ + │ self-attn │ │ cross-attn │ + └──────┬──────┘ └────────┬─────────┘ + │ │ + └────────┬───────────┘ + │ + ▼ stride 2 + ┌─────────────────────┐ + │ W=12, K tokens │ + └─────────┬───────────┘ + │ + ... (repeat log₂(W) times) + │ + ▼ + ┌─────────────────────┐ + │ W=1, K tokens │ ← full phase accumulated + └─────────┬───────────┘ + │ + ▼ + ┌─────────────────────┐ + │ Pool → (B, d) │ + │ summary embedding │ + └─────────────────────┘ +``` + +Each pyramid level can contain multiple (MHA → FF) sublayers for the +intra-window step before the inter-window cross-attention and stride. +The number of sublayers per level (`n_intra`) is a hyperparameter — +more sublayers allow deeper feature extraction per level at the cost +of compute. + +--- + +## 4. The Two Attention Mechanisms + +### 4.1 Intra-Window Self-Attention (Standard) + +Processes peaks within the same time window. Learns to classify signal +vs noise peaks, extract local spectral features, and build per-token +representations for downstream frequency matching. + +``` + Window w: [peak₀, peak₁, peak₂, ..., peak_K] + + ┌───────────────────┐ + │ Self-Attention │ + │ │ + peak_i ────Q─────┤ A_ij · V_j ├──── peak_i' + peak_j ────K,V───┤ (standard sum) │ + └───────────────────┘ + + Q, K, V all derived from feature stream f. + Phase stream p passes through unchanged. +``` + +Standard multi-head attention. Each window processed independently. +Complexity: O(K²) per window, O(W·K²) total. + +### 4.2 Inter-Window Cross-Attention (Complex Phase) + +Links tokens between adjacent windows. **This is where the magic happens.** + +The inter-window step has two stages: first, multiple cross-attention +rounds refine the feature representations using phase information; +then, a final step performs the actual coherent phase accumulation. + +#### Stage 1: Feature refinement with phase information (N_inter rounds) + +Phase angles are spectrally embedded — `(cos φ, sin φ, cos 2φ, sin 2φ, +..., cos Mφ, sin Mφ)` — and concatenated into the feature stream for +attention computation. Higher harmonics of the phase provide finer +angular resolution in the dot product, allowing the attention to +distinguish peaks whose phases are close but not identical. This also +balances the feature space: phase gets comparable dimensionality to +frequency features. + +``` + ┌─────────────────────────────────────────────────────┐ + │ Augmented feature vector for attention (stage 1): │ + │ │ + │ f_start, f_end, fourier(f), fourier(log A), t │ + │ + fourier(φ_start): cos(φ), sin(φ), cos(2φ),.. │ + │ + fourier(φ_end): cos(φ), sin(φ), cos(2φ),.. │ + │ │ + └─────────────────────────────────────────────────────┘ + + repeat N_inter times: + cross-attention between w ↔ w+1 (standard, features only) + features updated via residual + FF + + After N_inter rounds, the feature stream encodes which tokens + across windows are phase-coherent — noise peaks get suppressed, + signal tracks get reinforced. +``` + +#### Stage 2: Phase accumulation (single final step) + +``` + Window w: Window w+1: + ┌──────────┐ ┌──────────┐ + │ f_w, p_w │ │f_w+1,p_w+1│ + └────┬─────┘ └────┬──────┘ + │ │ + │ ┌───────────┴────────────┐ + └───►│ Cross-Attention │ + │ │ + │ A_ij from refined │ + │ features (well- │ + │ informed by phase) │ + │ │ + │ Phase accumulation: │ + │ Δp = conj(p_end_w) │ + │ ⊗ p_start_w+1 │ + │ │ + │ p'_w+1 = normalize( │ + │ Σ_j A_ij · Δp_ij) │ + │ │ + └────────────────────────┘ +``` + +The attention weights A_ij are now informed by multiple rounds of +feature exchange *including phase information*. Tokens that showed +phase continuity in stage 1 get high attention; noise peaks that +matched nothing get suppressed. + +The phase accumulation uses the raw `(cos φ, sin φ)` vectors — +higher spectral modes of φ are only used for computing attention +weights, not for the complex multiplication itself. + +When token i in window w+1 attends strongly to token j in window w +(same spectral track, high frequency, amplitude, and phase match): + +``` + Δp = conj(p_end_j) ⊗ p_start_i = (cos Δφ, sin Δφ) +``` + +This Δφ is the phase jump at the window boundary. For a clean signal, +Δφ ≈ 0 (phase continuity). The magnitude of the aggregated vector +gives **coherence** — a built-in SNR indicator. + +--- + +## 5. Phase Accumulation via Hierarchical Striding + +The key insight: after each inter-window layer, surviving tokens carry +phase information spanning their window's time range. Striding by 2 +doubles the effective span each layer. + +``` + Layer 0: Phase spans 1 window + ───┬───┬───┬───┬───┬───┬───┬─── + w0 w1 w2 w3 w4 w5 w6 w7 W = 8 + + inter-window cross-attention (w ↔ w+1) + stride 2: keep even windows + + Layer 1: Phase spans 2 windows + ───┬───────┬───────┬───────┬─── + w0 w2 w4 w6 W = 4 + + inter-window cross-attention + stride 2 + + Layer 2: Phase spans 4 windows + ───┬───────────────┬─────────────── + w0 w4 W = 2 + + inter-window cross-attention + stride 2 + + Layer 3: Phase spans 8 windows (full signal) + ───┬─────────────────────────────── + w0 W = 1 + + + Phase evolution across layers for one spectral track: + + Layer 0: p₀──Δφ₀₁──p₁──Δφ₁₂──p₂──Δφ₂₃──p₃──Δφ₃₄──p₄ + │ │ + Layer 1: p₀───┴──Δφ₀₂──────p₂───────┴──Δφ₂₄──────p₄ + │ │ + Layer 2: p₀────────┴────Δφ₀₄──────────p₄──┘ + │ + Layer 3: p₀───────────────┴── = (cos Σ Δφ, sin Σ Δφ) + └── total accumulated phase +``` + +After log₂(W) layers, each token carries the **total phase integral** +for its spectral track. This is the matched filter output. + +--- + +## 6. Complex Phase Operations + +All phase operations use 2D real arithmetic (no complex dtype needed): + +```python +def complex_mul(a_cos, a_sin, b_cos, b_sin): + """Multiply two unit complex numbers: a ⊗ b.""" + return (a_cos * b_cos - a_sin * b_sin, + a_cos * b_sin + a_sin * b_cos) + +def complex_conj_mul(a_cos, a_sin, b_cos, b_sin): + """Phase difference: conj(a) ⊗ b = (cos(θb-θa), sin(θb-θa)).""" + return (a_cos * b_cos + a_sin * b_sin, + a_cos * b_sin - a_sin * b_cos) + +def complex_normalize(c_cos, c_sin): + """Project back onto unit circle.""" + mag = (c_cos**2 + c_sin**2).sqrt().clamp(min=1e-8) + return c_cos / mag, c_sin / mag +``` + +The magnitude before normalization is the **coherence**: +- |c| ≈ 1 → attended phases agree → real signal +- |c| ≈ 0 → attended phases random → noise + +This can be fed back into the feature stream as a confidence signal. + +--- + +## 7. Detailed Layer Pseudocode + +```python +class PhaseFormerLayer(nn.Module): + """One level of the PhaseFormer pyramid. + + Each level: + 1. Intra-window self-attention (n_intra rounds) + 2. Inter-window cross-attention with phase features (n_inter rounds) + 3. Phase accumulation via complex multiply (single step) + 4. Stride 2 (implicit — outputs have W' = W // 2) + """ + + def __init__(self, d_model, n_heads, d_ff, n_intra=1, n_inter=1, + n_phase_harmonics=4): + self.n_phase_harmonics = n_phase_harmonics + # Phase feature dimension: n_harmonics * 2 (cos,sin) * 2 (start,end) + d_phase = n_phase_harmonics * 2 * 2 + + # Intra-window self-attention (repeated n_intra times) + self.intra_blocks = nn.ModuleList([ + nn.ModuleDict({ + 'norm1': LayerNorm(d_model), + 'attn': MultiHeadAttention(d_model, n_heads), + 'norm2': LayerNorm(d_model), + 'ff': FeedForward(d_model, d_ff), + }) + for _ in range(n_intra) + ]) + + # Inter-window cross-attention with phase features (n_inter rounds) + # Input is d_model + d_phase (features augmented with phase) + d_aug = d_model + d_phase + self.inter_blocks = nn.ModuleList([ + nn.ModuleDict({ + 'norm1': LayerNorm(d_aug), + 'q': nn.Linear(d_aug, d_model), + 'k': nn.Linear(d_aug, d_model), + 'v': nn.Linear(d_aug, d_model), + 'norm2': LayerNorm(d_model), + 'ff': FeedForward(d_model, d_ff), + }) + for _ in range(n_inter) + ]) + + # Final cross-attention for phase accumulation + self.phase_q = nn.Linear(d_aug, d_model) + self.phase_k = nn.Linear(d_aug, d_model) + self.phase_norm = LayerNorm(d_aug) + + def _phase_features(self, ps_cos, ps_sin, pe_cos, pe_sin): + """Spectral embedding of phase angles for attention. + + Returns (cos φ, sin φ, cos 2φ, sin 2φ, ...) for both + start and end phases, concatenated. + """ + parts = [] + for m in range(1, self.n_phase_harmonics + 1): + if m == 1: + parts += [ps_cos, ps_sin, pe_cos, pe_sin] + else: + # cos(mφ) and sin(mφ) via Chebyshev recurrence + # or direct: angle = m * atan2(sin, cos) + ps_angle = torch.atan2(ps_sin, ps_cos) + pe_angle = torch.atan2(pe_sin, pe_cos) + parts += [torch.cos(m * ps_angle), torch.sin(m * ps_angle), + torch.cos(m * pe_angle), torch.sin(m * pe_angle)] + return torch.stack(parts, dim=-1) # (..., n_harmonics * 4) + + def forward(self, f, ps_cos, ps_sin, pe_cos, pe_sin): + """ + f: (B, W, K, d_model) feature stream + p_*_cos/sin: (B, W, K) phase streams (unit complex) + + Returns same shapes with W' = W // 2 (strided). + """ + B, W, K, d = f.shape + + # ─── Intra-window self-attention ─────────────────────── + f_flat = f.reshape(B * W, K, d) + for block in self.intra_blocks: + f_flat = f_flat + block['attn']( + block['norm1'](f_flat)) + f_flat = f_flat + block['ff']( + block['norm2'](f_flat)) + f = f_flat.reshape(B, W, K, d) + + # ─── Inter-window cross-attention (feature refinement) ─ + f_even = f[:, 0::2] # (B, W//2, K, d) + f_odd = f[:, 1::2] # (B, W//2, K, d) + W2 = f_even.shape[1] + + # Build phase features for attention + pf_even = self._phase_features( + ps_cos[:, 0::2], ps_sin[:, 0::2], + pe_cos[:, 0::2], pe_sin[:, 0::2]) # (B, W2, K, d_phase) + pf_odd = self._phase_features( + ps_cos[:, 1::2], ps_sin[:, 1::2], + pe_cos[:, 1::2], pe_sin[:, 1::2]) + + for block in self.inter_blocks: + # Augment features with phase for Q, K computation + f_even_aug = torch.cat([f_even, pf_even], dim=-1) + f_odd_aug = torch.cat([f_odd, pf_odd], dim=-1) + + fe = f_even_aug.reshape(B * W2, K, -1) + fo = f_odd_aug.reshape(B * W2, K, -1) + + q = block['q'](block['norm1'](fe)) + k = block['k'](fo) + v = block['v'](fo) + + scale = d ** 0.5 + A = (q @ k.transpose(-1, -2) / scale).softmax(dim=-1) + + # Update even-window features (standard residual) + f_cross = (A @ v).reshape(B, W2, K, d) + f_even = f_even + f_cross + f_even = f_even + self.inter_blocks[0]['ff']( + block['norm2'](f_even.reshape(B * W2, K, d)) + ).reshape(B, W2, K, d) + + # ─── Phase accumulation (single final step) ─────────── + f_even_aug = torch.cat([f_even, pf_even], dim=-1) + f_odd_aug = torch.cat([f_odd, pf_odd], dim=-1) + + fe = f_even_aug.reshape(B * W2, K, -1) + fo = f_odd_aug.reshape(B * W2, K, -1) + + q = self.phase_q(self.phase_norm(fe)) + k = self.phase_k(fo) + A = (q @ k.transpose(-1, -2) / (d ** 0.5)).softmax(dim=-1) + + # Complex conjugate multiply: conj(p_end_odd) ⊗ p_start_even + pe_cos_o = pe_cos[:, 1::2].reshape(B * W2, K) + pe_sin_o = pe_sin[:, 1::2].reshape(B * W2, K) + ps_cos_e = ps_cos[:, 0::2].reshape(B * W2, K) + ps_sin_e = ps_sin[:, 0::2].reshape(B * W2, K) + + # Phase difference for each (even_i, odd_j) pair + delta_cos = (ps_cos_e.unsqueeze(2) * pe_cos_o.unsqueeze(1) + + ps_sin_e.unsqueeze(2) * pe_sin_o.unsqueeze(1)) + delta_sin = (ps_sin_e.unsqueeze(2) * pe_cos_o.unsqueeze(1) + - ps_cos_e.unsqueeze(2) * pe_sin_o.unsqueeze(1)) + + # Attention-weighted aggregation + agg_cos = (A * delta_cos).sum(dim=-1) + agg_sin = (A * delta_sin).sum(dim=-1) + + # Coherence + normalize + coherence = (agg_cos**2 + agg_sin**2).sqrt() + mag = coherence.clamp(min=1e-8) + ps_cos_out = (agg_cos / mag).reshape(B, W2, K) + ps_sin_out = (agg_sin / mag).reshape(B, W2, K) + pe_cos_out = pe_cos[:, 0::2] + pe_sin_out = pe_sin[:, 0::2] + + return f_even, ps_cos_out, ps_sin_out, \ + pe_cos_out, pe_sin_out, coherence.reshape(B, W2, K) +``` + +--- + +## 8. Full Architecture + +```python +class PhaseFormer(nn.Module): + """ + Hierarchical dual-stream transformer for coherent phase accumulation. + + Input: (B, W, K, 5) raw spectral tokens + Output: (B, d_out) summary embedding for downstream estimation + """ + + def __init__(self, d_model=64, n_heads=4, d_ff=256, + n_freq=3, amp_scale=10.0, n_intra=1, + max_layers=8): + # Token embedding (SpectralTokenEmbed-like) + self.embed = TokenEmbedding(n_freq, amp_scale) + self.input_proj = nn.LazyLinear(d_model) + + # Pyramid layers: one per halving step + self.layers = nn.ModuleList([ + PhaseFormerLayer(d_model, n_heads, d_ff, n_intra) + for _ in range(max_layers) + ]) + + # Final pooling + self.output_proj = nn.Linear(d_model + 2, d_out) + # │ └── accumulated phase (cos, sin) + # └────────── feature stream + + def forward(self, raw_tokens): + B, W, K, _ = raw_tokens.shape + + # Embed tokens → feature + phase streams + f, ps_cos, ps_sin, pe_cos, pe_sin = self.embed(raw_tokens) + f = self.input_proj(f) # (B, W, K, d_model) + + # Pad W to next power of 2 if needed + ... + + # Pyramid: log₂(W) layers + coherences = [] + layer_idx = 0 + while f.shape[1] > 1: + f, ps_cos, ps_sin, pe_cos, pe_sin, coh = \ + self.layers[layer_idx](f, ps_cos, ps_sin, pe_cos, pe_sin) + coherences.append(coh) + layer_idx += 1 + + # f: (B, 1, K, d_model), phase: (B, 1, K) + # Concatenate accumulated phase into feature stream + f_final = torch.cat([ + f.squeeze(1), # (B, K, d_model) + ps_cos.squeeze(1).unsqueeze(-1), # (B, K, 1) + ps_sin.squeeze(1).unsqueeze(-1), # (B, K, 1) + ], dim=-1) # (B, K, d_model+2) + + # Pool over K peaks → (B, d_model+2) + summary = f_final.mean(dim=1) + + return self.output_proj(summary) # (B, d_out) +``` + +--- + +## 9. Complexity Analysis + +``` + Standard Transformer: O(W² · K²) per layer + PhaseFormer (per layer): O(W_l · K²) intra + inter + PhaseFormer (total): O(W · K² · log₂W) + + ┌────────────┬────────────┬──────────────────────────┐ + │ Component │ Per layer │ Example (W=50, K=100) │ + ├────────────┼────────────┼──────────────────────────┤ + │ Intra-attn │ W_l · K² │ 50 · 10,000 = 500K │ + │ Inter-attn │ W_l · K² │ 50 · 10,000 = 500K │ + │ Total/layer│ 2W_l · K² │ 1M │ + │ All layers │ Σ 2W_l·K² │ ≈ 2M (geometric sum) │ + ├────────────┼────────────┼──────────────────────────┤ + │ Standard │ (WK)² = 25M per layer × L layers │ + └────────────┴────────────┴──────────────────────────┘ + + Speedup: ~10-50x for typical configurations. +``` + +--- + +## 10. Why This Works for Gravitational Waves + +``` + ┌──────────────────────────────────────────────────────┐ + │ LISA Signal │ + │ │ + │ Source 1: f₁(t) ───── slowly chirping ─────► │ + │ Source 2: f₂(t) ──── different chirp rate ──► │ + │ Source 3: f₃(t) ─── overlapping in freq ───► │ + │ + noise │ + │ │ + │ Standard transformer: must learn to separate │ + │ sources, track them, AND accumulate phase │ + │ — all through generic attention. Hard. │ + │ │ + │ PhaseFormer: frequency + amplitude attention │ + │ separates sources automatically. Phase │ + │ accumulation is structural. Network only learns │ + │ spectral matching — the easy part. │ + │ │ + └──────────────────────────────────────────────────────┘ + + Phase precision for frequency estimation: + + δf ~ 1 / (2π · T_obs · SNR) + + Matched filter achieves this by integrating phase over T_obs. + PhaseFormer does the same through log₂(W) layers of complex + phase accumulation — structurally guaranteed, not learned. + + ┌───────────────────────────────────────────┐ + │ │ + │ "The architecture IS the physics" │ + │ │ + └───────────────────────────────────────────┘ +``` + +--- + +## 11. Implementation Plan + +``` + Phase 1: Single-source prototype (falcon example 06) + ├── Implement PhaseFormerLayer in examples/06_spectral_analysis/src/ + ├── Test on current chirp signal (N=100k) + ├── Compare posterior width to CRB and standard TransformerEmbedding + └── Validate: inspect accumulated phases, check coherence values + + Phase 2: Multi-source (fuge integration) + ├── Move PhaseFormer to fuge.nn + ├── Test with 2-3 overlapping chirps + ├── Verify frequency-based track separation + └── Benchmark scaling with K (peaks per window) + + Phase 3: LISA scale + ├── Long observations (W > 1000) + ├── Hundreds of peaks per window + ├── Multiple simultaneous sources + ├── Upstream harmonic peak detector for EMRI-like signals + └── Integration with LISA data pipeline +``` + +--- + +## 12. Open Questions + +- **Phase stream gradient flow**: The complex multiply + normalize chain + may have vanishing/exploding gradients through many layers. May need + gradient clipping or residual connections in the phase stream. + +- **Odd W handling**: When W is odd, the stride-2 merge drops one window. + Pad to next power of 2, or handle remainder explicitly? + +- **Coherence feedback**: Should coherence (|aggregated phase|) be fed + back into the feature stream? This would let the network learn to + trust high-coherence tracks and ignore noise. + +- **Learnable scales**: Should the Fourier frequency/amplitude scales be + learnable rather than fixed geometric? May help adapt to + signal-specific distributions. + +- **Phase stream initialization**: Currently p_start and p_end come + directly from the tokenizer. Should they be projected/rotated by a + learned transformation first? + +- **Number of intra-window sublayers**: How many (MHA → FF) blocks per + pyramid level? More depth helps feature extraction but costs compute. + May need more at early levels (many peaks to sort) than later levels + (fewer, better-refined tokens). + +- **Uncertainty-damped spectral modes (Debye-Waller damping)**: The peak + detector can estimate measurement uncertainty σ_f for each frequency. + Higher Fourier modes in the spectral embedding should then be damped: + + ``` + fourier_i(f) = sin(s_i · f) · exp(-s_i² · σ_f² / 2) + ``` + + This is the exact result of marginalizing `cos(s·(f + σ·ε))` over + Gaussian noise ε — the expected value picks up a damping factor + `exp(-s²σ²/2)`. Physically: don't encode fine frequency structure + that the measurement can't resolve. + + Analogous to the **Debye-Waller factor** in crystallography, where + thermal vibrations damp high-order Bragg peaks by the same factor. + Here, measurement noise plays the role of thermal motion, and the + spectral embedding modes play the role of diffraction orders. + + This provides natural regularization: noisy peaks only contribute + through low-frequency modes (coarse matching), while precise peaks + retain their full multi-scale resolution. Acts as a physics-informed + alternative to dropout — the damping is set by actual measurement + precision, not a tuning parameter. + +``` + ┌────────────────────────────────────────────────────┐ + │ │ + │ ╱╲ Signal goes in │ + │ ╱ ╲ │ + │ ╱ ╲ Phases accumulate │ + │ ╱ ╱╲ ╲ │ + │ ╱ ╱ ╲ ╲ Sources separate │ + │ ╱ ╱ ╲ ╲ │ + │ ╱ ╱ ╱╲ ╲ ╲ Noise falls away │ + │ ╱ ╱ ╱ ╲ ╲ ╲ │ + │ ▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔ │ + │ K tokens, total accumulated phase │ + │ CRB-level precision on f₀ │ + │ │ + │ Physics in, posteriors out. │ + │ │ + └────────────────────────────────────────────────────┘ +``` diff --git a/examples/06_spectral_analysis/config.yml b/examples/06_spectral_analysis/config.yml new file mode 100644 index 0000000..be4e027 --- /dev/null +++ b/examples/06_spectral_analysis/config.yml @@ -0,0 +1,138 @@ +# ============================================================================= +# Falcon Example Configuration: 06_spectral_analysis (tokenized variant) +# ----------------------------------------------------------------------------- +# Uses Signal + Noise + Data decomposition with an intermediate "tokens" node +# that caches the expensive STFT tokenization. Training operates on compact +# token tensors instead of re-running the STFT every epoch. +# +# Graph: theta -> y (clean signal) +# \ +# -> x (observed = y + n) -> tokens +# / +# n (noise) +# +# NOTE: This config will currently FAIL during proposal/posterior sampling +# because falcon's inference graph does not include intermediate deterministic +# nodes (tokens has no estimator and no evidence). +# ============================================================================= + +logging: + wandb: + enabled: false + project: falcon_examples + group: 06_spectral_analysis + dir: ${run_dir} + local: + enabled: true + dir: ${paths.graph} + +paths: + import: "./src" + buffer: ${run_dir}/sim_dir + graph: ${run_dir}/graph_dir + samples: ${run_dir}/samples_dir + +buffer: + min_samples: 4092 + max_samples: 4092 + validation_samples: 256 + simulate_count: 256 + simulate_when_full: true + simulate_interval: 1 + store_fraction: 0.0 + +graph: + theta: + evidence: [tokens] + sample_chunk_size: 8 + + simulator: + _target_: falcon.priors.Product + priors: # 10x CRB at N=100k + - ['normal', 2.75e-3, 1e-7] # f0 + - ['normal', 1.0, 1e-5] # chirp_mass + - ['normal', 1.5, 1e-1] # harmonic_decay + + estimator: + _target_: falcon.estimators.Gaussian + loop: + num_epochs: 500 + batch_size: 64 + early_stop_patience: 32 + network: + hidden_dim: 512 + num_layers: 3 + momentum: 0.1 + min_var: 1.0e-25 + eig_update_freq: 1 + embedding: + _target_: fuge.nn.TransformerEmbedding + d_model: 64 + n_heads: 4 + n_layers: 2 + d_ff: 256 + dropout: 0.0 + _input_: + _target_: falcon.embeddings.RunningNorm + momentum: 0.01 + dim: [0, 1] + output_dtype: float32 + _input_: + _target_: embeddings.SpectralTokenEmbed + n_freq: 10 + phase_mode: boundary + _input_: [tokens] + optimizer: + lr: 0.005 + lr_decay_factor: 0.5 + scheduler_patience: 16 + inference: + gamma: 1.0 + discard_samples: false + log_ratio_threshold: -20.0 + + ray: + num_gpus: 0.2 + + y: + parents: [theta] + sample_chunk_size: 8 + simulator: + _target_: model.Signal + N: 100000 + ray: + num_gpus: 0.2 + + n: + simulator: + _target_: model.Noise + N: 100000 + sigma: 1.0 + + x: + parents: [y, n] + sample_chunk_size: 8 + simulator: + _target_: model.Data + observed: "./data/obs.npz['x']" + ray: + num_gpus: 0.2 + + tokens: + parents: [x] + sample_chunk_size: 8 + simulator: + _target_: embeddings.Tokenizer + k: 10000 # 2024 + n_peaks: 3 + n_dlnf: 11 + dlnf_min: 0.0 + dlnf_max: 0.05 + ray: + num_gpus: 0.2 + +sample: + prior: + n: 1000 + posterior: + n: 1000 diff --git a/examples/06_spectral_analysis/config2.yml b/examples/06_spectral_analysis/config2.yml new file mode 100644 index 0000000..cd2e424 --- /dev/null +++ b/examples/06_spectral_analysis/config2.yml @@ -0,0 +1,110 @@ +# ============================================================================= +# Falcon Example Configuration: 06_spectral_analysis (SVD variant) +# ----------------------------------------------------------------------------- +# Uses scaffolded streaming SVD embedding on raw signals. +# Clean signal (y) and noise (n) are scaffolds — available during training +# for whitening/PCA learning, but not at inference time. +# +# Graph: theta -> y (clean signal) +# \ +# -> x (observed = y + n) +# / +# n (noise) +# +# Prior ranges are 3x CRB at N=100k (narrow, for parameter estimation). +# Amortized only (no proposal resampling). +# ============================================================================= + +logging: + wandb: + enabled: false + project: falcon_examples + group: 06_spectral_analysis + dir: ${run_dir} + local: + enabled: true + dir: ${paths.graph} + +paths: + import: "./src" + buffer: ${run_dir}/sim_dir + graph: ${run_dir}/graph_dir + samples: ${run_dir}/samples_dir + +buffer: + min_samples: 4092 + max_samples: 4092 + validation_samples: 256 + simulate_count: 256 + simulate_when_full: true + simulate_interval: 1 + store_fraction: 0.0 + +graph: + theta: + evidence: [x] + scaffolds: [y, n] + + simulator: + _target_: falcon.priors.Product + priors: # 10x CRB at N=100k + - ['normal', 2.75e-3, 1e-7] # f0 + - ['normal', 1.0, 1e-5] # chirp_mass + - ['normal', 1.5, 1e-1] # harmonic_decay + + estimator: + _target_: falcon.estimators.Gaussian + loop: + num_epochs: 500 + batch_size: 64 + early_stop_patience: 32 + network: + hidden_dim: 128 + num_layers: 3 + momentum: 0.1 + min_var: 1.0e-25 + eig_update_freq: 1 + embedding: + _target_: embeddings.SVDEmbedding + N: 100000 + n_components: 64 + buffer_size: 256 + momentum: 0.1 + _input_: [x, y, n] + optimizer: + lr: 0.005 + lr_decay_factor: 0.5 + scheduler_patience: 16 + inference: + gamma: 1.0 + discard_samples: false + log_ratio_threshold: -20.0 + + ray: + num_gpus: 0.3 + + y: + parents: [theta] + simulator: + _target_: model.Signal + N: 100000 + ray: + num_gpus: 0.3 + + n: + simulator: + _target_: model.Noise + N: 100000 + sigma: 1.0 + + x: + parents: [y, n] + simulator: + _target_: model.Data + observed: "./data/obs.npz['x']" + +sample: + prior: + n: 1000 + posterior: + n: 1000 diff --git a/examples/06_spectral_analysis/config3.yml b/examples/06_spectral_analysis/config3.yml new file mode 100644 index 0000000..ac57f4b --- /dev/null +++ b/examples/06_spectral_analysis/config3.yml @@ -0,0 +1,106 @@ +# ============================================================================= +# Falcon Example Configuration: 06_spectral_analysis (N=256, no embedding) +# ----------------------------------------------------------------------------- +# Minimal config with 256 data bins — raw signal fed directly into Gaussian +# estimator (no embedding network needed at this data size). +# +# Graph: theta -> y (clean signal) +# \ +# -> x (observed = y + n) +# / +# n (noise) +# +# Prior ranges are ±10σ CRB at N=256. +# ============================================================================= + +logging: + wandb: + enabled: false + project: falcon_examples + group: 06_spectral_analysis + dir: ${run_dir} + local: + enabled: true + dir: ${paths.graph} + +paths: + import: "./src" + buffer: ${run_dir}/sim_dir + graph: ${run_dir}/graph_dir + samples: ${run_dir}/samples_dir + +buffer: + min_samples: 32000 + max_samples: 32000 + validation_samples: 256 + simulate_count: 256 + simulate_when_full: false + simulate_interval: 1 + store_fraction: 0.0 + +graph: + theta: + evidence: [x] + + simulator: + _target_: falcon.priors.Product + priors: # Normal, 5σ CRB at N=256 + # - ['normal', 2.7500000000e-3, 4.9540464808e-8] # f0 + # - ['normal', 1.0000000000, 7.0304073517e-5] # chirp_mass + #- ['normal', 1.5000000000, 2.9271998214e-1] # harmonic_decay + - ['normal', 2.7500000200e-3, 4.9540464808e-8] # f0 + - ['normal', 1.0000200000, 7.0304073517e-5] # chirp_mass + - ['normal', 1.6000000000, 2.9271998214e-1] # harmonic_decay + + estimator: + _target_: falcon.estimators.Gaussian + loop: + num_epochs: 500 + batch_size: 256 + early_stop_patience: 16 + network: + hidden_dim: 512 + num_layers: 3 + momentum: 0.1 + min_var: 1.0e-25 + eig_update_freq: 1 + embedding: + _target_: torch.nn.Identity + _input_: x + optimizer: + lr: 0.005 + lr_decay_factor: 0.5 + scheduler_patience: 16 + inference: + gamma: 0.1 + discard_samples: false + log_ratio_threshold: -20.0 + + ray: + num_gpus: 0.3 + + y: + parents: [theta] + simulator: + _target_: model.Signal + N: 256 + ray: + num_gpus: 0.3 + + n: + simulator: + _target_: model.Noise + N: 256 + sigma: 1.0 + + x: + parents: [y, n] + simulator: + _target_: model.Data + observed: "./data/obs_256.npz['x']" + +sample: + prior: + n: 1000 + posterior: + n: 1000 diff --git a/examples/06_spectral_analysis/config4.yml b/examples/06_spectral_analysis/config4.yml new file mode 100644 index 0000000..afb6287 --- /dev/null +++ b/examples/06_spectral_analysis/config4.yml @@ -0,0 +1,112 @@ +# ============================================================================= +# Falcon Example Configuration: 06_spectral_analysis (N=256, SVD embedding) +# ----------------------------------------------------------------------------- +# Like config3 but with scaffolded SVD embedding re-enabled. +# Clean signal (y) and noise (n) are scaffolds — available during training +# for whitening/PCA learning, but not at inference time. +# +# Graph: theta -> y (clean signal) +# \ +# -> x (observed = y + n) +# / +# n (noise) +# +# Prior ranges are Normal, 5σ CRB at N=256. +# ============================================================================= + +logging: + wandb: + enabled: false + project: falcon_examples + group: 06_spectral_analysis + dir: ${run_dir} + local: + enabled: true + dir: ${paths.graph} + +paths: + import: "./src" + buffer: ${run_dir}/sim_dir + graph: ${run_dir}/graph_dir + samples: ${run_dir}/samples_dir + +buffer: + min_samples: 4000 + max_samples: 4000 + validation_samples: 256 + simulate_count: 256 + simulate_when_full: true + simulate_interval: 1 + store_fraction: 0.0 + +graph: + theta: + evidence: [x] + scaffolds: [y, n] + + simulator: + _target_: falcon.priors.Product + priors: # Normal, 5σ CRB at N=256 + # - ['normal', 2.7500000000e-3, 4.9540464808e-8] # f0 + # - ['normal', 1.0000000000, 7.0304073517e-5] # chirp_mass + #- ['normal', 1.5000000000, 2.9271998214e-1] # harmonic_decay + - ['normal', 2.7500000200e-3, 4.9540464808e-8] # f0 + - ['normal', 1.0000200000, 7.0304073517e-5] # chirp_mass + - ['normal', 1.6000000000, 2.9271998214e-1] # harmonic_decay + + estimator: + _target_: falcon.estimators.Gaussian + loop: + num_epochs: 500 + batch_size: 256 + early_stop_patience: 16 + network: + hidden_dim: 512 + num_layers: 3 + momentum: 0.1 + min_var: 1.0e-25 + eig_update_freq: 1 + embedding: + _target_: embeddings.SVDEmbedding + N: 256 + n_components: 64 + buffer_size: 512 + momentum: 0.1 + _input_: [x, y, n] + optimizer: + lr: 0.005 + lr_decay_factor: 0.5 + scheduler_patience: 16 + inference: + gamma: 1.0 + discard_samples: false + log_ratio_threshold: -20.0 + + ray: + num_gpus: 0.3 + + y: + parents: [theta] + simulator: + _target_: model.Signal + N: 256 + ray: + num_gpus: 0.3 + + n: + simulator: + _target_: model.Noise + N: 256 + sigma: 1.0 + + x: + parents: [y, n] + simulator: + _target_: model.Data + observed: "./data/obs_256.npz['x']" + +sample: + prior: + n: 1000 + posterior: + n: 1000 diff --git a/examples/06_spectral_analysis/config_combined.yml b/examples/06_spectral_analysis/config_combined.yml new file mode 100644 index 0000000..c7f8e14 --- /dev/null +++ b/examples/06_spectral_analysis/config_combined.yml @@ -0,0 +1,147 @@ +# ============================================================================= +# Falcon Example Configuration: 06_spectral_analysis (combined variant) +# ----------------------------------------------------------------------------- +# Combines both embedding streams: +# - SVD path: raw signal (x,y,n) -> SVDEmbedding -> (B, 64) +# - Token path: x -> Tokenizer -> SpectralTokenEmbed -> RunningNorm +# -> TransformerEmbedding -> (B, 64) +# Both streams are concatenated via Concat -> (B, 128) -> RunningNorm +# +# Graph: theta -> y (clean signal) +# \ +# -> x (observed = y + n) -> tokens +# / +# n (noise) +# ============================================================================= + +logging: + wandb: + enabled: false + project: falcon_examples + group: 06_spectral_analysis + dir: ${run_dir} + local: + enabled: true + dir: ${paths.graph} + +paths: + import: "./src" + buffer: ${run_dir}/sim_dir + graph: ${run_dir}/graph_dir + samples: ${run_dir}/samples_dir + +buffer: + min_samples: 4092 + max_samples: 4092 + validation_samples: 256 + simulate_count: 256 + simulate_when_full: true + simulate_interval: 1 + store_fraction: 0.01 + +graph: + theta: + evidence: [tokens, x] + scaffolds: [y, n] + sample_chunk_size: 8 + + simulator: + _target_: falcon.priors.Product + priors: # 10x CRB at N=100k + - ['normal', 2.75e-3, 1e-1] # f0 + - ['normal', 1.0, 1e-1] # chirp_mass + - ['normal', 1.5, 1e-1] # harmonic_decay + + estimator: + _target_: falcon.estimators.Gaussian + loop: + num_epochs: 500 + batch_size: 64 + early_stop_patience: 32 + network: + hidden_dim: 512 + num_layers: 3 + momentum: 0.1 + min_var: 1.0e-25 + eig_update_freq: 1 + embedding: + _target_: embeddings.Concat + _input_: + # Stream 1: SVD on raw signal + - _target_: embeddings.SVDEmbedding + N: 100000 + n_components: 64 + buffer_size: 256 + momentum: 0.1 + _input_: [x, y, n] + # Stream 2: Spectral tokens through transformer + - _target_: fuge.nn.TransformerEmbedding + d_model: 64 + n_heads: 4 + n_layers: 2 + d_ff: 256 + dropout: 0.0 + _input_: + _target_: falcon.embeddings.RunningNorm + momentum: 0.01 + dim: [0, 1] + output_dtype: float32 + _input_: + _target_: embeddings.SpectralTokenEmbed + n_freq: 10 + phase_mode: boundary + _input_: [tokens] + optimizer: + lr: 0.005 + lr_decay_factor: 0.5 + scheduler_patience: 16 + inference: + gamma: 1.0 + discard_samples: false + log_ratio_threshold: -20.0 + + ray: + num_gpus: 0.2 + + y: + parents: [theta] + sample_chunk_size: 8 + simulator: + _target_: model.Signal + N: 100000 + ray: + num_gpus: 0.2 + + n: + simulator: + _target_: model.Noise + N: 100000 + sigma: 1.0 + + x: + parents: [y, n] + sample_chunk_size: 8 + simulator: + _target_: model.Data + observed: "./data/obs.npz['x']" + ray: + num_gpus: 0.2 + + tokens: + parents: [x] + sample_chunk_size: 8 + simulator: + _target_: embeddings.Tokenizer + k: 10000 + n_peaks: 3 + n_dlnf: 11 + dlnf_min: 0.0 + dlnf_max: 0.05 + ray: + num_gpus: 0.2 + +sample: + prior: + n: 1000 + posterior: + n: 1000 diff --git a/examples/06_spectral_analysis/data/generate_obs.py b/examples/06_spectral_analysis/data/generate_obs.py new file mode 100644 index 0000000..c752dbf --- /dev/null +++ b/examples/06_spectral_analysis/data/generate_obs.py @@ -0,0 +1,39 @@ +"""Generate a mock chirp observation at known ground-truth parameters. + +Saves the noisy signal to obs.npz for use as the observed data in falcon. + +Ground truth: f0=2.75e-3 Hz, chirp_mass=1.0, harmonic_decay=1.5 +""" + +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) +from chirp import chirp_signal + +TRUE_F0 = 2.75e-3 +TRUE_CHIRP_MASS = 1.0 +TRUE_HARMONIC_DECAY = 1.5 + +signal = chirp_signal( + f0=TRUE_F0, + chirp_mass=TRUE_CHIRP_MASS, + t_c=1e6, + A0=5.0, + harmonic_decay=TRUE_HARMONIC_DECAY, + n_harmonics=4, + N=1000_000, +) + +rng = np.random.default_rng(42) +noise = rng.standard_normal(len(signal)) +observation = signal + noise + +np.savez( + "obs_1M.npz", + x=observation, + true_theta=np.array([TRUE_F0, TRUE_CHIRP_MASS, TRUE_HARMONIC_DECAY]), +) +print(f"Saved obs.npz ({len(observation)} samples, SNR ~ {np.std(signal)/1.0:.1f})") diff --git a/examples/06_spectral_analysis/data/obs.npz b/examples/06_spectral_analysis/data/obs.npz new file mode 100644 index 0000000..e93ce93 Binary files /dev/null and b/examples/06_spectral_analysis/data/obs.npz differ diff --git a/examples/06_spectral_analysis/data/obs_1M.npz b/examples/06_spectral_analysis/data/obs_1M.npz new file mode 100644 index 0000000..460ae26 Binary files /dev/null and b/examples/06_spectral_analysis/data/obs_1M.npz differ diff --git a/examples/06_spectral_analysis/data/obs_256.npz b/examples/06_spectral_analysis/data/obs_256.npz new file mode 100644 index 0000000..ad16efc Binary files /dev/null and b/examples/06_spectral_analysis/data/obs_256.npz differ diff --git a/examples/06_spectral_analysis/make_animation.py b/examples/06_spectral_analysis/make_animation.py new file mode 100644 index 0000000..57d4bb7 --- /dev/null +++ b/examples/06_spectral_analysis/make_animation.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python +"""Generate animated GIF of simulated spectra from buffer dumps. + +Usage: python make_animation.py [RUN_PATH] (default: outputs/latest) + +Reads the buffer dumps (samples_dir/buffer/*.npz), computes the power +spectrum of each signal via FFT, and produces an animated GIF showing +how the simulated spectra evolve as training progresses. +""" + +import sys +from pathlib import Path + +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.animation import FuncAnimation, PillowWriter + +# --------------------------------------------------------------------------- +# Resolve run path +# --------------------------------------------------------------------------- +run_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("outputs/latest") +buffer_dir = run_path / "samples_dir" / "buffer" +if not buffer_dir.exists(): + print(f"No buffer directory found at {buffer_dir}") + sys.exit(1) + +dump_files = sorted(buffer_dir.glob("*.npz")) +n_dumps = len(dump_files) +print(f"Found {n_dumps} buffer dumps in {buffer_dir}") + +# --------------------------------------------------------------------------- +# Load a subset of dumps for the animation (evenly spaced) +# --------------------------------------------------------------------------- +N_FRAMES = min(120, n_dumps) +indices = np.linspace(0, n_dumps - 1, N_FRAMES, dtype=int) +selected_files = [dump_files[i] for i in indices] + +# Load first to get signal length and set up frequency axis +d0 = np.load(selected_files[0]) +signal_len = d0["x.value"].shape[0] +freqs = np.fft.rfftfreq(signal_len) # normalized frequency + +# Precompute all spectra +print("Computing power spectra...") +spectra = [] +thetas = [] +for f in selected_files: + d = np.load(f) + signal = d["x.value"] + theta = d["theta.value"] + # Power spectral density via FFT + fft_vals = np.fft.rfft(signal) + psd = np.abs(fft_vals) ** 2 / signal_len + spectra.append(psd) + thetas.append(theta) + +spectra = np.array(spectra) +thetas = np.array(thetas) + +# Find frequency range with signal (skip DC, focus on where peaks are) +# Use log-scale PSD, find range where max spectrum has interesting content +max_psd = spectra.max(axis=0) +threshold = max_psd.max() * 1e-6 +interesting = np.where(max_psd > threshold)[0] +freq_lo = max(1, interesting[0] - 10) +freq_hi = min(len(freqs) - 1, interesting[-1] + 100) + +# --------------------------------------------------------------------------- +# Build animation +# --------------------------------------------------------------------------- +fig, (ax_spec, ax_params) = plt.subplots(2, 1, figsize=(10, 7), + gridspec_kw={"height_ratios": [3, 1]}) +fig.subplots_adjust(hspace=0.35) + +# Spectrum axes +ax_spec.set_xlim(freqs[freq_lo], freqs[freq_hi]) +psd_slice = spectra[:, freq_lo:freq_hi] +ymin = max(psd_slice[psd_slice > 0].min() * 0.1, 1e-10) +ymax = psd_slice.max() * 10 +ax_spec.set_ylim(ymin, ymax) +ax_spec.set_yscale("log") +ax_spec.set_xlabel("Normalized frequency") +ax_spec.set_ylabel("Power spectral density") +line, = ax_spec.plot([], [], lw=0.5, color="steelblue") +title = ax_spec.set_title("") + +# Parameter scatter axes — show all parameters up to current frame +param_labels = [r"$f_0$", r"$\mathcal{M}$", r"$\alpha$"] +colors = ["#e74c3c", "#2ecc71", "#3498db"] +scatters = [] +for i in range(3): + ax_p = ax_params if i == 0 else ax_params.twinx() + if i == 2: + ax_p.spines["right"].set_position(("axes", 1.08)) + sc = ax_p.scatter([], [], s=8, alpha=0.5, color=colors[i], label=param_labels[i]) + ax_p.set_ylabel(param_labels[i], color=colors[i], fontsize=10) + ax_p.tick_params(axis="y", labelcolor=colors[i]) + scatters.append((ax_p, sc)) + +ax_params.set_xlabel("Sample index") +ax_params.set_xlim(0, n_dumps) +# Set y-ranges from data +for i, (ax_p, sc) in enumerate(scatters): + margin = (thetas[:, i].max() - thetas[:, i].min()) * 0.1 or 1e-8 + ax_p.set_ylim(thetas[:, i].min() - margin, thetas[:, i].max() + margin) + +fig.legend([s for _, s in scatters], param_labels, loc="upper right", + fontsize=8, markerscale=2) + + +def init(): + line.set_data([], []) + title.set_text("") + for _, sc in scatters: + sc.set_offsets(np.empty((0, 2))) + return [line, title] + [sc for _, sc in scatters] + + +def update(frame): + idx = frame + # Update spectrum + line.set_data(freqs[freq_lo:freq_hi], spectra[idx, freq_lo:freq_hi]) + sample_idx = indices[idx] + title.set_text(f"Sample {sample_idx}/{n_dumps - 1} | " + f"$f_0$={thetas[idx, 0]:.6e} " + f"$\\mathcal{{M}}$={thetas[idx, 1]:.6f} " + f"$\\alpha$={thetas[idx, 2]:.4f}") + + # Update parameter scatter — show all points up to current frame + for i, (ax_p, sc) in enumerate(scatters): + offsets = np.column_stack([indices[:idx + 1], thetas[:idx + 1, i]]) + sc.set_offsets(offsets) + + return [line, title] + [sc for _, sc in scatters] + + +print(f"Rendering {N_FRAMES} frames...") +anim = FuncAnimation(fig, update, frames=N_FRAMES, init_func=init, + blit=True, interval=80) + +out_path = run_path / "spectra_animation.gif" +anim.save(str(out_path), writer=PillowWriter(fps=12)) +print(f"Saved animation to {out_path}") +plt.close() diff --git a/examples/06_spectral_analysis/make_fisher_estimate.py b/examples/06_spectral_analysis/make_fisher_estimate.py new file mode 100644 index 0000000..138a2b4 --- /dev/null +++ b/examples/06_spectral_analysis/make_fisher_estimate.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python +"""Compute Fisher/CRB posterior estimate from config and observation. + +Usage: + python make_fisher_estimate.py --obs data/obs.npz + python make_fisher_estimate.py --config-name config2 --obs data/obs.npz + +Reads signal parameters (N, t_c, A0, n_harmonics, noise_sigma) and prior +bounds from a config file, loads ground-truth parameters and observation +from an NPZ file, then computes the Fisher information matrix via JAX +autodiff to produce the Cramér-Rao bound (CRB) posterior estimate. + +Outputs: + - CRB summary table (printed) + - Corner plot of CRB Gaussian samples with prior bounds (saved as PNG) +""" + +import argparse +import sys +from pathlib import Path + +import jax +jax.config.update("jax_enable_x64", True) +import jax.numpy as jnp +import numpy as np +import matplotlib.pyplot as plt +import corner +from omegaconf import OmegaConf + +sys.path.insert(0, str(Path(__file__).resolve().parent / "src")) +from chirp import _chirp_impl + +PARAM_NAMES = [r"$f_0$", r"$\mathcal{M}$", r"$\alpha$"] +PARAM_NAMES_PLAIN = ["f0", "chirp_mass", "harmonic_decay"] + + +def make_fisher_fn(t_c, A0, n_harmonics, N): + """Create a JIT-compiled Fisher matrix function for given signal parameters.""" + T_OBS = 0.9 * t_c + + @jax.jit + def compute_fisher(params, sigma): + def signal(p): + return _chirp_impl(p[0], p[1], t_c, A0, p[2], n_harmonics, N, T_OBS) + + J = jax.jacfwd(signal)(params) # (N, 3) + return J.T @ J / sigma ** 2 # (3, 3) + + return compute_fisher + + +def main(): + parser = argparse.ArgumentParser(description="Compute Fisher/CRB posterior estimate") + parser.add_argument("--config-name", type=str, default="config", + help="Config file stem (default: config -> config.yml)") + parser.add_argument("--obs", type=str, required=True, + help="Path to observation NPZ file (must contain 'true_theta')") + args = parser.parse_args() + + script_dir = Path(__file__).resolve().parent + config_path = script_dir / f"{args.config_name}.yml" + obs_path = Path(args.obs) + + if not config_path.exists(): + sys.exit(f"Config not found: {config_path}") + if not obs_path.exists(): + sys.exit(f"Observation file not found: {obs_path}") + + # ── Load config ────────────────────────────────────────────────── + cfg = OmegaConf.load(config_path) + graph = cfg.graph + + # Signal params from y (Signal node) or x (Simulator node) + if "y" in graph and "Signal" in graph.y.simulator.get("_target_", ""): + sig_cfg = graph.y.simulator + else: + sig_cfg = graph.x.simulator + + N = int(sig_cfg.get("N", 100_000)) + t_c = float(sig_cfg.get("t_c", 1e6)) + A0 = float(sig_cfg.get("A0", 5.0)) + n_harmonics = int(sig_cfg.get("n_harmonics", 4)) + + # Noise sigma from n (Noise node) or x (Simulator node) + if "n" in graph: + sigma = float(graph.n.simulator.get("sigma", 1.0)) + else: + sigma = float(graph.x.simulator.get("noise_sigma", 1.0)) + + # Prior bounds from theta + prior_bounds = [] + for p in graph.theta.simulator.priors: + prior_bounds.append((float(p[1]), float(p[2]))) + + print(f"Config: {config_path.name}") + print(f"Signal: N={N:,}, t_c={t_c:.0e}, A0={A0}, n_harmonics={n_harmonics}") + print(f"Noise sigma: {sigma}") + print(f"Prior bounds: {prior_bounds}") + + # ── Load observation ───────────────────────────────────────────── + obs = np.load(obs_path) + true_theta = obs["true_theta"] + print(f"\nGround truth: f0={true_theta[0]:.6e}, M={true_theta[1]:.6f}, " + f"alpha={true_theta[2]:.4f}") + + # ── Fisher matrix ──────────────────────────────────────────────── + print(f"\nComputing Fisher matrix (JAX autodiff, N={N:,})...") + compute_fisher = make_fisher_fn(t_c, A0, n_harmonics, N) + params_jax = jnp.array(true_theta, dtype=jnp.float64) + F = np.array(compute_fisher(params_jax, sigma)) + cov_crb = np.linalg.inv(F) + std_crb = np.sqrt(np.diag(cov_crb)) + + print(f"\nFisher matrix:\n{F}") + print(f"\nCRB covariance:\n{cov_crb}") + + # ── Summary table ──────────────────────────────────────────────── + print(f"\n{'Parameter':<20} {'True':>14} {'CRB std':>14} " + f"{'Prior width':>14} {'Prior/CRB':>10}") + print("-" * 76) + for i, name in enumerate(PARAM_NAMES_PLAIN): + pw = prior_bounds[i][1] - prior_bounds[i][0] + ratio = pw / std_crb[i] + print(f"{name:<20} {true_theta[i]:>14.6e} {std_crb[i]:>14.6e} " + f"{pw:>14.6e} {ratio:>10.1f}x") + + # ── Corner plot ────────────────────────────────────────────────── + rng = np.random.default_rng(0) + crb_samples = rng.multivariate_normal(true_theta, cov_crb, size=10_000) + + fig = corner.corner( + crb_samples, + labels=PARAM_NAMES, + color="C0", + hist_kwargs=dict(density=True, alpha=0.6), + plot_datapoints=False, + plot_density=False, + levels=(0.68, 0.95), + truths=true_theta, + truth_color="k", + show_titles=True, + ) + + # Show prior bounds as shaded regions on diagonal + axes = np.array(fig.axes).reshape(3, 3) + for i in range(3): + ax = axes[i, i] + lo, hi = prior_bounds[i] + ax.axvspan(lo, hi, alpha=0.15, color="C2", zorder=0) + + # CRB std text in upper-right area + info_lines = [f"N = {N:,}, σ_noise = {sigma}"] + for i, name in enumerate(PARAM_NAMES_PLAIN): + info_lines.append(f"{name} = {true_theta[i]:.6e} (σ_CRB = {std_crb[i]:.4e})") + fig.text(0.98, 0.72, "\n".join(info_lines), fontsize=8, fontfamily="monospace", + verticalalignment="top", horizontalalignment="right", + bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5)) + + from matplotlib.lines import Line2D + from matplotlib.patches import Patch + legend_handles = [ + Line2D([], [], color="C0", label="CRB (Fisher)"), + Patch(facecolor="C2", alpha=0.3, label="Prior range"), + Line2D([], [], color="k", ls="--", label="Ground truth"), + ] + fig.legend(handles=legend_handles, loc="upper right", fontsize=11, frameon=True) + + out_name = f"fisher_crb_{args.config_name}.png" + out_path = script_dir / out_name + fig.savefig(out_path, dpi=150) + print(f"\nSaved {out_path}") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/examples/06_spectral_analysis/make_plots.py b/examples/06_spectral_analysis/make_plots.py new file mode 100644 index 0000000..f3cd28c --- /dev/null +++ b/examples/06_spectral_analysis/make_plots.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python +"""Compare falcon posterior with Fisher/Cramér-Rao analytical bound. + +Usage: python make_plots.py [RUN_PATH] (default: outputs/latest) + +Computes the 3-parameter Fisher information matrix for the EMRI signal +at the ground-truth parameters using JAX autodiff, then plots the falcon +posterior samples alongside the CRB Gaussian in a corner plot. +""" + +import sys +from pathlib import Path + +import jax +jax.config.update("jax_enable_x64", True) +import jax.numpy as jnp +import numpy as np +import matplotlib.pyplot as plt +import corner + +import falcon +sys.path.insert(0, str(Path(__file__).resolve().parent / "src")) +from chirp import _chirp_impl + +PARAM_NAMES = [r"$f_0$", r"$\mathcal{M}$", r"$\alpha$"] +PARAM_UNITS = ["Hz", "", ""] + + +# ===================================================================== +# Fisher information matrix (3-parameter) +# ===================================================================== + +def make_fisher_fn(t_c, A0, n_harmonics, N): + """Create a JIT-compiled Fisher matrix function for given signal parameters.""" + T_OBS = 0.9 * t_c + + @jax.jit + def compute_fisher(params, sigma): + """Full 3x3 Fisher matrix at given parameter values. + + Parameters + ---------- + params : jnp.ndarray, shape (3,) + [f0, chirp_mass, harmonic_decay] + sigma : float + Noise standard deviation. + + Returns + ------- + fisher : jnp.ndarray, shape (3, 3) + """ + def signal(p): + return _chirp_impl(p[0], p[1], t_c, A0, p[2], n_harmonics, N, T_OBS) + + J = jax.jacfwd(signal)(params) # (N, 3) + return J.T @ J / sigma ** 2 # (3, 3) + + return compute_fisher + + +# ===================================================================== +# Main +# ===================================================================== + +def main(): + run_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("outputs/latest") + script_dir = Path(__file__).resolve().parent + + # ── Load falcon posterior ──────────────────────────────────────── + run = falcon.load_run(run_path) + theta = run.samples.posterior.stacked["theta"] # (n_samples, 3) + print(f"Loaded {len(theta)} posterior samples from {run_path}") + + # ── Ground truth ──────────────────────────────────────────────── + obs_path = script_dir / "data" / "obs.npz" + obs = np.load(obs_path) + true_theta = obs["true_theta"] # (3,) + print(f"Ground truth: f0={true_theta[0]:.6e}, M={true_theta[1]:.6f}, " + f"alpha={true_theta[2]:.4f}") + + # ── Signal parameters from config ──────────────────────────────── + sim_cfg = run.config.graph.x.simulator + sigma = float(sim_cfg.get("noise_sigma", 1.0)) + N = int(sim_cfg.get("N", 100_000)) + t_c = float(sim_cfg.get("t_c", 1e6)) + A0 = float(sim_cfg.get("A0", 5.0)) + n_harmonics = int(sim_cfg.get("n_harmonics", 4)) + + # ── Fisher matrix ─────────────────────────────────────────────── + print(f"\nComputing Fisher matrix (JAX autodiff, N={N:,})...") + compute_fisher = make_fisher_fn(t_c, A0, n_harmonics, N) + params_jax = jnp.array(true_theta, dtype=jnp.float64) + F = np.array(compute_fisher(params_jax, sigma)) + cov_crb = np.linalg.inv(F) + std_crb = np.sqrt(np.diag(cov_crb)) + + print(f"\nFisher matrix:\n{F}") + print(f"\nCRB covariance:\n{cov_crb}") + + # ── Comparison table ──────────────────────────────────────────── + std_post = theta.std(axis=0) + mean_post = theta.mean(axis=0) + + print(f"\n{'Parameter':<20} {'True':>12} {'Post. mean':>12} " + f"{'Post. std':>12} {'CRB std':>12} {'Ratio':>8}") + print("-" * 80) + names_plain = ["f0", "chirp_mass", "harmonic_decay"] + for i, name in enumerate(names_plain): + ratio = std_post[i] / std_crb[i] + print(f"{name:<20} {true_theta[i]:>12.6e} {mean_post[i]:>12.6e} " + f"{std_post[i]:>12.6e} {std_crb[i]:>12.6e} {ratio:>8.2f}x") + + # ── Draw CRB Gaussian samples for overlay ─────────────────────── + rng = np.random.default_rng(0) + crb_samples = rng.multivariate_normal(true_theta, cov_crb, size=len(theta)) + + # ── Corner plot ───────────────────────────────────────────────── + fig = corner.corner( + crb_samples, + labels=PARAM_NAMES, + color="C0", + hist_kwargs=dict(density=True, alpha=0.4), + plot_datapoints=False, + plot_density=False, + levels=(0.68, 0.95), + truths=true_theta, + truth_color="k", + show_titles=False, + ) + + corner.corner( + theta, + fig=fig, + color="C1", + hist_kwargs=dict(density=True, alpha=0.6), + plot_datapoints=False, + plot_density=False, + levels=(0.68, 0.95), + ) + + # Legend + from matplotlib.lines import Line2D + legend_handles = [ + Line2D([], [], color="C0", label="CRB (Fisher)"), + Line2D([], [], color="C1", label="Falcon posterior"), + Line2D([], [], color="k", ls="--", label="Ground truth"), + ] + fig.legend(handles=legend_handles, loc="upper right", fontsize=11, + frameon=True) + + out_path = run_path / "fisher_comparison.png" + fig.savefig(out_path, dpi=150) + print(f"\nSaved {out_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/06_spectral_analysis/make_plots2.py b/examples/06_spectral_analysis/make_plots2.py new file mode 100644 index 0000000..5f5c2b2 --- /dev/null +++ b/examples/06_spectral_analysis/make_plots2.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python +"""Compare falcon posterior with Fisher/Cramér-Rao analytical bound. + +For config2.yml (SVD scaffolding variant): signal params are on the y node, +noise sigma on the n node. + +Usage: python make_plots2.py [RUN_PATH] (default: outputs/latest) +""" + +import sys +from pathlib import Path + +import jax +jax.config.update("jax_enable_x64", True) +import jax.numpy as jnp +import numpy as np +import matplotlib.pyplot as plt +import corner + +import falcon +sys.path.insert(0, str(Path(__file__).resolve().parent / "src")) +from chirp import _chirp_impl + +PARAM_NAMES = [r"$f_0$", r"$\mathcal{M}$", r"$\alpha$"] +PARAM_UNITS = ["Hz", "", ""] + + +# ===================================================================== +# Fisher information matrix (3-parameter) +# ===================================================================== + +def make_fisher_fn(t_c, A0, n_harmonics, N): + """Create a JIT-compiled Fisher matrix function for given signal parameters.""" + T_OBS = 0.9 * t_c + + @jax.jit + def compute_fisher(params, sigma): + def signal(p): + return _chirp_impl(p[0], p[1], t_c, A0, p[2], n_harmonics, N, T_OBS) + + J = jax.jacfwd(signal)(params) # (N, 3) + return J.T @ J / sigma ** 2 # (3, 3) + + return compute_fisher + + +# ===================================================================== +# Main +# ===================================================================== + +def main(): + run_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("outputs/latest") + script_dir = Path(__file__).resolve().parent + + # ── Load falcon posterior ──────────────────────────────────────── + run = falcon.load_run(run_path) + theta = run.samples.posterior.stacked["theta"] # (n_samples, 3) + print(f"Loaded {len(theta)} posterior samples from {run_path}") + + # ── Ground truth ──────────────────────────────────────────────── + obs_path = script_dir / "data" / "obs.npz" + obs = np.load(obs_path) + true_theta = obs["true_theta"] # (3,) + print(f"Ground truth: f0={true_theta[0]:.6e}, M={true_theta[1]:.6f}, " + f"alpha={true_theta[2]:.4f}") + + # ── Signal parameters from config (config2 graph structure) ───── + sig_cfg = run.config.graph.y.simulator + N = int(sig_cfg.get("N", 100_000)) + t_c = float(sig_cfg.get("t_c", 1e6)) + A0 = float(sig_cfg.get("A0", 5.0)) + n_harmonics = int(sig_cfg.get("n_harmonics", 4)) + + noise_cfg = run.config.graph.n.simulator + sigma = float(noise_cfg.get("sigma", 1.0)) + + # ── Fisher matrix ─────────────────────────────────────────────── + print(f"\nComputing Fisher matrix (JAX autodiff, N={N:,})...") + compute_fisher = make_fisher_fn(t_c, A0, n_harmonics, N) + params_jax = jnp.array(true_theta, dtype=jnp.float64) + F = np.array(compute_fisher(params_jax, sigma)) + cov_crb = np.linalg.inv(F) + std_crb = np.sqrt(np.diag(cov_crb)) + + print(f"\nFisher matrix:\n{F}") + print(f"\nCRB covariance:\n{cov_crb}") + + # ── Comparison table ──────────────────────────────────────────── + std_post = theta.std(axis=0) + mean_post = theta.mean(axis=0) + + print(f"\n{'Parameter':<20} {'True':>12} {'Post. mean':>12} " + f"{'Post. std':>12} {'CRB std':>12} {'Ratio':>8}") + print("-" * 80) + names_plain = ["f0", "chirp_mass", "harmonic_decay"] + for i, name in enumerate(names_plain): + ratio = std_post[i] / std_crb[i] + print(f"{name:<20} {true_theta[i]:>12.6e} {mean_post[i]:>12.6e} " + f"{std_post[i]:>12.6e} {std_crb[i]:>12.6e} {ratio:>8.2f}x") + + # ── Draw CRB Gaussian samples for overlay ─────────────────────── + rng = np.random.default_rng(0) + crb_samples = rng.multivariate_normal(true_theta, cov_crb, size=len(theta)) + + # ── Corner plot ───────────────────────────────────────────────── + # Use 5-sigma CRB range centered on truth to avoid degenerate axes + plot_range = [ + (true_theta[i] - 5 * std_crb[i], true_theta[i] + 5 * std_crb[i]) + for i in range(len(true_theta)) + ] + + fig = corner.corner( + crb_samples, + labels=PARAM_NAMES, + color="C0", + hist_kwargs=dict(density=True, alpha=0.4), + plot_datapoints=False, + plot_density=False, + levels=(0.68, 0.95), + truths=true_theta, + truth_color="k", + show_titles=False, + range=plot_range, + ) + + corner.corner( + theta, + fig=fig, + color="C1", + hist_kwargs=dict(density=True, alpha=0.6), + plot_datapoints=False, + plot_density=False, + levels=(0.68, 0.95), + range=plot_range, + ) + + # Legend + from matplotlib.lines import Line2D + legend_handles = [ + Line2D([], [], color="C0", label="CRB (Fisher)"), + Line2D([], [], color="C1", label="Falcon posterior"), + Line2D([], [], color="k", ls="--", label="Ground truth"), + ] + fig.legend(handles=legend_handles, loc="upper right", fontsize=11, + frameon=True) + + out_path = run_path / "fisher_comparison.png" + fig.savefig(out_path, dpi=150) + print(f"\nSaved {out_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/06_spectral_analysis/make_prior_samples_plot.py b/examples/06_spectral_analysis/make_prior_samples_plot.py new file mode 100644 index 0000000..c47fe5c --- /dev/null +++ b/examples/06_spectral_analysis/make_prior_samples_plot.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python +"""Plot prior samples: corner plot of theta (top) and signal traces (bottom). + +Usage: + python make_prior_samples_plot.py [--run-dir outputs/latest] + +Loads prior samples from a completed falcon run and produces a single PNG with: + - Upper half: corner plot of the 3 theta parameters (f0, chirp_mass, harmonic_decay) + - Middle: 5 random signal traces in time domain (gray: noisy x, orange: clean y) + - Lower: same 5 samples in Fourier space (power spectral density) +""" + +import argparse +import sys +from pathlib import Path + +import numpy as np +import matplotlib.pyplot as plt +import corner + +import falcon + +PARAM_NAMES = [r"$f_0$", r"$\mathcal{M}$", r"$\alpha$"] + + +def main(): + parser = argparse.ArgumentParser(description="Plot prior samples (theta corner + signal traces)") + parser.add_argument("--run-dir", type=str, default="outputs/latest", + help="Path to the falcon run directory (default: outputs/latest)") + args = parser.parse_args() + + run_path = Path(args.run_dir) + script_dir = Path(__file__).resolve().parent + + # ── Load prior samples ─────────────────────────────────────────── + run = falcon.load_run(run_path) + theta = run.samples.prior.stacked["theta"] # (n_samples, 3) + x_list = run.samples.prior["x"] # list of (N,) arrays + y_list = run.samples.prior["y"] # list of (N,) arrays (clean signal) + print(f"Loaded {len(theta)} prior theta samples from {run_path}") + print(f"Loaded {len(x_list)} prior x samples, {len(y_list)} prior y samples") + + # ── Ground truth (if available) ────────────────────────────────── + obs_path = script_dir / "data" / "obs.npz" + true_theta = None + if obs_path.exists(): + true_theta = np.load(obs_path)["true_theta"] + + # ── Read signal params for time axis ───────────────────────────── + # Signal length from data; time axis from y (Signal) or x (Simulator) config + N_sig = len(x_list[0]) + cfg_graph = run.config.graph + sig_cfg = (cfg_graph.y.simulator if "y" in cfg_graph else cfg_graph.x.simulator) + t_c = float(sig_cfg.get("t_c", 1e6)) + T_obs = 0.9 * t_c + + # ── Select random traces ───────────────────────────────────────── + rng = np.random.default_rng(0) + n_traces = min(1, len(x_list)) + indices = rng.choice(len(x_list), size=n_traces, replace=False) + t = np.linspace(0, T_obs, N_sig) + dt = T_obs / (N_sig - 1) + freqs = np.fft.rfftfreq(N_sig, d=dt) + + # ── Build figure with subfigures ───────────────────────────────── + fig = plt.figure(figsize=(8, 14)) + subfigs = fig.subfigures(3, 1, height_ratios=[3, 1.2, 1.2], hspace=0.05) + + # -- Upper: corner plot (let corner create its own axes) -- + corner.corner( + theta, + labels=PARAM_NAMES, + fig=subfigs[0], + color="C0", + hist_kwargs=dict(density=True, alpha=0.6), + plot_datapoints=False, + plot_density=False, + levels=(0.68, 0.95), + truths=true_theta, + truth_color="k", + show_titles=True, + ) + + # -- Middle: time-domain traces (gray = noisy x, orange = clean y) -- + ax_time = subfigs[1].subplots() + for idx in indices: + ax_time.plot(t, x_list[idx], alpha=0.3, linewidth=0.4, color="gray") + for idx in indices: + ax_time.plot(t, y_list[idx], alpha=0.7, linewidth=0.6, color="C1") + ax_time.set_xlabel("Time (s)") + ax_time.set_ylabel("Amplitude") + ax_time.set_title(f"{n_traces} random prior samples (gray: x, orange: clean signal)") + + # -- Lower: Fourier-domain (same traces) -- + ax_freq = subfigs[2].subplots() + for idx in indices: + amp_x = np.abs(np.fft.rfft(x_list[idx])) / N_sig + ax_freq.plot(freqs, amp_x, alpha=0.3, linewidth=0.4, color="gray") + for idx in indices: + amp_y = np.abs(np.fft.rfft(y_list[idx])) / N_sig + ax_freq.plot(freqs, amp_y, alpha=0.7, linewidth=0.6, color="C1") + ax_freq.set_xlabel("Frequency (Hz)") + ax_freq.set_ylabel(r"Amplitude $|A|$") + ax_freq.set_yscale("log") + ax_freq.set_title("Fourier amplitudes (same samples)") + + # ── Save ───────────────────────────────────────────────────────── + out_path = run_path / "prior_samples.png" + fig.savefig(out_path, dpi=150, bbox_inches="tight") + print(f"Saved {out_path}") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/examples/06_spectral_analysis/make_spectrum_animation.py b/examples/06_spectral_analysis/make_spectrum_animation.py new file mode 100644 index 0000000..4603bc0 --- /dev/null +++ b/examples/06_spectral_analysis/make_spectrum_animation.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python +"""Generate animated GIF of training progression from buffer dumps. + +Usage: python make_spectrum_animation.py [RUN_PATH] (default: outputs/latest) + +Top panel: noise-free power spectrum for the current sample's theta, +with the observation's true spectrum as reference. +Bottom panels: evolving 2D scatter of theta samples with true values +and CRB error ellipses. +""" + +import sys +from pathlib import Path + +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.animation import FuncAnimation, PillowWriter +from matplotlib.patches import Ellipse +sys.path.insert(0, str(Path(__file__).resolve().parent / "src")) +from chirp import chirp_signal + +# --------------------------------------------------------------------------- +# Resolve run path and config +# --------------------------------------------------------------------------- +script_dir = Path(__file__).resolve().parent +run_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("outputs/latest") +buffer_dir = run_path / "samples_dir" / "buffer" +if not buffer_dir.exists(): + print(f"No buffer directory found at {buffer_dir}") + sys.exit(1) + +dump_files = sorted(buffer_dir.glob("*.npz")) +n_dumps = len(dump_files) +print(f"Found {n_dumps} buffer dumps in {buffer_dir}") + +# Signal parameters from config (with defaults matching model.py) +from omegaconf import OmegaConf +cfg = OmegaConf.load(run_path / "config.yml") +sim_cfg = cfg.graph.x.simulator +T_C = float(sim_cfg.get("t_c", 1e6)) +A0 = float(sim_cfg.get("A0", 5.0)) +N_HARMONICS = int(sim_cfg.get("n_harmonics", 4)) +N_SIG = int(sim_cfg.get("N", 100_000)) +NOISE_SIGMA = float(sim_cfg.get("noise_sigma", 1.0)) + +# --------------------------------------------------------------------------- +# Load observation and true spectrum +# --------------------------------------------------------------------------- +obs = np.load(script_dir / "data" / "obs.npz") +true_theta = obs["true_theta"] + +# Compute true (noise-free) spectrum — first call triggers JIT +print("Computing reference spectrum (JIT warmup)...") +true_signal = chirp_signal( + f0=float(true_theta[0]), chirp_mass=float(true_theta[1]), t_c=T_C, A0=A0, + harmonic_decay=float(true_theta[2]), n_harmonics=N_HARMONICS, N=N_SIG, +) +true_signal = np.asarray(true_signal) +freqs = np.fft.rfftfreq(N_SIG) +true_psd = np.abs(np.fft.rfft(true_signal)) ** 2 / N_SIG + +# --------------------------------------------------------------------------- +# Load subset of dumps +# --------------------------------------------------------------------------- +N_FRAMES = min(200, n_dumps) +indices = np.linspace(0, n_dumps - 1, N_FRAMES, dtype=int) +selected_files = [dump_files[i] for i in indices] + +print("Loading theta values...") +all_thetas = np.array([np.load(f)["theta.value"] for f in selected_files]) + +# --------------------------------------------------------------------------- +# Precompute all noise-free spectra +# --------------------------------------------------------------------------- +print(f"Computing {N_FRAMES} noise-free spectra...") +all_psd = np.empty((N_FRAMES, len(freqs))) +for idx in range(N_FRAMES): + theta = all_thetas[idx] + sig = chirp_signal( + f0=float(theta[0]), chirp_mass=float(theta[1]), t_c=T_C, A0=A0, + harmonic_decay=float(theta[2]), n_harmonics=N_HARMONICS, N=N_SIG, + ) + all_psd[idx] = np.abs(np.fft.rfft(np.asarray(sig))) ** 2 / N_SIG +print("Done.") + +# --------------------------------------------------------------------------- +# Find interesting frequency range +# --------------------------------------------------------------------------- +max_psd = np.maximum(all_psd.max(axis=0), true_psd) +threshold = max_psd.max() * 1e-6 +interesting = np.where(max_psd > threshold)[0] +freq_lo = max(1, interesting[0] - 10) +freq_hi = min(len(freqs) - 1, interesting[-1] + 50) +freq_slice = slice(freq_lo, freq_hi) + +# --------------------------------------------------------------------------- +# Try to compute CRB ellipses +# --------------------------------------------------------------------------- +cov_crb = None +try: + import jax + jax.config.update("jax_enable_x64", True) + import jax.numpy as jnp + from chirp import _chirp_impl + + T_OBS = 0.9 * T_C + + @jax.jit + def compute_fisher(params, sigma): + def signal(p): + return _chirp_impl(p[0], p[1], T_C, A0, p[2], N_HARMONICS, N_SIG, T_OBS) + J = jax.jacfwd(signal)(params) + return J.T @ J / sigma ** 2 + + print("Computing Fisher matrix for CRB ellipses...") + F = np.array(compute_fisher(jnp.array(true_theta, dtype=jnp.float64), NOISE_SIGMA)) + cov_crb = np.linalg.inv(F) +except Exception as e: + print(f"Skipping CRB ellipses: {e}") + +# --------------------------------------------------------------------------- +# Build figure +# --------------------------------------------------------------------------- +PAIRS = [(0, 1), (0, 2), (1, 2)] +PAIR_LABELS = [ + (r"$f_0$", r"$\mathcal{M}$"), + (r"$f_0$", r"$\alpha$"), + (r"$\mathcal{M}$", r"$\alpha$"), +] + +fig = plt.figure(figsize=(14, 8)) +gs = fig.add_gridspec(2, 3, height_ratios=[1.2, 1], hspace=0.35, wspace=0.35, + left=0.07, right=0.97, top=0.93, bottom=0.08) + +# -- Top: spectrum (full width) -- +ax_spec = fig.add_subplot(gs[0, :]) +ax_spec.set_xlim(freqs[freq_lo], freqs[freq_hi]) + +psd_slice = np.concatenate([all_psd[:, freq_slice], true_psd[freq_slice][None, :]]) +ymin = max(psd_slice[psd_slice > 0].min() * 0.3, 1e-12) +ymax = psd_slice.max() * 5 +ax_spec.set_ylim(ymin, ymax) +ax_spec.set_yscale("log") +ax_spec.set_xlabel("Normalized frequency") +ax_spec.set_ylabel("Power spectral density") + +# Reference spectrum (true theta) +ax_spec.plot(freqs[freq_slice], true_psd[freq_slice], color="k", + lw=1.0, alpha=0.4, label="True signal", zorder=1) +# Current sample spectrum (animated) +line, = ax_spec.plot([], [], lw=1.0, color="C1", zorder=2, label="Sample") +ax_spec.legend(loc="upper right", fontsize=9) +title = ax_spec.set_title("") + +# -- Bottom: parameter scatter plots -- +ax_params = [fig.add_subplot(gs[1, i]) for i in range(3)] + + +def draw_crb_ellipse(ax, pair_idx): + if cov_crb is None: + return + i, j = PAIRS[pair_idx] + sub_cov = cov_crb[np.ix_([i, j], [i, j])] + eigvals, eigvecs = np.linalg.eigh(sub_cov) + angle = np.degrees(np.arctan2(eigvecs[1, 0], eigvecs[0, 0])) + for n_sigma in [1, 2]: + w = 2 * n_sigma * np.sqrt(eigvals[0]) + h = 2 * n_sigma * np.sqrt(eigvals[1]) + ell = Ellipse(xy=(true_theta[i], true_theta[j]), width=w, height=h, + angle=angle, edgecolor="C0", facecolor="none", + ls="--", lw=1.2, alpha=0.6) + ax.add_patch(ell) + + +param_scatters = [] +for pi, ax in enumerate(ax_params): + i, j = PAIRS[pi] + sc = ax.scatter([], [], s=4, alpha=0.4, color="C1", rasterized=True) + ax.axvline(true_theta[i], color="k", ls=":", lw=0.8, alpha=0.5) + ax.axhline(true_theta[j], color="k", ls=":", lw=0.8, alpha=0.5) + ax.plot(true_theta[i], true_theta[j], "k+", ms=10, mew=1.5, zorder=5) + draw_crb_ellipse(ax, pi) + ax.set_xlabel(PAIR_LABELS[pi][0]) + ax.set_ylabel(PAIR_LABELS[pi][1]) + + margin_x = max((all_thetas[:, i].max() - all_thetas[:, i].min()) * 0.15, 1e-10) + margin_y = max((all_thetas[:, j].max() - all_thetas[:, j].min()) * 0.15, 1e-10) + ax.set_xlim(all_thetas[:, i].min() - margin_x, all_thetas[:, i].max() + margin_x) + ax.set_ylim(all_thetas[:, j].min() - margin_y, all_thetas[:, j].max() + margin_y) + param_scatters.append(sc) + +# --------------------------------------------------------------------------- +# Animation +# --------------------------------------------------------------------------- + +def init(): + line.set_data([], []) + for sc in param_scatters: + sc.set_offsets(np.empty((0, 2))) + title.set_text("") + return [line, title] + param_scatters + + +def update(frame): + # Spectrum + line.set_data(freqs[freq_slice], all_psd[frame, freq_slice]) + + # Parameter scatter — accumulate + for pi, sc in enumerate(param_scatters): + i, j = PAIRS[pi] + offsets = np.column_stack([all_thetas[:frame + 1, i], + all_thetas[:frame + 1, j]]) + sc.set_offsets(offsets) + + theta = all_thetas[frame] + title.set_text( + f"Sample {indices[frame]}/{n_dumps - 1} " + f"$f_0$={theta[0]:.6e} " + f"$\\mathcal{{M}}$={theta[1]:.6f} " + f"$\\alpha$={theta[2]:.4f}" + ) + return [line, title] + param_scatters + + +print(f"Rendering {N_FRAMES} frames...") +anim = FuncAnimation(fig, update, frames=N_FRAMES, init_func=init, + blit=True, interval=60) + +out_path = run_path / "animation.gif" +anim.save(str(out_path), writer=PillowWriter(fps=15)) +print(f"Saved animation to {out_path}") +plt.close() diff --git a/examples/06_spectral_analysis/src/chirp.py b/examples/06_spectral_analysis/src/chirp.py new file mode 100644 index 0000000..f47989a --- /dev/null +++ b/examples/06_spectral_analysis/src/chirp.py @@ -0,0 +1,105 @@ +"""Multi-harmonic chirp signal generator. + +Auto-differentiable w.r.t. all continuous parameters via JAX. + +Parameters +---------- +f0 : initial frequency (Hz) +chirp_mass : controls how fast the frequency sweeps up (arbitrary units) +t_c : coalescence time (s); must be > T_obs +A0 : overall amplitude scale +harmonic_decay : exponential decay rate for higher harmonics; + amplitude of k-th harmonic ~ exp(-harmonic_decay * (k - 1)) + +The waveform is built as: + h(t) = sum_{k=1}^{n_harmonics} A_k(t) * cos(k * phi(t)) + +where phi(t) = 2*pi * cumulative_trapezoid(f(t)) and + + f(t) = f0 * (1 - t / t_c)^(-3/8 * chirp_mass) (PN-inspired chirp) + A(t) = A0 * (1 - t / t_c)^(-1/4) (amplitude growth) +""" + +import functools + +import jax +import jax.numpy as jnp +import numpy as np + + +def chirp_signal( + f0: float, + chirp_mass: float, + t_c: float, + A0: float, + harmonic_decay: float, + n_harmonics: int = 4, + N: int = 1_000_000, + T_obs: float | None = None, +) -> np.ndarray: + """Generate a multi-harmonic chirp signal. + + Parameters + ---------- + f0 : float + Initial frequency (Hz). + chirp_mass : float + Chirp mass parameter controlling frequency evolution rate. + Higher values produce faster chirps. Units are arbitrary; + it enters as f(t) = f0 * (1 - t / t_c)^(-3/8 * chirp_mass). + t_c : float + Time of coalescence (s). Must exceed T_obs. + A0 : float + Overall amplitude scale. + harmonic_decay : float + Exponential decay factor for harmonic amplitudes. + Harmonic k has amplitude A(t) * exp(-harmonic_decay * (k-1)). + n_harmonics : int + Number of harmonics to include (1 = fundamental only). + N : int + Number of time-domain samples. + T_obs : float or None + Observation duration (s). Defaults to 0.9 * t_c. + + Returns + ------- + h : numpy.ndarray, shape (N,) + Time-domain signal. + """ + if T_obs is None: + T_obs = 0.9 * t_c + h = _chirp_impl(f0, chirp_mass, t_c, A0, harmonic_decay, n_harmonics, N, T_obs) + return np.asarray(h) + + +@functools.partial(jax.jit, static_argnums=(5, 6)) +def _chirp_impl(f0, chirp_mass, t_c, A0, harmonic_decay, n_harmonics, N, T_obs): + """JIT-compiled core. n_harmonics and N are static (integers).""" + t = jnp.linspace(0.0, T_obs, N) + dt = T_obs / (N - 1) + + # Dimensionless time to coalescence: tau in (0, 1] + tau = 1.0 - t / t_c + + # --- Frequency evolution (PN-inspired) --- + freq_exponent = -3.0 / 8.0 * chirp_mass + f_t = f0 * tau ** freq_exponent + + # --- Amplitude evolution --- + A_t = A0 * tau ** (-0.25) + + # --- Phase via trapezoidal cumulative integration --- + trapezoids = (f_t[:-1] + f_t[1:]) * 0.5 * dt + phase = jnp.concatenate([jnp.zeros(1), jnp.cumsum(trapezoids)]) + phase = 2.0 * jnp.pi * phase + + # --- Sum over harmonics --- + def _add_harmonic(h, k): + amp_k = A_t * jnp.exp(-harmonic_decay * (k - 1)) + h = h + amp_k * jnp.cos(k * phase) + return h, None + + ks = jnp.arange(1, n_harmonics + 1) + h, _ = jax.lax.scan(_add_harmonic, jnp.zeros(N), ks) + + return h diff --git a/examples/06_spectral_analysis/src/embeddings.py b/examples/06_spectral_analysis/src/embeddings.py new file mode 100644 index 0000000..0b0625b --- /dev/null +++ b/examples/06_spectral_analysis/src/embeddings.py @@ -0,0 +1,202 @@ +"""Embedding modules for the spectral analysis example. + +- SVDEmbedding: Scaffolded streaming SVD embedding using fuge.StreamingPCA +- Tokenizer: Simulator wrapper around fuge.ToneTokenizer for use as a falcon node +- TokenEmbed: Adapter for fuge.ToneTokenEmbedding that returns only the tensor +- SpectralTokenEmbed: Spectral (Fourier feature) embedding for token frequencies +- Concat: Concatenation along last dimension for combining embedding streams +""" + +import math +import torch +import torch.nn as nn +from fuge.spectral import ToneTokenizer, ToneTokenEmbedding + + +class SVDEmbedding(nn.Module): + """Scaffolded SVD embedding for raw signals. + + During training, uses clean signal (y) and noise (n) scaffolds to + learn whitening and PCA. At inference, applies learned projection + to the noisy observation (x) only. + """ + + def __init__(self, N=100_000, n_components=32, buffer_size=256, momentum=0.1): + super().__init__() + from falcon.embeddings import DiagonalWhitener + from fuge.svd import StreamingPCA + self.N = N + self.n_components = n_components + self.whitener = DiagonalWhitener(N, momentum=momentum, track_mean=False) + self.projector = StreamingPCA(n_components=n_components, buffer_size=buffer_size, + momentum=momentum, shrinkage=True) + self.mlp = nn.Sequential( + nn.LazyLinear(128), nn.ReLU(), + nn.LazyLinear(64), + ) + self.double() + + def forward(self, x, y=None, n=None): + x = x.double() + if y is not None and self.training: + y, n = y.double(), n.double() + self.whitener.update(y + n) + white_y = self.whitener(y) + self.projector.update(white_y) + x1 = self.whitener(x) + coeffs = self.projector(x1) # (batch, n_components) + return self.mlp(coeffs) + + +class Tokenizer: + """Simulator wrapper around fuge.ToneTokenizer for use as a falcon node. + + Takes raw signals (numpy) and returns spectral tokens (numpy). + Used as a deterministic intermediate node between x and theta. + """ + + def __init__(self, k=1024, n_peaks=3, n_dlnf=11, dlnf_min=0.0, dlnf_max=0.05): + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.tokenizer = ToneTokenizer(k=k, n_peaks=n_peaks, n_dlnf=n_dlnf, + dlnf_min=dlnf_min, dlnf_max=dlnf_max + ).to(self.device).double() + + def simulate_batch(self, batch_size, x): + x_tensor = torch.as_tensor(x, dtype=torch.float64, device=self.device) + with torch.no_grad(): + tokens = self.tokenizer(x_tensor) + return tokens.cpu().numpy() + + +class TokenEmbed(nn.Module): + """Adapter for fuge.ToneTokenEmbedding that returns only the tensor. + + ToneTokenEmbedding._embed applies cos/sin/log1p feature transforms. + Normalization is handled by an explicit RunningNorm layer in the + embedding pipeline (configured in YAML), keeping feature transforms + separate from normalization — important for sequential SBI where + the data distribution shifts across rounds. + + Args: + phase_mode: "center" (n_embed=5) or "boundary" (n_embed=7). + mask_phases: Zero out phase features (for ablation). + """ + + def __init__(self, phase_mode="center", mask_phases=False): + super().__init__() + self.embed = ToneTokenEmbedding(phase_mode=phase_mode, mask_phases=mask_phases).double() + + def forward(self, raw_tokens): + """Embed raw tokens and return flattened sequence. + + Signal processing (feature extraction) runs in float64. + Output remains float64; dtype conversion is handled by + the downstream RunningNorm layer (output_dtype config). + + Args: + raw_tokens: Tensor of shape (B, W, K, 5) from ToneTokenizer. + + Returns: + Tensor of shape (B, W*K, n_embed), float64. + """ + raw_tokens = raw_tokens.detach().double() + embedded = self.embed._embed(raw_tokens) # (B, W, K, n_embed) + B, W, K, n_embed = embedded.shape + return embedded.reshape(B, W * K, n_embed) + + +class SpectralTokenEmbed(nn.Module): + """Spectral (Fourier feature) embedding for tone tokens. + + Replaces z-score normalized raw frequencies with multi-scale Fourier + features, enabling the transformer to distinguish frequencies at + multiple resolutions — from coarse harmonic structure down to fine + spectral splitting. + + Feature layout per token: + - f_start spectral: sin/cos at n_freq scales -> 2 * n_freq + - f_end spectral: sin/cos at n_freq scales -> 2 * n_freq + - amplitude: log1p(amp) -> 1 + - phase_start: cos(ps), sin(ps) -> 2 + - phase_end: cos(pe), sin(pe) -> 2 + Total n_embed = 4 * n_freq + 5 + + Args: + n_freq: Number of Fourier scales for frequency embedding. + Scales are pi * 2^i for i = 0..n_freq-1. Default 8. + phase_mode: "boundary" keeps both phase endpoints (default), + "center" averages them (n_embed = 4 * n_freq + 3). + """ + + def __init__(self, n_freq=8, phase_mode="boundary", amp_scale=10.0): + super().__init__() + self.n_freq = n_freq + self.phase_mode = phase_mode + self.amp_scale = amp_scale + # Geometrically spaced scales: pi, 2pi, 4pi, ..., pi*2^(n_freq-1) + scales = math.pi * (2.0 ** torch.arange(n_freq, dtype=torch.float64)) + self.register_buffer("scales", scales) + # raw f_start + f_end (2) + spectral (4*n_freq) + amp (1) + phases + time (1) + self.n_embed = 4 * n_freq + (8 if phase_mode == "boundary" else 6) + + def _freq_embed(self, f): + """Fourier features for a frequency value in [-1, 1]. + + Args: + f: (...) shaped tensor. + Returns: + (..., 2*n_freq) shaped tensor of [sin, cos] pairs. + """ + # f: (...) -> (..., n_freq) + scaled = f.unsqueeze(-1) * self.scales + return torch.cat([torch.sin(scaled), torch.cos(scaled)], dim=-1) + + def forward(self, raw_tokens): + """Embed raw tokens with spectral frequency features. + + Args: + raw_tokens: (B, W, K, 5) from ToneTokenizer. + Features: f_start[-1,1], f_end[-1,1], amp, ps[-pi,pi], pe[-pi,pi] + + Returns: + (B, W*K, n_embed) float64 tensor. + """ + raw_tokens = raw_tokens.detach().double() + B, W, K, _ = raw_tokens.shape + + f_start = raw_tokens[..., 0] + f_end = raw_tokens[..., 1] + amp = torch.log1p(raw_tokens[..., 2]) / self.amp_scale + ps = raw_tokens[..., 3] + pe = raw_tokens[..., 4] + + t = torch.linspace(-1.0, 1.0, W, device=raw_tokens.device, dtype=raw_tokens.dtype) + t = t[None, :, None].expand(B, W, K) # (B, W, K) + + parts = [ + f_start.unsqueeze(-1), # (B,W,K, 1) raw frequency + f_end.unsqueeze(-1), # (B,W,K, 1) raw frequency + self._freq_embed(f_start), # (B,W,K, 2*n_freq) + self._freq_embed(f_end), # (B,W,K, 2*n_freq) + amp.unsqueeze(-1), # (B,W,K, 1) + t.unsqueeze(-1), # (B,W,K, 1) time position + ] + + if self.phase_mode == "boundary": + parts += [ + torch.stack([torch.cos(ps), torch.sin(ps)], dim=-1), + torch.stack([torch.cos(pe), torch.sin(pe)], dim=-1), + ] + else: + phi = (ps + pe) / 2 + parts.append(torch.stack([torch.cos(phi), torch.sin(phi)], dim=-1)) + + embedded = torch.cat(parts, dim=-1) # (B, W, K, n_embed) + return embedded.reshape(B, W * K, self.n_embed) + + +class Concat(nn.Module): + """Concatenate multiple inputs along the last dimension.""" + + def forward(self, *inputs): + return torch.cat(inputs, dim=-1) diff --git a/examples/06_spectral_analysis/src/model.py b/examples/06_spectral_analysis/src/model.py new file mode 100644 index 0000000..065060f --- /dev/null +++ b/examples/06_spectral_analysis/src/model.py @@ -0,0 +1,66 @@ +""" +Simulator components for the spectral analysis example. + + - Signal: noise-free chirp waveform generation via chirp._chirp_impl + - Noise: independent Gaussian noise generator + - Data: combiner node (x = y + n) +""" + +import functools + +import jax +import jax.numpy as jnp +import numpy as np +from chirp import _chirp_impl + + +class Signal: + """Noise-free chirp signal generator. + + Generates chirping multi-harmonic signals parameterized by + (f0, chirp_mass, harmonic_decay). Used as intermediate node: theta → y. + """ + + def __init__(self, N=100_000, t_c=1e6, A0=5.0, n_harmonics=4): + self.N = N + self.t_c = t_c + self.A0 = A0 + self.n_harmonics = n_harmonics + + @functools.partial(jax.jit, static_argnums=(3, 4)) + def _generate_clean(f0, chirp_mass, harmonic_decay, + n_harmonics, N, t_c, A0, T_obs): + return jax.vmap( + lambda f, m, d: _chirp_impl(f, m, t_c, A0, d, n_harmonics, N, T_obs) + )(f0, chirp_mass, harmonic_decay) + + self._generate_clean = _generate_clean + + def simulate_batch(self, batch_size, theta): + T_obs = 0.9 * self.t_c + signals = self._generate_clean( + jnp.asarray(theta[:, 0]), jnp.asarray(theta[:, 1]), + jnp.asarray(theta[:, 2]), + self.n_harmonics, self.N, self.t_c, self.A0, T_obs, + ) + return np.asarray(signals, dtype=np.float64) + + +class Noise: + """Independent Gaussian noise generator (no parents).""" + + def __init__(self, N=100_000, sigma=1.0): + self.N = N + self.sigma = sigma + + def simulate(self): + return np.float64(np.random.randn(self.N) * self.sigma) + + +class Data: + """Combiner node: x = y + n.""" + + def simulate(self, y, n): + return y + n + + diff --git a/examples/06_spectral_analysis/summarize_posterior.py b/examples/06_spectral_analysis/summarize_posterior.py new file mode 100644 index 0000000..94e2f5f --- /dev/null +++ b/examples/06_spectral_analysis/summarize_posterior.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python +"""Print mean and std of posterior theta samples. + +Usage: + python summarize_posterior.py [--run-dir outputs/latest] +""" + +import argparse +from pathlib import Path + +import numpy as np +import falcon + +PARAM_NAMES = ["f0", "chirp_mass", "harmonic_decay"] + + +def main(): + parser = argparse.ArgumentParser(description="Summarize posterior samples") + parser.add_argument("--run-dir", type=str, default="outputs/latest") + args = parser.parse_args() + + run = falcon.load_run(args.run_dir) + theta = run.samples.posterior.stacked["theta"] + print(f"Loaded {len(theta)} posterior samples from {args.run_dir}") + + # Ground truth if available + script_dir = Path(__file__).resolve().parent + obs_path = script_dir / "data" / "obs.npz" + true_theta = None + if obs_path.exists(): + true_theta = np.load(obs_path)["true_theta"] + + mean = theta.mean(axis=0) + std = theta.std(axis=0) + + header = f"{'Parameter':<20} {'Mean':>14} {'Std':>14}" + if true_theta is not None: + header += f" {'True':>14} {'Bias (σ)':>10}" + print(header) + print("-" * len(header)) + + for i, name in enumerate(PARAM_NAMES): + line = f"{name:<20} {mean[i]:>14.6e} {std[i]:>14.6e}" + if true_theta is not None: + bias = (mean[i] - true_theta[i]) / std[i] + line += f" {true_theta[i]:>14.6e} {bias:>+10.2f}" + print(line) + + +if __name__ == "__main__": + main()