//! Unit tests for `lmm::kernel` and `lmm` (dispatch, workspace, BOBYQA config, joint Wald).
use super::kernel::{reml_gradient, reml_hessian, LmmDualScratch, LmmHyperScratch};
use super::*;
use crate::dual::Dual;
use crate::glmm::DerivStatus;
use crate::test_support::{extra_level_of_row, intercept_only_spec, model_atom};
use crate::{Family, Grouping, GroupingRelation, ModelSpec, ReStructure, Sizing};
/// Grammar: `<mult>n<add>`, e.g. `2n1` = 2·n+1, `1.5n1` = ⌈1.5·n⌉+1, `1n2` =
/// n+2. Result clamped to BOBYQA's legal `[n+2, (n+1)(n+2)/2]`: flat constants
/// (and small-n underflow) violate the bounds, so the hook clamps rather than
/// panic deep in `Bobyqa::new`.
#[test]
fn npt_formula_parses_and_clamps() {
assert_eq!(npt_from_formula("2n1", 36), Some(73));
assert_eq!(npt_from_formula("1.5n1", 36), Some(55)); // ⌈54⌉+1
assert_eq!(npt_from_formula("1n2", 36), Some(38));
assert_eq!(npt_from_formula("1.5n1", 2), Some(4)); // ⌈3⌉+1 = 4 = n+2 ✓
assert_eq!(npt_from_formula("3n0", 2), Some(6)); // 6 = (n+1)(n+2)/2 cap
assert_eq!(npt_from_formula("1n0", 3), Some(5)); // clamped up to n+2
assert_eq!(npt_from_formula("500n500", 8), Some(45)); // max_fun grammar reuses
// the parser; the CLAMP is
// npt-specific — see Step 3
assert_eq!(npt_from_formula("garbage", 8), None);
assert_eq!(npt_from_formula("73", 8), None); // flat constants rejected
}
#[test]
fn formula_eval_unclamped() {
assert_eq!(eval_formula("500n500", 8), Some(4500));
assert_eq!(eval_formula("2n1", 36), Some(73));
assert_eq!(eval_formula("n2", 8), None); // mult is mandatory: write 1n2
}
/// Deterministic pseudo-data (NR LCG), uniform in (−1, 1). NR = Press,
/// Teukolsky, Vetterling & Flannery (2007), *Numerical Recipes: The Art of
/// Scientific Computing*, 3rd ed., Cambridge University Press.
fn lcg(state: &mut u64) -> f64 {
*state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
(((*state >> 11) as f64) / ((1u64 << 53) as f64)) * 2.0 - 1.0
}
/// n=48, p=3 (intercept + x1 + x2), 6 clusters,
/// y = 0.5 + 0.4·x1 − 0.2·x2 + u_c + 0.8·e.
fn hand_dataset() -> (Mat<f64>, Vec<f64>, Vec<u32>) {
hand_dataset_sized(48, 6, false)
}
/// `hand_dataset`, generalized over row/cluster count and cluster balance —
/// the dual-call cost table's size axis needs this q_p=1 shape at 3,000/
/// 30,000 rows, and its balanced-collapse comparison needs an unbalanced
/// variant of the same shape (`precompute_balanced_collapse` only arms on a
/// balanced primary). `skew` mirrors the speed grid's own skew recipe
/// (`campaigns/speed-grid/prep.R`): ~20% of clusters ("heavy") draw ~80% of
/// the rows. The balanced branch (`skew = false`) assigns `i % nc` and draws
/// no extra `lcg` values, so `hand_dataset_sized(n, nc, false)`'s output is
/// bit-identical to `hand_dataset()`'s.
fn hand_dataset_sized(n: usize, nc: usize, skew: bool) -> (Mat<f64>, Vec<f64>, Vec<u32>) {
let mut st = 42u64;
let u_c: Vec<f64> = (0..nc).map(|_| 0.6 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 3);
let mut y = vec![0.0f64; n];
let mut ids = vec![0u32; n];
// n_heavy clusters (0..n_heavy) draw 80% of rows, the rest draw 20% --
// clamped strictly below nc so a "light" cluster always exists to fall
// back into. Final index is clamped to nc-1 rather than trusted from the
// float arithmetic, since r can round up to exactly 0.8 or 1.0's boundary.
let n_heavy = if nc < 2 {
0
} else {
((nc as f64 * 0.2).round() as usize).clamp(1, nc - 1)
};
for i in 0..n {
let c = if !skew || nc < 2 {
i % nc
} else {
let r = lcg(&mut st) + 0.5; // in [0, 1)
let raw = if r < 0.8 {
(r / 0.8) * n_heavy as f64
} else {
n_heavy as f64 + ((r - 0.8) / 0.2) * (nc - n_heavy) as f64
};
(raw as usize).min(nc - 1)
};
ids[i] = c as u32;
let x1 = lcg(&mut st);
let x2 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
x[(i, 2)] = x2;
y[i] = 0.5 + 0.4 * x1 - 0.2 * x2 + u_c[c] + 0.8 * lcg(&mut st);
}
(x, y, ids)
}
/// Same quantity, two factorizations — both return
/// log|V| + log|X'V⁻¹X| + (N−P)·log σ̂², so agreement is FP-level
/// (≤ 1e-9 rel), not up-to-a-constant. THE formulation proof, held on
/// every θ probed.
///
/// `PINNED_DEV` values are frozen reference deviances for `hand_dataset`
/// across θ, pinned so the cross-check between the collapse and general
/// REML paths in `reml_deviance` holds against a fixed reference.
#[test]
fn deviance_matches_pinned_across_theta() {
const PINNED_DEV: [(f64, f64); 9] = [
(0.0, -5.133651734340493e1),
(1e-4, -5.133651830952116e1),
(1e-2, -5.134617257834838e1),
(0.1, -5.22466385997638e1),
(0.5, -6.044958474290513e1),
(1.0, -6.2089881344871685e1),
(2.0, -5.8665868559293685e1),
(10.0, -4.390226095246822e1),
(100.0, -2.0933925045275885e1),
];
let (x, y, ids) = hand_dataset();
let mut suff = LmmSuffStats::new(3, 6);
suff.add_rows(x.as_ref(), &y, &ids);
let mut fit = LmmFitScratch::new(3, 6);
let mut fit_c = LmmFitScratch::new(3, 6);
assert!(precompute_balanced_collapse(&suff, &mut fit_c));
for &(theta, dev_ship) in &PINNED_DEV {
let dev_gen = reml_deviance(&[theta], &suff, &mut fit);
assert!(dev_gen.is_finite(), "θ={theta}");
let tol = 1e-9 * dev_ship.abs().max(1.0);
assert!(
(dev_ship - dev_gen).abs() <= tol,
"θ={theta}: pinned {dev_ship} vs general {dev_gen}"
);
// Collapse arm — reassociation band vs the general loop incl. θ=0.
let dev_c = reml_deviance(&[theta], &suff, &mut fit_c);
let band = 1e-9 * dev_gen.abs().max(1.0);
assert!(
(dev_c - dev_gen).abs() <= band,
"θ={theta}: collapse {dev_c} vs general {dev_gen}"
);
}
}
/// All scratch is overwritten per call — re-evaluating a θ after an
/// intervening different-θ call reproduces bit-identical deviance and σ̂².
#[test]
fn reml_deviance_overwrites_state() {
let (x, y, ids) = hand_dataset();
let mut suff = LmmSuffStats::new(3, 6);
suff.add_rows(x.as_ref(), &y, &ids);
let mut fit = LmmFitScratch::new(3, 6);
let dev_a = reml_deviance(&[1.0], &suff, &mut fit);
let sig_a = fit.sigma_sq;
let _ = reml_deviance(&[2.0], &suff, &mut fit);
let dev_b = reml_deviance(&[1.0], &suff, &mut fit);
let sig_b = fit.sigma_sq;
assert_eq!(dev_a, dev_b, "deviance(θ=1) must be reproducible");
assert_eq!(sig_a, sig_b, "σ̂²(θ=1) must be reproducible");
}
/// Correctness prerequisite for workspace reuse across simulation draws:
/// `suff.reset()` followed by a refill on a DIFFERENT dataset (same shape, different `y`)
/// must reproduce a freshly-constructed workspace's fit bit-for-bit. Same
/// buffers + same code path ⇒ identical float reassociation, so the
/// assertion is exact `==`, not a tolerance band. If this fails, `reset()`
/// (src/lmm/kernel.rs) leaves some `LmmSuffStats` field stale across datasets.
#[test]
fn reused_workspace_refill_matches_fresh() {
let (x, y_a, ids) = hand_dataset();
// B: same shape/ids as A, deterministically different y (constant shift
// + a fixed rescale) — not randomized, so the comparison stays exact.
let y_b: Vec<f64> = y_a.iter().map(|&v| 1.7 - 0.3 * v).collect();
let targets: Vec<u32> = vec![1, 2];
// Fresh workspace, fit B directly.
let mut ws_fresh = LmmWorkspace::new(3, 6);
ws_fresh.suff.reset();
ws_fresh.suff.add_rows(x.as_ref(), &y_b, &ids);
let fit_fresh = fit_lmm(&mut ws_fresh, &targets, None);
// Reused workspace: fit A first, reset, refill with B, fit again.
let mut ws_reused = LmmWorkspace::new(3, 6);
ws_reused.suff.reset();
ws_reused.suff.add_rows(x.as_ref(), &y_a, &ids);
let _ = fit_lmm(&mut ws_reused, &targets, None);
ws_reused.suff.reset();
ws_reused.suff.add_rows(x.as_ref(), &y_b, &ids);
let fit_reused = fit_lmm(&mut ws_reused, &targets, None);
assert_eq!(
fit_fresh.deviance, fit_reused.deviance,
"deviance must be bit-identical after reset+refill on new data"
);
assert_eq!(
ws_fresh.fit.betas, ws_reused.fit.betas,
"betas must be bit-identical after reset+refill on new data"
);
assert_eq!(
ws_fresh.fit.var_diag, ws_reused.fit.var_diag,
"var_diag must be bit-identical after reset+refill on new data"
);
}
/// Exercises the plateau policy: a `MaxFunReached` cap-out must still
/// report the honest finite endpoint
/// (β̂/σ̂²/SE/deviance), with `converged = false` and `boundary_hit == 2`
/// (not the accepted-boundary 1). Forces the cap by swapping in a solver
/// whose `max_fun` is the legal minimum (`npt + 1`) — one eval past the
/// initial model build, nowhere near this dataset's optimum — bypassing
/// `LMM_MAX_FUN_FORMULA` entirely so the test carries no process-env race.
#[test]
fn maxfun_cap_reports_honest_endpoint() {
let (x, y, ids) = hand_dataset();
let targets: Vec<u32> = vec![1, 2];
let mut ws = LmmWorkspace::new(3, 6);
ws.suff.add_rows(x.as_ref(), &y, &ids);
let n_theta = ws.theta.len();
let npt = 2 * n_theta + 1; // n_theta == 1 here: PRIMA's minimum npt
let config = {
let mut c = Config::new(n_theta);
c.npt = npt;
c.max_fun = npt + 1;
c
};
ws.solver = Bobyqa::new(n_theta, config).expect("legal minimal config");
let fit = fit_lmm(&mut ws, &targets, None);
assert!(!fit.converged, "capped fit must not report converged");
assert_eq!(
fit.boundary_hit, 2,
"capped fit must not migrate into the accepted-boundary code"
);
assert_eq!(
fit.pinned_components, 0,
"a capped endpoint is a point, not an accepted boundary"
);
assert!(
fit.deviance.is_finite(),
"plateau policy: capped endpoint must report a finite deviance"
);
assert!(
fit.sigma_sq.is_finite(),
"plateau policy: capped endpoint must report a finite sigma_sq"
);
assert!(
fit.joint_t_sq.is_finite(),
"plateau policy: capped endpoint must report a finite joint_t_sq"
);
for &tj in &targets {
assert!(
ws.fit.betas[tj as usize].is_finite(),
"plateau policy: capped endpoint must not NaN-fill beta"
);
}
assert!(fit.n_eval <= npt + 1, "n_eval must reflect the forced cap");
// Pinned values are the deterministic truncated-BOBYQA endpoint (hand_dataset,
// max_fun = npt+1) — a regression lock, not an external oracle: any solver-path
// change that moves the honest cap-out endpoint should fail this test.
let rel = |got: f64, want: f64| (got - want).abs() / want.abs().max(1e-12);
assert!(
rel(fit.deviance, -62.08988134487164) < 1e-6,
"deviance = {}",
fit.deviance
);
assert!(
rel(fit.sigma_sq, 0.16043347869402982) < 1e-6,
"sigma_sq = {}",
fit.sigma_sq
);
assert!(
rel(fit.joint_t_sq, 14.568949550460516) < 1e-6,
"joint_t_sq = {}",
fit.joint_t_sq
);
let want_betas = [
0.4691004480864937,
0.26391548909385104,
-0.33307894295165125,
];
for (j, &wb) in want_betas.iter().enumerate() {
assert!(
rel(ws.fit.betas[j], wb) < 1e-6,
"betas[{j}] = {}, want {}",
ws.fit.betas[j],
wb
);
}
}
/// End-to-end q=1 parity on the hand dataset: the general machine against
/// frozen literal endpoints, at the amended tolerances (rel 1e-4, abs floors
/// β̂ 1e-5 / stat 1e-4 — the measured Brent θ̂-placement-noise floor).
///
/// `PINNED_BETAS`/`PINNED_STATS`/`PINNED_JOINT_T_SQ` are frozen reference
/// outputs on `hand_dataset` (targets `[1, 2]`), pinned so `fit_lmm`'s
/// general θ-search agrees with a fixed reference at the amended tolerances
/// above.
#[test]
fn fit_matches_pinned_q1_endpoint_on_hand_dataset() {
const PINNED_BETAS: [f64; 3] = [
4.699333343472561e-1,
2.5793782757575945e-1,
-3.2576278857950314e-1,
];
const PINNED_STATS: [f64; 2] = [2.3733464018128823e0, 3.177463198588152e0];
const PINNED_JOINT_T_SQ: f64 = 1.3713656562170998e1;
let (x, y, ids) = hand_dataset();
let targets: Vec<u32> = vec![1, 2];
let mut ws = LmmWorkspace::new(3, 6);
ws.suff.add_rows(x.as_ref(), &y, &ids);
let fit = fit_lmm(&mut ws, &targets, None);
assert!(fit.converged);
assert!(fit.boundary_hit <= 1);
for (j, &want) in PINNED_BETAS.iter().enumerate() {
let (a, b) = (want, ws.fit.betas[j]);
let d = (a - b).abs();
assert!(
d <= 1e-5 || d <= 1e-4 * a.abs().max(b.abs()),
"β[{j}]: {a} vs {b}"
);
}
for (idx, &tj) in targets.iter().enumerate() {
let a = PINNED_STATS[idx];
let b = ws.fit.t_sq[tj as usize].sqrt();
let d = (a - b).abs();
assert!(
d <= 1e-4 || d <= 1e-4 * a.abs().max(b.abs()),
"stat[{tj}]: {a} vs {b}"
);
}
let (a, b) = (PINNED_JOINT_T_SQ, fit.joint_t_sq);
let d = (a - b).abs();
assert!(
d <= 1e-4 || d <= 1e-4 * a.abs().max(b.abs()),
"joint: {a} vs {b}"
);
}
/// Deterministic pin: y carries NO between-cluster signal by construction —
/// residuals alternate ±0.8 within each cluster with equal counts, so every
/// cluster's residual sum is exactly 0 and the REML deviance is minimized at
/// θ = 0. The fit must pin (boundary_hit == 1), write θ̂ = exactly 0.0, and
/// count as converged: zero variance is a legitimate boundary optimum, not
/// a failure to fit.
#[test]
fn zero_between_cluster_variance_pins_at_exactly_zero() {
let n = 48usize;
let n_clusters = 6usize;
let mut st = 7u64;
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut ids = vec![0u32; n];
for i in 0..n {
ids[i] = (i % n_clusters) as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
// i/n_clusters cycles 0..8 within each cluster: 4 even, 4 odd ⇒
// the ±0.8 residuals cancel exactly per cluster.
let e = if (i / n_clusters) % 2 == 0 { 0.8 } else { -0.8 };
y[i] = 0.5 + 0.4 * x1 + e;
}
let mut ws = LmmWorkspace::new(2, n_clusters);
ws.suff.add_rows(x.as_ref(), &y, &ids);
let fit = fit_lmm(&mut ws, &[1], None);
assert!(fit.converged);
assert_eq!(fit.boundary_hit, 1);
assert_eq!(ws.theta[0], 0.0, "pin must be exact 0.0, not merely small");
assert!(ws.fit.betas[1].is_finite());
}
/// Rank deficiency is DETECTED, not refused, at kernel level: x2 = 0.1·x1
/// (the scaled-duplicate fixture — exact duplicates can slip through faer's
/// llt grey zone) returns a fit whose `pivot` records the exhaustion and
/// names the offending column.
///
/// This kernel entry point is below the alias gate, which is what actually
/// handles a design like this: through `fit_cold`/`fit_warm` the duplicate
/// column is dropped before the solver runs and the caller gets a clean
/// `p−1` fit with `aliased[2]`, matching R. `fit_lmm` called directly does
/// not NaN-fill — it hands back the numbers together with the statistic
/// that condemns them, and the standard error it reports (~9e7 on a
/// coefficient of 11.5) is truthful.
#[test]
fn rank_deficient_design_is_flagged_not_refused() {
let n = 48usize;
let n_clusters = 6usize;
let mut st = 11u64;
let mut x = Mat::<f64>::zeros(n, 3);
let mut y = vec![0.0f64; n];
let mut ids = vec![0u32; n];
for i in 0..n {
ids[i] = (i % n_clusters) as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
x[(i, 2)] = 0.1 * x1; // 0.1-scaled duplicate → guaranteed non-convergence
y[i] = 0.5 + 0.4 * x1 + 0.8 * lcg(&mut st);
}
let mut ws = LmmWorkspace::new(3, n_clusters);
ws.suff.add_rows(x.as_ref(), &y, &ids);
let fit = fit_lmm(&mut ws, &[1, 2], None);
assert!(
fit.pivot < PIVOT_MIN,
"the duplicate column must be detected, got pivot {}",
fit.pivot
);
assert_eq!(
fit.pivot_col, 2,
"the LATER column of the duplicated pair is the one named"
);
// The SE is what makes the returned numbers safe: Var(β̂₂) ~ 8e15, so
// the coefficient is reported with an error eight orders larger than
// itself. That is the channel a caller reads, and it does not lie.
assert!(
ws.fit.var_diag[2].sqrt() > 1e6 * ws.fit.betas[2].abs(),
"β̂₂ = {} must carry an SE orders above it, got {}",
ws.fit.betas[2],
ws.fit.var_diag[2].sqrt()
);
}
/// A truth-started fit (`theta_start: Some`) reaches the same answer as the
/// blind fit on the same bytes — and Some(0.0) exercises the
/// THETA_TRUTH_FLOOR clamp rather than starting on the 0 boundary.
/// Bands are the amended floors: two BOBYQA runs from different
/// starts each place θ̂ within the rho_end band of the same minimum.
#[test]
fn theta_start_some_matches_blind_fit() {
let (x, y, ids) = hand_dataset();
let targets: Vec<u32> = vec![1, 2];
let mut ws_blind = LmmWorkspace::new(3, 6);
ws_blind.suff.add_rows(x.as_ref(), &y, &ids);
let blind = fit_lmm(&mut ws_blind, &targets, None);
assert!(blind.converged);
for start in [[0.0], [0.6]] {
let mut ws = LmmWorkspace::new(3, 6);
ws.suff.add_rows(x.as_ref(), &y, &ids);
let fit = fit_lmm(&mut ws, &targets, Some(&start));
assert!(fit.converged, "start {start:?}");
for j in 0..3 {
let (a, b) = (ws_blind.fit.betas[j], ws.fit.betas[j]);
let d = (a - b).abs();
assert!(
d <= 1e-5 || d <= 1e-4 * a.abs().max(b.abs()),
"start {start:?} β[{j}]: blind {a} vs started {b}"
);
}
}
}
/// Bounded-allocation warm-path check for `fit_lmm`.
/// Marked #[ignore] because dhat measures
/// process-wide allocations; `alloc_test_guard` serializes test bodies, but
/// libtest's own per-test thread spawn still needs `--test-threads=1`:
/// RAYON_NUM_THREADS=1 cargo test -p glmm --features alloc-tests lmm_fit_warm_path_bounded_alloc -- --ignored --test-threads=1
///
/// BOUND locks the measured warm-path block count. LmmWorkspace itself is
/// allocation-free across fits (Bobyqa::new is the only solver allocation,
/// done once). On the faer kernel the per-call blocks are `chol_lower`'s
/// single `MemBuffer` — ~1 per deviance evaluation (15.1–15.7 evals/fit at
/// rho_end 1e-6, the measured mean; `chol_lower` factors `l_out` in place
/// instead of calling `MatRef::llt`, dropping its own `Mat::zeros`
/// allocation), the same acceptance the shipped path's 26 blocks/call carry;
/// if a future faer version changes its Cholesky internals, update the
/// bound — do not relax it. A hand-rolled owned-kernel replacement for
/// faer's `llt` is not used: its wasm `f64::ln` takes a different ULP path
/// than the native build (the factorization itself is fine), which breaks
/// cross-platform bit-equality. The faer bound stays the locked steady
/// state.
#[cfg(feature = "alloc-tests")]
#[test]
#[ignore]
fn lmm_fit_warm_path_bounded_alloc() {
let _serial = crate::test_support::alloc_test_guard();
const N_CALLS: usize = 100;
const BOUND: u64 = 2700; // Measured 2500 (this machine) — ~25 blocks/fit of faer `llt` internals on the family-blocked q=1 path (one m×m tail llt per eval). `fit_lmm` does not allocate per fit (the diagonal_theta index map is cached once on LmmGroupings; the ranef recovery pass solves in the ranef_ux/ranef_rhs scratch fields), so this count is purely faer's Cholesky internals — faer-version/machine specific. q=1 deviance is byte-identical to the hand-rolled augmented-factor deviance (held by the lmm_parity corpus + golden_rng), so the eval trajectory is unchanged; `chol_lower` factors `l_out` in place, one heap block per eval instead of two, which keeps this count low. If faer changes its Cholesky internals, update — do not relax.
let (x, y, ids) = hand_dataset();
let targets: Vec<u32> = vec![1, 2];
let mut ws = LmmWorkspace::new(3, 6);
// Warmup drives one-time setup outside the profiler window.
ws.suff.reset();
ws.suff.add_rows(x.as_ref(), &y, &ids);
let _ = fit_lmm(&mut ws, &targets, None);
let profiler = dhat::Profiler::builder().testing().build();
for _ in 0..N_CALLS {
ws.suff.reset();
ws.suff.add_rows(x.as_ref(), &y, &ids);
let _ = fit_lmm(&mut ws, &targets, None);
}
let stats = dhat::HeapStats::get();
drop(profiler);
assert!(
stats.total_blocks <= BOUND,
"fit_lmm allocated {} blocks across {} warm-path calls (BOUND = {})",
stats.total_blocks,
N_CALLS,
BOUND
);
}
// -----------------------------------------------------------------------
// Multi-grouping: layout-true datasets, suff-stats, family-blocked
// deviance vs a brute-force n×n oracle, and end-to-end fits.
// -----------------------------------------------------------------------
/// Layout-true multi-grouping dataset: primary S=6, crossed I=4, nested
/// np=2 (optional), p=3, n = n_blocks·atom rows. Ids come from the
/// contract layout helpers — the same functions the workspace uses.
#[allow(clippy::type_complexity)]
fn multi_dataset(
with_nested: bool,
n_blocks: usize,
) -> (Mat<f64>, Vec<f64>, Vec<u32>, Vec<Vec<u32>>, ModelSpec) {
let mut cluster = intercept_only_spec(Sizing::FixedClusters { n_clusters: 6 });
cluster.re.as_mut().unwrap().extra_groupings.push(Grouping {
relation: GroupingRelation::Crossed { n_clusters: 4 },
slopes: vec![],
});
if with_nested {
cluster.re.as_mut().unwrap().extra_groupings.push(Grouping {
relation: GroupingRelation::NestedWithin { n_per_parent: 2 },
slopes: vec![],
});
}
let n = n_blocks * model_atom(&cluster);
let mut st = 99u64;
let u_p: Vec<f64> = (0..6).map(|_| 0.5 * lcg(&mut st)).collect();
let u_x: Vec<f64> = (0..4).map(|_| 0.4 * lcg(&mut st)).collect();
let u_n: Vec<f64> = (0..12).map(|_| 0.3 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 3);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let n_extras = cluster.re.as_ref().unwrap().extra_groupings.len();
let mut eids: Vec<Vec<u32>> = vec![vec![0u32; n]; n_extras];
for i in 0..n {
pid[i] = cluster.re.as_ref().unwrap().sizing.cluster_of_row(i) as u32;
#[allow(clippy::needless_range_loop)]
for g in 0..n_extras {
eids[g][i] = extra_level_of_row(&cluster, g, i) as u32;
}
let x1 = lcg(&mut st);
let x2 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
x[(i, 2)] = x2;
y[i] = 0.5 + 0.4 * x1 - 0.2 * x2
+ u_p[pid[i] as usize]
+ u_x[eids[0][i] as usize]
+ if with_nested {
u_n[eids[1][i] as usize]
} else {
0.0
}
+ 0.8 * lcg(&mut st);
}
(x, y, pid, eids, cluster)
}
/// `diagonal_theta` / `n_theta` / `k_family` at q_p ∈ {1, 2, 3} — locks
/// the column-major vech ordering. q_p=1 must reproduce the intercept-only
/// baseline values; q_p>1 tests the standalone slope branch (no extras, k_crossed=0).
#[test]
fn groupings_vech_layout() {
let sizing = Sizing::FixedClusters { n_clusters: 4 };
let base = intercept_only_spec(sizing.clone());
// q_p = 1 (intercept-only): shape must be unchanged.
let g1 = LmmGroupings::from_cluster_spec(&base, 40, &[]);
assert_eq!(g1.n_theta(), 1);
assert_eq!(g1.k_family(), 4); // 4 clusters × 1
assert_eq!(g1.diagonal_theta(), &[0][..]);
// q_p = 2 (1 slope): vech([σ_00, σ_10, σ_11]) length 3; diagonals at 0, 2.
let mut spec2 = base.clone();
spec2.re.as_mut().unwrap().slopes.push(1);
let g2 = LmmGroupings::from_cluster_spec(&spec2, 40, &[1]);
assert_eq!(g2.primary_q, 2);
assert_eq!(g2.n_theta(), 3); // 2·3/2 = 3
assert_eq!(g2.k_family(), 8); // 4 clusters × 2
assert_eq!(g2.k_total, 8);
assert_eq!(g2.diagonal_theta(), &[0, 2][..]); // off-diagonal vech[1]=1 excluded
// q_p = 3 (2 slopes): vech([σ_00, σ_10, σ_11, σ_20, σ_21, σ_22]) length 6; diagonals at 0, 3, 5.
let mut spec3 = base.clone();
spec3.re.as_mut().unwrap().slopes.push(1);
spec3.re.as_mut().unwrap().slopes.push(2);
let g3 = LmmGroupings::from_cluster_spec(&spec3, 40, &[1, 2]);
assert_eq!(g3.primary_q, 3);
assert_eq!(g3.n_theta(), 6); // 3·4/2 = 6
assert_eq!(g3.k_family(), 12); // 4 clusters × 3
assert_eq!(g3.k_total, 12);
assert_eq!(g3.diagonal_theta(), &[0, 3, 5][..]);
}
/// `vech(Σ)` (column-major, matching the θ layout) for one `q×q` block's
/// `vech(Λ)` — the quantity `canonicalize_pinned_blocks` must leave alone.
fn sigma_vech(theta: &[f64], q: usize) -> Vec<f64> {
let mut lam = vec![0.0f64; q * q];
primary_lambda(theta, q, &mut lam);
let mut out = Vec::new();
for c in 0..q {
for r in c..q {
let mut s = 0.0;
for k in 0..=c {
s += lam[r * q + k] * lam[c * q + k];
}
out.push(s);
}
}
out
}
/// `canonicalize_pinned_blocks`: a pinned diagonal's column folds into the
/// trailing diagonals with Σ preserved, a second pass is a no-op, and a block
/// with no pinned diagonal comes back BIT-identical (that last one is what
/// keeps every non-singular fit's dump unmoved).
#[test]
fn canonicalize_pinned_blocks_folds_preserves_sigma_and_is_idempotent() {
let sizing = Sizing::FixedClusters { n_clusters: 4 };
let base = intercept_only_spec(sizing);
// q_p = 2, vech(Λ) = [λ00, λ10, λ11]: both diagonals at 0 with the whole
// variance carried by the unidentified λ10 — the shape a pinned q=2 fit
// stops at.
let mut spec2 = base.clone();
spec2.re.as_mut().unwrap().slopes.push(1);
let g2 = LmmGroupings::from_cluster_spec(&spec2, 40, &[1]);
let mut th2 = [0.0, 0.8172902747418043, 0.0];
let sig2 = sigma_vech(&th2, 2);
assert!(canonicalize_pinned_blocks(&g2, &mut th2));
assert_eq!(th2[0], 0.0);
assert_eq!(th2[1], 0.0); // column 0 zeroed outright
assert!((th2[2] - 0.8172902747418043).abs() <= 1e-15); // mass on λ11
for (got, want) in sigma_vech(&th2, 2).iter().zip(&sig2) {
assert!((got - want).abs() <= 1e-15 * want.abs().max(1.0));
}
// The pinned diagonal now has nothing below it, which is what makes the
// boundary-score shortcut valid there.
assert!(!g2.diagonal_has_nonzero_below(0, &th2));
let again = th2;
assert!(!canonicalize_pinned_blocks(&g2, &mut th2));
assert_eq!(th2[0].to_bits(), again[0].to_bits());
assert_eq!(th2[1].to_bits(), again[1].to_bits());
assert_eq!(th2[2].to_bits(), again[2].to_bits());
// No pinned diagonal ⇒ strict no-op, bit-for-bit.
let mut th_live = [1.3, 0.37, 0.9];
let before = th_live;
assert!(!canonicalize_pinned_blocks(&g2, &mut th_live));
for (a, b) in th_live.iter().zip(&before) {
assert_eq!(a.to_bits(), b.to_bits());
}
// q_p = 3, vech(Λ) = [λ00, λ10, λ20, λ11, λ21, λ22]: the pinned leading
// column carries two live entries that have to fold into columns 1 and 2.
let mut spec3 = base.clone();
spec3.re.as_mut().unwrap().slopes.push(1);
spec3.re.as_mut().unwrap().slopes.push(2);
let g3 = LmmGroupings::from_cluster_spec(&spec3, 40, &[1, 2]);
let mut th3 = [0.0, 0.5, 0.25, 0.0, 0.3, 0.4];
let sig3 = sigma_vech(&th3, 3);
assert!(canonicalize_pinned_blocks(&g3, &mut th3));
assert_eq!(&th3[..3], &[0.0, 0.0, 0.0]); // pinned column 0 gone
assert!(th3[3] > PIN_THETA); // its variance re-emerges on λ11
for (got, want) in sigma_vech(&th3, 3).iter().zip(&sig3) {
assert!(
(got - want).abs() <= 1e-14 * want.abs().max(1.0),
"Σ moved: {:?} vs {:?}",
sigma_vech(&th3, 3),
sig3
);
}
assert!(!g3.diagonal_has_nonzero_below(0, &th3));
let again3 = th3;
assert!(!canonicalize_pinned_blocks(&g3, &mut th3));
for (a, b) in th3.iter().zip(&again3) {
assert_eq!(a.to_bits(), b.to_bits());
}
}
/// Suff-stats bookkeeping on a hand-checkable block: counts per RE column,
/// per-column sums, crossed cross-counts.
#[test]
fn suff_stats_multi_accumulators() {
let (x, y, pid, eids, cluster) = multi_dataset(true, 1); // one atom block, n=48
let g = LmmGroupings::from_cluster_spec(&cluster, 48, &[]);
assert_eq!(g.n_primary, 6);
assert_eq!(g.nested_per_parent, 2);
assert_eq!(g.k_family(), 18); // 6 + 6·2
assert_eq!(g.k_total, 22); // + 4 crossed
assert_eq!(g.n_theta(), 3);
let mut suff = LmmSuffStats::with_groupings(3, g);
suff.add_rows_multi(x.as_ref(), &y, &pid, &eids, None);
// One full-factorial block: every primary level has 8 rows, every
// child 4, every crossed level 12.
for f in 0..6 {
assert_eq!(suff.counts[f], 8.0);
}
for c in 6..18 {
assert_eq!(suff.counts[c], 4.0);
}
for b in 18..22 {
assert_eq!(suff.counts[b], 12.0);
}
// Crossed co-occurrence: each (primary, crossed) pair shares exactly
// 2 rows in a full factorial of 6·4·2.
assert_eq!(suff.zx[(0, 0)], 2.0);
assert_eq!(suff.zx[(5, 3)], 2.0);
// Same-factor crossed pairs never co-occur.
assert_eq!(suff.zx[(18, 1)], 0.0);
// Intercept column sum = row count per level.
assert!((suff.s[(0, 0)] - 8.0).abs() < 1e-12);
}
/// Textbook REML deviance on the explicit n×n V — the oracle for the
/// family-blocked elimination. dev = ln|V| + ln|X'V⁻¹X| + (N−P)·ln σ̂²,
/// σ̂² = (y'V⁻¹y − β̂'X'V⁻¹y)/(N−P). `groups[g]` = grouping g's global
/// level ids (primary first); `theta[g]` the matching component.
fn brute_force_deviance(theta: &[f64], x: &Mat<f64>, y: &[f64], groups: &[&[u32]]) -> f64 {
use faer::linalg::solvers::Solve;
let n = x.nrows();
let p = x.ncols();
let mut v = Mat::<f64>::zeros(n, n);
for i in 0..n {
v[(i, i)] = 1.0;
}
for (g, ids) in groups.iter().enumerate() {
let t2 = theta[g] * theta[g];
for i in 0..n {
for j in 0..n {
if ids[i] == ids[j] {
v[(i, j)] += t2;
}
}
}
}
let vc = v.as_ref().llt(faer::Side::Lower).unwrap();
let mut log_det_v = 0.0;
for i in 0..n {
log_det_v += vc.L()[(i, i)].ln();
}
let log_det_v = 2.0 * log_det_v;
let mut vix = (*x).clone();
vc.solve_in_place(vix.as_mut());
let mut viy = Mat::<f64>::zeros(n, 1);
for i in 0..n {
viy[(i, 0)] = y[i];
}
vc.solve_in_place(viy.as_mut());
let mut xtvix = Mat::<f64>::zeros(p, p);
let mut xtviy = vec![0.0; p];
for a in 0..p {
for b in 0..p {
let mut acc = 0.0;
for i in 0..n {
acc += x[(i, a)] * vix[(i, b)];
}
xtvix[(a, b)] = acc;
}
let mut acc = 0.0;
for i in 0..n {
acc += x[(i, a)] * viy[(i, 0)];
}
xtviy[a] = acc;
}
let kc = xtvix.as_ref().llt(faer::Side::Lower).unwrap();
let mut log_det_k = 0.0;
for a in 0..p {
log_det_k += kc.L()[(a, a)].ln();
}
let log_det_k = 2.0 * log_det_k;
let mut beta = Mat::<f64>::zeros(p, 1);
for a in 0..p {
beta[(a, 0)] = xtviy[a];
}
kc.solve_in_place(beta.as_mut());
let mut ytviy = 0.0;
for i in 0..n {
ytviy += y[i] * viy[(i, 0)];
}
let mut bxy = 0.0;
for a in 0..p {
bxy += beta[(a, 0)] * xtviy[a];
}
let df = (n - p) as f64;
let sigma_sq = (ytviy - bxy) / df;
log_det_v + log_det_k + df * sigma_sq.ln()
}
fn assert_deviance_matches_oracle(with_nested: bool, thetas: &[Vec<f64>]) {
let (x, y, pid, eids, cluster) = multi_dataset(with_nested, 2);
let n = x.nrows();
let g = LmmGroupings::from_cluster_spec(&cluster, n, &[]);
let mut suff = LmmSuffStats::with_groupings(3, g);
suff.add_rows_multi(x.as_ref(), &y, &pid, &eids, None);
let gref = LmmGroupings::from_cluster_spec(&cluster, n, &[]);
let mut fit = LmmFitScratch::with_groupings(3, &gref);
let mut fit_c = LmmFitScratch::with_groupings(3, &gref);
assert!(precompute_balanced_collapse(&suff, &mut fit_c));
// Oracle wants global ids per grouping.
let mut groups: Vec<&[u32]> = vec![&pid];
for e in &eids {
groups.push(e);
}
for th in thetas {
let dev = reml_deviance(th, &suff, &mut fit);
let oracle = brute_force_deviance(th, &x, &y, &groups);
assert!(dev.is_finite(), "θ={th:?}");
let tol = 1e-8 * oracle.abs().max(1.0);
assert!(
(dev - oracle).abs() <= tol,
"θ={th:?}: family-blocked {dev} vs oracle {oracle}"
);
// Collapse arm: same θ through the balanced path — reassociation
// band vs the loop, oracle band absolute.
let dev_c = reml_deviance(th, &suff, &mut fit_c);
let band = 1e-9 * dev.abs().max(1.0);
assert!(
(dev_c - dev).abs() <= band,
"θ={th:?}: collapse {dev_c} vs loop {dev}"
);
assert!(
(dev_c - oracle).abs() <= tol,
"θ={th:?}: collapse vs oracle"
);
}
}
#[test]
fn crossed_deviance_matches_brute_force() {
assert_deviance_matches_oracle(
false,
&[
vec![0.5, 0.3],
vec![1.0, 1.0],
vec![2.0, 0.1],
vec![0.0, 0.7],
vec![1e-3, 1e-3],
],
);
}
#[test]
fn crossed_plus_nested_deviance_matches_brute_force() {
assert_deviance_matches_oracle(
true,
&[
vec![0.5, 0.3, 0.2],
vec![1.0, 1.0, 1.0],
vec![0.0, 0.5, 0.9],
vec![2.0, 0.05, 0.4],
],
);
}
/// Unbalanced counts must take the general loop byte-for-byte: a failed
/// precompute leaves collapse_n_active = 0 and the eval path untouched.
#[test]
fn unbalanced_counts_fall_back_byte_identical() {
let (x, y, pid, eids, cluster) = multi_dataset(true, 2);
let n = x.nrows() - 1; // truncate one row — last cluster short
let g = LmmGroupings::from_cluster_spec(&cluster, x.nrows(), &[]);
let mut suff = LmmSuffStats::with_groupings(3, g);
let eids_t: Vec<Vec<u32>> = eids.iter().map(|e| e[..n].to_vec()).collect();
suff.add_rows_multi(x.as_ref().subrows(0, n), &y[..n], &pid[..n], &eids_t, None);
let gref = LmmGroupings::from_cluster_spec(&cluster, x.nrows(), &[]);
let mut fit_a = LmmFitScratch::with_groupings(3, &gref);
let mut fit_b = LmmFitScratch::with_groupings(3, &gref);
assert!(!precompute_balanced_collapse(&suff, &mut fit_b));
for th in [[0.5, 0.3, 0.2], [1.0, 1.0, 1.0], [0.0, 0.5, 0.9]] {
let a = reml_deviance(&th, &suff, &mut fit_a);
let b = reml_deviance(&th, &suff, &mut fit_b);
assert_eq!(a.to_bits(), b.to_bits(), "θ={th:?}");
}
}
/// Off-grid N under `FixedSize`: row 17 of 18 sits in cluster 4, so five
/// primary levels exist and the nested block must cover all five parents.
#[test]
fn fixed_size_off_grid_n_keeps_the_partial_trailing_cluster() {
let mut cluster = intercept_only_spec(Sizing::FixedSize { cluster_size: 4 });
cluster.re.as_mut().unwrap().extra_groupings.push(Grouping {
relation: GroupingRelation::NestedWithin { n_per_parent: 2 },
slopes: vec![],
});
let n = 18;
let sizing = &cluster.re.as_ref().unwrap().sizing;
assert_eq!(sizing.cluster_of_row(n - 1), 4);
assert_eq!(sizing.n_clusters_at(n), 5);
let g = LmmGroupings::from_cluster_spec(&cluster, n, &[]);
assert_eq!(g.n_primary, 5);
assert_eq!(g.n_primary * g.nested_per_parent, 10);
assert_eq!(g.k_total, 15);
}
/// Nested-only in Regime B — the path with NO crossed tail (zx is 0×0)
/// and parents that grow with N.
#[test]
fn nested_regime_b_deviance_matches_brute_force() {
let mut cluster = intercept_only_spec(Sizing::FixedSize { cluster_size: 8 });
cluster.re.as_mut().unwrap().extra_groupings.push(Grouping {
relation: GroupingRelation::NestedWithin { n_per_parent: 2 },
slopes: vec![],
});
let n = 4 * model_atom(&cluster); // 64
let mut st = 7u64;
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut cid = vec![0u32; n];
let u_p: Vec<f64> = (0..8).map(|_| 0.5 * lcg(&mut st)).collect();
let u_c: Vec<f64> = (0..16).map(|_| 0.3 * lcg(&mut st)).collect();
for i in 0..n {
pid[i] = cluster.re.as_ref().unwrap().sizing.cluster_of_row(i) as u32;
cid[i] = extra_level_of_row(&cluster, 0, i) as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5 + 0.4 * x1 + u_p[pid[i] as usize] + u_c[cid[i] as usize] + 0.8 * lcg(&mut st);
}
let g = LmmGroupings::from_cluster_spec(&cluster, n, &[]);
let mut suff = LmmSuffStats::with_groupings(2, g);
let eids = vec![cid.clone()];
suff.add_rows_multi(x.as_ref(), &y, &pid, &eids, None);
let gref = LmmGroupings::from_cluster_spec(&cluster, n, &[]);
let mut fit = LmmFitScratch::with_groupings(2, &gref);
let mut fit_c = LmmFitScratch::with_groupings(2, &gref);
assert!(precompute_balanced_collapse(&suff, &mut fit_c));
for th in [[0.6, 0.4], [1.0, 1.0], [0.2, 0.0], [0.0, 0.0]] {
let dev = reml_deviance(&th, &suff, &mut fit);
let oracle = brute_force_deviance(&th, &x, &y, &[&pid, &cid]);
let tol = 1e-8 * oracle.abs().max(1.0);
assert!((dev - oracle).abs() <= tol, "θ={th:?}: {dev} vs {oracle}");
// Collapse arm — reassociation band vs the loop incl. the θ=0 edge.
let dev_c = reml_deviance(&th, &suff, &mut fit_c);
let band = 1e-9 * dev.abs().max(1.0);
assert!(
(dev_c - dev).abs() <= band,
"θ={th:?}: collapse {dev_c} vs {dev}"
);
}
}
/// Balanced-collapse applicability: balanced intercept designs precompute,
/// slope groupings and unbalanced counts fall back.
#[test]
fn balanced_collapse_applicability() {
// Balanced: the regime-B nested dataset (atom-multiple by construction).
let mut cluster = intercept_only_spec(Sizing::FixedSize { cluster_size: 8 });
cluster.re.as_mut().unwrap().extra_groupings.push(Grouping {
relation: GroupingRelation::NestedWithin { n_per_parent: 2 },
slopes: vec![],
});
let n = 4 * model_atom(&cluster); // 64
let max_n = 2 * n; // workspace sized for a larger grid top — active PREFIX
let mut st = 7u64;
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut cid = vec![0u32; n];
for i in 0..n {
pid[i] = cluster.re.as_ref().unwrap().sizing.cluster_of_row(i) as u32;
cid[i] = extra_level_of_row(&cluster, 0, i) as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5 + 0.4 * x1 + 0.8 * lcg(&mut st);
}
let g = LmmGroupings::from_cluster_spec(&cluster, max_n, &[]);
let n_primary = g.n_primary;
let mut suff = LmmSuffStats::with_groupings(2, g);
let eids = vec![cid.clone()];
suff.add_rows_multi(x.as_ref(), &y, &pid, &eids, None);
let gref = LmmGroupings::from_cluster_spec(&cluster, max_n, &[]);
let mut fit = LmmFitScratch::<f64>::with_groupings(2, &gref);
assert!(precompute_balanced_collapse(&suff, &mut fit));
assert_eq!(fit.collapse_n_active, n / 8);
assert!(fit.collapse_n_active < n_primary); // genuinely a prefix
// Unbalanced: drop the last row — the trailing cluster is short.
let mut suff_u =
LmmSuffStats::with_groupings(2, LmmGroupings::from_cluster_spec(&cluster, max_n, &[]));
let eids_u = vec![cid[..n - 1].to_vec()];
suff_u.add_rows_multi(
x.as_ref().subrows(0, n - 1),
&y[..n - 1],
&pid[..n - 1],
&eids_u,
None,
);
assert!(!precompute_balanced_collapse(&suff_u, &mut fit));
assert_eq!(fit.collapse_n_active, 0);
// Slope path: never applicable — populated, balanced data, so the
// rejection is the q_p guard, not the empty-suff early-out (balanced
// slope counts would otherwise pass the count checks).
let (xs, ys, ids_s) = slope_dataset();
let gs = slope_groupings();
let mut suff_s = LmmSuffStats::with_groupings(2, slope_groupings());
suff_s.add_rows_multi(xs.as_ref(), &ys, &ids_s, &[], None);
let mut fit_s = LmmFitScratch::<f64>::with_groupings(2, &gs);
assert!(!precompute_balanced_collapse(&suff_s, &mut fit_s));
}
/// Balanced collapse with prior weights: constant w ≡ 2 preserves the
/// per-cluster `counts` equality (`counts[f] = 2·n_f`, still exactly equal
/// across the balanced prefix), so the collapse must STILL trigger — and
/// the collapse-taken weighted fit must reproduce the unweighted one's
/// β/SE/tau2 (θ̃ = √c·θ maps the weighted profiled deviance onto the
/// unweighted one; θ̂² scales by 1/c, σ̂² by c, tau2 = θ²σ̂² invariant).
/// Both fits take the collapse branch (asserted below), so agreement is a
/// numeric check of the collapse kernel consuming weighted Grams, not just
/// of the accumulator.
#[test]
fn balanced_collapse_weighted_fit_invariant() {
let (x, y, ids) = hand_dataset(); // balanced: 6 clusters × 8 rows
let n = x.nrows();
let targets: Vec<u32> = vec![1, 2];
let w = vec![2.0f64; n];
let mut ws_w = LmmWorkspace::new(3, 6);
ws_w.suff
.add_rows_multi(x.as_ref(), &y, &ids, &[], Some(&w));
assert!(
precompute_balanced_collapse(&ws_w.suff, &mut ws_w.fit),
"constant weights keep exact per-cluster counts equality"
);
assert_eq!(ws_w.fit.collapse_n_active, 6);
let fit_w = fit_lmm(&mut ws_w, &targets, None);
assert!(fit_w.converged);
let mut ws_u = LmmWorkspace::new(3, 6);
ws_u.suff.add_rows(x.as_ref(), &y, &ids);
let fit_u = fit_lmm(&mut ws_u, &targets, None);
assert!(fit_u.converged);
// Two independent BOBYQA runs agree to the rho_end floor, not machine
// precision — same 1e-6 relative band as the fit.rs invariance tests.
for j in 0..3 {
let (a, b) = (ws_u.fit.betas[j], ws_w.fit.betas[j]);
assert!(
(a - b).abs() / a.abs() < 1e-6,
"β[{j}] unweighted {a} vs w≡2 {b}"
);
}
for &tj in &targets {
let (a, b) = (
ws_u.fit.var_diag[tj as usize].sqrt(),
ws_w.fit.var_diag[tj as usize].sqrt(),
);
assert!(
(a - b).abs() / a < 1e-6,
"se[{tj}] unweighted {a} vs w≡2 {b}"
);
}
let (tu, tw) = (
ws_u.theta[0] * ws_u.theta[0] * fit_u.sigma_sq,
ws_w.theta[0] * ws_w.theta[0] * fit_w.sigma_sq,
);
assert!(
(tu - tw).abs() / tu < 1e-6,
"tau2 unweighted {tu} vs w≡2 {tw}"
);
}
/// Two crossed factors — the dense cross-factor coupling block.
#[test]
fn two_crossed_factors_deviance_matches_brute_force() {
let mut cluster = intercept_only_spec(Sizing::FixedClusters { n_clusters: 3 });
for k in [4u32, 2u32] {
cluster.re.as_mut().unwrap().extra_groupings.push(Grouping {
relation: GroupingRelation::Crossed { n_clusters: k },
slopes: vec![],
});
}
let n = 2 * model_atom(&cluster); // 48
let mut st = 21u64;
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut e0 = vec![0u32; n];
let mut e1 = vec![0u32; n];
let u_p: Vec<f64> = (0..3).map(|_| 0.5 * lcg(&mut st)).collect();
let u_a: Vec<f64> = (0..4).map(|_| 0.4 * lcg(&mut st)).collect();
let u_b: Vec<f64> = (0..2).map(|_| 0.3 * lcg(&mut st)).collect();
for i in 0..n {
pid[i] = cluster.re.as_ref().unwrap().sizing.cluster_of_row(i) as u32;
e0[i] = extra_level_of_row(&cluster, 0, i) as u32;
e1[i] = extra_level_of_row(&cluster, 1, i) as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5
+ 0.4 * x1
+ u_p[pid[i] as usize]
+ u_a[e0[i] as usize]
+ u_b[e1[i] as usize]
+ 0.8 * lcg(&mut st);
}
let g = LmmGroupings::from_cluster_spec(&cluster, n, &[]);
let mut suff = LmmSuffStats::with_groupings(2, g);
let eids = vec![e0.clone(), e1.clone()];
suff.add_rows_multi(x.as_ref(), &y, &pid, &eids, None);
let gref = LmmGroupings::from_cluster_spec(&cluster, n, &[]);
let mut fit = LmmFitScratch::with_groupings(2, &gref);
let mut fit_c = LmmFitScratch::with_groupings(2, &gref);
assert!(precompute_balanced_collapse(&suff, &mut fit_c));
for th in [[0.5, 0.4, 0.3], [1.0, 1.0, 1.0], [0.3, 0.0, 0.8]] {
let dev = reml_deviance(&th, &suff, &mut fit);
let oracle = brute_force_deviance(&th, &x, &y, &[&pid, &e0, &e1]);
let tol = 1e-8 * oracle.abs().max(1.0);
assert!((dev - oracle).abs() <= tol, "θ={th:?}: {dev} vs {oracle}");
// Collapse arm — reassociation band vs the loop.
let dev_c = reml_deviance(&th, &suff, &mut fit_c);
let band = 1e-9 * dev.abs().max(1.0);
assert!(
(dev_c - dev).abs() <= band,
"θ={th:?}: collapse {dev_c} vs {dev}"
);
}
}
/// The dual REML objective takes the balanced collapse the `f64` objective
/// takes: on a balanced design `precompute_balanced_collapse` arms on a
/// `Dual<4>` scratch, and the gradient it then produces still matches the dense
/// score — the collapse is a reassociation of the same criterion, not a
/// different one. On an unbalanced design it must refuse on both scalars alike.
#[test]
fn dual_scratch_arms_the_balanced_collapse() {
let (x, y, ids) = hand_dataset_sized(240, 6, false);
let groupings = LmmGroupings::single(6);
let p = 3;
let mut suff = LmmSuffStats::with_groupings(p, groupings.clone());
suff.add_rows_multi(x.as_ref(), &y, &ids, &[], None);
let n_theta = groupings.n_theta();
let mut fit_d = LmmFitScratch::<Dual<4>>::with_groupings(p, &groupings);
assert!(
precompute_balanced_collapse(&suff, &mut fit_d),
"balanced design must arm on a dual scratch"
);
assert_eq!(fit_d.collapse_n_active, 6);
// Same criterion: the armed dual gradient still matches the dense score.
let z = dense_z(&groupings, &ids, &[], &x, &[]);
let mut scratch = LmmDualScratch::for_groupings(n_theta, p, &groupings)
.expect("n_theta 1 is inside the lane set");
let mut rng = fixed_seed_theta_lmm("int1", 6001);
for _ in 0..5 {
let theta = rng.next_theta();
let mut grad = vec![0.0; n_theta];
assert!(matches!(
reml_gradient(&theta, &suff, &mut scratch, &mut grad),
DerivStatus::Ok(_)
));
let want = dense_reml_score(&theta, z.as_ref(), x.as_ref(), &y, &groupings);
for j in 0..n_theta {
let band = 1e-8 * want[j].abs().max(1.0);
assert!(
(grad[j] - want[j]).abs() <= band,
"coord {j}: dual {} vs dense {}",
grad[j],
want[j]
);
}
}
// Unbalanced: refused on both scalars, so the dual arm keeps the loop.
let (xs, ys, ids_s) = hand_dataset_sized(240, 6, true);
let mut suff_s = LmmSuffStats::with_groupings(p, groupings.clone());
suff_s.add_rows_multi(xs.as_ref(), &ys, &ids_s, &[], None);
let mut fit_s_f64 = LmmFitScratch::<f64>::with_groupings(p, &groupings);
let mut fit_s_dual = LmmFitScratch::<Dual<4>>::with_groupings(p, &groupings);
assert!(!precompute_balanced_collapse(&suff_s, &mut fit_s_f64));
assert!(!precompute_balanced_collapse(&suff_s, &mut fit_s_dual));
}
/// Per-component pin: items carry NO between-level signal by construction
/// (each item sees every subject equally, and the ±0.8 residual pattern is
/// block-constant so item means cancel exactly), while subjects carry a
/// real u_p. The crossed component must pin at exactly 0 (boundary_hit
/// == 1) with the primary component interior.
#[test]
fn zero_crossed_variance_pins_only_that_component() {
let s_cl = 4usize;
let i_cl = 3usize;
let mut cluster = intercept_only_spec(Sizing::FixedClusters {
n_clusters: s_cl as u32,
});
cluster.re.as_mut().unwrap().extra_groupings.push(Grouping {
relation: GroupingRelation::Crossed {
n_clusters: i_cl as u32,
},
slopes: vec![],
});
let n = 4 * model_atom(&cluster); // 48: 4 blocks ⇒ ±0.8 cancels per item
let mut st = 5u64;
let u_p: Vec<f64> = (0..s_cl).map(|_| 0.8 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut eid = vec![0u32; n];
for i in 0..n {
pid[i] = cluster.re.as_ref().unwrap().sizing.cluster_of_row(i) as u32;
eid[i] = extra_level_of_row(&cluster, 0, i) as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
let e = if (i / model_atom(&cluster)) % 2 == 0 {
0.8
} else {
-0.8
};
y[i] = 0.5 + 0.4 * x1 + u_p[pid[i] as usize] + e;
}
let mut ws = LmmWorkspace::for_cluster_spec(2, &cluster, n, &[]);
let eids = vec![eid];
ws.suff.add_rows_multi(x.as_ref(), &y, &pid, &eids, None);
let fit = fit_lmm(&mut ws, &[1], None);
assert!(fit.converged);
assert_eq!(fit.boundary_hit, 1);
assert_eq!(ws.theta[1], 0.0, "crossed component must pin at exact 0.0");
assert!(
ws.theta[0] > PIN_THETA,
"primary component must stay interior"
);
assert!(fit.joint_t_sq.is_finite());
}
/// End-to-end crossed+nested fit recovers the generating β within wide
/// sanity bands and produces finite Wald machinery — the L1 smoke for the
/// full multi-grouping pipeline (the statistical gates live in L3).
#[test]
fn crossed_nested_fit_recovers_betas() {
let (x, y, pid, eids, cluster) = multi_dataset(true, 4); // n = 192
let mut ws = LmmWorkspace::for_cluster_spec(3, &cluster, x.nrows(), &[]);
ws.suff.add_rows_multi(x.as_ref(), &y, &pid, &eids, None);
let fit = fit_lmm(&mut ws, &[1, 2], None);
assert!(fit.converged);
assert!((ws.fit.betas[1] - 0.4).abs() < 0.15);
assert!((ws.fit.betas[2] + 0.2).abs() < 0.15);
// Deterministic regression lock (lcg-seeded multi_dataset) alongside the
// planted-value recovers-check above, which documents intent.
assert!((ws.fit.betas[1] - 0.40829926961384383).abs() / 0.40829926961384383_f64.abs() < 1e-6);
assert!((ws.fit.betas[2] - -0.2916210839321183).abs() / 0.2916210839321183_f64.abs() < 1e-6);
assert!(ws.fit.t_sq[1].is_finite() && ws.fit.t_sq[2].is_finite());
assert!(fit.joint_t_sq.is_finite() && fit.joint_t_sq > 0.0);
assert_eq!(ws.theta.len(), 3);
}
/// General-path twin of lmm_fit_warm_path_bounded_alloc: crossed+nested
/// workspace. Per-call blocks are the tail-llt faer internals (the family
/// loop is hand-rolled, zero-alloc) — the same acceptance class as q=1.
/// Warm-started from a cold prime fit's fitted θ (the loop tier's production
/// pattern), matching the few-eval regime the production path runs.
#[cfg(feature = "alloc-tests")]
#[test]
#[ignore]
fn lmm_fit_general_warm_path_bounded_alloc() {
let _serial = crate::test_support::alloc_test_guard();
const N_CALLS: usize = 100;
const BOUND_GENERAL: u64 = 3700; // Measured 3500 (this machine) — ~35 blocks/fit truth-started (scaled rho + spec-derived start; the few-eval regime the production path runs). Per-eval faer `llt` internals only: the family loop is hand-rolled zero-alloc, the cached diagonal_theta map avoids a per-fit Vec allocation, and the ranef recovery pass solves in the ranef_ux/ranef_rhs scratch fields, so this count is faer-version/machine specific. `chol_lower` factors in place, which keeps this count low. If faer changes its Cholesky internals, update — do not relax.
let (x, y, pid, eids, cluster) = multi_dataset(true, 2);
let targets: Vec<u32> = vec![1, 2];
let mut ws = LmmWorkspace::for_cluster_spec(3, &cluster, x.nrows(), &[]);
ws.suff.reset();
ws.suff.add_rows_multi(x.as_ref(), &y, &pid, &eids, None);
// prime cold, then warm-start subsequent refits from the previous fit's fitted θ
// (the loop tier's production pattern).
let _ = fit_lmm(&mut ws, &targets, None);
let warm = ws.theta.clone();
let profiler = dhat::Profiler::builder().testing().build();
for _ in 0..N_CALLS {
ws.suff.reset();
ws.suff.add_rows_multi(x.as_ref(), &y, &pid, &eids, None);
let _ = fit_lmm(&mut ws, &targets, Some(&warm));
}
let stats = dhat::HeapStats::get();
drop(profiler);
assert!(
stats.total_blocks <= BOUND_GENERAL,
"general fit_lmm allocated {} blocks across {} warm-path calls (BOUND = {})",
stats.total_blocks,
N_CALLS,
BOUND_GENERAL
);
}
/// Crossed-slopes twin of the bounded-alloc gate: the blocked path's only
/// per-eval heap traffic is the faer `llt` internals (everything else lives in
/// `LmmFitScratch.blocked_*`, sized once). Same acceptance class as the other
/// general fits; faer-version/machine specific — if faer changes its Cholesky
/// internals, update the bound, do not relax.
#[cfg(feature = "alloc-tests")]
#[test]
#[ignore]
fn lmm_fit_crossed_slope_warm_path_bounded_alloc() {
let _serial = crate::test_support::alloc_test_guard();
const N_CALLS: usize = 100;
// Measured ~5004 (this machine, faer 0.x, 2026-09-05): the per-eval Cholesky
// runs through `chol_lower`'s own scratch buffer at `Par::Seq` instead
// of `MatRef::llt`, so it never touches faer's global rayon pool. ALL faer-
// internal — `reml_deviance_blocked` itself is zero-alloc (every buffer is
// in `blocked_*` scratch; only a stack `lam_g`). faer-version/machine
// specific; if faer changes its Cholesky internals, update — do not relax.
const BOUND: u64 = 55000;
let (x, y, pid, eid) = crossed_slope_golden_dataset();
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 8 },
slopes: vec![1],
extra_groupings: vec![Grouping {
relation: GroupingRelation::Crossed { n_clusters: 6 },
slopes: vec![1],
}],
}),
};
let mut ws = LmmWorkspace::for_cluster_spec_ext(2, &cluster, x.nrows(), &[1], &[vec![1]]);
ws.suff.reset();
ws.suff
.add_rows_multi(x.as_ref(), &y, &pid, std::slice::from_ref(&eid), None);
// prime cold, then warm-start subsequent refits from the previous fit's fitted θ
// (the loop tier's production pattern).
let _ = fit_lmm(&mut ws, &[1], None);
let warm = ws.theta.clone();
let profiler = dhat::Profiler::builder().testing().build();
for _ in 0..N_CALLS {
ws.suff.reset();
ws.suff
.add_rows_multi(x.as_ref(), &y, &pid, std::slice::from_ref(&eid), None);
let _ = fit_lmm(&mut ws, &[1], Some(&warm));
}
let stats = dhat::HeapStats::get();
drop(profiler);
assert!(
stats.total_blocks <= BOUND,
"crossed-slope fit_lmm allocated {} blocks across {} warm-path calls (BOUND = {})",
stats.total_blocks,
N_CALLS,
BOUND
);
}
// -----------------------------------------------------------------------
// Standalone primary slopes: q_p×q_p primary block, oracle deviance,
// diagonal-only pin. Data lives on the engine's f32 plane (mirrors the scalar
// oracle convention); the brute force widens the identical bytes to f64, so
// the 1e-8 match is exact, not modulo an f32↔f64 roundtrip.
// -----------------------------------------------------------------------
/// n=64, p=2 (intercept + x1), 8 clusters, y carries u₀ + u₁·x1.
fn slope_dataset() -> (Mat<f64>, Vec<f64>, Vec<u32>) {
slope_dataset_sized(64, 8)
}
/// `slope_dataset`, generalized over row/cluster count — the dual-call cost
/// table's size axis needs this slope-primary shape at 3,000/30,000 rows.
/// No balance parameter: `precompute_balanced_collapse` requires a q_p=1
/// primary (`kernel.rs`'s `g.primary_q != 1` guard), and this shape's
/// primary carries a slope (q_p=2), so it never arms the collapse regardless
/// of cluster balance — a skewed variant would tell the cost table nothing
/// the balanced one doesn't already show.
fn slope_dataset_sized(n: usize, nc: usize) -> (Mat<f64>, Vec<f64>, Vec<u32>) {
let mut st = 71u64;
let u0: Vec<f64> = (0..nc).map(|_| 0.5 * lcg(&mut st)).collect();
let u1: Vec<f64> = (0..nc).map(|_| 0.3 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut ids = vec![0u32; n];
for i in 0..n {
let c = i % nc;
ids[i] = c as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5 + 0.4 * x1 + u0[c] + u1[c] * x1 + 0.8 * lcg(&mut st);
}
(x, y, ids)
}
/// A crossed-tail LMM dataset with the shape the speed grid's `cross*` cells
/// have: an intercept (or intercept+slope) primary at `nc` levels plus
/// `n_extra` crossed factors of `levels` levels each, `p = 2`. The tail the
/// deviance factors is `t_dim = n_extra·levels + p + 1` — 93 at
/// `(3, 30)` and 153 at `(5, 30)`, matching `cross4_g3000p5` and
/// `cross6_g30000p5`. `skew` uses the same ~20%-of-clusters-draw-~80%-of-rows
/// recipe as `hand_dataset_sized`, applied to the PRIMARY only: the collapse's
/// balance test reads the primary counts (and the nested-child counts, of which
/// there are none here), so skewing the crossed factors would not disarm it.
/// `slope` puts a random slope on x-column 1, taking the primary to `q_p = 2`
/// — that shape never arms the collapse, and exists only to push `n_theta` to
/// the `D12` rung (`3 + n_extra`).
fn cross_dataset_sized(
n: usize,
nc: usize,
n_extra: usize,
levels: usize,
slope: bool,
skew: bool,
) -> (Mat<f64>, Vec<f64>, Vec<u32>, Vec<Vec<u32>>) {
let mut st = 3307u64;
let u_p: Vec<f64> = (0..nc).map(|_| 0.5 * lcg(&mut st)).collect();
let u_s: Vec<f64> = (0..nc).map(|_| 0.3 * lcg(&mut st)).collect();
let u_e: Vec<Vec<f64>> = (0..n_extra)
.map(|_| (0..levels).map(|_| 0.4 * lcg(&mut st)).collect())
.collect();
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut eids = vec![vec![0u32; n]; n_extra];
let n_heavy = if nc < 2 {
0
} else {
((nc as f64 * 0.2).round() as usize).clamp(1, nc - 1)
};
for i in 0..n {
let c = if !skew || nc < 2 {
i % nc
} else {
let r = lcg(&mut st) + 0.5; // in [0, 1)
let raw = if r < 0.8 {
(r / 0.8) * n_heavy as f64
} else {
n_heavy as f64 + ((r - 0.8) / 0.2) * (nc - n_heavy) as f64
};
(raw as usize).min(nc - 1)
};
pid[i] = c as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
let mut yi = 0.5 + 0.4 * x1 + u_p[c] + 0.8 * lcg(&mut st);
if slope {
yi += u_s[c] * x1;
}
for g in 0..n_extra {
// Each factor gets its own stride so the crossed levels are not all
// aligned with each other (aligned factors would be one factor).
let l = (i / (g + 1)) % levels;
eids[g][i] = l as u32;
yi += u_e[g][l];
}
y[i] = yi;
}
(x, y, pid, eids)
}
/// The `LmmGroupings` matching `cross_dataset_sized`'s layout. `n_theta` is
/// `1 + n_extra` without a slope and `3 + n_extra` with one.
fn cross_groupings(
nc: usize,
n_extra: usize,
levels: usize,
slope: bool,
n: usize,
) -> LmmGroupings {
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters {
n_clusters: nc as u32,
},
slopes: if slope { vec![0] } else { vec![] },
extra_groupings: (0..n_extra)
.map(|_| Grouping {
relation: GroupingRelation::Crossed {
n_clusters: levels as u32,
},
slopes: vec![],
})
.collect(),
}),
};
let slope_cols: &[usize] = if slope { &[1] } else { &[] };
LmmGroupings::from_cluster_spec(&cluster, n, slope_cols)
}
/// n=96, p=3 (intercept + x1 + x2), 8 clusters, y carries u₀ + u₁·x1 + u₂·x2.
fn multislope_dataset() -> (Mat<f64>, Vec<f64>, Vec<u32>) {
let (n, nc) = (96usize, 8usize);
let mut st = 91u64;
let u0: Vec<f64> = (0..nc).map(|_| 0.5 * lcg(&mut st)).collect();
let u1: Vec<f64> = (0..nc).map(|_| 0.3 * lcg(&mut st)).collect();
let u2: Vec<f64> = (0..nc).map(|_| 0.25 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 3);
let mut y = vec![0.0f64; n];
let mut ids = vec![0u32; n];
for i in 0..n {
let c = i % nc;
ids[i] = c as u32;
let (x1, x2) = (lcg(&mut st), lcg(&mut st));
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
x[(i, 2)] = x2;
y[i] = 0.5 + 0.4 * x1 + 0.2 * x2 + u0[c] + u1[c] * x1 + u2[c] * x2 + 0.8 * lcg(&mut st);
}
(x, y, ids)
}
/// Textbook REML deviance with a q×q D over the slope columns of Z_p.
/// `theta` is the column-major vech of Λ (q×q lower-tri); D_rel = ΛΛ′
/// (σ-relative); V = I + Z·D_rel·Z′ with Z_i = [1, x[i, slope_cols]]. The
/// f32 data is widened to f64 so the oracle reads the same bytes the suff
/// stats accumulated.
fn brute_force_slope_deviance(
theta: &[f64],
x: &Mat<f64>,
y: &[f64],
ids: &[u32],
slope_cols: &[usize],
q: usize,
) -> f64 {
use faer::linalg::solvers::Solve;
let (n, p) = (x.nrows(), x.ncols());
// Λ (q×q lower-tri) from column-major vech, then D = ΛΛ′.
let mut lam = vec![0.0f64; q * q];
let mut t = 0;
for c in 0..q {
for r in c..q {
lam[r * q + c] = theta[t];
t += 1;
}
}
let mut d = vec![0.0f64; q * q];
for i in 0..q {
for j in 0..q {
let mut s = 0.0;
for k in 0..q {
s += lam[i * q + k] * lam[j * q + k];
}
d[i * q + j] = s;
}
}
let zrow = |i: usize| -> Vec<f64> {
let mut z = vec![1.0];
for &sc in slope_cols {
z.push(x[(i, sc)]);
}
z
};
let mut v = Mat::<f64>::zeros(n, n);
for i in 0..n {
v[(i, i)] += 1.0;
}
for i in 0..n {
let zi = zrow(i);
for j in 0..n {
if ids[i] == ids[j] {
let zj = zrow(j);
let mut acc = 0.0;
for a in 0..q {
for b in 0..q {
acc += zi[a] * d[a * q + b] * zj[b];
}
}
v[(i, j)] += acc;
}
}
}
// REML profile (unchanged from the scalar oracle): ldv + ldk + df·ln s².
let vc = v.as_ref().llt(faer::Side::Lower).unwrap();
let mut ldv = 0.0;
for i in 0..n {
ldv += vc.L()[(i, i)].ln();
}
let ldv = 2.0 * ldv;
let mut vix = (*x).clone();
vc.solve_in_place(vix.as_mut());
let mut viy = Mat::<f64>::zeros(n, 1);
for i in 0..n {
viy[(i, 0)] = y[i];
}
vc.solve_in_place(viy.as_mut());
let mut xtvix = Mat::<f64>::zeros(p, p);
let mut xtviy = vec![0.0; p];
for aa in 0..p {
for bb in 0..p {
let mut s = 0.0;
for i in 0..n {
s += x[(i, aa)] * vix[(i, bb)];
}
xtvix[(aa, bb)] = s;
}
let mut s = 0.0;
for i in 0..n {
s += x[(i, aa)] * viy[(i, 0)];
}
xtviy[aa] = s;
}
let kc = xtvix.as_ref().llt(faer::Side::Lower).unwrap();
let mut ldk = 0.0;
for aa in 0..p {
ldk += kc.L()[(aa, aa)].ln();
}
let ldk = 2.0 * ldk;
let mut beta = Mat::<f64>::zeros(p, 1);
for aa in 0..p {
beta[(aa, 0)] = xtviy[aa];
}
kc.solve_in_place(beta.as_mut());
let mut ytviy = 0.0;
for i in 0..n {
ytviy += y[i] * viy[(i, 0)];
}
let mut bxy = 0.0;
for aa in 0..p {
bxy += beta[(aa, 0)] * xtviy[aa];
}
let df = (n - p) as f64;
let s2 = (ytviy - bxy) / df;
ldv + ldk + df * s2.ln()
}
fn slope_groupings() -> LmmGroupings {
slope_groupings_sized(64, 8)
}
/// `slope_groupings`, generalized over row/cluster count to match
/// `slope_dataset_sized`.
fn slope_groupings_sized(n: usize, nc: usize) -> LmmGroupings {
// nc primary clusters, one slope on x_full col 1; no extras.
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters {
n_clusters: nc as u32,
},
slopes: vec![0],
extra_groupings: vec![],
}),
};
LmmGroupings::from_cluster_spec(&cluster, n, &[1])
}
fn multislope_groupings() -> LmmGroupings {
// 8 primary clusters, two slopes on x_full cols 1,2; no extras.
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 8 },
slopes: vec![0, 1],
extra_groupings: vec![],
}),
};
LmmGroupings::from_cluster_spec(&cluster, 96, &[1, 2])
}
#[test]
fn slope_deviance_matches_brute_force() {
let (x, y, ids) = slope_dataset();
let mut suff = LmmSuffStats::with_groupings(2, slope_groupings());
suff.add_rows_multi(x.as_ref(), &y, &ids, &[], None);
let mut fit = LmmFitScratch::with_groupings(2, &slope_groupings());
// θ = vech(Λ), q=2: [λ₀₀, λ₁₀, λ₁₁].
for th in [
vec![1.0, 0.0, 1.0],
vec![0.5, 0.2, 0.4],
vec![2.0, -0.5, 0.7],
vec![1e-3, 1e-3, 1e-3],
// θ at THETA_HI (BOBYQA's box upper bound): the per-family Crout
// pivot product must stay finite here — a product accumulated
// across all families instead of reset per family would overflow.
vec![THETA_HI, 0.0, THETA_HI],
] {
let dev = reml_deviance(&th, &suff, &mut fit);
let oracle = brute_force_slope_deviance(&th, &x, &y, &ids, &[1], 2);
assert!(dev.is_finite(), "θ={th:?}");
assert!(
(dev - oracle).abs() <= 1e-8 * oracle.abs().max(1.0),
"θ={th:?}: {dev} vs {oracle}"
);
}
}
#[test]
fn multislope_deviance_matches_brute_force() {
let (x, y, ids) = multislope_dataset();
let mut suff = LmmSuffStats::with_groupings(3, multislope_groupings());
suff.add_rows_multi(x.as_ref(), &y, &ids, &[], None);
let mut fit = LmmFitScratch::with_groupings(3, &multislope_groupings());
// θ = vech(Λ), q=3: [λ₀₀, λ₁₀, λ₂₀, λ₁₁, λ₂₁, λ₂₂].
for th in [
vec![1.0, 0.0, 0.0, 1.0, 0.0, 1.0],
vec![0.6, 0.2, -0.1, 0.4, 0.15, 0.3],
vec![1.5, -0.4, 0.3, 0.7, -0.2, 0.5],
] {
let dev = reml_deviance(&th, &suff, &mut fit);
let oracle = brute_force_slope_deviance(&th, &x, &y, &ids, &[1, 2], 3);
assert!(dev.is_finite(), "θ={th:?}");
assert!(
(dev - oracle).abs() <= 1e-8 * oracle.abs().max(1.0),
"θ={th:?}: {dev} vs {oracle}"
);
}
}
/// End-to-end single-slope fit recovers the planted structure within BOBYQA
/// bands and pins nothing on a well-identified design.
#[test]
fn slope_fit_converges_interior() {
let (x, y, ids) = slope_dataset();
let mut ws = LmmWorkspace::with_groupings(2, slope_groupings());
ws.suff.add_rows_multi(x.as_ref(), &y, &ids, &[], None);
let fit = fit_lmm(&mut ws, &[1], None);
assert!(fit.converged);
// Planted [intercept 0.5, slope 0.4]; small (n=64, 8 clusters) REML draw
// recovers ≈[0.46, 0.20] — directionally correct, finite-sample attenuated.
// Pin sign + a band tight enough to catch a sign flip, a collapse to 0, or a
// blow-up (mere `is_finite` passed any of those).
assert!(
(0.2..0.8).contains(&ws.fit.betas[0]),
"intercept {}",
ws.fit.betas[0]
);
assert!(
(0.05..0.6).contains(&ws.fit.betas[1]),
"slope {}",
ws.fit.betas[1]
);
// Deterministic regression lock alongside the bands above.
assert!((ws.fit.betas[0] - 0.46265883331118085).abs() / 0.46265883331118085_f64.abs() < 1e-6);
assert!((ws.fit.betas[1] - 0.20152611939449563).abs() / 0.20152611939449563_f64.abs() < 1e-6);
assert_eq!(fit.pinned_components & !0b11, 0); // only 2 components exist
}
/// End-to-end two-slope fit: 3 components (intercept + 2 slopes), interior.
#[test]
fn multislope_fit_converges_interior() {
let (x, y, ids) = multislope_dataset();
let mut ws = LmmWorkspace::with_groupings(3, multislope_groupings());
ws.suff.add_rows_multi(x.as_ref(), &y, &ids, &[], None);
let fit = fit_lmm(&mut ws, &[1, 2], None);
assert!(fit.converged);
// Planted [0.5, 0.4, 0.2]; recovered ≈[0.51, 0.64, 0.28]. Both slopes positive
// with β̂₁ > β̂₂ (planted ordering preserved) — pin that, so a β₁/β₂ swap or a
// scale collapse fails, which a bare finiteness check would not catch.
assert!(
(0.2..0.9).contains(&ws.fit.betas[0]),
"intercept {}",
ws.fit.betas[0]
);
assert!(
(0.2..1.1).contains(&ws.fit.betas[1]),
"slope x1 {}",
ws.fit.betas[1]
);
assert!(
(0.0..0.7).contains(&ws.fit.betas[2]),
"slope x2 {}",
ws.fit.betas[2]
);
assert!(
ws.fit.betas[1] > ws.fit.betas[2],
"x1 slope must exceed x2 slope"
);
// Deterministic regression lock alongside the bands above.
assert!((ws.fit.betas[0] - 0.5129839426148501).abs() / 0.5129839426148501_f64.abs() < 1e-6);
assert!((ws.fit.betas[1] - 0.6442611282130077).abs() / 0.6442611282130077_f64.abs() < 1e-6);
assert!((ws.fit.betas[2] - 0.28355377896623535).abs() / 0.28355377896623535_f64.abs() < 1e-6);
assert_eq!(fit.pinned_components & !0b111, 0); // only 3 components exist
}
/// The experimental two-stage warm restart must reach the same
/// optimum as single-stage on a well-behaved rung — stage 1 (npt = n+2,
/// rho_end 1e-3, measured correctness-safe on the validation corpus) finds the
/// basin, stage 2 (npt = 2n+1, shipped rho_end) refines from stage 1's point.
/// Uses the multislope fixture (n_theta = 6) so the shipped mid-npt formula
/// (`n_theta >= 3`) is the one exercised by the single-stage comparator.
#[test]
fn two_stage_matches_single_stage_optimum() {
let (x, y, ids) = multislope_dataset();
let mut ws1 = LmmWorkspace::with_groupings(3, multislope_groupings());
ws1.suff.add_rows_multi(x.as_ref(), &y, &ids, &[], None);
let targets = [1u32, 2];
let f1 = fit_lmm(&mut ws1, &targets, None);
let (x, y, ids) = multislope_dataset();
let mut ws2 = LmmWorkspace::with_groupings(3, multislope_groupings());
ws2.suff.add_rows_multi(x.as_ref(), &y, &ids, &[], None);
let f2 = fit_lmm_two_stage(&mut ws2, &targets, None);
assert!(f2.converged);
assert!(
(f1.deviance - f2.deviance).abs() < 1e-6,
"two-stage must land on the same optimum: {} vs {}",
f1.deviance,
f2.deviance
);
assert!(f2.n_eval > 0);
}
/// Slope-variance collapse pins the SLOPE component (bit 1), not the
/// intercept. x1 is a within-cluster antithetic ±1 pattern that carries a
/// real fixed slope but ZERO cluster-varying slope, and the residual is a
/// ±0.8 period-4 quadrature block (+,+,−,− against x1's +,−,+,−) so every
/// cluster has Σ resid = 0 AND Σ x1·resid = 0 exactly — the REML
/// slope-variance MLE is 0, so λ₁₁ pins (bit 1) while the planted u₀ keeps
/// λ₀₀ interior. (The original lockstep ±0.8 pattern made resid ≡ 0.8·x1 —
/// collinear with the slope covariate, so σ̂²→0 once large θ₀ absorbed the
/// exactly-identified cluster means, the deviance ran unbounded to the θ₀
/// box bound, and the λ₁₁ pin rode FP noise on the degenerate surface; the
/// quadrature pattern keeps σ̂² positive and θ̂₀ genuinely interior.) Large
/// balanced design (16 clusters × 16 rows) so finite-sample REML does not
/// overfit a spurious slope RE the way a small noisy draw does.
#[test]
fn zero_slope_variance_pins_slope_component() {
let (nc, per) = (16usize, 16usize);
let n = nc * per;
let mut st = 5u64;
let u0: Vec<f64> = (0..nc).map(|_| 0.6 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut ids = vec![0u32; n];
#[allow(clippy::needless_range_loop)]
for c in 0..nc {
for k in 0..per {
let i = c * per + k;
ids[i] = c as u32;
// x1: identical antithetic pattern in every cluster (±1
// alternating) — no between-cluster slope signal.
let x1 = if k % 2 == 0 { 1.0 } else { -1.0 };
// residual: ±0.8 period-4 quadrature against x1, so per cluster
// Σ x1·resid = 0 AND Σ resid = 0 (no slope/intercept RE pull
// from the noise; only the planted u₀ moves intercepts).
let e = if (k / 2) % 2 == 0 { 0.8 } else { -0.8 };
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5 + 0.4 * x1 + u0[c] + e;
}
}
let mut ws = LmmWorkspace::with_groupings(
2,
LmmGroupings::from_cluster_spec(
&ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters {
n_clusters: nc as u32,
},
slopes: vec![0],
extra_groupings: vec![],
}),
},
n,
&[1],
),
);
ws.suff.add_rows_multi(x.as_ref(), &y, &ids, &[], None);
let fit = fit_lmm(&mut ws, &[1], None);
assert!(fit.converged);
assert!(
ws.theta[2] == 0.0,
"slope λ₁₁ must pin to exactly 0, got {:e}",
ws.theta[2]
);
assert!(fit.pinned_components & 0b10 != 0, "slope component bit set");
assert!(
ws.theta[0] > PIN_THETA,
"intercept component must stay interior"
);
assert!(
ws.theta[0] < THETA_HI,
"intercept component must be off the box bound"
);
}
// -----------------------------------------------------------------------
// Composition: primary slope (1 + x1 | g) co-existing with an
// intercept-only crossed (1 | item) / nested (1 | g:sub) extra. The
// family-blocked deviance must match a brute-force V = I + Z_p D_p Z_p′ +
// τ_e² Z_e Z_e′. Data on the f32 plane (the suff-stats input convention);
// the oracle widens the identical bytes, so the 1e-8 match is exact.
// -----------------------------------------------------------------------
/// n=80, p=2 (intercept + x1), 8 primary clusters crossed with 5 items;
/// y carries u₀ + u₁·x1 (primary) + v (item intercept).
fn composed_dataset() -> (Mat<f64>, Vec<f64>, Vec<u32>, Vec<u32>) {
let (n, nc, ni) = (80usize, 8usize, 5usize);
let mut st = 41u64;
let u0: Vec<f64> = (0..nc).map(|_| 0.5 * lcg(&mut st)).collect();
let u1: Vec<f64> = (0..nc).map(|_| 0.3 * lcg(&mut st)).collect();
let v: Vec<f64> = (0..ni).map(|_| 0.4 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let (mut pid, mut iid) = (vec![0u32; n], vec![0u32; n]);
for i in 0..n {
let (c, it) = (i % nc, i % ni);
pid[i] = c as u32;
iid[i] = it as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5 + 0.4 * x1 + u0[c] + u1[c] * x1 + v[it] + 0.8 * lcg(&mut st);
}
(x, y, pid, iid)
}
/// primary (1 + x1 | g), crossed (1 | item); slope on x_full col 1.
fn composed_groupings() -> LmmGroupings {
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 8 },
slopes: vec![0],
extra_groupings: vec![Grouping {
relation: GroupingRelation::Crossed { n_clusters: 5 },
slopes: vec![],
}],
}),
};
LmmGroupings::from_cluster_spec(&cluster, 80, &[1])
}
// --- θ-layout generalization (scalar → vech ranges) ---
/// One intercept-only primary + one crossed grouping of RE width `q_g`
/// (intercept + `q_g−1` slopes), expressed through the slope machinery — the
/// θ-layout fixture. Slope columns are placeholders (layout reads only
/// `slopes.len()`).
fn groupings_primary1_crossed_qg(q_g: usize) -> LmmGroupings {
let slopes: Vec<crate::ColumnId> = (0..q_g - 1).map(|k| (k + 1) as u32).collect();
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 8 },
slopes: vec![],
extra_groupings: vec![Grouping {
relation: GroupingRelation::Crossed { n_clusters: 5 },
slopes,
}],
}),
};
LmmGroupings::from_cluster_spec(&cluster, 80, &[])
}
#[test]
fn extra_qg1_theta_layout_matches_scalar() {
// Intercept-only crossed factor through the slope machinery = the plain
// scalar layout: one primary scalar + one extra scalar.
let g = groupings_primary1_crossed_qg(1);
assert_eq!(g.n_theta(), 1 + 1);
assert_eq!(g.crossed[0].vech_start, 1);
assert_eq!(g.crossed[0].q, 1);
assert!(!g.extra_slopes_any);
}
#[test]
fn extra_qg2_theta_packs_vech3() {
let g = groupings_primary1_crossed_qg(2);
assert_eq!(g.crossed[0].q, 2);
assert_eq!(g.n_theta(), 1 + 3); // primary scalar + vech(2×2)=3
assert!(g.extra_slopes_any);
// The extra block's two diagonal θ indices are vech_start (=1) and
// vech_start + 2 (=3) under the column-major lower-tri convention.
let diag = &g.diagonal_theta;
assert!(diag.contains(&1) && diag.contains(&3));
// Off-diagonal λ₁₀ at index 2 is NOT a diagonal (signed box).
assert!(!diag.contains(&2));
}
// --- Extra-slope sufficient statistics ---
/// Brute-force the `s` columns for a crossed factor carrying a slope: the
/// intercept subcol is Σ_{rows∈level} [X y]; the slope subcol is Σ x_slope·[X y].
#[test]
fn extra_crossed_slope_s_columns_match_bruteforce() {
let n = 6usize;
let p = 3; // [1, x1, x2]
let xd = [
(0.5, -0.2),
(-0.3, 0.7),
(0.9, 0.1),
(-0.6, -0.4),
(0.2, 0.8),
(0.4, -0.5),
];
let cluster_ids = [0u32, 1, 0, 1, 0, 1];
let crossed_ids = [0u32, 1, 2, 0, 1, 2];
let y = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
let mut x = Mat::<f64>::zeros(n, p);
for i in 0..n {
x[(i, 0)] = 1.0;
x[(i, 1)] = xd[i].0;
x[(i, 2)] = xd[i].1;
}
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 2 },
slopes: vec![],
extra_groupings: vec![Grouping {
relation: GroupingRelation::Crossed { n_clusters: 3 },
slopes: vec![1],
}],
}),
};
// crossed slope on x_full col 1.
let g = LmmGroupings::from_cluster_spec_ext(&cluster, n, &[], &[vec![1]]);
// crossed block: q_g=2, offset = prim_width = 2 (n_primary=2, q_p=1).
assert_eq!(g.extra_offsets[0], 2);
assert_eq!(g.extra_q[0], 2);
let mut suff = LmmSuffStats::with_groupings(p, g);
suff.add_rows_multi(x.as_ref(), &y, &cluster_ids, &[crossed_ids.to_vec()], None);
let m = p + 1;
for c in 0..3usize {
let icol = 2 + c * 2;
let scol = icol + 1;
let mut s_int = vec![0.0; m];
let mut s_slope = vec![0.0; m];
for i in 0..n {
if crossed_ids[i] as usize == c {
let w = [x[(i, 0)], x[(i, 1)], x[(i, 2)], y[i]];
let x1 = x[(i, 1)];
for j in 0..m {
s_int[j] += w[j];
s_slope[j] += x1 * w[j];
}
}
}
for j in 0..m {
assert!(
(suff.s[(j, icol)] - s_int[j]).abs() < 1e-12,
"intercept col level {c} row {j}: got {} want {}",
suff.s[(j, icol)],
s_int[j]
);
assert!(
(suff.s[(j, scol)] - s_slope[j]).abs() < 1e-12,
"slope col level {c} row {j}: got {} want {}",
suff.s[(j, scol)],
s_slope[j]
);
}
// counts only on the intercept subcol.
let n_c = crossed_ids.iter().filter(|&&l| l as usize == c).count() as f64;
assert_eq!(suff.counts[icol], n_c);
assert_eq!(suff.counts[scol], 0.0);
}
}
/// REML deviance on the explicit n×n V for the composed model: the 2×2
/// primary slope block (D_p = ΛΛ′ over [1, x1]) PLUS the extra-grouping
/// intercept block (θ_e² when the extra ids match). The f32 data is widened
/// to f64 so the oracle reads the same bytes the suff stats accumulated.
/// `eid` is the extra grouping's level id per row (item, or nested child).
/// θ = [primary vech λ₀₀, λ₁₀, λ₁₁ ; extra scalar θ_e].
fn brute_force_composed_deviance(
theta: &[f64],
x: &Mat<f64>,
y: &[f64],
pid: &[u32],
eid: &[u32],
) -> f64 {
let n = x.nrows();
let (a, b, c) = (theta[0], theta[1], theta[2]);
// D_p = ΛΛ′, Λ = [[a,0],[b,c]] (column-major vech).
let (d00, d01, d11) = (a * a, a * b, b * b + c * c);
let te2 = theta[3] * theta[3];
let mut v = Mat::<f64>::zeros(n, n);
for i in 0..n {
v[(i, i)] += 1.0;
}
for i in 0..n {
for j in 0..n {
if pid[i] == pid[j] {
let (zi1, zj1) = (x[(i, 1)], x[(j, 1)]);
v[(i, j)] += d00 + d01 * (zi1 + zj1) + d11 * zi1 * zj1;
}
if eid[i] == eid[j] {
v[(i, j)] += te2;
}
}
}
reml_profile_from_v(&v, x, y)
}
/// REML profiled deviance from an explicit n×n marginal V (in residual-σ²
/// units): `log|V| + log|XᵀV⁻¹X| + (N−P)·log σ̂²`. The shared V→deviance back
/// end for every brute-force oracle (composed, crossed-slope, …).
fn reml_profile_from_v(v: &Mat<f64>, x: &Mat<f64>, y: &[f64]) -> f64 {
use faer::linalg::solvers::Solve;
let (n, p) = (x.nrows(), x.ncols());
let vc = v.as_ref().llt(faer::Side::Lower).unwrap();
let mut ldv = 0.0;
for i in 0..n {
ldv += vc.L()[(i, i)].ln();
}
let ldv = 2.0 * ldv;
let mut vix = (*x).clone();
vc.solve_in_place(vix.as_mut());
let mut viy = Mat::<f64>::zeros(n, 1);
for i in 0..n {
viy[(i, 0)] = y[i];
}
vc.solve_in_place(viy.as_mut());
let mut xtvix = Mat::<f64>::zeros(p, p);
let mut xtviy = vec![0.0; p];
for aa in 0..p {
for bb in 0..p {
let mut s = 0.0;
for i in 0..n {
s += x[(i, aa)] * vix[(i, bb)];
}
xtvix[(aa, bb)] = s;
}
let mut s = 0.0;
for i in 0..n {
s += x[(i, aa)] * viy[(i, 0)];
}
xtviy[aa] = s;
}
let kc = xtvix.as_ref().llt(faer::Side::Lower).unwrap();
let mut ldk = 0.0;
for aa in 0..p {
ldk += kc.L()[(aa, aa)].ln();
}
let ldk = 2.0 * ldk;
let mut beta = Mat::<f64>::zeros(p, 1);
for aa in 0..p {
beta[(aa, 0)] = xtviy[aa];
}
kc.solve_in_place(beta.as_mut());
let mut ytviy = 0.0;
for i in 0..n {
ytviy += y[i] * viy[(i, 0)];
}
let mut bxy = 0.0;
for aa in 0..p {
bxy += beta[(aa, 0)] * xtviy[aa];
}
let df = (n - p) as f64;
let s2 = (ytviy - bxy) / df;
ldv + ldk + df * s2.ln()
}
/// Brute-force REML deviance for a CROSSED-SLOPE model
/// `y ~ x1 + (1+x1 | primary) + (1+x1 | crossed)`: V = I + Z_p D_p Z_pᵀ +
/// Z_e D_e Z_eᵀ, each D a 2×2 from its vech θ over [1, x1]. θ =
/// [primary vech (3) ; crossed vech (3)].
fn brute_force_crossed_slope_deviance(
theta: &[f64],
x: &Mat<f64>,
y: &[f64],
pid: &[u32],
eid: &[u32],
) -> f64 {
let n = x.nrows();
let (ap, bp, cp) = (theta[0], theta[1], theta[2]);
let (dp00, dp01, dp11) = (ap * ap, ap * bp, bp * bp + cp * cp);
let (ae, be, ce) = (theta[3], theta[4], theta[5]);
let (de00, de01, de11) = (ae * ae, ae * be, be * be + ce * ce);
let mut v = Mat::<f64>::zeros(n, n);
for i in 0..n {
v[(i, i)] += 1.0;
}
for i in 0..n {
for j in 0..n {
let (zi, zj) = (x[(i, 1)], x[(j, 1)]);
if pid[i] == pid[j] {
v[(i, j)] += dp00 + dp01 * (zi + zj) + dp11 * zi * zj;
}
if eid[i] == eid[j] {
v[(i, j)] += de00 + de01 * (zi + zj) + de11 * zi * zj;
}
}
}
reml_profile_from_v(&v, x, y)
}
/// Slope + crossed: the composed deviance matches the brute-force oracle to
/// 1e-8 — the slope-composition gate. zx_slope carries the slope↔crossed
/// coupling; the primary 2×2 block and the item intercept block are coupled
/// through the shared family-blocked tail.
#[test]
fn composed_deviance_matches_brute_force() {
let (x, y, pid, iid) = composed_dataset();
let mut suff = LmmSuffStats::with_groupings(2, composed_groupings());
suff.add_rows_multi(x.as_ref(), &y, &pid, std::slice::from_ref(&iid), None); // item ids as the single extra grouping
let mut fit = LmmFitScratch::with_groupings(2, &composed_groupings());
// θ = [λ₀₀, λ₁₀, λ₁₁, θ_c].
for th in [
vec![1.0, 0.0, 1.0, 0.5],
vec![0.6, 0.2, 0.4, 0.3],
vec![1.5, -0.4, 0.7, 0.8],
] {
let dev = reml_deviance(&th, &suff, &mut fit);
let oracle = brute_force_composed_deviance(&th, &x, &y, &pid, &iid);
assert!(dev.is_finite(), "θ={th:?}");
assert!(
(dev - oracle).abs() <= 1e-8 * oracle.abs().max(1.0),
"θ={th:?}: {dev} vs {oracle}"
);
}
}
/// CROSSED SLOPES (the headline lme4-agreement case): `y ~ x1 + (1+x1 | primary)
/// + (1+x1 | item)` — both grouping factors carry a random slope on x1, so the
/// gated blocked path runs. Deviance must match the explicit-V oracle to 1e-7
/// across θ, including the primary-slope↔crossed-slope coupling (the x1²
/// weighted co-occurrence) the blocked `zx` fill captures.
#[test]
fn crossed_slope_deviance_matches_brute_force() {
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 5 },
slopes: vec![1],
extra_groupings: vec![Grouping {
relation: GroupingRelation::Crossed { n_clusters: 4 },
slopes: vec![1],
}],
}),
};
let n = 60; // atom = 5·4 = 20 ⇒ 3 balanced blocks
let mut st = 91u64;
let u0p: Vec<f64> = (0..5).map(|_| 0.5 * lcg(&mut st)).collect();
let u1p: Vec<f64> = (0..5).map(|_| 0.3 * lcg(&mut st)).collect();
let u0e: Vec<f64> = (0..4).map(|_| 0.4 * lcg(&mut st)).collect();
let u1e: Vec<f64> = (0..4).map(|_| 0.3 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut eid = vec![0u32; n];
for i in 0..n {
let par = cluster.re.as_ref().unwrap().sizing.cluster_of_row(i);
let item = extra_level_of_row(&cluster, 0, i) as usize;
pid[i] = par as u32;
eid[i] = item as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5
+ 0.4 * x1
+ u0p[par]
+ u1p[par] * x1
+ u0e[item]
+ u1e[item] * x1
+ 0.8 * lcg(&mut st);
}
let g = LmmGroupings::from_cluster_spec_ext(&cluster, n, &[1], &[vec![1]]);
assert!(g.extra_slopes_any, "must route to the blocked path");
let mut suff = LmmSuffStats::with_groupings(2, g);
suff.add_rows_multi(x.as_ref(), &y, &pid, &[eid.clone()], None);
let gref = LmmGroupings::from_cluster_spec_ext(&cluster, n, &[1], &[vec![1]]);
let mut fit = LmmFitScratch::with_groupings(2, &gref);
// θ = [primary vech (λ₀₀,λ₁₀,λ₁₁) ; crossed vech (λ₀₀,λ₁₀,λ₁₁)].
for th in [
vec![1.0, 0.0, 1.0, 1.0, 0.0, 1.0],
vec![0.7, 0.2, 0.5, 0.6, 0.1, 0.4],
vec![1.3, -0.3, 0.6, 0.9, -0.2, 0.5],
] {
let dev = reml_deviance(&th, &suff, &mut fit);
let oracle = brute_force_crossed_slope_deviance(&th, &x, &y, &pid, &eid);
assert!(dev.is_finite(), "θ={th:?}");
assert!(
(dev - oracle).abs() <= 1e-7 * oracle.abs().max(1.0),
"θ={th:?}: {dev} vs {oracle}"
);
}
}
/// NESTED SLOPES (the nested-slope defect): `y ~ x1 + (1+x1 | grp) + (1+x1 | class)`,
/// class nested in grp — both grouping factors carry a random slope on x1, so
/// the gated blocked path runs with a nested factor of q_n = 2. Before the
/// fix the blocked path assembled the nested children intercept-only (scalar
/// θ_n), diverging to NaN. The marginal V is grouping-agnostic (Σ_g Z_g D_g Z_gᵀ
/// over rows sharing a level id), so the crossed-slope oracle is reused with the
/// GLOBAL nested child id as the extra level. Matches the explicit-V oracle to
/// 1e-7 across θ.
#[test]
fn nested_slope_deviance_matches_brute_force() {
let n_per_parent = 3u32;
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 5 },
slopes: vec![1],
extra_groupings: vec![Grouping {
relation: GroupingRelation::NestedWithin { n_per_parent },
slopes: vec![1],
}],
}),
};
let n = 60; // atom = primary 5 · nested 3 = 15 ⇒ 4 balanced blocks
let n_child = 5 * n_per_parent as usize;
let mut st = 137u64;
let u0p: Vec<f64> = (0..5).map(|_| 0.5 * lcg(&mut st)).collect();
let u1p: Vec<f64> = (0..5).map(|_| 0.3 * lcg(&mut st)).collect();
let u0e: Vec<f64> = (0..n_child).map(|_| 0.4 * lcg(&mut st)).collect();
let u1e: Vec<f64> = (0..n_child).map(|_| 0.3 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut eid = vec![0u32; n];
for i in 0..n {
let par = cluster.re.as_ref().unwrap().sizing.cluster_of_row(i);
let child = extra_level_of_row(&cluster, 0, i); // GLOBAL child id
pid[i] = par as u32;
eid[i] = child as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5
+ 0.4 * x1
+ u0p[par]
+ u1p[par] * x1
+ u0e[child]
+ u1e[child] * x1
+ 0.8 * lcg(&mut st);
}
let g = LmmGroupings::from_cluster_spec_ext(&cluster, n, &[1], &[vec![1]]);
assert!(g.extra_slopes_any, "must route to the blocked path");
assert!(g.nested.is_some(), "must carry a nested factor");
let mut suff = LmmSuffStats::with_groupings(2, g);
suff.add_rows_multi(x.as_ref(), &y, &pid, &[eid.clone()], None);
let gref = LmmGroupings::from_cluster_spec_ext(&cluster, n, &[1], &[vec![1]]);
let mut fit = LmmFitScratch::with_groupings(2, &gref);
// θ = [primary vech (λ₀₀,λ₁₀,λ₁₁) ; nested vech (λ₀₀,λ₁₀,λ₁₁)].
for th in [
vec![1.0, 0.0, 1.0, 1.0, 0.0, 1.0],
vec![0.7, 0.2, 0.5, 0.6, 0.1, 0.4],
vec![1.3, -0.3, 0.6, 0.9, -0.2, 0.5],
] {
let dev = reml_deviance(&th, &suff, &mut fit);
let oracle = brute_force_crossed_slope_deviance(&th, &x, &y, &pid, &eid);
assert!(dev.is_finite(), "θ={th:?}");
assert!(
(dev - oracle).abs() <= 1e-7 * oracle.abs().max(1.0),
"θ={th:?}: {dev} vs {oracle}"
);
}
}
/// End-to-end NESTED-SLOPE fit: guards against a mis-assembled blocked
/// objective, which makes BOBYQA diverge to NaN (`converged = false`) on
/// every seed. With the correct objective the full θ-search must converge to
/// a finite interior fit. Asserts `converged`, no numerical failure
/// (`boundary_hit != 2`), finite θ̂/σ̂², and β̂ recovered near the planted
/// [0.5, 0.4].
#[test]
fn nested_slope_fit_converges() {
let n_per_parent = 3u32;
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 5 },
slopes: vec![1],
extra_groupings: vec![Grouping {
relation: GroupingRelation::NestedWithin { n_per_parent },
slopes: vec![1],
}],
}),
};
let n = 120; // atom = 5·3 = 15 ⇒ 8 balanced blocks
let n_child = 5 * n_per_parent as usize;
let mut st = 137u64;
let u0p: Vec<f64> = (0..5).map(|_| 0.5 * lcg(&mut st)).collect();
let u1p: Vec<f64> = (0..5).map(|_| 0.3 * lcg(&mut st)).collect();
let u0e: Vec<f64> = (0..n_child).map(|_| 0.4 * lcg(&mut st)).collect();
let u1e: Vec<f64> = (0..n_child).map(|_| 0.3 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut eid = vec![0u32; n];
for i in 0..n {
let par = cluster.re.as_ref().unwrap().sizing.cluster_of_row(i);
let child = extra_level_of_row(&cluster, 0, i);
pid[i] = par as u32;
eid[i] = child as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5
+ 0.4 * x1
+ u0p[par]
+ u1p[par] * x1
+ u0e[child]
+ u1e[child] * x1
+ 0.8 * lcg(&mut st);
}
let mut ws = LmmWorkspace::for_cluster_spec_ext(2, &cluster, n, &[1], &[vec![1]]);
ws.suff.reset();
ws.suff
.add_rows_multi(x.as_ref(), &y, &pid, &[eid.clone()], None);
let fit = fit_lmm(&mut ws, &[1], None);
assert!(fit.converged, "nested-slope fit must converge");
assert_ne!(fit.boundary_hit, 2, "must not be a numerical (NaN) failure");
assert!(
fit.sigma_sq.is_finite() && fit.sigma_sq > 0.0,
"σ̂² {}",
fit.sigma_sq
);
assert!(ws.theta.iter().all(|t| t.is_finite()), "θ̂ {:?}", ws.theta);
assert!(
(0.2..0.8).contains(&ws.fit.betas[0]),
"intercept {}",
ws.fit.betas[0]
);
assert!(
(0.1..0.7).contains(&ws.fit.betas[1]),
"slope {}",
ws.fit.betas[1]
);
// Deterministic regression lock (seed 137) alongside the wide recovers-check above.
assert!((ws.fit.betas[0] - 0.6209080774915476).abs() / 0.6209080774915476_f64.abs() < 1e-6);
assert!((ws.fit.betas[1] - 0.257915422474595).abs() / 0.257915422474595_f64.abs() < 1e-6);
}
/// General brute-force REML deviance: V = I + Σ_g Z_g D_g Z_gᵀ where each
/// factor `(ids, vech)` contributes a 2×2 D over [1, x1] (an intercept-only
/// factor passes `[θ, 0, 0]`). Used for the multi-crossed-factor oracle.
fn brute_force_slopes_deviance(x: &Mat<f64>, y: &[f64], factors: &[(&[u32], [f64; 3])]) -> f64 {
let n = x.nrows();
let mut v = Mat::<f64>::zeros(n, n);
for i in 0..n {
v[(i, i)] += 1.0;
}
for &(ids, vech) in factors {
let (a, b, c) = (vech[0], vech[1], vech[2]);
let (d00, d01, d11) = (a * a, a * b, b * b + c * c);
for i in 0..n {
for j in 0..n {
if ids[i] == ids[j] {
let (zi, zj) = (x[(i, 1)], x[(j, 1)]);
v[(i, j)] += d00 + d01 * (zi + zj) + d11 * zi * zj;
}
}
}
}
reml_profile_from_v(&v, x, y)
}
/// TWO crossed factors with slopes:
/// `y ~ x1 + (1 | primary) + (1+x1 | c1) + (1+x1 | c2)`. Exercises the
/// crossed↔crossed slope coupling (c1's slope column against c2's, the x1²
/// weighted co-occurrence between two distinct crossed factors) — the part
/// neither the composed nor single-crossed test reaches. Matches the
/// explicit-V oracle to 1e-7.
#[test]
fn two_crossed_slopes_deviance_matches_brute_force() {
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 3 },
slopes: vec![], // primary intercept-only
extra_groupings: vec![
Grouping {
relation: GroupingRelation::Crossed { n_clusters: 3 },
slopes: vec![1],
},
Grouping {
relation: GroupingRelation::Crossed { n_clusters: 3 },
slopes: vec![1],
},
],
}),
};
let n = 54; // atom = 3·3·3 = 27 ⇒ 2 blocks
let mut st = 73u64;
let up: Vec<f64> = (0..3).map(|_| 0.45 * lcg(&mut st)).collect();
let u0a: Vec<f64> = (0..3).map(|_| 0.4 * lcg(&mut st)).collect();
let u1a: Vec<f64> = (0..3).map(|_| 0.3 * lcg(&mut st)).collect();
let u0b: Vec<f64> = (0..3).map(|_| 0.35 * lcg(&mut st)).collect();
let u1b: Vec<f64> = (0..3).map(|_| 0.28 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut c1 = vec![0u32; n];
let mut c2 = vec![0u32; n];
for i in 0..n {
let par = cluster.re.as_ref().unwrap().sizing.cluster_of_row(i);
let a = extra_level_of_row(&cluster, 0, i) as usize;
let b = extra_level_of_row(&cluster, 1, i) as usize;
pid[i] = par as u32;
c1[i] = a as u32;
c2[i] = b as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5
+ 0.4 * x1
+ up[par]
+ u0a[a]
+ u1a[a] * x1
+ u0b[b]
+ u1b[b] * x1
+ 0.8 * lcg(&mut st);
}
let g = LmmGroupings::from_cluster_spec_ext(&cluster, n, &[], &[vec![1], vec![1]]);
assert!(g.extra_slopes_any);
let mut suff = LmmSuffStats::with_groupings(2, g);
suff.add_rows_multi(x.as_ref(), &y, &pid, &[c1.clone(), c2.clone()], None);
let gref = LmmGroupings::from_cluster_spec_ext(&cluster, n, &[], &[vec![1], vec![1]]);
let mut fit = LmmFitScratch::with_groupings(2, &gref);
// θ = [primary scalar ; c1 vech (3) ; c2 vech (3)].
for th in [
vec![1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0],
vec![0.6, 0.7, 0.2, 0.4, 0.6, -0.1, 0.35],
vec![0.8, 1.2, -0.3, 0.5, 0.9, 0.25, 0.45],
] {
let dev = reml_deviance(&th, &suff, &mut fit);
let oracle = brute_force_slopes_deviance(
&x,
&y,
&[
(&pid, [th[0], 0.0, 0.0]),
(&c1, [th[1], th[2], th[3]]),
(&c2, [th[4], th[5], th[6]]),
],
);
assert!(dev.is_finite(), "θ={th:?}");
assert!(
(dev - oracle).abs() <= 1e-7 * oracle.abs().max(1.0),
"θ={th:?}: {dev} vs {oracle}"
);
}
}
/// Deterministic crossed-slope dataset for the lme4 golden: 8 primary
/// clusters × 6 crossed levels × 2 reps (n=96),
/// `y = 1.0 + 0.8·x1 + u0p + u1p·x1 + u0e + u1e·x1 + ε`. The Rust generator is
/// the source of truth; `dump_crossed_slope_golden_csv` writes it for the R
/// `lme4::lmer` reference whose fit is frozen in `GOLDEN_LME4_*`.
fn crossed_slope_golden_dataset() -> (Mat<f64>, Vec<f64>, Vec<u32>, Vec<u32>) {
let (n_prim, n_cross, n) = (8usize, 6usize, 96usize);
let mut st = 20260629u64;
let u0p: Vec<f64> = (0..n_prim).map(|_| 0.7 * lcg(&mut st)).collect();
let u1p: Vec<f64> = (0..n_prim).map(|_| 0.5 * lcg(&mut st)).collect();
let u0e: Vec<f64> = (0..n_cross).map(|_| 0.6 * lcg(&mut st)).collect();
let u1e: Vec<f64> = (0..n_cross).map(|_| 0.4 * lcg(&mut st)).collect();
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut eid = vec![0u32; n];
for i in 0..n {
let pp = i % n_prim; // FixedClusters primary: i % n_clusters
let ee = (i / n_prim) % n_cross; // crossed: (i / n_prim) % n_cross
pid[i] = pp as u32;
eid[i] = ee as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] =
1.0 + 0.8 * x1 + u0p[pp] + u1p[pp] * x1 + u0e[ee] + u1e[ee] * x1 + 0.5 * lcg(&mut st);
}
(x, y, pid, eid)
}
/// Run once (`cargo test -p ... dump_crossed_slope_golden_csv -- --ignored`)
/// to regenerate the CSV the R reference reads. Not a normal test.
#[test]
#[ignore]
fn dump_crossed_slope_golden_csv() {
// Serialized under alloc-tests so its allocations can't land in a
// concurrent dhat profiler window on an `-- --ignored` run.
#[cfg(feature = "alloc-tests")]
let _serial = crate::test_support::alloc_test_guard();
let (x, y, pid, eid) = crossed_slope_golden_dataset();
let mut s = String::from("x1,y,pid,eid\n");
for i in 0..y.len() {
s.push_str(&format!("{},{},{},{}\n", x[(i, 1)], y[i], pid[i], eid[i]));
}
std::fs::write("/tmp/crossed_slope_golden.csv", s).unwrap();
}
/// L3 golden: `glmm`'s crossed-slope fit must reproduce `lme4::lmer`'s REML fit
/// of `y ~ x1 + (1+x1|pid) + (1+x1|eid)` on the committed dataset — fixed
/// effects, residual σ², and both 2×2 RE covariances. Frozen from
/// `/tmp/golden_fit.R` (lme4 1.1, bobyqa). Recovered D_g = σ̂²·Λ_gΛ_gᵀ from θ̂.
#[test]
fn crossed_slope_fit_matches_lme4_golden() {
// lme4 golden (REML, bobyqa).
const G_BETA0: f64 = 1.0582083262;
const G_BETA1: f64 = 0.6334043248;
const G_SIGMA2: f64 = 0.0921249591;
const G_PID_V0: f64 = 0.1406815355; // var(intercept)
const G_PID_V1: f64 = 0.1237856496; // var(x1)
const G_PID_COV: f64 = 0.0127473486;
const G_EID_V0: f64 = 0.1828301299;
const G_EID_V1: f64 = 0.0396985129;
const G_EID_COV: f64 = -0.0456611171;
let (x, y, pid, eid) = crossed_slope_golden_dataset();
let n = y.len();
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 8 },
slopes: vec![1],
extra_groupings: vec![Grouping {
relation: GroupingRelation::Crossed { n_clusters: 6 },
slopes: vec![1],
}],
}),
};
let mut ws = LmmWorkspace::for_cluster_spec_ext(2, &cluster, n, &[1], &[vec![1]]);
ws.suff.reset();
ws.suff
.add_rows_multi(x.as_ref(), &y, &pid, std::slice::from_ref(&eid), None);
let fit = fit_lmm(&mut ws, &[1], None);
assert!(fit.converged, "golden fit must converge");
let s2 = fit.sigma_sq;
// Fixed effects + residual variance.
assert!(
(ws.fit.betas[0] - G_BETA0).abs() < 1e-4,
"β0 {} vs {G_BETA0}",
ws.fit.betas[0]
);
assert!(
(ws.fit.betas[1] - G_BETA1).abs() < 1e-4,
"β1 {} vs {G_BETA1}",
ws.fit.betas[1]
);
assert!(
(s2 - G_SIGMA2).abs() <= 1e-3 * G_SIGMA2,
"σ² {s2} vs {G_SIGMA2}"
);
// D_g = σ̂²·Λ_gΛ_gᵀ from θ̂ (primary vech θ[0..3], crossed vech θ[3..6]).
let dblock = |t: &[f64]| {
let (a, b, c) = (t[0], t[1], t[2]);
(s2 * a * a, s2 * (b * b + c * c), s2 * a * b) // (v0, v1, cov)
};
let (pv0, pv1, pcov) = dblock(&ws.theta[0..3]);
let (ev0, ev1, ecov) = dblock(&ws.theta[3..6]);
let close = |got: f64, want: f64, name: &str| {
assert!(
(got - want).abs() <= 2e-3 * want.abs().max(1e-3),
"{name}: {got} vs {want}"
);
};
close(pv0, G_PID_V0, "pid var0");
close(pv1, G_PID_V1, "pid var1");
close(pcov, G_PID_COV, "pid cov");
close(ev0, G_EID_V0, "eid var0");
close(ev1, G_EID_V1, "eid var1");
close(ecov, G_EID_COV, "eid cov");
}
/// Slope + NESTED: `(1 + x1 | g) + (1 | g:sub)` — the composed deviance with
/// a nested child tail (vs the crossed tail above). Exercises the
/// primary-slope↔child off-diagonal (read from `s`) and the shifted nested
/// offset `q_p·n_primary + f·np + c`. The nested child ids are globalized
/// (parent·np + within) — the workspace layout the contract helpers produce.
#[test]
fn composed_nested_deviance_matches_brute_force() {
// 8 primary clusters × 2 children each, fixed-size 8 ⇒ 64 rows / 4 blocks.
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedSize { cluster_size: 8 },
slopes: vec![0],
extra_groupings: vec![Grouping {
relation: GroupingRelation::NestedWithin { n_per_parent: 2 },
slopes: vec![],
}],
}),
};
let n = 4 * model_atom(&cluster); // 64
let mut st = 47u64;
let u0: Vec<f64> = (0..8).map(|_| 0.5 * lcg(&mut st)).collect();
let u1: Vec<f64> = (0..8).map(|_| 0.3 * lcg(&mut st)).collect();
let u_c: Vec<f64> = (0..16).map(|_| 0.35 * lcg(&mut st)).collect(); // 8 parents × 2 children
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut cid = vec![0u32; n]; // globalized child id (parent·np + within)
for i in 0..n {
let par = cluster.re.as_ref().unwrap().sizing.cluster_of_row(i);
let child = extra_level_of_row(&cluster, 0, i); // already globalized par·np + within
pid[i] = par as u32;
cid[i] = child as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5 + 0.4 * x1 + u0[par] + u1[par] * x1 + u_c[child] + 0.8 * lcg(&mut st);
}
let g = LmmGroupings::from_cluster_spec(&cluster, n, &[1]);
let mut suff = LmmSuffStats::with_groupings(2, g);
suff.add_rows_multi(x.as_ref(), &y, &pid, &[cid.clone()], None);
let gref = LmmGroupings::from_cluster_spec(&cluster, n, &[1]);
let mut fit = LmmFitScratch::with_groupings(2, &gref);
// The brute-force oracle is V-shape-agnostic: the nested child block adds
// θ_n² when the (globalized) child ids match — same form as the crossed.
for th in [
vec![1.0, 0.0, 1.0, 0.5],
vec![0.7, 0.25, 0.5, 0.4],
vec![1.3, -0.3, 0.6, 0.2],
] {
let dev = reml_deviance(&th, &suff, &mut fit);
let oracle = brute_force_composed_deviance(&th, &x, &y, &pid, &cid);
assert!(dev.is_finite(), "θ={th:?}");
assert!(
(dev - oracle).abs() <= 1e-8 * oracle.abs().max(1.0),
"θ={th:?}: {dev} vs {oracle}"
);
}
}
/// Bounded-allocation twin — the standalone slope workspace
/// allocates only faer `llt` internals on the warm `fit_lmm` loop, the same
/// acceptance class as the q=1 / general twins. `alloc_test_guard` serializes
/// test bodies, but libtest's own per-test thread spawn still needs
/// `--test-threads=1`:
/// RAYON_NUM_THREADS=1 cargo test -p glmm --features alloc-tests lmm_fit_slope_warm_path_bounded_alloc -- --ignored --test-threads=1
#[cfg(feature = "alloc-tests")]
#[test]
#[ignore]
fn lmm_fit_slope_warm_path_bounded_alloc() {
let _serial = crate::test_support::alloc_test_guard();
const N_CALLS: usize = 100;
const BOUND_SLOPE: u64 = 5500; // Measured 5200 (this machine) — ~52 blocks/fit of faer `llt` internals (one m×m tail llt per eval × ~54 evals on the blind 3-D q_p=2 surface; the family loop + primary Λ/Gram are zero-alloc scratch, the cached diagonal_theta map avoids a per-fit Vec allocation, and the ranef recovery pass solves in the ranef_ux/ranef_rhs scratch fields). Higher total than q=1's 2500 only via the larger blind eval count, not a richer per-eval alloc — faer-version/machine specific. `chol_lower` factors in place, which keeps this count low. If faer's Cholesky internals change, update — do not relax.
let (x, y, ids) = slope_dataset();
let targets: Vec<u32> = vec![1];
let mut ws = LmmWorkspace::with_groupings(2, slope_groupings());
ws.suff.reset();
ws.suff.add_rows_multi(x.as_ref(), &y, &ids, &[], None);
let _ = fit_lmm(&mut ws, &targets, None);
let profiler = dhat::Profiler::builder().testing().build();
for _ in 0..N_CALLS {
ws.suff.reset();
ws.suff.add_rows_multi(x.as_ref(), &y, &ids, &[], None);
let _ = fit_lmm(&mut ws, &targets, None);
}
let stats = dhat::HeapStats::get();
drop(profiler);
assert!(
stats.total_blocks <= BOUND_SLOPE,
"slope fit_lmm allocated {} blocks across {} warm-path calls (BOUND = {})",
stats.total_blocks,
N_CALLS,
BOUND_SLOPE
);
}
// -----------------------------------------------------------------------
// Task 10: LMM REML dual gradient/Hessian vs. `dense_reml_score`, a
// hand-written verification reference sharing no code with the blocked
// kernel. Gate fixtures reuse the existing int1/nest2/q2s/q3s LMM shapes
// already built above in this module.
// -----------------------------------------------------------------------
/// Dense `n × k` RE design Z for [`dense_reml_score`]. LEVEL-major column
/// order in the primary block (`f·q + c`) — matching that function's own Λ
/// packing, NOT `LmmGroupings`'s own scattered `d·n_primary + f` convention
/// (see `dense_reml_score`'s doc comment on Λ). Extra groupings (`q_g == 1`
/// only, matching every gate fixture) get one one-hot column per level,
/// declaration order, right after the primary block. Z and Λ agreeing on this
/// ordering is what keeps `V = I + ZΛΛᵀZᵀ` correct; the two orderings agree
/// up to a simultaneous column permutation of Z and Λ, so using an
/// independent one here (rather than `LmmGroupings`'s own) is the point —
/// the reference shares no code path with the kernel.
fn dense_z(
groupings: &LmmGroupings,
primary_ids: &[u32],
extra_ids: &[Vec<u32>],
x: &Mat<f64>,
slope_cols: &[usize],
) -> Mat<f64> {
let n = primary_ids.len();
let q = groupings.primary_q;
let s = groupings.n_primary;
let prim_width = q * s;
let extra_widths: Vec<usize> = extra_ids
.iter()
.map(|ids| ids.iter().copied().max().map_or(0, |m| m as usize + 1))
.collect();
let k = prim_width + extra_widths.iter().sum::<usize>();
let mut z = Mat::<f64>::zeros(n, k);
for i in 0..n {
let f = primary_ids[i] as usize;
z[(i, f * q)] = 1.0;
for (d, &sc) in slope_cols.iter().enumerate() {
z[(i, f * q + 1 + d)] = x[(i, sc)];
}
}
let mut off = prim_width;
for (e, ids) in extra_ids.iter().enumerate() {
for i in 0..n {
z[(i, off + ids[i] as usize)] = 1.0;
}
off += extra_widths[e];
}
z
}
/// Dense REML score, the verification reference for `reml_gradient`.
///
/// **The criterion it differentiates**, matching `reml_deviance`'s own
/// normalization comment (`src/lmm/kernel.rs:739–742`) exactly:
///
/// `d(θ) = log|V| + log|XᵀV⁻¹X| + (N−P)·log σ̂²`, `V = I + ZΛΛᵀZᵀ`
///
/// Two things about that statement decide the score's form and both are easy to
/// get wrong. First, `V` is the σ²-FREE relative covariance — σ² is not a
/// parameter of `V`, it is profiled out as `σ̂² = yᵀPy/(N−P)` and re-enters only
/// through the third term. Second, the criterion carries NO additive Gaussian
/// REML constant: no `(N−P)log 2π`, no `+(N−P)`. lme4's `REMLcrit` does carry
/// it, which is why the speed-grid's `analyze.R` adds it back before comparing.
/// A reference that includes it would still have the right SCORE (a constant
/// differentiates away) but would not let the deviance values be compared, so
/// write the criterion without it and say so.
///
/// With `P = V⁻¹ − V⁻¹X(XᵀV⁻¹X)⁻¹XᵀV⁻¹`, `∂(log|V| + log|XᵀV⁻¹X|)/∂θ_j =
/// tr(P V_j)` and `∂(yᵀPy)/∂θ_j = −yᵀP V_j P y`, so
///
/// `∂d/∂θ_j = tr(P V_j) − (N−P)·(yᵀP V_j P y)/(yᵀPy)`
///
/// (Harville 1977 eq. 4; the same identity AI-REML differentiates — Gilmour,
/// Thompson & Cullis 1995 §2). The `(N−P)/(yᵀPy)` factor is the profiled σ̂²'s
/// contribution; dropping it gives the score of the criterion at fixed σ² = 1,
/// which is a different function and will miss the 1e-8 band everywhere.
/// `Py = Pr` with `r = y − Xβ̂_GLS`, since `PX = 0`.
///
/// Every matrix is formed densely and factored with faer: O(n³) is fine for a
/// test fixture at n ≤ 200, and the point is that this shares no code path with
/// the blocked kernel.
///
/// `z` is the dense `n × k` RE design, `x` the `n × p` fixed design, `theta` the
/// vech-packed Cholesky parameters. Returns `n_theta` entries in `theta`'s order.
#[cfg(test)]
fn dense_reml_score(
theta: &[f64],
z: MatRef<f64>,
x: MatRef<f64>,
y: &[f64],
groupings: &LmmGroupings,
) -> Vec<f64> {
use faer::linalg::solvers::Solve;
let (n, k) = (z.nrows(), z.ncols());
let p = x.ncols();
let q = groupings.primary_q;
let s = groupings.n_primary;
let n_theta = theta.len();
// 1. Λ (k×k) from theta. The packing is `primary_lambda`'s
// (`src/lmm/mod.rs:1051`), rebuilt here rather than called so the
// reference shares nothing with the kernel: theta is column-major vech
// of a LOWER-triangular q×q block, enumerated column c outer, row r ≥ c
// inner, i.e. theta[t] ↦ Λ_p[(r, c)] for t = 0, 1, 2, … over
// (c, r) = (0,0), (0,1), …, (0,q−1), (1,1), …, (q−1,q−1).
// Λ is that q×q block repeated on the diagonal once per primary level,
// followed by one scalar θ per extra grouping level (q_g == 1 on every
// extra in the gate cells), in declaration order.
let mut lam = Mat::<f64>::zeros(k, k);
let mut slot = vec![(0usize, 0usize); theta.len()]; // theta index → (r, c) in the q×q block
{
let mut t = 0;
for c in 0..q {
for r in c..q {
slot[t] = (r, c);
t += 1;
}
}
}
let prim_theta = q * (q + 1) / 2;
for f in 0..s {
for (t, &(r, c)) in slot[..prim_theta].iter().enumerate() {
lam[(f * q + r, f * q + c)] = theta[t];
}
}
// Extra groupings: one scalar θ per level, own diagonal band — matches
// `dense_z`'s column order (primary block first, then each extra grouping
// in declaration order). q_g == 1 on every gate fixture.
let mut extras: Vec<(usize, usize, usize)> = Vec::new(); // (vech_start, n_levels, lam_col_offset)
let mut off = q * s;
if let Some(nf) = groupings.nested {
debug_assert_eq!(nf.q, 1, "dense_reml_score only covers q_g == 1 extras");
let n_levels = s * groupings.nested_per_parent;
extras.push((nf.vech_start, n_levels, off));
off += n_levels;
}
for cf in &groupings.crossed {
debug_assert_eq!(cf.q, 1, "dense_reml_score only covers q_g == 1 extras");
extras.push((cf.vech_start, cf.n_levels, off));
off += cf.n_levels;
}
debug_assert_eq!(off, k, "z's columns must match the groupings layout");
for &(vs, n_levels, lam_off) in &extras {
for l in 0..n_levels {
lam[(lam_off + l, lam_off + l)] = theta[vs];
}
}
// 2. V = I + Z Λ Λᵀ Zᵀ. σ² is NOT here — it profiles out (see above).
let mut zl = Mat::<f64>::zeros(n, k);
for i in 0..n {
for c in 0..k {
let mut acc = 0.0;
for mm in 0..k {
acc += z[(i, mm)] * lam[(mm, c)];
}
zl[(i, c)] = acc;
}
}
let mut v = Mat::<f64>::zeros(n, n);
for i in 0..n {
v[(i, i)] = 1.0;
}
for i in 0..n {
for j in 0..n {
let mut acc = 0.0;
for c in 0..k {
acc += zl[(i, c)] * zl[(j, c)];
}
v[(i, j)] += acc;
}
}
// 3. Vi applied to X, Z, y through one Cholesky factorization — Vi is
// never formed as an explicit n×n inverse; XtViX and its own Cholesky
// solve the GLS step.
let vc = v
.as_ref()
.llt(faer::Side::Lower)
.expect("V must be SPD at a valid θ draw");
let mut vix = Mat::<f64>::zeros(n, p);
for i in 0..n {
for a in 0..p {
vix[(i, a)] = x[(i, a)];
}
}
vc.solve_in_place(vix.as_mut());
let mut viz = Mat::<f64>::zeros(n, k);
for i in 0..n {
for c in 0..k {
viz[(i, c)] = z[(i, c)];
}
}
vc.solve_in_place(viz.as_mut());
let mut viy = Mat::<f64>::zeros(n, 1);
for i in 0..n {
viy[(i, 0)] = y[i];
}
vc.solve_in_place(viy.as_mut());
let mut xtvix = Mat::<f64>::zeros(p, p);
for a in 0..p {
for b in 0..p {
let mut acc = 0.0;
for i in 0..n {
acc += x[(i, a)] * vix[(i, b)];
}
xtvix[(a, b)] = acc;
}
}
let mut xtviz = Mat::<f64>::zeros(p, k);
for a in 0..p {
for c in 0..k {
let mut acc = 0.0;
for i in 0..n {
acc += x[(i, a)] * viz[(i, c)];
}
xtviz[(a, c)] = acc;
}
}
let mut xtviy = Mat::<f64>::zeros(p, 1);
for a in 0..p {
let mut acc = 0.0;
for i in 0..n {
acc += x[(i, a)] * viy[(i, 0)];
}
xtviy[(a, 0)] = acc;
}
let kc = xtvix
.as_ref()
.llt(faer::Side::Lower)
.expect("X'ViX must be SPD (X full column rank)");
// 4. β̂_GLS = (X'ViX)⁻¹X'Viy; Py = Vi r with r = y − Xβ̂ (Vi is linear, so
// Vi r = Vi y − Vi X β̂ needs no extra solve). X'Vi r = 0 by
// construction (the GLS normal equations), which is exactly why
// Py = Pr = Vi r carries no projector correction — the doc comment's
// "Py = Pr … since PX = 0".
let mut beta_hat = xtviy.clone();
kc.solve_in_place(beta_hat.as_mut());
let mut py = Mat::<f64>::zeros(n, 1);
for i in 0..n {
let mut acc = viy[(i, 0)];
for b in 0..p {
acc -= vix[(i, b)] * beta_hat[(b, 0)];
}
py[(i, 0)] = acc;
}
let mut ypy = 0.0;
for i in 0..n {
ypy += y[i] * py[(i, 0)];
}
let df = (n - p) as f64;
// Z'PZ, reduced to k×k by the same Woodbury-style identity:
// Z'PZ = Z'ViZ − (X'ViZ)'(X'ViX)⁻¹(X'ViZ)
// so the per-θ_j trace below never needs an n×n P (see the cyclic-trace
// note at the score loop).
let mut ztviz = Mat::<f64>::zeros(k, k);
for a in 0..k {
for b in 0..k {
let mut acc = 0.0;
for i in 0..n {
acc += z[(i, a)] * viz[(i, b)];
}
ztviz[(a, b)] = acc;
}
}
let mut w = xtviz.clone(); // (X'ViX)⁻¹ X'ViZ, p×k
kc.solve_in_place(w.as_mut());
let mut ztpz = Mat::<f64>::zeros(k, k);
for a in 0..k {
for b in 0..k {
let mut acc = 0.0;
for i in 0..p {
acc += xtviz[(i, a)] * w[(i, b)];
}
ztpz[(a, b)] = ztviz[(a, b)] - acc;
}
}
let mut zpy = vec![0.0f64; k];
for a in 0..k {
let mut acc = 0.0;
for i in 0..n {
acc += z[(i, a)] * py[(i, 0)];
}
zpy[a] = acc;
}
// 5. Per θ_j: ∂Λ/∂θ_j is the indicator matrix E_j — ones at
// (f·q+r_j, f·q+c_j) for every primary level f, zero everywhere else
// (an extra grouping's θ_j puts ones on its own diagonal band
// instead). Then V_j = Z(E_jΛᵀ + ΛE_jᵀ)Zᵀ and
// score_j = tr(P V_j) − (n − p)·(pyᵀ V_j py) / ypy.
// By trace cyclicity tr(P V_j) = tr((Z'PZ)(E_jΛᵀ + ΛE_jᵀ)) — a k×k
// trace, not an n×n one — and pyᵀV_jpy = zpyᵀ(E_jΛᵀ + ΛE_jᵀ)zpy, a
// k-length quadratic form; `n` never reappears past `ztpz`/`zpy` above.
// Note `(n − p)/ypy` is `1/σ̂²`: the profiled scale, not a fixed 1.
let mut score = vec![0.0f64; n_theta];
for j in 0..n_theta {
let mut ej = Mat::<f64>::zeros(k, k);
if j < prim_theta {
let (r, c) = slot[j];
for f in 0..s {
ej[(f * q + r, f * q + c)] = 1.0;
}
} else {
for &(vs, n_levels, lam_off) in &extras {
if vs == j {
for l in 0..n_levels {
ej[(lam_off + l, lam_off + l)] = 1.0;
}
}
}
}
// mid = E_j Λᵀ + Λ E_jᵀ (k×k)
let mut mid = Mat::<f64>::zeros(k, k);
for a in 0..k {
for b in 0..k {
let mut acc = 0.0;
for c in 0..k {
acc += ej[(a, c)] * lam[(b, c)]; // E_j Λᵀ
acc += lam[(a, c)] * ej[(b, c)]; // Λ E_jᵀ
}
mid[(a, b)] = acc;
}
}
let mut tr = 0.0;
for a in 0..k {
for b in 0..k {
tr += ztpz[(a, b)] * mid[(b, a)];
}
}
let mut quad = 0.0;
for a in 0..k {
for b in 0..k {
quad += zpy[a] * mid[(a, b)] * zpy[b];
}
}
score[j] = tr - df * quad / ypy;
}
score
}
/// Ten fixed-seed θ draws per LMM shape for the REML gradient/Hessian gate.
/// Diagonal vech(Λ) lanes drawn positive (mirrors GLMM's own
/// `FixedSeedTheta`, `src/glmm/tests.rs`); off-diagonal lanes unconstrained —
/// `D = ΛΛ'` is PD for any off-diagonal value, only the diagonal needs to
/// stay positive.
struct FixedSeedThetaLmm {
state: u64,
n_theta: usize,
diag: Vec<usize>,
}
impl FixedSeedThetaLmm {
fn next_theta(&mut self) -> Vec<f64> {
let mut out = vec![0.0f64; self.n_theta];
for (j, out_j) in out.iter_mut().enumerate() {
let r = lcg(&mut self.state);
*out_j = if self.diag.contains(&j) {
0.3 + 0.4 * (r + 0.5) // in [0.3, 0.7]
} else {
0.5 * r // in [-0.25, 0.25]
};
}
out
}
}
/// θ-layout diagonal indices per gate shape — column-major vech, so `q2s`
/// (q=2) is `[0, 2]` and `q3s` (q=3) is `[0, 3, 5]`; `nest2`'s two scalars
/// (primary + one nested extra) are both diagonal. `cross3`/`cross5` have an
/// intercept-only (q=1) primary plus `n_extra` scalar crossed factors, so
/// every one of their `1 + n_extra` entries is diagonal. `cross6s`'s q=2
/// slope primary puts one off-diagonal entry (index 1) ahead of its six
/// scalar crossed factors.
fn fixed_seed_theta_lmm(shape: &str, state: u64) -> FixedSeedThetaLmm {
match shape {
"int1" => FixedSeedThetaLmm {
state,
n_theta: 1,
diag: vec![0],
},
"nest2" => FixedSeedThetaLmm {
state,
n_theta: 2,
diag: vec![0, 1],
},
"q2s" => FixedSeedThetaLmm {
state,
n_theta: 3,
diag: vec![0, 2],
},
"q3s" => FixedSeedThetaLmm {
state,
n_theta: 6,
diag: vec![0, 3, 5],
},
"cross3" => FixedSeedThetaLmm {
state,
n_theta: 4,
diag: vec![0, 1, 2, 3],
},
"cross5" => FixedSeedThetaLmm {
state,
n_theta: 6,
diag: (0..6).collect(),
},
// q_p=2 primary's vech is [λ00, λ10, λ11] (indices 0..2), so index 1
// is the only off-diagonal entry; indices 3..9 are the six crossed
// factors' scalar variances, all diagonal.
"cross6s" => FixedSeedThetaLmm {
state,
n_theta: 9,
diag: vec![0, 2, 3, 4, 5, 6, 7, 8],
},
other => panic!("unknown LMM gate shape {other}"),
}
}
/// `lmm_gate_fixture`'s `nest2` construction, generalized over row count and
/// primary cluster count (`n`, `nc` vary independently, `n` a multiple of
/// `nc` — child count is always `2*nc`). Preserves the original row→cluster
/// mapping exactly: the primary stays `Sizing::FixedSize { cluster_size: n /
/// nc }` (contiguous blocks of `n / nc` rows, `cluster_of_row(i) = i /
/// cluster_size` — `src/spec.rs`), NOT `FixedClusters` (`i % n_clusters`,
/// round-robin) — the two only happen to agree on `pid[i]` when `n / nc == 1`.
/// Same seed `5101`, same draw order as the original bespoke construction, so
/// `nest_dataset_sized(64, 8)` reproduces it row for row (`cluster_size =
/// 64/8 = 8`, matching the original `FixedSize { cluster_size: 8 }`).
fn nest_dataset_sized(n: usize, nc: usize) -> (Mat<f64>, Vec<f64>, Vec<u32>, Vec<u32>) {
assert!(
n % nc == 0,
"nest_dataset_sized: n={n} must be a multiple of nc={nc} for every cluster to keep the same size"
);
let mut cluster = intercept_only_spec(Sizing::FixedSize {
cluster_size: (n / nc) as u32,
});
cluster.re.as_mut().unwrap().extra_groupings.push(Grouping {
relation: GroupingRelation::NestedWithin { n_per_parent: 2 },
slopes: vec![],
});
let mut st = 5101u64;
let mut x = Mat::<f64>::zeros(n, 2);
let mut y = vec![0.0f64; n];
let mut pid = vec![0u32; n];
let mut cid = vec![0u32; n];
let u_p: Vec<f64> = (0..nc).map(|_| 0.5 * lcg(&mut st)).collect();
let u_c: Vec<f64> = (0..2 * nc).map(|_| 0.3 * lcg(&mut st)).collect();
for i in 0..n {
pid[i] = cluster.re.as_ref().unwrap().sizing.cluster_of_row(i) as u32;
cid[i] = extra_level_of_row(&cluster, 0, i) as u32;
let x1 = lcg(&mut st);
x[(i, 0)] = 1.0;
x[(i, 1)] = x1;
y[i] = 0.5 + 0.4 * x1 + u_p[pid[i] as usize] + u_c[cid[i] as usize] + 0.8 * lcg(&mut st);
}
(x, y, pid, cid)
}
/// One (x, y, primary ids, extra ids, groupings, primary-slope x-cols, p)
/// gate fixture per shape name. `int1` reuses `hand_dataset`; `q2s`/`q3s`
/// reuse `slope_dataset`/`slope_groupings` and `multislope_dataset`/
/// `multislope_groupings` (defined above in this module); `nest2` mirrors
/// `nested_regime_b_deviance_matches_brute_force`'s own dataset shape (same
/// family — 8 primary clusters × 2 nested children, n=64 — an independent
/// draw, own seed).
#[allow(clippy::type_complexity)]
fn lmm_gate_fixture(
shape: &str,
) -> (
Mat<f64>,
Vec<f64>,
Vec<u32>,
Vec<Vec<u32>>,
LmmGroupings,
Vec<usize>,
usize,
) {
match shape {
"int1" => {
let (x, y, ids) = hand_dataset();
(x, y, ids, vec![], LmmGroupings::single(6), vec![], 3)
}
"nest2" => {
let (x, y, pid, cid) = nest_dataset_sized(64, 8);
let mut cluster = intercept_only_spec(Sizing::FixedSize { cluster_size: 8 });
cluster.re.as_mut().unwrap().extra_groupings.push(Grouping {
relation: GroupingRelation::NestedWithin { n_per_parent: 2 },
slopes: vec![],
});
let groupings = LmmGroupings::from_cluster_spec(&cluster, 64, &[]);
(x, y, pid, vec![cid], groupings, vec![], 2)
}
"q2s" => {
let (x, y, ids) = slope_dataset();
(x, y, ids, vec![], slope_groupings(), vec![1], 2)
}
"q3s" => {
let (x, y, ids) = multislope_dataset();
(x, y, ids, vec![], multislope_groupings(), vec![1, 2], 3)
}
other => panic!("unknown LMM gate shape {other}"),
}
}
/// `lmm_gate_fixture`'s `int1`/`q2s`/`nest2` shapes, generalized over row
/// count, primary cluster count, and (where a slope primary does not already
/// rule it out) cluster balance — the dual-call cost table's size axis. Also
/// serves the new crossed shapes `cross3`/`cross5`/`cross6s` (see
/// `cross_dataset_sized`). `q3s` is left at its fixed size: it has a slope
/// primary, like `q2s`, and tells the balanced-collapse story nothing `q2s`
/// doesn't already. `q2s`/`nest2`/`cross6s` never take `skew = true` — `q2s`
/// and `cross6s` have a slope primary (see `slope_dataset_sized`'s doc
/// comment), and `nest2` has no skewed variant defined.
#[allow(clippy::type_complexity)]
fn lmm_gate_fixture_sized(
shape: &str,
n: usize,
nc: usize,
skew: bool,
) -> (
Mat<f64>,
Vec<f64>,
Vec<u32>,
Vec<Vec<u32>>,
LmmGroupings,
Vec<usize>,
usize,
) {
match shape {
"int1" => {
let (x, y, ids) = hand_dataset_sized(n, nc, skew);
(x, y, ids, vec![], LmmGroupings::single(nc), vec![], 3)
}
"q2s" => {
assert!(
!skew,
"q2s has a slope primary and never arms the balanced collapse; no skewed variant exists"
);
let (x, y, ids) = slope_dataset_sized(n, nc);
(x, y, ids, vec![], slope_groupings_sized(n, nc), vec![1], 2)
}
"nest2" => {
assert!(
!skew,
"nest2 has no skewed variant — see nest_dataset_sized"
);
let (x, y, pid, cid) = nest_dataset_sized(n, nc);
let mut cluster = intercept_only_spec(Sizing::FixedSize {
cluster_size: (n / nc) as u32,
});
cluster.re.as_mut().unwrap().extra_groupings.push(Grouping {
relation: GroupingRelation::NestedWithin { n_per_parent: 2 },
slopes: vec![],
});
let groupings = LmmGroupings::from_cluster_spec(&cluster, n, &[]);
(x, y, pid, vec![cid], groupings, vec![], 2)
}
"cross3" => {
let (x, y, pid, eids) = cross_dataset_sized(n, nc, 3, 30, false, skew);
let groupings = cross_groupings(nc, 3, 30, false, n);
(x, y, pid, eids, groupings, vec![], 2)
}
"cross5" => {
let (x, y, pid, eids) = cross_dataset_sized(n, nc, 5, 30, false, skew);
let groupings = cross_groupings(nc, 5, 30, false, n);
(x, y, pid, eids, groupings, vec![], 2)
}
"cross6s" => {
assert!(!skew, "cross6s has a slope primary; no skewed variant");
let (x, y, pid, eids) = cross_dataset_sized(n, nc, 6, 30, true, false);
let groupings = cross_groupings(nc, 6, 30, true, n);
(x, y, pid, eids, groupings, vec![1], 2)
}
other => panic!("lmm_gate_fixture_sized: unsupported shape {other}"),
}
}
/// Gate (a): `reml_gradient` vs. `dense_reml_score`, ten fixed-seed θ draws
/// per shape (`{int1, nest2, q2s, q3s}`), band `1e-8` relative. The LMM
/// criterion has no PIRLS inside it — it is a closed-form evaluation of the
/// suff stats — so there is no mode error and no tolerance to tighten: this
/// band is tighter than the GLMM gates on purpose, and a miss here is a real
/// bug rather than a convergence artefact.
#[test]
fn reml_gradient_matches_dense_score_per_shape() {
for shape in ["int1", "nest2", "q2s", "q3s"] {
let (x, y, primary_ids, extra_ids, groupings, slope_cols, p) = lmm_gate_fixture(shape);
let mut suff = LmmSuffStats::with_groupings(p, groupings.clone());
suff.add_rows_multi(x.as_ref(), &y, &primary_ids, &extra_ids, None);
let z = dense_z(&groupings, &primary_ids, &extra_ids, &x, &slope_cols);
let n_theta = groupings.n_theta();
let mut scratch =
LmmDualScratch::for_groupings(n_theta, p, &groupings).unwrap_or_else(|| {
panic!("{shape}: n_theta {n_theta} exceeds the instantiated lane set")
});
if shape == "q3s" {
assert!(
matches!(scratch, LmmDualScratch::D6(_)),
"n_theta 6 must take the D6 rung, not D8"
);
}
let mut rng = fixed_seed_theta_lmm(shape, 6001);
debug_assert_eq!(
rng.n_theta, n_theta,
"{shape}: seed rng/groupings n_theta mismatch"
);
for draw in 0..10 {
let theta = rng.next_theta();
let mut grad = vec![0.0; n_theta];
match reml_gradient(&theta, &suff, &mut scratch, &mut grad) {
DerivStatus::Ok(_) => {}
DerivStatus::NotConverged => {
panic!("{shape} draw {draw} θ={theta:?}: reml_gradient NotConverged")
}
DerivStatus::Unsupported => {
panic!("{shape} draw {draw} θ={theta:?}: reml_gradient Unsupported")
}
}
let want = dense_reml_score(&theta, z.as_ref(), x.as_ref(), &y, &groupings);
for j in 0..n_theta {
let band = 1e-8 * want[j].abs().max(1.0);
assert!(
(grad[j] - want[j]).abs() <= band,
"{shape} draw {draw} θ={theta:?} coord {j}: dual {} vs dense {}",
grad[j],
want[j]
);
}
}
}
}
/// A chunked gradient is bit-for-bit a single wide pass. `q3s` has n_θ = 6, so
/// a `Dual<6>` scratch seeds it in ONE pass while a `Dual<4>` scratch seeds it
/// in two (coordinates 0..4, then 4..6). Forward-mode lanes are independent, so
/// every entry must be EXACTLY equal — not close. This is the property the
/// chunked gradient rests on; a difference means some operation in `dual.rs`
/// mixes lanes.
#[test]
fn chunked_gradient_is_bit_identical_to_one_pass() {
let (x, y, primary_ids, extra_ids, groupings, _slope_cols, p) = lmm_gate_fixture("q3s");
let mut suff = LmmSuffStats::with_groupings(p, groupings.clone());
suff.add_rows_multi(x.as_ref(), &y, &primary_ids, &extra_ids, None);
let n_theta = groupings.n_theta();
assert_eq!(n_theta, 6);
let mut one = LmmDualScratch::D6(LmmFitScratch::with_groupings(p, &groupings));
let mut two = LmmDualScratch::D4(LmmFitScratch::with_groupings(p, &groupings));
let mut rng = fixed_seed_theta_lmm("q3s", 6001);
for draw in 0..10 {
let theta = rng.next_theta();
let mut g1 = vec![0.0; n_theta];
let mut g2 = vec![0.0; n_theta];
assert!(matches!(
reml_gradient(&theta, &suff, &mut one, &mut g1),
DerivStatus::Ok(_)
));
assert!(matches!(
reml_gradient(&theta, &suff, &mut two, &mut g2),
DerivStatus::Ok(_)
));
assert_eq!(g1, g2, "draw {draw} θ={theta:?}: one pass vs two chunks");
}
}
/// The chunked gradient reaches a shape no single instantiated rung covers:
/// a `q_p = 8` primary has n_θ = 36, three `Dual<12>` passes. Checked against
/// `dense_reml_score`, the same independent textbook criterion
/// `reml_gradient_matches_dense_score_per_shape` uses, at the same `1e-8`
/// relative band — the REML criterion has no PIRLS inside it, so there is no
/// mode error to loosen the band for.
#[test]
fn chunked_gradient_reaches_n_theta_36() {
let (nc, n, p) = (30usize, 300usize, 8usize);
let slope_cols: Vec<usize> = (1..8).collect();
let mut st = 8801u64;
let u: Vec<Vec<f64>> = (0..8)
.map(|_| (0..nc).map(|_| 0.4 * lcg(&mut st)).collect())
.collect();
let mut x = Mat::<f64>::zeros(n, p);
let mut y = vec![0.0f64; n];
let mut ids = vec![0u32; n];
for i in 0..n {
let c = i % nc;
ids[i] = c as u32;
x[(i, 0)] = 1.0;
let mut yi = 0.5 + u[0][c] + 0.8 * lcg(&mut st);
for d in 1..p {
let xd = lcg(&mut st);
x[(i, d)] = xd;
yi += 0.2 * xd + u[d][c] * xd;
}
y[i] = yi;
}
let cluster = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters {
n_clusters: nc as u32,
},
slopes: (0..7).collect(),
extra_groupings: vec![],
}),
};
let groupings = LmmGroupings::from_cluster_spec(&cluster, n, &slope_cols);
let n_theta = groupings.n_theta();
assert_eq!(n_theta, 36, "q_p = 8 must give vech length 36");
let mut suff = LmmSuffStats::with_groupings(p, groupings.clone());
suff.add_rows_multi(x.as_ref(), &y, &ids, &[], None);
let z = dense_z(&groupings, &ids, &[], &x, &slope_cols);
let mut scratch = LmmDualScratch::for_groupings(n_theta, p, &groupings)
.expect("the gradient ladder covers every n_theta");
assert!(
matches!(scratch, LmmDualScratch::D12(_)),
"top rung expected"
);
// Diagonal-positive θ; the diagonal indices come from the groupings' own
// cached map rather than a hand-written vech offset list.
let mut theta = vec![0.15_f64; n_theta];
for &d in groupings.diagonal_theta() {
theta[d] = 0.6;
}
let mut grad = vec![0.0; n_theta];
assert!(matches!(
reml_gradient(&theta, &suff, &mut scratch, &mut grad),
DerivStatus::Ok(_)
));
let want = dense_reml_score(&theta, z.as_ref(), x.as_ref(), &y, &groupings);
for j in 0..n_theta {
let band = 1e-8 * want[j].abs().max(1.0);
assert!(
(grad[j] - want[j]).abs() <= band,
"coord {j}: chunked {} vs dense {}",
grad[j],
want[j]
);
}
}
/// Gate (b): `reml_hessian` vs. central FD of `reml_gradient` (step `1e-5`,
/// band `1e-6` relative), same shapes, plus exact `hess[(i,j)] ==
/// hess[(j,i)]` symmetry. `q3s` reaches `n_θ = 6`, so `N = 6` is the largest
/// lane count either gate exercises.
#[test]
fn reml_hessian_matches_fd_of_gradient_per_shape() {
for shape in ["int1", "nest2", "q2s", "q3s"] {
let (x, y, primary_ids, extra_ids, groupings, _slope_cols, p) = lmm_gate_fixture(shape);
let mut suff = LmmSuffStats::with_groupings(p, groupings.clone());
suff.add_rows_multi(x.as_ref(), &y, &primary_ids, &extra_ids, None);
let n_theta = groupings.n_theta();
let mut dual_scratch = LmmDualScratch::for_groupings(n_theta, p, &groupings)
.unwrap_or_else(|| {
panic!("{shape}: n_theta {n_theta} exceeds the instantiated lane set")
});
let mut hyper_scratch = LmmHyperScratch::for_groupings(n_theta, p, &groupings)
.unwrap_or_else(|| {
panic!("{shape}: n_theta {n_theta} exceeds the instantiated lane set")
});
let mut rng = fixed_seed_theta_lmm(shape, 6501);
for draw in 0..10 {
let theta = rng.next_theta();
let mut grad = vec![0.0; n_theta];
let mut hess = Mat::<f64>::zeros(n_theta, n_theta);
match reml_hessian(&theta, &suff, &mut hyper_scratch, &mut grad, &mut hess) {
DerivStatus::Ok(_) => {}
DerivStatus::NotConverged => {
panic!("{shape} draw {draw} θ={theta:?}: reml_hessian NotConverged")
}
DerivStatus::Unsupported => {
panic!("{shape} draw {draw} θ={theta:?}: reml_hessian Unsupported")
}
}
for i in 0..n_theta {
for j in 0..n_theta {
assert_eq!(
hess[(i, j)],
hess[(j, i)],
"{shape} draw {draw} θ={theta:?}: hess not symmetric at ({i},{j})"
);
}
}
let h = 1e-5;
for kcoord in 0..n_theta {
let mut tp = theta.clone();
tp[kcoord] += h;
let mut gp = vec![0.0; n_theta];
match reml_gradient(&tp, &suff, &mut dual_scratch, &mut gp) {
DerivStatus::Ok(_) => {}
DerivStatus::NotConverged => {
panic!(
"{shape} draw {draw} coord {kcoord} (+h): reml_gradient NotConverged"
)
}
DerivStatus::Unsupported => {
panic!("{shape} draw {draw} coord {kcoord} (+h): reml_gradient Unsupported")
}
}
let mut tm = theta.clone();
tm[kcoord] -= h;
let mut gm = vec![0.0; n_theta];
match reml_gradient(&tm, &suff, &mut dual_scratch, &mut gm) {
DerivStatus::Ok(_) => {}
DerivStatus::NotConverged => {
panic!(
"{shape} draw {draw} coord {kcoord} (-h): reml_gradient NotConverged"
)
}
DerivStatus::Unsupported => {
panic!("{shape} draw {draw} coord {kcoord} (-h): reml_gradient Unsupported")
}
}
for j in 0..n_theta {
let fd = (gp[j] - gm[j]) / (2.0 * h);
let band = 1e-6 * fd.abs().max(1.0);
assert!(
(hess[(j, kcoord)] - fd).abs() <= band,
"{shape} draw {draw} θ={theta:?} hess[{j},{kcoord}]: {} vs fd {fd}",
hess[(j, kcoord)]
);
}
}
}
}
}
/// LMM twin of `glmm::tests::glmm_dual_call_cost_table` — per-call wall cost
/// of the plain `f64` REML deviance vs. its `Dual<N>` gradient vs. its
/// `HyperDual<N,H>` gradient+Hessian, at fixed θ, across `n_theta` rungs
/// **and** a row-count axis. An instrument, not a gate.
///
/// Rung coverage: unlike the GLMM side, the lane pick here is on `n_theta`
/// alone (`run_reml_hessian` refuses `n_theta > N`; `run_reml_gradient` has
/// no refusal — above the top rung it runs `⌈n_theta / 12⌉` passes of
/// `Dual<12>` instead), not `n_theta + p`, so padding `p` cannot move the
/// rung — the shape itself
/// must carry more θ. `int1`/`nest2`/`q2s`/`q3s` give n_θ = 1, 2, 3, 6 — D4
/// through D6. `cross3`/`cross5` add n_θ = 4, 6 (D4, D6); `cross6s`
/// (n_θ = 9, a slope primary plus six crossed factors) is the first fixture
/// in this file to reach the D12 rung. `reml_gradient`/`reml_hessian` also
/// refuse `extra_slopes_any` shapes (`kernel.rs`), but none of these are —
/// `nest2`'s and the `cross*` shapes' extra groupings carry no slope, so they
/// are unaffected. `n_theta` is structural, like GLMM's `m`, so it does not
/// move with row count — every size tier below repeats the same rung.
///
/// Size and balance coverage: a profile on the speed grid found the
/// mechanism this table exists to isolate — on a **balanced** design the
/// `f64` REML objective takes `precompute_balanced_collapse`'s fast path
/// (`kernel.rs:370`, gated on `primary_q == 1`: a plain-intercept primary),
/// skipping the whole per-family loop and the syrk that `f64` would
/// otherwise share with `Dual`/`HyperDual`; both the dual and hyper scalars
/// arm the same collapse (`reml_gradient`/`reml_hessian` call
/// `precompute_balanced_collapse` on every call), so the balanced/skewed
/// split isolates the collapse's own saving on both sides. A table with
/// only balanced rows cannot show this, because it never isolates the
/// collapse's own saving from the dual/hyper overhead. So every q_p=1
/// primary shape (`int1`, `nest2`, `cross3`, `cross5`) is generalized over
/// both size and a `balance` column: `bal` (unchanged `i % nc` assignment)
/// vs. `skew` (~20% of clusters draw ~80% of rows, mirroring
/// `campaigns/speed-grid/prep.R`'s own skew recipe). The `armed` column
/// prints whether `precompute_balanced_collapse` actually took the fast path
/// on that row's `f64` fit, so a row cannot silently claim a collapse it did
/// not take. The predicate is the same for both scalars — it reads `suff`
/// alone, not the scratch type — so one `armed` column speaks for all three.
/// `q2s`/`q3s`/`cross6s` (slope primaries) never arm the collapse
/// regardless of balance, so they get the size axis (where generalized)
/// without a skew variant — see `slope_dataset_sized`'s doc comment. `q3s`
/// keeps only its original fixed-size row: like `q2s`, it has a slope
/// primary and so would add no new information. New size tiers use
/// `nc = n / 20`: `prep.R`'s `per_group = 20` is the one rows-per-cluster
/// value its own generator uses at every one of its three size magnitudes,
/// so it is the ratio mirrored here rather than invented.
///
/// Timing method: adaptive repetition count — see
/// `glmm::tests::glmm_dual_call_cost_table`'s doc comment for the full
/// rationale (fixed `R = 20` is both too slow for a 30,000-row `HyperDual`
/// cell and too few reps for the small fixtures' noise floor). One probe
/// call sizes the rest to `BUDGET_NS`, floored at `REP_FLOOR` total calls;
/// the count actually used is printed per quantity.
///
/// Run: `cargo test -p glmm --release lmm_dual_call_cost_table -- --ignored --nocapture`
#[test]
#[ignore]
fn lmm_dual_call_cost_table() {
// Serialized under alloc-tests so its allocations can't land in a
// concurrent dhat profiler window on an `-- --ignored` run.
#[cfg(feature = "alloc-tests")]
let _serial = crate::test_support::alloc_test_guard();
use std::time::Instant;
const BUDGET_NS: u128 = 3_000_000_000;
const REP_FLOOR: usize = 3;
const REP_CAP: usize = 200_000; // bounds the loop for cells so cheap the budget implies absurd rep counts
/// Returns `(min_ns_over_the_timed_calls, reps_used, impractical)`.
fn adaptive_call_ns(mut call: impl FnMut()) -> (u128, usize, bool) {
let t0 = Instant::now();
call();
let probe = t0.elapsed().as_nanos().max(1);
if probe > BUDGET_NS {
return (probe, 1, true);
}
let reps = ((BUDGET_NS / probe).max(REP_FLOOR as u128) as usize).min(REP_CAP);
let mut best = u128::MAX;
for _ in 0..reps - 1 {
let t = Instant::now();
call();
best = best.min(t.elapsed().as_nanos());
}
(best, reps, false)
}
// Slope-primary shapes (q2s, q3s) never arm the collapse (`primary_q ==
// 1` guard), so their balance is reported as "n/a" rather than "bal" —
// the label would otherwise imply a comparison this shape can't make.
fn balance_label(shape: &str, skew: bool) -> &'static str {
match shape {
"int1" | "nest2" | "cross3" | "cross5" => {
if skew {
"skew"
} else {
"bal"
}
}
_ => "n/a",
}
}
println!(
"engine\tfixture\tbalance\tarmed\tn_rows\tn_theta\tp\tm\trung\tf64_reps\tdual_reps\thyper_reps\tf64_ns\tdual_ns\thyper_ns\tdual_over_f64\thyper_over_f64"
);
// `None` = the original fixed fixture (unchanged, via `lmm_gate_fixture`);
// `Some((n, nc, skew))` = a size-axis tier via `lmm_gate_fixture_sized`.
type SizeTier = (usize, usize, bool);
let cells: &[(&str, Option<SizeTier>)] = &[
("int1", None),
("int1", Some((3000, 150, false))),
("int1", Some((3000, 150, true))),
("int1", Some((30_000, 1500, false))),
("int1", Some((30_000, 1500, true))),
("nest2", None),
("q2s", None),
("q2s", Some((3000, 150, false))),
("q2s", Some((30_000, 1500, false))),
("q3s", None),
("nest2", Some((3000, 150, false))),
("nest2", Some((30_000, 1500, false))),
("cross3", Some((3000, 600, false))),
("cross3", Some((3000, 600, true))),
("cross5", Some((30_000, 6000, false))),
("cross5", Some((30_000, 6000, true))),
("cross6s", Some((3000, 600, false))),
];
for &(shape, tier) in cells {
let (x, y, primary_ids, extra_ids, groupings, _slope_cols, p) = match tier {
None => lmm_gate_fixture(shape),
Some((n, nc, skew)) => lmm_gate_fixture_sized(shape, n, nc, skew),
};
let skew = tier.map(|(_, _, s)| s).unwrap_or(false);
let balance = balance_label(shape, skew);
let n_rows = x.nrows();
let mut suff = LmmSuffStats::with_groupings(p, groupings.clone());
suff.add_rows_multi(x.as_ref(), &y, &primary_ids, &extra_ids, None);
let n_theta = groupings.n_theta();
let mut rng = fixed_seed_theta_lmm(shape, 9001);
let theta = rng.next_theta();
let mut fit_f64 = LmmFitScratch::<f64>::with_groupings(p, &groupings);
// The production `f64` path arms the collapse once per fit
// (`fit_lmm_impl`), so a table that times `reml_deviance::<f64>` on a
// bare scratch times the fallback loop and the balanced/skewed rows say
// nothing. `armed` is printed once so a row cannot silently claim a
// collapse it did not take.
let armed = precompute_balanced_collapse(&suff, &mut fit_f64);
// `reml_gradient`/`reml_hessian` call `precompute_balanced_collapse`
// on every call (production arms once per gradient call), so the
// `f64` timed closure re-arms it here too — otherwise this side pays
// zero precomputes per timed call while the dual/hyper side pays one,
// and `dual_over_f64`/`hyper_over_f64` would overstate the dual cost
// on exactly the balanced high-`t_dim` cells this table exists to
// measure.
let (f64_ns, f64_reps, f64_impractical) = adaptive_call_ns(|| {
precompute_balanced_collapse(&suff, &mut fit_f64);
let _ = reml_deviance(&theta, &suff, &mut fit_f64);
});
let dual_scratch = LmmDualScratch::for_groupings(n_theta, p, &groupings);
let hyper_scratch = LmmHyperScratch::for_groupings(n_theta, p, &groupings);
let (dual_ns_str, dual_over, dual_reps, dual_impractical) = match dual_scratch {
Some(mut scratch) => {
let mut grad = vec![0.0; n_theta];
let mut status = false;
let (dual_ns, dual_reps, dual_impractical) = adaptive_call_ns(|| {
status = matches!(
reml_gradient(&theta, &suff, &mut scratch, &mut grad),
DerivStatus::Ok(_)
);
});
if status {
(
dual_ns.to_string(),
format!("{:.2}", dual_ns as f64 / f64_ns as f64),
dual_reps,
dual_impractical,
)
} else {
(
"NA".to_string(),
"NA".to_string(),
dual_reps,
dual_impractical,
)
}
}
None => ("NA".to_string(), "NA".to_string(), 0, false),
};
let (hyper_ns_str, hyper_over, hyper_reps, hyper_impractical) = match hyper_scratch {
Some(mut scratch) => {
let mut grad = vec![0.0; n_theta];
let mut hess = Mat::<f64>::zeros(n_theta, n_theta);
let mut status = false;
let (hyper_ns, hyper_reps, hyper_impractical) = adaptive_call_ns(|| {
status = matches!(
reml_hessian(&theta, &suff, &mut scratch, &mut grad, &mut hess),
DerivStatus::Ok(_)
);
});
if status {
(
hyper_ns.to_string(),
format!("{:.2}", hyper_ns as f64 / f64_ns as f64),
hyper_reps,
hyper_impractical,
)
} else {
(
"NA".to_string(),
"NA".to_string(),
hyper_reps,
hyper_impractical,
)
}
}
None => ("NA".to_string(), "NA".to_string(), 0, false),
};
let rung = match n_theta {
0..=4 => "D4",
5 => "D5",
6 => "D6",
7..=8 => "D8",
9..=12 => "D12",
_ => "NA",
};
let armed_str = if armed { "yes" } else { "no" };
println!(
"lmm\t{shape}\t{balance}\t{armed_str}\t{n_rows}\t{n_theta}\t{p}\t{n_theta}\t{rung}\t{f64_reps}\t{dual_reps}\t{hyper_reps}\t{f64_ns}\t{dual_ns_str}\t{hyper_ns_str}\t{dual_over}\t{hyper_over}"
);
if f64_impractical {
println!("# note: {shape}/{balance}/n={n_rows} f64 probe alone exceeded the {BUDGET_NS}ns budget; reporting a single call, not a minimum");
}
if dual_impractical {
println!("# note: {shape}/{balance}/n={n_rows} dual probe alone exceeded the {BUDGET_NS}ns budget; reporting a single call, not a minimum");
}
if hyper_impractical {
println!("# note: {shape}/{balance}/n={n_rows} hyper probe alone exceeded the {BUDGET_NS}ns budget; reporting a single call, not a minimum");
}
}
}