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    gam_problem::validate_log_strengths(lambdas.iter().copied())
151        .map_err(|error| format!("initial_log_lambdas: {error}"))?;
152    Ok(lambdas)
153}
154
155fn fitted_log_lambdas(lambdas: &Array1<f64>, context: &str) -> Result<Array1<f64>, String> {
156    lambdas
157        .iter()
158        .copied()
159        .enumerate()
160        .map(|(coordinate, value)| {
161            gam_problem::checked_log_strength(value)
162                .map_err(|error| format!("{context} coordinate {coordinate}: {error}"))
163        })
164        .collect::<Result<Vec<_>, _>>()
165        .map(Array1::from_vec)
166}
167
168pub(crate) fn build_two_block_exact_joint_setup(
169    data: ArrayView2<'_, f64>,
170    meanspec: &TermCollectionSpec,
171    noisespec: &TermCollectionSpec,
172    mean_penalties: usize,
173    noise_penalties: usize,
174    extra_rho0: &[f64],
175    rho0_override: Option<&Array1<f64>>,
176    kappa_options: &SpatialLengthScaleOptimizationOptions,
177) -> Result<ExactJointHyperSetup, gam_terms::basis::BasisError> {
178    // GAMLSS-specific part: assemble the rho seed in [mean | noise | extra]
179    // penalty order, honoring a caller override when it matches the layout.
180    let rho_dim = mean_penalties + noise_penalties + extra_rho0.len();
181    let mut rho0vec = Array1::<f64>::zeros(rho_dim);
182    if let Some(rho0) = rho0_override.filter(|rho0| rho0.len() == rho_dim) {
183        rho0vec.assign(rho0);
184    } else {
185        for (i, &rho_init) in extra_rho0.iter().enumerate() {
186            rho0vec[mean_penalties + noise_penalties + i] = rho_init;
187        }
188    }
189
190    // Generic part: per-block log(kappa) seed/bounds and exact-joint assembly,
191    // with the two linear predictors (mean, noise) in theta order.
192    build_location_scale_exact_joint_setup(data, &[meanspec, noisespec], rho0vec, kappa_options)
193}
194
195pub(crate) fn gaussian_location_scalewarm_start(
196    y: &Array1<f64>,
197    weights: &Array1<f64>,
198    mu_block: &ParameterBlockSpec,
199    log_sigma_block: &ParameterBlockSpec,
200    ridge_floor: f64,
201    mean_beta_hint: Option<&Array1<f64>>,
202    noise_beta_hint: Option<&Array1<f64>>,
203) -> Result<(Array1<f64>, Array1<f64>, f64), String> {
204    let betamu = if let Some(beta) = mean_beta_hint {
205        beta.clone()
206    } else {
207        solve_penalizedweighted_projection(
208            &mu_block.design,
209            &mu_block.offset,
210            y,
211            weights,
212            &mu_block.penalties,
213            &mu_block.initial_log_lambdas,
214            ridge_floor,
215        )?
216    };
217    let mut mu_hat = mu_block.solver_design().matrixvectormultiply(&betamu);
218    mu_hat += mu_block.solver_offset();
219    let mut weighted_ss = 0.0;
220    let mut weight_sum = 0.0;
221    for i in 0..y.len() {
222        let wi = weights[i].max(0.0);
223        let resid = y[i] - mu_hat[i];
224        weighted_ss += wi * resid * resid;
225        weight_sum += wi;
226    }
227    if !weighted_ss.is_finite() || !weight_sum.is_finite() || weight_sum <= 0.0 {
228        return Err(
229            "gaussian location-scale warm start could not estimate residual scale".to_string(),
230        );
231    }
232    // Warm-start σ̂ must clear the logb floor so the inverse link
233    //   η = log(σ − b)
234    // is finite. Use a relative cushion above b so the warm-start is in the
235    // smooth interior of the link domain.
236    let sigma_hat = (weighted_ss / weight_sum)
237        .sqrt()
238        .max(LOGB_SIGMA_FLOOR * 1.5);
239    let beta_log_sigma = if let Some(beta) = noise_beta_hint {
240        beta.clone()
241    } else {
242        let eta_sigma = (sigma_hat - LOGB_SIGMA_FLOOR).ln();
243        let sigma_target = Array1::from_elem(y.len(), eta_sigma);
244        solve_penalizedweighted_projection(
245            &log_sigma_block.design,
246            &log_sigma_block.offset,
247            &sigma_target,
248            weights,
249            &log_sigma_block.penalties,
250            &log_sigma_block.initial_log_lambdas,
251            ridge_floor,
252        )?
253    };
254    Ok((betamu, beta_log_sigma, sigma_hat))
255}
256
257/// Total output count for every two-block location-scale family in this
258/// module (mu/log_sigma or threshold/log_sigma). The wiggle variants add a
259/// third zero-channel block but still drive only two output channels.
260pub(crate) const LOCATION_SCALE_N_OUTPUTS: usize = 2;
261
262/// Construct a fully wired location-scale parameter block.
263///
264/// This is the **only** way to build a LocationScale `ParameterBlockSpec` in
265/// this module — by construction the `AdditiveBlockJacobian` callback is
266/// always installed, so the channel-aware identifiability audit cannot be
267/// silently bypassed by a future `build_blocks` impl that forgets to wire
268/// the callback at the tail (re-introducing #319).
269///
270/// `own_output` is the zero-based output channel this block drives
271/// (e.g. 0 for `mu`/`threshold`, 1 for `log_sigma`). `n_family_outputs` is
272/// fixed at [`LOCATION_SCALE_N_OUTPUTS`] for every two-block family here
273/// but is exposed so the helper composes cleanly with any future
274/// k-block extension.
275pub(crate) fn build_location_scale_block(
276    name: impl Into<String>,
277    design: DesignMatrix,
278    offset: Array1<f64>,
279    penalties: Vec<PenaltyMatrix>,
280    nullspace_dims: Vec<usize>,
281    initial_log_lambdas: Array1<f64>,
282    initial_beta: Option<Array1<f64>>,
283    own_output: usize,
284    n_family_outputs: usize,
285    caller: &str,
286) -> Result<ParameterBlockSpec, String> {
287    if own_output >= n_family_outputs {
288        return Err(format!(
289            "{caller}: own_output={own_output} >= n_family_outputs={n_family_outputs}"
290        ));
291    }
292    let mut spec = ParameterBlockSpec {
293        name: name.into(),
294        design,
295        offset,
296        penalties,
297        nullspace_dims,
298        initial_log_lambdas,
299        initial_beta,
300        gauge_priority: 100,
301        jacobian_callback: None,
302        stacked_design: None,
303        stacked_offset: None,
304    };
305    let dense = spec.effective_design(caller)?;
306    spec.jacobian_callback = Some(std::sync::Arc::new(AdditiveBlockJacobian {
307        design: dense,
308        own_output,
309        n_family_outputs,
310    }));
311    Ok(spec)
312}
313
314/// Construct the wiggle block that accompanies a two-block location-scale
315/// family. The wiggle modulates the inverse link nonlinearly and
316/// contributes no linear effective Jacobian — the installed callback
317/// therefore exposes a zero `(n × p_w)` design under
318/// `n_family_outputs = LOCATION_SCALE_N_OUTPUTS`.
319pub(crate) fn build_location_scale_wiggle_block(
320    name: impl Into<String>,
321    design: DesignMatrix,
322    offset: Array1<f64>,
323    penalties: Vec<PenaltyMatrix>,
324    nullspace_dims: Vec<usize>,
325    initial_log_lambdas: Array1<f64>,
326    initial_beta: Option<Array1<f64>>,
327    n_rows: usize,
328) -> Result<ParameterBlockSpec, String> {
329    let p_w = design.ncols();
330    let mut spec = ParameterBlockSpec {
331        name: name.into(),
332        design,
333        offset,
334        penalties,
335        nullspace_dims,
336        initial_log_lambdas,
337        initial_beta,
338        gauge_priority: 100,
339        jacobian_callback: None,
340        stacked_design: None,
341        stacked_offset: None,
342    };
343    spec.jacobian_callback = Some(std::sync::Arc::new(AdditiveBlockJacobian {
344        design: ndarray::Array2::<f64>::zeros((n_rows, p_w)),
345        own_output: 0,
346        n_family_outputs: LOCATION_SCALE_N_OUTPUTS,
347    }));
348    Ok(spec)
349}
350
351pub(crate) fn prepared_gaussian_log_sigma_design(
352    mu_design: &DesignMatrix,
353    log_sigma_design: &DesignMatrix,
354) -> Result<DesignMatrix, String> {
355    if mu_design.nrows() != log_sigma_design.nrows() {
356        return Err(GamlssError::DimensionMismatch {
357            reason: format!(
358                "gaussian log-sigma design row mismatch: mean rows={}, log_sigma rows={}",
359                mu_design.nrows(),
360                log_sigma_design.nrows()
361            ),
362        }
363        .into());
364    }
365    // Gaussian location-scale remains identifiable even when μ and log σ use
366    // the same covariate basis:
367    //
368    //   L(μ, η) = 0.5 * Σ_i [ (y_i - μ_i)^2 exp(-2η_i) + 2η_i ],
369    //   μ = X_μ β_μ,  η = X_σ β_σ.
370    //
371    // Shared columns are not a frame mismatch. β_μ and β_σ enter through
372    // different sufficient statistics (residual and residual²), so replacing
373    // X_σ with (I - P_{X_μ}) X_σ would impose an extra constraint and can
374    // erase real heteroscedastic signal when the two blocks share a basis.
375    Ok(log_sigma_design.clone())
376}
377
378pub(crate) fn identified_binomial_log_sigma_design(
379    threshold_design: &TermCollectionDesign,
380    log_sigma_design: &TermCollectionDesign,
381    weights: &Array1<f64>,
382) -> Result<DesignMatrix, String> {
383    let non_intercept_start = log_sigma_design
384        .intercept_range
385        .end
386        .min(log_sigma_design.design.ncols());
387    let transform = build_scale_deviation_transform_design(
388        &threshold_design.design,
389        &log_sigma_design.design,
390        weights,
391        non_intercept_start,
392    )?;
393    build_scale_deviation_operator(
394        threshold_design.design.clone(),
395        log_sigma_design.design.clone(),
396        &transform,
397    )
398}
399
400pub(crate) fn identity_penalty(dim: usize) -> Array2<f64> {
401    let mut penalty = Array2::<f64>::zeros((dim, dim));
402    for i in 0..dim {
403        penalty[[i, i]] = 1.0;
404    }
405    penalty
406}
407
408pub(crate) fn append_binomial_log_sigma_shrinkage_penalty_design(
409    design: &mut TermCollectionDesign,
410) {
411    let p = design.design.ncols();
412    design
413        .penalties
414        .push(BlockwisePenalty::new(0..p, identity_penalty(p)));
415    // Identity penalty penalizes the full space → nullspace dimension is 0.
416    design.nullspace_dims.push(0);
417    design.penaltyinfo.push(PenaltyBlockInfo {
418        global_index: design.penaltyinfo.len(),
419        termname: Some("log_sigma_shrinkage".to_string()),
420        penalty: ActivePenaltyInfo {
421            source: PenaltySource::Other("shrinkage".to_string()),
422            original_index: 0,
423            effective_rank: p,
424            normalization_scale: 1.0,
425            kronecker_factors: None,
426        },
427    });
428}
429
430/// Build the (mean, log-σ) parameter-block pair for a Gaussian location-scale
431/// family. Shared verbatim by the non-wiggle and wiggle Gaussian builders so the
432/// scale-block construction — prepared log-σ design, formula-native penalties,
433/// and the joint Gaussian warm start — lives in exactly one place. Callers
434/// supply the per-block log-λ vectors sliced from their own layout (two-block vs
435/// with-wiggle) and append any extra blocks.
436pub(crate) fn build_gaussian_mean_and_scale_blocks(
437    y: &Array1<f64>,
438    weights: &Array1<f64>,
439    mean_design: &TermCollectionDesign,
440    noise_design: &TermCollectionDesign,
441    mean_offset: &Array1<f64>,
442    noise_offset: &Array1<f64>,
443    mean_log_lambdas: Array1<f64>,
444    noise_log_lambdas: Array1<f64>,
445    mean_beta_hint: Option<Array1<f64>>,
446    noise_beta_hint: Option<Array1<f64>>,
447    context: &str,
448) -> Result<(ParameterBlockSpec, ParameterBlockSpec), String> {
449    let mean_offset = mean_design
450        .compose_offset(mean_offset.view(), &format!("{context}: mu"))
451        .map_err(|error| error.to_string())?;
452    let noise_offset = noise_design
453        .compose_offset(noise_offset.view(), &format!("{context}: log_sigma"))
454        .map_err(|error| error.to_string())?;
455    let mut meanspec = build_location_scale_block(
456        "mu",
457        mean_design.design.clone(),
458        mean_offset,
459        mean_design.penalties_as_penalty_matrix(),
460        mean_design.nullspace_dims.clone(),
461        mean_log_lambdas,
462        mean_beta_hint,
463        0,
464        LOCATION_SCALE_N_OUTPUTS,
465        &format!("{context}: mu"),
466    )?;
467    let prepared_noise_design =
468        prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design)?;
469    // The formula-native penalty topology is authoritative. Smooth terms carry
470    // their own REML-selected null-space penalty when `double_penalty=true`
471    // (the default), while an explicit `double_penalty=false` remains a real
472    // opt-out. In particular the global log-σ intercept is likelihood-identified
473    // and must stay unpenalized: adding a Gaussian-only projector over the joint
474    // null space placed a data-scale-dependent prior on the overall σ level and
475    // introduced an extra smoothing coordinate absent from the formula (#1561).
476    let mut noisespec = build_location_scale_block(
477        "log_sigma",
478        prepared_noise_design,
479        noise_offset,
480        noise_design.penalties_as_penalty_matrix(),
481        noise_design.nullspace_dims.clone(),
482        noise_log_lambdas,
483        noise_beta_hint,
484        1,
485        LOCATION_SCALE_N_OUTPUTS,
486        &format!("{context}: log_sigma"),
487    )?;
488    if meanspec.initial_beta.is_none() || noisespec.initial_beta.is_none() {
489        let (betamu0, beta_ls0, _) = gaussian_location_scalewarm_start(
490            y,
491            weights,
492            &meanspec,
493            &noisespec,
494            1e-10,
495            meanspec.initial_beta.as_ref(),
496            noisespec.initial_beta.as_ref(),
497        )?;
498        if meanspec.initial_beta.is_none() {
499            meanspec.initial_beta = Some(betamu0);
500        }
501        if noisespec.initial_beta.is_none() {
502            noisespec.initial_beta = Some(beta_ls0);
503        }
504    }
505    Ok((meanspec, noisespec))
506}
507
508/// Build the (threshold, log-σ) parameter-block pair for a Binomial
509/// location-scale family. Shared by the non-wiggle and wiggle Binomial builders;
510/// mirrors [`build_gaussian_mean_and_scale_blocks`] but with the binomial-
511/// identified log-σ design, the link-aware joint warm start, and the same
512/// REML-selected full-span scale shrinkage penalty.
513pub(crate) fn build_binomial_threshold_and_scale_blocks(
514    y: &Array1<f64>,
515    weights: &Array1<f64>,
516    link_kind: &InverseLink,
517    mean_design: &TermCollectionDesign,
518    noise_design: &TermCollectionDesign,
519    mean_offset: &Array1<f64>,
520    noise_offset: &Array1<f64>,
521    mean_log_lambdas: Array1<f64>,
522    noise_log_lambdas: Array1<f64>,
523    mean_beta_hint: Option<Array1<f64>>,
524    noise_beta_hint: Option<Array1<f64>>,
525    context: &str,
526) -> Result<(ParameterBlockSpec, ParameterBlockSpec), String> {
527    let mean_offset = mean_design
528        .compose_offset(mean_offset.view(), &format!("{context}: threshold"))
529        .map_err(|error| error.to_string())?;
530    let noise_offset = noise_design
531        .compose_offset(noise_offset.view(), &format!("{context}: log_sigma"))
532        .map_err(|error| error.to_string())?;
533    let identifiednoise_design =
534        identified_binomial_log_sigma_design(mean_design, noise_design, weights)?;
535    let p_noise = identifiednoise_design.ncols();
536    let mut log_sigma_penalty_matrices: Vec<PenaltyMatrix> =
537        noise_design.penalties_as_penalty_matrix();
538    log_sigma_penalty_matrices.push(PenaltyMatrix::Dense(identity_penalty(p_noise)));
539    let mut thresholdspec = build_location_scale_block(
540        "threshold",
541        mean_design.design.clone(),
542        mean_offset,
543        mean_design.penalties_as_penalty_matrix(),
544        vec![],
545        mean_log_lambdas,
546        mean_beta_hint,
547        0,
548        LOCATION_SCALE_N_OUTPUTS,
549        &format!("{context}: threshold"),
550    )?;
551    let mut log_sigmaspec = build_location_scale_block(
552        "log_sigma",
553        identifiednoise_design,
554        noise_offset,
555        log_sigma_penalty_matrices,
556        vec![],
557        noise_log_lambdas,
558        noise_beta_hint,
559        1,
560        LOCATION_SCALE_N_OUTPUTS,
561        &format!("{context}: log_sigma"),
562    )?;
563    if thresholdspec.initial_beta.is_none() || log_sigmaspec.initial_beta.is_none() {
564        let (beta_t0, beta_ls0) = binomial_location_scalewarm_start(
565            y,
566            weights,
567            link_kind,
568            &thresholdspec,
569            &log_sigmaspec,
570            thresholdspec.initial_beta.as_ref(),
571            log_sigmaspec.initial_beta.as_ref(),
572        )?;
573        if thresholdspec.initial_beta.is_none() {
574            thresholdspec.initial_beta = Some(beta_t0);
575        }
576        if log_sigmaspec.initial_beta.is_none() {
577            log_sigmaspec.initial_beta = Some(beta_ls0);
578        }
579    }
580    Ok((thresholdspec, log_sigmaspec))
581}
582
583/// Convert a wiggle block's `PenaltySpec`s into the `PenaltyMatrix` list the
584/// location-scale wiggle block expects. Shared by the Gaussian and Binomial
585/// wiggle builders, which previously inlined the identical match.
586pub(crate) fn wiggle_block_penalty_matrices(
587    wiggle_block: &ParameterBlockInput,
588) -> Vec<PenaltyMatrix> {
589    let p_wiggle = wiggle_block.design.ncols();
590    wiggle_block
591        .penalties
592        .iter()
593        .map(|spec| match spec {
594            crate::model_types::PenaltySpec::Block {
595                local, col_range, ..
596            } => PenaltyMatrix::Blockwise {
597                local: local.clone(),
598                col_range: col_range.clone(),
599                total_dim: p_wiggle,
600            },
601            crate::model_types::PenaltySpec::Dense(m)
602            | crate::model_types::PenaltySpec::DenseWithMean { matrix: m, .. } => {
603                PenaltyMatrix::Dense(m.clone())
604            }
605        })
606        .collect()
607}
608
609pub(crate) fn binomial_location_scale_link_eta_from_probability(
610    link_kind: &InverseLink,
611    probability: f64,
612) -> Result<f64, String> {
613    let target = probability.clamp(1e-6, 1.0 - 1e-6);
614    match link_kind {
615        InverseLink::Standard(StandardLink::Logit) => Ok((target / (1.0 - target)).ln()),
616        InverseLink::Standard(StandardLink::Probit) => standard_normal_quantile(target)
617            .map_err(|err| format!("failed to invert probit warm-start probability: {err}")),
618        InverseLink::Standard(StandardLink::CLogLog) => Ok((-((1.0 - target).ln())).ln()),
619        other => Err(GamlssError::UnsupportedConfiguration { reason: format!(
620            "binomial location-scale warm start requires logit, probit, or cloglog link, got {other:?}"
621        ) }.into()),
622    }
623}
624
625pub(crate) fn weighted_binomial_prevalence(
626    y: &Array1<f64>,
627    weights: &Array1<f64>,
628) -> Result<f64, String> {
629    if y.len() != weights.len() {
630        return Err(GamlssError::DimensionMismatch { reason: format!(
631            "binomial location-scale warm start dimension mismatch: y has length {}, weights have length {}",
632            y.len(),
633            weights.len()
634        ) }.into());
635    }
636    let mut weight_sum = 0.0;
637    let mut success_sum = 0.0;
638    for (&yi, &wi) in y.iter().zip(weights.iter()) {
639        if !yi.is_finite() {
640            return Err(GamlssError::NonFinite {
641                reason: format!(
642                    "binomial location-scale warm start encountered non-finite response {yi}"
643                ),
644            }
645            .into());
646        }
647        if !wi.is_finite() || wi < 0.0 {
648            return Err(GamlssError::InvalidInput {
649                reason: format!(
650                    "binomial location-scale warm start requires finite non-negative weights; weight={wi}"
651                ),
652            }
653            .into());
654        }
655        if wi > 0.0 {
656            weight_sum += wi;
657            success_sum += wi * yi;
658        }
659    }
660    if !weight_sum.is_finite() || weight_sum <= 0.0 {
661        return Err(
662            "binomial location-scale warm start requires positive total weight".to_string(),
663        );
664    }
665    Ok(success_sum / weight_sum)
666}
667
668pub(crate) fn project_constant_eta_into_block(
669    block: &ParameterBlockSpec,
670    weights: &Array1<f64>,
671    eta: f64,
672) -> Result<Array1<f64>, String> {
673    let target_eta = Array1::from_elem(block.design.nrows(), eta);
674    solve_penalizedweighted_projection(
675        &block.design,
676        &block.offset,
677        &target_eta,
678        weights,
679        &block.penalties,
680        &block.initial_log_lambdas,
681        1e-10,
682    )
683}
684
685// Deterministic warm start for the binomial location-scale model. This stays
686// out of the optimizer: it projects a prevalence-matched threshold and neutral
687// log-sigma value into the actual penalized block spaces.
688pub(crate) fn binomial_location_scalewarm_start(
689    y: &Array1<f64>,
690    weights: &Array1<f64>,
691    link_kind: &InverseLink,
692    threshold_block: &ParameterBlockSpec,
693    log_sigma_block: &ParameterBlockSpec,
694    mean_beta_hint: Option<&Array1<f64>>,
695    noise_beta_hint: Option<&Array1<f64>>,
696) -> Result<(Array1<f64>, Array1<f64>), String> {
697    if let (Some(mean_beta), Some(noise_beta)) = (mean_beta_hint, noise_beta_hint) {
698        return Ok((mean_beta.clone(), noise_beta.clone()));
699    }
700
701    let beta_threshold = match mean_beta_hint {
702        Some(beta) => beta.clone(),
703        None => {
704            let prevalence = weighted_binomial_prevalence(y, weights)?;
705            let eta = binomial_location_scale_link_eta_from_probability(link_kind, prevalence)?;
706            project_constant_eta_into_block(threshold_block, weights, eta)?
707        }
708    };
709    let beta_log_sigma = match noise_beta_hint {
710        Some(beta) => beta.clone(),
711        None => project_constant_eta_into_block(log_sigma_block, weights, 0.0)?,
712    };
713    Ok((beta_threshold, beta_log_sigma))
714}
715
716#[derive(Clone)]
717pub(crate) struct BinomialMeanWiggleSpec {
718    pub y: Array1<f64>,
719    pub weights: Array1<f64>,
720    pub link_kind: InverseLink,
721    pub wiggle_knots: Array1<f64>,
722    pub wiggle_degree: usize,
723    pub eta_block: ParameterBlockInput,
724    pub wiggle_block: ParameterBlockInput,
725}
726
727#[derive(Clone)]
728pub struct GaussianLocationScaleTermSpec {
729    pub y: Array1<f64>,
730    pub weights: Array1<f64>,
731    pub meanspec: TermCollectionSpec,
732    pub log_sigmaspec: TermCollectionSpec,
733    pub mean_offset: Array1<f64>,
734    pub log_sigma_offset: Array1<f64>,
735}
736
737#[derive(Clone)]
738pub struct GaussianLocationScaleWiggleTermSpec {
739    pub y: Array1<f64>,
740    pub weights: Array1<f64>,
741    pub meanspec: TermCollectionSpec,
742    pub log_sigmaspec: TermCollectionSpec,
743    pub mean_offset: Array1<f64>,
744    pub log_sigma_offset: Array1<f64>,
745    pub wiggle_knots: Array1<f64>,
746    pub wiggle_degree: usize,
747    pub wiggle_block: ParameterBlockInput,
748}
749
750#[derive(Clone)]
751pub struct BinomialLocationScaleTermSpec {
752    pub y: Array1<f64>,
753    pub weights: Array1<f64>,
754    pub link_kind: InverseLink,
755    pub thresholdspec: TermCollectionSpec,
756    pub log_sigmaspec: TermCollectionSpec,
757    pub threshold_offset: Array1<f64>,
758    pub log_sigma_offset: Array1<f64>,
759}
760
761#[derive(Clone)]
762pub struct BinomialLocationScaleWiggleTermSpec {
763    pub y: Array1<f64>,
764    pub weights: Array1<f64>,
765    pub link_kind: InverseLink,
766    pub thresholdspec: TermCollectionSpec,
767    pub log_sigmaspec: TermCollectionSpec,
768    pub threshold_offset: Array1<f64>,
769    pub log_sigma_offset: Array1<f64>,
770    pub wiggle_knots: Array1<f64>,
771    pub wiggle_degree: usize,
772    pub wiggle_block: ParameterBlockInput,
773}
774
775#[derive(Clone, Debug)]
776pub struct BlockwiseTermFitResult {
777    pub fit: UnifiedFitResult,
778    pub meanspec_resolved: TermCollectionSpec,
779    pub noisespec_resolved: TermCollectionSpec,
780    pub mean_design: TermCollectionDesign,
781    pub noise_design: TermCollectionDesign,
782}
783
784pub(crate) struct BlockwiseTermFitResultParts {
785    pub fit: UnifiedFitResult,
786    pub meanspec_resolved: TermCollectionSpec,
787    pub noisespec_resolved: TermCollectionSpec,
788    pub mean_design: TermCollectionDesign,
789    pub noise_design: TermCollectionDesign,
790}
791
792pub struct BlockwiseTermWiggleFitResult {
793    pub fit: BlockwiseTermFitResult,
794    pub wiggle_knots: Array1<f64>,
795    pub wiggle_degree: usize,
796}
797
798pub struct BinomialMeanWiggleTermFitResult {
799    pub fit: UnifiedFitResult,
800    pub resolvedspec: TermCollectionSpec,
801    pub design: TermCollectionDesign,
802    pub wiggle_knots: Array1<f64>,
803    pub wiggle_degree: usize,
804    /// Standard I-spline warp coefficients `β_w` for the saved-model predict
805    /// runtime when frozen-basis de-aliasing engaged (#1596). Observation-space
806    /// residualization preserves this coefficient chart, so the fit and predict
807    /// runtime consume the same non-negative vector.
808    pub saved_warp_beta: Option<Vec<f64>>,
809    /// Frozen-index mean-coordinate shift `s = β_frozen_source − β_saved` for the
810    /// predict runtime (#2141). Predict evaluates the warp basis at
811    /// `X·(β_saved + s) = η̂` (the frozen index the fit pinned `B` at) instead of
812    /// the de-aliased base predictor `X·β_saved`, reproducing the fitted `q`.
813    pub saved_index_shift: Option<Vec<f64>>,
814}
815
816pub(crate) struct BlockwiseTermWiggleFitResultParts {
817    pub fit: BlockwiseTermFitResult,
818    pub wiggle_knots: Array1<f64>,
819    pub wiggle_degree: usize,
820}
821
822pub(crate) fn validate_term_collection_design(
823    label: &str,
824    design: &TermCollectionDesign,
825) -> Result<(), String> {
826    let p = design.design.ncols();
827    let n = design.design.nrows();
828    for rows in exact_design_row_chunks(n, p) {
829        let chunk = design
830            .design
831            .try_row_chunk(rows)
832            .map_err(|e| format!("{label}.design row chunk materialization failed: {e}"))?;
833        validate_all_finite_estimation(&format!("{label}.design"), chunk.iter().copied())
834            .map_err(|e| e.to_string())?;
835    }
836    if design.nullspace_dims.len() != design.penalties.len() {
837        return Err(GamlssError::DimensionMismatch {
838            reason: format!(
839                "{label}.nullspace_dims length mismatch: got {}, expected {}",
840                design.nullspace_dims.len(),
841                design.penalties.len()
842            ),
843        }
844        .into());
845    }
846    if design.penaltyinfo.len() != design.penalties.len() {
847        return Err(GamlssError::DimensionMismatch {
848            reason: format!(
849                "{label}.penaltyinfo length mismatch: got {}, expected {}",
850                design.penaltyinfo.len(),
851                design.penalties.len()
852            ),
853        }
854        .into());
855    }
856    for (idx, bp) in design.penalties.iter().enumerate() {
857        validate_all_finite_estimation(
858            &format!("{label}.penalties[{idx}]"),
859            bp.local.iter().copied(),
860        )
861        .map_err(|e| e.to_string())?;
862        if bp.col_range.end > p {
863            return Err(GamlssError::DimensionMismatch {
864                reason: format!(
865                    "{label}.penalties[{idx}] col_range {}..{} exceeds design width {}",
866                    bp.col_range.start, bp.col_range.end, p
867                ),
868            }
869            .into());
870        }
871    }
872    if let Some(bounds) = design.coefficient_lower_bounds.as_ref() {
873        if bounds.len() != p {
874            return Err(GamlssError::ConstraintViolation {
875                reason: format!(
876                    "{label}.coefficient_lower_bounds length mismatch: got {}, expected {p}",
877                    bounds.len()
878                ),
879            }
880            .into());
881        }
882        for (idx, &bound) in bounds.iter().enumerate() {
883            if !(bound.is_finite() || bound == f64::NEG_INFINITY) {
884                return Err(GamlssError::NonFinite { reason: format!(
885                    "{label}.coefficient_lower_bounds[{idx}] must be finite or -inf, got {bound}",
886                ) }.into());
887            }
888        }
889    }
890    if let Some(constraints) = design.linear_constraints.as_ref() {
891        validate_all_finite_estimation(
892            &format!("{label}.linear_constraints.a"),
893            constraints.a.iter().copied(),
894        )
895        .map_err(|e| e.to_string())?;
896        validate_all_finite_estimation(
897            &format!("{label}.linear_constraints.b"),
898            constraints.b.iter().copied(),
899        )
900        .map_err(|e| e.to_string())?;
901        if constraints.a.ncols() != p {
902            return Err(GamlssError::DimensionMismatch {
903                reason: format!(
904                    "{label}.linear_constraints.a column mismatch: got {}, expected {p}",
905                    constraints.a.ncols()
906                ),
907            }
908            .into());
909        }
910        if constraints.a.nrows() != constraints.b.len() {
911            return Err(GamlssError::DimensionMismatch {
912                reason: format!(
913                    "{label}.linear_constraints row mismatch: a has {}, b has {}",
914                    constraints.a.nrows(),
915                    constraints.b.len()
916                ),
917            }
918            .into());
919        }
920    }
921    if design.intercept_range.start > design.intercept_range.end || design.intercept_range.end > p {
922        return Err(GamlssError::ConstraintViolation {
923            reason: format!(
924                "{label}.intercept_range out of bounds: {:?} for {} columns",
925                design.intercept_range, p
926            ),
927        }
928        .into());
929    }
930    Ok(())
931}
932
933impl BlockwiseTermFitResult {
934    pub(crate) fn try_from_parts(parts: BlockwiseTermFitResultParts) -> Result<Self, String> {
935        let BlockwiseTermFitResultParts {
936            fit,
937            meanspec_resolved,
938            noisespec_resolved,
939            mean_design,
940            noise_design,
941        } = parts;
942
943        fit.validate_numeric_finiteness()
944            .map_err(|e| format!("{e}"))?;
945        if fit.block_states.len() < 2 {
946            return Err(GamlssError::DimensionMismatch {
947                reason: format!(
948                    "BlockwiseTermFitResult requires at least 2 block states, got {}",
949                    fit.block_states.len()
950                ),
951            }
952            .into());
953        }
954        validate_term_collection_design("blockwise_term.mean_design", &mean_design)?;
955        validate_term_collection_design("blockwise_term.noise_design", &noise_design)?;
956        if mean_design.design.nrows() != noise_design.design.nrows() {
957            return Err(GamlssError::DimensionMismatch {
958                reason: format!(
959                    "BlockwiseTermFitResult row mismatch: mean_design={}, noise_design={}",
960                    mean_design.design.nrows(),
961                    noise_design.design.nrows()
962                ),
963            }
964            .into());
965        }
966        if fit.block_states[0].beta.len() != mean_design.design.ncols() {
967            return Err(GamlssError::DimensionMismatch {
968                reason: format!(
969                    "BlockwiseTermFitResult mean beta length mismatch: got {}, expected {}",
970                    fit.block_states[0].beta.len(),
971                    mean_design.design.ncols()
972                ),
973            }
974            .into());
975        }
976        if fit.block_states[1].beta.len() != noise_design.design.ncols() {
977            return Err(GamlssError::DimensionMismatch {
978                reason: format!(
979                    "BlockwiseTermFitResult noise beta length mismatch: got {}, expected {}",
980                    fit.block_states[1].beta.len(),
981                    noise_design.design.ncols()
982                ),
983            }
984            .into());
985        }
986        if fit.block_states[0].eta.len() != mean_design.design.nrows() {
987            return Err(GamlssError::DimensionMismatch {
988                reason: format!(
989                    "BlockwiseTermFitResult mean eta length mismatch: got {}, expected {}",
990                    fit.block_states[0].eta.len(),
991                    mean_design.design.nrows()
992                ),
993            }
994            .into());
995        }
996        if fit.block_states[1].eta.len() != noise_design.design.nrows() {
997            return Err(GamlssError::DimensionMismatch {
998                reason: format!(
999                    "BlockwiseTermFitResult noise eta length mismatch: got {}, expected {}",
1000                    fit.block_states[1].eta.len(),
1001                    noise_design.design.nrows()
1002                ),
1003            }
1004            .into());
1005        }
1006
1007        Ok(Self {
1008            fit,
1009            meanspec_resolved,
1010            noisespec_resolved,
1011            mean_design,
1012            noise_design,
1013        })
1014    }
1015
1016    pub(crate) fn validate_numeric_finiteness(&self) -> Result<(), String> {
1017        Self::try_from_parts(BlockwiseTermFitResultParts {
1018            fit: self.fit.clone(),
1019            meanspec_resolved: self.meanspec_resolved.clone(),
1020            noisespec_resolved: self.noisespec_resolved.clone(),
1021            mean_design: self.mean_design.clone(),
1022            noise_design: self.noise_design.clone(),
1023        })
1024        .map(|_| ())
1025    }
1026}
1027
1028impl BlockwiseTermWiggleFitResult {
1029    pub(crate) fn try_from_parts(parts: BlockwiseTermWiggleFitResultParts) -> Result<Self, String> {
1030        let BlockwiseTermWiggleFitResultParts {
1031            fit,
1032            wiggle_knots,
1033            wiggle_degree,
1034        } = parts;
1035
1036        fit.validate_numeric_finiteness()
1037            .map_err(|e| e.to_string())?;
1038        if fit.fit.block_states.len() < 3 {
1039            return Err(GamlssError::DimensionMismatch {
1040                reason: format!(
1041                    "BlockwiseTermWiggleFitResult requires at least 3 block states, got {}",
1042                    fit.fit.block_states.len()
1043                ),
1044            }
1045            .into());
1046        }
1047        if wiggle_knots.is_empty() {
1048            return Err(GamlssError::UnsupportedConfiguration {
1049                reason: "BlockwiseTermWiggleFitResult requires non-empty wiggle_knots".to_string(),
1050            }
1051            .into());
1052        }
1053        validate_all_finite_estimation(
1054            "blockwise_term_wiggle.wiggle_knots",
1055            wiggle_knots.iter().copied(),
1056        )
1057        .map_err(|e| e.to_string())?;
1058
1059        Ok(Self {
1060            fit,
1061            wiggle_knots,
1062            wiggle_degree,
1063        })
1064    }
1065}
1066
1067pub struct BinomialLocationScaleFitResult {
1068    pub fit: BlockwiseTermFitResult,
1069    pub wiggle_knots: Option<Array1<f64>>,
1070    pub wiggle_degree: Option<usize>,
1071    pub beta_link_wiggle: Option<Vec<f64>>,
1072}
1073
1074pub struct GaussianLocationScaleFitResult {
1075    pub fit: BlockwiseTermFitResult,
1076    pub wiggle_knots: Option<Array1<f64>>,
1077    pub wiggle_degree: Option<usize>,
1078    pub beta_link_wiggle: Option<Vec<f64>>,
1079    /// Response standardization factor applied internally during fitting.
1080    ///
1081    /// The Gaussian location-scale path fits on `y / response_scale` so the
1082    /// fixed log-σ soft floor `LOGB_SIGMA_FLOOR = 0.01` is *operationally*
1083    /// scale-relative (1 % of the response spread) rather than absolute,
1084    /// keeping κ = dlogσ/dη ≈ 1 across the realistic σ range and informing the
1085    /// scale block like gamlss. The returned coefficient `blocks`, `beta`, and
1086    /// link-wiggle knots/coefficients are already mapped back to **raw response
1087    /// units** (the Location/Mean block scaled by `response_scale`, the Scale
1088    /// block intercept shifted by `+ln(response_scale)`), so downstream
1089    /// reconstruction `μ = X_mean·β` comes out in raw units with no further
1090    /// rescaling.
1091    ///
1092    /// The σ reconstruction, however, **must scale the floor too** to stay
1093    /// response-scale-equivariant (#884):
1094    ///
1095    /// ```text
1096    /// σ = response_scale·LOGB_SIGMA_FLOOR + exp(X_scale·β)
1097    ///   = response_scale·(LOGB_SIGMA_FLOOR + exp(η_internal)).
1098    /// ```
1099    ///
1100    /// The intercept shift carries only the `exp(η)` term; reconstructing with a
1101    /// raw `LOGB_SIGMA_FLOOR` instead of `response_scale·LOGB_SIGMA_FLOOR` leaves
1102    /// the non-equivariant residual `LOGB_SIGMA_FLOOR·(1 − response_scale)`.
1103    ///
1104    /// This field records the factor that was applied for transparency,
1105    /// covariance bookkeeping, and the equivariant σ-floor reconstruction; it is
1106    /// `1.0` when no standardization was needed (degenerate constant response).
1107    pub response_scale: f64,
1108}
1109
1110/// Exact coefficient-frame map for the frozen-basis binomial mean-wiggle
1111/// de-aliasing step.
1112///
1113/// The joint solver sees `[X, B - XA]` and returns coordinates
1114/// `(beta_mean_solver, beta_w)`. Saved prediction deliberately uses `[X, B]`,
1115/// so the reported coordinates are
1116///
1117/// ```text
1118/// beta_mean_saved = beta_mean_solver - A beta_w
1119/// beta_w_saved    = beta_w.
1120/// ```
1121///
1122/// This is one linear section with the cross-block lift
1123/// `M = [[I, -A], [0, I]]`. Keeping it as a [`gam_problem::Gauge`] makes the
1124/// same map authoritative for coefficients, covariance, and the active
1125/// geometry lineage used by saved-model ALO.
1126fn binomial_mean_wiggle_saved_frame_gauge(
1127    alias: &Array2<f64>,
1128    mean_width: usize,
1129    wiggle_width: usize,
1130) -> Result<gam_problem::Gauge, String> {
1131    if alias.dim() != (mean_width, wiggle_width) {
1132        return Err(format!(
1133            "binomial mean-wiggle de-alias map is {}x{}, expected {mean_width}x{wiggle_width}",
1134            alias.nrows(),
1135            alias.ncols(),
1136        ));
1137    }
1138    let total_width = mean_width
1139        .checked_add(wiggle_width)
1140        .ok_or_else(|| "binomial mean-wiggle coefficient dimension overflows usize".to_string())?;
1141    let mut transform = Array2::<f64>::eye(total_width);
1142    for row in 0..mean_width {
1143        for column in 0..wiggle_width {
1144            transform[[row, mean_width + column]] = -alias[[row, column]];
1145        }
1146    }
1147    let gauge = gam_problem::Gauge::from_t(
1148        transform,
1149        &[mean_width, wiggle_width],
1150        &[mean_width, wiggle_width],
1151    );
1152    gauge.validate().map_err(|reason| {
1153        format!("binomial mean-wiggle saved coefficient gauge is invalid: {reason}")
1154    })?;
1155    Ok(gauge)
1156}
1157
1158fn binomial_mean_wiggle_saved_geometry(
1159    geometry: &gam_solve::model_types::FitGeometry,
1160    saved_frame: &gam_problem::Gauge,
1161) -> Result<gam_solve::model_types::FitGeometry, String> {
1162    let mut saved_geometry = geometry.clone();
1163    saved_geometry.coefficient_gauge = geometry
1164        .coefficient_gauge
1165        .left_compose(saved_frame)
1166        .map_err(|reason| {
1167            format!(
1168                "binomial mean-wiggle active geometry cannot compose with its exact saved-result gauge: {reason}"
1169            )
1170        })?;
1171    Ok(saved_geometry)
1172}
1173
1174fn binomial_mean_wiggle_saved_covariance(
1175    covariance: &Array2<f64>,
1176    saved_frame: &gam_problem::Gauge,
1177    label: &str,
1178) -> Result<Array2<f64>, String> {
1179    let expected = saved_frame.reduced_total();
1180    if covariance.dim() != (expected, expected) {
1181        return Err(format!(
1182            "binomial mean-wiggle {label} is {}x{}; exact saved-result gauge requires {expected}x{expected} solver-frame coordinates",
1183            covariance.nrows(),
1184            covariance.ncols(),
1185        ));
1186    }
1187    if let Some(((row, column), value)) = covariance
1188        .indexed_iter()
1189        .find(|(_, value)| !value.is_finite())
1190    {
1191        return Err(format!(
1192            "binomial mean-wiggle {label} is non-finite at ({row}, {column}): {value}"
1193        ));
1194    }
1195    let saved = saved_frame.lift_covariance(covariance);
1196    if let Some(((row, column), value)) = saved.indexed_iter().find(|(_, value)| !value.is_finite())
1197    {
1198        return Err(format!(
1199            "binomial mean-wiggle saved-frame {label} is non-finite at ({row}, {column}): {value}"
1200        ));
1201    }
1202    Ok(saved)
1203}
1204
1205/// Atomically move a converged frozen-basis mean-wiggle fit from the solver's
1206/// residualized-design coordinates into the saved prediction coordinates.
1207///
1208/// The penalized Hessian remains in its exact active solver coordinates;
1209/// composing its coefficient gauge records how raw saved rows pull back into
1210/// that frame. Covariances, in contrast, push forward through the saved-frame
1211/// map. No dimension mismatch is ignorable: returning a partially transformed
1212/// fit would make point estimates and uncertainty describe different models.
1213fn finalize_binomial_mean_wiggle_saved_frame(
1214    fit: &mut UnifiedFitResult,
1215    alias: &Array2<f64>,
1216    mean_design: &Array2<f64>,
1217    mean_offset: &Array1<f64>,
1218) -> Result<(), String> {
1219    use gam_problem::BlockRole;
1220
1221    if fit.blocks.len() != 2
1222        || fit.blocks[0].role != BlockRole::Mean
1223        || fit.blocks[1].role != BlockRole::LinkWiggle
1224    {
1225        return Err(format!(
1226            "binomial mean-wiggle saved-frame finalization requires fitted blocks [Mean, LinkWiggle], got {:?}",
1227            fit.blocks
1228                .iter()
1229                .map(|block| block.role)
1230                .collect::<Vec<_>>()
1231        ));
1232    }
1233    if fit.block_states.len() != 2 {
1234        return Err(format!(
1235            "binomial mean-wiggle saved-frame finalization requires two fitted block states, got {}",
1236            fit.block_states.len(),
1237        ));
1238    }
1239    if mean_offset.len() != mean_design.nrows() {
1240        return Err(format!(
1241            "binomial mean-wiggle mean offset has {} rows, expected {}",
1242            mean_offset.len(),
1243            mean_design.nrows(),
1244        ));
1245    }
1246
1247    let mean_width = fit.blocks[0].beta.len();
1248    let wiggle_width = fit.blocks[1].beta.len();
1249    if mean_design.ncols() != mean_width {
1250        return Err(format!(
1251            "binomial mean-wiggle saved mean design has {} columns, expected fitted width {mean_width}",
1252            mean_design.ncols(),
1253        ));
1254    }
1255    for block_index in 0..2 {
1256        if fit.block_states[block_index].beta != fit.blocks[block_index].beta {
1257            return Err(format!(
1258                "binomial mean-wiggle fitted block {block_index} and block-state coefficients disagree before saved-frame finalization"
1259            ));
1260        }
1261    }
1262    let total_width = mean_width
1263        .checked_add(wiggle_width)
1264        .ok_or_else(|| "binomial mean-wiggle coefficient dimension overflows usize".to_string())?;
1265    if fit.beta.len() != total_width {
1266        return Err(format!(
1267            "binomial mean-wiggle flat coefficient vector has width {}, expected {total_width}",
1268            fit.beta.len(),
1269        ));
1270    }
1271    if fit.beta.slice(s![0..mean_width]) != fit.blocks[0].beta
1272        || fit.beta.slice(s![mean_width..total_width]) != fit.blocks[1].beta
1273    {
1274        return Err(
1275            "binomial mean-wiggle flat and block coefficient vectors disagree before saved-frame finalization"
1276                .to_string(),
1277        );
1278    }
1279
1280    let saved_frame = binomial_mean_wiggle_saved_frame_gauge(alias, mean_width, wiggle_width)?;
1281    let saved_blocks =
1282        saved_frame.lift_block_betas(&[fit.blocks[0].beta.clone(), fit.blocks[1].beta.clone()]);
1283    let saved_mean_eta = mean_design.dot(&saved_blocks[0]) + mean_offset;
1284    let mut saved_beta = Array1::<f64>::zeros(total_width);
1285    saved_beta
1286        .slice_mut(s![0..mean_width])
1287        .assign(&saved_blocks[0]);
1288    saved_beta
1289        .slice_mut(s![mean_width..total_width])
1290        .assign(&saved_blocks[1]);
1291
1292    let saved_conditional = fit
1293        .covariance_conditional
1294        .as_ref()
1295        .map(|covariance| {
1296            binomial_mean_wiggle_saved_covariance(
1297                covariance,
1298                &saved_frame,
1299                "conditional covariance",
1300            )
1301        })
1302        .transpose()?;
1303    let saved_corrected = fit
1304        .covariance_corrected
1305        .as_ref()
1306        .map(|covariance| {
1307            binomial_mean_wiggle_saved_covariance(covariance, &saved_frame, "corrected covariance")
1308        })
1309        .transpose()?;
1310    let saved_geometry = binomial_mean_wiggle_saved_geometry(
1311        fit.geometry.as_ref().ok_or_else(|| {
1312            "binomial mean-wiggle fit is missing its exact active geometry".to_string()
1313        })?,
1314        &saved_frame,
1315    )?;
1316
1317    let mut saved_inference = fit.inference.clone();
1318    if let Some(inference) = saved_inference.as_mut() {
1319        if inference.beta_covariance.is_none() && inference.beta_standard_errors.is_some() {
1320            return Err(
1321                "binomial mean-wiggle inference has conditional standard errors without their covariance"
1322                    .to_string(),
1323            );
1324        }
1325        if inference.beta_covariance_corrected.is_none()
1326            && inference.beta_standard_errors_corrected.is_some()
1327        {
1328            return Err(
1329                "binomial mean-wiggle inference has corrected standard errors without their covariance"
1330                    .to_string(),
1331            );
1332        }
1333        if let Some(covariance) = inference.beta_covariance.take() {
1334            let covariance = binomial_mean_wiggle_saved_covariance(
1335                covariance.as_array(),
1336                &saved_frame,
1337                "inference conditional covariance",
1338            )?;
1339            if inference.beta_standard_errors.is_some() {
1340                inference.beta_standard_errors = Some(
1341                    gam_problem::se_from_covariance(&covariance).map_err(|reason| {
1342                        format!(
1343                            "binomial mean-wiggle saved conditional standard errors are invalid: {reason}"
1344                        )
1345                    })?,
1346                );
1347            }
1348            inference.beta_covariance = Some(covariance.into());
1349        }
1350        if let Some(covariance) = inference.beta_covariance_corrected.take() {
1351            let covariance = binomial_mean_wiggle_saved_covariance(
1352                &covariance,
1353                &saved_frame,
1354                "inference corrected covariance",
1355            )?;
1356            if inference.beta_standard_errors_corrected.is_some() {
1357                inference.beta_standard_errors_corrected = Some(
1358                    gam_problem::se_from_covariance(&covariance).map_err(|reason| {
1359                        format!(
1360                            "binomial mean-wiggle saved corrected standard errors are invalid: {reason}"
1361                        )
1362                    })?,
1363                );
1364            }
1365            inference.beta_covariance_corrected = Some(covariance);
1366        }
1367        if let Some(covariance) = inference.beta_covariance_frequentist.take() {
1368            inference.beta_covariance_frequentist = Some(binomial_mean_wiggle_saved_covariance(
1369                &covariance,
1370                &saved_frame,
1371                "frequentist covariance",
1372            )?);
1373        }
1374        if let Some(correction) = inference.smoothing_correction.take() {
1375            inference.smoothing_correction = Some(binomial_mean_wiggle_saved_covariance(
1376                &correction,
1377                &saved_frame,
1378                "smoothing covariance correction",
1379            )?);
1380        }
1381    }
1382
1383    fit.blocks[0].beta = saved_blocks[0].clone();
1384    fit.blocks[1].beta = saved_blocks[1].clone();
1385    fit.block_states[0].beta = saved_blocks[0].clone();
1386    fit.block_states[0].eta = saved_mean_eta;
1387    fit.block_states[1].beta = saved_blocks[1].clone();
1388    fit.beta = saved_beta;
1389    fit.covariance_conditional = saved_conditional;
1390    fit.covariance_corrected = saved_corrected;
1391    fit.geometry = Some(saved_geometry);
1392    fit.inference = saved_inference;
1393    Ok(())
1394}
1395
1396#[cfg(test)]
1397mod binomial_mean_wiggle_saved_frame_tests {
1398    use super::*;
1399    use ndarray::array;
1400
1401    #[test]
1402    fn cross_block_dealias_composes_non_square_geometry_and_pushes_covariance() {
1403        let alias = array![[2.0], [-0.5]];
1404        let saved_frame = binomial_mean_wiggle_saved_frame_gauge(&alias, 2, 1)
1405            .expect("valid cross-block de-alias map");
1406
1407        // The canonical solver retained one of two Mean directions plus the
1408        // LinkWiggle direction: active(2) -> solver raw(3) is rectangular.
1409        let active_to_solver = gam_problem::Gauge::from_t(
1410            array![[1.0, 0.0], [0.0, 0.0], [0.0, 1.0]],
1411            &[2, 1],
1412            &[1, 1],
1413        );
1414        let active_hessian = array![[7.0, 1.5], [1.5, 4.0]];
1415        let geometry = gam_solve::model_types::FitGeometry {
1416            coefficient_gauge: active_to_solver,
1417            penalized_hessian: active_hessian.clone().into(),
1418            working: None,
1419        };
1420        let saved_geometry = binomial_mean_wiggle_saved_geometry(&geometry, &saved_frame)
1421            .expect("non-square active geometry composes through saved frame");
1422
1423        assert_eq!(
1424            saved_geometry.coefficient_gauge.t_full,
1425            array![[1.0, -2.0], [0.0, 0.5], [0.0, 1.0]],
1426        );
1427        assert_eq!(
1428            saved_geometry.penalized_hessian.as_array(),
1429            &active_hessian,
1430            "precision stays in the canonical active frame",
1431        );
1432
1433        let solver_covariance = Array2::<f64>::eye(3);
1434        let saved_covariance = binomial_mean_wiggle_saved_covariance(
1435            &solver_covariance,
1436            &saved_frame,
1437            "test covariance",
1438        )
1439        .expect("covariance pushes into saved frame");
1440        assert_eq!(
1441            saved_covariance,
1442            array![[5.0, -1.0, -2.0], [-1.0, 1.25, 0.5], [-2.0, 0.5, 1.0]],
1443            "the -A cross block must alter both Mean variance and Mean/Wiggle covariance",
1444        );
1445    }
1446}
1447
1448/// Fit the binomial mean link-wiggle model. The observation-space de-aliasing
1449/// preserves the standard I-spline coefficient coordinate, which is returned
1450/// for the saved-model predict runtime.
1451pub(crate) fn fit_binomial_mean_wiggle(
1452    spec: BinomialMeanWiggleSpec,
1453    options: &BlockwiseFitOptions,
1454) -> Result<(UnifiedFitResult, Option<Vec<f64>>, Option<Vec<f64>>), String> {
1455    let n = spec.y.len();
1456    validate_len_match("weights vs y", n, spec.weights.len())?;
1457    validateweights(&spec.weights, "fit_binomial_mean_wiggle")?;
1458    validate_binomial_response(&spec.y, "fit_binomial_mean_wiggle")?;
1459    validate_blockrows("eta", n, &spec.eta_block)?;
1460    validate_blockrows("wiggle", n, &spec.wiggle_block)?;
1461    if matches!(
1462        spec.link_kind,
1463        InverseLink::Standard(StandardLink::Identity)
1464    ) {
1465        return Err(GamlssError::UnsupportedConfiguration {
1466            reason: "fit_binomial_mean_wiggle does not support identity link".to_string(),
1467        }
1468        .into());
1469    }
1470    gam_terms::inference::formula_dsl::require_binomial_inverse_link_supports_joint_wiggle(
1471        &spec.link_kind,
1472        "fit_binomial_mean_wiggle",
1473    )?;
1474    if spec.wiggle_degree < 2 {
1475        return Err(GamlssError::ConstraintViolation {
1476            reason: format!(
1477                "fit_binomial_mean_wiggle: wiggle_degree must be >= 2, got {}",
1478                spec.wiggle_degree
1479            ),
1480        }
1481        .into());
1482    }
1483    let minimum_knots = minimum_monotone_wiggle_knot_count(spec.wiggle_degree)?;
1484    if spec.wiggle_knots.len() < minimum_knots {
1485        return Err(GamlssError::DimensionMismatch { reason: format!(
1486            "fit_binomial_mean_wiggle: wiggle_knots length {} is too short for degree {} (need at least {})",
1487            spec.wiggle_knots.len(),
1488            spec.wiggle_degree,
1489            minimum_knots
1490        ) }.into());
1491    }
1492
1493    // ----- Frozen-basis Gauss-Newton link-warp fit (#1596) -----
1494    //
1495    // The warp basis `B(η)` is frozen at the current index `η̂` so that
1496    // `q = η + B(η̂)·β_w` is linear in `(β_η, β_w)` (`∂q/∂η = 1`). To keep the
1497    // mean block `X` full and identifiable we fit the warp through the
1498    // observation-space residualized design `B⊥ = (I - P_X)B(η̂)`. We re-freeze
1499    // at the refit `η̂` until the caller's outer convergence policy certifies the
1500    // fixed point.
1501    let x_dense: Array2<f64> = spec.eta_block.design.to_dense();
1502    let (pilot_beta, pilot_eta): (Array1<f64>, Array1<f64>) = {
1503        let pilot_beta = spec.eta_block.initial_beta.clone().ok_or_else(|| {
1504            "fit_binomial_mean_wiggle: eta block carries no pilot β to seed the \
1505             frozen-basis warp index"
1506                .to_string()
1507        })?;
1508        if x_dense.ncols() != pilot_beta.len() {
1509            return Err(GamlssError::DimensionMismatch {
1510                reason: format!(
1511                    "fit_binomial_mean_wiggle: eta design has {} columns but pilot β has {} \
1512                     coefficients",
1513                    x_dense.ncols(),
1514                    pilot_beta.len()
1515                ),
1516            }
1517            .into());
1518        }
1519        let mut eta = x_dense.dot(&pilot_beta);
1520        eta += &spec.eta_block.offset;
1521        (pilot_beta, eta)
1522    };
1523
1524    // Original (full-width) warp penalties / nullspace metadata, captured before
1525    // `spec.wiggle_block` is consumed. The residualized block keeps the same
1526    // coefficient coordinate and therefore the same penalties.
1527    let wiggle_penalties_full = spec.wiggle_block.penalties.clone();
1528    let wiggle_nullspace_dims = spec.wiggle_block.nullspace_dims.clone();
1529    if !wiggle_nullspace_dims.is_empty()
1530        && wiggle_nullspace_dims.len() != wiggle_penalties_full.len()
1531    {
1532        return Err(GamlssError::DimensionMismatch {
1533            reason: format!(
1534                "fit_binomial_mean_wiggle: wiggle block has {} penalties but {} nullspace dimensions",
1535                wiggle_penalties_full.len(),
1536                wiggle_nullspace_dims.len()
1537            ),
1538        }
1539        .into());
1540    }
1541    let wiggle_log_lambdas = spec.wiggle_block.initial_log_lambdas.clone();
1542    let wiggle_beta_initial = spec.wiggle_block.initial_beta.clone();
1543    let eta_block_input = spec.eta_block.clone();
1544
1545    let family = BinomialMeanWiggleFamily {
1546        y: spec.y,
1547        weights: spec.weights,
1548        link_kind: spec.link_kind,
1549        wiggle_knots: spec.wiggle_knots,
1550        wiggle_degree: spec.wiggle_degree,
1551        policy: gam_runtime::resource::ResourcePolicy::default_library(),
1552        frozen_warp_design: None,
1553    };
1554
1555    // Build the de-aliased warp block at a frozen index.  The identifiable
1556    // warp is the part of `B(η̂)` outside the mean column space:
1557    //
1558    //     B⊥ = (I - P_X) B = B - X A,     A = (XᵀX)^+ XᵀB.
1559    //
1560    // The previous implementation used `Z = null(XᵀB)` and fitted `B Z`.
1561    // That is too strong: when the warp basis has no coefficient combination
1562    // exactly orthogonal to `X` (for example a two-column flexible link beside
1563    // an intercept+slope mean), it drops every warp coefficient even though
1564    // the nonlinear columns of `B(η̂)` have a nonzero residual after projection
1565    // onto `X`.  Residualizing in observation space removes only the truly
1566    // mean-aliased component and leaves the curved, identifiable link-shape
1567    // signal available to the joint solve.
1568    //
1569    // The returned `A` is used after fitting: because the inner problem used
1570    // `Xβ + (B - XA)β_w`, while prediction reconstructs the saved warp as
1571    // `Xβ_saved + Bβ_w`, we save `β_saved = β - Aβ_w`.
1572    let build_dealiased = |frozen: &Array1<f64>,
1573                           beta_hint: Option<&Array1<f64>>,
1574                           log_lambda_hint: Option<&Array1<f64>>|
1575     -> Result<
1576        (
1577            ParameterBlockInput,
1578            Array2<f64>,
1579            std::sync::Arc<Array2<f64>>,
1580        ),
1581        String,
1582    > {
1583        use faer::Side;
1584        use gam_linalg::faer_ndarray::FaerEigh;
1585
1586        let b_full = family.wiggle_design(frozen.view())?;
1587        let xtx = x_dense.t().dot(&x_dense);
1588        let xtb = x_dense.t().dot(&b_full);
1589        let (evals, evecs) = xtx
1590            .eigh(Side::Lower)
1591            .map_err(|e| format!("frozen-basis warp de-aliasing mean QR failed: {e}"))?;
1592        let max_eval = evals.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1593        let cutoff = 1.0e3 * f64::EPSILON * (xtx.nrows().max(1) as f64) * max_eval.max(1.0);
1594        let mut alias = Array2::<f64>::zeros((x_dense.ncols(), b_full.ncols()));
1595        for k in 0..evals.len() {
1596            let lam = evals[k];
1597            if !lam.is_finite() || lam.abs() <= cutoff {
1598                continue;
1599            }
1600            let uk = evecs.column(k);
1601            let uk_xtb = uk.t().dot(&xtb);
1602            for i in 0..alias.nrows() {
1603                for j in 0..alias.ncols() {
1604                    alias[[i, j]] += uk[i] * uk_xtb[j] / lam;
1605                }
1606            }
1607        }
1608        let bda = &b_full - &x_dense.dot(&alias);
1609        let max_b = b_full.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1610        let max_resid = bda.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1611        let resid_tol =
1612            1.0e3 * f64::EPSILON * (bda.nrows().max(bda.ncols()).max(1) as f64) * max_b.max(1.0);
1613        if max_resid <= resid_tol {
1614            return Err("frozen-basis warp de-aliasing left no identifiable warp \
1615                        direction (the mean block already spans the warp in \
1616                        observation space)"
1617                .to_string());
1618        }
1619        let penalties: Vec<crate::model_types::PenaltySpec> = wiggle_penalties_full
1620            .iter()
1621            .map(|p| {
1622                let s = penalty_spec_to_dense(p, b_full.ncols())?;
1623                Ok(crate::model_types::PenaltySpec::Dense(s))
1624            })
1625            .collect::<Result<_, String>>()?;
1626        let q = bda.ncols();
1627        let initial_beta = match beta_hint {
1628            Some(beta) if beta.len() == q => Some(beta.clone()),
1629            Some(beta) => {
1630                return Err(format!(
1631                    "frozen-basis warp warm start has {} coefficients but the realized basis has {q}",
1632                    beta.len()
1633                ));
1634            }
1635            None => Some(Array1::zeros(q)),
1636        };
1637        let block = ParameterBlockInput {
1638            design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(bda.clone())),
1639            offset: Array1::zeros(frozen.len()),
1640            penalties,
1641            nullspace_dims: wiggle_nullspace_dims.clone(),
1642            initial_log_lambdas: log_lambda_hint
1643                .cloned()
1644                .or_else(|| wiggle_log_lambdas.clone()),
1645            initial_beta,
1646        };
1647        Ok((block, alias, std::sync::Arc::new(bda)))
1648    };
1649
1650    // Outer Gauss-Newton / backfitting loop over the frozen warp index. The
1651    // smoothing parameters remain continuously REML/LAML-optimized at every
1652    // step. A fit is returned only after the frozen index is a certified fixed
1653    // point; exhaustion returns non-convergence evidence instead of minting a
1654    // fit from the last iterate. The exact coefficient source used to build the
1655    // accepted frozen basis is retained for prediction (#2141).
1656    if options.outer_max_iter == 0 || !options.outer_tol.is_finite() || options.outer_tol <= 0.0 {
1657        return Err(GamlssError::InvalidInput {
1658            reason: format!(
1659                "fit_binomial_mean_wiggle requires positive outer convergence policy; outer_max_iter={}, outer_tol={}",
1660                options.outer_max_iter, options.outer_tol
1661            ),
1662        }
1663        .into());
1664    }
1665    let mut frozen_source_beta = pilot_beta;
1666    let mut frozen_eta = pilot_eta;
1667    let mut eta_block_warm = eta_block_input.clone();
1668    let mut wiggle_beta_warm = wiggle_beta_initial;
1669    let mut wiggle_log_lambda_warm = wiggle_log_lambdas.clone();
1670    let mut converged: Option<(UnifiedFitResult, Array2<f64>, Array1<f64>)> = None;
1671    let mut last_delta = f64::INFINITY;
1672    let mut last_scale = 1.0_f64;
1673    for _outer in 0..options.outer_max_iter {
1674        let (wiggle_block, alias, bda) = build_dealiased(
1675            &frozen_eta,
1676            wiggle_beta_warm.as_ref(),
1677            wiggle_log_lambda_warm.as_ref(),
1678        )?;
1679        let eta_penalty_count = eta_block_warm.penalties.len();
1680        let wiggle_penalty_count = wiggle_block.penalties.len();
1681        let blocks = vec![
1682            eta_block_warm.clone().intospec("eta")?,
1683            wiggle_block.intospec("wiggle")?,
1684        ];
1685        let mut fam = family.clone();
1686        fam.frozen_warp_design = Some(bda);
1687        let fit = fit_custom_family(&fam, &blocks, options).map_err(|e| e.to_string())?;
1688        let mean_state = fit
1689            .block_states
1690            .get(BinomialMeanWiggleFamily::BLOCK_ETA)
1691            .ok_or_else(|| {
1692                "fit_binomial_mean_wiggle: frozen-basis refit did not expose a fitted eta block"
1693                    .to_string()
1694            })?;
1695        if mean_state.eta.len() != frozen_eta.len()
1696            || mean_state.beta.len() != frozen_source_beta.len()
1697        {
1698            return Err(GamlssError::DimensionMismatch {
1699                reason: "fit_binomial_mean_wiggle: frozen-basis refit returned an incompatible eta block"
1700                    .to_string(),
1701            }
1702            .into());
1703        }
1704        let new_eta = mean_state.eta.clone();
1705        let new_source_beta = mean_state.beta.clone();
1706        let new_wiggle_beta = fit
1707            .block_states
1708            .get(BinomialMeanWiggleFamily::BLOCK_WIGGLE)
1709            .map(|state| state.beta.clone())
1710            .ok_or_else(|| {
1711                "fit_binomial_mean_wiggle: frozen-basis refit did not expose a fitted wiggle block"
1712                    .to_string()
1713            })?;
1714        last_scale = frozen_eta
1715            .iter()
1716            .chain(new_eta.iter())
1717            .map(|value| value.abs())
1718            .fold(1.0_f64, f64::max);
1719        last_delta = new_eta
1720            .iter()
1721            .zip(frozen_eta.iter())
1722            .map(|(a, b)| (a - b).abs())
1723            .fold(0.0_f64, f64::max);
1724        if last_delta <= options.outer_tol * last_scale {
1725            converged = Some((fit, alias, frozen_source_beta));
1726            break;
1727        }
1728
1729        let expected_log_lambdas = eta_penalty_count + wiggle_penalty_count;
1730        if fit.log_lambdas.len() != expected_log_lambdas {
1731            return Err(GamlssError::DimensionMismatch {
1732                reason: format!(
1733                    "fit_binomial_mean_wiggle: refit returned {} log-lambdas for {expected_log_lambdas} penalties",
1734                    fit.log_lambdas.len()
1735                ),
1736            }
1737            .into());
1738        }
1739        eta_block_warm.initial_beta = Some(new_source_beta.clone());
1740        eta_block_warm.initial_log_lambdas =
1741            Some(fit.log_lambdas.slice(s![0..eta_penalty_count]).to_owned());
1742        wiggle_beta_warm = Some(new_wiggle_beta);
1743        wiggle_log_lambda_warm = Some(
1744            fit.log_lambdas
1745                .slice(s![eta_penalty_count..expected_log_lambdas])
1746                .to_owned(),
1747        );
1748        frozen_source_beta = new_source_beta;
1749        frozen_eta = new_eta;
1750    }
1751    let (mut fit, last_alias, frozen_source_beta) = converged.ok_or_else(|| {
1752        GamlssError::NumericalFailure {
1753            reason: format!(
1754                "fit_binomial_mean_wiggle frozen-index fixed point did not converge in {} outer iterations: delta={last_delta:.3e}, scale={last_scale:.3e}, tolerance={:.3e}",
1755                options.outer_max_iter,
1756                options.outer_tol * last_scale,
1757            ),
1758        }
1759        .to_string()
1760    })?;
1761    // Capture the mean coefficients whose linear predictor is the *frozen index*
1762    // `η̂` the warp basis `B(η̂)` was pinned at (#2141). The reported deviance is
1763    // evaluated with `q = X·β_saved + B(η̂)·β_w`, so `predict` must re-evaluate the
1764    // warp basis at `X·β_frozen_source` (= η̂), NOT at the de-aliased base
1765    // predictor `X·β_saved`. On the failing data those differ by the identifiable
1766    // de-alias projection `X·A·β_w`, so predict-at-`β_saved` reconstructs a
1767    // *different* link than the fit used. We persist the shift
1768    // `s = β_frozen_source − β_saved` (a mean-coordinate vector) so predict can form the
1769    // frozen index `X·(β_saved + s) = η̂` and reproduce the fitted `q` exactly.
1770    // The solver coefficient is already the standard I-spline coefficient:
1771    // observation-space residualization changed the design, not its coefficient
1772    // chart. The family imposes β_w ≥ 0 during the continuously optimized
1773    // constrained REML/LAML fit. Since B' is an M-spline basis with non-negative
1774    // values, dq/dη = 1 + B'(η)·β_w ≥ 1 for every η, including between
1775    // knots; no post-fit sampling or smoothing-parameter ladder is needed.
1776    let saved_warp_beta = fit
1777        .block_states
1778        .get(BinomialMeanWiggleFamily::BLOCK_WIGGLE)
1779        .map(|state| state.beta.to_vec())
1780        .ok_or_else(|| {
1781            "fit_binomial_mean_wiggle: converged fit is missing its LinkWiggle block state"
1782                .to_string()
1783        })?;
1784    validate_monotone_wiggle_beta_nonnegative(
1785        &saved_warp_beta,
1786        "fit_binomial_mean_wiggle saved warp",
1787    )?;
1788    finalize_binomial_mean_wiggle_saved_frame(
1789        &mut fit,
1790        &last_alias,
1791        &x_dense,
1792        &eta_block_input.offset,
1793    )?;
1794    // The frozen-index shift `s = β_frozen_source − β_saved` for the predict
1795    // runtime (#2141). `β_saved` is the just-de-aliased mean block; adding
1796    // `X·s` to the predict base predictor recovers the frozen warp index `η̂`.
1797    // Only meaningful when a warp actually engaged (`saved_warp_beta` present).
1798    let saved_mean_state = fit
1799        .block_states
1800        .get(BinomialMeanWiggleFamily::BLOCK_ETA)
1801        .ok_or_else(|| {
1802            "fit_binomial_mean_wiggle: finalized fit is missing its Mean block state".to_string()
1803        })?;
1804    if frozen_source_beta.len() != saved_mean_state.beta.len() {
1805        return Err(format!(
1806            "fit_binomial_mean_wiggle: frozen-index source has {} coefficients, but saved Mean block has {}",
1807            frozen_source_beta.len(),
1808            saved_mean_state.beta.len(),
1809        ));
1810    }
1811    let saved_index_shift = Some((&frozen_source_beta - &saved_mean_state.beta).to_vec());
1812    Ok((fit, Some(saved_warp_beta), saved_index_shift))
1813}
1814
1815/// Densify a wiggle-block penalty spec to its full `p×p` matrix for the
1816/// observation-space de-aliasing path (#1596). The link-warp block carries only
1817/// `Dense`/`DenseWithMean` difference (and optional ridge) penalties.
1818fn penalty_spec_to_dense(
1819    spec: &crate::model_types::PenaltySpec,
1820    p: usize,
1821) -> Result<Array2<f64>, String> {
1822    use crate::model_types::PenaltySpec;
1823    match spec {
1824        PenaltySpec::Dense(m) | PenaltySpec::DenseWithMean { matrix: m, .. } => {
1825            if m.nrows() != p || m.ncols() != p {
1826                return Err(format!(
1827                    "frozen-basis warp penalty must be {p}x{p}, got {}x{}",
1828                    m.nrows(),
1829                    m.ncols()
1830                ));
1831            }
1832            Ok(m.clone())
1833        }
1834        PenaltySpec::Block {
1835            local, col_range, ..
1836        } => {
1837            let mut full = Array2::<f64>::zeros((p, p));
1838            if col_range.end > p || local.nrows() != col_range.len() {
1839                return Err("frozen-basis warp penalty block range out of bounds".to_string());
1840            }
1841            full.slice_mut(s![col_range.clone(), col_range.clone()])
1842                .assign(local);
1843            Ok(full)
1844        }
1845    }
1846}
1847
1848pub(crate) trait LocationScaleFamilyBuilder {
1849    type Family: CustomFamily + Clone + Send + Sync + 'static;
1850
1851    fn meanspec(&self) -> &TermCollectionSpec;
1852    fn noisespec(&self) -> &TermCollectionSpec;
1853
1854    fn build_blocks(
1855        &self,
1856        theta: &Array1<f64>,
1857        mean_design: &TermCollectionDesign,
1858        noise_design: &TermCollectionDesign,
1859        mean_beta_hint: Option<Array1<f64>>,
1860        noise_beta_hint: Option<Array1<f64>>,
1861    ) -> Result<Vec<ParameterBlockSpec>, String>;
1862
1863    fn build_family(
1864        &self,
1865        mean_design: &TermCollectionDesign,
1866        noise_design: &TermCollectionDesign,
1867    ) -> Self::Family;
1868
1869    fn extract_primary_betas(
1870        &self,
1871        fit: &UnifiedFitResult,
1872    ) -> Result<(Array1<f64>, Array1<f64>), String>;
1873
1874    fn mean_penalty_count(&self, mean_design: &TermCollectionDesign) -> usize {
1875        mean_design.penalties.len()
1876    }
1877
1878    fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
1879        noise_design.penalties.len()
1880    }
1881
1882    fn exact_spatial_joint_supported(&self) -> bool {
1883        false
1884    }
1885
1886    fn require_exact_spatial_joint(&self) -> bool {
1887        false
1888    }
1889
1890    fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
1891        crate::seeding::SeedRiskProfile::GeneralizedLinear
1892    }
1893
1894    fn extra_rho0(&self) -> Result<Array1<f64>, String> {
1895        Ok(Array1::zeros(0))
1896    }
1897
1898    fn build_psiderivative_blocks(
1899        &self,
1900        arr: ndarray::ArrayView2<'_, f64>,
1901        term_spec: &TermCollectionSpec,
1902        term_spec2: &TermCollectionSpec,
1903        term_design: &TermCollectionDesign,
1904        term_design2: &TermCollectionDesign,
1905    ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String>;
1906}
1907
1908pub(crate) fn fit_location_scale_terms<B: LocationScaleFamilyBuilder>(
1909    data: ndarray::ArrayView2<'_, f64>,
1910    builder: B,
1911    options: &BlockwiseFitOptions,
1912    kappa_options: &SpatialLengthScaleOptimizationOptions,
1913) -> Result<BlockwiseTermFitResult, String> {
1914    // Large-n location-scale fits keep the caller's explicit Hessian request.
1915    // The unified REML evaluator chooses a dense or matrix-free exact
1916    // representation from the realized (n, p, K) work model, so there is no
1917    // large-scale downgrade to BFGS here.
1918
1919    let mut mean_beta_hint: Option<Array1<f64>> = None;
1920    let mut noise_beta_hint: Option<Array1<f64>> = None;
1921    let extra_rho0 = builder.extra_rho0()?;
1922
1923    let mean_boot_design =
1924        build_term_collection_design(data, builder.meanspec()).map_err(|e| e.to_string())?;
1925    let noise_boot_design =
1926        build_term_collection_design(data, builder.noisespec()).map_err(|e| e.to_string())?;
1927    let mean_bootspec = freeze_term_collection_from_design(builder.meanspec(), &mean_boot_design)
1928        .map_err(|e| e.to_string())?;
1929    let noise_bootspec =
1930        freeze_term_collection_from_design(builder.noisespec(), &noise_boot_design)
1931            .map_err(|e| e.to_string())?;
1932
1933    let require_exact_spatial_joint = builder.require_exact_spatial_joint();
1934    let analytic_joint_derivatives_check = if builder.exact_spatial_joint_supported() {
1935        builder
1936            .build_psiderivative_blocks(
1937                data,
1938                &mean_bootspec,
1939                &noise_bootspec,
1940                &mean_boot_design,
1941                &noise_boot_design,
1942            )
1943            .map(|_| ())
1944    } else {
1945        Err(
1946            "analytic spatial psi derivatives are unavailable for this location-scale family"
1947                .to_string(),
1948        )
1949    };
1950    let analytic_joint_derivatives_available = analytic_joint_derivatives_check.is_ok();
1951    if require_exact_spatial_joint {
1952        analytic_joint_derivatives_check.map_err(|err| {
1953            format!("exact two-block spatial path requires analytic psi derivatives: {err}")
1954        })?;
1955    }
1956    let mean_penalty_count = builder.mean_penalty_count(&mean_boot_design);
1957    let noise_penalty_count = builder.noise_penalty_count(&noise_boot_design);
1958
1959    // Honor an explicit user-supplied `length_scale=X` on every spatial term
1960    // in both the mean and noise blocks: when every term is κ-locked (no
1961    // anisotropy, no per-axis ψ contrasts), the joint-spatial outer optimizer
1962    // has nothing to optimize. Routing through it anyway wraps the full
1963    // two-block coefficient solve inside an unnecessary outer loop where
1964    // each evaluation runs the inner Newton from scratch. This is the same
1965    // short-circuit the Bernoulli marginal-slope entry point performs at
1966    // bernoulli_marginal_slope.rs:16432-16442; mirroring it here makes the
1967    // GAMLSS path skip straight to the `(!enabled || log_kappa_dim == 0)`
1968    // fast path in `optimize_spatial_length_scale_exact_joint`.
1969    let mut effective_kappa_options = kappa_options.clone();
1970    if effective_kappa_options.enabled
1971        && gam_terms::smooth::all_spatial_terms_kappa_fixed(&mean_bootspec)
1972        && gam_terms::smooth::all_spatial_terms_kappa_fixed(&noise_bootspec)
1973    {
1974        log::info!(
1975            "[GAMLSS spatial] disabling κ/ψ optimization: every spatial term in \
1976             both blocks has an explicit length_scale and no anisotropy; \
1977             user-supplied kernel scale is fixed"
1978        );
1979        effective_kappa_options.enabled = false;
1980    }
1981    let kappa_options: &SpatialLengthScaleOptimizationOptions = &effective_kappa_options;
1982
1983    // Macro to invoke the exact-joint spatial optimizer with shared closures.
1984    // The exact path evaluates the full profiled/Laplace objective over
1985    // theta = [rho, psi] with the real joint Hessian required by NewtonTR/ARC.
1986    macro_rules! run_exact_joint_spatial {
1987        () => {{
1988            let joint_setup = build_two_block_exact_joint_setup(
1989                data,
1990                builder.meanspec(),
1991                builder.noisespec(),
1992                mean_penalty_count,
1993                noise_penalty_count,
1994                extra_rho0.as_slice().unwrap_or(&[]),
1995                None,
1996                kappa_options,
1997            )
1998            .map_err(|error| error.to_string())?;
1999            let mean_terms = spatial_length_scale_term_indices(builder.meanspec());
2000            let noise_terms = spatial_length_scale_term_indices(builder.noisespec());
2001            let mean_beta_hint_cell = std::cell::RefCell::new(mean_beta_hint.clone());
2002            let noise_beta_hint_cell = std::cell::RefCell::new(noise_beta_hint.clone());
2003            let hyper_warm_start_cell =
2004                std::cell::RefCell::new(None::<CustomFamilyWarmStart>);
2005            // Two-block GAMLSS/location-scale joint likelihoods have a
2006            // β-dependent cross-block Hessian (the (μ,log σ) / (t,log σ)
2007            // off-diagonal blocks involve residual/response scalars that
2008            // shift when β moves). The Wood-Fasiolo structural property
2009            // `H^{-1/2} B_k H^{-1/2} ≽ 0` plus parameter-independent
2010            // nullspace — the mathematical basis for EFS convergence —
2011            // fails here, so EFS/HybridEFS must be excluded at plan time
2012            // rather than retried as a silent first attempt that stalls
2013            // for hundreds of seconds before the runner falls back.
2014            let gamlss_disable_fixed_point = true;
2015            let outer_policy = {
2016                // GAMLSS spatial path: psi_dim = log_kappa_dim + auxiliary_dim,
2017                // matching the (theta_dim - rho_dim) decomposition the
2018                // optimizer uses internally. Build realized ParameterBlockSpecs
2019                // at the seed rho so the family's own cost model — which
2020                // multiplies coefficient-gradient / coefficient-Hessian
2021                // per-row cost by the joint outer-coordinate dimension and
2022                // total p — produces honest `predicted_*_work` estimates.
2023                // Previously this fed `predicted_*_work: 0` to the planner,
2024                // which then ungated dense outer Hessian work that costs
2025                // hundreds of seconds per eval at large scale (see
2026                // `OuterDerivativePolicy::OUTER_HESSIAN_WORK_BUDGET`).
2027                let theta_seed = joint_setup.theta0();
2028                let rho_dim = joint_setup.rho_dim();
2029                let psi_dim = theta_seed.len() - rho_dim;
2030                let rho_seed = theta_seed.slice(s![..rho_dim]).to_owned();
2031                let policy_blocks_res = builder.build_blocks(
2032                    &rho_seed,
2033                    &mean_boot_design,
2034                    &noise_boot_design,
2035                    mean_beta_hint_cell.borrow().clone(),
2036                    noise_beta_hint_cell.borrow().clone(),
2037                );
2038                let mut policy = match policy_blocks_res {
2039                    Ok(policy_blocks) => {
2040                        let policy_family =
2041                            builder.build_family(&mean_boot_design, &noise_boot_design);
2042                        crate::custom_family::CustomFamily::outer_derivative_policy(
2043                            &policy_family,
2044                            &policy_blocks,
2045                            psi_dim,
2046                            options,
2047                        )
2048                    }
2049                    Err(err) => {
2050                        // Block construction at the seed should not fail for
2051                        // any in-tree family, but if it does, fall back to a
2052                        // policy that names the capability honestly and
2053                        // declines to predict cost. Setting work to
2054                        // `u128::MAX` routes the planner through gradient-only
2055                        // BFGS (the universal Hessian-work budget is
2056                        // saturating, so a sentinel is fine here).
2057                        log::warn!(
2058                            "[GAMLSS spatial] failed to realize policy blocks at seed rho ({err}); \
2059                             routing outer optimizer through gradient-only BFGS"
2060                        );
2061                        let capability = if analytic_joint_derivatives_available {
2062                            crate::custom_family::ExactOuterDerivativeOrder::Second
2063                        } else {
2064                            crate::custom_family::ExactOuterDerivativeOrder::First
2065                        };
2066                        crate::custom_family::OuterDerivativePolicy {
2067                            capability,
2068                            predicted_gradient_work: u128::MAX,
2069                            predicted_hessian_work: u128::MAX,
2070                            // No GAMLSS family today overrides its
2071                            // outer-only `_with_options` hooks to consume
2072                            // `outer_score_subsample`; staged-κ would
2073                            // build pilot masks the family then ignores.
2074                            subsample_capable: false,
2075                        }
2076                    }
2077                };
2078                if !analytic_joint_derivatives_available {
2079                    // Capability must not exceed what the analytic derivatives
2080                    // path can supply — the macro's hyper evaluator returns
2081                    // an error otherwise.
2082                    policy.capability =
2083                        crate::custom_family::ExactOuterDerivativeOrder::First;
2084                }
2085                policy
2086            };
2087            optimize_spatial_length_scale_exact_joint(
2088                data,
2089                &[builder.meanspec().clone(), builder.noisespec().clone()],
2090                &[mean_terms, noise_terms],
2091                kappa_options,
2092                &joint_setup,
2093                builder.exact_spatial_seed_risk_profile(),
2094                analytic_joint_derivatives_available,
2095                analytic_joint_derivatives_available,
2096                gamlss_disable_fixed_point,
2097                None,
2098                outer_policy,
2099                |theta,
2100                 specs: &[TermCollectionSpec],
2101                 designs: &[TermCollectionDesign],
2102                 provenance| {
2103                    assert_eq!(
2104                        specs.len(),
2105                        2,
2106                        "joint spatial closure expects exactly two block specs (mean, noise); got {}",
2107                        specs.len(),
2108                    );
2109                    assert_eq!(
2110                        designs.len(),
2111                        2,
2112                        "joint spatial closure expects exactly two block designs (mean, noise); got {}",
2113                        designs.len(),
2114                    );
2115                    let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
2116                    let fit = {
2117                        let blocks = builder.build_blocks(
2118                            &rho,
2119                            &designs[0],
2120                            &designs[1],
2121                            mean_beta_hint_cell.borrow().clone(),
2122                            noise_beta_hint_cell.borrow().clone(),
2123                        )?;
2124                        if mean_beta_hint_cell.borrow().is_none()
2125                            && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
2126                        {
2127                            *mean_beta_hint_cell.borrow_mut() = Some(beta);
2128                        }
2129                        if noise_beta_hint_cell.borrow().is_none()
2130                            && let Some(beta) =
2131                                blocks.get(1).and_then(|block| block.initial_beta.clone())
2132                        {
2133                            *noise_beta_hint_cell.borrow_mut() = Some(beta);
2134                        }
2135                        let family = builder.build_family(&designs[0], &designs[1]);
2136                        // Branch on whether the κ optimizer drives rho.
2137                        //
2138                        // * `log_kappa_dim() > 0 && kappa_options.enabled` ⇒
2139                        //   the outer (ρ, ψ) optimizer is active and
2140                        //   passes each candidate ρ to this closure;
2141                        //   the inner fit must hold log-lambdas fixed
2142                        //   at the supplied ρ so the outer derivative
2143                        //   has a well-defined directional gradient.
2144                        //
2145                        // * Otherwise (κ disabled via the locked-κ
2146                        //   short-circuit, or no spatial terms at all)
2147                        //   the fast path in
2148                        //   `optimize_spatial_length_scale_exact_joint`
2149                        //   calls this closure exactly once at
2150                        //   `theta = theta0`; ρ must still be optimized
2151                        //   from data because the user never pinned it.
2152                        //   `fit_custom_family` performs the joint
2153                        //   ρ + coefficient REML fit at the user's
2154                        //   (now-fixed) kernel scale, which is the
2155                        //   intended behaviour when `length_scale=…` is
2156                        //   set on every spatial term.
2157                        if joint_setup.log_kappa_dim() > 0 && kappa_options.enabled {
2158                            let (certified_outer, mode) = match provenance {
2159                                SpatialFitProvenance::Certified { outer, mode } => (outer, mode),
2160                                SpatialFitProvenance::NoOuterOptimization => {
2161                                    return Err(
2162                                        "active GAMLSS spatial optimization returned no certified outer provenance"
2163                                            .to_string(),
2164                                    );
2165                                }
2166                            };
2167                            let exact_options =
2168                                crate::outer_subsample::exact_outer_options_for_row_set(
2169                                    options,
2170                                    &crate::row_kernel::RowSet::All,
2171                                );
2172                            fit_custom_family_fixed_log_lambdas_from_owned_mode(
2173                                &family,
2174                                &blocks,
2175                                &exact_options,
2176                                mode,
2177                                theta,
2178                                certified_outer,
2179                            )?
2180                        } else {
2181                            fit_custom_family(&family, &blocks, options)?
2182                        }
2183                    };
2184                    let (mean_beta, noise_beta) = builder.extract_primary_betas(&fit)?;
2185                    mean_beta_hint = Some(mean_beta);
2186                    noise_beta_hint = Some(noise_beta);
2187                    *mean_beta_hint_cell.borrow_mut() = mean_beta_hint.clone();
2188                    *noise_beta_hint_cell.borrow_mut() = noise_beta_hint.clone();
2189                    Ok(fit)
2190                },
2191                |theta,
2192                 specs: &[TermCollectionSpec],
2193                 designs: &[TermCollectionDesign],
2194                 eval_mode,
2195                 row_set: &crate::row_kernel::RowSet| {
2196                    use gam_problem::EvalMode;
2197                    if !analytic_joint_derivatives_available {
2198                        return Err(
2199                            "analytic spatial psi derivatives are unavailable for this exact two-block path"
2200                                .to_string(),
2201                        );
2202                    }
2203                    let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
2204                    let blocks = builder.build_blocks(
2205                        &rho,
2206                        &designs[0],
2207                        &designs[1],
2208                        mean_beta_hint_cell.borrow().clone(),
2209                        noise_beta_hint_cell.borrow().clone(),
2210                    )?;
2211                    if mean_beta_hint_cell.borrow().is_none()
2212                        && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
2213                    {
2214                        *mean_beta_hint_cell.borrow_mut() = Some(beta);
2215                    }
2216                    if noise_beta_hint_cell.borrow().is_none()
2217                        && let Some(beta) = blocks.get(1).and_then(|block| block.initial_beta.clone())
2218                    {
2219                        *noise_beta_hint_cell.borrow_mut() = Some(beta);
2220                    }
2221                    let family = builder.build_family(&designs[0], &designs[1]);
2222                    let psiderivative_blocks = builder.build_psiderivative_blocks(
2223                        data,
2224                        &specs[0],
2225                        &specs[1],
2226                        &designs[0],
2227                        &designs[1],
2228                    )?;
2229                    let hyper_layout = crate::custom_family::CustomFamilyHyperLayout::new(
2230                        psiderivative_blocks,
2231                        Vec::new(),
2232                        theta.slice(s![joint_setup.rho_dim()..]).to_owned(),
2233                    )?;
2234                    let warm_start = hyper_warm_start_cell.borrow().clone();
2235                    // Forward the κ-staging row set to the family by installing it
2236                    // on the canonical `outer_score_subsample` option. Inner-PIRLS
2237                    // and final covariance still run on full data (the per-row
2238                    // weight is consulted only by outer-only paths inside the
2239                    // family). When the staging schedule is full-data the option
2240                    // stays `None` and the call is equivalent to the prior path.
2241                    let eval_options =
2242                        crate::outer_subsample::exact_outer_options_for_row_set(options, row_set);
2243                    let owned = evaluate_custom_family_joint_hyper_owned(
2244                        &family,
2245                        &blocks,
2246                        &eval_options,
2247                        &rho,
2248                        &hyper_layout,
2249                        warm_start.as_ref(),
2250                        eval_mode,
2251                    )?;
2252                    *hyper_warm_start_cell.borrow_mut() = Some(owned.result.warm_start.clone());
2253                    if !owned.result.inner_converged {
2254                        return Err(
2255                            "exact two-block spatial inner solve did not converge".to_string(),
2256                        );
2257                    }
2258                    if matches!(eval_mode, EvalMode::ValueGradientHessian)
2259                        && !owned.result.outer_hessian.is_analytic()
2260                    {
2261                        return Err(
2262                            "exact two-block spatial objective requires a full joint [rho, psi] hessian"
2263                            .to_string(),
2264                        );
2265                    }
2266                    Ok(ExactJointEvaluation {
2267                        objective: owned.result.objective,
2268                        gradient: owned.result.gradient,
2269                        hessian: owned.result.outer_hessian,
2270                        mode: owned.mode,
2271                    })
2272                },
2273                |theta,
2274                 specs: &[TermCollectionSpec],
2275                 designs: &[TermCollectionDesign],
2276                 row_set: &crate::row_kernel::RowSet| {
2277                    if !analytic_joint_derivatives_available {
2278                        return Err(
2279                            "analytic spatial psi derivatives are unavailable for this exact two-block path"
2280                                .to_string(),
2281                        );
2282                    }
2283                    let rho = theta.slice(s![..joint_setup.rho_dim()]).to_owned();
2284                    let blocks = builder.build_blocks(
2285                        &rho,
2286                        &designs[0],
2287                        &designs[1],
2288                        mean_beta_hint_cell.borrow().clone(),
2289                        noise_beta_hint_cell.borrow().clone(),
2290                    )?;
2291                    if mean_beta_hint_cell.borrow().is_none()
2292                        && let Some(beta) = blocks.first().and_then(|block| block.initial_beta.clone())
2293                    {
2294                        *mean_beta_hint_cell.borrow_mut() = Some(beta);
2295                    }
2296                    if noise_beta_hint_cell.borrow().is_none()
2297                        && let Some(beta) = blocks.get(1).and_then(|block| block.initial_beta.clone())
2298                    {
2299                        *noise_beta_hint_cell.borrow_mut() = Some(beta);
2300                    }
2301                    let family = builder.build_family(&designs[0], &designs[1]);
2302                    let psiderivative_blocks = builder.build_psiderivative_blocks(
2303                        data,
2304                        &specs[0],
2305                        &specs[1],
2306                        &designs[0],
2307                        &designs[1],
2308                    )?;
2309                    let hyper_layout = crate::custom_family::CustomFamilyHyperLayout::new(
2310                        psiderivative_blocks,
2311                        Vec::new(),
2312                        theta.slice(s![joint_setup.rho_dim()..]).to_owned(),
2313                    )?;
2314                    let warm_start = hyper_warm_start_cell.borrow().clone();
2315                    let eval_options =
2316                        crate::outer_subsample::exact_outer_options_for_row_set(options, row_set);
2317                    let owned = evaluate_custom_family_joint_hyper_efs_owned(
2318                        &family,
2319                        &blocks,
2320                        &eval_options,
2321                        &rho,
2322                        &hyper_layout,
2323                        warm_start.as_ref(),
2324                    )?;
2325                    *hyper_warm_start_cell.borrow_mut() = Some(owned.result.warm_start.clone());
2326                    if !owned.result.inner_converged {
2327                        return Err(
2328                            "exact two-block spatial EFS inner solve did not converge".to_string(),
2329                        );
2330                    }
2331                    Ok(ExactJointEfsEvaluation {
2332                        evaluation: owned.result.efs_eval,
2333                        mode: owned.mode,
2334                    })
2335                },
2336                |_beta: &Array1<f64>| Ok(gam_solve::rho_optimizer::SeedOutcome::NoSlot),
2337            )
2338        }};
2339    }
2340
2341    let mut solved = run_exact_joint_spatial!()
2342        .map_err(|err| format!("exact two-block spatial optimization failed: {err}"))?;
2343
2344    let expected_noise_penalty_count = builder.noise_penalty_count(&solved.designs[1]);
2345    let actual_noise_penalty_count = solved.designs[1].penalties.len();
2346    if expected_noise_penalty_count > actual_noise_penalty_count {
2347        if expected_noise_penalty_count != actual_noise_penalty_count + 1 {
2348            return Err(GamlssError::UnsupportedConfiguration {
2349                reason: format!(
2350                    "location-scale result noise design expected {} penalties after augmentation, got {} before augmentation",
2351                    expected_noise_penalty_count, actual_noise_penalty_count
2352                ),
2353            }
2354            .into());
2355        }
2356        append_binomial_log_sigma_shrinkage_penalty_design(&mut solved.designs[1]);
2357    }
2358
2359    BlockwiseTermFitResult::try_from_parts(BlockwiseTermFitResultParts {
2360        fit: solved.fit,
2361        meanspec_resolved: solved.resolved_specs.remove(0),
2362        noisespec_resolved: solved.resolved_specs.remove(0),
2363        mean_design: solved.designs.remove(0),
2364        noise_design: solved.designs.remove(0),
2365    })
2366}
2367
2368pub(crate) struct GaussianLocationScaleTermBuilder {
2369    pub(crate) y: Array1<f64>,
2370    pub(crate) weights: Array1<f64>,
2371    pub(crate) meanspec: TermCollectionSpec,
2372    pub(crate) noisespec: TermCollectionSpec,
2373    pub(crate) mean_offset: Array1<f64>,
2374    pub(crate) noise_offset: Array1<f64>,
2375}
2376
2377impl LocationScaleFamilyBuilder for GaussianLocationScaleTermBuilder {
2378    type Family = GaussianLocationScaleFamily;
2379
2380    fn meanspec(&self) -> &TermCollectionSpec {
2381        &self.meanspec
2382    }
2383
2384    fn noisespec(&self) -> &TermCollectionSpec {
2385        &self.noisespec
2386    }
2387
2388    fn exact_spatial_joint_supported(&self) -> bool {
2389        true
2390    }
2391
2392    fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
2393        crate::seeding::SeedRiskProfile::GaussianLocationScale
2394    }
2395
2396    fn build_blocks(
2397        &self,
2398        theta: &Array1<f64>,
2399        mean_design: &TermCollectionDesign,
2400        noise_design: &TermCollectionDesign,
2401        mean_beta_hint: Option<Array1<f64>>,
2402        noise_beta_hint: Option<Array1<f64>>,
2403    ) -> Result<Vec<ParameterBlockSpec>, String> {
2404        let layout = GamlssLambdaLayout::two_block(
2405            mean_design.penalties.len(),
2406            self.noise_penalty_count(noise_design),
2407        );
2408        layout.validate_theta_len(theta.len(), "gaussian location-scale")?;
2409        let (meanspec, noisespec) = build_gaussian_mean_and_scale_blocks(
2410            &self.y,
2411            &self.weights,
2412            mean_design,
2413            noise_design,
2414            &self.mean_offset,
2415            &self.noise_offset,
2416            layout.mean_from(theta),
2417            layout.noise_from(theta),
2418            mean_beta_hint,
2419            noise_beta_hint,
2420            "GaussianLocationScale::build_blocks",
2421        )?;
2422        Ok(vec![meanspec, noisespec])
2423    }
2424
2425    fn build_family(
2426        &self,
2427        mean_design: &TermCollectionDesign,
2428        noise_design: &TermCollectionDesign,
2429    ) -> Self::Family {
2430        let preparednoise_design =
2431            prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design)
2432                .expect("prepared Gaussian log-sigma design should match block construction");
2433        GaussianLocationScaleFamily {
2434            y: self.y.clone(),
2435            weights: self.weights.clone(),
2436            mu_design: Some(mean_design.design.clone()),
2437            log_sigma_design: Some(preparednoise_design),
2438            policy: gam_runtime::resource::ResourcePolicy::default_library(),
2439            cached_row_scalars: std::sync::RwLock::new(None),
2440        }
2441    }
2442
2443    fn extract_primary_betas(
2444        &self,
2445        fit: &UnifiedFitResult,
2446    ) -> Result<(Array1<f64>, Array1<f64>), String> {
2447        let mean_beta = fit
2448            .block_states
2449            .get(GaussianLocationScaleFamily::BLOCK_MU)
2450            .ok_or_else(|| "missing Gaussian mu block state".to_string())?
2451            .beta
2452            .clone();
2453        let noise_beta = fit
2454            .block_states
2455            .get(GaussianLocationScaleFamily::BLOCK_LOG_SIGMA)
2456            .ok_or_else(|| "missing Gaussian log_sigma block state".to_string())?
2457            .beta
2458            .clone();
2459        Ok((mean_beta, noise_beta))
2460    }
2461
2462    fn build_psiderivative_blocks(
2463        &self,
2464        data: ndarray::ArrayView2<'_, f64>,
2465        meanspec_resolved: &TermCollectionSpec,
2466        noisespec_resolved: &TermCollectionSpec,
2467        mean_design: &TermCollectionDesign,
2468        noise_design: &TermCollectionDesign,
2469    ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2470        let mean_derivs =
2471            build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2472                .ok_or_else(|| "missing Gaussian mean spatial psi derivatives".to_string())?;
2473        let noise_derivs =
2474            build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2475                .ok_or_else(|| "missing Gaussian log-sigma spatial psi derivatives".to_string())?;
2476        Ok(vec![mean_derivs, noise_derivs])
2477    }
2478}
2479
2480pub(crate) struct GaussianLocationScaleWiggleTermBuilder {
2481    pub(crate) y: Array1<f64>,
2482    pub(crate) weights: Array1<f64>,
2483    pub(crate) meanspec: TermCollectionSpec,
2484    pub(crate) noisespec: TermCollectionSpec,
2485    pub(crate) mean_offset: Array1<f64>,
2486    pub(crate) noise_offset: Array1<f64>,
2487    pub(crate) wiggle_knots: Array1<f64>,
2488    pub(crate) wiggle_degree: usize,
2489    pub(crate) wiggle_block: ParameterBlockInput,
2490}
2491
2492impl LocationScaleFamilyBuilder for GaussianLocationScaleWiggleTermBuilder {
2493    type Family = GaussianLocationScaleWiggleFamily;
2494
2495    fn meanspec(&self) -> &TermCollectionSpec {
2496        &self.meanspec
2497    }
2498
2499    fn noisespec(&self) -> &TermCollectionSpec {
2500        &self.noisespec
2501    }
2502
2503    fn exact_spatial_joint_supported(&self) -> bool {
2504        true
2505    }
2506
2507    fn exact_spatial_seed_risk_profile(&self) -> crate::seeding::SeedRiskProfile {
2508        crate::seeding::SeedRiskProfile::GaussianLocationScale
2509    }
2510
2511    fn require_exact_spatial_joint(&self) -> bool {
2512        true
2513    }
2514
2515    fn extra_rho0(&self) -> Result<Array1<f64>, String> {
2516        initial_log_lambdas_orzeros(&self.wiggle_block)
2517    }
2518
2519    fn build_blocks(
2520        &self,
2521        theta: &Array1<f64>,
2522        mean_design: &TermCollectionDesign,
2523        noise_design: &TermCollectionDesign,
2524        mean_beta_hint: Option<Array1<f64>>,
2525        noise_beta_hint: Option<Array1<f64>>,
2526    ) -> Result<Vec<ParameterBlockSpec>, String> {
2527        let layout = GamlssLambdaLayout::withwiggle(
2528            mean_design.penalties.len(),
2529            self.noise_penalty_count(noise_design),
2530            self.wiggle_block.penalties.len(),
2531        );
2532        layout.validate_theta_len(theta.len(), "gaussian location-scale wiggle")?;
2533        let (mut meanspec, mut noisespec) = build_gaussian_mean_and_scale_blocks(
2534            &self.y,
2535            &self.weights,
2536            mean_design,
2537            noise_design,
2538            &self.mean_offset,
2539            &self.noise_offset,
2540            layout.mean_from(theta),
2541            layout.noise_from(theta),
2542            mean_beta_hint,
2543            noise_beta_hint,
2544            "GaussianLocationScaleWiggle::build_blocks",
2545        )?;
2546        // Keep the dynamic full-width wiggle basis safe from a canonical-gauge
2547        // column drop: route the shared level/intercept alias onto the
2548        // column-reducible mean and log-sigma blocks by giving them a lower
2549        // gauge priority than the wiggle block's fixed 100 (see the binomial
2550        // wiggle path and `build_location_scale_wiggle_block`).
2551        meanspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2552        noisespec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2553        let n_rows = meanspec.design.nrows();
2554        let wigglespec = build_location_scale_wiggle_block(
2555            "wiggle",
2556            self.wiggle_block.design.clone(),
2557            self.wiggle_block.offset.clone(),
2558            wiggle_block_penalty_matrices(&self.wiggle_block),
2559            self.wiggle_block.nullspace_dims.clone(),
2560            layout.wiggle_from(theta),
2561            self.wiggle_block.initial_beta.clone(),
2562            n_rows,
2563        )?;
2564        Ok(vec![meanspec, noisespec, wigglespec])
2565    }
2566
2567    fn build_family(
2568        &self,
2569        mean_design: &TermCollectionDesign,
2570        noise_design: &TermCollectionDesign,
2571    ) -> Self::Family {
2572        let preparednoise_design =
2573            prepared_gaussian_log_sigma_design(&mean_design.design, &noise_design.design).expect(
2574                "prepared Gaussian log-sigma design should match wiggle block construction",
2575            );
2576        GaussianLocationScaleWiggleFamily {
2577            y: self.y.clone(),
2578            weights: self.weights.clone(),
2579            mu_design: Some(mean_design.design.clone()),
2580            log_sigma_design: Some(preparednoise_design),
2581            wiggle_knots: self.wiggle_knots.clone(),
2582            wiggle_degree: self.wiggle_degree,
2583            policy: gam_runtime::resource::ResourcePolicy::default_library(),
2584            cached_row_scalars: std::sync::RwLock::new(None),
2585        }
2586    }
2587
2588    fn extract_primary_betas(
2589        &self,
2590        fit: &UnifiedFitResult,
2591    ) -> Result<(Array1<f64>, Array1<f64>), String> {
2592        let mean_beta = fit
2593            .block_states
2594            .get(GaussianLocationScaleWiggleFamily::BLOCK_MU)
2595            .ok_or_else(|| "missing Gaussian wiggle mu block state".to_string())?
2596            .beta
2597            .clone();
2598        let noise_beta = fit
2599            .block_states
2600            .get(GaussianLocationScaleWiggleFamily::BLOCK_LOG_SIGMA)
2601            .ok_or_else(|| "missing Gaussian wiggle log_sigma block state".to_string())?
2602            .beta
2603            .clone();
2604        Ok((mean_beta, noise_beta))
2605    }
2606
2607    fn build_psiderivative_blocks(
2608        &self,
2609        data: ndarray::ArrayView2<'_, f64>,
2610        meanspec_resolved: &TermCollectionSpec,
2611        noisespec_resolved: &TermCollectionSpec,
2612        mean_design: &TermCollectionDesign,
2613        noise_design: &TermCollectionDesign,
2614    ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2615        let mean_derivs =
2616            build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?.ok_or_else(
2617                || "missing Gaussian wiggle mean spatial psi derivatives".to_string(),
2618            )?;
2619        let noise_derivs =
2620            build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2621                .ok_or_else(|| {
2622                    "missing Gaussian wiggle log-sigma spatial psi derivatives".to_string()
2623                })?;
2624        Ok(vec![mean_derivs, noise_derivs, Vec::new()])
2625    }
2626}
2627
2628pub(crate) struct BinomialLocationScaleTermBuilder {
2629    pub(crate) y: Array1<f64>,
2630    pub(crate) weights: Array1<f64>,
2631    pub(crate) link_kind: InverseLink,
2632    pub(crate) meanspec: TermCollectionSpec,
2633    pub(crate) noisespec: TermCollectionSpec,
2634    pub(crate) mean_offset: Array1<f64>,
2635    pub(crate) noise_offset: Array1<f64>,
2636}
2637
2638impl LocationScaleFamilyBuilder for BinomialLocationScaleTermBuilder {
2639    type Family = BinomialLocationScaleFamily;
2640
2641    fn meanspec(&self) -> &TermCollectionSpec {
2642        &self.meanspec
2643    }
2644
2645    fn noisespec(&self) -> &TermCollectionSpec {
2646        &self.noisespec
2647    }
2648
2649    fn exact_spatial_joint_supported(&self) -> bool {
2650        true
2651    }
2652
2653    fn require_exact_spatial_joint(&self) -> bool {
2654        true
2655    }
2656
2657    fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
2658        noise_design.penalties.len() + 1
2659    }
2660
2661    fn build_blocks(
2662        &self,
2663        theta: &Array1<f64>,
2664        mean_design: &TermCollectionDesign,
2665        noise_design: &TermCollectionDesign,
2666        mean_beta_hint: Option<Array1<f64>>,
2667        noise_beta_hint: Option<Array1<f64>>,
2668    ) -> Result<Vec<ParameterBlockSpec>, String> {
2669        let layout = GamlssLambdaLayout::two_block(
2670            mean_design.penalties.len(),
2671            self.noise_penalty_count(noise_design),
2672        );
2673        layout.validate_theta_len(theta.len(), "binomial location-scale")?;
2674        let (thresholdspec, log_sigmaspec) = build_binomial_threshold_and_scale_blocks(
2675            &self.y,
2676            &self.weights,
2677            &self.link_kind,
2678            mean_design,
2679            noise_design,
2680            &self.mean_offset,
2681            &self.noise_offset,
2682            layout.mean_from(theta),
2683            layout.noise_from(theta),
2684            mean_beta_hint,
2685            noise_beta_hint,
2686            "BinomialLocationScale::build_blocks",
2687        )?;
2688        Ok(vec![thresholdspec, log_sigmaspec])
2689    }
2690
2691    fn build_family(
2692        &self,
2693        mean_design: &TermCollectionDesign,
2694        noise_design: &TermCollectionDesign,
2695    ) -> Self::Family {
2696        let identifiednoise_design =
2697            identified_binomial_log_sigma_design(mean_design, noise_design, &self.weights)
2698                .expect("identified binomial log-sigma design");
2699        BinomialLocationScaleFamily {
2700            y: self.y.clone(),
2701            weights: self.weights.clone(),
2702            link_kind: self.link_kind.clone(),
2703            threshold_design: Some(mean_design.design.clone()),
2704            log_sigma_design: Some(identifiednoise_design),
2705            policy: gam_runtime::resource::ResourcePolicy::default_library(),
2706        }
2707    }
2708
2709    fn extract_primary_betas(
2710        &self,
2711        fit: &UnifiedFitResult,
2712    ) -> Result<(Array1<f64>, Array1<f64>), String> {
2713        let mean_beta = fit
2714            .block_states
2715            .get(BinomialLocationScaleFamily::BLOCK_T)
2716            .ok_or_else(|| "missing Binomial threshold block state".to_string())?
2717            .beta
2718            .clone();
2719        let noise_beta = fit
2720            .block_states
2721            .get(BinomialLocationScaleFamily::BLOCK_LOG_SIGMA)
2722            .ok_or_else(|| "missing Binomial log_sigma block state".to_string())?
2723            .beta
2724            .clone();
2725        Ok((mean_beta, noise_beta))
2726    }
2727
2728    fn build_psiderivative_blocks(
2729        &self,
2730        data: ndarray::ArrayView2<'_, f64>,
2731        meanspec_resolved: &TermCollectionSpec,
2732        noisespec_resolved: &TermCollectionSpec,
2733        mean_design: &TermCollectionDesign,
2734        noise_design: &TermCollectionDesign,
2735    ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2736        let mean_derivs =
2737            build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2738                .ok_or_else(|| "missing threshold spatial psi derivatives".to_string())?;
2739        let noise_derivs =
2740            build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2741                .ok_or_else(|| "missing log_sigma spatial psi derivatives".to_string())?;
2742        Ok(vec![mean_derivs, noise_derivs])
2743    }
2744}
2745
2746pub(crate) struct BinomialLocationScaleWiggleTermBuilder {
2747    pub(crate) y: Array1<f64>,
2748    pub(crate) weights: Array1<f64>,
2749    pub(crate) link_kind: InverseLink,
2750    pub(crate) meanspec: TermCollectionSpec,
2751    pub(crate) noisespec: TermCollectionSpec,
2752    pub(crate) mean_offset: Array1<f64>,
2753    pub(crate) noise_offset: Array1<f64>,
2754    pub(crate) wiggle_knots: Array1<f64>,
2755    pub(crate) wiggle_degree: usize,
2756    pub(crate) wiggle_block: ParameterBlockInput,
2757}
2758
2759impl LocationScaleFamilyBuilder for BinomialLocationScaleWiggleTermBuilder {
2760    type Family = BinomialLocationScaleWiggleFamily;
2761
2762    fn meanspec(&self) -> &TermCollectionSpec {
2763        &self.meanspec
2764    }
2765
2766    fn noisespec(&self) -> &TermCollectionSpec {
2767        &self.noisespec
2768    }
2769
2770    fn exact_spatial_joint_supported(&self) -> bool {
2771        true
2772    }
2773
2774    fn require_exact_spatial_joint(&self) -> bool {
2775        true
2776    }
2777
2778    fn extra_rho0(&self) -> Result<Array1<f64>, String> {
2779        initial_log_lambdas_orzeros(&self.wiggle_block)
2780    }
2781
2782    fn noise_penalty_count(&self, noise_design: &TermCollectionDesign) -> usize {
2783        noise_design.penalties.len() + 1
2784    }
2785
2786    fn build_blocks(
2787        &self,
2788        theta: &Array1<f64>,
2789        mean_design: &TermCollectionDesign,
2790        noise_design: &TermCollectionDesign,
2791        mean_beta_hint: Option<Array1<f64>>,
2792        noise_beta_hint: Option<Array1<f64>>,
2793    ) -> Result<Vec<ParameterBlockSpec>, String> {
2794        let layout = GamlssLambdaLayout::withwiggle(
2795            mean_design.penalties.len(),
2796            self.noise_penalty_count(noise_design),
2797            self.wiggle_block.penalties.len(),
2798        );
2799        layout.validate_theta_len(theta.len(), "wiggle location-scale")?;
2800        let (mut thresholdspec, mut log_sigmaspec) = build_binomial_threshold_and_scale_blocks(
2801            &self.y,
2802            &self.weights,
2803            &self.link_kind,
2804            mean_design,
2805            noise_design,
2806            &self.mean_offset,
2807            &self.noise_offset,
2808            layout.mean_from(theta),
2809            layout.noise_from(theta),
2810            mean_beta_hint,
2811            noise_beta_hint,
2812            "BinomialLocationScaleWiggle::build_blocks",
2813        )?;
2814        // The dynamic monotone wiggle basis is regenerated at full raw width
2815        // every inner iteration and asserts `x.ncols() == spec.design.ncols()`
2816        // in `block_geometry`, so it cannot tolerate a canonical-gauge column
2817        // drop. The level/intercept direction the I-spline shares with the
2818        // threshold block must therefore be routed onto the threshold (and the
2819        // log-sigma) block, whose static designs are column-reducible and
2820        // lifted back via the canonical per-block transform `T`. Give both
2821        // non-wiggle blocks a lower gauge priority than the wiggle block (which
2822        // `build_location_scale_wiggle_block` fixes at 100) so the shared-level
2823        // alias drop lands on them and leaves the dynamic wiggle basis full
2824        // width — mirroring the binomial mean-wiggle path.
2825        thresholdspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2826        log_sigmaspec.gauge_priority = LINK_WIGGLE_GAUGE_PRIORITY;
2827        let n_rows = thresholdspec.design.nrows();
2828        let wigglespec = build_location_scale_wiggle_block(
2829            "wiggle",
2830            self.wiggle_block.design.clone(),
2831            self.wiggle_block.offset.clone(),
2832            wiggle_block_penalty_matrices(&self.wiggle_block),
2833            vec![],
2834            layout.wiggle_from(theta),
2835            self.wiggle_block.initial_beta.clone(),
2836            n_rows,
2837        )?;
2838        Ok(vec![thresholdspec, log_sigmaspec, wigglespec])
2839    }
2840
2841    fn build_family(
2842        &self,
2843        mean_design: &TermCollectionDesign,
2844        noise_design: &TermCollectionDesign,
2845    ) -> Self::Family {
2846        let identifiednoise_design =
2847            identified_binomial_log_sigma_design(mean_design, noise_design, &self.weights)
2848                .expect("identified binomial log-sigma design should match block construction");
2849        BinomialLocationScaleWiggleFamily {
2850            y: self.y.clone(),
2851            weights: self.weights.clone(),
2852            link_kind: self.link_kind.clone(),
2853            threshold_design: Some(mean_design.design.clone()),
2854            log_sigma_design: Some(identifiednoise_design),
2855            wiggle_knots: self.wiggle_knots.clone(),
2856            wiggle_degree: self.wiggle_degree,
2857            policy: gam_runtime::resource::ResourcePolicy::default_library(),
2858        }
2859    }
2860
2861    fn extract_primary_betas(
2862        &self,
2863        fit: &UnifiedFitResult,
2864    ) -> Result<(Array1<f64>, Array1<f64>), String> {
2865        let mean_beta = fit
2866            .block_states
2867            .get(BinomialLocationScaleWiggleFamily::BLOCK_T)
2868            .ok_or_else(|| "missing Binomial wiggle threshold block state".to_string())?
2869            .beta
2870            .clone();
2871        let noise_beta = fit
2872            .block_states
2873            .get(BinomialLocationScaleWiggleFamily::BLOCK_LOG_SIGMA)
2874            .ok_or_else(|| "missing Binomial wiggle log_sigma block state".to_string())?
2875            .beta
2876            .clone();
2877        Ok((mean_beta, noise_beta))
2878    }
2879
2880    fn build_psiderivative_blocks(
2881        &self,
2882        data: ndarray::ArrayView2<'_, f64>,
2883        meanspec_resolved: &TermCollectionSpec,
2884        noisespec_resolved: &TermCollectionSpec,
2885        mean_design: &TermCollectionDesign,
2886        noise_design: &TermCollectionDesign,
2887    ) -> Result<Vec<Vec<CustomFamilyBlockPsiDerivative>>, String> {
2888        let mean_derivs =
2889            build_block_spatial_psi_derivatives(data, meanspec_resolved, mean_design)?
2890                .ok_or_else(|| "missing threshold spatial psi derivatives".to_string())?;
2891        let noise_derivs =
2892            build_block_spatial_psi_derivatives(data, noisespec_resolved, noise_design)?
2893                .ok_or_else(|| "missing log_sigma spatial psi derivatives".to_string())?;
2894        // The wiggle block has no direct spatial design matrix of its own in the
2895        // term builder. Spatial psi moves the wiggle family only through the
2896        // realized threshold/log-sigma designs, which in turn perturb q0 and the
2897        // realized wiggle basis B(q0). The exact joint wiggle psi hooks consume
2898        // those threshold/log-sigma derivative payloads and reconstruct the full
2899        // flattened likelihood-side [rho, psi] calculus internally, so the
2900        // wiggle block intentionally contributes no direct CustomFamilyBlockPsiDerivative
2901        // entries here.
2902        Ok(vec![mean_derivs, noise_derivs, Vec::new()])
2903    }
2904}
2905
2906pub(crate) fn fit_gaussian_location_scale_terms(
2907    data: ndarray::ArrayView2<'_, f64>,
2908    spec: GaussianLocationScaleTermSpec,
2909    options: &BlockwiseFitOptions,
2910    kappa_options: &SpatialLengthScaleOptimizationOptions,
2911) -> Result<BlockwiseTermFitResult, String> {
2912    validate_gaussian_location_scale_termspec(data, &spec, "fit_gaussian_location_scale_terms")?;
2913    fit_location_scale_terms(
2914        data,
2915        GaussianLocationScaleTermBuilder {
2916            y: spec.y,
2917            weights: spec.weights,
2918            meanspec: spec.meanspec,
2919            noisespec: spec.log_sigmaspec,
2920            mean_offset: spec.mean_offset,
2921            noise_offset: spec.log_sigma_offset,
2922        },
2923        options,
2924        kappa_options,
2925    )
2926}
2927
2928pub(crate) fn fit_gaussian_location_scalewiggle_terms(
2929    data: ndarray::ArrayView2<'_, f64>,
2930    spec: GaussianLocationScaleWiggleTermSpec,
2931    options: &BlockwiseFitOptions,
2932    kappa_options: &SpatialLengthScaleOptimizationOptions,
2933) -> Result<BlockwiseTermFitResult, String> {
2934    validate_gaussian_location_scalewiggle_termspec(
2935        data,
2936        &spec,
2937        "fit_gaussian_location_scalewiggle_terms",
2938    )?;
2939    fit_location_scale_terms(
2940        data,
2941        GaussianLocationScaleWiggleTermBuilder {
2942            y: spec.y,
2943            weights: spec.weights,
2944            meanspec: spec.meanspec,
2945            noisespec: spec.log_sigmaspec,
2946            mean_offset: spec.mean_offset,
2947            noise_offset: spec.log_sigma_offset,
2948            wiggle_knots: spec.wiggle_knots,
2949            wiggle_degree: spec.wiggle_degree,
2950            wiggle_block: spec.wiggle_block,
2951        },
2952        options,
2953        kappa_options,
2954    )
2955}
2956
2957pub(crate) fn select_gaussian_location_scale_link_wiggle_basis_from_pilot(
2958    pilot: &BlockwiseTermFitResult,
2959    wiggle_cfg: &WiggleBlockConfig,
2960    wiggle_penalty_orders: &[usize],
2961) -> Result<SelectedWiggleBasis, String> {
2962    let q_seed = pilot
2963        .fit
2964        .block_states
2965        .first()
2966        .ok_or_else(|| "pilot Gaussian wiggle fit is missing mean block".to_string())?
2967        .eta
2968        .view();
2969    select_wiggle_basis_from_seed(q_seed, wiggle_cfg, wiggle_penalty_orders)
2970}
2971
2972pub(crate) fn fit_gaussian_location_scale_terms_with_selected_wiggle(
2973    data: ndarray::ArrayView2<'_, f64>,
2974    spec: GaussianLocationScaleTermSpec,
2975    selected_wiggle_basis: SelectedWiggleBasis,
2976    options: &BlockwiseFitOptions,
2977    kappa_options: &SpatialLengthScaleOptimizationOptions,
2978) -> Result<BlockwiseTermWiggleFitResult, String> {
2979    let SelectedWiggleBasis {
2980        knots: wiggle_knots,
2981        degree: wiggle_degree,
2982        block: wiggle_block,
2983        ..
2984    } = selected_wiggle_basis;
2985    let solved = fit_gaussian_location_scalewiggle_terms(
2986        data,
2987        GaussianLocationScaleWiggleTermSpec {
2988            y: spec.y,
2989            weights: spec.weights,
2990            meanspec: spec.meanspec,
2991            log_sigmaspec: spec.log_sigmaspec,
2992            mean_offset: spec.mean_offset,
2993            log_sigma_offset: spec.log_sigma_offset,
2994            wiggle_knots: wiggle_knots.clone(),
2995            wiggle_degree,
2996            wiggle_block,
2997        },
2998        options,
2999        kappa_options,
3000    )?;
3001
3002    BlockwiseTermWiggleFitResult::try_from_parts(BlockwiseTermWiggleFitResultParts {
3003        fit: solved,
3004        wiggle_knots,
3005        wiggle_degree,
3006    })
3007}
3008
3009pub(crate) fn fit_binomial_location_scale_terms(
3010    data: ndarray::ArrayView2<'_, f64>,
3011    spec: BinomialLocationScaleTermSpec,
3012    options: &BlockwiseFitOptions,
3013    kappa_options: &SpatialLengthScaleOptimizationOptions,
3014) -> Result<BlockwiseTermFitResult, String> {
3015    validate_binomial_location_scale_termspec(data, &spec, "fit_binomial_location_scale_terms")?;
3016    fit_location_scale_terms(
3017        data,
3018        BinomialLocationScaleTermBuilder {
3019            y: spec.y,
3020            weights: spec.weights,
3021            link_kind: spec.link_kind,
3022            meanspec: spec.thresholdspec,
3023            noisespec: spec.log_sigmaspec,
3024            mean_offset: spec.threshold_offset,
3025            noise_offset: spec.log_sigma_offset,
3026        },
3027        options,
3028        kappa_options,
3029    )
3030}
3031
3032pub(crate) fn fit_binomial_location_scalewiggle_terms(
3033    data: ndarray::ArrayView2<'_, f64>,
3034    spec: BinomialLocationScaleWiggleTermSpec,
3035    options: &BlockwiseFitOptions,
3036    kappa_options: &SpatialLengthScaleOptimizationOptions,
3037) -> Result<BlockwiseTermFitResult, String> {
3038    validate_binomial_location_scalewiggle_termspec(
3039        data,
3040        &spec,
3041        "fit_binomial_location_scalewiggle_terms",
3042    )?;
3043    fit_location_scale_terms(
3044        data,
3045        BinomialLocationScaleWiggleTermBuilder {
3046            y: spec.y,
3047            weights: spec.weights,
3048            link_kind: spec.link_kind,
3049            meanspec: spec.thresholdspec,
3050            noisespec: spec.log_sigmaspec,
3051            mean_offset: spec.threshold_offset,
3052            noise_offset: spec.log_sigma_offset,
3053            wiggle_knots: spec.wiggle_knots,
3054            wiggle_degree: spec.wiggle_degree,
3055            wiggle_block: spec.wiggle_block,
3056        },
3057        options,
3058        kappa_options,
3059    )
3060}
3061
3062pub(crate) fn select_binomial_location_scale_link_wiggle_basis_from_pilot(
3063    pilot: &BlockwiseTermFitResult,
3064    wiggle_cfg: &WiggleBlockConfig,
3065    wiggle_penalty_orders: &[usize],
3066) -> Result<SelectedWiggleBasis, String> {
3067    let eta_t = pilot
3068        .fit
3069        .block_states
3070        .first()
3071        .ok_or_else(|| "pilot fit is missing threshold block".to_string())?
3072        .eta
3073        .view();
3074    let eta_ls = pilot
3075        .fit
3076        .block_states
3077        .get(1)
3078        .ok_or_else(|| "pilot fit is missing log_sigma block".to_string())?
3079        .eta
3080        .view();
3081    let sigma = eta_ls.mapv(safe_exp);
3082    let q_seed = Array1::from_iter(eta_t.iter().zip(sigma.iter()).map(|(&t, &s)| -t / s));
3083    select_wiggle_basis_from_seed(q_seed.view(), wiggle_cfg, wiggle_penalty_orders)
3084}
3085
3086pub(crate) fn fit_binomial_location_scale_terms_with_selected_wiggle(
3087    data: ndarray::ArrayView2<'_, f64>,
3088    spec: BinomialLocationScaleTermSpec,
3089    selected_wiggle_basis: SelectedWiggleBasis,
3090    options: &BlockwiseFitOptions,
3091    kappa_options: &SpatialLengthScaleOptimizationOptions,
3092) -> Result<BlockwiseTermWiggleFitResult, String> {
3093    let SelectedWiggleBasis {
3094        knots: wiggle_knots,
3095        degree: wiggle_degree,
3096        block: wiggle_block,
3097        ..
3098    } = selected_wiggle_basis;
3099    let solved = fit_binomial_location_scalewiggle_terms(
3100        data,
3101        BinomialLocationScaleWiggleTermSpec {
3102            y: spec.y,
3103            weights: spec.weights,
3104            link_kind: spec.link_kind,
3105            thresholdspec: spec.thresholdspec,
3106            log_sigmaspec: spec.log_sigmaspec,
3107            threshold_offset: spec.threshold_offset,
3108            log_sigma_offset: spec.log_sigma_offset,
3109            wiggle_knots: wiggle_knots.clone(),
3110            wiggle_degree,
3111            wiggle_block,
3112        },
3113        options,
3114        kappa_options,
3115    )?;
3116
3117    BlockwiseTermWiggleFitResult::try_from_parts(BlockwiseTermWiggleFitResultParts {
3118        fit: solved,
3119        wiggle_knots,
3120        wiggle_degree,
3121    })
3122}
3123
3124pub(crate) fn select_binomial_mean_link_wiggle_basis_from_pilot(
3125    pilot_design: &TermCollectionDesign,
3126    pilot_fit: &UnifiedFitResult,
3127    wiggle_cfg: &WiggleBlockConfig,
3128    wiggle_penalty_orders: &[usize],
3129) -> Result<SelectedWiggleBasis, String> {
3130    let q_seed = pilot_design
3131        .apply(pilot_fit.beta.view())
3132        .map_err(|error| error.to_string())?;
3133    select_wiggle_basis_from_seed(q_seed.view(), wiggle_cfg, wiggle_penalty_orders)
3134}
3135
3136pub(crate) fn fit_binomial_mean_wiggle_terms_with_selected_basis(
3137    data: ndarray::ArrayView2<'_, f64>,
3138    pilot_spec: &TermCollectionSpec,
3139    pilot_design: &TermCollectionDesign,
3140    pilot_fit: &UnifiedFitResult,
3141    y: &Array1<f64>,
3142    weights: &Array1<f64>,
3143    link_kind: InverseLink,
3144    selected_wiggle_basis: SelectedWiggleBasis,
3145    options: &BlockwiseFitOptions,
3146    kappa_options: &SpatialLengthScaleOptimizationOptions,
3147) -> Result<BinomialMeanWiggleTermFitResult, String> {
3148    const RHO_BOUND: f64 = 12.0;
3149
3150    validate_term_weights(
3151        data,
3152        y.len(),
3153        weights,
3154        "fit_binomial_mean_wiggle_terms_with_selected_basis",
3155    )?;
3156    validate_binomial_response(y, "fit_binomial_mean_wiggle_terms_with_selected_basis")?;
3157
3158    // Large-n binomial mean-wiggle fits keep the caller's explicit Hessian
3159    // request. The unified evaluator chooses the scalable exact representation
3160    // (dense for small work, operator HVP for large work) instead of routing to
3161    // gradient-only BFGS by observation count.
3162
3163    let SelectedWiggleBasis {
3164        knots: wiggle_knots,
3165        degree: wiggle_degree,
3166        block: wiggle_block,
3167        ..
3168    } = selected_wiggle_basis;
3169
3170    let spatial_terms = spatial_length_scale_term_indices(pilot_spec);
3171    if spatial_terms.is_empty() {
3172        let (fit, saved_warp_beta, saved_index_shift) = fit_binomial_mean_wiggle(
3173            BinomialMeanWiggleSpec {
3174                y: y.clone(),
3175                weights: weights.clone(),
3176                link_kind,
3177                wiggle_knots: wiggle_knots.clone(),
3178                wiggle_degree,
3179                eta_block: ParameterBlockInput {
3180                    design: pilot_design.design.clone(),
3181                    offset: pilot_design.affine_offset.clone(),
3182                    penalties: pilot_design
3183                        .penalties
3184                        .iter()
3185                        .map(crate::model_types::PenaltySpec::from_blockwise_ref)
3186                        .collect(),
3187                    nullspace_dims: vec![],
3188                    initial_log_lambdas: Some(fitted_log_lambdas(
3189                        &pilot_fit.lambdas,
3190                        "binomial mean-wiggle pilot lambda",
3191                    )?),
3192                    initial_beta: Some(pilot_fit.beta.clone()),
3193                },
3194                wiggle_block,
3195            },
3196            options,
3197        )?;
3198        return Ok(BinomialMeanWiggleTermFitResult {
3199            fit,
3200            resolvedspec: pilot_spec.clone(),
3201            design: pilot_design.clone(),
3202            wiggle_knots,
3203            wiggle_degree,
3204            saved_warp_beta,
3205            saved_index_shift,
3206        });
3207    }
3208
3209    let dims_per_term = spatial_dims_per_term(pilot_spec, &spatial_terms);
3210    let log_kappa0 =
3211        SpatialLogKappaCoords::from_length_scales_aniso(pilot_spec, &spatial_terms, kappa_options)
3212            .reseed_from_data(data, pilot_spec, &spatial_terms, kappa_options)
3213            .map_err(|error| error.to_string())?;
3214    let log_kappa_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
3215        data,
3216        pilot_spec,
3217        &spatial_terms,
3218        &dims_per_term,
3219        kappa_options,
3220    )
3221    .map_err(|error| error.to_string())?;
3222    let log_kappa_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
3223        data,
3224        pilot_spec,
3225        &spatial_terms,
3226        &dims_per_term,
3227        kappa_options,
3228    )
3229    .map_err(|error| error.to_string())?;
3230    // Project seed onto bounds; spec.length_scale is a hint, not a constraint.
3231    let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
3232
3233    let eta_penalty_count = pilot_design.penalties.len();
3234    let wiggle_penalty_count = initial_log_lambdas_orzeros(&wiggle_block)?.len();
3235    let rho_dim = eta_penalty_count + wiggle_penalty_count;
3236    let baseline_resolvedspec = log_kappa0
3237        .apply_tospec(pilot_spec, &spatial_terms)
3238        .map_err(|e| e.to_string())?;
3239    let baseline_design =
3240        build_term_collection_design(data, &baseline_resolvedspec).map_err(|e| e.to_string())?;
3241    let baseline_fit = fit_binomial_mean_wiggle(
3242        BinomialMeanWiggleSpec {
3243            y: y.clone(),
3244            weights: weights.clone(),
3245            link_kind: link_kind.clone(),
3246            wiggle_knots: wiggle_knots.clone(),
3247            wiggle_degree,
3248            eta_block: ParameterBlockInput {
3249                design: baseline_design.design.clone(),
3250                offset: baseline_design.affine_offset.clone(),
3251                penalties: baseline_design
3252                    .penalties
3253                    .iter()
3254                    .map(crate::model_types::PenaltySpec::from_blockwise_ref)
3255                    .collect(),
3256                nullspace_dims: vec![],
3257                initial_log_lambdas: Some(fitted_log_lambdas(
3258                    &pilot_fit.lambdas,
3259                    "binomial mean-wiggle pilot lambda",
3260                )?),
3261                initial_beta: Some(pilot_fit.beta.clone()),
3262            },
3263            wiggle_block: wiggle_block.clone(),
3264        },
3265        options,
3266    )?
3267    .0;
3268    let baseline_log_lambdas = fitted_log_lambdas(
3269        &baseline_fit.lambdas,
3270        "binomial mean-wiggle baseline lambda",
3271    )?;
3272    if baseline_log_lambdas.len() != rho_dim {
3273        return Err(GamlssError::DimensionMismatch {
3274            reason: format!(
3275                "baseline binomial mean-wiggle fit returned {} log-lambdas, expected {rho_dim}",
3276                baseline_log_lambdas.len()
3277            ),
3278        }
3279        .into());
3280    }
3281    let baseline_eta_beta = baseline_fit
3282        .block_states
3283        .get(BinomialMeanWiggleFamily::BLOCK_ETA)
3284        .ok_or_else(|| "baseline binomial mean-wiggle fit missing eta block".to_string())?
3285        .beta
3286        .clone();
3287    let baseline_wiggle_beta = Some(
3288        baseline_fit
3289            .block_states
3290            .get(BinomialMeanWiggleFamily::BLOCK_WIGGLE)
3291            .ok_or_else(|| "baseline binomial mean-wiggle fit missing wiggle block".to_string())?
3292            .beta
3293            .clone(),
3294    );
3295    let theta_dim = rho_dim + log_kappa0.len();
3296    let mut theta0 = Array1::<f64>::zeros(theta_dim);
3297    theta0
3298        .slice_mut(s![0..rho_dim])
3299        .assign(&baseline_log_lambdas);
3300    theta0
3301        .slice_mut(s![rho_dim..theta_dim])
3302        .assign(log_kappa0.as_array());
3303
3304    let mut lower = Array1::<f64>::from_elem(theta_dim, -RHO_BOUND);
3305    let mut upper = Array1::<f64>::from_elem(theta_dim, RHO_BOUND);
3306    lower
3307        .slice_mut(s![rho_dim..theta_dim])
3308        .assign(log_kappa_lower.as_array());
3309    upper
3310        .slice_mut(s![rho_dim..theta_dim])
3311        .assign(log_kappa_upper.as_array());
3312
3313    let pilot_spec_cloned = pilot_spec.clone();
3314    let pilot_beta = baseline_eta_beta;
3315    let wiggle_design = wiggle_block.design.clone();
3316    let wiggle_offset = wiggle_block.offset.clone();
3317    let wiggle_penalties = wiggle_block.penalties.clone();
3318    let wiggle_initial_beta = baseline_wiggle_beta;
3319    let wiggle_knots_cloned = wiggle_knots.clone();
3320    let y_cloned = y.clone();
3321    let weights_cloned = weights.clone();
3322    let link_kind_cloned = link_kind.clone();
3323    let outer_family = BinomialMeanWiggleFamily {
3324        y: y_cloned.clone(),
3325        weights: weights_cloned.clone(),
3326        link_kind: link_kind_cloned.clone(),
3327        wiggle_knots: wiggle_knots_cloned.clone(),
3328        wiggle_degree,
3329        policy: gam_runtime::resource::ResourcePolicy::default_library(),
3330        // The spatial joint-κ path keeps the dynamic warp basis (#1596 frozen
3331        // basis applies to the non-spatial `fit_binomial_mean_wiggle` loop).
3332        frozen_warp_design: None,
3333    };
3334    let screening_cap = Arc::new(AtomicUsize::new(0));
3335    let mut outer_options = options.clone();
3336    outer_options.screening_max_inner_iterations = Some(Arc::clone(&screening_cap));
3337    struct MeanWiggleOuterState {
3338        pub(crate) warm_cache: Option<crate::custom_family::CustomFamilyWarmStart>,
3339        pub(crate) last_eval: Option<(
3340            Array1<f64>,
3341            f64,
3342            Array1<f64>,
3343            gam_problem::HessianValue,
3344            crate::custom_family::CustomFamilyWarmStart,
3345        )>,
3346    }
3347
3348    let build_realized_blocks = |theta: &Array1<f64>| -> Result<
3349        (
3350            TermCollectionSpec,
3351            TermCollectionDesign,
3352            Vec<ParameterBlockSpec>,
3353            Vec<CustomFamilyBlockPsiDerivative>,
3354        ),
3355        String,
3356    > {
3357        let log_kappa =
3358            SpatialLogKappaCoords::from_theta_tail_with_dims(theta, rho_dim, dims_per_term.clone());
3359        let resolvedspec = log_kappa
3360            .apply_tospec(&pilot_spec_cloned, &spatial_terms)
3361            .map_err(|e| e.to_string())?;
3362        let design =
3363            build_term_collection_design(data, &resolvedspec).map_err(|e| e.to_string())?;
3364        let eta_derivs = build_block_spatial_psi_derivatives(data, &resolvedspec, &design)?
3365            .ok_or_else(|| {
3366                "missing eta spatial psi derivatives for binomial mean wiggle".to_string()
3367            })?;
3368        let blocks = vec![
3369            ParameterBlockSpec {
3370                name: "eta".to_string(),
3371                design: design.design.clone(),
3372                offset: design.affine_offset.clone(),
3373                penalties: design.penalties_as_penalty_matrix(),
3374                nullspace_dims: vec![],
3375                initial_log_lambdas: theta.slice(s![0..eta_penalty_count]).to_owned(),
3376                initial_beta: Some(pilot_beta.clone()),
3377                // Lower gauge priority on the static eta design: it yields the
3378                // shared level/intercept direction to the dynamic full-width
3379                // wiggle I-spline block (see fit_binomial_mean_wiggle).
3380                gauge_priority: LINK_WIGGLE_GAUGE_PRIORITY,
3381                jacobian_callback: None,
3382                stacked_design: None,
3383                stacked_offset: None,
3384            },
3385            ParameterBlockSpec {
3386                name: "wiggle".to_string(),
3387                design: wiggle_design.clone(),
3388                offset: wiggle_offset.clone(),
3389                penalties: {
3390                    let p_wiggle = wiggle_design.ncols();
3391                    wiggle_penalties
3392                        .iter()
3393                        .map(|spec| match spec {
3394                            crate::model_types::PenaltySpec::Block {
3395                                local, col_range, ..
3396                            } => PenaltyMatrix::Blockwise {
3397                                local: local.clone(),
3398                                col_range: col_range.clone(),
3399                                total_dim: p_wiggle,
3400                            },
3401                            crate::model_types::PenaltySpec::Dense(m)
3402                            | crate::model_types::PenaltySpec::DenseWithMean {
3403                                matrix: m, ..
3404                            } => PenaltyMatrix::Dense(m.clone()),
3405                        })
3406                        .collect()
3407                },
3408                nullspace_dims: vec![],
3409                initial_log_lambdas: theta.slice(s![eta_penalty_count..rho_dim]).to_owned(),
3410                initial_beta: wiggle_initial_beta.clone(),
3411                gauge_priority: DEFAULT_GAUGE_PRIORITY,
3412                jacobian_callback: None,
3413                stacked_design: None,
3414                stacked_offset: None,
3415            },
3416        ];
3417        Ok((resolvedspec, design, blocks, eta_derivs))
3418    };
3419
3420    let build_eval = |theta: &Array1<f64>,
3421                      warm_cache: Option<&crate::custom_family::CustomFamilyWarmStart>,
3422                      need_hessian: bool|
3423     -> Result<
3424        (
3425            crate::custom_family::CustomFamilyJointHyperResult,
3426            TermCollectionSpec,
3427            TermCollectionDesign,
3428        ),
3429        String,
3430    > {
3431        let (resolvedspec, design, blocks, eta_derivs) = build_realized_blocks(theta)?;
3432        let hyper_layout = crate::custom_family::CustomFamilyHyperLayout::new(
3433            vec![eta_derivs, Vec::new()],
3434            Vec::new(),
3435            theta.slice(s![rho_dim..]).to_owned(),
3436        )?;
3437        let eval = evaluate_custom_family_joint_hyper(
3438            &outer_family,
3439            &blocks,
3440            &outer_options,
3441            &theta.slice(s![0..rho_dim]).to_owned(),
3442            &hyper_layout,
3443            warm_cache,
3444            if need_hessian {
3445                gam_problem::EvalMode::ValueGradientHessian
3446            } else {
3447                gam_problem::EvalMode::ValueAndGradient
3448            },
3449        )?;
3450        Ok((eval, resolvedspec, design))
3451    };
3452
3453    let build_efs = |theta: &Array1<f64>,
3454                     warm_cache: Option<&crate::custom_family::CustomFamilyWarmStart>|
3455     -> Result<crate::custom_family::CustomFamilyJointHyperEfsResult, String> {
3456        let (_, _, blocks, eta_derivs) = build_realized_blocks(theta)?;
3457        let hyper_layout = crate::custom_family::CustomFamilyHyperLayout::new(
3458            vec![eta_derivs, Vec::new()],
3459            Vec::new(),
3460            theta.slice(s![rho_dim..]).to_owned(),
3461        )?;
3462        evaluate_custom_family_joint_hyper_efs(
3463            &outer_family,
3464            &blocks,
3465            &outer_options,
3466            &theta.slice(s![0..rho_dim]).to_owned(),
3467            &hyper_layout,
3468            warm_cache,
3469        )
3470        .map_err(|e| e.to_string())
3471    };
3472
3473    use crate::model_types::EstimationError;
3474    use gam_problem::{DeclaredHessianForm, Derivative, OuterEval};
3475    use gam_solve::rho_optimizer::OuterEvalOrder;
3476
3477    // Exact first-order AND second-order [rho, psi] calculus is available
3478    // for all inverse links via the shared jet formulas plus the generic
3479    // exact-Newton D_βH / D²_βH closures routed through
3480    // evaluate_custom_family_joint_hyper -> joint_outer_evaluate ->
3481    // BorrowedJointDerivProvider. This enables the analytic-Hessian outer
3482    // plan for REML optimization instead of the downgraded gradient-only
3483    // outer strategies.
3484    //
3485    // Spatial log-kappa coordinates are ψ (design-moving) dimensions because
3486    // they rebuild the spatial basis and penalties at each outer proposal.
3487    let analytic_outer_hessian_available = true;
3488    let mut seed_heuristic = theta0.to_vec();
3489    for value in &mut seed_heuristic[..rho_dim] {
3490        *value = value.exp();
3491    }
3492    let problem = gam_solve::rho_optimizer::OuterProblem::new(theta_dim)
3493        .with_gradient(Derivative::Analytic)
3494        .with_hessian(if analytic_outer_hessian_available {
3495            DeclaredHessianForm::Either
3496        } else {
3497            DeclaredHessianForm::Unavailable
3498        })
3499        .with_psi_dim(theta_dim - rho_dim)
3500        .with_tolerance(options.outer_tol)
3501        .with_max_iter(options.outer_max_iter)
3502        .with_bounds(lower.clone(), upper.clone())
3503        .with_initial_rho(theta0.clone())
3504        .with_seed_config(crate::seeding::SeedConfig {
3505            max_seeds: 4,
3506            seed_budget: 2,
3507            risk_profile: crate::seeding::SeedRiskProfile::GeneralizedLinear,
3508            num_auxiliary_trailing: theta_dim - rho_dim,
3509            ..Default::default()
3510        })
3511        .with_screening_cap(Arc::clone(&screening_cap))
3512        .with_rho_bound(12.0)
3513        .with_heuristic_lambdas(seed_heuristic);
3514
3515    let eval_outer = |state: &mut MeanWiggleOuterState,
3516                      theta: &Array1<f64>,
3517                      order: OuterEvalOrder|
3518     -> Result<OuterEval, EstimationError> {
3519        if let Some((cached_theta, cached_cost, cached_grad, cached_hess, cached_warm)) =
3520            &state.last_eval
3521            && cached_theta == theta
3522            && (!matches!(order, OuterEvalOrder::ValueGradientHessian)
3523                || matches!(
3524                    cached_hess,
3525                    gam_problem::HessianValue::Dense(_) | gam_problem::HessianValue::Operator(_)
3526                ))
3527        {
3528            state.warm_cache = Some(cached_warm.clone());
3529            return Ok(OuterEval {
3530                cost: *cached_cost,
3531                gradient: cached_grad.clone(),
3532                hessian: cached_hess.clone(),
3533                inner_beta_hint: None,
3534            });
3535        }
3536        let need_hessian = matches!(order, OuterEvalOrder::ValueGradientHessian)
3537            && analytic_outer_hessian_available;
3538        let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), need_hessian)
3539            .map_err(EstimationError::InvalidInput)?;
3540        if !eval.inner_converged {
3541            state.warm_cache = Some(eval.warm_start);
3542            crate::bail_invalid_estim!(
3543                "binomial mean-wiggle exact spatial inner solve did not converge"
3544            );
3545        }
3546        let hessian_result = eval.outer_hessian.clone();
3547        state.last_eval = Some((
3548            theta.clone(),
3549            eval.objective,
3550            eval.gradient.clone(),
3551            eval.outer_hessian.clone(),
3552            eval.warm_start.clone(),
3553        ));
3554        state.warm_cache = Some(eval.warm_start);
3555        Ok(OuterEval {
3556            cost: eval.objective,
3557            gradient: eval.gradient,
3558            hessian: hessian_result,
3559            inner_beta_hint: None,
3560        })
3561    };
3562
3563    let mut obj = problem.build_objective_with_screening_proxy(
3564        MeanWiggleOuterState {
3565            warm_cache: None,
3566            last_eval: None,
3567        },
3568        |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3569            if let Some((cached_theta, cached_cost, _, _, cached_warm)) = &state.last_eval
3570                && cached_theta == theta
3571            {
3572                state.warm_cache = Some(cached_warm.clone());
3573                return Ok(*cached_cost);
3574            }
3575            let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), false)
3576                .map_err(EstimationError::InvalidInput)?;
3577            if !eval.inner_converged {
3578                state.warm_cache = Some(eval.warm_start);
3579                crate::bail_invalid_estim!(
3580                    "binomial mean-wiggle exact spatial cost inner solve did not converge"
3581                        .to_string(),
3582                );
3583            }
3584            state.warm_cache = Some(eval.warm_start);
3585            Ok(eval.objective)
3586        },
3587        |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3588            eval_outer(
3589                state,
3590                theta,
3591                if analytic_outer_hessian_available {
3592                    OuterEvalOrder::ValueGradientHessian
3593                } else {
3594                    OuterEvalOrder::ValueAndGradient
3595                },
3596            )
3597        },
3598        |state: &mut MeanWiggleOuterState, theta: &Array1<f64>, order: OuterEvalOrder| {
3599            eval_outer(state, theta, order)
3600        },
3601        Some(|state: &mut MeanWiggleOuterState| {
3602            state.warm_cache = None;
3603            state.last_eval = None;
3604        }),
3605        Some(|state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3606            let eval = build_efs(theta, state.warm_cache.as_ref())
3607                .map_err(EstimationError::InvalidInput)?;
3608            if !eval.inner_converged {
3609                state.warm_cache = Some(eval.warm_start);
3610                crate::bail_invalid_estim!(
3611                    "binomial mean-wiggle exact spatial EFS inner solve did not converge"
3612                        .to_string(),
3613                );
3614            }
3615            state.warm_cache = Some(eval.warm_start);
3616            Ok(eval.efs_eval)
3617        }),
3618        // Seed-screening ranking proxy (#969). The cost closure above
3619        // hard-errors on a non-converged inner solve — correct for
3620        // line-search costs, but under the screening cap (wired into the
3621        // outer options and installed by the cascade) the inner solve is
3622        // truncated BY DESIGN, so screening through it rejects every seed
3623        // — the all-seeds-rejected front-door genus. Screening only RANKS
3624        // candidates: the truncated solve's penalized objective is the
3625        // ranking signal; convergence is demanded of the selected seed's
3626        // full-budget fit, not of capped probes.
3627        |state: &mut MeanWiggleOuterState, theta: &Array1<f64>| {
3628            if let Some((cached_theta, cached_cost, _, _, cached_warm)) = &state.last_eval
3629                && cached_theta == theta
3630            {
3631                state.warm_cache = Some(cached_warm.clone());
3632                return Ok(*cached_cost);
3633            }
3634            let (eval, _, _) = build_eval(theta, state.warm_cache.as_ref(), false)
3635                .map_err(EstimationError::InvalidInput)?;
3636            state.warm_cache = Some(eval.warm_start);
3637            Ok(eval.objective)
3638        },
3639    );
3640
3641    let outer = problem
3642        .run(&mut obj, "binomial mean wiggle exact spatial hyper")
3643        .map_err(|e| e.to_string())?;
3644    if !outer.converged {
3645        return Err(GamlssError::NumericalFailure { reason: format!(
3646            "binomial mean wiggle exact spatial hyper did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
3647            outer.iterations,
3648            outer.final_value,
3649            outer.final_grad_norm_report(),
3650        ) }.into());
3651    }
3652    let theta_star = outer.rho;
3653
3654    let log_kappa =
3655        SpatialLogKappaCoords::from_theta_tail_with_dims(&theta_star, rho_dim, dims_per_term);
3656    let resolvedspec = log_kappa
3657        .apply_tospec(&pilot_spec_cloned, &spatial_terms)
3658        .map_err(|e| e.to_string())?;
3659    let design = build_term_collection_design(data, &resolvedspec).map_err(|e| e.to_string())?;
3660    let resolvedspec =
3661        freeze_term_collection_from_design(&resolvedspec, &design).map_err(|e| e.to_string())?;
3662    let fit = fit_binomial_mean_wiggle(
3663        BinomialMeanWiggleSpec {
3664            y: y_cloned,
3665            weights: weights_cloned,
3666            link_kind: link_kind_cloned,
3667            wiggle_knots: wiggle_knots.clone(),
3668            wiggle_degree,
3669            eta_block: ParameterBlockInput {
3670                design: design.design.clone(),
3671                offset: design.affine_offset.clone(),
3672                penalties: design
3673                    .penalties
3674                    .iter()
3675                    .map(crate::model_types::PenaltySpec::from_blockwise_ref)
3676                    .collect(),
3677                nullspace_dims: vec![],
3678                initial_log_lambdas: Some(theta_star.slice(s![0..eta_penalty_count]).to_owned()),
3679                initial_beta: Some(pilot_beta),
3680            },
3681            wiggle_block: ParameterBlockInput {
3682                design: wiggle_design,
3683                offset: wiggle_offset,
3684                penalties: wiggle_penalties,
3685                nullspace_dims: vec![],
3686                initial_log_lambdas: Some(
3687                    theta_star.slice(s![eta_penalty_count..rho_dim]).to_owned(),
3688                ),
3689                initial_beta: wiggle_initial_beta,
3690            },
3691        },
3692        options,
3693    )?;
3694    let (fit, saved_warp_beta, saved_index_shift) = fit;
3695
3696    Ok(BinomialMeanWiggleTermFitResult {
3697        fit,
3698        resolvedspec,
3699        design,
3700        wiggle_knots,
3701        wiggle_degree,
3702        saved_warp_beta,
3703        saved_index_shift,
3704    })
3705}