glmm 0.3.2

Standalone f64 GLMM fit kernels (OLS, GLM, LMM, GLMM) in pure Rust on faer — the validation-pinned numerics from the MCPower engine.
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
//! LMM (`Family::Gaussian`, `re: Some`) dispatch — marshals `fit_warm`'s
//! inputs into the dense LMM workspace and calls `fit_lmm`. The numerical
//! kernel lives in `src/lmm/mod.rs`; this module only builds the workspace,
//! accumulates sufficient statistics, and maps `LmmFit` back to `Fit`.

use crate::lmm::{fit_lmm, LmmFit, LmmGroupings, LmmWorkspace};
// Used only by the test-only `fit_mle`/`fit_lmm_into` baselines. The stable core
// reaches the LMM path via `accumulate_lmm_rows`/`lmm_run_on`, which take neither.
#[cfg(test)]
use crate::{ModelSpec, StartValues};

use super::common::{
    assemble_varcorr, fill_se_by_predictor, nan_vcov, vcov_from_chol, FitDiagnostics,
};
// Used only by the test-only `fit_mle` baseline, which builds its own throwaway
// `Mat` — the stable core reaches the LMM path via `accumulate_lmm_rows`, which
// takes a caller-filled `x_mat` directly.
#[cfg(test)]
use super::common::{to_col_major, warm_theta};
use super::{Fit, FitOptions};

// ---------------------------------------------------------------------------
// LMM dispatch (Estimator::Mle)
// ---------------------------------------------------------------------------

/// Accumulates the θ-independent sufficient statistics (reset + add_rows_multi)
/// from an already-column-major `x_mat`. Single-sourced across `fit_mle` and
/// `with_lmm_objective`'s dense arm, which differ only in `weights` and in
/// whether the guard is written out at the call site. Takes `x_mat` rather than row-major `x` so the `fit_on` hot
/// path can fill a build-once `n_max`-sized buffer and pass a `subrows(0, n)`
/// view in, instead of allocating a fresh `Mat` every call; the three other
/// callers (`fit_mle`, `build_lmm_seam_ws`, `refit_lmm`) each build their own
/// throwaway `Mat` via `to_col_major` at the call site.
#[allow(clippy::too_many_arguments)] // marshals the kernel's (x_mat, y, n, p, ids…) surface
pub(super) fn accumulate_lmm_rows(
    ws: &mut LmmWorkspace,
    x_mat: faer::MatRef<'_, f64>,
    y: &[f64],
    n: usize,
    p: usize,
    cluster_ids: &[u32],
    extra_ids: &[Vec<u32>],
    weights: Option<&[f64]>,
) {
    ws.suff.reset();
    // Before any row is accumulated: the RE design-column scales are per dataset, and
    // every slope value the accumulator stores is already divided by them. The
    // workspace is reused across draws with different `x`, so this is refreshed per
    // call, not cached at construction — and outside the `n > 0 && p > 0` guard,
    // because `SuffStats::reset` does not clear the scales and an empty draw would
    // otherwise leave the previous draw's in place.
    ws.suff.groupings.set_slope_scales(x_mat, weights);
    if n > 0 && p > 0 {
        ws.suff
            .add_rows_multi(x_mat, y, cluster_ids, extra_ids, weights);
    }
}

/// Borrowed result of [`lmm_run_on`]: the [`LmmFit`] summary plus the workspace
/// result slots the `Fit` assembly reads (β̂, per-predictor Var/z², the augmented
/// factor, θ̂, groupings, row count). Lifetime ties back to the [`LmmWorkspace`]
/// that owns the storage.
pub(crate) struct LmmResultView<'a> {
    fit: LmmFit,
    betas: &'a [f64],
    var_diag: &'a [f64],
    // Read only by the loop-tier `t_sq` accessor below (dead in a default build).
    #[allow(dead_code)]
    t_sq: &'a [f64],
    factor: faer::MatRef<'a, f64>,
    theta: &'a [f64],
    groupings: &'a LmmGroupings,
    n_rows: usize,
    /// Spherical conditional modes û at θ̂, or empty when the recovery did not
    /// run or did not succeed (see `LmmFitScratch::ranef_ok`).
    ranef_u: &'a [f64],
    /// Derivative diagnostics, computed in [`lmm_run_on`] (the one place the
    /// suff stats are in scope): the box-projected θ-gradient ∞-norm of the
    /// REML criterion at the accepted θ̂ (caller's θ units; NaN when no exact
    /// gradient — extra slopes, non-converged), and the
    /// per-θ-coordinate boundary score `½·∂²D/∂θ_jj²·s_j²` (NaN except at
    /// pinned diagonals). Length `n_theta`.
    kkt_grad_norm: f64,
    boundary_score: Vec<f64>,
}

// Loop-tier read accessors (via the `FitView`/`loop_advanced` surface); the
// stable path reads these slots through `lmm_view_to_fit` instead.
#[allow(dead_code)]
impl LmmResultView<'_> {
    /// Per-target Wald z², predictor-indexed length p — only the
    /// `target_indices` slots are written; a non-target slot reads 0.0 on a
    /// fresh workspace or a previous fit's value on a reused one.
    pub(crate) fn t_sq(&self) -> &[f64] {
        self.t_sq
    }
    /// Fixed-effect estimates β̂, predictor-indexed.
    pub(crate) fn betas(&self) -> &[f64] {
        self.betas
    }
    /// Per-predictor Var(β̂_j), predictor-indexed length p — only the
    /// `target_indices` slots are written; a non-target slot reads 0.0 on a
    /// fresh workspace or a previous fit's value on a reused one.
    pub(crate) fn var_diag(&self) -> &[f64] {
        self.var_diag
    }
    /// This route's [`FitDiagnostics`] — every field is real here: the dense LMM
    /// is the one route that reports θ boundary state, per-component pins, AND a
    /// recorded pivot. The pivot floor is this route's own ([`crate::lmm::PIVOT_MIN`]),
    /// which sits at the same 1e-12 as OLS/GLM but was calibrated separately and
    /// stays a separate constant for that reason.
    pub(crate) fn diagnostics(&self) -> FitDiagnostics {
        FitDiagnostics {
            converged: self.fit.converged,
            boundary_hit: self.fit.boundary_hit,
            pinned_components: self.fit.pinned_components,
            pivot: self.fit.pivot,
            pivot_col: self.fit.pivot_col,
            ill_conditioned: self.fit.pivot < crate::lmm::PIVOT_MIN,
            // No PIRLS on the Gaussian LMM path (profiled REML, no IRLS inner
            // loop) — always 0/false.
            pirls_exhausted: 0,
            final_pirls_exhausted: false,
            // No FD-Hessian SE pass on the Gaussian LMM path (no PIRLS to
            // perturb) — always false.
            hessian_fallback: false,
        }
    }
    /// Joint Wald-χ² over the target set (the omnibus significance read).
    pub(crate) fn joint_t_sq(&self) -> f64 {
        self.fit.joint_t_sq
    }
    /// Objective evaluations the θ-solve spent.
    pub(crate) fn n_eval(&self) -> usize {
        self.fit.n_eval
    }
    /// Residual variance σ̂² (the LMM dispersion).
    pub(crate) fn dispersion(&self) -> f64 {
        self.fit.sigma_sq
    }
    /// Fitted θ̂ vech (primary block then extras, column-major lower-triangular),
    /// in the solver's INTERNAL RE column scale — `FitView::theta` divides it back
    /// into the design's own units for the warm-start carry.
    pub(crate) fn theta(&self) -> &[f64] {
        self.theta
    }
    /// Grouping structure, for the internal RE column scales θ̂ carries.
    pub(crate) fn groupings(&self) -> &LmmGroupings {
        self.groupings
    }
}

/// Runs the θ-solve on a pre-accumulated workspace and returns a borrowed view
/// of the results. Caller contract: `ws.suff` holds the accumulated rows
/// ([`accumulate_lmm_rows`]). Warm start threads `theta` only — the LMM β is
/// solved exactly given θ, so a β start is irrelevant; `None` (cold) uses the
/// kernel's THETA0 blind start. `want_score` is `FitOptions::boundary_score`:
/// the hyper-dual REML Hessian behind the pinned-component score runs only
/// when asked; the gradient (and `kkt_grad_norm`) runs regardless.
pub(crate) fn lmm_run_on<'a>(
    ws: &'a mut LmmWorkspace,
    target_indices: &[u32],
    theta_start: Option<&[f64]>,
    want_score: bool,
) -> LmmResultView<'a> {
    let fit = fit_lmm(ws, target_indices, theta_start);

    // Derivative diagnostics, from the exact dual-number REML derivatives
    // (`reml_gradient`/`reml_hessian`). Same box projection, the same opt-in
    // for the score, and the same `½·H_jj` boundary-score identity as the GLMM
    // block in `glmm::fit_glmm` — the θ box is the same one
    // (`blind_theta_and_bounds`), the scales come from the same
    // `theta_row_scales`, the gradient taking them once and the score twice.
    // Both scratches live on the workspace (`LmmWorkspace::dual_scratch` /
    // `hyper_scratch`), built on the first request and reused by every later
    // fit on the same workspace — a warm refit allocates nothing here. They are
    // separate slots because the gradient and the score resolve to different
    // scalar types. `Unsupported` (extra slopes) leaves NaN, and so does an
    // absent score scratch: `LmmHyperScratch::for_groupings` is `None` above
    // `n_theta = 12`, where the gradient's own scratch still resolves because
    // `reml_gradient` chunks.
    let n_theta = ws.suff.groupings.n_theta();
    let mut kkt_grad_norm = f64::NAN;
    let mut boundary_score = vec![f64::NAN; n_theta];
    if fit.converged {
        let p = ws.suff.m - 1;
        // Owned, so it survives the two `&mut ws` destructures below; `diag`
        // borrows the groupings and has to be re-derived inside each.
        let sc = ws.suff.groupings.theta_row_scales();
        if ws.dual_scratch.is_none() {
            ws.dual_scratch =
                crate::lmm::LmmDualScratch::for_groupings(n_theta, p, &ws.suff.groupings)
                    .map(Box::new);
        }
        {
            let LmmWorkspace {
                suff,
                theta,
                dual_scratch,
                ..
            } = &mut *ws;
            if let Some(scratch) = dual_scratch.as_deref_mut() {
                let g = &suff.groupings;
                let diag = g.diagonal_theta();
                let mut grad = vec![0.0; n_theta];
                let st = crate::lmm::reml_gradient(&theta[..n_theta], suff, scratch, &mut grad);
                if matches!(st, crate::glmm::DerivStatus::Ok(_)) {
                    let mut acc = 0.0_f64;
                    for j in 0..n_theta {
                        let gj = grad[j];
                        let is_diag = diag.contains(&j);
                        let (lo, hi) = if is_diag {
                            (0.0, crate::lmm::THETA_HI)
                        } else {
                            (-crate::lmm::THETA_HI, crate::lmm::THETA_HI)
                        };
                        // Projected gradient on a box, in the internal θ̃ where the
                        // box lives; back-mapped by ×s_j (∂D/∂θ = s·∂D/∂θ̃) — see
                        // the GLMM twin for the derivation.
                        let pj = if theta[j] <= lo {
                            gj.min(0.0)
                        } else if theta[j] >= hi {
                            gj.max(0.0)
                        } else {
                            gj
                        };
                        acc = acc.max((pj * sc[j]).abs());
                    }
                    kkt_grad_norm = acc;
                }
            }
        }
        if fit.pinned_components != 0 && want_score {
            if ws.hyper_scratch.is_none() {
                ws.hyper_scratch =
                    crate::lmm::LmmHyperScratch::for_groupings(n_theta, p, &ws.suff.groupings)
                        .map(Box::new);
            }
            let LmmWorkspace {
                suff,
                theta,
                hyper_scratch,
                ..
            } = &mut *ws;
            if let Some(scratch) = hyper_scratch.as_deref_mut() {
                let g = &suff.groupings;
                let diag = g.diagonal_theta();
                let mut grad = vec![0.0; n_theta];
                let mut hess = faer::Mat::<f64>::zeros(n_theta, n_theta);
                let st = crate::lmm::reml_hessian(
                    &theta[..n_theta],
                    suff,
                    scratch,
                    &mut grad,
                    &mut hess,
                );
                if matches!(st, crate::glmm::DerivStatus::Ok(_)) {
                    for (kk, &ti) in diag.iter().enumerate() {
                        if kk < u64::BITS as usize
                            && (fit.pinned_components >> kk) & 1 == 1
                            && !g.diagonal_has_nonzero_below(kk, &theta[..n_theta])
                        {
                            // s = θ_jj², θ̃ = sc·θ ⇒ dD/ds = sc²·dD/ds̃: the
                            // score takes the SQUARE of the scale. Valid only
                            // where the deviance is even in θ_jj at θ_jj = 0
                            // — see `LmmGroupings::diagonal_has_nonzero_below`
                            // for why a non-zero entry below the diagonal
                            // breaks that. `canonicalize_pinned_blocks` runs
                            // after the pin loop and zeroes that column, so
                            // this gate never fires; it is defensive only.
                            boundary_score[ti] = 0.5 * hess[(ti, ti)] * sc[ti] * sc[ti];
                        }
                    }
                }
            }
        }
    }

    LmmResultView {
        ranef_u: if ws.fit.ranef_ok {
            &ws.fit.ranef_u[..ws.suff.groupings.k_total]
        } else {
            &[]
        },
        fit,
        betas: &ws.fit.betas,
        var_diag: &ws.fit.var_diag,
        t_sq: &ws.fit.t_sq,
        factor: ws.fit.factor.as_ref(),
        theta: &ws.theta,
        groupings: &ws.suff.groupings,
        n_rows: ws.suff.n_rows,
        kkt_grad_norm,
        boundary_score,
    }
}

/// Maps an [`LmmResultView`] to the full stable `Fit`: SE from `var_diag`, tau2
/// from θ̂, varcorr, vcov from the augmented factor, the REML criterion loglik,
/// the conditional modes + per-row means, and — when `opts.weights` is set — the
/// −Σlog wᵢ weighted-deviance correction.
///
/// `x`/`ids` are here for `fitted` alone (`Xβ̂ + Zb̂` needs the design rows and
/// the per-row levels), the same way `ols_view_to_fit` already takes `x`/`y`.
pub(crate) fn lmm_view_to_fit(
    view: &LmmResultView<'_>,
    x: &[f64],
    ids: &crate::GroupIds,
    n: usize,
    p: usize,
    opts: &FitOptions,
) -> Fit {
    let lmm_fit = &view.fit;
    let diag = view.diagnostics();
    // view.betas: length p, all fixed effects.
    // view.var_diag: length p, predictor-indexed (LME/LMM are predictor-indexed, unlike OLS).
    let beta = view.betas.to_vec();
    let sigma_sq = lmm_fit.sigma_sq;

    let mut se = vec![f64::NAN; p];
    fill_se_by_predictor(view.var_diag, &opts.target_indices, &mut se);

    // tau2[k] = theta[k]^2 * sigma_sq — the k-th variance component in original scale.
    // θ̂ holds the fitted Cholesky parameters; diagonal entries satisfy
    // theta[k] = sqrt(tau_k / sigma_sq), so theta[k]^2 * sigma_sq = tau_k.
    // Gated on the finite endpoint (`deviance`), not `converged` — the plateau
    // policy: a `MaxFunReached` cap-out reports an honest θ̂ (its finite
    // endpoint) with `converged == false` rather than NaN-filling.
    let has_endpoint = lmm_fit.deviance.is_finite();
    // θ̂ is in the solver's internal RE units (`LmmGroupings::set_slope_scales`);
    // `theta_row_scales` divides each entry back into the design's own units
    // before it is squared, so `tau2` is a variance in the user's scale.
    let theta_scales = view.groupings.theta_row_scales();
    let tau2: Vec<f64> = if has_endpoint {
        view.theta
            .iter()
            .zip(theta_scales.iter())
            .map(|(&t, &s)| (t / s) * (t / s) * sigma_sq)
            .collect()
    } else {
        view.theta.iter().map(|_| f64::NAN).collect()
    };

    // varcorr[g] = vech(σ̂²·Λ̂_gΛ̂_g') per grouping — the q≥2-valid
    // covariance; equals tau2's diagonal only at the (0,0) primary entry.
    let varcorr = if has_endpoint {
        assemble_varcorr(view.theta, view.groupings, sigma_sq)
    } else {
        vec![]
    };

    // Var(β̂) = σ̂²·(X'V⁻¹X)⁻¹ from the top-left p×p block of the augmented
    // factor — the same `L_XX` and the same σ̂² the `var_diag` forward solve
    // above uses, so `vcov`'s diagonal is `var_diag`. Gated on the finite
    // endpoint like `tau2`/`varcorr`: the degenerate return NaN-fills
    // `var_diag` and `sigma_sq` together, so `se` and `vcov` go NaN together.
    let vcov = if has_endpoint {
        vcov_from_chol(view.factor, p, &opts.target_indices, sigma_sq)
    } else {
        nan_vcov(p)
    };

    // Equal to `n` on every real path; diverges only in the degenerate
    // n > 0 && p == 0 case, where `accumulate_lmm_rows` skips `add_rows_multi`
    // (and so never increments `n_rows`) while `n` still reflects the caller's
    // row count.
    let n_rows = view.n_rows;
    let n_theta = view.theta.len();

    // Level counts depend only on the design, so they are reported on every
    // fit; a header line ("Number of obs: …, groups: …") prints on non-converged
    // fits too. `fitted`/`ranef` stay gated on `converged`: a conditional mode
    // at a non-converged θ̂ is not a BLUP of anything (the same rule the GLMM
    // paths apply). `ranef_u` is empty when the recovery did not run.
    let ranef_levels = super::common::ranef_level_counts(view.groupings);
    let (fitted, ranef) = if lmm_fit.converged && !view.ranef_u.is_empty() {
        let ranef = super::common::assemble_ranef_sparse(view.theta, view.groupings, view.ranef_u);
        let fitted = super::common::lmm_fitted(
            x,
            n,
            p,
            view.betas,
            &ranef,
            view.groupings,
            &ids.primary,
            &ids.extra,
            opts.offset.as_deref(),
        );
        (fitted, ranef)
    } else {
        (vec![], vec![])
    };

    let mut diagnostics = super::common::materialize_diagnostics(&diag, p, &varcorr);
    // Derivative diagnostics, computed in `lmm_run_on` and carried on the
    // view (they do not pass through `FitDiagnostics` — same routing rule as
    // the GLMM side, see `fit/glmm.rs`). Both are NaN/empty already when the
    // fit did not converge (`lmm_run_on` gates the computation on it).
    diagnostics.kkt_grad_norm = view.kkt_grad_norm;
    // `pinned_scores` is keyed to `diagonal_theta` order like `pinned_flags`,
    // so collapse the θ-coordinate buffer onto the diagonals first.
    let diag_scores: Vec<f64> = view
        .groupings
        .diagonal_theta()
        .iter()
        .map(|&ti| view.boundary_score[ti])
        .collect();
    diagnostics.boundary_score = super::common::pinned_scores(&diag_scores, &varcorr);
    let mut fit = Fit {
        beta,
        se,
        vcov,
        tau2,
        // REML σ̂² (NaN-filled by the kernel alongside var_diag on the
        // degenerate path, so this stays honest without an endpoint gate).
        dispersion: sigma_sq,
        diagnostics,
        varcorr,
        stddev_se: vec![], // LMM has no Hessian SE machinery
        n_eval: lmm_fit.n_eval,
        #[cfg(feature = "counters")]
        counters: lmm_fit.counters,
        deviance: lmm_fit.deviance,
        // REML criterion on the logLik scale, from the base (unweighted) deviance;
        // the weighted correction below rewrites both together (using `n`, not
        // `n_rows` — see note above `n_rows`).
        loglik: super::common::lmm_loglik(lmm_fit.deviance, n_rows, p),
        df: if has_endpoint { p + n_theta + 1 } else { 0 },
        reml: true,
        fitted,
        ranef,
        ranef_levels,
    };
    fit.diagnostics.singular = fit.diagnostics.singular
        || fit.has_negligible_component(&super::common::re_scale_grid(view.groupings));

    // Weighted Gaussian log-density carries +½Σlog wᵢ per row; on the −2ℓ scale
    // the REML deviance gains −Σlog wᵢ (θ-independent — added post-optimization,
    // argmin unchanged), and the criterion-scale loglik is recomputed from the
    // corrected deviance. Matches lme4's weighted REMLcrit up to the additive
    // constant the engine strips from its deviance convention (documented on
    // `lmm::reml_deviance`). This is the single site every caller reaches, so
    // no caller applies the correction itself.
    if let Some(w) = &opts.weights {
        fit.deviance -= w.iter().map(|v| v.ln()).sum::<f64>();
        fit.loglik = super::common::lmm_loglik(fit.deviance, n, p);
    }
    fit
}

/// θ-solve + `Fit` assembly from a pre-built, already-accumulated workspace.
/// Composes [`lmm_run_on`] + [`lmm_view_to_fit`]. Caller contract: `ws.suff`
/// holds the accumulated rows (`accumulate_lmm_rows`, reset + add_rows_multi per
/// dataset). The weighted-deviance correction is inside `lmm_view_to_fit`, so
/// callers must NOT re-apply it. Only the test-gated `fit_mle` and `refit_lmm`
/// baselines compose through this; the core calls `lmm_run_on`/`lmm_view_to_fit`
/// directly.
#[cfg(test)]
pub(super) fn fit_lmm_into(
    ws: &mut LmmWorkspace,
    x: &[f64],
    ids: &crate::GroupIds,
    n: usize,
    p: usize,
    opts: &FitOptions,
    start: Option<&StartValues>,
) -> Fit {
    let view = lmm_run_on(
        ws,
        &opts.target_indices,
        warm_theta(start),
        opts.boundary_score,
    );
    lmm_view_to_fit(&view, x, ids, n, p, opts)
}
/// LMM dispatch adapter. `cluster_ids`/`extra_ids` are the per-row level ids from
/// the entry's [`GroupIds`]; `model` is the sizing-corrected spec (counts derived
/// from those ids). Slope x-columns come from `model.re`. Mirrors `fit_glmm`.
/// Test-only baseline since the stable path dispatches through the unified core
/// ([`super::core::fit_on`]) over `accumulate_lmm_rows`/`lmm_run_on`.
#[cfg(test)]
#[allow(clippy::too_many_arguments)] // marshals the kernel's (x, y, n, p, spec, ids…) surface
pub(super) fn fit_mle(
    x: &[f64],
    y: &[f64],
    n: usize,
    p: usize,
    model: &ModelSpec,
    cluster_ids: &[u32],
    extra_ids: &[Vec<u32>],
    start: Option<&StartValues>,
    opts: &FitOptions,
) -> Fit {
    let re = model
        .re
        .as_ref()
        .expect("fit_mle requires a mixed model (re: Some)");
    // slope_cols: x column indices for the primary RE slopes (empty = intercept-only)
    let slope_cols: Vec<usize> = re.slopes.iter().map(|&c| c as usize).collect();
    // Extra-grouping slope x-columns, declaration order. On the standalone path the
    // ModelSpec's slope columns ARE x-matrix indices (unlike MCPower, which resolves
    // them separately), so they are read directly here.
    let extra_slope_cols: Vec<Vec<usize>> = re
        .extra_groupings
        .iter()
        .map(|g| g.slopes.iter().map(|&c| c as usize).collect())
        .collect();

    // Build workspace — allocates solver, suff-stats, fit scratch for this model shape
    let mut ws = LmmWorkspace::for_cluster_spec_ext(p, model, n, &slope_cols, &extra_slope_cols);
    // Identity-link offset is an exact y-shift before accumulation (the suff
    // stats are the only place raw y enters) — mirrors `fit_ols` / the sparse
    // `fit_mle_sparse`; change together.
    let y_shifted: Vec<f64>;
    let y_eff: &[f64] = match &opts.offset {
        Some(o) => {
            y_shifted = y.iter().zip(o).map(|(&yi, &oi)| yi - oi).collect();
            &y_shifted
        }
        None => y,
    };
    let x_mat = to_col_major(x, n, p);
    accumulate_lmm_rows(
        &mut ws,
        x_mat.as_ref().subrows(0, n),
        y_eff,
        n,
        p,
        cluster_ids,
        extra_ids,
        opts.weights.as_deref(),
    );

    let ids = crate::GroupIds {
        primary: cluster_ids.to_vec(),
        extra: extra_ids.to_vec(),
    };
    fit_lmm_into(&mut ws, x, &ids, n, p, opts, start)
}
/// Test-only forced-NoZ entry (mirror of forcing `fit_mle_sparse` directly):
/// the NoZ↔Sparse cross-checks and timed sweeps must reach the dense kernel
/// regardless of where `classify_design`'s performance boundary sits — going
/// through `fit_cold` would compare Sparse against itself on any cell the
/// router sends to Sparse. Takes the sizing-corrected spec, like `fit_mle`.
#[cfg(test)]
#[allow(clippy::too_many_arguments)] // marshals the same surface as fit_mle
pub(crate) fn fit_mle_noz_pub(
    x: &[f64],
    y: &[f64],
    n: usize,
    p: usize,
    sized: &ModelSpec,
    cluster_ids: &[u32],
    extra_ids: &[Vec<u32>],
    start: Option<&StartValues>,
    opts: &FitOptions,
) -> Fit {
    fit_mle(x, y, n, p, sized, cluster_ids, extra_ids, start, opts)
}