use super::{build_workspace, fit_on};
use crate::fit::{spec_sized_from_ids_pub, Perm};
use crate::test_support::assert_near;
use crate::{
fit_cold, BinomialLink, Family, FitOptions, GroupIds, Grouping, GroupingRelation, ModelSpec,
ReStructure, Sizing,
};
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
}
fn ols_case() -> (
Vec<f64>,
Vec<f64>,
usize,
usize,
ModelSpec,
GroupIds,
FitOptions,
) {
let n = 30;
let p = 3;
let mut x = Vec::with_capacity(n * p);
let mut y = Vec::with_capacity(n);
for i in 0..n {
let a = i as f64;
let b = ((i * 5) % 7) as f64 - 3.0;
x.extend_from_slice(&[1.0, a, b]);
y.push(0.7 + 1.1 * a - 0.5 * b + ((i % 3) as f64 - 1.0));
}
let model = ModelSpec {
family: Family::Gaussian,
re: None,
};
let opts = FitOptions {
target_indices: vec![1, 2],
..FitOptions::default()
};
(x, y, n, p, model, GroupIds::default(), opts)
}
fn glm_case() -> (
Vec<f64>,
Vec<f64>,
usize,
usize,
ModelSpec,
GroupIds,
FitOptions,
) {
let n = 30;
let p = 2;
let mut st = 7u64;
let mut x = Vec::with_capacity(n * p);
let mut y = Vec::with_capacity(n);
for _ in 0..n {
let x1 = 0.4 * lcg(&mut st);
x.extend_from_slice(&[1.0, x1]);
let eta: f64 = 0.5 + 0.6 * x1;
y.push(eta.exp().round());
}
let model = ModelSpec {
family: Family::Poisson {
link: crate::PoissonLink::Log,
},
re: None,
};
let opts = FitOptions {
target_indices: vec![0, 1],
..FitOptions::default()
};
(x, y, n, p, model, GroupIds::default(), opts)
}
fn lmm_intercept_case() -> (
Vec<f64>,
Vec<f64>,
usize,
usize,
ModelSpec,
GroupIds,
FitOptions,
) {
let n_clusters = 6usize;
let per = 8usize;
let n = n_clusters * per;
let p = 2usize;
let mut st = 13u64;
let mut x = vec![0.0f64; n * p];
let mut y = vec![0.0f64; n];
let mut ids_v = vec![0u32; n];
for i in 0..n {
ids_v[i] = (i % n_clusters) as u32;
st = st
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
let x1 = ((st >> 33) as f64 / (1u64 << 31) as f64) - 1.0;
x[i * 2] = 1.0;
x[i * 2 + 1] = x1;
let re = 0.3 * ((ids_v[i] as f64) - (n_clusters as f64) / 2.0);
y[i] = 0.5 + 0.4 * x1 + re + 0.1 * ((i % 5) as f64 - 2.0);
}
let model = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 1 },
slopes: vec![],
extra_groupings: vec![],
}),
};
let ids = GroupIds {
primary: ids_v,
extra: vec![],
};
let opts = FitOptions {
target_indices: vec![0, 1],
..FitOptions::default()
};
(x, y, n, p, model, ids, opts)
}
fn glmm_binomial_intercept_case() -> (
Vec<f64>,
Vec<f64>,
usize,
usize,
ModelSpec,
GroupIds,
FitOptions,
) {
let n_clusters = 8usize;
let per = 16usize;
let n = n_clusters * per;
let p = 2usize;
let mut x = vec![0.0f64; n * p];
let mut y = vec![0.0f64; n];
let mut ids_v = vec![0u32; n];
for i in 0..n {
let c = i / per;
ids_v[i] = c as u32;
let x1 = ((i % per) as f64) / (per as f64) - 0.5;
x[i * 2] = 1.0;
x[i * 2 + 1] = x1;
let eta = -0.2 + 0.9 * x1 + 0.4 * (c as f64 - 3.5);
y[i] = if (eta + ((i * 7 % 11) as f64 - 5.0) * 0.15) > 0.0 {
1.0
} else {
0.0
};
}
let model = ModelSpec {
family: Family::Binomial {
link: BinomialLink::Logit,
},
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 1 },
slopes: vec![],
extra_groupings: vec![],
}),
};
let ids = GroupIds {
primary: ids_v,
extra: vec![],
};
let opts = FitOptions {
target_indices: vec![1],
..FitOptions::default()
};
(x, y, n, p, model, ids, opts)
}
fn crossed_extra_glmm_case() -> (
Vec<f64>,
Vec<f64>,
usize,
usize,
ModelSpec,
GroupIds,
FitOptions,
) {
let g1 = 6usize; let g2 = 4usize; let n = 96usize;
let p = 2usize;
let mut x = vec![0.0f64; n * p];
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 c1 = i % g1;
let c2 = (i / g1) % g2;
pid[i] = c1 as u32;
eid[i] = c2 as u32;
let x1 = ((i % 8) as f64) / 8.0 - 0.5;
x[i * 2] = 1.0;
x[i * 2 + 1] = x1;
let eta = -0.1 + 0.8 * x1 + 0.3 * (c1 as f64 - 2.5) + 0.25 * (c2 as f64 - 1.5);
y[i] = if (eta + ((i * 7 % 11) as f64 - 5.0) * 0.15) > 0.0 {
1.0
} else {
0.0
};
}
let model = ModelSpec {
family: Family::Binomial {
link: BinomialLink::Logit,
},
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 1 },
slopes: vec![],
extra_groupings: vec![Grouping {
relation: GroupingRelation::Crossed { n_clusters: 1 },
slopes: vec![],
}],
}),
};
let ids = GroupIds {
primary: pid,
extra: vec![eid],
};
let opts = FitOptions {
target_indices: vec![1],
..FitOptions::default()
};
(x, y, n, p, model, ids, opts)
}
fn nested_unbalanced_case() -> (
Vec<f64>,
Vec<f64>,
usize,
usize,
ModelSpec,
GroupIds,
FitOptions,
) {
let children_per_parent = [3usize, 2, 2, 2];
let rows_per_child = 4usize;
let p = 2usize;
let mut st = 7u64;
let mut x = Vec::new();
let mut y = Vec::new();
let mut pid = Vec::new();
let mut eid = Vec::new();
for (parent, &kids) in children_per_parent.iter().enumerate() {
let u_p = 0.5 * lcg(&mut st);
for child in 0..kids {
let u_c = 0.3 * lcg(&mut st);
for _ in 0..rows_per_child {
let x1 = lcg(&mut st);
x.extend_from_slice(&[1.0, x1]);
y.push(0.4 + 0.6 * x1 + u_p + u_c + 0.2 * lcg(&mut st));
pid.push(parent as u32);
eid.push((3 * parent + child) as u32);
}
}
}
let n = y.len();
let model = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 1 },
slopes: vec![],
extra_groupings: vec![Grouping {
relation: GroupingRelation::NestedWithin { n_per_parent: 1 },
slopes: vec![],
}],
}),
};
let ids = GroupIds {
primary: pid,
extra: vec![eid],
};
let opts = FitOptions {
target_indices: vec![0, 1],
..FitOptions::default()
};
(x, y, n, p, model, ids, opts)
}
#[test]
fn fitview_accessors_match_fit_for_ols() {
let (x, y, n, p, model, ids, opts) = ols_case();
let cold = fit_cold(&x, &y, n, p, &model, &ids, &opts);
let mut ws = build_workspace(&model, Perm::IDENTITY, n, p, &opts);
let v = fit_on(&mut ws, &x, &y, &ids, None, &opts);
assert_eq!(v.converged(), cold.converged());
assert_eq!(v.t_sq().len(), opts.target_indices.len());
assert_eq!(v.betas().len(), p);
}
#[test]
fn fitview_diagnostics_agree_with_materialized_fit() {
let (x, y, n, p, model, ids, opts) = ols_case();
let mut ws = build_workspace(&model, Perm::IDENTITY, n, p, &opts);
let d = fit_on(&mut ws, &x, &y, &ids, None, &opts).diagnostics();
let cold = fit_cold(&x, &y, n, p, &model, &ids, &opts);
assert_eq!(d.converged, cold.converged());
assert_eq!(d.boundary_hit, 0, "OLS has no θ to pin");
assert!(!cold.singular(), "OLS reports no variance component");
assert!(!d.ill_conditioned);
assert!(d.pivot > crate::ols::PIVOT_MIN, "pivot {}", d.pivot);
let (x, y, n, p, model, ids, opts) = lmm_intercept_case();
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let mut ws = build_workspace(&sized, perm, n, p, &opts);
let d = fit_on(&mut ws, &x, &y, &ids, None, &opts).diagnostics();
let cold = fit_cold(&x, &y, n, p, &model, &ids, &opts);
assert_eq!(d.converged, cold.converged());
assert!(!d.ill_conditioned);
assert!(d.pivot > crate::lmm::PIVOT_MIN, "pivot {}", d.pivot);
assert!(d.boundary_hit != 1 || cold.singular());
let (x, y, n, p, model, ids, opts) = crossed_extra_case(vec![1]);
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let mut ws = build_workspace(&sized, perm, n, p, &opts);
assert!(ws.is_prebuilt());
let d = fit_on(&mut ws, &x, &y, &ids, None, &opts).diagnostics();
let cold = fit_cold(&x, &y, n, p, &model, &ids, &opts);
assert_eq!(d.converged, cold.converged());
assert_eq!(d.boundary_hit == 1, cold.singular());
assert!(!d.ill_conditioned && d.pivot.is_nan());
}
#[test]
fn fitview_diagnostics_flag_a_rank_deficient_lmm_draw() {
const D: f64 = 1e-7;
let (n, n_clusters, p) = (48usize, 6usize, 3usize);
let mut st = 11u64;
let mut x = vec![0.0f64; n * p];
let mut y = vec![0.0f64; n];
let mut ids_v = vec![0u32; n];
for i in 0..n {
ids_v[i] = (i % n_clusters) as u32;
let x1 = lcg(&mut st);
x[i * p] = 1.0;
x[i * p + 1] = x1;
x[i * p + 2] = 0.1 * x1 * (1.0 + D * if i % 2 == 0 { 1.0 } else { -1.0 });
y[i] = 0.5 + 0.4 * x1 + 0.8 * lcg(&mut st);
}
let model = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters {
n_clusters: n_clusters as u32,
},
slopes: vec![],
extra_groupings: vec![],
}),
};
let ids = GroupIds {
primary: ids_v,
extra: vec![],
};
let opts = FitOptions {
target_indices: vec![1, 2],
..FitOptions::default()
};
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let mut ws = build_workspace(&sized, perm, n, p, &opts);
assert!(ws.is_lmm_dense());
let d = fit_on(&mut ws, &x, &y, &ids, None, &opts).diagnostics();
assert!(d.converged, "the near-duplicate design is still computable");
assert!(
d.ill_conditioned,
"a near-duplicate column must raise the flag, pivot was {}",
d.pivot
);
assert!(d.pivot < crate::lmm::PIVOT_MIN, "pivot {}", d.pivot);
assert_eq!(
d.pivot_col, 2,
"the LATER column of the duplicated pair is the one named"
);
let cold = fit_cold(&x, &y, n, p, &model, &ids, &opts);
assert_eq!(cold.aliased(), vec![false, false, true]);
}
#[test]
fn fitview_diagnostics_flag_a_weighted_collinear_ols_fit() {
let (n, p, split) = (60usize, 3usize, 40usize);
const WSMALL: f64 = 1e-11;
let mut x = Vec::with_capacity(n * p);
let mut y = Vec::with_capacity(n);
let mut w = Vec::with_capacity(n);
for i in 0..n {
let a = ((i * 13) % 17) as f64 - 8.0;
let delta = if i < split { 0.0 } else { 1.0 };
x.extend_from_slice(&[1.0, a, a + delta]);
y.push(0.5 + 1.3 * a + 0.477 * (a + delta) + ((i % 3) as f64 - 1.0));
w.push(if i < split { 1.0 } else { WSMALL });
}
let model = ModelSpec {
family: Family::Gaussian,
re: None,
};
let opts = FitOptions {
target_indices: vec![0, 1, 2],
weights: Some(w),
..FitOptions::default()
};
let ids = GroupIds::default();
let mut ws = build_workspace(&model, Perm::IDENTITY, n, p, &opts);
assert!(ws.is_ols());
let d = fit_on(&mut ws, &x, &y, &ids, None, &opts).diagnostics();
assert!(d.converged, "the fit is computable and must be returned");
assert!(
d.ill_conditioned,
"the weighted collinearity must raise the flag, pivot was {}",
d.pivot
);
assert!(d.pivot < crate::ols::PIVOT_MIN, "pivot {}", d.pivot);
assert_eq!(d.pivot_col, 2);
let cold = fit_cold(&x, &y, n, p, &model, &ids, &opts);
assert!(cold.converged() && cold.aliased() == vec![false; p]);
}
#[cfg(feature = "alloc-tests")]
#[test]
#[ignore]
fn fitview_diagnostics_zero_alloc() {
let _serial = crate::test_support::alloc_test_guard();
const N_CALLS: usize = 1000;
let (xo, yo, no, po, mo, ido, oo) = ols_case();
let mut ws_ols = build_workspace(&mo, Perm::IDENTITY, no, po, &oo);
let (xl, yl, nl, pl, ml, idl, ol) = lmm_intercept_case();
let (sized_l, idl, perm) = spec_sized_from_ids_pub(&ml, &idl);
let mut ws_lmm = build_workspace(&sized_l, perm, nl, pl, &ol);
let (xs, ys, ns, ps, ms, ids_s, os) = crossed_extra_case(vec![1]);
let (sized_s, ids_s, perm) = spec_sized_from_ids_pub(&ms, &ids_s);
let mut ws_sparse = build_workspace(&sized_s, perm, ns, ps, &os);
let v_ols = fit_on(&mut ws_ols, &xo, &yo, &ido, None, &oo);
let v_lmm = fit_on(&mut ws_lmm, &xl, &yl, &idl, None, &ol);
let v_sparse = fit_on(&mut ws_sparse, &xs, &ys, &ids_s, None, &os);
let profiler = dhat::Profiler::builder().testing().build();
for _ in 0..N_CALLS {
std::hint::black_box(std::hint::black_box(&v_ols).diagnostics());
std::hint::black_box(std::hint::black_box(&v_lmm).diagnostics());
std::hint::black_box(std::hint::black_box(&v_sparse).diagnostics());
}
let stats = dhat::HeapStats::get();
drop(profiler);
assert_eq!(
stats.total_blocks, 0,
"FitView::diagnostics allocated {} blocks across {} reads per arm",
stats.total_blocks, N_CALLS
);
}
#[test]
fn build_workspace_routes_fixed_only_to_ols_and_mixed_to_lmm() {
let opts = FitOptions::default();
let fixed = ModelSpec {
family: Family::Gaussian,
re: None,
};
let ws_fixed = build_workspace(&fixed, Perm::IDENTITY, 100, 3, &opts);
assert!(ws_fixed.is_ols());
let (_, _, n, p, mixed, ids, _) = lmm_intercept_case();
let (sized, _ids, perm) = spec_sized_from_ids_pub(&mixed, &ids);
let ws_mix = build_workspace(&sized, perm, n, p, &FitOptions::default());
assert!(ws_mix.is_lmm_dense());
let glm = ModelSpec {
family: Family::Poisson {
link: crate::PoissonLink::Log,
},
re: None,
};
assert!(build_workspace(&glm, Perm::IDENTITY, 50, 2, &opts).is_glm());
let nb = ModelSpec {
family: Family::NegativeBinomial {
link: crate::NegBinomialLink::Log,
},
re: None,
};
assert!(build_workspace(&nb, Perm::IDENTITY, 50, 2, &opts).is_prebuilt());
}
#[test]
fn single_intercept_gaussian_routes_to_bobyqa_not_brent() {
let (_x, _y, n, p, model, ids, opts) = lmm_intercept_case();
let (sized, _ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let ws = build_workspace(&sized, perm, n, p, &opts);
assert!(ws.is_lmm_dense());
}
#[test]
fn fit_on_ols_reused_ws_near_identical_to_fit_cold() {
let (x, y, n, p, model, ids, opts) = ols_case();
let cold = fit_cold(&x, &y, n, p, &model, &ids, &opts);
let mut ws = build_workspace(&model, Perm::IDENTITY, n, p, &opts);
let f1 = fit_on(&mut ws, &x, &y, &ids, None, &opts).into_fit(&x, &y, &ids, n, p, &model, &opts);
let f2 = fit_on(&mut ws, &x, &y, &ids, None, &opts).into_fit(&x, &y, &ids, n, p, &model, &opts);
assert_near(&cold.beta, &f1.beta, "beta vs fit_cold");
assert_near(&cold.se, &f1.se, "se vs fit_cold");
assert_near(&f1.beta, &f2.beta, "beta first vs second reuse");
}
#[test]
fn fit_on_reused_ws_near_identical_to_fit_cold_lmm() {
let (x, y, n, p, model, ids, opts) = lmm_intercept_case();
let cold = fit_cold(&x, &y, n, p, &model, &ids, &opts);
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let mut ws = build_workspace(&sized, perm, n, p, &opts);
let f1 = fit_on(&mut ws, &x, &y, &ids, None, &opts).into_fit(&x, &y, &ids, n, p, &model, &opts);
assert_near(&cold.beta, &f1.beta, "beta vs fit_cold");
assert_near(&cold.tau2, &f1.tau2, "tau2 vs fit_cold");
assert_near(
&[cold.dispersion],
&[f1.dispersion],
"dispersion vs fit_cold",
);
let y2: Vec<f64> = y
.iter()
.enumerate()
.map(|(i, &v)| v + 0.3 * ((i % 7) as f64 - 3.0))
.collect();
let cold2 = fit_cold(&x, &y2, n, p, &model, &ids, &opts);
let f2 =
fit_on(&mut ws, &x, &y2, &ids, None, &opts).into_fit(&x, &y2, &ids, n, p, &model, &opts);
assert_near(&cold2.beta, &f2.beta, "draw-2 beta vs fit_cold");
assert_near(&cold2.tau2, &f2.tau2, "draw-2 tau2 vs fit_cold");
assert!(
(f1.beta[0] - f2.beta[0]).abs() > 1e-6,
"draws are degenerate"
);
}
#[test]
fn fit_on_glmm_dense_matches_fit_cold() {
let (x, y, n, p, model, ids, opts) = glmm_binomial_intercept_case();
let cold = fit_cold(&x, &y, n, p, &model, &ids, &opts);
assert!(cold.converged());
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let mut ws = build_workspace(&sized, perm, n, p, &opts);
assert!(ws.is_glmm_dense());
let f1 = fit_on(&mut ws, &x, &y, &ids, None, &opts).into_fit(&x, &y, &ids, n, p, &model, &opts);
assert_near(&cold.beta, &f1.beta, "beta vs fit_cold");
assert_near(&cold.se, &f1.se, "se vs fit_cold");
let y2: Vec<f64> = y
.iter()
.enumerate()
.map(|(i, &v)| if i % 5 == 0 { 1.0 - v } else { v })
.collect();
let cold2 = fit_cold(&x, &y2, n, p, &model, &ids, &opts);
assert!(cold2.converged());
let f2 =
fit_on(&mut ws, &x, &y2, &ids, None, &opts).into_fit(&x, &y2, &ids, n, p, &model, &opts);
assert_near(&cold2.beta, &f2.beta, "draw-2 beta vs fit_cold");
assert_near(&cold2.se, &f2.se, "draw-2 se vs fit_cold");
}
#[test]
#[should_panic(expected = "shape")]
fn fit_on_panics_on_level_count_mismatch() {
let (x, y, n, p, model, ids, opts) = lmm_intercept_case();
let (sized, _ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let mut ws = build_workspace(&sized, perm, n, p, &opts);
let fewer = GroupIds {
primary: vec![0u32; n],
extra: vec![],
};
let _ = fit_on(&mut ws, &x, &y, &fewer, None, &opts);
}
fn crossed_extra_case(
extra_slopes: Vec<u32>,
) -> (
Vec<f64>,
Vec<f64>,
usize,
usize,
ModelSpec,
GroupIds,
FitOptions,
) {
let g1 = 8usize; let g2 = 5usize; let n = 80usize;
let p = 2usize;
let mut st = 99u64;
let u1: Vec<f64> = (0..g1).map(|_| 0.5 * lcg(&mut st)).collect();
let s2: Vec<f64> = (0..g2).map(|_| 0.4 * lcg(&mut st)).collect();
let mut x = vec![0.0f64; n * p];
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 c1 = i % g1;
let c2 = i % g2;
pid[i] = c1 as u32;
eid[i] = c2 as u32;
let x1 = lcg(&mut st);
x[i * 2] = 1.0;
x[i * 2 + 1] = x1;
y[i] = 0.5 + 0.7 * x1 + u1[c1] + s2[c2] * x1 + 0.2 * lcg(&mut st);
}
let model = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 1 },
slopes: vec![],
extra_groupings: vec![Grouping {
relation: GroupingRelation::Crossed { n_clusters: 1 },
slopes: extra_slopes,
}],
}),
};
let ids = GroupIds {
primary: pid,
extra: vec![eid],
};
let opts = FitOptions {
target_indices: vec![0, 1],
..FitOptions::default()
};
(x, y, n, p, model, ids, opts)
}
#[test]
fn fit_on_sparse_matches_fit_cold() {
let (x, y, n, p, model, ids, opts) = crossed_extra_case(vec![1]);
let cold = fit_cold(&x, &y, n, p, &model, &ids, &opts);
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let mut ws = build_workspace(&sized, perm, n, p, &opts);
assert!(ws.is_prebuilt()); let via =
fit_on(&mut ws, &x, &y, &ids, None, &opts).into_fit(&x, &y, &ids, n, p, &model, &opts);
assert_near(&cold.beta, &via.beta, "sparse Level 1 beta vs fit_cold");
assert_near(&cold.se, &via.se, "sparse Level 1 se vs fit_cold");
}
#[test]
#[should_panic(expected = "extra grouping")]
fn fit_on_panics_on_extra_level_count_overflow() {
let (x, y, n, p, model, ids, opts) = crossed_extra_case(vec![1]);
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let mut ws = build_workspace(&sized, perm, n, p, &opts);
let mut more = ids.into_owned();
more.extra[0][0] = *more.extra[0].iter().max().unwrap() + 1;
let _ = fit_on(&mut ws, &x, &y, &more, None, &opts);
}
#[test]
#[should_panic(expected = "extra grouping")]
fn fit_on_panics_on_dense_lmm_extra_level_count_overflow() {
let (x, y, n, p, model, ids, opts) = crossed_extra_case(vec![]);
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let mut ws = build_workspace(&sized, perm, n, p, &opts);
assert!(ws.is_lmm_dense());
let mut more = ids.into_owned();
more.extra[0][0] = *more.extra[0].iter().max().unwrap() + 1;
let _ = fit_on(&mut ws, &x, &y, &more, None, &opts);
}
#[test]
#[should_panic(expected = "extra grouping")]
fn fit_on_panics_on_dense_glmm_extra_level_count_overflow() {
let (x, y, n, p, model, ids, opts) = crossed_extra_glmm_case();
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let mut ws = build_workspace(&sized, perm, n, p, &opts);
assert!(ws.is_glmm_dense());
let mut more = ids.into_owned();
more.extra[0][0] = *more.extra[0].iter().max().unwrap() + 1;
let _ = fit_on(&mut ws, &x, &y, &more, None, &opts);
}
#[test]
fn fit_on_accepts_nested_draw_below_capacity() {
let (x, y, n, p, model, ids, opts) = nested_unbalanced_case();
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let mut ws = build_workspace(&sized, perm, n, p, &opts);
assert_eq!(ws.build_extra_capacity, vec![12]); let used = *ids.extra[0].iter().max().unwrap() as usize + 1;
assert_eq!(used, 11); let _ = fit_on(&mut ws, &x, &y, &ids, None, &opts);
}
#[test]
#[should_panic(expected = "target count is frozen at build")]
fn fit_on_panics_on_grown_target_count() {
let (x, y, n, p, model, ids, opts) = ols_case();
let mut ws = build_workspace(&model, Perm::IDENTITY, n, p, &opts);
let wider = FitOptions {
target_indices: vec![0, 1, 2],
..opts
};
let _ = fit_on(&mut ws, &x, &y, &ids, None, &wider);
}
#[test]
#[should_panic(expected = "weights presence is frozen at build")]
fn fit_on_panics_when_a_weighted_glmm_workspace_is_reused_unweighted() {
let (x, y, n, p, model, ids, opts) = glmm_binomial_intercept_case();
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
let w: Vec<f64> = (0..n).map(|i| 1.0 + (i % 3) as f64).collect();
let opts_w = FitOptions {
weights: Some(w),
..opts.clone()
};
let mut ws = build_workspace(&sized, perm, n, p, &opts_w);
let _ = fit_on(&mut ws, &x, &y, &ids, None, &opts_w);
let _ = fit_on(&mut ws, &x, &y, &ids, None, &opts);
}
#[test]
#[should_panic(expected = "at most one NestedWithin")]
fn build_workspace_rejects_two_nested_groupings() {
let (_, _, n, p, model, _, opts) = nested_unbalanced_case();
let mut sized = model;
let nested = sized.re.as_ref().unwrap().extra_groupings[0].clone();
sized.re.as_mut().unwrap().extra_groupings.push(nested);
let _ = build_workspace(&sized, Perm::IDENTITY, n, p, &opts);
}
#[test]
fn fit_on_varying_n_below_n_max_matches_fit_cold() {
let (full_x, full_y, n_max, p, model, _ids, opts) = ols_case();
let mut ws = build_workspace(&model, Perm::IDENTITY, n_max, p, &opts);
for &n in &[10usize, 21, n_max] {
let x = &full_x[..n * p];
let y = &full_y[..n];
let ids = GroupIds::default();
let cold = fit_cold(x, y, n, p, &model, &ids, &opts);
let via =
fit_on(&mut ws, x, y, &ids, None, &opts).into_fit(x, y, &ids, n, p, &model, &opts);
assert_near(&cold.beta, &via.beta, &format!("n={n} beta"));
assert_near(&cold.se, &via.se, &format!("n={n} se"));
}
}
#[test]
fn fit_on_weighted_reuse_matches_fit_cold() {
let (x, y, n, p, model, ids, mut opts_w) = ols_case();
let w: Vec<f64> = (0..n).map(|i| 1.0 + (i % 4) as f64).collect();
opts_w.weights = Some(w);
let opts_unit = FitOptions {
weights: Some(vec![1.0; n]),
..opts_w.clone()
};
let mut ws = build_workspace(&model, Perm::IDENTITY, n, p, &opts_w);
let cold_w = fit_cold(&x, &y, n, p, &model, &ids, &opts_w);
let via_w =
fit_on(&mut ws, &x, &y, &ids, None, &opts_w).into_fit(&x, &y, &ids, n, p, &model, &opts_w);
assert_near(&cold_w.beta, &via_w.beta, "weighted beta");
let cold_u = fit_cold(&x, &y, n, p, &model, &ids, &opts_unit);
let via_u = fit_on(&mut ws, &x, &y, &ids, None, &opts_unit)
.into_fit(&x, &y, &ids, n, p, &model, &opts_unit);
assert_near(&cold_u.beta, &via_u.beta, "unit-weight-after-weighted beta");
assert_near(&cold_u.se, &via_u.se, "unit-weight-after-weighted se");
}
#[test]
fn fit_on_ols_smaller_then_larger_matches_fit_cold() {
let (full_x, full_y, n_max, p, model, _ids, opts) = ols_case();
let mut ws = build_workspace(&model, Perm::IDENTITY, n_max, p, &opts);
for &n in &[12usize, n_max] {
let x = &full_x[..n * p];
let y = &full_y[..n];
let ids = GroupIds::default();
let cold = fit_cold(x, y, n, p, &model, &ids, &opts);
let via =
fit_on(&mut ws, x, y, &ids, None, &opts).into_fit(x, y, &ids, n, p, &model, &opts);
assert_near(&cold.beta, &via.beta, &format!("n={n} beta"));
assert_near(&cold.se, &via.se, &format!("n={n} se"));
}
}
#[test]
fn fit_on_ols_larger_then_smaller_matches_fit_cold() {
let (full_x, full_y, n_max, p, model, _ids, opts) = ols_case();
let mut ws = build_workspace(&model, Perm::IDENTITY, n_max, p, &opts);
for &n in &[n_max, 12usize] {
let x = &full_x[..n * p];
let y = &full_y[..n];
let ids = GroupIds::default();
let cold = fit_cold(x, y, n, p, &model, &ids, &opts);
let via =
fit_on(&mut ws, x, y, &ids, None, &opts).into_fit(x, y, &ids, n, p, &model, &opts);
assert_near(&cold.beta, &via.beta, &format!("n={n} beta"));
assert_near(&cold.se, &via.se, &format!("n={n} se"));
}
}
#[test]
fn fit_on_glm_smaller_then_larger_matches_fit_cold() {
let (full_x, full_y, n_max, p, model, _ids, opts) = glm_case();
let mut ws = build_workspace(&model, Perm::IDENTITY, n_max, p, &opts);
assert!(ws.is_glm());
for &n in &[12usize, n_max] {
let x = &full_x[..n * p];
let y = &full_y[..n];
let ids = GroupIds::default();
let cold = fit_cold(x, y, n, p, &model, &ids, &opts);
let via =
fit_on(&mut ws, x, y, &ids, None, &opts).into_fit(x, y, &ids, n, p, &model, &opts);
assert_near(&cold.beta, &via.beta, &format!("n={n} beta"));
assert_near(&cold.se, &via.se, &format!("n={n} se"));
}
}
#[test]
fn fit_on_glm_larger_then_smaller_matches_fit_cold() {
let (full_x, full_y, n_max, p, model, _ids, opts) = glm_case();
let mut ws = build_workspace(&model, Perm::IDENTITY, n_max, p, &opts);
for &n in &[n_max, 12usize] {
let x = &full_x[..n * p];
let y = &full_y[..n];
let ids = GroupIds::default();
let cold = fit_cold(x, y, n, p, &model, &ids, &opts);
let via =
fit_on(&mut ws, x, y, &ids, None, &opts).into_fit(x, y, &ids, n, p, &model, &opts);
assert_near(&cold.beta, &via.beta, &format!("n={n} beta"));
assert_near(&cold.se, &via.se, &format!("n={n} se"));
}
}
#[test]
fn fit_on_lmm_dense_smaller_then_larger_matches_fit_cold() {
let (full_x, full_y, n_max, p, model, ids_full, opts) = lmm_intercept_case();
let (sized, ids_full, perm) = spec_sized_from_ids_pub(&model, &ids_full);
let mut ws = build_workspace(&sized, perm, n_max, p, &opts);
assert!(ws.is_lmm_dense());
for &n in &[24usize, n_max] {
let x = &full_x[..n * p];
let y = &full_y[..n];
let ids = GroupIds {
primary: ids_full.primary[..n].to_vec(),
extra: vec![],
};
let cold = fit_cold(x, y, n, p, &model, &ids, &opts);
let via =
fit_on(&mut ws, x, y, &ids, None, &opts).into_fit(x, y, &ids, n, p, &model, &opts);
assert_near(&cold.beta, &via.beta, &format!("n={n} beta"));
assert_near(&cold.tau2, &via.tau2, &format!("n={n} tau2"));
}
}
#[test]
fn fit_on_lmm_dense_larger_then_smaller_matches_fit_cold() {
let (full_x, full_y, n_max, p, model, ids_full, opts) = lmm_intercept_case();
let (sized, ids_full, perm) = spec_sized_from_ids_pub(&model, &ids_full);
let mut ws = build_workspace(&sized, perm, n_max, p, &opts);
for &n in &[n_max, 24usize] {
let x = &full_x[..n * p];
let y = &full_y[..n];
let ids = GroupIds {
primary: ids_full.primary[..n].to_vec(),
extra: vec![],
};
let cold = fit_cold(x, y, n, p, &model, &ids, &opts);
let via =
fit_on(&mut ws, x, y, &ids, None, &opts).into_fit(x, y, &ids, n, p, &model, &opts);
assert_near(&cold.beta, &via.beta, &format!("n={n} beta"));
assert_near(&cold.tau2, &via.tau2, &format!("n={n} tau2"));
}
}
#[test]
fn fit_on_ols_scaled_x_gate_matches_has_weights() {
let (x, y, n, p, model, ids, opts) = ols_case();
assert!(opts.weights.is_none());
let w: Vec<f64> = (0..n).map(|i| 1.0 + (i % 3) as f64).collect();
let opts_w = FitOptions {
weights: Some(w),
..opts.clone()
};
let cold_u = fit_cold(&x, &y, n, p, &model, &ids, &opts);
let mut ws_u = build_workspace(&model, Perm::IDENTITY, n, p, &opts);
let via_u =
fit_on(&mut ws_u, &x, &y, &ids, None, &opts).into_fit(&x, &y, &ids, n, p, &model, &opts);
assert_near(&cold_u.beta, &via_u.beta, "unweighted beta");
assert_near(&cold_u.se, &via_u.se, "unweighted se");
let cold_w = fit_cold(&x, &y, n, p, &model, &ids, &opts_w);
let mut ws_w = build_workspace(&model, Perm::IDENTITY, n, p, &opts_w);
let via_w = fit_on(&mut ws_w, &x, &y, &ids, None, &opts_w)
.into_fit(&x, &y, &ids, n, p, &model, &opts_w);
assert_near(&cold_w.beta, &via_w.beta, "weighted beta");
assert_near(&cold_w.se, &via_w.se, "weighted se");
}
#[test]
fn fit_on_lmm_dense_offset_round_trip_varying_n() {
let (full_x, full_y, n_max, p, model, ids_full, mut opts) = lmm_intercept_case();
let offset: Vec<f64> = (0..n_max).map(|i| 0.05 * ((i % 4) as f64 - 1.5)).collect();
opts.offset = Some(offset.clone());
let (sized, ids_full, perm) = spec_sized_from_ids_pub(&model, &ids_full);
let mut ws = build_workspace(&sized, perm, n_max, p, &opts);
for &n in &[24usize, n_max, 24usize] {
let x = &full_x[..n * p];
let y = &full_y[..n];
let ids = GroupIds {
primary: ids_full.primary[..n].to_vec(),
extra: vec![],
};
let opts_n = FitOptions {
offset: Some(offset[..n].to_vec()),
..opts.clone()
};
let cold = fit_cold(x, y, n, p, &model, &ids, &opts_n);
let via =
fit_on(&mut ws, x, y, &ids, None, &opts).into_fit(x, y, &ids, n, p, &model, &opts_n);
assert_near(&cold.beta, &via.beta, &format!("n={n} beta"));
assert_near(&cold.tau2, &via.tau2, &format!("n={n} tau2"));
}
}
#[cfg(feature = "alloc-tests")]
fn lmm_slope_case() -> (
Vec<f64>,
Vec<f64>,
usize,
usize,
ModelSpec,
GroupIds,
FitOptions,
) {
let n_clusters = 6usize;
let per = 10usize;
let n = n_clusters * per;
let p = 2usize;
let mut st = 271u64;
let mut x = vec![0.0f64; n * p];
let mut y = vec![0.0f64; n];
let mut ids_v = vec![0u32; n];
for i in 0..n {
let c = i % n_clusters;
ids_v[i] = c as u32;
let x1 = 7.0 * lcg(&mut st);
x[i * 2] = 1.0;
x[i * 2 + 1] = x1;
let u = 0.3 * ((c as f64) - (n_clusters as f64) / 2.0);
let s = 0.05 * ((c as f64) - 2.0);
y[i] = 0.5 + 0.4 * x1 + u + s * x1 + 0.1 * lcg(&mut st);
}
let model = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 1 },
slopes: vec![1],
extra_groupings: vec![],
}),
};
let ids = GroupIds {
primary: ids_v,
extra: vec![],
};
let opts = FitOptions {
target_indices: vec![0, 1],
..FitOptions::default()
};
(x, y, n, p, model, ids, opts)
}
#[cfg(feature = "alloc-tests")]
fn reordered_crossed_case() -> (
Vec<f64>,
Vec<f64>,
usize,
usize,
ModelSpec,
GroupIds,
FitOptions,
) {
let g1 = 4usize; let g2 = 9usize; let n = 72usize;
let p = 2usize;
let mut st = 4242u64;
let u1: Vec<f64> = (0..g1).map(|_| 0.5 * lcg(&mut st)).collect();
let u2: Vec<f64> = (0..g2).map(|_| 0.4 * lcg(&mut st)).collect();
let mut x = vec![0.0f64; n * p];
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 (c1, c2) = (i % g1, i % g2);
pid[i] = c1 as u32;
eid[i] = c2 as u32;
let x1 = lcg(&mut st);
x[i * 2] = 1.0;
x[i * 2 + 1] = x1;
y[i] = 0.5 + 0.7 * x1 + u1[c1] + u2[c2] + 0.2 * lcg(&mut st);
}
let model = ModelSpec {
family: Family::Gaussian,
re: Some(ReStructure {
sizing: Sizing::FixedClusters { n_clusters: 1 },
slopes: vec![],
extra_groupings: vec![Grouping {
relation: GroupingRelation::Crossed { n_clusters: 1 },
slopes: vec![],
}],
}),
};
let ids = GroupIds {
primary: pid,
extra: vec![eid],
};
let opts = FitOptions {
target_indices: vec![0, 1],
..FitOptions::default()
};
(x, y, n, p, model, ids, opts)
}
#[cfg(feature = "alloc-tests")]
#[test]
#[ignore]
fn fit_on_theta_marshalling_bounded_alloc() {
let _serial = crate::test_support::alloc_test_guard();
const N_CALLS: usize = 100;
const BOUND: u64 = 6200;
let (xs, ys, ns, ps, ms, ids_s, os) = lmm_slope_case();
let (sized_s, ids_s, perm_s) = spec_sized_from_ids_pub(&ms, &ids_s);
assert!(perm_s.is_identity());
let mut ws_scaled = build_workspace(&sized_s, perm_s, ns, ps, &os);
assert!(ws_scaled.is_lmm_dense());
let (xr, yr, nr, pr, mr, ids_r, or) = reordered_crossed_case();
let (sized_r, ids_r, perm_r) = spec_sized_from_ids_pub(&mr, &ids_r);
assert!(
!perm_r.is_identity(),
"the warm-start permutation arm needs a reordering workspace"
);
let mut ws_reordered = build_workspace(&sized_r, perm_r, nr, pr, &or);
assert!(ws_reordered.is_lmm_dense());
let start_s = {
let v = fit_on(&mut ws_scaled, &xs, &ys, &ids_s, None, &os);
assert!(
v.theta().iter().all(|t| t.is_finite()) && v.theta().len() == 3,
"the scaled arm must carry a q=2 θ"
);
crate::StartValues {
beta: v.betas().to_vec(),
theta: v.theta().to_vec(),
}
};
let start_r = {
let v = fit_on(&mut ws_reordered, &xr, &yr, &ids_r, None, &or);
crate::StartValues {
beta: v.betas().to_vec(),
theta: v.theta().to_vec(),
}
};
let profiler = dhat::Profiler::builder().testing().build();
for _ in 0..N_CALLS {
let v = fit_on(&mut ws_scaled, &xs, &ys, &ids_s, Some(&start_s), &os);
std::hint::black_box(std::hint::black_box(&v).theta());
let v = fit_on(&mut ws_reordered, &xr, &yr, &ids_r, Some(&start_r), &or);
std::hint::black_box(std::hint::black_box(&v).theta());
}
let stats = dhat::HeapStats::get();
drop(profiler);
assert!(
stats.total_blocks <= BOUND,
"fit_on allocated {} blocks across {} warm-path calls per arm (BOUND = {})",
stats.total_blocks,
N_CALLS,
BOUND
);
}
#[cfg(feature = "alloc-tests")]
#[test]
#[ignore]
fn lmm_dual_scratch_built_once_per_workspace() {
let _serial = crate::test_support::alloc_test_guard();
const N_CALLS: usize = 20;
const BOUND: u64 = 704;
let (x, y, n, p, model, ids, opts) = lmm_slope_case();
let (sized, ids, perm) = spec_sized_from_ids_pub(&model, &ids);
assert!(perm.is_identity());
let mut ws = build_workspace(&sized, perm, n, p, &opts);
assert!(ws.is_lmm_dense());
let start = {
let v = fit_on(&mut ws, &x, &y, &ids, None, &opts);
assert!(
v.converged(),
"the gate needs a converged fit — the derivative block is skipped otherwise"
);
crate::StartValues {
beta: v.betas().to_vec(),
theta: v.theta().to_vec(),
}
};
let profiler = dhat::Profiler::builder().testing().build();
for _ in 0..N_CALLS {
let v = fit_on(&mut ws, &x, &y, &ids, Some(&start), &opts);
std::hint::black_box(std::hint::black_box(&v).theta());
}
let stats = dhat::HeapStats::get();
drop(profiler);
assert!(
stats.total_blocks <= BOUND,
"fit_on allocated {} blocks across {} warm-path calls (BOUND = {})",
stats.total_blocks,
N_CALLS,
BOUND
);
}
#[cfg(feature = "counters")]
#[test]
fn eval_counters_count_evals_after_last_improvement() {
use crate::counters::{EvalCounters, Stage};
let mut c = EvalCounters::new();
for obj in [5.0, 4.0, 4.5, 3.0, 3.5, 3.5, 3.5] {
c.record_eval(Stage::Two, obj);
}
assert_eq!(c.stage_evals[Stage::Two as usize], 7);
assert_eq!(c.stage_last_improve[Stage::Two as usize], 4);
assert_eq!(c.evals_after_last_improve(Stage::Two), 3);
assert_eq!(c.evals_after_last_improve(Stage::One), 0);
}
#[cfg(feature = "counters")]
#[test]
fn eval_counters_bucket_pirls_iterations_per_eval() {
use crate::counters::EvalCounters;
let mut c = EvalCounters::new();
c.set_pirls_iters(3);
c.commit_pirls_iters();
c.set_pirls_iters(5);
c.commit_pirls_iters();
c.set_pirls_iters(3);
c.commit_pirls_iters();
assert_eq!(c.pirls_hist[3], 2);
assert_eq!(c.pirls_hist[5], 1);
assert_eq!(c.pirls_hist.iter().sum::<u32>(), 3);
}
#[cfg(feature = "counters")]
#[test]
fn eval_counters_accumulate_agq_nodes() {
use crate::counters::EvalCounters;
let mut c = EvalCounters::new();
c.record_agq_eval(8 * 7);
c.record_agq_eval(8 * 7);
assert_eq!(c.agq_evals, 2);
assert_eq!(c.agq_node_evals, 112);
}
#[cfg(feature = "counters")]
#[test]
fn fit_carries_zeroed_counters_on_closed_form_routes() {
use crate::counters::Stage;
let (x, y, n, p, model, ids, opts) = ols_case();
let f = fit_cold(&x, &y, n, p, &model, &ids, &opts);
assert_eq!(f.n_eval, 0, "OLS runs no optimizer");
assert_eq!(f.counters.stage_evals, [0, 0]);
assert_eq!(f.counters.evals_after_last_improve(Stage::Two), 0);
assert_eq!(f.counters.agq_evals, 0);
assert_eq!(f.counters.pirls_hist.iter().sum::<u32>(), 0);
}