Skip to main content

gam_models/gamlss/gaussian/
binomial_mean_wiggle.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)]
8pub struct BinomialMeanWiggleFamily {
9    pub y: Array1<f64>,
10    pub weights: Array1<f64>,
11    pub link_kind: InverseLink,
12    pub wiggle_knots: Array1<f64>,
13    pub wiggle_degree: usize,
14    /// Resource policy threaded into PsiDesignMap construction during
15    /// exact-Newton joint psi evaluation. Defaults to
16    /// `ResourcePolicy::default_library()` when the family is built without
17    /// an explicit policy.
18    pub policy: gam_runtime::resource::ResourcePolicy,
19    /// The **frozen, identifiable** warp design
20    /// `B⊥ = (I - P_X) B(η̂)` used for the duration of one inner joint-Newton
21    /// solve (#1596). `None` preserves the original dynamic-basis behaviour;
22    /// `Some(B⊥)` switches the family into the
23    /// frozen-basis Gauss-Newton mode.
24    ///
25    /// **Why frozen.** The fully-coupled `q = η + B(η)·β_w` model regenerates the
26    /// monotone I-spline basis at the *moving* `η` every cycle. The trust-region
27    /// quadratic model freezes `B` at the cycle-start `η`, but the line search
28    /// re-evaluates the objective with `B` rebuilt at the trial `η`; for any step
29    /// that moves `η` the actual reduction diverges from the model, the trust
30    /// radius collapses, and the constrained KKT certificate refuses every
31    /// iterate (`active_set_incomplete`) even when the optimal warp is flat.
32    /// Freezing `B(η̂)` makes `q = η + B⊥·β_w` linear in `(β_η, β_w)` with
33    /// `∂q/∂η = 1` and no `∂B/∂η` chain term — a well-conditioned two-block GLM
34    /// that certifies. The caller re-freezes at the refit `η̂` and returns only
35    /// after the caller-supplied outer convergence policy certifies the resulting
36    /// Gauss-Newton fixed point (`fit_binomial_mean_wiggle`).
37    ///
38    /// **Why observation-space residualized (identifiable).** A monotone
39    /// I-spline of the linear predictor `η` can represent mean-block directions,
40    /// so the raw `B(η̂)` columns alias `X`. The caller fits
41    /// `B⊥ = B - X A`, with `A = (XᵀX)^+XᵀB`, removing only that aliased
42    /// observation-space component while preserving the standard I-spline
43    /// coefficient coordinate. Consequently the exact structural constraint is
44    /// still simply `β_w ≥ 0`; prediction reconstructs `B(η_new)·β_w` and
45    /// compensates the saved mean coefficient by `-Aβ_w`.
46    pub frozen_warp_design: Option<Arc<Array2<f64>>>,
47}
48
49pub(crate) struct BinomialMeanWiggleGeometry {
50    pub(crate) basis: Array2<f64>,
51    pub(crate) basis_d1: Array2<f64>,
52    pub(crate) basis_d2: Array2<f64>,
53    pub(crate) basis_d3: Array2<f64>,
54    pub(crate) dq_dq0: Array1<f64>,
55    pub(crate) d2q_dq02: Array1<f64>,
56    pub(crate) d3q_dq03: Array1<f64>,
57    pub(crate) d4q_dq04: Array1<f64>,
58}
59
60pub(crate) struct BinomialMeanWiggleJointPsiDirection {
61    pub(crate) x_eta_psi: Option<Array2<f64>>,
62    pub(crate) z_eta_psi: Array1<f64>,
63}
64
65impl BinomialMeanWiggleFamily {
66    pub const BLOCK_ETA: usize = 0;
67    pub const BLOCK_WIGGLE: usize = 1;
68
69    pub(crate) fn wiggle_basiswith_options(
70        &self,
71        q0: ArrayView1<'_, f64>,
72        options: BasisOptions,
73    ) -> Result<Array2<f64>, String> {
74        monotone_wiggle_basis_with_derivative_order(
75            q0,
76            &self.wiggle_knots,
77            self.wiggle_degree,
78            options.derivative_order,
79        )
80    }
81
82    pub(crate) fn wiggle_design(&self, q0: ArrayView1<'_, f64>) -> Result<Array2<f64>, String> {
83        self.wiggle_basiswith_options(q0, BasisOptions::value())
84    }
85
86    pub(crate) fn wiggle_dq_dq0(
87        &self,
88        q0: ArrayView1<'_, f64>,
89        beta_link_wiggle: ArrayView1<'_, f64>,
90    ) -> Result<Array1<f64>, String> {
91        let d_constrained = self.wiggle_basiswith_options(q0, BasisOptions::first_derivative())?;
92        if d_constrained.ncols() != beta_link_wiggle.len() {
93            return Err(GamlssError::DimensionMismatch { reason: format!(
94                "wiggle derivative/beta mismatch: basis has {} columns but beta_link_wiggle has {} coefficients",
95                d_constrained.ncols(),
96                beta_link_wiggle.len()
97            ) }.into());
98        }
99        Ok(d_constrained.dot(&beta_link_wiggle) + 1.0)
100    }
101
102    pub(crate) fn wiggle_d2q_dq02(
103        &self,
104        q0: ArrayView1<'_, f64>,
105        beta_link_wiggle: ArrayView1<'_, f64>,
106    ) -> Result<Array1<f64>, String> {
107        let d2 = self.wiggle_basiswith_options(q0, BasisOptions::second_derivative())?;
108        if d2.ncols() != beta_link_wiggle.len() {
109            return Err(GamlssError::DimensionMismatch { reason: format!(
110                "wiggle second-derivative/beta mismatch: basis has {} columns but beta_link_wiggle has {} coefficients",
111                d2.ncols(),
112                beta_link_wiggle.len()
113            ) }.into());
114        }
115        Ok(d2.dot(&beta_link_wiggle))
116    }
117
118    pub(crate) fn wiggle_d3basis_constrained(
119        &self,
120        q0: ArrayView1<'_, f64>,
121    ) -> Result<Array2<f64>, String> {
122        monotone_wiggle_basis_with_derivative_order(q0, &self.wiggle_knots, self.wiggle_degree, 3)
123    }
124
125    pub(crate) fn wiggle_d3q_dq03(
126        &self,
127        q0: ArrayView1<'_, f64>,
128        beta_link_wiggle: ArrayView1<'_, f64>,
129    ) -> Result<Array1<f64>, String> {
130        let d3 = self.wiggle_d3basis_constrained(q0)?;
131        if d3.ncols() != beta_link_wiggle.len() {
132            return Err(GamlssError::DimensionMismatch { reason: format!(
133                "wiggle third-derivative/beta mismatch: basis has {} columns but beta_link_wiggle has {} coefficients",
134                d3.ncols(),
135                beta_link_wiggle.len()
136            ) }.into());
137        }
138        Ok(d3.dot(&beta_link_wiggle))
139    }
140
141    pub(crate) fn wiggle_d4q_dq04(
142        &self,
143        q0: ArrayView1<'_, f64>,
144        beta_link_wiggle: ArrayView1<'_, f64>,
145    ) -> Result<Array1<f64>, String> {
146        let d4 = monotone_wiggle_basis_with_derivative_order(
147            q0,
148            &self.wiggle_knots,
149            self.wiggle_degree,
150            4,
151        )?;
152        if d4.ncols() != beta_link_wiggle.len() {
153            return Err(GamlssError::DimensionMismatch { reason: format!(
154                "wiggle fourth-derivative/beta mismatch: basis has {} columns but beta_link_wiggle has {} coefficients",
155                d4.ncols(),
156                beta_link_wiggle.len()
157            ) }.into());
158        }
159        Ok(d4.dot(&beta_link_wiggle))
160    }
161
162    pub(crate) fn wiggle_geometry(
163        &self,
164        q0: ArrayView1<'_, f64>,
165        beta_link_wiggle: ArrayView1<'_, f64>,
166    ) -> Result<BinomialMeanWiggleGeometry, String> {
167        // Frozen-basis (#1596): the warp offset `s = B⊥·β_w` uses the pinned,
168        // identifiable design `B⊥ = (I-P_X)B(η̂)`, so it is a per-row constant w.r.t.
169        // the *live* linear predictor `η`. Then `q = η + s` gives `∂q/∂η = 1`
170        // exactly and every higher derivative of `q` in `η` vanishes, and the
171        // `∂B/∂η` chain bases drop out (the warp basis does not move with `η`).
172        // The value basis `B⊥` — the column block carrying `∂q/∂β_w` — is the
173        // only surviving geometry term.
174        if let Some(frozen) = self.frozen_warp_design.as_ref() {
175            let n = frozen.nrows();
176            let pw = frozen.ncols();
177            return Ok(BinomialMeanWiggleGeometry {
178                basis: frozen.as_ref().clone(),
179                basis_d1: Array2::zeros((n, pw)),
180                basis_d2: Array2::zeros((n, pw)),
181                basis_d3: Array2::zeros((n, pw)),
182                dq_dq0: Array1::ones(n),
183                d2q_dq02: Array1::zeros(n),
184                d3q_dq03: Array1::zeros(n),
185                d4q_dq04: Array1::zeros(n),
186            });
187        }
188        let basis = self.wiggle_design(q0)?;
189        let basis_d1 = self.wiggle_basiswith_options(q0, BasisOptions::first_derivative())?;
190        let basis_d2 = self.wiggle_basiswith_options(q0, BasisOptions::second_derivative())?;
191        let basis_d3 = self.wiggle_d3basis_constrained(q0)?;
192        let dq_dq0 = self.wiggle_dq_dq0(q0, beta_link_wiggle)?;
193        let d2q_dq02 = self.wiggle_d2q_dq02(q0, beta_link_wiggle)?;
194        let d3q_dq03 = self.wiggle_d3q_dq03(q0, beta_link_wiggle)?;
195        let d4q_dq04 = self.wiggle_d4q_dq04(q0, beta_link_wiggle)?;
196        Ok(BinomialMeanWiggleGeometry {
197            basis,
198            basis_d1,
199            basis_d2,
200            basis_d3,
201            dq_dq0,
202            d2q_dq02,
203            d3q_dq03,
204            d4q_dq04,
205        })
206    }
207
208    pub(crate) fn neglog_q_derivatives(
209        &self,
210        y: f64,
211        weight: f64,
212        q: f64,
213    ) -> Result<(f64, f64, f64), String> {
214        let jet = inverse_link_jet_for_inverse_link(&self.link_kind, q)
215            .map_err(|e| format!("fixed-link wiggle inverse-link evaluation failed: {e}"))?;
216        // Pass μ RAW: the dispatch returns the exact q-derivatives of the
217        // evaluated loss for every representable μ in (0,1) and handles the
218        // saturated boundary itself. See binomial_location_scalerow (#948).
219        Ok(binomial_neglog_q_derivatives_dispatch(
220            y,
221            weight,
222            q,
223            jet.mu,
224            jet.d1,
225            jet.d2,
226            jet.d3,
227            &self.link_kind,
228        ))
229    }
230
231    pub(crate) fn neglog_q_fourth_derivative(
232        &self,
233        y: f64,
234        weight: f64,
235        q: f64,
236    ) -> Result<f64, String> {
237        let jet = inverse_link_jet_for_inverse_link(&self.link_kind, q)
238            .map_err(|e| format!("fixed-link wiggle inverse-link evaluation failed: {e}"))?;
239        // Pass μ RAW — see neglog_q_derivatives above (#948).
240        binomial_neglog_q_fourth_derivative_dispatch(
241            y,
242            weight,
243            q,
244            jet.mu,
245            jet.d1,
246            jet.d2,
247            jet.d3,
248            &self.link_kind,
249        )
250    }
251
252    pub(crate) fn dense_eta_design_fromspecs<'a>(
253        &self,
254        specs: &'a [ParameterBlockSpec],
255    ) -> Result<Cow<'a, Array2<f64>>, String> {
256        if specs.len() != 2 {
257            return Err(GamlssError::DimensionMismatch {
258                reason: format!(
259                    "BinomialMeanWiggleFamily expects 2 specs, got {}",
260                    specs.len()
261                ),
262            }
263            .into());
264        }
265        Ok(match specs[Self::BLOCK_ETA].design.as_dense_ref() {
266            Some(d) => Cow::Borrowed(d),
267            None => Cow::Owned(
268                specs[Self::BLOCK_ETA]
269                    .design
270                    .try_to_dense_with_policy(
271                        &self.policy.material_policy(),
272                        "BinomialMeanWiggle dense_eta_design_fromspecs eta",
273                    )
274                    .map_err(|e| e.to_string())?
275                    .as_ref()
276                    .clone(),
277            ),
278        })
279    }
280
281    pub(crate) fn exact_newton_joint_psi_direction(
282        &self,
283        block_states: &[ParameterBlockState],
284        derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
285        psi_index: usize,
286        x_eta: &Array2<f64>,
287    ) -> Result<Option<BinomialMeanWiggleJointPsiDirection>, String> {
288        validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
289        if derivative_blocks.len() != 2 {
290            return Err(GamlssError::DimensionMismatch { reason: format!(
291                "BinomialMeanWiggleFamily joint psi direction expects 2 derivative block lists, got {}",
292                derivative_blocks.len()
293            ) }.into());
294        }
295        let n = self.y.len();
296        let p_eta = x_eta.ncols();
297        let beta_eta = &block_states[Self::BLOCK_ETA].beta;
298        let mut global = 0usize;
299        for (block_idx, block_derivs) in derivative_blocks.iter().enumerate() {
300            for deriv in block_derivs {
301                if global == psi_index {
302                    if block_idx != Self::BLOCK_ETA {
303                        return Ok(None);
304                    }
305                    let x_eta_psi_map = resolve_custom_family_x_psi_map(
306                        deriv,
307                        n,
308                        p_eta,
309                        0..n,
310                        "BinomialMeanWiggleFamily eta",
311                        &self.policy,
312                    )?;
313                    let x_eta_psi = x_eta_psi_map.row_chunk(0..n)?;
314                    let z_eta_psi = x_eta_psi.dot(beta_eta);
315                    return Ok(Some(BinomialMeanWiggleJointPsiDirection {
316                        x_eta_psi: Some(x_eta_psi),
317                        z_eta_psi,
318                    }));
319                }
320                global += 1;
321            }
322        }
323        Ok(None)
324    }
325
326    pub(crate) fn exact_newton_joint_psi_action(
327        &self,
328        block_states: &[ParameterBlockState],
329        derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
330        psi_index: usize,
331        p_eta: usize,
332    ) -> Result<Option<(CustomFamilyPsiDesignAction, Array1<f64>)>, String> {
333        validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
334        if derivative_blocks.len() != 2 {
335            return Err(GamlssError::DimensionMismatch { reason: format!(
336                "BinomialMeanWiggleFamily joint psi action expects 2 derivative block lists, got {}",
337                derivative_blocks.len()
338            ) }.into());
339        }
340        let n = self.y.len();
341        let beta_eta = &block_states[Self::BLOCK_ETA].beta;
342        let mut global = 0usize;
343        for (block_idx, block_derivs) in derivative_blocks.iter().enumerate() {
344            for deriv in block_derivs {
345                if global == psi_index {
346                    if block_idx != Self::BLOCK_ETA {
347                        return Ok(None);
348                    }
349                    let action = match CustomFamilyPsiDesignAction::from_first_derivative(
350                        deriv,
351                        n,
352                        p_eta,
353                        0..n,
354                        "BinomialMeanWiggleFamily eta",
355                    ) {
356                        Ok(action) => action,
357                        Err(_) => return Ok(None),
358                    };
359                    let z_eta_psi = action.forward_mul(beta_eta.view());
360                    return Ok(Some((action, z_eta_psi)));
361                }
362                global += 1;
363            }
364        }
365        Ok(None)
366    }
367
368    pub(crate) fn bmw_static_hessian_operator(
369        &self,
370        block_states: &[ParameterBlockState],
371        x_eta_arc: Arc<Array2<f64>>,
372    ) -> Result<Arc<RowCoeffOperator>, String> {
373        validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
374        let eta = &block_states[Self::BLOCK_ETA].eta;
375        let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
376        let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
377        let n = self.y.len();
378        if eta.len() != n || etaw.len() != n || self.weights.len() != n {
379            return Err(GamlssError::DimensionMismatch {
380                reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
381            }
382            .into());
383        }
384        let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
385        let p_eta = x_eta_arc.ncols();
386        let pw = geom.basis.ncols();
387        let mut coeff_eta = Array1::<f64>::zeros(n);
388        let mut coeff_etaw_b = Array1::<f64>::zeros(n);
389        let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
390        let mut coeff_ww = Array1::<f64>::zeros(n);
391        for row in 0..n {
392            let q = eta[row] + etaw[row];
393            let (m1, m2, _) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
394            let a = geom.dq_dq0[row];
395            let b = geom.d2q_dq02[row];
396            coeff_eta[row] = hessian_coeff_fromobjective_q_terms(m1, m2, a, a, b);
397            coeff_etaw_b[row] = m2 * a;
398            coeff_etaw_d1[row] = m1;
399            coeff_ww[row] = m2;
400        }
401        Ok(Arc::new(RowCoeffOperator::from_directions(
402            vec![p_eta, pw],
403            vec![
404                (0, x_eta_arc),
405                (1, Arc::new(geom.basis)),
406                (1, Arc::new(geom.basis_d1)),
407            ],
408            vec![
409                (0, 0, coeff_eta),
410                (0, 1, coeff_etaw_b),
411                (0, 2, coeff_etaw_d1),
412                (1, 1, coeff_ww),
413            ],
414            n,
415        )))
416    }
417
418    pub(crate) fn bmw_directional_operator(
419        &self,
420        block_states: &[ParameterBlockState],
421        x_eta_arc: Arc<Array2<f64>>,
422        d_beta_flat: &Array1<f64>,
423    ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
424        validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
425        let eta = &block_states[Self::BLOCK_ETA].eta;
426        let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
427        let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
428        let n = self.y.len();
429        if eta.len() != n || etaw.len() != n || self.weights.len() != n {
430            return Err(GamlssError::DimensionMismatch {
431                reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
432            }
433            .into());
434        }
435        let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
436        let p_eta = x_eta_arc.ncols();
437        let pw = geom.basis.ncols();
438        let total = p_eta + pw;
439        if d_beta_flat.len() != total {
440            return Err(GamlssError::DimensionMismatch {
441                reason: format!(
442                    "BinomialMeanWiggleFamily joint d_beta length mismatch: got {}, expected {}",
443                    d_beta_flat.len(),
444                    total
445                ),
446            }
447            .into());
448        }
449        let u_eta = d_beta_flat.slice(s![0..p_eta]).to_owned();
450        let uw = d_beta_flat.slice(s![p_eta..total]).to_owned();
451        let xi = fast_av(x_eta_arc.as_ref(), &u_eta);
452        let phi = fast_av(&geom.basis, &uw);
453        let basis1_u = fast_av(&geom.basis_d1, &uw);
454        let basis2_u = fast_av(&geom.basis_d2, &uw);
455
456        let mut coeff_eta = Array1::<f64>::zeros(n);
457        let mut coeff_etaw_b = Array1::<f64>::zeros(n);
458        let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
459        let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
460        let mut coeff_ww_bb = Array1::<f64>::zeros(n);
461        let mut coeff_ww_db = Array1::<f64>::zeros(n);
462        for row in 0..n {
463            let q = eta[row] + etaw[row];
464            let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
465            let a = geom.dq_dq0[row];
466            let b = geom.d2q_dq02[row];
467            let c = geom.d3q_dq03[row];
468            let q_u = a * xi[row] + phi[row];
469            let a_u = b * xi[row] + basis1_u[row];
470            let b_u = c * xi[row] + basis2_u[row];
471            coeff_eta[row] = directionalhessian_coeff_fromobjective_q_terms(
472                m1, m2, m3, q_u, a, a, b, a_u, a_u, b_u,
473            );
474            coeff_etaw_b[row] = m3 * q_u * a + m2 * a_u;
475            coeff_etaw_d1[row] = m2 * (a * xi[row] + q_u);
476            coeff_etaw_d2[row] = m1 * xi[row];
477            coeff_ww_bb[row] = m3 * q_u;
478            coeff_ww_db[row] = m2 * xi[row];
479        }
480        Ok(Some(Arc::new(RowCoeffOperator::from_directions(
481            vec![p_eta, pw],
482            vec![
483                (0, x_eta_arc),
484                (1, Arc::new(geom.basis)),
485                (1, Arc::new(geom.basis_d1)),
486                (1, Arc::new(geom.basis_d2)),
487            ],
488            vec![
489                (0, 0, coeff_eta),
490                (0, 1, coeff_etaw_b),
491                (0, 2, coeff_etaw_d1),
492                (0, 3, coeff_etaw_d2),
493                (1, 1, coeff_ww_bb),
494                (1, 2, coeff_ww_db),
495            ],
496            n,
497        ))))
498    }
499
500    pub(crate) fn bmw_second_directional_operator(
501        &self,
502        block_states: &[ParameterBlockState],
503        x_eta_arc: Arc<Array2<f64>>,
504        d_beta_u_flat: &Array1<f64>,
505        d_beta_v_flat: &Array1<f64>,
506    ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
507        validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
508        let eta = &block_states[Self::BLOCK_ETA].eta;
509        let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
510        let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
511        let n = self.y.len();
512        if eta.len() != n || etaw.len() != n || self.weights.len() != n {
513            return Err(GamlssError::DimensionMismatch {
514                reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
515            }
516            .into());
517        }
518        let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
519        let p_eta = x_eta_arc.ncols();
520        let pw = geom.basis.ncols();
521        let total = p_eta + pw;
522        if d_beta_u_flat.len() != total || d_beta_v_flat.len() != total {
523            return Err(GamlssError::DimensionMismatch { reason: format!(
524                "BinomialMeanWiggleFamily joint second d_beta length mismatch: got {} and {}, expected {}",
525                d_beta_u_flat.len(),
526                d_beta_v_flat.len(),
527                total
528            ) }.into());
529        }
530        let u_eta = d_beta_u_flat.slice(s![0..p_eta]).to_owned();
531        let v_eta = d_beta_v_flat.slice(s![0..p_eta]).to_owned();
532        let uw = d_beta_u_flat.slice(s![p_eta..total]).to_owned();
533        let vw = d_beta_v_flat.slice(s![p_eta..total]).to_owned();
534
535        let xi_u = fast_av(x_eta_arc.as_ref(), &u_eta);
536        let xi_v = fast_av(x_eta_arc.as_ref(), &v_eta);
537        let phi_u = fast_av(&geom.basis, &uw);
538        let phi_v = fast_av(&geom.basis, &vw);
539        let b1u = fast_av(&geom.basis_d1, &uw);
540        let b1v = fast_av(&geom.basis_d1, &vw);
541        let b2u = fast_av(&geom.basis_d2, &uw);
542        let b2v = fast_av(&geom.basis_d2, &vw);
543        let b3u = fast_av(&geom.basis_d3, &uw);
544        let b3v = fast_av(&geom.basis_d3, &vw);
545
546        let mut coeff_eta = Array1::<f64>::zeros(n);
547        let mut coeff_etaw_b = Array1::<f64>::zeros(n);
548        let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
549        let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
550        let mut coeff_etaw_d3 = Array1::<f64>::zeros(n);
551        let mut coeff_ww_bb = Array1::<f64>::zeros(n);
552        let mut coeff_ww_db = Array1::<f64>::zeros(n);
553        let mut coeff_ww_ddb = Array1::<f64>::zeros(n);
554        let mut coeff_ww_dd = Array1::<f64>::zeros(n);
555
556        for row in 0..n {
557            let q = eta[row] + etaw[row];
558            let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
559            let m4 = self.neglog_q_fourth_derivative(self.y[row], self.weights[row], q)?;
560            let a = geom.dq_dq0[row];
561            let b = geom.d2q_dq02[row];
562            let c = geom.d3q_dq03[row];
563            let d = geom.d4q_dq04[row];
564
565            let q_u = a * xi_u[row] + phi_u[row];
566            let a_u = b * xi_u[row] + b1u[row];
567            let b_u = c * xi_u[row] + b2u[row];
568            let q_v = a * xi_v[row] + phi_v[row];
569            let a_v = b * xi_v[row] + b1v[row];
570            let b_v = c * xi_v[row] + b2v[row];
571            let q_uv = b * xi_u[row] * xi_v[row] + b1u[row] * xi_v[row] + b1v[row] * xi_u[row];
572            let a_uv = c * xi_u[row] * xi_v[row] + b2u[row] * xi_v[row] + b2v[row] * xi_u[row];
573            let b_uv = d * xi_u[row] * xi_v[row] + b3u[row] * xi_v[row] + b3v[row] * xi_u[row];
574
575            coeff_eta[row] = second_directionalhessian_coeff_fromobjective_q_terms(
576                m1, m2, m3, m4, q_u, q_v, q_uv, a, a, b, a_u, a_v, a_u, a_v, a_uv, a_uv, b_u, b_v,
577                b_uv,
578            );
579            let d2_c_b = m4 * q_u * q_v * a + m3 * (q_uv * a + q_u * a_v + q_v * a_u) + m2 * a_uv;
580            let dc_b_u = m3 * q_u * a + m2 * a_u;
581            let dc_b_v = m3 * q_v * a + m2 * a_v;
582            let c_b_static = m2 * a;
583            let d2_c_b1 = m3 * q_u * q_v + m2 * q_uv;
584            let dc_b1_u = m2 * q_u;
585            let dc_b1_v = m2 * q_v;
586
587            coeff_etaw_b[row] = d2_c_b;
588            coeff_etaw_d1[row] = dc_b_u * xi_v[row] + dc_b_v * xi_u[row] + d2_c_b1;
589            coeff_etaw_d2[row] =
590                c_b_static * xi_u[row] * xi_v[row] + dc_b1_u * xi_v[row] + dc_b1_v * xi_u[row];
591            coeff_etaw_d3[row] = m1 * xi_u[row] * xi_v[row];
592
593            let dw = m2;
594            let dw_u = m3 * q_u;
595            let dw_v = m3 * q_v;
596            let dw_uv = m4 * q_u * q_v + m3 * q_uv;
597            let xixj = xi_u[row] * xi_v[row];
598            coeff_ww_bb[row] = dw_uv;
599            coeff_ww_db[row] = dw_v * xi_u[row] + dw_u * xi_v[row];
600            coeff_ww_ddb[row] = dw * xixj;
601            coeff_ww_dd[row] = 2.0 * dw * xixj;
602        }
603
604        Ok(Some(Arc::new(RowCoeffOperator::from_directions(
605            vec![p_eta, pw],
606            vec![
607                (0, x_eta_arc),
608                (1, Arc::new(geom.basis)),
609                (1, Arc::new(geom.basis_d1)),
610                (1, Arc::new(geom.basis_d2)),
611                (1, Arc::new(geom.basis_d3)),
612            ],
613            vec![
614                (0, 0, coeff_eta),
615                (0, 1, coeff_etaw_b),
616                (0, 2, coeff_etaw_d1),
617                (0, 3, coeff_etaw_d2),
618                (0, 4, coeff_etaw_d3),
619                (1, 1, coeff_ww_bb),
620                (1, 2, coeff_ww_db),
621                (1, 3, coeff_ww_ddb),
622                (2, 2, coeff_ww_dd),
623            ],
624            n,
625        ))))
626    }
627
628    /// Build the [`BlockEffectiveJacobian`] for block `block_idx`.
629    ///
630    /// `BinomialMeanWiggle` has a single location output (n_outputs = 1):
631    /// - block 0 (eta):    output 0 = design rows
632    /// - block 1 (wiggle): all zeros (nonlinear link modulation)
633    pub fn block_effective_jacobian(
634        specs: &[ParameterBlockSpec],
635        block_idx: usize,
636    ) -> Result<Box<dyn BlockEffectiveJacobian>, String> {
637        crate::block_layout::block_jacobian::AdditiveWiggleBlockLayout {
638            family: "BinomialMeanWiggleFamily",
639            n_outputs: 1,
640            additive_blocks: &[Self::BLOCK_ETA],
641            wiggle_block: Some(Self::BLOCK_WIGGLE),
642        }
643        .block_effective_jacobian(specs, block_idx)
644    }
645}
646
647impl CustomFamily for BinomialMeanWiggleFamily {
648    fn exact_newton_joint_hessian_beta_dependent(&self) -> bool {
649        true
650    }
651
652    /// The binomial mean link-wiggle refit must NOT carry the full-span
653    /// Jeffreys/Firth augmentation, for the same structural reason
654    /// `GaussianLocationScaleWiggleFamily` opts out (#684–#688) — and the
655    /// binomial wiggle hits it harder. This is a *second-stage* refit: the
656    /// pilot binomial mean fit has already converged through the ordinary
657    /// PIRLS path (which is itself un-Firthed unless the user opts in — the
658    /// standard binomial fit logs `firth=false` / `jeffreys_logdet=none`), so
659    /// the wiggle refit only adds a *penalized*, *monotone-constrained*
660    /// I-spline link-shape correction `q = η + B(η)·β_w` around an
661    /// already-finite mode. Two failure modes follow from leaving the term on
662    /// (default `true`):
663    ///
664    /// 1. **Phantom stationarity residual.** When `H_pen` is full-rank and
665    ///    well-conditioned (the normal case — e.g. `cond ≈ 5.5e2` on the #872
666    ///    pure-probit repro) the Jeffreys gate smooth-steps the curvature
667    ///    `H_Φ → 0`, but the matching score `∇Φ` does not vanish in lock-step,
668    ///    so it leaks a nonzero `|∇L − Sβ + ∇Φ|` into the inner joint-Newton
669    ///    KKT residual. The certificate then refuses every iterate and the
670    ///    outer REML rejects all seeds (exactly the #684–#688 abort signature).
671    /// 2. **Saturation barrier / divergence.** `−Φ = −½log|I_J|` is folded into
672    ///    the objective and `∇Φ ∝ I_J⁻¹` into the gradient. The I-spline warp
673    ///    can drive the binomial linear predictor toward saturation, where the
674    ///    reduced Fisher information `I_J` goes singular: `−Φ → +∞` and
675    ///    `∇Φ → ∞`. The augmented objective grows a barrier that the joint
676    ///    Newton diverges into — the #872 repro runs the full 1200-cycle budget
677    ///    with the augmented objective pinned at ~4.6e9 and the augmented
678    ///    residual at ~5.8e9 while the plain data gradient is only ~2.3e2,
679    ///    aborting the documented `link(type=flexible(...)) + linkwiggle(...)`
680    ///    fit.
681    ///
682    /// Separation robustness is not lost: the wiggle block carries both a
683    /// difference penalty (λ selected by REML) and a hard non-negativity
684    /// constraint, and the underlying mean is fit by the pilot; a penalized,
685    /// constrained refit around a finite pilot mode does not run away to
686    /// `β → ∞` the way an unpenalized MLE can. Turning the term off here makes
687    /// the wiggle refit consistent with the un-Firthed pilot and removes the
688    /// phantom residual that blocked convergence.
689    fn joint_jeffreys_term_required(&self) -> bool {
690        false
691    }
692
693    fn coefficient_hessian_cost(&self, specs: &[ParameterBlockSpec]) -> u64 {
694        // The mean-wiggle Hessian is exposed as a row-coefficient operator,
695        // so the hot representation cost is one Θ(n · (p_eta + p_w)) HVP
696        // rather than dense Θ(n · (p_eta + p_w)^2) assembly.
697        let p_total = specs
698            .iter()
699            .map(|s| s.design.ncols() as u64)
700            .fold(0u64, |acc, p| acc.saturating_add(p));
701        (self.y.len() as u64).saturating_mul(p_total.max(1))
702    }
703
704    fn block_linear_constraints(
705        &self,
706        _: &[ParameterBlockState],
707        block_idx: usize,
708        spec: &ParameterBlockSpec,
709    ) -> Result<Option<ConstraintSet>, String> {
710        if block_idx != Self::BLOCK_WIGGLE {
711            return Ok(None);
712        }
713        // Frozen-basis residualization preserves the original I-spline
714        // coefficient coordinate, so the same exact non-negative cone applies
715        // in both dynamic and frozen modes. For β_w ≥ 0, the M-spline
716        // derivative gives dq/dη = 1 + B'(η)·β_w ≥ 1 everywhere.
717        Ok(monotone_wiggle_nonnegative_constraints(spec.design.ncols()))
718    }
719
720    fn post_update_block_beta(
721        &self,
722        _: &[ParameterBlockState],
723        block_idx: usize,
724        block_spec: &ParameterBlockSpec,
725        beta: Array1<f64>,
726    ) -> Result<Array1<f64>, String> {
727        assert!(!block_spec.name.is_empty());
728        if block_idx != Self::BLOCK_WIGGLE {
729            return Ok(beta);
730        }
731        let beta = project_monotone_wiggle_beta_nonnegative(beta);
732        validate_monotone_wiggle_beta_nonnegative(&beta, "BinomialMeanWiggleFamily post-update")?;
733        Ok(beta)
734    }
735
736    fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
737        validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
738        let eta = &block_states[Self::BLOCK_ETA].eta;
739        let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
740        let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
741        let n = self.y.len();
742        if eta.len() != n || etaw.len() != n || self.weights.len() != n {
743            return Err(GamlssError::DimensionMismatch {
744                reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
745            }
746            .into());
747        }
748        // Frozen-basis (#1596): with `B⊥` pinned, `q = η + B⊥·β_w` so
749        // `∂q/∂η = 1` exactly (the warp offset is constant in the live `η`).
750        // Otherwise `∂q/∂η = 1 + B'(η)·β_w`, the dynamic warp slope.
751        let dq_dq0 = if self.frozen_warp_design.is_some() {
752            Array1::<f64>::ones(n)
753        } else {
754            self.wiggle_dq_dq0(eta.view(), betaw.view())?
755        };
756        if dq_dq0.len() != n {
757            return Err(GamlssError::DimensionMismatch {
758                reason: format!(
759                    "BinomialMeanWiggleFamily dq/dq0 length mismatch: got {}, expected {}",
760                    dq_dq0.len(),
761                    n
762                ),
763            }
764            .into());
765        }
766
767        // Certify the entire batch before constructing any working array.  The
768        // q-space NLL derivative stack is the single source of truth for both
769        // the objective and the working geometry; probability/variance and
770        // derivative floors would define a different loss in the tails.
771        let mut rows = Vec::with_capacity(n);
772        for i in 0..n {
773            let q = eta[i] + etaw[i];
774            if !eta[i].is_finite() || !etaw[i].is_finite() || !q.is_finite() {
775                return Err(GamlssError::RowGeometryUnrepresentable {
776                    row: i,
777                    quantity: "binomial mean-wiggle predictor q",
778                    eta: eta[i],
779                    value: q,
780                }
781                .into());
782            }
783            let yi = self.y[i];
784            let wi = self.weights[i];
785            if !yi.is_finite() || !(0.0..=1.0).contains(&yi) {
786                return Err(GamlssError::InvalidInput {
787                    reason: format!(
788                        "BinomialMeanWiggleFamily requires y in [0, 1]; found y[{i}]={yi}"
789                    ),
790                }
791                .into());
792            }
793            if !wi.is_finite() || wi < 0.0 {
794                return Err(GamlssError::InvalidInput {
795                    reason: format!(
796                        "BinomialMeanWiggleFamily requires finite non-negative weights; found weight[{i}]={wi}"
797                    ),
798                }
799                .into());
800            }
801            let slope = dq_dq0[i];
802            if !slope.is_finite() {
803                return Err(GamlssError::RowGeometryUnrepresentable {
804                    row: i,
805                    quantity: "binomial mean-wiggle warp slope",
806                    eta: eta[i],
807                    value: slope,
808                }
809                .into());
810            }
811            if wi == 0.0 {
812                rows.push((0.0, eta[i], 0.0, etaw[i], 0.0));
813                continue;
814            }
815            let jet = inverse_link_jet_for_inverse_link(&self.link_kind, q)
816                .map_err(|e| format!("fixed-link wiggle inverse-link evaluation failed: {e}"))?;
817            let row_ll =
818                binomial_location_scale_log_likelihood(yi, wi, q, &self.link_kind, jet.mu)?;
819            // The objective and its gradient stay EXACT: `row_ll` is the true
820            // NLL and `m1` its exact q-space score. The working *curvature*,
821            // however, is the EXPECTED Fisher information
822            //   f = w·(μ'(q))² / (μ(1−μ)) ≥ 0,
823            // NOT the observed second derivative
824            //   m2 = −(ℓ''(μ)·μ'² + ℓ'(μ)·μ'').
825            // For a non-canonical binomial link (cauchit / loglog / …) the
826            // observed `m2` can go NEGATIVE at a row whose fitted probability
827            // disagrees with its response — cauchit at q≈1.248 with y=0 gives
828            // m2=−0.23 — because the ℓ'·μ'' term is unbounded below there. That
829            // is a genuine property of the observed Hessian, not an
830            // unrepresentable row, so keying representability off `m2 > 0`
831            // wrongly refused every such link (gam#2155). Fisher scoring uses the
832            // expected information, which is ≥ 0 for every interior μ, so the
833            // per-block working weight `X'WX` is PSD by construction and the
834            // inner solve and the outer REML `log|X'WX + S|` term are
835            // well-defined for every binomial link. This mirrors the
836            // location-scale sibling's expected-information choice
837            // (gam#1020 / gam#2353); the fixed point (score = 0) is identical to
838            // the observed-Newton one, only the iteration geometry changes.
839            let (m1, _, _) = binomial_neglog_q_derivatives_dispatch(
840                yi,
841                wi,
842                q,
843                jet.mu,
844                jet.d1,
845                jet.d2,
846                jet.d3,
847                &self.link_kind,
848            );
849            let (fisher, _, _) =
850                binomial_expected_q_information_derivatives(wi, jet.mu, jet.d1, jet.d2, jet.d3);
851            for (quantity, value, nonnegative) in [
852                ("binomial mean-wiggle row log likelihood", row_ll, false),
853                ("binomial mean-wiggle q score", m1, false),
854                ("binomial mean-wiggle expected q information", fisher, true),
855            ] {
856                if !value.is_finite() || (nonnegative && value < 0.0) {
857                    return Err(GamlssError::RowGeometryUnrepresentable {
858                        row: i,
859                        quantity,
860                        eta: q,
861                        value,
862                    }
863                    .into());
864                }
865            }
866            // η block working geometry — dead when the warp slope vanishes or the
867            // row carries no expected information (saturated μ).
868            let (z_eta_i, w_eta_i) = if slope == 0.0 || fisher == 0.0 {
869                (eta[i], 0.0)
870            } else {
871                let weight = fisher * slope * slope;
872                if !weight.is_finite() || weight <= 0.0 {
873                    return Err(GamlssError::RowGeometryUnrepresentable {
874                        row: i,
875                        quantity: "binomial mean-wiggle eta working weight",
876                        eta: eta[i],
877                        value: weight,
878                    }
879                    .into());
880                }
881                let response = eta[i] - m1 / (fisher * slope);
882                if !response.is_finite() {
883                    return Err(GamlssError::RowGeometryUnrepresentable {
884                        row: i,
885                        quantity: "binomial mean-wiggle eta working response",
886                        eta: eta[i],
887                        value: response,
888                    }
889                    .into());
890                }
891                (response, weight)
892            };
893            // wiggle block working geometry — dead when the row carries no
894            // expected information.
895            let (z_wiggle_i, w_wiggle_i) = if fisher == 0.0 {
896                (etaw[i], 0.0)
897            } else {
898                let z_wiggle_i = etaw[i] - m1 / fisher;
899                if !z_wiggle_i.is_finite() {
900                    return Err(GamlssError::RowGeometryUnrepresentable {
901                        row: i,
902                        quantity: "binomial mean-wiggle wiggle working response",
903                        eta: etaw[i],
904                        value: z_wiggle_i,
905                    }
906                    .into());
907                }
908                (z_wiggle_i, fisher)
909            };
910            rows.push((row_ll, z_eta_i, w_eta_i, z_wiggle_i, w_wiggle_i));
911        }
912
913        let mut ll = 0.0;
914        for (i, row) in rows.iter().enumerate() {
915            ll += row.0;
916            if !ll.is_finite() {
917                return Err(GamlssError::RowGeometryUnrepresentable {
918                    row: i,
919                    quantity: "binomial mean-wiggle cumulative log likelihood",
920                    eta: eta[i],
921                    value: ll,
922                }
923                .into());
924            }
925        }
926        let z_eta = Array1::from_iter(rows.iter().map(|row| row.1));
927        let w_eta = Array1::from_iter(rows.iter().map(|row| row.2));
928        let z_wiggle = Array1::from_iter(rows.iter().map(|row| row.3));
929        let w_wiggle = Array1::from_iter(rows.iter().map(|row| row.4));
930
931        Ok(FamilyEvaluation {
932            log_likelihood: ll,
933            blockworking_sets: vec![
934                BlockWorkingSet::diagonal_checked(z_eta, w_eta)?,
935                BlockWorkingSet::diagonal_checked(z_wiggle, w_wiggle)?,
936            ],
937        })
938    }
939
940    fn exact_newton_joint_gradient_evaluation(
941        &self,
942        block_states: &[ParameterBlockState],
943        specs: &[ParameterBlockSpec],
944    ) -> Result<Option<ExactNewtonJointGradientEvaluation>, String> {
945        // Assemble the exact joint score from the per-block IRLS working sets
946        // (X_bᵀ(w⊙(z−η)) per block), the same source of truth the inner
947        // joint-Newton RHS uses — consistent with the family's explicit joint
948        // Hessian and matching FD of the log-likelihood.
949        let eval = self.evaluate(block_states)?;
950        gamlss_joint_gradient_from_working_sets(&eval, specs, block_states).map(Some)
951    }
952
953    fn block_geometry(
954        &self,
955        block_states: &[ParameterBlockState],
956        spec: &ParameterBlockSpec,
957    ) -> Result<(DesignMatrix, Array1<f64>), String> {
958        if spec.name != "wiggle" {
959            return Ok((spec.design.clone(), spec.offset.clone()));
960        }
961        if block_states.is_empty() {
962            return Err(GamlssError::UnsupportedConfiguration {
963                reason: "wiggle geometry requires eta block".to_string(),
964            }
965            .into());
966        }
967        let eta = &block_states[Self::BLOCK_ETA].eta;
968        if eta.len() != self.y.len() {
969            return Err(GamlssError::DimensionMismatch {
970                reason: "BinomialMeanWiggleFamily eta size mismatch".to_string(),
971            }
972            .into());
973        }
974        // Frozen-basis (#1596): return the pinned, identifiable warp design
975        // `B⊥ = (I-P_X)B(η̂)` rather than the live `B(η)`. `B⊥` is constant across
976        // inner cycles, so the engine rebuilds the *same* matrix every cycle —
977        // the death-spiral source (a basis that moves under the line search) is
978        // gone, while the dynamic-geometry plumbing is preserved unchanged.
979        let x = match self.frozen_warp_design.as_ref() {
980            Some(frozen) => frozen.as_ref().clone(),
981            None => self.wiggle_design(eta.view())?,
982        };
983        if x.ncols() != spec.design.ncols() {
984            return Err(GamlssError::DimensionMismatch {
985                reason: format!(
986                    "dynamic wiggle design col mismatch: got {}, expected {}",
987                    x.ncols(),
988                    spec.design.ncols()
989                ),
990            }
991            .into());
992        }
993        let nrows = x.nrows();
994        Ok((
995            DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(x)),
996            Array1::zeros(nrows),
997        ))
998    }
999
1000    fn block_geometry_is_dynamic(&self) -> bool {
1001        true
1002    }
1003
1004    fn exact_newton_joint_hessian_workspace(
1005        &self,
1006        block_states: &[ParameterBlockState],
1007        specs: &[ParameterBlockSpec],
1008    ) -> Result<Option<Arc<dyn ExactNewtonJointHessianWorkspace>>, String> {
1009        let x_eta = self.dense_eta_design_fromspecs(specs)?.into_owned();
1010        let workspace = BinomialMeanWiggleHessianWorkspace::new(
1011            self.clone(),
1012            block_states.to_vec(),
1013            specs.to_vec(),
1014            x_eta,
1015        )?;
1016        Ok(Some(Arc::new(workspace)))
1017    }
1018
1019    fn inner_coefficient_hessian_hvp_available(&self, specs: &[ParameterBlockSpec]) -> bool {
1020        self.dense_eta_design_fromspecs(specs).is_ok()
1021    }
1022
1023    fn exact_newton_joint_hessian_with_specs(
1024        &self,
1025        block_states: &[ParameterBlockState],
1026        specs: &[ParameterBlockSpec],
1027    ) -> Result<Option<Array2<f64>>, String> {
1028        validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1029        let x_eta = self.dense_eta_design_fromspecs(specs)?;
1030        let eta = &block_states[Self::BLOCK_ETA].eta;
1031        let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1032        let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1033        let n = self.y.len();
1034        if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1035            return Err(GamlssError::DimensionMismatch {
1036                reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1037            }
1038            .into());
1039        }
1040        let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1041        let p_eta = x_eta.ncols();
1042        let pw = geom.basis.ncols();
1043        let mut coeff_eta = Array1::<f64>::zeros(n);
1044        let mut coeff_etaw_b = Array1::<f64>::zeros(n);
1045        let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
1046        let mut coeff_ww = Array1::<f64>::zeros(n);
1047        for row in 0..n {
1048            let q = eta[row] + etaw[row];
1049            let (m1, m2, _) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1050            let a = geom.dq_dq0[row];
1051            let b = geom.d2q_dq02[row];
1052            coeff_eta[row] = hessian_coeff_fromobjective_q_terms(m1, m2, a, a, b);
1053            coeff_etaw_b[row] = m2 * a;
1054            coeff_etaw_d1[row] = m1;
1055            coeff_ww[row] = m2;
1056        }
1057        let h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
1058        let h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
1059            + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?;
1060        let h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww)?;
1061        assert_eq!(h_eta_eta.nrows(), p_eta);
1062        assert_eq!(h_ww.nrows(), pw);
1063        Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
1064            &h_eta_eta, &h_eta_w, &h_ww,
1065        )))
1066    }
1067
1068    fn exact_newton_joint_hessian_directional_derivative_with_specs(
1069        &self,
1070        block_states: &[ParameterBlockState],
1071        specs: &[ParameterBlockSpec],
1072        d_beta_flat: &Array1<f64>,
1073    ) -> Result<Option<Array2<f64>>, String> {
1074        validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1075        let x_eta = self.dense_eta_design_fromspecs(specs)?;
1076        let eta = &block_states[Self::BLOCK_ETA].eta;
1077        let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1078        let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1079        let n = self.y.len();
1080        if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1081            return Err(GamlssError::DimensionMismatch {
1082                reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1083            }
1084            .into());
1085        }
1086        let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1087        let p_eta = x_eta.ncols();
1088        let pw = geom.basis.ncols();
1089        if d_beta_flat.len() != p_eta + pw {
1090            return Err(GamlssError::DimensionMismatch {
1091                reason: format!(
1092                    "BinomialMeanWiggleFamily joint d_beta length mismatch: got {}, expected {}",
1093                    d_beta_flat.len(),
1094                    p_eta + pw
1095                ),
1096            }
1097            .into());
1098        }
1099        let u_eta = d_beta_flat.slice(s![0..p_eta]).to_owned();
1100        let uw = d_beta_flat.slice(s![p_eta..p_eta + pw]).to_owned();
1101        let xi = x_eta.dot(&u_eta);
1102        let phi = geom.basis.dot(&uw);
1103        let basis1_u = geom.basis_d1.dot(&uw);
1104        let basis2_u = geom.basis_d2.dot(&uw);
1105
1106        let mut coeff_eta = Array1::<f64>::zeros(n);
1107        let mut coeff_etaw_b = Array1::<f64>::zeros(n);
1108        let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
1109        let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
1110        let mut coeff_ww_bb = Array1::<f64>::zeros(n);
1111        let mut coeff_ww_db = Array1::<f64>::zeros(n);
1112        for row in 0..n {
1113            let q = eta[row] + etaw[row];
1114            let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1115            let a = geom.dq_dq0[row];
1116            let b = geom.d2q_dq02[row];
1117            let c = geom.d3q_dq03[row];
1118            let q_u = a * xi[row] + phi[row];
1119            let a_u = b * xi[row] + basis1_u[row];
1120            let b_u = c * xi[row] + basis2_u[row];
1121            coeff_eta[row] = directionalhessian_coeff_fromobjective_q_terms(
1122                m1, m2, m3, q_u, a, a, b, a_u, a_u, b_u,
1123            );
1124            coeff_etaw_b[row] = m3 * q_u * a + m2 * a_u;
1125            coeff_etaw_d1[row] = m2 * (a * xi[row] + q_u);
1126            coeff_etaw_d2[row] = m1 * xi[row];
1127            coeff_ww_bb[row] = m3 * q_u;
1128            coeff_ww_db[row] = m2 * xi[row];
1129        }
1130
1131        let d_h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
1132        let d_h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
1133            + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?
1134            + &xt_diag_y_dense(&x_eta, &coeff_etaw_d2, &geom.basis_d2)?;
1135        let a_ww = xt_diag_y_dense(&geom.basis_d1, &coeff_ww_db, &geom.basis)?;
1136        let d_h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww_bb)? + &a_ww + a_ww.t();
1137        Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
1138            &d_h_eta_eta,
1139            &d_h_eta_w,
1140            &d_h_ww,
1141        )))
1142    }
1143
1144    /// Exact second-order directional derivative D²H[u,v] of the joint Hessian
1145    /// for the BinomialMeanWiggle two-block model (eta, wiggle).
1146    ///
1147    /// # Mathematical derivation
1148    ///
1149    /// The negative log-likelihood Hessian element for indices (a, b) in the
1150    /// joint coefficient vector is:
1151    ///
1152    ///   H_ab = m2 * q_a * q_b + m1 * q_ab
1153    ///
1154    /// where m_k = d^k F / dq^k (k-th derivative of the negative log-likelihood
1155    /// w.r.t. the effective predictor q), q_a = dq/d(beta_a), and q_ab =
1156    /// d²q/(d(beta_a) d(beta_b)).
1157    ///
1158    /// The effective predictor is q = q0 + w(q0) where q0 = X_eta * beta_eta
1159    /// and w(q0) = B(q0) * beta_w is the link wiggle.  Write:
1160    ///   a = dq/dq0 = 1 + B'·beta_w       (geometry first derivative)
1161    ///   b = d²q/dq0² = B''·beta_w         (geometry second derivative)
1162    ///   c = d³q/dq0³ = B'''·beta_w        (geometry third derivative)
1163    ///   d = d⁴q/dq0⁴ = B''''·beta_w       (geometry fourth derivative)
1164    ///
1165    /// For a perturbation direction u = (u_eta, u_w), the chain-rule
1166    /// perturbations are:
1167    ///   q_u   = a·xi_u + phi_u             (first-order predictor perturbation)
1168    ///   a_u   = b·xi_u + basis1_u          (perturbation of geometry factor a)
1169    ///   b_u   = c·xi_u + basis2_u          (perturbation of geometry factor b)
1170    ///   c_u   = d·xi_u + basis3_u          (perturbation of geometry factor c)
1171    ///
1172    /// where xi_u = X_eta·u_eta, phi_u = B·u_w, basis_k_u = B^(k)·u_w.
1173    ///
1174    /// Mixed second-order perturbations (u,v) are:
1175    ///   q_uv  = b·xi_u·xi_v + basis1_u·xi_v + basis1_v·xi_u
1176    ///   a_uv  = c·xi_u·xi_v + basis2_u·xi_v + basis2_v·xi_u
1177    ///   b_uv  = d·xi_u·xi_v + basis3_u·xi_v + basis3_v·xi_u
1178    ///
1179    /// ## Block decomposition
1180    ///
1181    /// **eta-eta block** (X_eta' diag(coeff) X_eta):
1182    ///   The Hessian element for eta indices (i,j) factors as
1183    ///     H(eta_i, eta_j) = [m2·a² + m1·b] · x_eta(i)·x_eta(j)
1184    ///   so D²H_eta_eta[u,v] = X_eta' diag(coeff_eta) X_eta
1185    ///   where coeff_eta uses `second_directionalhessian_coeff_fromobjective_q_terms`
1186    ///   with q_a=a, q_b=a, q_ab=b and their chain-rule perturbations.
1187    ///
1188    /// **eta-w block** (X_eta' diag(...) [B, B', B'', B''']):
1189    ///   The static Hessian is:
1190    ///     H(eta_i, w_j) = (m2·a)·x_eta(i)·B_j + m1·x_eta(i)·B'_j
1191    ///   Taking D²[u,v] requires differentiating both the scalar coefficients
1192    ///   (m2·a, m1) and the basis matrices (B, B' depend on q0 via the chain
1193    ///   rule dB_j/du = B'_j·xi_u).  The full product rule gives four basis-matrix
1194    ///   tiers: B, B', B'', B'''.
1195    ///
1196    /// **w-w block** (B' diag(...) B, etc.):
1197    ///   The static Hessian is H(w_i, w_j) = m2·B_i·B_j.
1198    ///   D²[u,v] expands via the product rule on m2, B_i, B_j, each of which
1199    ///   depends on beta through q and q0.  This gives terms involving
1200    ///   B·B, B'·B, B'·B', and B''·B (all symmetrised).
1201    fn exact_newton_joint_hessian_second_directional_derivative_with_specs(
1202        &self,
1203        block_states: &[ParameterBlockState],
1204        specs: &[ParameterBlockSpec],
1205        d_beta_u_flat: &Array1<f64>,
1206        d_beta_v_flat: &Array1<f64>,
1207    ) -> Result<Option<Array2<f64>>, String> {
1208        validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1209        let x_eta = self.dense_eta_design_fromspecs(specs)?;
1210        let eta = &block_states[Self::BLOCK_ETA].eta;
1211        let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1212        let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1213        let n = self.y.len();
1214        if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1215            return Err(GamlssError::DimensionMismatch {
1216                reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1217            }
1218            .into());
1219        }
1220        let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1221        let p_eta = x_eta.ncols();
1222        let pw = geom.basis.ncols();
1223        let total = p_eta + pw;
1224        if d_beta_u_flat.len() != total || d_beta_v_flat.len() != total {
1225            return Err(GamlssError::DimensionMismatch { reason: format!(
1226                "BinomialMeanWiggleFamily joint second d_beta length mismatch: got {} and {}, expected {}",
1227                d_beta_u_flat.len(),
1228                d_beta_v_flat.len(),
1229                total
1230            ) }.into());
1231        }
1232
1233        // Split directions into eta and wiggle components.
1234        let u_eta = d_beta_u_flat.slice(s![0..p_eta]).to_owned();
1235        let v_eta = d_beta_v_flat.slice(s![0..p_eta]).to_owned();
1236        let uw = d_beta_u_flat.slice(s![p_eta..total]).to_owned();
1237        let vw = d_beta_v_flat.slice(s![p_eta..total]).to_owned();
1238
1239        // Per-row linear-predictor perturbations from each direction.
1240        let xi_u = x_eta.dot(&u_eta); // eta perturbation in direction u
1241        let xi_v = x_eta.dot(&v_eta); // eta perturbation in direction v
1242        let phi_u = geom.basis.dot(&uw); // direct wiggle basis, direction u
1243        let phi_v = geom.basis.dot(&vw); // direct wiggle basis, direction v
1244        let b1u = geom.basis_d1.dot(&uw); // first-derivative basis, direction u
1245        let b1v = geom.basis_d1.dot(&vw);
1246        let b2u = geom.basis_d2.dot(&uw); // second-derivative basis, direction u
1247        let b2v = geom.basis_d2.dot(&vw);
1248        let b3u = geom.basis_d3.dot(&uw); // third-derivative basis, direction u
1249        let b3v = geom.basis_d3.dot(&vw);
1250
1251        // Per-row chain-rule perturbations of q, a = dq/dq0, b = d²q/dq0²:
1252        //   q_u = a·xi_u + phi_u
1253        //   a_u = b·xi_u + basis1_u
1254        //   b_u = c·xi_u + basis2_u
1255        //   c_u = d·xi_u + basis3_u
1256        // Mixed second-order perturbations:
1257        //   q_uv = b·xi_u·xi_v + basis1_u·xi_v + basis1_v·xi_u
1258        //   a_uv = c·xi_u·xi_v + basis2_u·xi_v + basis2_v·xi_u
1259        //   b_uv = d·xi_u·xi_v + basis3_u·xi_v + basis3_v·xi_u
1260
1261        // Scaled basis matrices for the cross-product terms in the w-w and eta-w
1262        // blocks (same pattern as GaussianLocationScaleWiggleFamily).
1263        let basis_u = scale_matrix_rows(&geom.basis_d1, &xi_u)?; // dB/du = B'·xi_u
1264        let basis_v = scale_matrix_rows(&geom.basis_d1, &xi_v)?; // dB/dv = B'·xi_v
1265        let basis_uv = scale_matrix_rows(&geom.basis_d2, &(&xi_u * &xi_v))?; // d²B/dudv = B''·xi_u·xi_v
1266        // Per-row coefficient arrays for assembling the block-matrix products.
1267        let mut coeff_eta = Array1::<f64>::zeros(n);
1268
1269        // Coefficients for the eta-w block: X_eta' diag(c_*) M where M ∈ {B, B', B'', B'''}
1270        //
1271        // The static cross-Hessian is:
1272        //   H(eta_i, w_j) = (m2·a)·x_i·B_j + m1·x_i·B'_j
1273        // where B_j and B'_j are row evaluations of basis column j.
1274        //
1275        // Write C_B = m2·a (scalar coefficient multiplying B in the cross block)
1276        // and   C_B1 = m1  (scalar coefficient multiplying B' in the cross block).
1277        //
1278        // Product rule on C_B·B:
1279        //   d(C_B·B)/du = (dC_B/du)·B + C_B·B'·xi_u
1280        //   d²(C_B·B)/dudv = (d²C_B/dudv)·B + (dC_B/du)·B'·xi_v
1281        //                   + (dC_B/dv)·B'·xi_u + C_B·B''·xi_u·xi_v
1282        //
1283        // Product rule on C_B1·B':
1284        //   d²(C_B1·B')/dudv = (d²C_B1/dudv)·B' + (dC_B1/du)·B''·xi_v
1285        //                     + (dC_B1/dv)·B''·xi_u + C_B1·B'''·xi_u·xi_v
1286        //
1287        // Derivatives of the scalar coefficients:
1288        //   C_B  = m2·a
1289        //   dC_B/du  = m3·q_u·a + m2·a_u
1290        //   dC_B/dv  = m3·q_v·a + m2·a_v
1291        //   d²C_B/dudv = m4·q_u·q_v·a + m3·(q_uv·a + q_u·a_v + q_v·a_u) + m2·a_uv
1292        //
1293        //   C_B1 = m1
1294        //   dC_B1/du = m2·q_u
1295        //   dC_B1/dv = m2·q_v
1296        //   d²C_B1/dudv = m3·q_u·q_v + m2·q_uv
1297        //
1298        // Grouping by basis-matrix tier:
1299        //   B:   d²C_B/dudv
1300        //   B':  (dC_B/du)·xi_v + (dC_B/dv)·xi_u + d²C_B1/dudv
1301        //   B'': C_B·xi_u·xi_v + (dC_B1/du)·xi_v + (dC_B1/dv)·xi_u
1302        //   B''': C_B1·xi_u·xi_v
1303        let mut coeff_etaw_b = Array1::<f64>::zeros(n);
1304        let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
1305        let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
1306        let mut coeff_etaw_d3 = Array1::<f64>::zeros(n);
1307
1308        // Coefficients for the w-w block.
1309        //
1310        // The static w-w Hessian is:
1311        //   H(w_i, w_j) = m2·B_i·B_j
1312        //
1313        // Note: there is no m1·q_ij term because d²q/(d(beta_w_i) d(beta_w_j)) = 0
1314        // (the basis vectors B_i enter q linearly in beta_w).
1315        //
1316        // Product rule on m2·B_i·B_j, treating each factor as depending on beta:
1317        //   d²(m2·B_i·B_j)/dudv
1318        //     = (d²m2/dudv)·B_i·B_j                        → B'diag B  (symmetrised)
1319        //     + (dm2/du)·(B'_i·xi_v·B_j + B_i·B'_j·xi_v)  → dw_u terms
1320        //     + (dm2/dv)·(B'_i·xi_u·B_j + B_i·B'_j·xi_u)  → dw_v terms
1321        //     + m2·(B''_i·xi_u·xi_v·B_j + B'_i·xi_u·B'_j·xi_v
1322        //          + B'_i·xi_v·B'_j·xi_u + B_i·B''_j·xi_u·xi_v)
1323        //
1324        // where dm2/du = m3·q_u, dm2/dv = m3·q_v, d²m2/dudv = m4·q_u·q_v + m3·q_uv.
1325        //
1326        // Following the Gaussian LS wiggle pattern, we express this via:
1327        //   xt_diag_x_dense(B, dw_uv)                    — coeff: d²m2
1328        //   xt_diag_y_dense(basis_u, dw_v, B) + transpose — dB/du weighted by dm2/dv
1329        //   xt_diag_y_dense(basis_v, dw_u, B) + transpose — dB/dv weighted by dm2/du
1330        //   xt_diag_y_dense(basis_uv, w, B) + transpose   — d²B/dudv weighted by m2
1331        //   xt_diag_y_dense(basis_u, w, basis_v) + transpose — dB/du·dB/dv weighted by m2
1332        let mut dw = Array1::<f64>::zeros(n);
1333        let mut dw_u = Array1::<f64>::zeros(n);
1334        let mut dw_v = Array1::<f64>::zeros(n);
1335        let mut dw_uv = Array1::<f64>::zeros(n);
1336
1337        for row in 0..n {
1338            let q = eta[row] + etaw[row];
1339            let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1340            let m4 = self.neglog_q_fourth_derivative(self.y[row], self.weights[row], q)?;
1341            let a = geom.dq_dq0[row];
1342            let b = geom.d2q_dq02[row];
1343            let c = geom.d3q_dq03[row];
1344            let d = geom.d4q_dq04[row];
1345
1346            // Chain-rule perturbations in direction u.
1347            let q_u = a * xi_u[row] + phi_u[row];
1348            let a_u = b * xi_u[row] + b1u[row];
1349            let b_u = c * xi_u[row] + b2u[row];
1350
1351            // Chain-rule perturbations in direction v.
1352            let q_v = a * xi_v[row] + phi_v[row];
1353            let a_v = b * xi_v[row] + b1v[row];
1354            let b_v = c * xi_v[row] + b2v[row];
1355
1356            // Mixed second-order perturbations.
1357            let q_uv = b * xi_u[row] * xi_v[row] + b1u[row] * xi_v[row] + b1v[row] * xi_u[row];
1358            let a_uv = c * xi_u[row] * xi_v[row] + b2u[row] * xi_v[row] + b2v[row] * xi_u[row];
1359            let b_uv = d * xi_u[row] * xi_v[row] + b3u[row] * xi_v[row] + b3v[row] * xi_u[row];
1360
1361            // ── eta-eta block ──
1362            // H(eta_i, eta_j) uses q_a = a, q_b = a, q_ab = b (absorbing x_eta
1363            // into the matrix product).  The perturbations of these geometric
1364            // quantities are: dq_a/du = a_u, dq_b/du = a_u (since q_a = q_b = a),
1365            // dq_ab/du = b_u (since q_ab = b), and analogously for v.
1366            coeff_eta[row] = second_directionalhessian_coeff_fromobjective_q_terms(
1367                m1, m2, m3, m4, q_u, q_v, q_uv, a, a, b, // q_a, q_b, q_ab
1368                a_u, a_v, // dq_a_u, dq_a_v
1369                a_u, a_v, // dq_b_u, dq_b_v  (q_b = a so same perturbation)
1370                a_uv, a_uv, // d2q_a_uv, d2q_b_uv
1371                b_u, b_v,  // dq_ab_u, dq_ab_v  (q_ab = b)
1372                b_uv, // d2q_ab_uv
1373            );
1374
1375            // ── eta-w block coefficients ──
1376            // See the derivation in the docstring above.  We group by which basis
1377            // matrix tier (B, B', B'', B''') the coefficient multiplies.
1378
1379            // d²(m2·a)/dudv
1380            let d2_c_b = m4 * q_u * q_v * a + m3 * (q_uv * a + q_u * a_v + q_v * a_u) + m2 * a_uv;
1381            // d(m2·a)/du and d(m2·a)/dv
1382            let dc_b_u = m3 * q_u * a + m2 * a_u;
1383            let dc_b_v = m3 * q_v * a + m2 * a_v;
1384            // m2·a (static coefficient for B in the cross block)
1385            let c_b_static = m2 * a;
1386            // d²(m1)/dudv
1387            let d2_c_b1 = m3 * q_u * q_v + m2 * q_uv;
1388            // d(m1)/du and d(m1)/dv
1389            let dc_b1_u = m2 * q_u;
1390            let dc_b1_v = m2 * q_v;
1391
1392            coeff_etaw_b[row] = d2_c_b;
1393            coeff_etaw_d1[row] = dc_b_u * xi_v[row] + dc_b_v * xi_u[row] + d2_c_b1;
1394            coeff_etaw_d2[row] =
1395                c_b_static * xi_u[row] * xi_v[row] + dc_b1_u * xi_v[row] + dc_b1_v * xi_u[row];
1396            coeff_etaw_d3[row] = m1 * xi_u[row] * xi_v[row];
1397
1398            // ── w-w block coefficients ──
1399            // The w-w static Hessian coefficient is m2 (for B'diag B).
1400            dw[row] = m2;
1401            dw_u[row] = m3 * q_u;
1402            dw_v[row] = m3 * q_v;
1403            dw_uv[row] = m4 * q_u * q_v + m3 * q_uv;
1404        }
1405
1406        // ── Assemble eta-eta block ──
1407        let d2_h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
1408
1409        // ── Assemble eta-w block ──
1410        // The second-order directional derivative of the cross block H_eta_w is:
1411        //   d²H_eta_w[u,v] = X_eta' diag(coeff_etaw_b)  B
1412        //                   + X_eta' diag(coeff_etaw_d1) B'
1413        //                   + X_eta' diag(coeff_etaw_d2) B''
1414        //                   + X_eta' diag(coeff_etaw_d3) B'''
1415        let d2_h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
1416            + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?
1417            + &xt_diag_y_dense(&x_eta, &coeff_etaw_d2, &geom.basis_d2)?
1418            + &xt_diag_y_dense(&x_eta, &coeff_etaw_d3, &geom.basis_d3)?;
1419
1420        // ── Assemble w-w block ──
1421        // Following the Gaussian LS wiggle pattern (lines 6351-6363), the w-w
1422        // second directional derivative is assembled from scaled basis products:
1423        //
1424        //   d²(m2·B_i·B_j)/dudv decomposition:
1425        //     (d²m2)     · B_i·B_j        → xt_diag_x(B, dw_uv)
1426        //     (dm2/du)   · dB_j/dv · B_i  → xt_diag_y(basis_v, dw_u, B) + transpose
1427        //     (dm2/dv)   · dB_j/du · B_i  → xt_diag_y(basis_u, dw_v, B) + transpose
1428        //     m2 · d²B_j/dudv · B_i       → xt_diag_y(basis_uv, dw, B) + transpose
1429        //     m2 · dB_i/du · dB_j/dv      → xt_diag_y(basis_u, dw, basis_v) + transpose
1430        let a_ab = xt_diag_y_dense(&basis_uv, &dw, &geom.basis)?;
1431        let a_ij = xt_diag_y_dense(&basis_u, &dw, &basis_v)?;
1432        let a_iwj = xt_diag_y_dense(&basis_u, &dw_v, &geom.basis)?;
1433        let a_jwi = xt_diag_y_dense(&basis_v, &dw_u, &geom.basis)?;
1434        let d2_h_ww = &a_ab
1435            + &a_ab.t()
1436            + &a_ij
1437            + a_ij.t()
1438            + &a_iwj
1439            + a_iwj.t()
1440            + &a_jwi
1441            + a_jwi.t()
1442            + &xt_diag_x_dense(&geom.basis, &dw_uv)?;
1443
1444        Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
1445            &d2_h_eta_eta,
1446            &d2_h_eta_w,
1447            &d2_h_ww,
1448        )))
1449    }
1450
1451    fn exact_newton_joint_psi_terms(
1452        &self,
1453        block_states: &[ParameterBlockState],
1454        specs: &[ParameterBlockSpec],
1455        hyper_layout: &crate::custom_family::CustomFamilyHyperLayout,
1456        psi_index: usize,
1457    ) -> Result<Option<gam_problem::ExactNewtonJointPsiTerms>, String> {
1458        if hyper_layout.family_axis_count() != 0 {
1459            return Err("BinomialMeanWiggleFamily does not declare family-owned hyper axes"
1460                .to_string());
1461        }
1462        let derivative_blocks = hyper_layout.design_derivative_blocks();
1463        validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1464        if derivative_blocks.len() != 2 {
1465            return Err(GamlssError::DimensionMismatch { reason: format!(
1466                "BinomialMeanWiggleFamily joint psi terms expect 2 derivative block lists, got {}",
1467                derivative_blocks.len()
1468            ) }.into());
1469        }
1470        let x_eta = self.dense_eta_design_fromspecs(specs)?;
1471        let eta = &block_states[Self::BLOCK_ETA].eta;
1472        let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1473        let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1474        let n = self.y.len();
1475        if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1476            return Err(GamlssError::DimensionMismatch {
1477                reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1478            }
1479            .into());
1480        }
1481        let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1482        let p_eta = x_eta.ncols();
1483        let pw = geom.basis.ncols();
1484        let implicit_dir =
1485            self.exact_newton_joint_psi_action(block_states, derivative_blocks, psi_index, p_eta)?;
1486        let dense_dir = if implicit_dir.is_none() {
1487            self.exact_newton_joint_psi_direction(
1488                block_states,
1489                derivative_blocks,
1490                psi_index,
1491                &x_eta,
1492            )?
1493        } else {
1494            None
1495        };
1496        let z_eta_psi = if let Some((_, ref z_eta_psi)) = implicit_dir {
1497            z_eta_psi
1498        } else if let Some(ref dir_a) = dense_dir {
1499            &dir_a.z_eta_psi
1500        } else {
1501            return Ok(None);
1502        };
1503
1504        let mut objective_psi = 0.0;
1505        let mut score_eta_xa = Array1::<f64>::zeros(n);
1506        let mut score_eta_x = Array1::<f64>::zeros(n);
1507        let mut score_w_b = Array1::<f64>::zeros(n);
1508        let mut score_w_d1 = Array1::<f64>::zeros(n);
1509
1510        let mut coeff_eta_eta_xx = Array1::<f64>::zeros(n);
1511        let mut coeff_eta_eta_xa_x = Array1::<f64>::zeros(n);
1512        let mut coeff_eta_w_xa_b = Array1::<f64>::zeros(n);
1513        let mut coeff_eta_w_x_b = Array1::<f64>::zeros(n);
1514        let mut coeff_eta_w_x_d1 = Array1::<f64>::zeros(n);
1515        let mut coeff_eta_w_xa_d1 = Array1::<f64>::zeros(n);
1516        let mut coeff_eta_w_x_d2 = Array1::<f64>::zeros(n);
1517        let mut coeff_ww_bb = Array1::<f64>::zeros(n);
1518        let mut coeff_ww_db = Array1::<f64>::zeros(n);
1519
1520        for row in 0..n {
1521            let q = eta[row] + etaw[row];
1522            let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1523            let z_a = z_eta_psi[row];
1524            let a = geom.dq_dq0[row];
1525            let b = geom.d2q_dq02[row];
1526            let c = geom.d3q_dq03[row];
1527            let q_a = a * z_a;
1528
1529            objective_psi += m1 * q_a;
1530
1531            score_eta_xa[row] = m1 * a;
1532            score_eta_x[row] = m2 * q_a * a + m1 * b * z_a;
1533            score_w_b[row] = m2 * q_a;
1534            score_w_d1[row] = m1 * z_a;
1535
1536            coeff_eta_eta_xx[row] =
1537                m3 * q_a * a * a + m2 * (2.0 * a * b * z_a + q_a * b) + m1 * c * z_a;
1538            coeff_eta_eta_xa_x[row] = m2 * a * a + m1 * b;
1539            coeff_eta_w_xa_b[row] = m2 * a;
1540            coeff_eta_w_x_b[row] = m3 * q_a * a + m2 * b * z_a;
1541            coeff_eta_w_x_d1[row] = m2 * (a * z_a + q_a);
1542            coeff_eta_w_xa_d1[row] = m1;
1543            coeff_eta_w_x_d2[row] = m1 * z_a;
1544            coeff_ww_bb[row] = m3 * q_a;
1545            coeff_ww_db[row] = m2 * z_a;
1546        }
1547
1548        let score_w = gam_linalg::faer_ndarray::fast_atv(&geom.basis, &score_w_b)
1549            + gam_linalg::faer_ndarray::fast_atv(&geom.basis_d1, &score_w_d1);
1550
1551        if let Some((action, _)) = implicit_dir {
1552            let score_eta = action.transpose_mul(score_eta_xa.view())
1553                + gam_linalg::faer_ndarray::fast_atv(x_eta.as_ref(), &score_eta_x);
1554            let score_psi = binomial_pack_mean_wiggle_joint_score(&score_eta, &score_w);
1555            let x_eta_arc = shared_dense_arc(x_eta.as_ref());
1556            let basis_arc = Arc::new(geom.basis.clone());
1557            let basis_d1_arc = Arc::new(geom.basis_d1.clone());
1558            let basis_d2_arc = Arc::new(geom.basis_d2.clone());
1559            let zeros = Array1::<f64>::zeros(n);
1560            let operator = CustomFamilyJointPsiOperator::new(
1561                p_eta + pw,
1562                vec![
1563                    CustomFamilyJointDesignChannel::new(
1564                        0..p_eta,
1565                        Arc::clone(&x_eta_arc),
1566                        Some(action),
1567                    ),
1568                    CustomFamilyJointDesignChannel::new(
1569                        p_eta..p_eta + pw,
1570                        Arc::clone(&basis_arc),
1571                        None,
1572                    ),
1573                    CustomFamilyJointDesignChannel::new(
1574                        p_eta..p_eta + pw,
1575                        Arc::clone(&basis_d1_arc),
1576                        None,
1577                    ),
1578                    CustomFamilyJointDesignChannel::new(
1579                        p_eta..p_eta + pw,
1580                        Arc::clone(&basis_d2_arc),
1581                        None,
1582                    ),
1583                ],
1584                vec![
1585                    CustomFamilyJointDesignPairContribution::new(
1586                        0,
1587                        0,
1588                        coeff_eta_eta_xa_x.clone(),
1589                        coeff_eta_eta_xx.clone(),
1590                    ),
1591                    CustomFamilyJointDesignPairContribution::new(
1592                        0,
1593                        1,
1594                        coeff_eta_w_xa_b.clone(),
1595                        coeff_eta_w_x_b.clone(),
1596                    ),
1597                    CustomFamilyJointDesignPairContribution::new(
1598                        1,
1599                        0,
1600                        coeff_eta_w_xa_b.clone(),
1601                        coeff_eta_w_x_b.clone(),
1602                    ),
1603                    CustomFamilyJointDesignPairContribution::new(
1604                        0,
1605                        2,
1606                        coeff_eta_w_xa_d1.clone(),
1607                        coeff_eta_w_x_d1.clone(),
1608                    ),
1609                    CustomFamilyJointDesignPairContribution::new(
1610                        2,
1611                        0,
1612                        coeff_eta_w_xa_d1.clone(),
1613                        coeff_eta_w_x_d1.clone(),
1614                    ),
1615                    CustomFamilyJointDesignPairContribution::new(
1616                        0,
1617                        3,
1618                        zeros.clone(),
1619                        coeff_eta_w_x_d2.clone(),
1620                    ),
1621                    CustomFamilyJointDesignPairContribution::new(
1622                        3,
1623                        0,
1624                        zeros.clone(),
1625                        coeff_eta_w_x_d2.clone(),
1626                    ),
1627                    CustomFamilyJointDesignPairContribution::new(
1628                        1,
1629                        1,
1630                        zeros.clone(),
1631                        coeff_ww_bb.clone(),
1632                    ),
1633                    CustomFamilyJointDesignPairContribution::new(
1634                        2,
1635                        1,
1636                        zeros.clone(),
1637                        coeff_ww_db.clone(),
1638                    ),
1639                    CustomFamilyJointDesignPairContribution::new(1, 2, zeros, coeff_ww_db.clone()),
1640                ],
1641            );
1642            return Ok(Some(gam_problem::ExactNewtonJointPsiTerms {
1643                objective_psi,
1644                score_psi,
1645                hessian_psi: Array2::zeros((0, 0)),
1646                hessian_psi_operator: Some(std::sync::Arc::new(operator)),
1647            }));
1648        }
1649
1650        let dir_a =
1651            dense_dir.expect("dense psi direction should exist when implicit direction is absent");
1652        let x_eta_psi = dir_a
1653            .x_eta_psi
1654            .as_ref()
1655            .expect("dense eta psi design should exist when implicit direction is absent");
1656        let score_psi = binomial_pack_mean_wiggle_joint_score(
1657            &(gam_linalg::faer_ndarray::fast_atv(x_eta_psi, &score_eta_xa)
1658                + gam_linalg::faer_ndarray::fast_atv(x_eta.as_ref(), &score_eta_x)),
1659            &score_w,
1660        );
1661        let a_eta_eta = xt_diag_y_dense(x_eta_psi, &coeff_eta_eta_xa_x, &x_eta)?;
1662        let h_eta_eta = &a_eta_eta + &a_eta_eta.t() + &xt_diag_x_dense(&x_eta, &coeff_eta_eta_xx)?;
1663        let h_eta_w = xt_diag_y_dense(x_eta_psi, &coeff_eta_w_xa_b, &geom.basis)?
1664            + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_b, &geom.basis)?
1665            + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_d1, &geom.basis_d1)?
1666            + &xt_diag_y_dense(x_eta_psi, &coeff_eta_w_xa_d1, &geom.basis_d1)?
1667            + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_d2, &geom.basis_d2)?;
1668        let a_ww = xt_diag_y_dense(&geom.basis_d1, &coeff_ww_db, &geom.basis)?;
1669        let h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww_bb)? + &a_ww + a_ww.t();
1670
1671        Ok(Some(gam_problem::ExactNewtonJointPsiTerms {
1672            objective_psi,
1673            score_psi,
1674            hessian_psi: binomial_pack_mean_wiggle_joint_symmetrichessian(
1675                &h_eta_eta, &h_eta_w, &h_ww,
1676            ),
1677            hessian_psi_operator: None,
1678        }))
1679    }
1680}
1681
1682pub(crate) struct BinomialMeanWiggleHessianWorkspace {
1683    pub(crate) family: BinomialMeanWiggleFamily,
1684    pub(crate) block_states: Vec<ParameterBlockState>,
1685    /// The block specs the workspace was built against. Held so the workspace
1686    /// can answer `joint_gradient_evaluation` from its own captured state; the
1687    /// designs are `Arc`-backed, so this clone is O(1) in the matrix data.
1688    pub(crate) specs: Vec<ParameterBlockSpec>,
1689    pub(crate) x_eta: Arc<Array2<f64>>,
1690    pub(crate) hessian_operator: Arc<RowCoeffOperator>,
1691}
1692
1693impl BinomialMeanWiggleHessianWorkspace {
1694    pub(crate) fn new(
1695        family: BinomialMeanWiggleFamily,
1696        block_states: Vec<ParameterBlockState>,
1697        specs: Vec<ParameterBlockSpec>,
1698        x_eta: Array2<f64>,
1699    ) -> Result<Self, String> {
1700        let x_eta = Arc::new(x_eta);
1701        let hessian_operator = family.bmw_static_hessian_operator(&block_states, x_eta.clone())?;
1702        Ok(Self {
1703            family,
1704            block_states,
1705            specs,
1706            x_eta,
1707            hessian_operator,
1708        })
1709    }
1710}
1711
1712impl ExactNewtonJointHessianWorkspace for BinomialMeanWiggleHessianWorkspace {
1713    /// The terminal constrained-posterior assembly places the ambient centre at
1714    /// `beta_unc = beta_hat - Sigma * grad(l_p)`, so it needs the exact
1715    /// likelihood score at the certified mode. This family answers
1716    /// `exact_newton_joint_gradient_evaluation` analytically, which makes the
1717    /// inner solve skip the `FamilyEvaluation` whose working sets would
1718    /// otherwise carry that score forward, leaving this workspace as the sole
1719    /// retained source at assembly. Without this method it exposes none and a
1720    /// converged constrained fit is refused outright (#2474).
1721    ///
1722    /// Reporting it here does not re-evaluate a live family: the workspace the
1723    /// assembly holds is built at the converged states (fresh, by the
1724    /// returned-mode curvature certificate) and owns its own family clone, and
1725    /// the delegate is the same working-set identity the inner joint-Newton RHS
1726    /// uses, so the score reported is the one the mode was certified with.
1727    fn joint_gradient_evaluation(
1728        &self,
1729    ) -> Result<Option<ExactNewtonJointGradientEvaluation>, String> {
1730        self.family
1731            .exact_newton_joint_gradient_evaluation(&self.block_states, &self.specs)
1732    }
1733
1734    fn warm_up_outer_caches_for_mode(
1735        &self,
1736        eval_mode: gam_problem::EvalMode,
1737    ) -> Result<(), String> {
1738        match eval_mode {
1739            gam_problem::EvalMode::ValueOnly
1740            | gam_problem::EvalMode::ValueAndGradient
1741            | gam_problem::EvalMode::ValueGradientHessian => Ok(()),
1742        }
1743    }
1744
1745    fn hessian_matvec_available(&self) -> bool {
1746        true
1747    }
1748
1749    fn hessian_matvec(&self, v: &Array1<f64>) -> Result<Option<Array1<f64>>, String> {
1750        Ok(Some(gam_problem::HyperOperator::mul_vec(
1751            self.hessian_operator.as_ref(),
1752            v,
1753        )))
1754    }
1755
1756    fn hessian_diagonal(&self) -> Result<Option<Array1<f64>>, String> {
1757        // The source resolver requires a finite diagonal alongside the HVP to
1758        // build an operator curvature source; `None` here made every
1759        // joint-workspace fit of this family die at the inner-solve boundary
1760        // with "supplied no inner-solve curvature source" (#2299 link-wiggle
1761        // gate). The static operator's diagonal is exact and O(n·(p_η+p_w)).
1762        Ok(Some(self.hessian_operator.diagonal()))
1763    }
1764
1765    fn directional_derivative(
1766        &self,
1767        d_beta_flat: &Array1<f64>,
1768    ) -> Result<Option<Array2<f64>>, String> {
1769        Ok(self
1770            .directional_derivative_operator(d_beta_flat)?
1771            .map(|operator| operator.to_dense()))
1772    }
1773
1774    fn directional_derivative_operator(
1775        &self,
1776        d_beta_flat: &Array1<f64>,
1777    ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
1778        self.family
1779            .bmw_directional_operator(&self.block_states, self.x_eta.clone(), d_beta_flat)
1780    }
1781
1782    fn second_directional_derivative(
1783        &self,
1784        d_beta_u_flat: &Array1<f64>,
1785        d_beta_v_flat: &Array1<f64>,
1786    ) -> Result<Option<Array2<f64>>, String> {
1787        Ok(self
1788            .second_directional_derivative_operator(d_beta_u_flat, d_beta_v_flat)?
1789            .map(|operator| operator.to_dense()))
1790    }
1791
1792    fn second_directional_derivative_operator(
1793        &self,
1794        d_beta_u: &Array1<f64>,
1795        d_beta_v: &Array1<f64>,
1796    ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
1797        self.family.bmw_second_directional_operator(
1798            &self.block_states,
1799            self.x_eta.clone(),
1800            d_beta_u,
1801            d_beta_v,
1802        )
1803    }
1804}
1805
1806impl CustomFamilyGenerative for BinomialMeanWiggleFamily {
1807    fn generativespec(
1808        &self,
1809        block_states: &[ParameterBlockState],
1810    ) -> Result<GenerativeSpec, String> {
1811        validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1812        let eta = &block_states[Self::BLOCK_ETA].eta;
1813        let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1814        if eta.len() != self.y.len() || etaw.len() != self.y.len() {
1815            return Err(GamlssError::DimensionMismatch {
1816                reason: "BinomialMeanWiggleFamily generative size mismatch".to_string(),
1817            }
1818            .into());
1819        }
1820        let mean = gamlss_rowwise_map_result(self.y.len(), |i| {
1821            let jet = inverse_link_jet_for_inverse_link(&self.link_kind, eta[i] + etaw[i])
1822                .map_err(|e| format!("fixed-link wiggle inverse-link evaluation failed: {e}"))?;
1823            Ok(jet.mu)
1824        })?;
1825        Ok(GenerativeSpec {
1826            mean,
1827            noise: NoiseModel::Bernoulli,
1828        })
1829    }
1830}
1831
1832#[cfg(test)]
1833mod exact_frozen_monotonicity_tests {
1834    use super::*;
1835
1836    fn frozen_family_and_wiggle_spec() -> (BinomialMeanWiggleFamily, ParameterBlockSpec) {
1837        let n = 4;
1838        let p = 3;
1839        let frozen = Array2::<f64>::zeros((n, p));
1840        let family = BinomialMeanWiggleFamily {
1841            y: Array1::zeros(n),
1842            weights: Array1::ones(n),
1843            link_kind: InverseLink::Standard(StandardLink::Logit),
1844            wiggle_knots: Array1::linspace(-1.0, 1.0, 8),
1845            wiggle_degree: 3,
1846            policy: gam_runtime::resource::ResourcePolicy::default_library(),
1847            frozen_warp_design: Some(Arc::new(frozen.clone())),
1848        };
1849        let spec = ParameterBlockSpec {
1850            name: "wiggle".to_string(),
1851            design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(frozen)),
1852            offset: Array1::zeros(n),
1853            penalties: vec![],
1854            nullspace_dims: vec![],
1855            initial_log_lambdas: Array1::zeros(0),
1856            initial_beta: None,
1857            gauge_priority: 100,
1858            jacobian_callback: None,
1859            stacked_design: None,
1860            stacked_offset: None,
1861        };
1862        (family, spec)
1863    }
1864
1865    #[test]
1866    fn frozen_warp_keeps_exact_nonnegative_i_spline_cone() {
1867        let (family, spec) = frozen_family_and_wiggle_spec();
1868        let constraints = family
1869            .block_linear_constraints(&[], BinomialMeanWiggleFamily::BLOCK_WIGGLE, &spec)
1870            .expect("frozen constraint construction")
1871            .expect("wiggle block must be constrained")
1872            .to_dense()
1873            .expect("wiggle cone is a small dense system");
1874        assert_eq!(constraints.a, Array2::<f64>::eye(3));
1875        assert_eq!(constraints.b, Array1::<f64>::zeros(3));
1876
1877        let solver_slop = Array1::from_vec(vec![
1878            -0.5 * crate::wiggle::MONOTONE_WIGGLE_ACTIVE_SET_TOL,
1879            0.2,
1880            0.0,
1881        ]);
1882        let projected = family
1883            .post_update_block_beta(
1884                &[],
1885                BinomialMeanWiggleFamily::BLOCK_WIGGLE,
1886                &spec,
1887                solver_slop,
1888            )
1889            .expect("active-set slop projects onto the exact cone");
1890        assert_eq!(projected[0], 0.0);
1891
1892        let material_violation = Array1::from_vec(vec![
1893            -2.0 * crate::wiggle::MONOTONE_WIGGLE_ACTIVE_SET_TOL,
1894            0.2,
1895            0.0,
1896        ]);
1897        assert!(
1898            family
1899                .post_update_block_beta(
1900                    &[],
1901                    BinomialMeanWiggleFamily::BLOCK_WIGGLE,
1902                    &spec,
1903                    material_violation,
1904                )
1905                .is_err(),
1906            "a material negative I-spline coefficient must be rejected"
1907        );
1908    }
1909}