gam-sae 0.3.148

Sparse-autoencoder latent-manifold terms for the gam penalized-likelihood engine
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
//! Tests for the in-frame curved cascade (`inframe_curved.rs`): planted low-rank
//! curved recovery at `p = 2048` with memory orders below the dense path, parity
//! with the full-`p` fit when the frame contains the truth, and rejection of a
//! region with no curved structure.

use ndarray::{Array2, Array3};

use super::atom::{SaeAtomBasisKind, SaeManifoldAtom};
use super::inframe_curved::{
    ChartOccupancyStatus, CurvedRegion, InFrameCurvedConfig, WeightFrameOccupancy,
    activate_residual_frame, dense_ambient_radial_reference, fit_inframe_curved_regions,
    fit_inframe_curved_weight_frame_catalog, inframe_curved_region_prediction,
    residual_span_frame,
};
use super::weight_frame_catalog::{
    WeightFrameCatalogConfig, WeightFrameMatrix, WeightFrameSource,
    frame_catalog_from_weight_matrices,
};

/// Deterministic LCG in `[-1, 1)` so the tests are reproducible without an RNG
/// dependency.
struct Lcg(u64);
impl Lcg {
    fn new(seed: u64) -> Self {
        Self(seed | 1)
    }
    fn next_unit(&mut self) -> f64 {
        // Numerical Recipes LCG constants.
        self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
        let bits = (self.0 >> 11) as f64 / (1u64 << 53) as f64; // [0,1)
        2.0 * bits - 1.0
    }
    fn normal(&mut self) -> f64 {
        // Box–Muller from two uniforms in (0,1).
        let u1 = (self.next_unit() * 0.5 + 0.5).max(1.0e-12);
        let u2 = self.next_unit() * 0.5 + 0.5;
        (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
    }
}

/// Build a `p × r` column-orthonormal ambient embedding by Gram–Schmidt on
/// random Gaussian columns.
fn random_orthonormal(p: usize, r: usize, seed: u64) -> Array2<f64> {
    let mut rng = Lcg::new(seed);
    let mut q = Array2::<f64>::zeros((p, r));
    for col in 0..r {
        let mut v: Vec<f64> = (0..p).map(|_| rng.normal()).collect();
        for prev in 0..col {
            let mut dot = 0.0;
            for i in 0..p {
                dot += v[i] * q[[i, prev]];
            }
            for i in 0..p {
                v[i] -= dot * q[[i, prev]];
            }
        }
        let mut norm = 0.0;
        for i in 0..p {
            norm += v[i] * v[i];
        }
        norm = norm.sqrt().max(1.0e-12);
        for i in 0..p {
            q[[i, col]] = v[i] / norm;
        }
    }
    q
}

/// Plant a curved (spherical-shell) structure of intrinsic dimension `r_true`
/// inside a `p`-dimensional residual: latent points on a noisy shell of radius
/// ~1, embedded through a random orthonormal frame, plus tiny ambient noise.
fn planted_curved_residual(
    n: usize,
    p: usize,
    r_true: usize,
    shell_noise: f64,
    ambient_noise: f64,
    seed: u64,
) -> (Array2<f64>, Array2<f64>) {
    let q = random_orthonormal(p, r_true, seed);
    let mut rng = Lcg::new(seed ^ 0x9E3779B97F4A7C15);
    let mut latent = Array2::<f64>::zeros((n, r_true));
    for i in 0..n {
        let mut v: Vec<f64> = (0..r_true).map(|_| rng.normal()).collect();
        let mut norm = 0.0;
        for x in &v {
            norm += x * x;
        }
        norm = norm.sqrt().max(1.0e-12);
        let radius = 1.0 + shell_noise * rng.normal();
        for x in &mut v {
            *x = radius * *x / norm;
        }
        for j in 0..r_true {
            latent[[i, j]] = v[j];
        }
    }
    // Ambient residual = latent @ Qᵀ + ambient noise.
    let mut residual = Array2::<f64>::zeros((n, p));
    for i in 0..n {
        for j in 0..p {
            let mut acc = 0.0;
            for k in 0..r_true {
                acc += latent[[i, k]] * q[[j, k]];
            }
            residual[[i, j]] = acc + ambient_noise * rng.normal();
        }
    }
    (residual, q)
}

fn project_matrix_onto_frame(matrix: &Array2<f64>, frame: &Array2<f64>) -> Array2<f64> {
    frame.dot(&frame.t().dot(matrix))
}

fn relative_frobenius(a: &Array2<f64>, b: &Array2<f64>) -> f64 {
    let mut diff = 0.0;
    let mut denom = 0.0;
    for (x, y) in a.iter().zip(b.iter()) {
        let d = x - y;
        diff += d * d;
        denom += y * y;
    }
    (diff / denom.max(1.0e-30)).sqrt()
}

#[test]
fn weight_frame_catalog_spans_component_column_images() {
    let p = 24;
    let q_ov = random_orthonormal(p, 2, 1001);
    let q_mlp = random_orthonormal(p, 2, 1002);

    let mut w_v = Array2::<f64>::zeros((2, p));
    for j in 0..p {
        w_v[[0, j]] = 0.7 * (j as f64 + 1.0);
        w_v[[1, j]] = if j % 2 == 0 { 1.0 } else { -0.5 };
    }
    let ov = q_ov.dot(&w_v);

    let mut down_coeff = Array2::<f64>::zeros((2, 5));
    for j in 0..5 {
        down_coeff[[0, j]] = 1.0 + j as f64;
        down_coeff[[1, j]] = if j % 2 == 0 { 2.0 } else { -1.0 };
    }
    let w_down = q_mlp.dot(&down_coeff);

    let components = vec![
        WeightFrameMatrix::attention_head_ov(3, 14, q_ov.view(), w_v.view()).expect("OV builds"),
        WeightFrameMatrix::mlp_down_projection(3, w_down.view()),
    ];
    let catalog = frame_catalog_from_weight_matrices(
        &components,
        &WeightFrameCatalogConfig {
            frame_rank_min: 2,
            frame_rank_max: 2,
            ..Default::default()
        },
    )
    .expect("catalog builds");

    assert_eq!(catalog.entries().len(), 2);
    assert_eq!(
        catalog.entries()[0].source,
        WeightFrameSource::AttentionHeadOv { layer: 3, head: 14 }
    );
    assert_eq!(
        catalog.entries()[1].source,
        WeightFrameSource::MlpDownProjection { layer: 3 }
    );

    let ov_frame = catalog.entries()[0].frame.frame().to_owned();
    let mlp_frame = catalog.entries()[1].frame.frame().to_owned();
    let ov_projected = project_matrix_onto_frame(&ov, &ov_frame);
    let mlp_projected = project_matrix_onto_frame(&w_down, &mlp_frame);
    assert!(
        relative_frobenius(&ov_projected, &ov) < 1.0e-10,
        "OV catalog frame must span exactly the OV column image"
    );
    assert!(
        relative_frobenius(&mlp_projected, &w_down) < 1.0e-10,
        "MLP catalog frame must span exactly the W_down column image"
    );
}

#[test]
fn weight_sourced_atom_fit_is_tagged_with_component_source() {
    let n = 240;
    let p = 96;
    let r = 4;
    let (residual, q) = planted_curved_residual(n, p, r, 0.02, 0.0, 1414);
    let mut w_v = Array2::<f64>::zeros((r, p));
    for j in 0..p {
        for k in 0..r {
            w_v[[k, j]] = ((j + 1 + k) as f64).sin();
        }
    }
    let components =
        vec![WeightFrameMatrix::attention_head_ov(2, 14, q.view(), w_v.view()).expect("OV")];
    let catalog = frame_catalog_from_weight_matrices(
        &components,
        &WeightFrameCatalogConfig {
            frame_rank_min: r,
            frame_rank_max: r,
            ..Default::default()
        },
    )
    .expect("catalog");
    let config = InFrameCurvedConfig {
        frame_rank_min: r,
        frame_rank_max: r,
        min_rows: 16,
        alpha: 1.0,
        ..Default::default()
    };
    let result = fit_inframe_curved_weight_frame_catalog(
        residual.view(),
        &catalog,
        &[WeightFrameOccupancy {
            frame_index: 0,
            rows: (0..n).collect(),
            basis_size: 5,
        }],
        n,
        &config,
    )
    .expect("weight-frame fit");

    assert_eq!(result.records.len(), 1);
    assert_eq!(result.records[0].occupancy_status, ChartOccupancyStatus::Occupied);
    assert_eq!(
        result.records[0].frame_source,
        Some(WeightFrameSource::AttentionHeadOv { layer: 2, head: 14 }),
        "occupied atom record must carry native mechanism attribution"
    );
    if !result.curved_prediction.regions().is_empty() {
        assert_eq!(
            result.curved_prediction.regions()[0].frame_source(),
            Some(&WeightFrameSource::AttentionHeadOv { layer: 2, head: 14 })
        );
    }
}

#[test]
fn zero_occupancy_weight_frame_is_reported_chartable_unoccupied() {
    let n = 160;
    let p = 64;
    let r = 3;
    let (residual, q_used) = planted_curved_residual(n, p, r, 0.03, 0.0, 5150);
    let q_unused = random_orthonormal(p, r, 5151);
    let mut coeff = Array2::<f64>::zeros((r, p));
    for j in 0..p {
        for k in 0..r {
            coeff[[k, j]] = ((j + 2 * k + 1) as f64).cos();
        }
    }
    let components = vec![
        WeightFrameMatrix::attention_head_ov(6, 1, q_used.view(), coeff.view()).expect("OV"),
        WeightFrameMatrix::mlp_down_projection(6, q_unused.view()),
    ];
    let catalog = frame_catalog_from_weight_matrices(
        &components,
        &WeightFrameCatalogConfig {
            frame_rank_min: r,
            frame_rank_max: r,
            ..Default::default()
        },
    )
    .expect("catalog");
    let config = InFrameCurvedConfig {
        frame_rank_min: r,
        frame_rank_max: r,
        min_rows: 16,
        alpha: 1.0,
        ..Default::default()
    };
    let result = fit_inframe_curved_weight_frame_catalog(
        residual.view(),
        &catalog,
        &[
            WeightFrameOccupancy {
                frame_index: 0,
                rows: (0..n).collect(),
                basis_size: 4,
            },
            WeightFrameOccupancy {
                frame_index: 1,
                rows: Vec::new(),
                basis_size: 4,
            },
        ],
        n,
        &config,
    )
    .expect("weight-frame atlas");

    assert_eq!(result.records.len(), 2);
    let unoccupied = &result.records[1];
    assert_eq!(
        unoccupied.occupancy_status,
        ChartOccupancyStatus::ChartableUnoccupied,
        "unused weight frame remains in the atlas as chartable but unoccupied here"
    );
    assert_eq!(
        unoccupied.frame_source,
        Some(WeightFrameSource::MlpDownProjection { layer: 6 })
    );
    assert_eq!(unoccupied.evidence.n_rows, 0);
    assert!(!unoccupied.evidence.accepted);
}

#[test]
fn planted_low_rank_curved_recovered_inframe_p2048() {
    let n = 1200;
    let p = 2048;
    let r_true = 6;
    let (residual, _q) = planted_curved_residual(n, p, r_true, 0.02, 0.0, 42);
    let m = 8usize; // atom basis size

    let config = InFrameCurvedConfig {
        frame_rank_min: 2,
        frame_rank_max: 16,
        crossfit_folds: 4,
        min_rows: 32,
        ..Default::default()
    };
    let region = CurvedRegion {
        rows: (0..n).collect(),
        basis_size: m,
    };

    let start = std::time::Instant::now();
    let result = fit_inframe_curved_regions(residual.view(), &[region], n, &config)
        .expect("in-frame curved fit");
    let elapsed = start.elapsed();

    // Frame learned at (or above) the true intrinsic rank, well below p.
    let rec = &result.records[0];
    assert!(
        rec.frame_rank >= r_true && rec.frame_rank <= 16,
        "frame rank {} should recover the intrinsic rank {r_true} within the band",
        rec.frame_rank
    );
    assert!(
        rec.frame_rank < p,
        "frame rank {} must be far below ambient p={p}",
        rec.frame_rank
    );

    // The curved chart beats the linear baseline out-of-sample and is accepted.
    assert!(
        rec.evidence.deviance_gain > 0.0,
        "curved chart should improve held-out deviance, got {}",
        rec.evidence.deviance_gain
    );
    assert_eq!(result.accepted_regions, vec![0], "planted region accepted");

    // MEASURED memory ledger: border and covariance orders below the dense path.
    let ledger = &result.ledger;
    assert_eq!(ledger.dense_border_coeffs, m * p);
    assert_eq!(ledger.inframe_border_coeffs, m * rec.frame_rank);
    let border_shrink = ledger.border_shrink();
    let cov_shrink = ledger.cov_shrink();
    assert!(
        border_shrink >= (p as f64 / 16.0),
        "border must shrink by ~p/r; got {border_shrink:.1}x (dense={} inframe={})",
        ledger.dense_border_coeffs,
        ledger.inframe_border_coeffs
    );
    assert!(
        cov_shrink >= 10_000.0,
        "posterior covariance must shrink by (p/r)²; got {cov_shrink:.0}x"
    );
    // Dense per-atom covariance is multi-GB; in-frame is well under a MB.
    assert!(
        ledger.dense_cov_bytes >= 1_000_000_000,
        "dense (M·p)² covariance should be ~GB, got {} bytes",
        ledger.dense_cov_bytes
    );
    assert!(
        ledger.inframe_cov_bytes <= 4_000_000,
        "in-frame (M·r)² covariance should be well under a MB, got {} bytes",
        ledger.inframe_cov_bytes
    );

    // The p=2048 in-frame fit is tractable in wall-clock terms (the whole point).
    assert!(
        elapsed.as_secs_f64() < 30.0,
        "in-frame fit at p={p} took {elapsed:?}; should be fast"
    );

    assert_eq!(result.curved_prediction.n_rows(), n);
    assert_eq!(result.curved_prediction.output_dim(), p);
    assert_eq!(
        result.curved_prediction.inframe_entries(),
        n * rec.frame_rank,
        "hot prediction storage must be the accepted region's N_g x r image"
    );
    assert_eq!(
        result.curved_prediction.accepted_ambient_entries_if_eager(),
        n * p,
        "the eager ambient atom image would have been N_g x p"
    );
    assert!(
        result.curved_prediction.inframe_entries()
            < result.curved_prediction.accepted_ambient_entries_if_eager(),
        "curved prediction must stay in-frame on the hot path"
    );

    // The lifted curved prediction reconstructs the planted shell when explicitly
    // materialized for verification (rows nonzero, radius ~1).
    let pred = result.curved_prediction.materialize_ambient();
    let mut mean_norm = 0.0;
    for i in 0..n {
        let mut ss = 0.0;
        for j in 0..p {
            ss += pred[[i, j]] * pred[[i, j]];
        }
        mean_norm += ss.sqrt();
    }
    mean_norm /= n as f64;
    assert!(
        (mean_norm - 1.0).abs() < 0.25,
        "reconstructed shell radius {mean_norm:.3} should be ~1"
    );
}

#[test]
fn inframe_matches_dense_full_p_when_frame_contains_truth() {
    // Residual lying EXACTLY in an r_true-dim subspace (no ambient noise): the
    // learned frame contains the truth, so the in-frame radial fit and the
    // full-p dense radial fit must agree to machine precision.
    let n = 300;
    let p = 256;
    let r_true = 4;
    let (residual, _q) = planted_curved_residual(n, p, r_true, 0.1, 0.0, 7);

    let config = InFrameCurvedConfig {
        frame_rank_min: r_true,
        frame_rank_max: r_true,
        min_rows: 16,
        ..Default::default()
    };
    let rows: Vec<usize> = (0..n).collect();

    let inframe_pred = inframe_curved_region_prediction(residual.view(), &rows, &config)
        .expect("in-frame prediction")
        .expect("frame learned");
    let rank = inframe_pred.frame_rank();
    assert_eq!(rank, r_true, "frame rank pinned to the true intrinsic rank");
    assert_eq!(
        inframe_pred.inframe_entries(),
        n * r_true,
        "single-region prediction stays in the learned r-frame"
    );

    let dense_pred =
        dense_ambient_radial_reference(residual.view(), config.whitening_ridge).expect("dense fit");
    let inframe_ambient = inframe_pred.materialize_ambient();

    // Both are n×p ambient predictions of the SAME chart; compare Frobenius.
    let mut diff = 0.0;
    let mut denom = 0.0;
    for i in 0..n {
        for j in 0..p {
            let d = inframe_ambient[[i, j]] - dense_pred[[i, j]];
            diff += d * d;
            denom += dense_pred[[i, j]] * dense_pred[[i, j]];
        }
    }
    let rel = (diff / denom.max(1.0e-30)).sqrt();
    assert!(
        rel < 1.0e-6,
        "in-frame and dense full-p radial fits must match on the frame's span; \
         relative Frobenius diff {rel:.3e}"
    );
}

#[test]
fn residual_span_frame_is_the_production_hook_low_rank_and_spans_truth() {
    // The production seam (`residual_span_frame`) must (a) return a frame whose
    // rank is far below p — so a curved atom carrying it flips the arrow-Schur
    // onto its M·r frames_engaged path instead of the dense M·p Hessian — and
    // (b) span the planted subspace, so the in-frame fit loses no signal. The
    // second property is checked structurally: the residual projected onto the
    // frame and lifted back must reconstruct the (exactly low-rank) residual.
    let n = 400;
    let p = 1024;
    let r_true = 6;
    let (residual, _q) = planted_curved_residual(n, p, r_true, 0.05, 0.0, 2130);
    let config = InFrameCurvedConfig {
        frame_rank_min: r_true,
        frame_rank_max: 16,
        min_rows: 16,
        ..Default::default()
    };
    let rows: Vec<usize> = (0..n).collect();

    let frame = residual_span_frame(residual.view(), &rows, &config)
        .expect("frame learns")
        .expect("beneficial low-rank frame exists for a planted low-rank residual");
    let r = frame.rank();
    assert!(
        r >= r_true && r <= 16 && r < p,
        "seam frame rank {r} should recover the intrinsic rank {r_true} and stay far below p={p}"
    );

    // The frame spans the truth: R (U Uᵀ) ≈ R for an exactly-low-rank residual.
    let u = frame.frame().to_owned(); // p × r
    let mut diff = 0.0;
    let mut denom = 0.0;
    for i in 0..n {
        // z_i = R_i · U  (length r); lifted = z_i · Uᵀ  (length p).
        let mut z = vec![0.0; r];
        for (k, zk) in z.iter_mut().enumerate() {
            let mut acc = 0.0;
            for j in 0..p {
                acc += residual[[i, j]] * u[[j, k]];
            }
            *zk = acc;
        }
        for j in 0..p {
            let mut lifted = 0.0;
            for (k, &zk) in z.iter().enumerate() {
                lifted += zk * u[[j, k]];
            }
            let d = residual[[i, j]] - lifted;
            diff += d * d;
            denom += residual[[i, j]] * residual[[i, j]];
        }
    }
    let rel = (diff / denom.max(1e-30)).sqrt();
    assert!(
        rel < 1e-6,
        "seam frame must span the planted subspace (residual reconstructs through U Uᵀ); rel={rel:.3e}"
    );

    // A full-rank (isotropic) residual admits no beneficial low-rank frame, so
    // the seam returns None and the caller correctly leaves that region dense.
    let mut rng = Lcg::new(9001);
    let mut iso = Array2::<f64>::zeros((64, 8));
    for i in 0..64 {
        for j in 0..8 {
            iso[[i, j]] = rng.normal();
        }
    }
    let tight = InFrameCurvedConfig {
        frame_rank_min: 2,
        frame_rank_max: 4,
        rank_cutoff: 1e-9, // count every direction ⇒ numerical rank fills the width
        ..Default::default()
    };
    let iso_rows: Vec<usize> = (0..64).collect();
    let got = residual_span_frame(iso.view(), &iso_rows, &tight).expect("runs");
    // rank_max=4 < p=8 so a frame is still returned, but it must be a strict
    // low-rank projection (r <= 4), never the full width.
    if let Some(f) = got {
        assert!(f.rank() <= 4 && f.rank() < 8, "seam frame must stay strictly low-rank");
    }
}

#[test]
fn activate_residual_frame_installs_factored_decoder_and_engages_frames() {
    // The one-call wiring hook must (a) install a low-rank decoder_frame learned
    // from the residual, and (b) leave the decoder EXACTLY factored as B = C·Uᵀ
    // (B == (B U) Uᵀ) so the factored arrow-Schur C-solve converges — the same
    // invariant maybe_activate_decoder_frame enforces, but with the frame sourced
    // from the residual span (no dense fit).
    let n = 200;
    let p = 128;
    let m = 3usize; // atom basis size
    let r_true = 5;
    let (residual, _q) = planted_curved_residual(n, p, r_true, 0.05, 0.0, 4242);

    // A minimal atom whose decoder is generic full-rank (M×p); activation must
    // project it onto the residual frame.
    let mut rng = Lcg::new(77);
    let mut decoder = Array2::<f64>::zeros((m, p));
    for a in 0..m {
        for j in 0..p {
            decoder[[a, j]] = rng.normal();
        }
    }
    let basis_values = Array2::<f64>::zeros((1, m));
    let basis_jacobian = Array3::<f64>::zeros((1, m, 1));
    let smooth_penalty = Array2::<f64>::eye(m);
    let mut atom = SaeManifoldAtom::new(
        "seam",
        SaeAtomBasisKind::Periodic,
        1,
        basis_values,
        basis_jacobian,
        decoder,
        smooth_penalty,
    )
    .expect("atom builds");
    assert!(atom.decoder_frame.is_none(), "starts on the full-p path");

    let config = InFrameCurvedConfig {
        frame_rank_min: r_true,
        frame_rank_max: 16,
        min_rows: 16,
        ..Default::default()
    };
    let rows: Vec<usize> = (0..n).collect();
    let r = activate_residual_frame(&mut atom, residual.view(), &rows, &config)
        .expect("activation runs")
        .expect("beneficial low-rank frame installed");
    assert!(r >= r_true && r < p, "installed frame rank {r} low-rank vs p={p}");
    let frame = atom.decoder_frame.as_ref().expect("frame installed");
    assert_eq!(frame.rank(), r);

    // Decoder is now exactly factored: B == (B U) Uᵀ (projection is idempotent).
    let u = frame.frame().to_owned();
    let mut reproj = atom.decoder_coefficients.dot(&u).dot(&u.t());
    reproj -= &atom.decoder_coefficients;
    let mut fro = 0.0;
    for v in reproj.iter() {
        fro += v * v;
    }
    assert!(
        fro.sqrt() < 1e-9,
        "activated decoder must satisfy B = (B U) Uᵀ exactly; residual {:.3e}",
        fro.sqrt()
    );
}

#[test]
fn linear_structure_is_not_promoted_to_curved() {
    // The meaningful negative control for the curved gate is LINEAR structure,
    // not isotropic noise. (An isotropic Gaussian blob has genuine radial
    // concentration: its points cluster near a shell of radius ~√r, so the
    // 1-parameter radial chart is a *better* one-dimensional summary than a
    // rank-1 line and the gate honestly accepts it — that is not a false
    // positive, it is real held-out deviance the linear stage did not capture.)
    //
    // What the gate must NOT do is promote structure a LINEAR model already
    // explains. Here the residual is rank-1: every row is a scalar multiple of a
    // single ambient direction (plus negligible noise). Rank-1 PCA reconstructs
    // it exactly; the radial chart, which forces every row onto one radius,
    // destroys the sign/magnitude along the line and reconstructs strictly
    // worse. So the curved refinement loses out-of-sample and must be rejected —
    // the linear lane owns this region.
    // Exactly rank-1 (no ambient noise): with even tiny noise the fine spectral
    // cutoff would count the noise directions and the whitening floor would
    // amplify them into a spurious isotropic blob. An exact rank-1 residual keeps
    // every extra frame direction at zero projection, so the radial chart can
    // only collapse the single real axis onto ±one radius — strictly worse than
    // the exact rank-1 linear reconstruction.
    let n = 800;
    let p = 512;
    let dir = random_orthonormal(p, 1, 314);
    let mut rng = Lcg::new(2130);
    let mut residual = Array2::<f64>::zeros((n, p));
    for i in 0..n {
        let s = rng.normal(); // signed coordinate along the single linear axis
        for j in 0..p {
            residual[[i, j]] = s * dir[[j, 0]];
        }
    }

    let config = InFrameCurvedConfig {
        frame_rank_min: 4,
        frame_rank_max: 8,
        min_rows: 32,
        ..Default::default()
    };
    let region = CurvedRegion {
        rows: (0..n).collect(),
        basis_size: 8,
    };
    let result = fit_inframe_curved_regions(residual.view(), &[region], n, &config)
        .expect("fit runs on linear structure");
    assert!(
        result.accepted_regions.is_empty(),
        "purely linear (rank-1) structure must NOT be promoted to a curved atom; \
         deviance_gain={} margin={}",
        result.records[0].evidence.deviance_gain,
        result.records[0].evidence.margin
    );
}

#[test]
fn ledger_shrink_matches_reviewer_frontier_shape() {
    // The reviewer's headline shape: M=8, p=4096, r=16 ⇒ border shrinks 256×,
    // per-atom covariance shrinks 65536× (8.6 GB → 131 KB). Assert the ledger
    // arithmetic reproduces those numbers exactly for a single accepted region.
    let n = 600;
    let p = 4096;
    let r_target = 16;
    // Plant exactly r_target intrinsic directions so the frame lands at r=16.
    let (residual, _q) = planted_curved_residual(n, p, r_target, 0.02, 0.0, 99);
    let config = InFrameCurvedConfig {
        frame_rank_min: r_target,
        frame_rank_max: r_target,
        min_rows: 32,
        ..Default::default()
    };
    let region = CurvedRegion {
        rows: (0..n).collect(),
        basis_size: 8,
    };
    let result =
        fit_inframe_curved_regions(residual.view(), &[region], n, &config).expect("fit");
    assert_eq!(result.records[0].frame_rank, r_target);
    if result.accepted_regions.is_empty() {
        // Even if the gate is conservative on this synthetic draw, the per-record
        // border/cov arithmetic is what we are asserting; recompute from record.
        let rec = &result.records[0];
        assert_eq!(rec.dense_border_coeffs, 8 * p);
        assert_eq!(rec.inframe_border_coeffs, 8 * r_target);
    } else {
        let ledger = &result.ledger;
        assert_eq!(ledger.dense_border_coeffs, 8 * p);
        assert_eq!(ledger.inframe_border_coeffs, 8 * r_target);
        // (8·4096)² · 8 bytes = 8.59 GB ; (8·16)² · 8 bytes = 131072 bytes.
        assert_eq!(ledger.dense_cov_bytes, (8 * p) * (8 * p) * 8);
        assert_eq!(ledger.inframe_cov_bytes, (8 * r_target) * (8 * r_target) * 8);
        assert_eq!(ledger.inframe_cov_bytes, 131_072);
        assert!((ledger.border_shrink() - 256.0).abs() < 1.0e-9);
        assert!((ledger.cov_shrink() - 65_536.0).abs() < 1.0e-6);
    }
}

#[test]
fn accepted_curved_prediction_hot_path_stays_in_r_frame() {
    let n = 180;
    let p = 128;
    let r_true = 4;
    let (residual, _q) = planted_curved_residual(n, p, r_true, 0.02, 0.0, 974);
    let config = InFrameCurvedConfig {
        frame_rank_min: r_true,
        frame_rank_max: r_true,
        min_rows: 16,
        ..Default::default()
    };
    let region = CurvedRegion {
        rows: (0..n).collect(),
        basis_size: 5,
    };
    let result =
        fit_inframe_curved_regions(residual.view(), &[region], n, &config).expect("fit");
    assert_eq!(result.accepted_regions, vec![0], "planted curved region accepted");
    let prediction = &result.curved_prediction;
    assert_eq!(prediction.regions().len(), 1);
    assert_eq!(prediction.regions()[0].frame_rank(), r_true);
    assert_eq!(
        prediction.inframe_entries(),
        n * r_true,
        "accepted atom image must be stored as N_g x r"
    );
    assert_eq!(
        prediction.accepted_ambient_entries_if_eager(),
        n * p,
        "the forbidden eager atom image would be N_g x p"
    );
    assert!(
        prediction.inframe_entries() * (p / r_true) <= prediction.accepted_ambient_entries_if_eager(),
        "in-frame storage should scale by r instead of p"
    );

    let slice = prediction.materialize_rows(&[0, n / 2, n - 1]);
    assert_eq!(slice.dim(), (3, p));
    let mut slice_energy = 0.0;
    for value in slice.iter() {
        slice_energy += value * value;
    }
    assert!(
        slice_energy > 0.0,
        "ambient lifting happens only for the requested residual slice"
    );
}