Skip to main content

gam_models/gamlss/
builders.rs

1// Real concern-organized submodule of the gamlss family stack.
2// Cross-module items are re-exported flat through the parent (`gamlss.rs`),
3// so `use super::*;` makes the sibling-concern symbols this module references
4// resolve through the parent namespace.
5use super::*;
6
7#[derive(Clone, Copy)]
8pub(crate) struct GamlssLambdaLayout {
9    pub(crate) k_mean: usize,
10    pub(crate) k_noise: usize,
11    pub(crate) kwiggle: usize,
12}
13
14impl GamlssLambdaLayout {
15    pub(crate) fn two_block(k_mean: usize, k_noise: usize) -> Self {
16        Self {
17            k_mean,
18            k_noise,
19            kwiggle: 0,
20        }
21    }
22
23    pub(crate) fn withwiggle(k_mean: usize, k_noise: usize, kwiggle: usize) -> Self {
24        Self {
25            k_mean,
26            k_noise,
27            kwiggle,
28        }
29    }
30
31    pub(crate) fn total(self) -> usize {
32        self.k_mean + self.k_noise + self.kwiggle
33    }
34
35    pub(crate) fn noise_start(self) -> usize {
36        self.k_mean
37    }
38
39    pub(crate) fn noise_end(self) -> usize {
40        self.k_mean + self.k_noise
41    }
42
43    pub(crate) fn wiggle_start(self) -> usize {
44        self.k_mean + self.k_noise
45    }
46
47    pub(crate) fn wiggle_end(self) -> usize {
48        self.k_mean + self.k_noise + self.kwiggle
49    }
50
51    pub(crate) fn validate_theta_len(self, theta_len: usize, context: &str) -> Result<(), String> {
52        let needed = self.total();
53        if theta_len < needed {
54            return Err(GamlssError::DimensionMismatch {
55                reason: format!(
56                    "{context} theta too short: got {}, need at least {}",
57                    theta_len, needed
58                ),
59            }
60            .into());
61        }
62        Ok(())
63    }
64
65    pub(crate) fn mean_from(self, theta: &Array1<f64>) -> Array1<f64> {
66        theta.slice(s![0..self.k_mean]).to_owned()
67    }
68
69    pub(crate) fn noise_from(self, theta: &Array1<f64>) -> Array1<f64> {
70        theta
71            .slice(s![self.noise_start()..self.noise_end()])
72            .to_owned()
73    }
74
75    pub(crate) fn wiggle_from(self, theta: &Array1<f64>) -> Array1<f64> {
76        theta
77            .slice(s![self.wiggle_start()..self.wiggle_end()])
78            .to_owned()
79    }
80}
81
82#[derive(Clone, Copy)]
83pub(crate) struct GamlssBetaLayout {
84    pub(crate) pt: usize,
85    pub(crate) pls: usize,
86    pub(crate) pw: usize,
87}
88
89impl GamlssBetaLayout {
90    pub(crate) fn withwiggle(pt: usize, pls: usize, pw: usize) -> Self {
91        Self { pt, pls, pw }
92    }
93
94    pub(crate) fn total(self) -> usize {
95        self.pt + self.pls + self.pw
96    }
97
98    pub(crate) fn split_three(
99        self,
100        flat: &Array1<f64>,
101        context: &str,
102    ) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>), String> {
103        if flat.len() != self.total() {
104            return Err(GamlssError::DimensionMismatch {
105                reason: format!(
106                    "{context} length mismatch: got {}, expected {}",
107                    flat.len(),
108                    self.total()
109                ),
110            }
111            .into());
112        }
113        Ok((
114            flat.slice(s![0..self.pt]).to_owned(),
115            flat.slice(s![self.pt..self.pt + self.pls]).to_owned(),
116            flat.slice(s![self.pt + self.pls..self.total()]).to_owned(),
117        ))
118    }
119}
120
121#[derive(Clone, Debug)]
122pub struct FamilyMetadata {
123    pub name: &'static str,
124    pub parameternames: &'static [&'static str],
125    pub parameter_links: &'static [ParameterLink],
126}
127
128pub(crate) const DEFAULT_GAUGE_PRIORITY: u8 = 100;
129
130pub(crate) const LINK_WIGGLE_GAUGE_PRIORITY: u8 = 80;
131
132pub(crate) fn initial_log_lambdas_orzeros(
133    block: &ParameterBlockInput,
134) -> Result<Array1<f64>, String> {
135    let k = block.penalties.len();
136    let lambdas = block
137        .initial_log_lambdas
138        .clone()
139        .unwrap_or_else(|| Array1::<f64>::zeros(k));
140    if lambdas.len() != k {
141        return Err(GamlssError::DimensionMismatch {
142            reason: format!(
143                "initial_log_lambdas length mismatch: got {}, expected {}",
144                lambdas.len(),
145                k
146            ),
147        }
148        .into());
149    }
150    Ok(lambdas)
151}
152
153pub(crate) fn build_two_block_exact_joint_setup(
154    data: ArrayView2<'_, f64>,
155    meanspec: &TermCollectionSpec,
156    noisespec: &TermCollectionSpec,
157    mean_penalties: usize,
158    noise_penalties: usize,
159    extra_rho0: &[f64],
160    rho0_override: Option<&Array1<f64>>,
161    kappa_options: &SpatialLengthScaleOptimizationOptions,
162) -> ExactJointHyperSetup {
163    // GAMLSS-specific part: assemble the rho seed in [mean | noise | extra]
164    // penalty order, honoring a caller override when it matches the layout.
165    let rho_dim = mean_penalties + noise_penalties + extra_rho0.len();
166    let mut rho0vec = Array1::<f64>::zeros(rho_dim);
167    if let Some(rho0) = rho0_override.filter(|rho0| rho0.len() == rho_dim) {
168        rho0vec.assign(rho0);
169    } else {
170        for (i, &rho_init) in extra_rho0.iter().enumerate() {
171            rho0vec[mean_penalties + noise_penalties + i] = rho_init;
172        }
173    }
174
175    // Generic part: per-block log(kappa) seed/bounds and exact-joint assembly,
176    // with the two linear predictors (mean, noise) in theta order.
177    build_location_scale_exact_joint_setup(data, &[meanspec, noisespec], rho0vec, kappa_options)
178}
179
180pub(crate) fn gaussian_location_scalewarm_start(
181    y: &Array1<f64>,
182    weights: &Array1<f64>,
183    mu_block: &ParameterBlockSpec,
184    log_sigma_block: &ParameterBlockSpec,
185    ridge_floor: f64,
186    mean_beta_hint: Option<&Array1<f64>>,
187    noise_beta_hint: Option<&Array1<f64>>,
188) -> Result<(Array1<f64>, Array1<f64>, f64), String> {
189    let betamu = if let Some(beta) = mean_beta_hint {
190        beta.clone()
191    } else {
192        solve_penalizedweighted_projection(
193            &mu_block.design,
194            &mu_block.offset,
195            y,
196            weights,
197            &mu_block.penalties,
198            &mu_block.initial_log_lambdas,
199            ridge_floor,
200        )?
201    };
202    let mut mu_hat = mu_block.solver_design().matrixvectormultiply(&betamu);
203    mu_hat += mu_block.solver_offset();
204    let mut weighted_ss = 0.0;
205    let mut weight_sum = 0.0;
206    for i in 0..y.len() {
207        let wi = weights[i].max(0.0);
208        let resid = y[i] - mu_hat[i];
209        weighted_ss += wi * resid * resid;
210        weight_sum += wi;
211    }
212    if !weighted_ss.is_finite() || !weight_sum.is_finite() || weight_sum <= 0.0 {
213        return Err(
214            "gaussian location-scale warm start could not estimate residual scale".to_string(),
215        );
216    }
217    // Warm-start σ̂ must clear the logb floor so the inverse link
218    //   η = log(σ − b)
219    // is finite. Use a relative cushion above b so the warm-start is in the
220    // smooth interior of the link domain.
221    let sigma_hat = (weighted_ss / weight_sum)
222        .sqrt()
223        .max(LOGB_SIGMA_FLOOR * 1.5);
224    let beta_log_sigma = if let Some(beta) = noise_beta_hint {
225        beta.clone()
226    } else {
227        let eta_sigma = (sigma_hat - LOGB_SIGMA_FLOOR).ln();
228        let sigma_target = Array1::from_elem(y.len(), eta_sigma);
229        solve_penalizedweighted_projection(
230            &log_sigma_block.design,
231            &log_sigma_block.offset,
232            &sigma_target,
233            weights,
234            &log_sigma_block.penalties,
235            &log_sigma_block.initial_log_lambdas,
236            ridge_floor,
237        )?
238    };
239    Ok((betamu, beta_log_sigma, sigma_hat))
240}
241
242/// Total output count for every two-block location-scale family in this
243/// module (mu/log_sigma or threshold/log_sigma). The wiggle variants add a
244/// third zero-channel block but still drive only two output channels.
245pub(crate) const LOCATION_SCALE_N_OUTPUTS: usize = 2;
246
247/// Construct a fully wired location-scale parameter block.
248///
249/// This is the **only** way to build a LocationScale `ParameterBlockSpec` in
250/// this module — by construction the `AdditiveBlockJacobian` callback is
251/// always installed, so the channel-aware identifiability audit cannot be
252/// silently bypassed by a future `build_blocks` impl that forgets to wire
253/// the callback at the tail (re-introducing #319).
254///
255/// `own_output` is the zero-based output channel this block drives
256/// (e.g. 0 for `mu`/`threshold`, 1 for `log_sigma`). `n_family_outputs` is
257/// fixed at [`LOCATION_SCALE_N_OUTPUTS`] for every two-block family here
258/// but is exposed so the helper composes cleanly with any future
259/// k-block extension.
260pub(crate) fn build_location_scale_block(
261    name: impl Into<String>,
262    design: DesignMatrix,
263    offset: Array1<f64>,
264    penalties: Vec<PenaltyMatrix>,
265    nullspace_dims: Vec<usize>,
266    initial_log_lambdas: Array1<f64>,
267    initial_beta: Option<Array1<f64>>,
268    own_output: usize,
269    n_family_outputs: usize,
270    caller: &str,
271) -> Result<ParameterBlockSpec, String> {
272    if own_output >= n_family_outputs {
273        return Err(format!(
274            "{caller}: own_output={own_output} >= n_family_outputs={n_family_outputs}"
275        ));
276    }
277    let mut spec = ParameterBlockSpec {
278        name: name.into(),
279        design,
280        offset,
281        penalties,
282        nullspace_dims,
283        initial_log_lambdas,
284        initial_beta,
285        gauge_priority: 100,
286        jacobian_callback: None,
287        stacked_design: None,
288        stacked_offset: None,
289    };
290    let dense = spec.effective_design(caller)?;
291    spec.jacobian_callback = Some(std::sync::Arc::new(AdditiveBlockJacobian {
292        design: dense,
293        own_output,
294        n_family_outputs,
295    }));
296    Ok(spec)
297}
298
299/// Construct the wiggle block that accompanies a two-block location-scale
300/// family. The wiggle modulates the inverse link nonlinearly and
301/// contributes no linear effective Jacobian — the installed callback
302/// therefore exposes a zero `(n × p_w)` design under
303/// `n_family_outputs = LOCATION_SCALE_N_OUTPUTS`.
304pub(crate) fn build_location_scale_wiggle_block(
305    name: impl Into<String>,
306    design: DesignMatrix,
307    offset: Array1<f64>,
308    penalties: Vec<PenaltyMatrix>,
309    nullspace_dims: Vec<usize>,
310    initial_log_lambdas: Array1<f64>,
311    initial_beta: Option<Array1<f64>>,
312    n_rows: usize,
313) -> Result<ParameterBlockSpec, String> {
314    let p_w = design.ncols();
315    let mut spec = ParameterBlockSpec {
316        name: name.into(),
317        design,
318        offset,
319        penalties,
320        nullspace_dims,
321        initial_log_lambdas,
322        initial_beta,
323        gauge_priority: 100,
324        jacobian_callback: None,
325        stacked_design: None,
326        stacked_offset: None,
327    };
328    spec.jacobian_callback = Some(std::sync::Arc::new(AdditiveBlockJacobian {
329        design: ndarray::Array2::<f64>::zeros((n_rows, p_w)),
330        own_output: 0,
331        n_family_outputs: LOCATION_SCALE_N_OUTPUTS,
332    }));
333    Ok(spec)
334}
335
336pub(crate) fn prepared_gaussian_log_sigma_design(
337    mu_design: &DesignMatrix,
338    log_sigma_design: &DesignMatrix,
339) -> Result<DesignMatrix, String> {
340    if mu_design.nrows() != log_sigma_design.nrows() {
341        return Err(GamlssError::DimensionMismatch {
342            reason: format!(
343                "gaussian log-sigma design row mismatch: mean rows={}, log_sigma rows={}",
344                mu_design.nrows(),
345                log_sigma_design.nrows()
346            ),
347        }
348        .into());
349    }
350    // Gaussian location-scale remains identifiable even when μ and log σ use
351    // the same covariate basis:
352    //
353    //   L(μ, η) = 0.5 * Σ_i [ (y_i - μ_i)^2 exp(-2η_i) + 2η_i ],
354    //   μ = X_μ β_μ,  η = X_σ β_σ.
355    //
356    // Shared columns are not a frame mismatch. β_μ and β_σ enter through
357    // different sufficient statistics (residual and residual²), so replacing
358    // X_σ with (I - P_{X_μ}) X_σ would impose an extra constraint and can
359    // erase real heteroscedastic signal when the two blocks share a basis.
360    Ok(log_sigma_design.clone())
361}
362
363pub(crate) fn identified_binomial_log_sigma_design(
364    threshold_design: &TermCollectionDesign,
365    log_sigma_design: &TermCollectionDesign,
366    weights: &Array1<f64>,
367) -> Result<DesignMatrix, String> {
368    let non_intercept_start = log_sigma_design
369        .intercept_range
370        .end
371        .min(log_sigma_design.design.ncols());
372    let transform = build_scale_deviation_transform_design(
373        &threshold_design.design,
374        &log_sigma_design.design,
375        weights,
376        non_intercept_start,
377    )?;
378    build_scale_deviation_operator(
379        threshold_design.design.clone(),
380        log_sigma_design.design.clone(),
381        &transform,
382    )
383}
384
385pub(crate) fn identity_penalty(dim: usize) -> Array2<f64> {
386    let mut penalty = Array2::<f64>::zeros((dim, dim));
387    for i in 0..dim {
388        penalty[[i, i]] = 1.0;
389    }
390    penalty
391}
392
393/// Orthogonal projector `P₀ = U₀U₀ᵀ` onto the joint null space of the supplied
394/// penalty blocks over a `dim`-column coefficient space.
395///
396/// Used as the log-σ *shrinkage* penalty for the Gaussian location-scale scale
397/// block. The smooth's own wiggliness penalty already governs its range space
398/// (the curvature directions REML trades off against fit); its null space —
399/// the constant + low-order polynomial log-σ trend that carries the dominant
400/// heteroscedastic signal — is left unpenalized and is only weakly identified
401/// in the coupled (μ, log σ) likelihood, which lets the inner Newton wander
402/// (#1073's "flat/ill-conditioned surface"). A *full-space* identity ridge
403/// fixed that instability but DOUBLE-penalized the range space: REML then drove
404/// the shrinkage λ up, crushing the genuine heteroscedastic curve back to a
405/// constant σ (the underfit this issue reports). Penalizing the null space
406/// ALONE keeps the weakly-identified polynomial trend from blowing up without
407/// touching the wiggliness directions the smooth penalty already controls —
408/// exactly mgcv's `select = TRUE` null-space penalty.
409///
410/// When the supplied penalties already span the whole space (null space empty),
411/// the projector is the zero matrix and the shrinkage term is inert; when there
412/// are no penalties at all (e.g. a purely parametric log-σ design), the null
413/// space is everything and this returns the identity — recovering the previous
414/// full-space ridge exactly where it was the right thing to do.
415pub(crate) fn penalty_nullspace_projector(penalties: &[PenaltyMatrix], dim: usize) -> Array2<f64> {
416    use gam_linalg::faer_ndarray::FaerEigh;
417    use faer::Side;
418
419    if dim == 0 {
420        return Array2::<f64>::zeros((0, 0));
421    }
422    // Combined penalty S = Σ_k S_k over the dim-column scale space. Each block
423    // penalty is already expressed on this space (the scale design's columns).
424    let mut combined = Array2::<f64>::zeros((dim, dim));
425    for pen in penalties {
426        let dense = pen.to_dense();
427        assert_eq!(
428            dense.nrows(),
429            dim,
430            "scale penalty block dim {} != scale design cols {dim}",
431            dense.nrows()
432        );
433        if dense.nrows() == dim && dense.ncols() == dim {
434            combined += &dense;
435        }
436    }
437    // Symmetrize defensively (eigendecomposition assumes self-adjoint input).
438    let combined_sym = 0.5 * (&combined + &combined.t());
439    let (eigvals, eigvecs) = match combined_sym.eigh(Side::Lower) {
440        Ok(decomp) => decomp,
441        // A failed decomposition (degenerate / non-finite) should not silently
442        // drop the stabilizing shrinkage; fall back to the full-space ridge,
443        // which is the conservative (always-positive-definite) choice.
444        Err(_) => return identity_penalty(dim),
445    };
446    // Null space = eigenvectors whose eigenvalue is ≈ 0 relative to the largest.
447    // The combined wiggliness penalty's range-space eigenvalues are O(1) after
448    // basis normalization, so a relative floor cleanly separates the genuine
449    // null directions (constant / low-order polynomial) from the penalized
450    // curvature directions.
451    let max_eig = eigvals.iter().cloned().fold(0.0_f64, f64::max);
452    let tol = (max_eig * 1e-8).max(1e-12);
453    let mut projector = Array2::<f64>::zeros((dim, dim));
454    for (j, &lambda) in eigvals.iter().enumerate() {
455        if lambda <= tol {
456            let v = eigvecs.column(j);
457            // Accumulate v vᵀ into the projector.
458            for a in 0..dim {
459                let va = v[a];
460                if va == 0.0 {
461                    continue;
462                }
463                for b in 0..dim {
464                    projector[[a, b]] += va * v[b];
465                }
466            }
467        }
468    }
469    projector
470}
471
472pub(crate) fn append_binomial_log_sigma_shrinkage_penalty_design(
473    design: &mut TermCollectionDesign,
474) {
475    let p = design.design.ncols();
476    design
477        .penalties
478        .push(BlockwisePenalty::new(0..p, identity_penalty(p)));
479    // Identity penalty penalizes the full space → nullspace dimension is 0.
480    design.nullspace_dims.push(0);
481    design.penaltyinfo.push(PenaltyBlockInfo {
482        global_index: design.penaltyinfo.len(),
483        termname: Some("log_sigma_shrinkage".to_string()),
484        penalty: PenaltyInfo {
485            source: PenaltySource::Other("shrinkage".to_string()),
486            original_index: 0,
487            active: true,
488            effective_rank: p,
489            dropped_reason: None,
490            nullspace_dim_hint: 0,
491            normalization_scale: 1.0,
492            kronecker_factors: None,
493        },
494    });
495}
496
497/// Build the (mean, log-σ) parameter-block pair for a Gaussian location-scale
498/// family. Shared verbatim by the non-wiggle and wiggle Gaussian builders so the
499/// scale-block construction — prepared log-σ design, the REML-selected full-span
500/// shrinkage penalty on the scale nullspace, and the joint Gaussian warm start —
501/// lives in exactly one place. Callers supply the per-block log-λ vectors sliced
502/// from their own layout (two-block vs with-wiggle) and append any extra blocks.
503pub(crate) fn build_gaussian_mean_and_scale_blocks(
504    y: &Array1<f64>,
505    weights: &Array1<f64>,
506    mean_design: &TermCollectionDesign,
507    noise_design: &TermCollectionDesign,
508    mean_offset: &Array1<f64>,
509    noise_offset: &Array1<f64>,
510    mean_log_lambdas: Array1<f64>,
511    noise_log_lambdas: Array1<f64>,
512    mean_beta_hint: Option<Array1<f64>>,
513    noise_beta_hint: Option<Array1<f64>>,
514    context: &str,
515) -> Result<(ParameterBlockSpec, ParameterBlockSpec), String> {
516    let mut meanspec = build_location_scale_block(
517        "mu",
518        mean_design.design.clone(),
519        mean_offset.clone(),
520        mean_design.penalties_as_penalty_matrix(),
521        mean_design.nullspace_dims.clone(),
522        mean_log_lambdas,
523        mean_beta_hint,
524        0,
525        LOCATION_SCALE_N_OUTPUTS,
526        &format!("{context}: mu"),
527    )?;
528    let prepared_noise_design =
529        prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design)?;
530    let p_noise = prepared_noise_design.ncols();
531    let mut log_sigma_penalty_matrices = noise_design.penalties_as_penalty_matrix();
532    // Shrinkage penalty on the scale block's *null space only* (mgcv
533    // `select = TRUE`): it stabilizes the weakly-identified constant/polynomial
534    // log-σ trend without double-penalizing the wiggliness directions the
535    // smooth penalty already governs. A full-space identity here over-shrinks
536    // the genuine heteroscedastic curve back to a constant σ (#1073).
537    let shrinkage = penalty_nullspace_projector(&log_sigma_penalty_matrices, p_noise);
538    // The rank of an orthogonal projector equals its trace (P = P² for a projector,
539    // so trace(P) = trace(P²) = ||P||_F² = sum of squared singular values = rank).
540    // The diagonal-threshold test `diag[i] > 0.5` used previously was wrong: for a
541    // rank-d projector onto a low-dimensional subspace (e.g. d=2 null directions of
542    // a TP spline with p=10 columns), each diagonal entry is O(d/p) << 0.5, so the
543    // threshold always returned 0 — misreporting the shrinkage penalty as having
544    // zero penalized dimensions. Trace-based rank is exact for a symmetric
545    // idempotent matrix (rounded to the nearest integer to absorb floating-point
546    // rounding in the eigendecomposition).
547    let shrinkage_rank = (0..p_noise).map(|i| shrinkage[[i, i]]).sum::<f64>().round() as usize;
548    log_sigma_penalty_matrices.push(PenaltyMatrix::Dense(shrinkage));
549    let mut log_sigma_nullspace_dims = noise_design.nullspace_dims.clone();
550    // The null-space projector penalizes a rank-`shrinkage_rank` subspace, so
551    // the remaining unpenalized directions number `p_noise − shrinkage_rank`.
552    log_sigma_nullspace_dims.push(p_noise.saturating_sub(shrinkage_rank));
553    let mut noisespec = build_location_scale_block(
554        "log_sigma",
555        prepared_noise_design,
556        noise_offset.clone(),
557        log_sigma_penalty_matrices,
558        log_sigma_nullspace_dims,
559        noise_log_lambdas,
560        noise_beta_hint,
561        1,
562        LOCATION_SCALE_N_OUTPUTS,
563        &format!("{context}: log_sigma"),
564    )?;
565    if meanspec.initial_beta.is_none() || noisespec.initial_beta.is_none() {
566        let (betamu0, beta_ls0, _) = gaussian_location_scalewarm_start(
567            y,
568            weights,
569            &meanspec,
570            &noisespec,
571            1e-10,
572            meanspec.initial_beta.as_ref(),
573            noisespec.initial_beta.as_ref(),
574        )?;
575        if meanspec.initial_beta.is_none() {
576            meanspec.initial_beta = Some(betamu0);
577        }
578        if noisespec.initial_beta.is_none() {
579            noisespec.initial_beta = Some(beta_ls0);
580        }
581    }
582    Ok((meanspec, noisespec))
583}
584
585/// Build the (threshold, log-σ) parameter-block pair for a Binomial
586/// location-scale family. Shared by the non-wiggle and wiggle Binomial builders;
587/// mirrors [`build_gaussian_mean_and_scale_blocks`] but with the binomial-
588/// identified log-σ design, the link-aware joint warm start, and the same
589/// REML-selected full-span scale shrinkage penalty.
590pub(crate) fn build_binomial_threshold_and_scale_blocks(
591    y: &Array1<f64>,
592    weights: &Array1<f64>,
593    link_kind: &InverseLink,
594    mean_design: &TermCollectionDesign,
595    noise_design: &TermCollectionDesign,
596    mean_offset: &Array1<f64>,
597    noise_offset: &Array1<f64>,
598    mean_log_lambdas: Array1<f64>,
599    noise_log_lambdas: Array1<f64>,
600    mean_beta_hint: Option<Array1<f64>>,
601    noise_beta_hint: Option<Array1<f64>>,
602    context: &str,
603) -> Result<(ParameterBlockSpec, ParameterBlockSpec), String> {
604    let identifiednoise_design =
605        identified_binomial_log_sigma_design(mean_design, noise_design, weights)?;
606    let p_noise = identifiednoise_design.ncols();
607    let mut log_sigma_penalty_matrices: Vec<PenaltyMatrix> =
608        noise_design.penalties_as_penalty_matrix();
609    log_sigma_penalty_matrices.push(PenaltyMatrix::Dense(identity_penalty(p_noise)));
610    let mut thresholdspec = build_location_scale_block(
611        "threshold",
612        mean_design.design.clone(),
613        mean_offset.clone(),
614        mean_design.penalties_as_penalty_matrix(),
615        vec![],
616        mean_log_lambdas,
617        mean_beta_hint,
618        0,
619        LOCATION_SCALE_N_OUTPUTS,
620        &format!("{context}: threshold"),
621    )?;
622    let mut log_sigmaspec = build_location_scale_block(
623        "log_sigma",
624        identifiednoise_design,
625        noise_offset.clone(),
626        log_sigma_penalty_matrices,
627        vec![],
628        noise_log_lambdas,
629        noise_beta_hint,
630        1,
631        LOCATION_SCALE_N_OUTPUTS,
632        &format!("{context}: log_sigma"),
633    )?;
634    if thresholdspec.initial_beta.is_none() || log_sigmaspec.initial_beta.is_none() {
635        let (beta_t0, beta_ls0) = binomial_location_scalewarm_start(
636            y,
637            weights,
638            link_kind,
639            &thresholdspec,
640            &log_sigmaspec,
641            thresholdspec.initial_beta.as_ref(),
642            log_sigmaspec.initial_beta.as_ref(),
643        )?;
644        if thresholdspec.initial_beta.is_none() {
645            thresholdspec.initial_beta = Some(beta_t0);
646        }
647        if log_sigmaspec.initial_beta.is_none() {
648            log_sigmaspec.initial_beta = Some(beta_ls0);
649        }
650    }
651    Ok((thresholdspec, log_sigmaspec))
652}
653
654/// Convert a wiggle block's `PenaltySpec`s into the `PenaltyMatrix` list the
655/// location-scale wiggle block expects. Shared by the Gaussian and Binomial
656/// wiggle builders, which previously inlined the identical match.
657pub(crate) fn wiggle_block_penalty_matrices(
658    wiggle_block: &ParameterBlockInput,
659) -> Vec<PenaltyMatrix> {
660    let p_wiggle = wiggle_block.design.ncols();
661    wiggle_block
662        .penalties
663        .iter()
664        .map(|spec| match spec {
665            crate::model_types::PenaltySpec::Block {
666                local, col_range, ..
667            } => PenaltyMatrix::Blockwise {
668                local: local.clone(),
669                col_range: col_range.clone(),
670                total_dim: p_wiggle,
671            },
672            crate::model_types::PenaltySpec::Dense(m)
673            | crate::model_types::PenaltySpec::DenseWithMean { matrix: m, .. } => {
674                PenaltyMatrix::Dense(m.clone())
675            }
676        })
677        .collect()
678}
679
680pub(crate) fn binomial_location_scale_link_eta_from_probability(
681    link_kind: &InverseLink,
682    probability: f64,
683) -> Result<f64, String> {
684    let target = probability.clamp(1e-6, 1.0 - 1e-6);
685    match link_kind {
686        InverseLink::Standard(StandardLink::Logit) => Ok((target / (1.0 - target)).ln()),
687        InverseLink::Standard(StandardLink::Probit) => standard_normal_quantile(target)
688            .map_err(|err| format!("failed to invert probit warm-start probability: {err}")),
689        InverseLink::Standard(StandardLink::CLogLog) => Ok((-((1.0 - target).ln())).ln()),
690        other => Err(GamlssError::UnsupportedConfiguration { reason: format!(
691            "binomial location-scale warm start requires logit, probit, or cloglog link, got {other:?}"
692        ) }.into()),
693    }
694}
695
696pub(crate) fn weighted_binomial_prevalence(
697    y: &Array1<f64>,
698    weights: &Array1<f64>,
699) -> Result<f64, String> {
700    if y.len() != weights.len() {
701        return Err(GamlssError::DimensionMismatch { reason: format!(
702            "binomial location-scale warm start dimension mismatch: y has length {}, weights have length {}",
703            y.len(),
704            weights.len()
705        ) }.into());
706    }
707    let mut weight_sum = 0.0;
708    let mut success_sum = 0.0;
709    for (&yi, &wi) in y.iter().zip(weights.iter()) {
710        if !yi.is_finite() {
711            return Err(GamlssError::NonFinite {
712                reason: format!(
713                    "binomial location-scale warm start encountered non-finite response {yi}"
714                ),
715            }
716            .into());
717        }
718        let weight = floor_positiveweight(wi, MIN_WEIGHT);
719        if weight > 0.0 {
720            weight_sum += weight;
721            success_sum += weight * yi;
722        }
723    }
724    if !weight_sum.is_finite() || weight_sum <= 0.0 {
725        return Err(
726            "binomial location-scale warm start requires positive total weight".to_string(),
727        );
728    }
729    Ok(success_sum / weight_sum)
730}
731
732pub(crate) fn project_constant_eta_into_block(
733    block: &ParameterBlockSpec,
734    weights: &Array1<f64>,
735    eta: f64,
736) -> Result<Array1<f64>, String> {
737    let target_eta = Array1::from_elem(block.design.nrows(), eta);
738    solve_penalizedweighted_projection(
739        &block.design,
740        &block.offset,
741        &target_eta,
742        weights,
743        &block.penalties,
744        &block.initial_log_lambdas,
745        1e-10,
746    )
747}
748
749// Deterministic warm start for the binomial location-scale model. This stays
750// out of the optimizer: it projects a prevalence-matched threshold and neutral
751// log-sigma value into the actual penalized block spaces.
752pub(crate) fn binomial_location_scalewarm_start(
753    y: &Array1<f64>,
754    weights: &Array1<f64>,
755    link_kind: &InverseLink,
756    threshold_block: &ParameterBlockSpec,
757    log_sigma_block: &ParameterBlockSpec,
758    mean_beta_hint: Option<&Array1<f64>>,
759    noise_beta_hint: Option<&Array1<f64>>,
760) -> Result<(Array1<f64>, Array1<f64>), String> {
761    if let (Some(mean_beta), Some(noise_beta)) = (mean_beta_hint, noise_beta_hint) {
762        return Ok((mean_beta.clone(), noise_beta.clone()));
763    }
764
765    let beta_threshold = match mean_beta_hint {
766        Some(beta) => beta.clone(),
767        None => {
768            let prevalence = weighted_binomial_prevalence(y, weights)?;
769            let eta = binomial_location_scale_link_eta_from_probability(link_kind, prevalence)?;
770            project_constant_eta_into_block(threshold_block, weights, eta)?
771        }
772    };
773    let beta_log_sigma = match noise_beta_hint {
774        Some(beta) => beta.clone(),
775        None => project_constant_eta_into_block(log_sigma_block, weights, 0.0)?,
776    };
777    Ok((beta_threshold, beta_log_sigma))
778}
779
780#[derive(Clone)]
781pub(crate) struct BinomialMeanWiggleSpec {
782    pub y: Array1<f64>,
783    pub weights: Array1<f64>,
784    pub link_kind: InverseLink,
785    pub wiggle_knots: Array1<f64>,
786    pub wiggle_degree: usize,
787    pub eta_block: ParameterBlockInput,
788    pub wiggle_block: ParameterBlockInput,
789}
790
791#[derive(Clone)]
792pub struct GaussianLocationScaleTermSpec {
793    pub y: Array1<f64>,
794    pub weights: Array1<f64>,
795    pub meanspec: TermCollectionSpec,
796    pub log_sigmaspec: TermCollectionSpec,
797    pub mean_offset: Array1<f64>,
798    pub log_sigma_offset: Array1<f64>,
799}
800
801#[derive(Clone)]
802pub struct GaussianLocationScaleWiggleTermSpec {
803    pub y: Array1<f64>,
804    pub weights: Array1<f64>,
805    pub meanspec: TermCollectionSpec,
806    pub log_sigmaspec: TermCollectionSpec,
807    pub mean_offset: Array1<f64>,
808    pub log_sigma_offset: Array1<f64>,
809    pub wiggle_knots: Array1<f64>,
810    pub wiggle_degree: usize,
811    pub wiggle_block: ParameterBlockInput,
812}
813
814#[derive(Clone)]
815pub struct BinomialLocationScaleTermSpec {
816    pub y: Array1<f64>,
817    pub weights: Array1<f64>,
818    pub link_kind: InverseLink,
819    pub thresholdspec: TermCollectionSpec,
820    pub log_sigmaspec: TermCollectionSpec,
821    pub threshold_offset: Array1<f64>,
822    pub log_sigma_offset: Array1<f64>,
823}
824
825#[derive(Clone)]
826pub struct BinomialLocationScaleWiggleTermSpec {
827    pub y: Array1<f64>,
828    pub weights: Array1<f64>,
829    pub link_kind: InverseLink,
830    pub thresholdspec: TermCollectionSpec,
831    pub log_sigmaspec: TermCollectionSpec,
832    pub threshold_offset: Array1<f64>,
833    pub log_sigma_offset: Array1<f64>,
834    pub wiggle_knots: Array1<f64>,
835    pub wiggle_degree: usize,
836    pub wiggle_block: ParameterBlockInput,
837}
838
839#[derive(Clone, Debug)]
840pub struct BlockwiseTermFitResult {
841    pub fit: UnifiedFitResult,
842    pub meanspec_resolved: TermCollectionSpec,
843    pub noisespec_resolved: TermCollectionSpec,
844    pub mean_design: TermCollectionDesign,
845    pub noise_design: TermCollectionDesign,
846}
847
848pub(crate) struct BlockwiseTermFitResultParts {
849    pub fit: UnifiedFitResult,
850    pub meanspec_resolved: TermCollectionSpec,
851    pub noisespec_resolved: TermCollectionSpec,
852    pub mean_design: TermCollectionDesign,
853    pub noise_design: TermCollectionDesign,
854}
855
856pub struct BlockwiseTermWiggleFitResult {
857    pub fit: BlockwiseTermFitResult,
858    pub wiggle_knots: Array1<f64>,
859    pub wiggle_degree: usize,
860}
861
862pub struct BinomialMeanWiggleTermFitResult {
863    pub fit: UnifiedFitResult,
864    pub resolvedspec: TermCollectionSpec,
865    pub design: TermCollectionDesign,
866    pub wiggle_knots: Array1<f64>,
867    pub wiggle_degree: usize,
868}
869
870pub(crate) struct BlockwiseTermWiggleFitResultParts {
871    pub fit: BlockwiseTermFitResult,
872    pub wiggle_knots: Array1<f64>,
873    pub wiggle_degree: usize,
874}
875
876pub(crate) fn validate_term_collection_design(
877    label: &str,
878    design: &TermCollectionDesign,
879) -> Result<(), String> {
880    let p = design.design.ncols();
881    let n = design.design.nrows();
882    for rows in exact_design_row_chunks(n, p) {
883        let chunk = design
884            .design
885            .try_row_chunk(rows)
886            .map_err(|e| format!("{label}.design row chunk materialization failed: {e}"))?;
887        validate_all_finite_estimation(&format!("{label}.design"), chunk.iter().copied())
888            .map_err(|e| e.to_string())?;
889    }
890    if design.nullspace_dims.len() != design.penalties.len() {
891        return Err(GamlssError::DimensionMismatch {
892            reason: format!(
893                "{label}.nullspace_dims length mismatch: got {}, expected {}",
894                design.nullspace_dims.len(),
895                design.penalties.len()
896            ),
897        }
898        .into());
899    }
900    if design.penaltyinfo.len() != design.penalties.len() {
901        return Err(GamlssError::DimensionMismatch {
902            reason: format!(
903                "{label}.penaltyinfo length mismatch: got {}, expected {}",
904                design.penaltyinfo.len(),
905                design.penalties.len()
906            ),
907        }
908        .into());
909    }
910    for (idx, bp) in design.penalties.iter().enumerate() {
911        validate_all_finite_estimation(
912            &format!("{label}.penalties[{idx}]"),
913            bp.local.iter().copied(),
914        )
915        .map_err(|e| e.to_string())?;
916        if bp.col_range.end > p {
917            return Err(GamlssError::DimensionMismatch {
918                reason: format!(
919                    "{label}.penalties[{idx}] col_range {}..{} exceeds design width {}",
920                    bp.col_range.start, bp.col_range.end, p
921                ),
922            }
923            .into());
924        }
925    }
926    if let Some(bounds) = design.coefficient_lower_bounds.as_ref() {
927        if bounds.len() != p {
928            return Err(GamlssError::ConstraintViolation {
929                reason: format!(
930                    "{label}.coefficient_lower_bounds length mismatch: got {}, expected {p}",
931                    bounds.len()
932                ),
933            }
934            .into());
935        }
936        for (idx, &bound) in bounds.iter().enumerate() {
937            if !(bound.is_finite() || bound == f64::NEG_INFINITY) {
938                return Err(GamlssError::NonFinite { reason: format!(
939                    "{label}.coefficient_lower_bounds[{idx}] must be finite or -inf, got {bound}",
940                ) }.into());
941            }
942        }
943    }
944    if let Some(constraints) = design.linear_constraints.as_ref() {
945        validate_all_finite_estimation(
946            &format!("{label}.linear_constraints.a"),
947            constraints.a.iter().copied(),
948        )
949        .map_err(|e| e.to_string())?;
950        validate_all_finite_estimation(
951            &format!("{label}.linear_constraints.b"),
952            constraints.b.iter().copied(),
953        )
954        .map_err(|e| e.to_string())?;
955        if constraints.a.ncols() != p {
956            return Err(GamlssError::DimensionMismatch {
957                reason: format!(
958                    "{label}.linear_constraints.a column mismatch: got {}, expected {p}",
959                    constraints.a.ncols()
960                ),
961            }
962            .into());
963        }
964        if constraints.a.nrows() != constraints.b.len() {
965            return Err(GamlssError::DimensionMismatch {
966                reason: format!(
967                    "{label}.linear_constraints row mismatch: a has {}, b has {}",
968                    constraints.a.nrows(),
969                    constraints.b.len()
970                ),
971            }
972            .into());
973        }
974    }
975    if design.intercept_range.start > design.intercept_range.end || design.intercept_range.end > p {
976        return Err(GamlssError::ConstraintViolation {
977            reason: format!(
978                "{label}.intercept_range out of bounds: {:?} for {} columns",
979                design.intercept_range, p
980            ),
981        }
982        .into());
983    }
984    Ok(())
985}
986
987impl BlockwiseTermFitResult {
988    pub(crate) fn try_from_parts(parts: BlockwiseTermFitResultParts) -> Result<Self, String> {
989        let BlockwiseTermFitResultParts {
990            fit,
991            meanspec_resolved,
992            noisespec_resolved,
993            mean_design,
994            noise_design,
995        } = parts;
996
997        fit.validate_numeric_finiteness()
998            .map_err(|e| format!("{e}"))?;
999        if fit.block_states.len() < 2 {
1000            return Err(GamlssError::DimensionMismatch {
1001                reason: format!(
1002                    "BlockwiseTermFitResult requires at least 2 block states, got {}",
1003                    fit.block_states.len()
1004                ),
1005            }
1006            .into());
1007        }
1008        validate_term_collection_design("blockwise_term.mean_design", &mean_design)?;
1009        validate_term_collection_design("blockwise_term.noise_design", &noise_design)?;
1010        if mean_design.design.nrows() != noise_design.design.nrows() {
1011            return Err(GamlssError::DimensionMismatch {
1012                reason: format!(
1013                    "BlockwiseTermFitResult row mismatch: mean_design={}, noise_design={}",
1014                    mean_design.design.nrows(),
1015                    noise_design.design.nrows()
1016                ),
1017            }
1018            .into());
1019        }
1020        if fit.block_states[0].beta.len() != mean_design.design.ncols() {
1021            return Err(GamlssError::DimensionMismatch {
1022                reason: format!(
1023                    "BlockwiseTermFitResult mean beta length mismatch: got {}, expected {}",
1024                    fit.block_states[0].beta.len(),
1025                    mean_design.design.ncols()
1026                ),
1027            }
1028            .into());
1029        }
1030        if fit.block_states[1].beta.len() != noise_design.design.ncols() {
1031            return Err(GamlssError::DimensionMismatch {
1032                reason: format!(
1033                    "BlockwiseTermFitResult noise beta length mismatch: got {}, expected {}",
1034                    fit.block_states[1].beta.len(),
1035                    noise_design.design.ncols()
1036                ),
1037            }
1038            .into());
1039        }
1040        if fit.block_states[0].eta.len() != mean_design.design.nrows() {
1041            return Err(GamlssError::DimensionMismatch {
1042                reason: format!(
1043                    "BlockwiseTermFitResult mean eta length mismatch: got {}, expected {}",
1044                    fit.block_states[0].eta.len(),
1045                    mean_design.design.nrows()
1046                ),
1047            }
1048            .into());
1049        }
1050        if fit.block_states[1].eta.len() != noise_design.design.nrows() {
1051            return Err(GamlssError::DimensionMismatch {
1052                reason: format!(
1053                    "BlockwiseTermFitResult noise eta length mismatch: got {}, expected {}",
1054                    fit.block_states[1].eta.len(),
1055                    noise_design.design.nrows()
1056                ),
1057            }
1058            .into());
1059        }
1060
1061        Ok(Self {
1062            fit,
1063            meanspec_resolved,
1064            noisespec_resolved,
1065            mean_design,
1066            noise_design,
1067        })
1068    }
1069
1070    pub(crate) fn validate_numeric_finiteness(&self) -> Result<(), String> {
1071        Self::try_from_parts(BlockwiseTermFitResultParts {
1072            fit: self.fit.clone(),
1073            meanspec_resolved: self.meanspec_resolved.clone(),
1074            noisespec_resolved: self.noisespec_resolved.clone(),
1075            mean_design: self.mean_design.clone(),
1076            noise_design: self.noise_design.clone(),
1077        })
1078        .map(|_| ())
1079    }
1080}
1081
1082impl BlockwiseTermWiggleFitResult {
1083    pub(crate) fn try_from_parts(parts: BlockwiseTermWiggleFitResultParts) -> Result<Self, String> {
1084        let BlockwiseTermWiggleFitResultParts {
1085            fit,
1086            wiggle_knots,
1087            wiggle_degree,
1088        } = parts;
1089
1090        fit.validate_numeric_finiteness()
1091            .map_err(|e| e.to_string())?;
1092        if fit.fit.block_states.len() < 3 {
1093            return Err(GamlssError::DimensionMismatch {
1094                reason: format!(
1095                    "BlockwiseTermWiggleFitResult requires at least 3 block states, got {}",
1096                    fit.fit.block_states.len()
1097                ),
1098            }
1099            .into());
1100        }
1101        if wiggle_knots.is_empty() {
1102            return Err(GamlssError::UnsupportedConfiguration {
1103                reason: "BlockwiseTermWiggleFitResult requires non-empty wiggle_knots".to_string(),
1104            }
1105            .into());
1106        }
1107        validate_all_finite_estimation(
1108            "blockwise_term_wiggle.wiggle_knots",
1109            wiggle_knots.iter().copied(),
1110        )
1111        .map_err(|e| e.to_string())?;
1112
1113        Ok(Self {
1114            fit,
1115            wiggle_knots,
1116            wiggle_degree,
1117        })
1118    }
1119}
1120
1121pub struct BinomialLocationScaleFitResult {
1122    pub fit: BlockwiseTermFitResult,
1123    pub wiggle_knots: Option<Array1<f64>>,
1124    pub wiggle_degree: Option<usize>,
1125    pub beta_link_wiggle: Option<Vec<f64>>,
1126}
1127
1128pub struct GaussianLocationScaleFitResult {
1129    pub fit: BlockwiseTermFitResult,
1130    pub wiggle_knots: Option<Array1<f64>>,
1131    pub wiggle_degree: Option<usize>,
1132    pub beta_link_wiggle: Option<Vec<f64>>,
1133    /// Response standardization factor applied internally during fitting.
1134    ///
1135    /// The Gaussian location-scale path fits on `y / response_scale` so the
1136    /// fixed log-σ soft floor `LOGB_SIGMA_FLOOR = 0.01` is *operationally*
1137    /// scale-relative (1 % of the response spread) rather than absolute,
1138    /// keeping κ = dlogσ/dη ≈ 1 across the realistic σ range and informing the
1139    /// scale block like gamlss. The returned coefficient `blocks`, `beta`, and
1140    /// link-wiggle knots/coefficients are already mapped back to **raw response
1141    /// units** (the Location/Mean block scaled by `response_scale`, the Scale
1142    /// block intercept shifted by `+ln(response_scale)`), so downstream
1143    /// reconstruction `μ = X_mean·β` comes out in raw units with no further
1144    /// rescaling.
1145    ///
1146    /// The σ reconstruction, however, **must scale the floor too** to stay
1147    /// response-scale-equivariant (#884):
1148    ///
1149    /// ```text
1150    /// σ = response_scale·LOGB_SIGMA_FLOOR + exp(X_scale·β)
1151    ///   = response_scale·(LOGB_SIGMA_FLOOR + exp(η_internal)).
1152    /// ```
1153    ///
1154    /// The intercept shift carries only the `exp(η)` term; reconstructing with a
1155    /// raw `LOGB_SIGMA_FLOOR` instead of `response_scale·LOGB_SIGMA_FLOOR` leaves
1156    /// the non-equivariant residual `LOGB_SIGMA_FLOOR·(1 − response_scale)`.
1157    ///
1158    /// This field records the factor that was applied for transparency,
1159    /// covariance bookkeeping, and the equivariant σ-floor reconstruction; it is
1160    /// `1.0` when no standardization was needed (degenerate constant response).
1161    pub response_scale: f64,
1162}
1163
1164pub(crate) fn fit_binomial_mean_wiggle(
1165    spec: BinomialMeanWiggleSpec,
1166    options: &BlockwiseFitOptions,
1167) -> Result<UnifiedFitResult, String> {
1168    let n = spec.y.len();
1169    validate_len_match("weights vs y", n, spec.weights.len())?;
1170    validateweights(&spec.weights, "fit_binomial_mean_wiggle")?;
1171    validate_binomial_response(&spec.y, "fit_binomial_mean_wiggle")?;
1172    validate_blockrows("eta", n, &spec.eta_block)?;
1173    validate_blockrows("wiggle", n, &spec.wiggle_block)?;
1174    if matches!(
1175        spec.link_kind,
1176        InverseLink::Standard(StandardLink::Identity)
1177    ) {
1178        return Err(GamlssError::UnsupportedConfiguration {
1179            reason: "fit_binomial_mean_wiggle does not support identity link".to_string(),
1180        }
1181        .into());
1182    }
1183    gam_terms::inference::formula_dsl::require_binomial_inverse_link_supports_joint_wiggle(
1184        &spec.link_kind,
1185        "fit_binomial_mean_wiggle",
1186    )?;
1187    if spec.wiggle_degree < 2 {
1188        return Err(GamlssError::ConstraintViolation {
1189            reason: format!(
1190                "fit_binomial_mean_wiggle: wiggle_degree must be >= 2, got {}",
1191                spec.wiggle_degree
1192            ),
1193        }
1194        .into());
1195    }
1196    let minimum_knots = minimum_monotone_wiggle_knot_count(spec.wiggle_degree)?;
1197    if spec.wiggle_knots.len() < minimum_knots {
1198        return Err(GamlssError::DimensionMismatch { reason: format!(
1199            "fit_binomial_mean_wiggle: wiggle_knots length {} is too short for degree {} (need at least {})",
1200            spec.wiggle_knots.len(),
1201            spec.wiggle_degree,
1202            minimum_knots
1203        ) }.into());
1204    }
1205
1206    let family = BinomialMeanWiggleFamily {
1207        y: spec.y,
1208        weights: spec.weights,
1209        link_kind: spec.link_kind,
1210        wiggle_knots: spec.wiggle_knots,
1211        wiggle_degree: spec.wiggle_degree,
1212        policy: gam_runtime::resource::ResourcePolicy::default_library(),
1213    };
1214    let blocks = vec![
1215        // The wiggle block is a DYNAMIC monotone I-spline basis that the
1216        // family regenerates at full (raw) width every inner iteration
1217        // (`block_geometry_is_dynamic` + the `x.ncols() == spec.design.ncols()`
1218        // assertion in `block_geometry`), so it cannot tolerate a physical
1219        // column drop. The level/intercept direction that the I-spline shares
1220        // with the eta block must therefore be yielded by the *eta* block,
1221        // whose static term-collection design is safely column-reducible (and
1222        // lifted back via the canonical per-block transform `T`). Give the eta
1223        // block the lower gauge priority so the canonical-gauge RRQR routes the
1224        // shared-level alias drop onto eta and leaves the dynamic wiggle basis
1225        // full-width.
1226        spec.eta_block
1227            .intospec_with_gauge_priority("eta", LINK_WIGGLE_GAUGE_PRIORITY)?,
1228        spec.wiggle_block.intospec("wiggle")?,
1229    ];
1230    fit_custom_family(&family, &blocks, options).map_err(|e| e.to_string())
1231}
1232
1233pub(crate) trait LocationScaleFamilyBuilder {
1234    type Family: CustomFamily + Clone + Send + Sync + 'static;
1235
1236    fn meanspec(&self) -> &TermCollectionSpec;
1237    fn noisespec(&self) -> &TermCollectionSpec;
1238
1239    fn build_blocks(
1240        &self,
1241        theta: &Array1<f64>,
1242        mean_design: &TermCollectionDesign,
1243        noise_design: &TermCollectionDesign,
1244        mean_beta_hint: Option<Array1<f64>>,
1245        noise_beta_hint: Option<Array1<f64>>,
1246    ) -> Result<Vec<ParameterBlockSpec>, String>;
1247
1248    fn build_family(
1249        &self,
1250        mean_design: &TermCollectionDesign,
1251        noise_design: &TermCollectionDesign,
1252    ) -> Self::Family;
1253
1254    fn extract_primary_betas(
1255        &self,
1256        fit: &UnifiedFitResult,
1257    ) -> Result<(Array1<f64>, Array1<f64>), String>;
1258
1259    fn mean_penalty_count(&self, mean_design: &TermCollectionDesign) -> usize {
1260        mean_design.penalties.len()
1261    }
1262
1263    fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
1264        noise_design.penalties.len()
1265    }
1266
1267    fn exact_spatial_joint_supported(&self) -> bool {
1268        false
1269    }
1270
1271    fn require_exact_spatial_joint(&self) -> bool {
1272        false
1273    }
1274
1275    fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
1276        crate::seeding::SeedRiskProfile::GeneralizedLinear
1277    }
1278
1279    fn extra_rho0(&self) -> Result<Array1<f64>, String> {
1280        Ok(Array1::zeros(0))
1281    }
1282
1283    fn build_psiderivative_blocks(
1284        &self,
1285        arr: ndarray::ArrayView2<'_, f64>,
1286        term_spec: &TermCollectionSpec,
1287        term_spec2: &TermCollectionSpec,
1288        term_design: &TermCollectionDesign,
1289        term_design2: &TermCollectionDesign,
1290    ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String>;
1291}
1292
1293pub(crate) fn fit_location_scale_terms<B: LocationScaleFamilyBuilder>(
1294    data: ndarray::ArrayView2<'_, f64>,
1295    builder: B,
1296    options: &BlockwiseFitOptions,
1297    kappa_options: &SpatialLengthScaleOptimizationOptions,
1298) -> Result<BlockwiseTermFitResult, String> {
1299    // Large-n location-scale fits keep the caller's explicit Hessian request.
1300    // The unified REML evaluator chooses a dense or matrix-free exact
1301    // representation from the realized (n, p, K) work model, so there is no
1302    // large-scale downgrade to BFGS here.
1303
1304    let mut mean_beta_hint: Option<Array1<f64>> = None;
1305    let mut noise_beta_hint: Option<Array1<f64>> = None;
1306    let extra_rho0 = builder.extra_rho0()?;
1307
1308    let mean_boot_design =
1309        build_term_collection_design(data, builder.meanspec()).map_err(|e| e.to_string())?;
1310    let noise_boot_design =
1311        build_term_collection_design(data, builder.noisespec()).map_err(|e| e.to_string())?;
1312    let mean_bootspec = freeze_term_collection_from_design(builder.meanspec(), &mean_boot_design)
1313        .map_err(|e| e.to_string())?;
1314    let noise_bootspec =
1315        freeze_term_collection_from_design(builder.noisespec(), &noise_boot_design)
1316            .map_err(|e| e.to_string())?;
1317
1318    let require_exact_spatial_joint = builder.require_exact_spatial_joint();
1319    let analytic_joint_derivatives_check = if builder.exact_spatial_joint_supported() {
1320        builder
1321            .build_psiderivative_blocks(
1322                data,
1323                &mean_bootspec,
1324                &noise_bootspec,
1325                &mean_boot_design,
1326                &noise_boot_design,
1327            )
1328            .map(|_| ())
1329    } else {
1330        Err(
1331            "analytic spatial psi derivatives are unavailable for this location-scale family"
1332                .to_string(),
1333        )
1334    };
1335    let analytic_joint_derivatives_available = analytic_joint_derivatives_check.is_ok();
1336    if require_exact_spatial_joint {
1337        analytic_joint_derivatives_check.map_err(|err| {
1338            format!("exact two-block spatial path requires analytic psi derivatives: {err}")
1339        })?;
1340    }
1341    let mean_penalty_count = builder.mean_penalty_count(&mean_boot_design);
1342    let noise_penalty_count = builder.noise_penalty_count(&noise_boot_design);
1343
1344    // Honor an explicit user-supplied `length_scale=X` on every spatial term
1345    // in both the mean and noise blocks: when every term is κ-locked (no
1346    // anisotropy, no per-axis ψ contrasts), the joint-spatial outer optimizer
1347    // has nothing to optimize. Routing through it anyway wraps the full
1348    // two-block coefficient solve inside an unnecessary outer loop where
1349    // each evaluation runs the inner Newton from scratch. This is the same
1350    // short-circuit the Bernoulli marginal-slope entry point performs at
1351    // bernoulli_marginal_slope.rs:16432-16442; mirroring it here makes the
1352    // GAMLSS path skip straight to the `(!enabled || log_kappa_dim == 0)`
1353    // fast path in `optimize_spatial_length_scale_exact_joint`.
1354    let mut effective_kappa_options = kappa_options.clone();
1355    if effective_kappa_options.enabled
1356        && gam_terms::smooth::all_spatial_terms_kappa_fixed(&mean_bootspec)
1357        && gam_terms::smooth::all_spatial_terms_kappa_fixed(&noise_bootspec)
1358    {
1359        log::info!(
1360            "[GAMLSS spatial] disabling κ/ψ optimization: every spatial term in \
1361             both blocks has an explicit length_scale and no anisotropy; \
1362             user-supplied kernel scale is fixed"
1363        );
1364        effective_kappa_options.enabled = false;
1365    }
1366    let kappa_options: &SpatialLengthScaleOptimizationOptions = &effective_kappa_options;
1367
1368    // Macro to invoke the exact-joint spatial optimizer with shared closures.
1369    // The exact path evaluates the full profiled/Laplace objective over
1370    // theta = [rho, psi] with the real joint Hessian required by NewtonTR/ARC.
1371    macro_rules! run_exact_joint_spatial {
1372        () => {{
1373            let joint_setup = build_two_block_exact_joint_setup(
1374                data,
1375                builder.meanspec(),
1376                builder.noisespec(),
1377                mean_penalty_count,
1378                noise_penalty_count,
1379                extra_rho0.as_slice().unwrap_or(&[]),
1380                None,
1381                kappa_options,
1382            );
1383            let mean_terms = spatial_length_scale_term_indices(builder.meanspec());
1384            let noise_terms = spatial_length_scale_term_indices(builder.noisespec());
1385            let mean_beta_hint_cell = std::cell::RefCell::new(mean_beta_hint.clone());
1386            let noise_beta_hint_cell = std::cell::RefCell::new(noise_beta_hint.clone());
1387            let hyper_warm_start_cell =
1388                std::cell::RefCell::new(None::<CustomFamilyWarmStart>);
1389            // Two-block GAMLSS/location-scale joint likelihoods have a
1390            // β-dependent cross-block Hessian (the (μ,log σ) / (t,log σ)
1391            // off-diagonal blocks involve residual/response scalars that
1392            // shift when β moves). The Wood-Fasiolo structural property
1393            // `H^{-1/2} B_k H^{-1/2} ≽ 0` plus parameter-independent
1394            // nullspace — the mathematical basis for EFS convergence —
1395            // fails here, so EFS/HybridEFS must be excluded at plan time
1396            // rather than retried as a silent first attempt that stalls
1397            // for hundreds of seconds before the runner falls back.
1398            let gamlss_disable_fixed_point = true;
1399            let outer_policy = {
1400                // GAMLSS spatial path: psi_dim = log_kappa_dim + auxiliary_dim,
1401                // matching the (theta_dim - rho_dim) decomposition the
1402                // optimizer uses internally. Build realized ParameterBlockSpecs
1403                // at the seed rho so the family's own cost model — which
1404                // multiplies coefficient-gradient / coefficient-Hessian
1405                // per-row cost by the joint outer-coordinate dimension and
1406                // total p — produces honest `predicted_*_work` estimates.
1407                // Previously this fed `predicted_*_work: 0` to the planner,
1408                // which then ungated dense outer Hessian work that costs
1409                // hundreds of seconds per eval at large scale (see
1410                // `OuterDerivativePolicy::OUTER_HESSIAN_WORK_BUDGET`).
1411                let theta_seed = joint_setup.theta0();
1412                let rho_dim = joint_setup.rho_dim();
1413                let psi_dim = theta_seed.len() - rho_dim;
1414                let rho_seed = theta_seed.slice(s![..rho_dim]).to_owned();
1415                let policy_blocks_res = builder.build_blocks(
1416                    &rho_seed,
1417                    &mean_boot_design,
1418                    &noise_boot_design,
1419                    mean_beta_hint_cell.borrow().clone(),
1420                    noise_beta_hint_cell.borrow().clone(),
1421                );
1422                let mut policy = match policy_blocks_res {
1423                    Ok(policy_blocks) => {
1424                        let policy_family =
1425                            builder.build_family(&mean_boot_design, &noise_boot_design);
1426                        crate::custom_family::CustomFamily::outer_derivative_policy(
1427                            &policy_family,
1428                            &policy_blocks,
1429                            psi_dim,
1430                            options,
1431                        )
1432                    }
1433                    Err(err) => {
1434                        // Block construction at the seed should not fail for
1435                        // any in-tree family, but if it does, fall back to a
1436                        // policy that names the capability honestly and
1437                        // declines to predict cost. Setting work to
1438                        // `u128::MAX` routes the planner through gradient-only
1439                        // BFGS (the universal Hessian-work budget is
1440                        // saturating, so a sentinel is fine here).
1441                        log::warn!(
1442                            "[GAMLSS spatial] failed to realize policy blocks at seed rho ({err}); \
1443                             routing outer optimizer through gradient-only BFGS"
1444                        );
1445                        let capability = if analytic_joint_derivatives_available {
1446                            crate::custom_family::ExactOuterDerivativeOrder::Second
1447                        } else {
1448                            crate::custom_family::ExactOuterDerivativeOrder::First
1449                        };
1450                        crate::custom_family::OuterDerivativePolicy {
1451                            capability,
1452                            predicted_gradient_work: u128::MAX,
1453                            predicted_hessian_work: u128::MAX,
1454                            // No GAMLSS family today overrides its
1455                            // outer-only `_with_options` hooks to consume
1456                            // `outer_score_subsample`; staged-κ would
1457                            // build pilot masks the family then ignores.
1458                            subsample_capable: false,
1459                        }
1460                    }
1461                };
1462                if !analytic_joint_derivatives_available {
1463                    // Capability must not exceed what the analytic derivatives
1464                    // path can supply — the macro's hyper evaluator returns
1465                    // an error otherwise.
1466                    policy.capability =
1467                        crate::custom_family::ExactOuterDerivativeOrder::First;
1468                }
1469                policy
1470            };
1471            optimize_spatial_length_scale_exact_joint(
1472                data,
1473                &[builder.meanspec().clone(), builder.noisespec().clone()],
1474                &[mean_terms, noise_terms],
1475                kappa_options,
1476                &joint_setup,
1477                builder.exact_spatial_seed_risk_profile(),
1478                analytic_joint_derivatives_available,
1479                analytic_joint_derivatives_available,
1480                gamlss_disable_fixed_point,
1481                None,
1482                outer_policy,
1483                |theta, specs: &[TermCollectionSpec], designs: &[TermCollectionDesign]| {
1484                    assert_eq!(
1485                        specs.len(),
1486                        2,
1487                        "joint spatial closure expects exactly two block specs (mean, noise); got {}",
1488                        specs.len(),
1489                    );
1490                    assert_eq!(
1491                        designs.len(),
1492                        2,
1493                        "joint spatial closure expects exactly two block designs (mean, noise); got {}",
1494                        designs.len(),
1495                    );
1496                    let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
1497                    let fit = {
1498                        let blocks = builder.build_blocks(
1499                            &rho,
1500                            &designs[0],
1501                            &designs[1],
1502                            mean_beta_hint_cell.borrow().clone(),
1503                            noise_beta_hint_cell.borrow().clone(),
1504                        )?;
1505                        if mean_beta_hint_cell.borrow().is_none()
1506                            && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
1507                        {
1508                            *mean_beta_hint_cell.borrow_mut() = Some(beta);
1509                        }
1510                        if noise_beta_hint_cell.borrow().is_none()
1511                            && let Some(beta) =
1512                                blocks.get(1).and_then(|block| block.initial_beta.clone())
1513                        {
1514                            *noise_beta_hint_cell.borrow_mut() = Some(beta);
1515                        }
1516                        let family = builder.build_family(&designs[0], &designs[1]);
1517                        // Branch on whether the κ optimizer drives rho.
1518                        //
1519                        // * `log_kappa_dim() > 0 && kappa_options.enabled` ⇒
1520                        //   the outer (ρ, ψ) optimizer is active and
1521                        //   passes each candidate ρ to this closure;
1522                        //   the inner fit must hold log-lambdas fixed
1523                        //   at the supplied ρ so the outer derivative
1524                        //   has a well-defined directional gradient.
1525                        //
1526                        // * Otherwise (κ disabled via the locked-κ
1527                        //   short-circuit, or no spatial terms at all)
1528                        //   the fast path in
1529                        //   `optimize_spatial_length_scale_exact_joint`
1530                        //   calls this closure exactly once at
1531                        //   `theta = theta0`; ρ must still be optimized
1532                        //   from data because the user never pinned it.
1533                        //   `fit_custom_family` performs the joint
1534                        //   ρ + coefficient REML fit at the user's
1535                        //   (now-fixed) kernel scale, which is the
1536                        //   intended behaviour when `length_scale=…` is
1537                        //   set on every spatial term.
1538                        if joint_setup.log_kappa_dim() > 0 && kappa_options.enabled {
1539                            let warm_start = hyper_warm_start_cell.borrow().clone();
1540                            fit_custom_family_fixed_log_lambdas(
1541                                &family,
1542                                &blocks,
1543                                options,
1544                                warm_start.as_ref(),
1545                                0,
1546                                None,
1547                                true,
1548                            )?
1549                        } else {
1550                            fit_custom_family(&family, &blocks, options)?
1551                        }
1552                    };
1553                    let (mean_beta, noise_beta) = builder.extract_primary_betas(&fit)?;
1554                    mean_beta_hint = Some(mean_beta);
1555                    noise_beta_hint = Some(noise_beta);
1556                    *mean_beta_hint_cell.borrow_mut() = mean_beta_hint.clone();
1557                    *noise_beta_hint_cell.borrow_mut() = noise_beta_hint.clone();
1558                    Ok(fit)
1559                },
1560                |theta,
1561                 specs: &[TermCollectionSpec],
1562                 designs: &[TermCollectionDesign],
1563                 eval_mode,
1564                 row_set: &crate::row_kernel::RowSet| {
1565                    use gam_problem::EvalMode;
1566                    if !analytic_joint_derivatives_available {
1567                        return Err(
1568                            "analytic spatial psi derivatives are unavailable for this exact two-block path"
1569                                .to_string(),
1570                        );
1571                    }
1572                    let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
1573                    let blocks = builder.build_blocks(
1574                        &rho,
1575                        &designs[0],
1576                        &designs[1],
1577                        mean_beta_hint_cell.borrow().clone(),
1578                        noise_beta_hint_cell.borrow().clone(),
1579                    )?;
1580                    if mean_beta_hint_cell.borrow().is_none()
1581                        && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
1582                    {
1583                        *mean_beta_hint_cell.borrow_mut() = Some(beta);
1584                    }
1585                    if noise_beta_hint_cell.borrow().is_none()
1586                        && let Some(beta) = blocks.get(1).and_then(|block| block.initial_beta.clone())
1587                    {
1588                        *noise_beta_hint_cell.borrow_mut() = Some(beta);
1589                    }
1590                    let family = builder.build_family(&designs[0], &designs[1]);
1591                    let psiderivative_blocks = builder.build_psiderivative_blocks(
1592                        data,
1593                        &specs[0],
1594                        &specs[1],
1595                        &designs[0],
1596                        &designs[1],
1597                    )?;
1598                    let warm_start = hyper_warm_start_cell.borrow().clone();
1599                    // Forward the κ-staging row set to the family by installing it
1600                    // on the canonical `outer_score_subsample` option. Inner-PIRLS
1601                    // and final covariance still run on full data (the per-row
1602                    // weight is consulted only by outer-only paths inside the
1603                    // family). When the staging schedule is full-data the option
1604                    // stays `None` and the call is equivalent to the prior path.
1605                    let eval_options = match row_set {
1606                        crate::row_kernel::RowSet::All => {
1607                            std::borrow::Cow::Borrowed(options)
1608                        }
1609                        crate::row_kernel::RowSet::Subsample {
1610                            rows,
1611                            n_full,
1612                        } => {
1613                            let subsample = crate::outer_subsample::
1614                                OuterScoreSubsample::from_weighted_rows(
1615                                    (**rows).clone(),
1616                                    *n_full,
1617                                    *n_full as u64,
1618                                );
1619                            let mut cloned = options.clone();
1620                            cloned.outer_score_subsample =
1621                                Some(std::sync::Arc::new(subsample));
1622                            std::borrow::Cow::Owned(cloned)
1623                        }
1624                    };
1625                    let eval = evaluate_custom_family_joint_hyper(
1626                        &family,
1627                        &blocks,
1628                        eval_options.as_ref(),
1629                        &rho,
1630                        &psiderivative_blocks,
1631                        warm_start.as_ref(),
1632                        eval_mode,
1633                    )?;
1634                    *hyper_warm_start_cell.borrow_mut() = Some(eval.warm_start.clone());
1635                    if !eval.inner_converged {
1636                        return Err(
1637                            "exact two-block spatial inner solve did not converge".to_string(),
1638                        );
1639                    }
1640                    if matches!(eval_mode, EvalMode::ValueGradientHessian)
1641                        && !eval.outer_hessian.is_analytic()
1642                    {
1643                        return Err(
1644                            "exact two-block spatial objective requires a full joint [rho, psi] hessian"
1645                                .to_string(),
1646                        );
1647                    }
1648                    Ok((eval.objective, eval.gradient, eval.outer_hessian))
1649                },
1650                |theta, specs: &[TermCollectionSpec], designs: &[TermCollectionDesign]| {
1651                    if !analytic_joint_derivatives_available {
1652                        return Err(
1653                            "analytic spatial psi derivatives are unavailable for this exact two-block path"
1654                                .to_string(),
1655                        );
1656                    }
1657                    let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
1658                    let blocks = builder.build_blocks(
1659                        &rho,
1660                        &designs[0],
1661                        &designs[1],
1662                        mean_beta_hint_cell.borrow().clone(),
1663                        noise_beta_hint_cell.borrow().clone(),
1664                    )?;
1665                    if mean_beta_hint_cell.borrow().is_none()
1666                        && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
1667                    {
1668                        *mean_beta_hint_cell.borrow_mut() = Some(beta);
1669                    }
1670                    if noise_beta_hint_cell.borrow().is_none()
1671                        && let Some(beta) = blocks.get(1).and_then(|block| block.initial_beta.clone())
1672                    {
1673                        *noise_beta_hint_cell.borrow_mut() = Some(beta);
1674                    }
1675                    let family = builder.build_family(&designs[0], &designs[1]);
1676                    let psiderivative_blocks = builder.build_psiderivative_blocks(
1677                        data,
1678                        &specs[0],
1679                        &specs[1],
1680                        &designs[0],
1681                        &designs[1],
1682                    )?;
1683                    let warm_start = hyper_warm_start_cell.borrow().clone();
1684                    let eval = evaluate_custom_family_joint_hyper_efs(
1685                        &family,
1686                        &blocks,
1687                        options,
1688                        &rho,
1689                        &psiderivative_blocks,
1690                        warm_start.as_ref(),
1691                    )?;
1692                    *hyper_warm_start_cell.borrow_mut() = Some(eval.warm_start.clone());
1693                    if !eval.inner_converged {
1694                        return Err(
1695                            "exact two-block spatial EFS inner solve did not converge".to_string(),
1696                        );
1697                    }
1698                    Ok(eval.efs_eval)
1699                },
1700                |_beta: &Array1<f64>| Ok(gam_solve::rho_optimizer::SeedOutcome::NoSlot),
1701            )
1702        }};
1703    }
1704
1705    let mut solved = run_exact_joint_spatial!()
1706        .map_err(|err| format!("exact two-block spatial optimization failed: {err}"))?;
1707
1708    let expected_noise_penalty_count = builder.noise_penalty_count(&solved.designs[1]);
1709    let actual_noise_penalty_count = solved.designs[1].penalties.len();
1710    if expected_noise_penalty_count > actual_noise_penalty_count {
1711        if expected_noise_penalty_count != actual_noise_penalty_count + 1 {
1712            return Err(GamlssError::UnsupportedConfiguration {
1713                reason: format!(
1714                    "location-scale result noise design expected {} penalties after augmentation, got {} before augmentation",
1715                    expected_noise_penalty_count, actual_noise_penalty_count
1716                ),
1717            }
1718            .into());
1719        }
1720        append_binomial_log_sigma_shrinkage_penalty_design(&mut solved.designs[1]);
1721    }
1722
1723    BlockwiseTermFitResult::try_from_parts(BlockwiseTermFitResultParts {
1724        fit: solved.fit,
1725        meanspec_resolved: solved.resolved_specs.remove(0),
1726        noisespec_resolved: solved.resolved_specs.remove(0),
1727        mean_design: solved.designs.remove(0),
1728        noise_design: solved.designs.remove(0),
1729    })
1730}
1731
1732pub(crate) struct GaussianLocationScaleTermBuilder {
1733    pub(crate) y: Array1<f64>,
1734    pub(crate) weights: Array1<f64>,
1735    pub(crate) meanspec: TermCollectionSpec,
1736    pub(crate) noisespec: TermCollectionSpec,
1737    pub(crate) mean_offset: Array1<f64>,
1738    pub(crate) noise_offset: Array1<f64>,
1739}
1740
1741impl LocationScaleFamilyBuilder for GaussianLocationScaleTermBuilder {
1742    type Family = GaussianLocationScaleFamily;
1743
1744    fn meanspec(&self) -> &TermCollectionSpec {
1745        &self.meanspec
1746    }
1747
1748    fn noisespec(&self) -> &TermCollectionSpec {
1749        &self.noisespec
1750    }
1751
1752    fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
1753        // Mirror the Binomial location-scale path: the log-sigma (scale)
1754        // block carries an extra nullspace shrinkage penalty so its
1755        // polynomial nullspace (constant log-sigma, plus the linear term for
1756        // tp/Duchon bases) is not left unpenalized. Without it, outer REML
1757        // optimizes lambda_sigma on a flat/ill-conditioned surface, which can
1758        // flatten the scale envelope (bad Pearson/CRPS/PIT/NLL) and diverge
1759        // the coupled inner Newton (log_sigma residual blows up, beta ->
1760        // infinity). The strength of this shrinkage is REML-selected.
1761        noise_design.penalties.len() + 1
1762    }
1763
1764    fn exact_spatial_joint_supported(&self) -> bool {
1765        true
1766    }
1767
1768    fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
1769        crate::seeding::SeedRiskProfile::GaussianLocationScale
1770    }
1771
1772    fn build_blocks(
1773        &self,
1774        theta: &Array1<f64>,
1775        mean_design: &TermCollectionDesign,
1776        noise_design: &TermCollectionDesign,
1777        mean_beta_hint: Option<Array1<f64>>,
1778        noise_beta_hint: Option<Array1<f64>>,
1779    ) -> Result<Vec<ParameterBlockSpec>, String> {
1780        let layout = GamlssLambdaLayout::two_block(
1781            mean_design.penalties.len(),
1782            self.noise_penalty_count(noise_design),
1783        );
1784        layout.validate_theta_len(theta.len(), "gaussian location-scale")?;
1785        let (meanspec, noisespec) = build_gaussian_mean_and_scale_blocks(
1786            &self.y,
1787            &self.weights,
1788            mean_design,
1789            noise_design,
1790            &self.mean_offset,
1791            &self.noise_offset,
1792            layout.mean_from(theta),
1793            layout.noise_from(theta),
1794            mean_beta_hint,
1795            noise_beta_hint,
1796            "GaussianLocationScale::build_blocks",
1797        )?;
1798        Ok(vec![meanspec, noisespec])
1799    }
1800
1801    fn build_family(
1802        &self,
1803        mean_design: &TermCollectionDesign,
1804        noise_design: &TermCollectionDesign,
1805    ) -> Self::Family {
1806        let preparednoise_design =
1807            prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design)
1808                .expect("prepared Gaussian log-sigma design should match block construction");
1809        GaussianLocationScaleFamily {
1810            y: self.y.clone(),
1811            weights: self.weights.clone(),
1812            mu_design: Some(mean_design.design.clone()),
1813            log_sigma_design: Some(preparednoise_design),
1814            policy: gam_runtime::resource::ResourcePolicy::default_library(),
1815            cached_row_scalars: std::sync::RwLock::new(None),
1816        }
1817    }
1818
1819    fn extract_primary_betas(
1820        &self,
1821        fit: &UnifiedFitResult,
1822    ) -> Result<(Array1<f64>, Array1<f64>), String> {
1823        let mean_beta = fit
1824            .block_states
1825            .get(GaussianLocationScaleFamily::BLOCK_MU)
1826            .ok_or_else(|| "missing Gaussian mu block state".to_string())?
1827            .beta
1828            .clone();
1829        let noise_beta = fit
1830            .block_states
1831            .get(GaussianLocationScaleFamily::BLOCK_LOG_SIGMA)
1832            .ok_or_else(|| "missing Gaussian log_sigma block state".to_string())?
1833            .beta
1834            .clone();
1835        Ok((mean_beta, noise_beta))
1836    }
1837
1838    fn build_psiderivative_blocks(
1839        &self,
1840        data: ndarray::ArrayView2<'_, f64>,
1841        meanspec_resolved: &TermCollectionSpec,
1842        noisespec_resolved: &TermCollectionSpec,
1843        mean_design: &TermCollectionDesign,
1844        noise_design: &TermCollectionDesign,
1845    ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
1846        let mean_derivs =
1847            build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
1848                .ok_or_else(|| "missing Gaussian mean spatial psi derivatives".to_string())?;
1849        let noise_derivs =
1850            build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
1851                .ok_or_else(|| "missing Gaussian log-sigma spatial psi derivatives".to_string())?;
1852        Ok(vec![mean_derivs, noise_derivs])
1853    }
1854}
1855
1856pub(crate) struct GaussianLocationScaleWiggleTermBuilder {
1857    pub(crate) y: Array1<f64>,
1858    pub(crate) weights: Array1<f64>,
1859    pub(crate) meanspec: TermCollectionSpec,
1860    pub(crate) noisespec: TermCollectionSpec,
1861    pub(crate) mean_offset: Array1<f64>,
1862    pub(crate) noise_offset: Array1<f64>,
1863    pub(crate) wiggle_knots: Array1<f64>,
1864    pub(crate) wiggle_degree: usize,
1865    pub(crate) wiggle_block: ParameterBlockInput,
1866}
1867
1868impl LocationScaleFamilyBuilder for GaussianLocationScaleWiggleTermBuilder {
1869    type Family = GaussianLocationScaleWiggleFamily;
1870
1871    fn meanspec(&self) -> &TermCollectionSpec {
1872        &self.meanspec
1873    }
1874
1875    fn noisespec(&self) -> &TermCollectionSpec {
1876        &self.noisespec
1877    }
1878
1879    fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
1880        // Same nullspace log-sigma shrinkage penalty as the non-wiggle
1881        // Gaussian builder; see GaussianLocationScaleTermBuilder.
1882        noise_design.penalties.len() + 1
1883    }
1884
1885    fn exact_spatial_joint_supported(&self) -> bool {
1886        true
1887    }
1888
1889    fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
1890        crate::seeding::SeedRiskProfile::GaussianLocationScale
1891    }
1892
1893    fn require_exact_spatial_joint(&self) -> bool {
1894        true
1895    }
1896
1897    fn extra_rho0(&self) -> Result<Array1<f64>, String> {
1898        initial_log_lambdas_orzeros(&self.wiggle_block)
1899    }
1900
1901    fn build_blocks(
1902        &self,
1903        theta: &Array1<f64>,
1904        mean_design: &TermCollectionDesign,
1905        noise_design: &TermCollectionDesign,
1906        mean_beta_hint: Option<Array1<f64>>,
1907        noise_beta_hint: Option<Array1<f64>>,
1908    ) -> Result<Vec<ParameterBlockSpec>, String> {
1909        let layout = GamlssLambdaLayout::withwiggle(
1910            mean_design.penalties.len(),
1911            self.noise_penalty_count(noise_design),
1912            self.wiggle_block.penalties.len(),
1913        );
1914        layout.validate_theta_len(theta.len(), "gaussian location-scale wiggle")?;
1915        let (mut meanspec, mut noisespec) = build_gaussian_mean_and_scale_blocks(
1916            &self.y,
1917            &self.weights,
1918            mean_design,
1919            noise_design,
1920            &self.mean_offset,
1921            &self.noise_offset,
1922            layout.mean_from(theta),
1923            layout.noise_from(theta),
1924            mean_beta_hint,
1925            noise_beta_hint,
1926            "GaussianLocationScaleWiggle::build_blocks",
1927        )?;
1928        // Keep the dynamic full-width wiggle basis safe from a canonical-gauge
1929        // column drop: route the shared level/intercept alias onto the
1930        // column-reducible mean and log-sigma blocks by giving them a lower
1931        // gauge priority than the wiggle block's fixed 100 (see the binomial
1932        // wiggle path and `build_location_scale_wiggle_block`).
1933        meanspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
1934        noisespec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
1935        let n_rows = meanspec.design.nrows();
1936        let wigglespec = build_location_scale_wiggle_block(
1937            "wiggle",
1938            self.wiggle_block.design.clone(),
1939            self.wiggle_block.offset.clone(),
1940            wiggle_block_penalty_matrices(&self.wiggle_block),
1941            self.wiggle_block.nullspace_dims.clone(),
1942            layout.wiggle_from(theta),
1943            self.wiggle_block.initial_beta.clone(),
1944            n_rows,
1945        )?;
1946        Ok(vec![meanspec, noisespec, wigglespec])
1947    }
1948
1949    fn build_family(
1950        &self,
1951        mean_design: &TermCollectionDesign,
1952        noise_design: &TermCollectionDesign,
1953    ) -> Self::Family {
1954        let preparednoise_design =
1955            prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design).expect(
1956                "prepared Gaussian log-sigma design should match wiggle block construction",
1957            );
1958        GaussianLocationScaleWiggleFamily {
1959            y: self.y.clone(),
1960            weights: self.weights.clone(),
1961            mu_design: Some(mean_design.design.clone()),
1962            log_sigma_design: Some(preparednoise_design),
1963            wiggle_knots: self.wiggle_knots.clone(),
1964            wiggle_degree: self.wiggle_degree,
1965            policy: gam_runtime::resource::ResourcePolicy::default_library(),
1966            cached_row_scalars: std::sync::RwLock::new(None),
1967        }
1968    }
1969
1970    fn extract_primary_betas(
1971        &self,
1972        fit: &UnifiedFitResult,
1973    ) -> Result<(Array1<f64>, Array1<f64>), String> {
1974        let mean_beta = fit
1975            .block_states
1976            .get(GaussianLocationScaleWiggleFamily::BLOCK_MU)
1977            .ok_or_else(|| "missing Gaussian wiggle mu block state".to_string())?
1978            .beta
1979            .clone();
1980        let noise_beta = fit
1981            .block_states
1982            .get(GaussianLocationScaleWiggleFamily::BLOCK_LOG_SIGMA)
1983            .ok_or_else(|| "missing Gaussian wiggle log_sigma block state".to_string())?
1984            .beta
1985            .clone();
1986        Ok((mean_beta, noise_beta))
1987    }
1988
1989    fn build_psiderivative_blocks(
1990        &self,
1991        data: ndarray::ArrayView2<'_, f64>,
1992        meanspec_resolved: &TermCollectionSpec,
1993        noisespec_resolved: &TermCollectionSpec,
1994        mean_design: &TermCollectionDesign,
1995        noise_design: &TermCollectionDesign,
1996    ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
1997        let mean_derivs =
1998            build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?.ok_or_else(
1999                || "missing Gaussian wiggle mean spatial psi derivatives".to_string(),
2000            )?;
2001        let noise_derivs =
2002            build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2003                .ok_or_else(|| {
2004                    "missing Gaussian wiggle log-sigma spatial psi derivatives".to_string()
2005                })?;
2006        Ok(vec![mean_derivs, noise_derivs, Vec::new()])
2007    }
2008}
2009
2010pub(crate) struct BinomialLocationScaleTermBuilder {
2011    pub(crate) y: Array1<f64>,
2012    pub(crate) weights: Array1<f64>,
2013    pub(crate) link_kind: InverseLink,
2014    pub(crate) meanspec: TermCollectionSpec,
2015    pub(crate) noisespec: TermCollectionSpec,
2016    pub(crate) mean_offset: Array1<f64>,
2017    pub(crate) noise_offset: Array1<f64>,
2018}
2019
2020impl LocationScaleFamilyBuilder for BinomialLocationScaleTermBuilder {
2021    type Family = BinomialLocationScaleFamily;
2022
2023    fn meanspec(&self) -> &TermCollectionSpec {
2024        &self.meanspec
2025    }
2026
2027    fn noisespec(&self) -> &TermCollectionSpec {
2028        &self.noisespec
2029    }
2030
2031    fn exact_spatial_joint_supported(&self) -> bool {
2032        true
2033    }
2034
2035    fn require_exact_spatial_joint(&self) -> bool {
2036        true
2037    }
2038
2039    fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
2040        noise_design.penalties.len() + 1
2041    }
2042
2043    fn build_blocks(
2044        &self,
2045        theta: &Array1<f64>,
2046        mean_design: &TermCollectionDesign,
2047        noise_design: &TermCollectionDesign,
2048        mean_beta_hint: Option<Array1<f64>>,
2049        noise_beta_hint: Option<Array1<f64>>,
2050    ) -> Result<Vec<ParameterBlockSpec>, String> {
2051        let layout = GamlssLambdaLayout::two_block(
2052            mean_design.penalties.len(),
2053            self.noise_penalty_count(noise_design),
2054        );
2055        layout.validate_theta_len(theta.len(), "binomial location-scale")?;
2056        let (thresholdspec, log_sigmaspec) = build_binomial_threshold_and_scale_blocks(
2057            &self.y,
2058            &self.weights,
2059            &self.link_kind,
2060            mean_design,
2061            noise_design,
2062            &self.mean_offset,
2063            &self.noise_offset,
2064            layout.mean_from(theta),
2065            layout.noise_from(theta),
2066            mean_beta_hint,
2067            noise_beta_hint,
2068            "BinomialLocationScale::build_blocks",
2069        )?;
2070        Ok(vec![thresholdspec, log_sigmaspec])
2071    }
2072
2073    fn build_family(
2074        &self,
2075        mean_design: &TermCollectionDesign,
2076        noise_design: &TermCollectionDesign,
2077    ) -> Self::Family {
2078        let identifiednoise_design =
2079            identified_binomial_log_sigma_design(mean_design, noise_design, &self.weights)
2080                .expect("identified binomial log-sigma design");
2081        BinomialLocationScaleFamily {
2082            y: self.y.clone(),
2083            weights: self.weights.clone(),
2084            link_kind: self.link_kind.clone(),
2085            threshold_design: Some(mean_design.design.clone()),
2086            log_sigma_design: Some(identifiednoise_design),
2087            policy: gam_runtime::resource::ResourcePolicy::default_library(),
2088        }
2089    }
2090
2091    fn extract_primary_betas(
2092        &self,
2093        fit: &UnifiedFitResult,
2094    ) -> Result<(Array1<f64>, Array1<f64>), String> {
2095        let mean_beta = fit
2096            .block_states
2097            .get(BinomialLocationScaleFamily::BLOCK_T)
2098            .ok_or_else(|| "missing Binomial threshold block state".to_string())?
2099            .beta
2100            .clone();
2101        let noise_beta = fit
2102            .block_states
2103            .get(BinomialLocationScaleFamily::BLOCK_LOG_SIGMA)
2104            .ok_or_else(|| "missing Binomial log_sigma block state".to_string())?
2105            .beta
2106            .clone();
2107        Ok((mean_beta, noise_beta))
2108    }
2109
2110    fn build_psiderivative_blocks(
2111        &self,
2112        data: ndarray::ArrayView2<'_, f64>,
2113        meanspec_resolved: &TermCollectionSpec,
2114        noisespec_resolved: &TermCollectionSpec,
2115        mean_design: &TermCollectionDesign,
2116        noise_design: &TermCollectionDesign,
2117    ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2118        let mean_derivs =
2119            build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2120                .ok_or_else(|| "missing threshold spatial psi derivatives".to_string())?;
2121        let noise_derivs =
2122            build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2123                .ok_or_else(|| "missing log_sigma spatial psi derivatives".to_string())?;
2124        Ok(vec![mean_derivs, noise_derivs])
2125    }
2126}
2127
2128pub(crate) struct BinomialLocationScaleWiggleTermBuilder {
2129    pub(crate) y: Array1<f64>,
2130    pub(crate) weights: Array1<f64>,
2131    pub(crate) link_kind: InverseLink,
2132    pub(crate) meanspec: TermCollectionSpec,
2133    pub(crate) noisespec: TermCollectionSpec,
2134    pub(crate) mean_offset: Array1<f64>,
2135    pub(crate) noise_offset: Array1<f64>,
2136    pub(crate) wiggle_knots: Array1<f64>,
2137    pub(crate) wiggle_degree: usize,
2138    pub(crate) wiggle_block: ParameterBlockInput,
2139}
2140
2141impl LocationScaleFamilyBuilder for BinomialLocationScaleWiggleTermBuilder {
2142    type Family = BinomialLocationScaleWiggleFamily;
2143
2144    fn meanspec(&self) -> &TermCollectionSpec {
2145        &self.meanspec
2146    }
2147
2148    fn noisespec(&self) -> &TermCollectionSpec {
2149        &self.noisespec
2150    }
2151
2152    fn exact_spatial_joint_supported(&self) -> bool {
2153        true
2154    }
2155
2156    fn require_exact_spatial_joint(&self) -> bool {
2157        true
2158    }
2159
2160    fn extra_rho0(&self) -> Result<Array1<f64>, String> {
2161        initial_log_lambdas_orzeros(&self.wiggle_block)
2162    }
2163
2164    fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
2165        noise_design.penalties.len() + 1
2166    }
2167
2168    fn build_blocks(
2169        &self,
2170        theta: &Array1<f64>,
2171        mean_design: &TermCollectionDesign,
2172        noise_design: &TermCollectionDesign,
2173        mean_beta_hint: Option<Array1<f64>>,
2174        noise_beta_hint: Option<Array1<f64>>,
2175    ) -> Result<Vec<ParameterBlockSpec>, String> {
2176        let layout = GamlssLambdaLayout::withwiggle(
2177            mean_design.penalties.len(),
2178            self.noise_penalty_count(noise_design),
2179            self.wiggle_block.penalties.len(),
2180        );
2181        layout.validate_theta_len(theta.len(), "wiggle location-scale")?;
2182        let (mut thresholdspec, mut log_sigmaspec) = build_binomial_threshold_and_scale_blocks(
2183            &self.y,
2184            &self.weights,
2185            &self.link_kind,
2186            mean_design,
2187            noise_design,
2188            &self.mean_offset,
2189            &self.noise_offset,
2190            layout.mean_from(theta),
2191            layout.noise_from(theta),
2192            mean_beta_hint,
2193            noise_beta_hint,
2194            "BinomialLocationScaleWiggle::build_blocks",
2195        )?;
2196        // The dynamic monotone wiggle basis is regenerated at full raw width
2197        // every inner iteration and asserts `x.ncols() == spec.design.ncols()`
2198        // in `block_geometry`, so it cannot tolerate a canonical-gauge column
2199        // drop. The level/intercept direction the I-spline shares with the
2200        // threshold block must therefore be routed onto the threshold (and the
2201        // log-sigma) block, whose static designs are column-reducible and
2202        // lifted back via the canonical per-block transform `T`. Give both
2203        // non-wiggle blocks a lower gauge priority than the wiggle block (which
2204        // `build_location_scale_wiggle_block` fixes at 100) so the shared-level
2205        // alias drop lands on them and leaves the dynamic wiggle basis full
2206        // width — mirroring the binomial mean-wiggle path.
2207        thresholdspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2208        log_sigmaspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2209        let n_rows = thresholdspec.design.nrows();
2210        let wigglespec = build_location_scale_wiggle_block(
2211            "wiggle",
2212            self.wiggle_block.design.clone(),
2213            self.wiggle_block.offset.clone(),
2214            wiggle_block_penalty_matrices(&self.wiggle_block),
2215            vec![],
2216            layout.wiggle_from(theta),
2217            self.wiggle_block.initial_beta.clone(),
2218            n_rows,
2219        )?;
2220        Ok(vec![thresholdspec, log_sigmaspec, wigglespec])
2221    }
2222
2223    fn build_family(
2224        &self,
2225        mean_design: &TermCollectionDesign,
2226        noise_design: &TermCollectionDesign,
2227    ) -> Self::Family {
2228        let identifiednoise_design =
2229            identified_binomial_log_sigma_design(mean_design, noise_design, &self.weights)
2230                .expect("identified binomial log-sigma design should match block construction");
2231        BinomialLocationScaleWiggleFamily {
2232            y: self.y.clone(),
2233            weights: self.weights.clone(),
2234            link_kind: self.link_kind.clone(),
2235            threshold_design: Some(mean_design.design.clone()),
2236            log_sigma_design: Some(identifiednoise_design),
2237            wiggle_knots: self.wiggle_knots.clone(),
2238            wiggle_degree: self.wiggle_degree,
2239            policy: gam_runtime::resource::ResourcePolicy::default_library(),
2240        }
2241    }
2242
2243    fn extract_primary_betas(
2244        &self,
2245        fit: &UnifiedFitResult,
2246    ) -> Result<(Array1<f64>, Array1<f64>), String> {
2247        let mean_beta = fit
2248            .block_states
2249            .get(BinomialLocationScaleWiggleFamily::BLOCK_T)
2250            .ok_or_else(|| "missing Binomial wiggle threshold block state".to_string())?
2251            .beta
2252            .clone();
2253        let noise_beta = fit
2254            .block_states
2255            .get(BinomialLocationScaleWiggleFamily::BLOCK_LOG_SIGMA)
2256            .ok_or_else(|| "missing Binomial wiggle log_sigma block state".to_string())?
2257            .beta
2258            .clone();
2259        Ok((mean_beta, noise_beta))
2260    }
2261
2262    fn build_psiderivative_blocks(
2263        &self,
2264        data: ndarray::ArrayView2<'_, f64>,
2265        meanspec_resolved: &TermCollectionSpec,
2266        noisespec_resolved: &TermCollectionSpec,
2267        mean_design: &TermCollectionDesign,
2268        noise_design: &TermCollectionDesign,
2269    ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2270        let mean_derivs =
2271            build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2272                .ok_or_else(|| "missing threshold spatial psi derivatives".to_string())?;
2273        let noise_derivs =
2274            build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2275                .ok_or_else(|| "missing log_sigma spatial psi derivatives".to_string())?;
2276        // The wiggle block has no direct spatial design matrix of its own in the
2277        // term builder. Spatial psi moves the wiggle family only through the
2278        // realized threshold/log-sigma designs, which in turn perturb q0 and the
2279        // realized wiggle basis B(q0). The exact joint wiggle psi hooks consume
2280        // those threshold/log-sigma derivative payloads and reconstruct the full
2281        // flattened likelihood-side [rho, psi] calculus internally, so the
2282        // wiggle block intentionally contributes no direct CustomFamilyBlockPsiDerivative
2283        // entries here.
2284        Ok(vec![mean_derivs, noise_derivs, Vec::new()])
2285    }
2286}
2287
2288pub(crate) fn fit_gaussian_location_scale_terms(
2289    data: ndarray::ArrayView2<'_, f64>,
2290    spec: GaussianLocationScaleTermSpec,
2291    options: &BlockwiseFitOptions,
2292    kappa_options: &SpatialLengthScaleOptimizationOptions,
2293) -> Result<BlockwiseTermFitResult, String> {
2294    validate_gaussian_location_scale_termspec(data, &spec, "fit_gaussian_location_scale_terms")?;
2295    fit_location_scale_terms(
2296        data,
2297        GaussianLocationScaleTermBuilder {
2298            y: spec.y,
2299            weights: spec.weights,
2300            meanspec: spec.meanspec,
2301            noisespec: spec.log_sigmaspec,
2302            mean_offset: spec.mean_offset,
2303            noise_offset: spec.log_sigma_offset,
2304        },
2305        options,
2306        kappa_options,
2307    )
2308}
2309
2310pub(crate) fn fit_gaussian_location_scalewiggle_terms(
2311    data: ndarray::ArrayView2<'_, f64>,
2312    spec: GaussianLocationScaleWiggleTermSpec,
2313    options: &BlockwiseFitOptions,
2314    kappa_options: &SpatialLengthScaleOptimizationOptions,
2315) -> Result<BlockwiseTermFitResult, String> {
2316    validate_gaussian_location_scalewiggle_termspec(
2317        data,
2318        &spec,
2319        "fit_gaussian_location_scalewiggle_terms",
2320    )?;
2321    fit_location_scale_terms(
2322        data,
2323        GaussianLocationScaleWiggleTermBuilder {
2324            y: spec.y,
2325            weights: spec.weights,
2326            meanspec: spec.meanspec,
2327            noisespec: spec.log_sigmaspec,
2328            mean_offset: spec.mean_offset,
2329            noise_offset: spec.log_sigma_offset,
2330            wiggle_knots: spec.wiggle_knots,
2331            wiggle_degree: spec.wiggle_degree,
2332            wiggle_block: spec.wiggle_block,
2333        },
2334        options,
2335        kappa_options,
2336    )
2337}
2338
2339pub(crate) fn select_gaussian_location_scale_link_wiggle_basis_from_pilot(
2340    pilot: &BlockwiseTermFitResult,
2341    wiggle_cfg: &WiggleBlockConfig,
2342    wiggle_penalty_orders: &[usize],
2343) -> Result<SelectedWiggleBasis, String> {
2344    let q_seed = pilot
2345        .fit
2346        .block_states
2347        .first()
2348        .ok_or_else(|| "pilot Gaussian wiggle fit is missing mean block".to_string())?
2349        .eta
2350        .view();
2351    select_wiggle_basis_from_seed(q_seed, wiggle_cfg, wiggle_penalty_orders)
2352}
2353
2354pub(crate) fn fit_gaussian_location_scale_terms_with_selected_wiggle(
2355    data: ndarray::ArrayView2<'_, f64>,
2356    spec: GaussianLocationScaleTermSpec,
2357    selected_wiggle_basis: SelectedWiggleBasis,
2358    options: &BlockwiseFitOptions,
2359    kappa_options: &SpatialLengthScaleOptimizationOptions,
2360) -> Result<BlockwiseTermWiggleFitResult, String> {
2361    let SelectedWiggleBasis {
2362        knots: wiggle_knots,
2363        degree: wiggle_degree,
2364        block: wiggle_block,
2365        ..
2366    } = selected_wiggle_basis;
2367    let solved = fit_gaussian_location_scalewiggle_terms(
2368        data,
2369        GaussianLocationScaleWiggleTermSpec {
2370            y: spec.y,
2371            weights: spec.weights,
2372            meanspec: spec.meanspec,
2373            log_sigmaspec: spec.log_sigmaspec,
2374            mean_offset: spec.mean_offset,
2375            log_sigma_offset: spec.log_sigma_offset,
2376            wiggle_knots: wiggle_knots.clone(),
2377            wiggle_degree,
2378            wiggle_block,
2379        },
2380        options,
2381        kappa_options,
2382    )?;
2383
2384    BlockwiseTermWiggleFitResult::try_from_parts(BlockwiseTermWiggleFitResultParts {
2385        fit: solved,
2386        wiggle_knots,
2387        wiggle_degree,
2388    })
2389}
2390
2391pub(crate) fn fit_binomial_location_scale_terms(
2392    data: ndarray::ArrayView2<'_, f64>,
2393    spec: BinomialLocationScaleTermSpec,
2394    options: &BlockwiseFitOptions,
2395    kappa_options: &SpatialLengthScaleOptimizationOptions,
2396) -> Result<BlockwiseTermFitResult, String> {
2397    validate_binomial_location_scale_termspec(data, &spec, "fit_binomial_location_scale_terms")?;
2398    fit_location_scale_terms(
2399        data,
2400        BinomialLocationScaleTermBuilder {
2401            y: spec.y,
2402            weights: spec.weights,
2403            link_kind: spec.link_kind,
2404            meanspec: spec.thresholdspec,
2405            noisespec: spec.log_sigmaspec,
2406            mean_offset: spec.threshold_offset,
2407            noise_offset: spec.log_sigma_offset,
2408        },
2409        options,
2410        kappa_options,
2411    )
2412}
2413
2414pub(crate) fn fit_binomial_location_scalewiggle_terms(
2415    data: ndarray::ArrayView2<'_, f64>,
2416    spec: BinomialLocationScaleWiggleTermSpec,
2417    options: &BlockwiseFitOptions,
2418    kappa_options: &SpatialLengthScaleOptimizationOptions,
2419) -> Result<BlockwiseTermFitResult, String> {
2420    validate_binomial_location_scalewiggle_termspec(
2421        data,
2422        &spec,
2423        "fit_binomial_location_scalewiggle_terms",
2424    )?;
2425    fit_location_scale_terms(
2426        data,
2427        BinomialLocationScaleWiggleTermBuilder {
2428            y: spec.y,
2429            weights: spec.weights,
2430            link_kind: spec.link_kind,
2431            meanspec: spec.thresholdspec,
2432            noisespec: spec.log_sigmaspec,
2433            mean_offset: spec.threshold_offset,
2434            noise_offset: spec.log_sigma_offset,
2435            wiggle_knots: spec.wiggle_knots,
2436            wiggle_degree: spec.wiggle_degree,
2437            wiggle_block: spec.wiggle_block,
2438        },
2439        options,
2440        kappa_options,
2441    )
2442}
2443
2444pub(crate) fn select_binomial_location_scale_link_wiggle_basis_from_pilot(
2445    pilot: &BlockwiseTermFitResult,
2446    wiggle_cfg: &WiggleBlockConfig,
2447    wiggle_penalty_orders: &[usize],
2448) -> Result<SelectedWiggleBasis, String> {
2449    let eta_t = pilot
2450        .fit
2451        .block_states
2452        .first()
2453        .ok_or_else(|| "pilot fit is missing threshold block".to_string())?
2454        .eta
2455        .view();
2456    let eta_ls = pilot
2457        .fit
2458        .block_states
2459        .get(1)
2460        .ok_or_else(|| "pilot fit is missing log_sigma block".to_string())?
2461        .eta
2462        .view();
2463    let sigma = eta_ls.mapv(safe_exp);
2464    let q_seed = Array1::from_iter(eta_t.iter().zip(sigma.iter()).map(|(&t, &s)| -t / s));
2465    select_wiggle_basis_from_seed(q_seed.view(), wiggle_cfg, wiggle_penalty_orders)
2466}
2467
2468pub(crate) fn fit_binomial_location_scale_terms_with_selected_wiggle(
2469    data: ndarray::ArrayView2<'_, f64>,
2470    spec: BinomialLocationScaleTermSpec,
2471    selected_wiggle_basis: SelectedWiggleBasis,
2472    options: &BlockwiseFitOptions,
2473    kappa_options: &SpatialLengthScaleOptimizationOptions,
2474) -> Result<BlockwiseTermWiggleFitResult, String> {
2475    let SelectedWiggleBasis {
2476        knots: wiggle_knots,
2477        degree: wiggle_degree,
2478        block: wiggle_block,
2479        ..
2480    } = selected_wiggle_basis;
2481    let solved = fit_binomial_location_scalewiggle_terms(
2482        data,
2483        BinomialLocationScaleWiggleTermSpec {
2484            y: spec.y,
2485            weights: spec.weights,
2486            link_kind: spec.link_kind,
2487            thresholdspec: spec.thresholdspec,
2488            log_sigmaspec: spec.log_sigmaspec,
2489            threshold_offset: spec.threshold_offset,
2490            log_sigma_offset: spec.log_sigma_offset,
2491            wiggle_knots: wiggle_knots.clone(),
2492            wiggle_degree,
2493            wiggle_block,
2494        },
2495        options,
2496        kappa_options,
2497    )?;
2498
2499    BlockwiseTermWiggleFitResult::try_from_parts(BlockwiseTermWiggleFitResultParts {
2500        fit: solved,
2501        wiggle_knots,
2502        wiggle_degree,
2503    })
2504}
2505
2506pub(crate) fn select_binomial_mean_link_wiggle_basis_from_pilot(
2507    pilot_design: &TermCollectionDesign,
2508    pilot_fit: &UnifiedFitResult,
2509    wiggle_cfg: &WiggleBlockConfig,
2510    wiggle_penalty_orders: &[usize],
2511) -> Result<SelectedWiggleBasis, String> {
2512    let q_seed = pilot_design.design.dot(&pilot_fit.beta);
2513    select_wiggle_basis_from_seed(q_seed.view(), wiggle_cfg, wiggle_penalty_orders)
2514}
2515
2516pub(crate) fn fit_binomial_mean_wiggle_terms_with_selected_basis(
2517    data: ndarray::ArrayView2<'_, f64>,
2518    pilot_spec: &TermCollectionSpec,
2519    pilot_design: &TermCollectionDesign,
2520    pilot_fit: &UnifiedFitResult,
2521    y: &Array1<f64>,
2522    weights: &Array1<f64>,
2523    link_kind: InverseLink,
2524    selected_wiggle_basis: SelectedWiggleBasis,
2525    options: &BlockwiseFitOptions,
2526    kappa_options: &SpatialLengthScaleOptimizationOptions,
2527) -> Result<BinomialMeanWiggleTermFitResult, String> {
2528    const RHO_BOUND: f64 = 12.0;
2529
2530    validate_term_weights(
2531        data,
2532        y.len(),
2533        weights,
2534        "fit_binomial_mean_wiggle_terms_with_selected_basis",
2535    )?;
2536    validate_binomial_response(y, "fit_binomial_mean_wiggle_terms_with_selected_basis")?;
2537
2538    // Large-n binomial mean-wiggle fits keep the caller's explicit Hessian
2539    // request. The unified evaluator chooses the scalable exact representation
2540    // (dense for small work, operator HVP for large work) instead of routing to
2541    // gradient-only BFGS by observation count.
2542
2543    let SelectedWiggleBasis {
2544        knots: wiggle_knots,
2545        degree: wiggle_degree,
2546        block: wiggle_block,
2547        ..
2548    } = selected_wiggle_basis;
2549
2550    let spatial_terms = spatial_length_scale_term_indices(pilot_spec);
2551    if spatial_terms.is_empty() {
2552        let fit = fit_binomial_mean_wiggle(
2553            BinomialMeanWiggleSpec {
2554                y: y.clone(),
2555                weights: weights.clone(),
2556                link_kind,
2557                wiggle_knots: wiggle_knots.clone(),
2558                wiggle_degree,
2559                eta_block: ParameterBlockInput {
2560                    design: pilot_design.design.clone(),
2561                    offset: Array1::zeros(y.len()),
2562                    penalties: pilot_design
2563                        .penalties
2564                        .iter()
2565                        .map(crate::model_types::PenaltySpec::from_blockwise_ref)
2566                        .collect(),
2567                    nullspace_dims: vec![],
2568                    initial_log_lambdas: Some(
2569                        pilot_fit
2570                            .lambdas
2571                            .mapv(|v| v.max(WARMSTART_LOG_LAMBDA_FLOOR).ln()),
2572                    ),
2573                    initial_beta: Some(pilot_fit.beta.clone()),
2574                },
2575                wiggle_block,
2576            },
2577            options,
2578        )?;
2579        return Ok(BinomialMeanWiggleTermFitResult {
2580            fit,
2581            resolvedspec: pilot_spec.clone(),
2582            design: pilot_design.clone(),
2583            wiggle_knots,
2584            wiggle_degree,
2585        });
2586    }
2587
2588    let dims_per_term = spatial_dims_per_term(pilot_spec, &spatial_terms);
2589    let log_kappa0 =
2590        SpatialLogKappaCoords::from_length_scales_aniso(pilot_spec, &spatial_terms, kappa_options)
2591            .reseed_from_data(data, pilot_spec, &spatial_terms, kappa_options);
2592    let log_kappa_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
2593        data,
2594        pilot_spec,
2595        &spatial_terms,
2596        &dims_per_term,
2597        kappa_options,
2598    );
2599    let log_kappa_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
2600        data,
2601        pilot_spec,
2602        &spatial_terms,
2603        &dims_per_term,
2604        kappa_options,
2605    );
2606    // Project seed onto bounds; spec.length_scale is a hint, not a constraint.
2607    let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
2608
2609    let eta_penalty_count = pilot_design.penalties.len();
2610    let wiggle_penalty_count = initial_log_lambdas_orzeros(&wiggle_block)?.len();
2611    let rho_dim = eta_penalty_count + wiggle_penalty_count;
2612    let baseline_resolvedspec = log_kappa0
2613        .apply_tospec(pilot_spec, &spatial_terms)
2614        .map_err(|e| e.to_string())?;
2615    let baseline_design =
2616        build_term_collection_design(data, &baseline_resolvedspec).map_err(|e| e.to_string())?;
2617    let baseline_fit = fit_binomial_mean_wiggle(
2618        BinomialMeanWiggleSpec {
2619            y: y.clone(),
2620            weights: weights.clone(),
2621            link_kind: link_kind.clone(),
2622            wiggle_knots: wiggle_knots.clone(),
2623            wiggle_degree,
2624            eta_block: ParameterBlockInput {
2625                design: baseline_design.design.clone(),
2626                offset: Array1::zeros(y.len()),
2627                penalties: baseline_design
2628                    .penalties
2629                    .iter()
2630                    .map(crate::model_types::PenaltySpec::from_blockwise_ref)
2631                    .collect(),
2632                nullspace_dims: vec![],
2633                initial_log_lambdas: Some(
2634                    pilot_fit
2635                        .lambdas
2636                        .mapv(|v| v.max(WARMSTART_LOG_LAMBDA_FLOOR).ln()),
2637                ),
2638                initial_beta: Some(pilot_fit.beta.clone()),
2639            },
2640            wiggle_block: wiggle_block.clone(),
2641        },
2642        options,
2643    )?;
2644    let baseline_log_lambdas = baseline_fit
2645        .lambdas
2646        .mapv(|v| v.max(WARMSTART_LOG_LAMBDA_FLOOR).ln());
2647    if baseline_log_lambdas.len() != rho_dim {
2648        return Err(GamlssError::DimensionMismatch {
2649            reason: format!(
2650                "baseline binomial mean-wiggle fit returned {} log-lambdas, expected {rho_dim}",
2651                baseline_log_lambdas.len()
2652            ),
2653        }
2654        .into());
2655    }
2656    let baseline_eta_beta = baseline_fit
2657        .block_states
2658        .get(BinomialMeanWiggleFamily::BLOCK_ETA)
2659        .ok_or_else(|| "baseline binomial mean-wiggle fit missing eta block".to_string())?
2660        .beta
2661        .clone();
2662    let baseline_wiggle_beta = Some(
2663        baseline_fit
2664            .block_states
2665            .get(BinomialMeanWiggleFamily::BLOCK_WIGGLE)
2666            .ok_or_else(|| "baseline binomial mean-wiggle fit missing wiggle block".to_string())?
2667            .beta
2668            .clone(),
2669    );
2670    let theta_dim = rho_dim + log_kappa0.len();
2671    let mut theta0 = Array1::<f64>::zeros(theta_dim);
2672    theta0
2673        .slice_mut(s![0..rho_dim])
2674        .assign(&baseline_log_lambdas);
2675    theta0
2676        .slice_mut(s![rho_dim..theta_dim])
2677        .assign(log_kappa0.as_array());
2678
2679    let mut lower = Array1::<f64>::from_elem(theta_dim, -RHO_BOUND);
2680    let mut upper = Array1::<f64>::from_elem(theta_dim, RHO_BOUND);
2681    lower
2682        .slice_mut(s![rho_dim..theta_dim])
2683        .assign(log_kappa_lower.as_array());
2684    upper
2685        .slice_mut(s![rho_dim..theta_dim])
2686        .assign(log_kappa_upper.as_array());
2687
2688    let pilot_spec_cloned = pilot_spec.clone();
2689    let pilot_beta = baseline_eta_beta;
2690    let wiggle_design = wiggle_block.design.clone();
2691    let wiggle_offset = wiggle_block.offset.clone();
2692    let wiggle_penalties = wiggle_block.penalties.clone();
2693    let wiggle_initial_beta = baseline_wiggle_beta;
2694    let wiggle_knots_cloned = wiggle_knots.clone();
2695    let y_cloned = y.clone();
2696    let weights_cloned = weights.clone();
2697    let link_kind_cloned = link_kind.clone();
2698    let outer_family = BinomialMeanWiggleFamily {
2699        y: y_cloned.clone(),
2700        weights: weights_cloned.clone(),
2701        link_kind: link_kind_cloned.clone(),
2702        wiggle_knots: wiggle_knots_cloned.clone(),
2703        wiggle_degree,
2704        policy: gam_runtime::resource::ResourcePolicy::default_library(),
2705    };
2706    let screening_cap = Arc::new(AtomicUsize::new(0));
2707    let mut outer_options = options.clone();
2708    outer_options.screening_max_inner_iterations = Some(Arc::clone(&screening_cap));
2709    struct MeanWiggleOuterState {
2710        pub(crate) warm_cache: Option<crate::custom_family::CustomFamilyWarmStart>,
2711        pub(crate) last_eval: Option<(
2712            Array1<f64>,
2713            f64,
2714            Array1<f64>,
2715            gam_problem::HessianResult,
2716            crate::custom_family::CustomFamilyWarmStart,
2717        )>,
2718    }
2719
2720    let build_realized_blocks = |theta: &Array1<f64>| -> Result<
2721        (
2722            TermCollectionSpec,
2723            TermCollectionDesign,
2724            Vec<ParameterBlockSpec>,
2725            Vec<CustomFamilyBlockPsiDerivative>,
2726        ),
2727        String,
2728    > {
2729        let log_kappa =
2730            SpatialLogKappaCoords::from_theta_tail_with_dims(theta, rho_dim, dims_per_term.clone());
2731        let resolvedspec = log_kappa
2732            .apply_tospec(&pilot_spec_cloned, &spatial_terms)
2733            .map_err(|e| e.to_string())?;
2734        let design =
2735            build_term_collection_design(data, &resolvedspec).map_err(|e| e.to_string())?;
2736        let eta_derivs = build_block_spatial_psi_derivatives(data, &resolvedspec, &design)?
2737            .ok_or_else(|| {
2738                "missing eta spatial psi derivatives for binomial mean wiggle".to_string()
2739            })?;
2740        let blocks = vec![
2741            ParameterBlockSpec {
2742                name: "eta".to_string(),
2743                design: design.design.clone(),
2744                offset: Array1::zeros(y_cloned.len()),
2745                penalties: design.penalties_as_penalty_matrix(),
2746                nullspace_dims: vec![],
2747                initial_log_lambdas: theta.slice(s![0..eta_penalty_count]).to_owned(),
2748                initial_beta: Some(pilot_beta.clone()),
2749                // Lower gauge priority on the static eta design: it yields the
2750                // shared level/intercept direction to the dynamic full-width
2751                // wiggle I-spline block (see fit_binomial_mean_wiggle).
2752                gauge_priority: LINK_WIGGLE_GAUGE_PRIORITY,
2753                jacobian_callback: None,
2754                stacked_design: None,
2755                stacked_offset: None,
2756            },
2757            ParameterBlockSpec {
2758                name: "wiggle".to_string(),
2759                design: wiggle_design.clone(),
2760                offset: wiggle_offset.clone(),
2761                penalties: {
2762                    let p_wiggle = wiggle_design.ncols();
2763                    wiggle_penalties
2764                        .iter()
2765                        .map(|spec| match spec {
2766                            crate::model_types::PenaltySpec::Block {
2767                                local, col_range, ..
2768                            } => PenaltyMatrix::Blockwise {
2769                                local: local.clone(),
2770                                col_range: col_range.clone(),
2771                                total_dim: p_wiggle,
2772                            },
2773                            crate::model_types::PenaltySpec::Dense(m)
2774                            | crate::model_types::PenaltySpec::DenseWithMean {
2775                                matrix: m, ..
2776                            } => PenaltyMatrix::Dense(m.clone()),
2777                        })
2778                        .collect()
2779                },
2780                nullspace_dims: vec![],
2781                initial_log_lambdas: theta.slice(s![eta_penalty_count..rho_dim]).to_owned(),
2782                initial_beta: wiggle_initial_beta.clone(),
2783                gauge_priority: DEFAULT_GAUGE_PRIORITY,
2784                jacobian_callback: None,
2785                stacked_design: None,
2786                stacked_offset: None,
2787            },
2788        ];
2789        Ok((resolvedspec, design, blocks, eta_derivs))
2790    };
2791
2792    let build_eval = |theta: &Array1<f64>,
2793                      warm_cache: Option<&crate::custom_family::CustomFamilyWarmStart>,
2794                      need_hessian: bool|
2795     -> Result<
2796        (
2797            crate::custom_family::CustomFamilyJointHyperResult,
2798            TermCollectionSpec,
2799            TermCollectionDesign,
2800        ),
2801        String,
2802    > {
2803        let (resolvedspec, design, blocks, eta_derivs) = build_realized_blocks(theta)?;
2804        let eval = evaluate_custom_family_joint_hyper(
2805            &outer_family,
2806            &blocks,
2807            &outer_options,
2808            &theta.slice(s![0..rho_dim]).to_owned(),
2809            &[eta_derivs, Vec::new()],
2810            warm_cache,
2811            if need_hessian {
2812                gam_problem::EvalMode::ValueGradientHessian
2813            } else {
2814                gam_problem::EvalMode::ValueAndGradient
2815            },
2816        )?;
2817        Ok((eval, resolvedspec, design))
2818    };
2819
2820    let build_efs = |theta: &Array1<f64>,
2821                     warm_cache: Option<&crate::custom_family::CustomFamilyWarmStart>|
2822     -> Result<crate::custom_family::CustomFamilyJointHyperEfsResult, String> {
2823        let (_, _, blocks, eta_derivs) = build_realized_blocks(theta)?;
2824        evaluate_custom_family_joint_hyper_efs(
2825            &outer_family,
2826            &blocks,
2827            &outer_options,
2828            &theta.slice(s![0..rho_dim]).to_owned(),
2829            &[eta_derivs, Vec::new()],
2830            warm_cache,
2831        )
2832        .map_err(|e| e.to_string())
2833    };
2834
2835    use crate::model_types::EstimationError;
2836    use gam_solve::rho_optimizer::OuterEvalOrder;
2837    use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
2838
2839    // Exact first-order AND second-order [rho, psi] calculus is available
2840    // for all inverse links via the shared jet formulas plus the generic
2841    // exact-Newton D_βH / D²_βH closures routed through
2842    // evaluate_custom_family_joint_hyper -> joint_outer_evaluate ->
2843    // BorrowedJointDerivProvider. This enables the analytic-Hessian outer
2844    // plan for REML optimization instead of the downgraded gradient-only
2845    // outer strategies.
2846    //
2847    // Spatial log-kappa coordinates are ψ (design-moving) dimensions because
2848    // they rebuild the spatial basis and penalties at each outer proposal.
2849    let analytic_outer_hessian_available = true;
2850    let mut seed_heuristic = theta0.to_vec();
2851    for value in &mut seed_heuristic[..rho_dim] {
2852        *value = value.exp();
2853    }
2854    let problem = gam_solve::rho_optimizer::OuterProblem::new(theta_dim)
2855        .with_gradient(Derivative::Analytic)
2856        .with_hessian(if analytic_outer_hessian_available {
2857            DeclaredHessianForm::Either
2858        } else {
2859            DeclaredHessianForm::Unavailable
2860        })
2861        .with_psi_dim(theta_dim - rho_dim)
2862        .with_tolerance(options.outer_tol)
2863        .with_max_iter(options.outer_max_iter)
2864        .with_bounds(lower.clone(), upper.clone())
2865        .with_initial_rho(theta0.clone())
2866        .with_seed_config(crate::seeding::SeedConfig {
2867            max_seeds: 4,
2868            seed_budget: 2,
2869            risk_profile: crate::seeding::SeedRiskProfile::GeneralizedLinear,
2870            num_auxiliary_trailing: theta_dim - rho_dim,
2871            ..Default::default()
2872        })
2873        .with_screening_cap(Arc::clone(&screening_cap))
2874        .with_rho_bound(12.0)
2875        .with_heuristic_lambdas(seed_heuristic);
2876
2877    let eval_outer = |state: &mut MeanWiggleOuterState,
2878                      theta: &Array1<f64>,
2879                      order: OuterEvalOrder|
2880     -> Result<OuterEval, EstimationError> {
2881        if let Some((cached_theta, cached_cost, cached_grad, cached_hess, cached_warm)) =
2882            &state.last_eval
2883            && cached_theta == theta
2884            && (!matches!(order, OuterEvalOrder::ValueGradientHessian)
2885                || matches!(
2886                    cached_hess,
2887                    gam_problem::HessianResult::Analytic(_)
2888                        | gam_problem::HessianResult::Operator(_)
2889                ))
2890        {
2891            state.warm_cache = Some(cached_warm.clone());
2892            return Ok(OuterEval {
2893                cost: *cached_cost,
2894                gradient: cached_grad.clone(),
2895                hessian: cached_hess.clone(),
2896                inner_beta_hint: None,
2897            });
2898        }
2899        let need_hessian = matches!(order, OuterEvalOrder::ValueGradientHessian)
2900            && analytic_outer_hessian_available;
2901        let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), need_hessian)
2902            .map_err(EstimationError::InvalidInput)?;
2903        if !eval.inner_converged {
2904            state.warm_cache = Some(eval.warm_start);
2905            crate::bail_invalid_estim!(
2906                "binomial mean-wiggle exact spatial inner solve did not converge"
2907            );
2908        }
2909        let hessian_result = eval.outer_hessian.clone();
2910        state.last_eval = Some((
2911            theta.clone(),
2912            eval.objective,
2913            eval.gradient.clone(),
2914            eval.outer_hessian.clone(),
2915            eval.warm_start.clone(),
2916        ));
2917        state.warm_cache = Some(eval.warm_start);
2918        Ok(OuterEval {
2919            cost: eval.objective,
2920            gradient: eval.gradient,
2921            hessian: hessian_result,
2922            inner_beta_hint: None,
2923        })
2924    };
2925
2926    let mut obj = problem.build_objective_with_screening_proxy(
2927        MeanWiggleOuterState {
2928            warm_cache: None,
2929            last_eval: None,
2930        },
2931        |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
2932            if let Some((cached_theta, cached_cost, _, _, cached_warm)) = &state.last_eval
2933                && cached_theta == theta
2934            {
2935                state.warm_cache = Some(cached_warm.clone());
2936                return Ok(*cached_cost);
2937            }
2938            let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), false)
2939                .map_err(EstimationError::InvalidInput)?;
2940            if !eval.inner_converged {
2941                state.warm_cache = Some(eval.warm_start);
2942                crate::bail_invalid_estim!(
2943                    "binomial mean-wiggle exact spatial cost inner solve did not converge"
2944                        .to_string(),
2945                );
2946            }
2947            state.warm_cache = Some(eval.warm_start);
2948            Ok(eval.objective)
2949        },
2950        |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
2951            eval_outer(
2952                state,
2953                theta,
2954                if analytic_outer_hessian_available {
2955                    OuterEvalOrder::ValueGradientHessian
2956                } else {
2957                    OuterEvalOrder::ValueAndGradient
2958                },
2959            )
2960        },
2961        |state: &mut MeanWiggleOuterState, theta: &Array1<f64>, order: OuterEvalOrder| {
2962            eval_outer(state, theta, order)
2963        },
2964        Some(|state: &mut MeanWiggleOuterState| {
2965            state.warm_cache = None;
2966            state.last_eval = None;
2967        }),
2968        Some(|state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
2969            let eval = build_efs(theta, state.warm_cache.as_ref())
2970                .map_err(EstimationError::InvalidInput)?;
2971            if !eval.inner_converged {
2972                state.warm_cache = Some(eval.warm_start);
2973                crate::bail_invalid_estim!(
2974                    "binomial mean-wiggle exact spatial EFS inner solve did not converge"
2975                        .to_string(),
2976                );
2977            }
2978            state.warm_cache = Some(eval.warm_start);
2979            Ok(eval.efs_eval)
2980        }),
2981        // Seed-screening ranking proxy (#969). The cost closure above
2982        // hard-errors on a non-converged inner solve — correct for
2983        // line-search costs, but under the screening cap (wired into the
2984        // outer options and installed by the cascade) the inner solve is
2985        // truncated BY DESIGN, so screening through it rejects every seed
2986        // — the all-seeds-rejected front-door genus. Screening only RANKS
2987        // candidates: the truncated solve's penalized objective is the
2988        // ranking signal; convergence is demanded of the selected seed's
2989        // full-budget fit, not of capped probes.
2990        |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
2991            if let Some((cached_theta, cached_cost, _, _, cached_warm)) = &state.last_eval
2992                && cached_theta == theta
2993            {
2994                state.warm_cache = Some(cached_warm.clone());
2995                return Ok(*cached_cost);
2996            }
2997            let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), false)
2998                .map_err(EstimationError::InvalidInput)?;
2999            state.warm_cache = Some(eval.warm_start);
3000            Ok(eval.objective)
3001        },
3002    );
3003
3004    let outer = problem
3005        .run(&mut obj, "binomial mean wiggle exact spatial hyper")
3006        .map_err(|e| e.to_string())?;
3007    if !outer.converged {
3008        return Err(GamlssError::NumericalFailure { reason: format!(
3009            "binomial mean wiggle exact spatial hyper did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
3010            outer.iterations,
3011            outer.final_value,
3012            outer.final_grad_norm_report(),
3013        ) }.into());
3014    }
3015    let theta_star = outer.rho;
3016
3017    let log_kappa =
3018        SpatialLogKappaCoords::from_theta_tail_with_dims(&theta_star, rho_dim, dims_per_term);
3019    let resolvedspec = log_kappa
3020        .apply_tospec(&pilot_spec_cloned, &spatial_terms)
3021        .map_err(|e| e.to_string())?;
3022    let design = build_term_collection_design(data, &resolvedspec).map_err(|e| e.to_string())?;
3023    let resolvedspec =
3024        freeze_term_collection_from_design(&resolvedspec, &design).map_err(|e| e.to_string())?;
3025    let fit = fit_binomial_mean_wiggle(
3026        BinomialMeanWiggleSpec {
3027            y: y_cloned,
3028            weights: weights_cloned,
3029            link_kind: link_kind_cloned,
3030            wiggle_knots: wiggle_knots.clone(),
3031            wiggle_degree,
3032            eta_block: ParameterBlockInput {
3033                design: design.design.clone(),
3034                offset: Array1::zeros(y.len()),
3035                penalties: design
3036                    .penalties
3037                    .iter()
3038                    .map(crate::model_types::PenaltySpec::from_blockwise_ref)
3039                    .collect(),
3040                nullspace_dims: vec![],
3041                initial_log_lambdas: Some(theta_star.slice(s![0..eta_penalty_count]).to_owned()),
3042                initial_beta: Some(pilot_beta),
3043            },
3044            wiggle_block: ParameterBlockInput {
3045                design: wiggle_design,
3046                offset: wiggle_offset,
3047                penalties: wiggle_penalties,
3048                nullspace_dims: vec![],
3049                initial_log_lambdas: Some(
3050                    theta_star.slice(s![eta_penalty_count..rho_dim]).to_owned(),
3051                ),
3052                initial_beta: wiggle_initial_beta,
3053            },
3054        },
3055        options,
3056    )?;
3057
3058    Ok(BinomialMeanWiggleTermFitResult {
3059        fit,
3060        resolvedspec,
3061        design,
3062        wiggle_knots,
3063        wiggle_degree,
3064    })
3065}