Skip to main content

gam_models/gamlss/
builders.rs

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