1use 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 pub policy: gam_runtime::resource::ResourcePolicy,
19 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 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 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 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 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 fn joint_jeffreys_term_required(&self) -> bool {
690 false
691 }
692
693 fn coefficient_hessian_cost(&self, specs: &[ParameterBlockSpec]) -> u64 {
694 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 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 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 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 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 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 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 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 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 =
1011 BinomialMeanWiggleHessianWorkspace::new(self.clone(), block_states.to_vec(), x_eta)?;
1012 Ok(Some(Arc::new(workspace)))
1013 }
1014
1015 fn inner_coefficient_hessian_hvp_available(&self, specs: &[ParameterBlockSpec]) -> bool {
1016 self.dense_eta_design_fromspecs(specs).is_ok()
1017 }
1018
1019 fn exact_newton_joint_hessian_with_specs(
1020 &self,
1021 block_states: &[ParameterBlockState],
1022 specs: &[ParameterBlockSpec],
1023 ) -> Result<Option<Array2<f64>>, String> {
1024 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1025 let x_eta = self.dense_eta_design_fromspecs(specs)?;
1026 let eta = &block_states[Self::BLOCK_ETA].eta;
1027 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1028 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1029 let n = self.y.len();
1030 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1031 return Err(GamlssError::DimensionMismatch {
1032 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1033 }
1034 .into());
1035 }
1036 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1037 let p_eta = x_eta.ncols();
1038 let pw = geom.basis.ncols();
1039 let mut coeff_eta = Array1::<f64>::zeros(n);
1040 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
1041 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
1042 let mut coeff_ww = Array1::<f64>::zeros(n);
1043 for row in 0..n {
1044 let q = eta[row] + etaw[row];
1045 let (m1, m2, _) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1046 let a = geom.dq_dq0[row];
1047 let b = geom.d2q_dq02[row];
1048 coeff_eta[row] = hessian_coeff_fromobjective_q_terms(m1, m2, a, a, b);
1049 coeff_etaw_b[row] = m2 * a;
1050 coeff_etaw_d1[row] = m1;
1051 coeff_ww[row] = m2;
1052 }
1053 let h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
1054 let h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
1055 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?;
1056 let h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww)?;
1057 assert_eq!(h_eta_eta.nrows(), p_eta);
1058 assert_eq!(h_ww.nrows(), pw);
1059 Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
1060 &h_eta_eta, &h_eta_w, &h_ww,
1061 )))
1062 }
1063
1064 fn exact_newton_joint_hessian_directional_derivative_with_specs(
1065 &self,
1066 block_states: &[ParameterBlockState],
1067 specs: &[ParameterBlockSpec],
1068 d_beta_flat: &Array1<f64>,
1069 ) -> Result<Option<Array2<f64>>, String> {
1070 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1071 let x_eta = self.dense_eta_design_fromspecs(specs)?;
1072 let eta = &block_states[Self::BLOCK_ETA].eta;
1073 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1074 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1075 let n = self.y.len();
1076 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1077 return Err(GamlssError::DimensionMismatch {
1078 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1079 }
1080 .into());
1081 }
1082 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1083 let p_eta = x_eta.ncols();
1084 let pw = geom.basis.ncols();
1085 if d_beta_flat.len() != p_eta + pw {
1086 return Err(GamlssError::DimensionMismatch {
1087 reason: format!(
1088 "BinomialMeanWiggleFamily joint d_beta length mismatch: got {}, expected {}",
1089 d_beta_flat.len(),
1090 p_eta + pw
1091 ),
1092 }
1093 .into());
1094 }
1095 let u_eta = d_beta_flat.slice(s![0..p_eta]).to_owned();
1096 let uw = d_beta_flat.slice(s![p_eta..p_eta + pw]).to_owned();
1097 let xi = x_eta.dot(&u_eta);
1098 let phi = geom.basis.dot(&uw);
1099 let basis1_u = geom.basis_d1.dot(&uw);
1100 let basis2_u = geom.basis_d2.dot(&uw);
1101
1102 let mut coeff_eta = Array1::<f64>::zeros(n);
1103 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
1104 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
1105 let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
1106 let mut coeff_ww_bb = Array1::<f64>::zeros(n);
1107 let mut coeff_ww_db = Array1::<f64>::zeros(n);
1108 for row in 0..n {
1109 let q = eta[row] + etaw[row];
1110 let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1111 let a = geom.dq_dq0[row];
1112 let b = geom.d2q_dq02[row];
1113 let c = geom.d3q_dq03[row];
1114 let q_u = a * xi[row] + phi[row];
1115 let a_u = b * xi[row] + basis1_u[row];
1116 let b_u = c * xi[row] + basis2_u[row];
1117 coeff_eta[row] = directionalhessian_coeff_fromobjective_q_terms(
1118 m1, m2, m3, q_u, a, a, b, a_u, a_u, b_u,
1119 );
1120 coeff_etaw_b[row] = m3 * q_u * a + m2 * a_u;
1121 coeff_etaw_d1[row] = m2 * (a * xi[row] + q_u);
1122 coeff_etaw_d2[row] = m1 * xi[row];
1123 coeff_ww_bb[row] = m3 * q_u;
1124 coeff_ww_db[row] = m2 * xi[row];
1125 }
1126
1127 let d_h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
1128 let d_h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
1129 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?
1130 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d2, &geom.basis_d2)?;
1131 let a_ww = xt_diag_y_dense(&geom.basis_d1, &coeff_ww_db, &geom.basis)?;
1132 let d_h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww_bb)? + &a_ww + a_ww.t();
1133 Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
1134 &d_h_eta_eta,
1135 &d_h_eta_w,
1136 &d_h_ww,
1137 )))
1138 }
1139
1140 fn exact_newton_joint_hessian_second_directional_derivative_with_specs(
1198 &self,
1199 block_states: &[ParameterBlockState],
1200 specs: &[ParameterBlockSpec],
1201 d_beta_u_flat: &Array1<f64>,
1202 d_beta_v_flat: &Array1<f64>,
1203 ) -> Result<Option<Array2<f64>>, String> {
1204 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1205 let x_eta = self.dense_eta_design_fromspecs(specs)?;
1206 let eta = &block_states[Self::BLOCK_ETA].eta;
1207 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1208 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1209 let n = self.y.len();
1210 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1211 return Err(GamlssError::DimensionMismatch {
1212 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1213 }
1214 .into());
1215 }
1216 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1217 let p_eta = x_eta.ncols();
1218 let pw = geom.basis.ncols();
1219 let total = p_eta + pw;
1220 if d_beta_u_flat.len() != total || d_beta_v_flat.len() != total {
1221 return Err(GamlssError::DimensionMismatch { reason: format!(
1222 "BinomialMeanWiggleFamily joint second d_beta length mismatch: got {} and {}, expected {}",
1223 d_beta_u_flat.len(),
1224 d_beta_v_flat.len(),
1225 total
1226 ) }.into());
1227 }
1228
1229 let u_eta = d_beta_u_flat.slice(s![0..p_eta]).to_owned();
1231 let v_eta = d_beta_v_flat.slice(s![0..p_eta]).to_owned();
1232 let uw = d_beta_u_flat.slice(s![p_eta..total]).to_owned();
1233 let vw = d_beta_v_flat.slice(s![p_eta..total]).to_owned();
1234
1235 let xi_u = x_eta.dot(&u_eta); let xi_v = x_eta.dot(&v_eta); let phi_u = geom.basis.dot(&uw); let phi_v = geom.basis.dot(&vw); let b1u = geom.basis_d1.dot(&uw); let b1v = geom.basis_d1.dot(&vw);
1242 let b2u = geom.basis_d2.dot(&uw); let b2v = geom.basis_d2.dot(&vw);
1244 let b3u = geom.basis_d3.dot(&uw); let b3v = geom.basis_d3.dot(&vw);
1246
1247 let basis_u = scale_matrix_rows(&geom.basis_d1, &xi_u)?; let basis_v = scale_matrix_rows(&geom.basis_d1, &xi_v)?; let basis_uv = scale_matrix_rows(&geom.basis_d2, &(&xi_u * &xi_v))?; let mut coeff_eta = Array1::<f64>::zeros(n);
1264
1265 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
1300 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
1301 let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
1302 let mut coeff_etaw_d3 = Array1::<f64>::zeros(n);
1303
1304 let mut dw = Array1::<f64>::zeros(n);
1329 let mut dw_u = Array1::<f64>::zeros(n);
1330 let mut dw_v = Array1::<f64>::zeros(n);
1331 let mut dw_uv = Array1::<f64>::zeros(n);
1332
1333 for row in 0..n {
1334 let q = eta[row] + etaw[row];
1335 let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1336 let m4 = self.neglog_q_fourth_derivative(self.y[row], self.weights[row], q)?;
1337 let a = geom.dq_dq0[row];
1338 let b = geom.d2q_dq02[row];
1339 let c = geom.d3q_dq03[row];
1340 let d = geom.d4q_dq04[row];
1341
1342 let q_u = a * xi_u[row] + phi_u[row];
1344 let a_u = b * xi_u[row] + b1u[row];
1345 let b_u = c * xi_u[row] + b2u[row];
1346
1347 let q_v = a * xi_v[row] + phi_v[row];
1349 let a_v = b * xi_v[row] + b1v[row];
1350 let b_v = c * xi_v[row] + b2v[row];
1351
1352 let q_uv = b * xi_u[row] * xi_v[row] + b1u[row] * xi_v[row] + b1v[row] * xi_u[row];
1354 let a_uv = c * xi_u[row] * xi_v[row] + b2u[row] * xi_v[row] + b2v[row] * xi_u[row];
1355 let b_uv = d * xi_u[row] * xi_v[row] + b3u[row] * xi_v[row] + b3v[row] * xi_u[row];
1356
1357 coeff_eta[row] = second_directionalhessian_coeff_fromobjective_q_terms(
1363 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, b_uv, );
1370
1371 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;
1377 let dc_b_u = m3 * q_u * a + m2 * a_u;
1379 let dc_b_v = m3 * q_v * a + m2 * a_v;
1380 let c_b_static = m2 * a;
1382 let d2_c_b1 = m3 * q_u * q_v + m2 * q_uv;
1384 let dc_b1_u = m2 * q_u;
1386 let dc_b1_v = m2 * q_v;
1387
1388 coeff_etaw_b[row] = d2_c_b;
1389 coeff_etaw_d1[row] = dc_b_u * xi_v[row] + dc_b_v * xi_u[row] + d2_c_b1;
1390 coeff_etaw_d2[row] =
1391 c_b_static * xi_u[row] * xi_v[row] + dc_b1_u * xi_v[row] + dc_b1_v * xi_u[row];
1392 coeff_etaw_d3[row] = m1 * xi_u[row] * xi_v[row];
1393
1394 dw[row] = m2;
1397 dw_u[row] = m3 * q_u;
1398 dw_v[row] = m3 * q_v;
1399 dw_uv[row] = m4 * q_u * q_v + m3 * q_uv;
1400 }
1401
1402 let d2_h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
1404
1405 let d2_h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
1412 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?
1413 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d2, &geom.basis_d2)?
1414 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d3, &geom.basis_d3)?;
1415
1416 let a_ab = xt_diag_y_dense(&basis_uv, &dw, &geom.basis)?;
1427 let a_ij = xt_diag_y_dense(&basis_u, &dw, &basis_v)?;
1428 let a_iwj = xt_diag_y_dense(&basis_u, &dw_v, &geom.basis)?;
1429 let a_jwi = xt_diag_y_dense(&basis_v, &dw_u, &geom.basis)?;
1430 let d2_h_ww = &a_ab
1431 + &a_ab.t()
1432 + &a_ij
1433 + a_ij.t()
1434 + &a_iwj
1435 + a_iwj.t()
1436 + &a_jwi
1437 + a_jwi.t()
1438 + &xt_diag_x_dense(&geom.basis, &dw_uv)?;
1439
1440 Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
1441 &d2_h_eta_eta,
1442 &d2_h_eta_w,
1443 &d2_h_ww,
1444 )))
1445 }
1446
1447 fn exact_newton_joint_psi_terms(
1448 &self,
1449 block_states: &[ParameterBlockState],
1450 specs: &[ParameterBlockSpec],
1451 hyper_layout: &crate::custom_family::CustomFamilyHyperLayout,
1452 psi_index: usize,
1453 ) -> Result<Option<gam_problem::ExactNewtonJointPsiTerms>, String> {
1454 if hyper_layout.family_axis_count() != 0 {
1455 return Err("BinomialMeanWiggleFamily does not declare family-owned hyper axes"
1456 .to_string());
1457 }
1458 let derivative_blocks = hyper_layout.design_derivative_blocks();
1459 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1460 if derivative_blocks.len() != 2 {
1461 return Err(GamlssError::DimensionMismatch { reason: format!(
1462 "BinomialMeanWiggleFamily joint psi terms expect 2 derivative block lists, got {}",
1463 derivative_blocks.len()
1464 ) }.into());
1465 }
1466 let x_eta = self.dense_eta_design_fromspecs(specs)?;
1467 let eta = &block_states[Self::BLOCK_ETA].eta;
1468 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1469 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1470 let n = self.y.len();
1471 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1472 return Err(GamlssError::DimensionMismatch {
1473 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1474 }
1475 .into());
1476 }
1477 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1478 let p_eta = x_eta.ncols();
1479 let pw = geom.basis.ncols();
1480 let implicit_dir =
1481 self.exact_newton_joint_psi_action(block_states, derivative_blocks, psi_index, p_eta)?;
1482 let dense_dir = if implicit_dir.is_none() {
1483 self.exact_newton_joint_psi_direction(
1484 block_states,
1485 derivative_blocks,
1486 psi_index,
1487 &x_eta,
1488 )?
1489 } else {
1490 None
1491 };
1492 let z_eta_psi = if let Some((_, ref z_eta_psi)) = implicit_dir {
1493 z_eta_psi
1494 } else if let Some(ref dir_a) = dense_dir {
1495 &dir_a.z_eta_psi
1496 } else {
1497 return Ok(None);
1498 };
1499
1500 let mut objective_psi = 0.0;
1501 let mut score_eta_xa = Array1::<f64>::zeros(n);
1502 let mut score_eta_x = Array1::<f64>::zeros(n);
1503 let mut score_w_b = Array1::<f64>::zeros(n);
1504 let mut score_w_d1 = Array1::<f64>::zeros(n);
1505
1506 let mut coeff_eta_eta_xx = Array1::<f64>::zeros(n);
1507 let mut coeff_eta_eta_xa_x = Array1::<f64>::zeros(n);
1508 let mut coeff_eta_w_xa_b = Array1::<f64>::zeros(n);
1509 let mut coeff_eta_w_x_b = Array1::<f64>::zeros(n);
1510 let mut coeff_eta_w_x_d1 = Array1::<f64>::zeros(n);
1511 let mut coeff_eta_w_xa_d1 = Array1::<f64>::zeros(n);
1512 let mut coeff_eta_w_x_d2 = Array1::<f64>::zeros(n);
1513 let mut coeff_ww_bb = Array1::<f64>::zeros(n);
1514 let mut coeff_ww_db = Array1::<f64>::zeros(n);
1515
1516 for row in 0..n {
1517 let q = eta[row] + etaw[row];
1518 let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1519 let z_a = z_eta_psi[row];
1520 let a = geom.dq_dq0[row];
1521 let b = geom.d2q_dq02[row];
1522 let c = geom.d3q_dq03[row];
1523 let q_a = a * z_a;
1524
1525 objective_psi += m1 * q_a;
1526
1527 score_eta_xa[row] = m1 * a;
1528 score_eta_x[row] = m2 * q_a * a + m1 * b * z_a;
1529 score_w_b[row] = m2 * q_a;
1530 score_w_d1[row] = m1 * z_a;
1531
1532 coeff_eta_eta_xx[row] =
1533 m3 * q_a * a * a + m2 * (2.0 * a * b * z_a + q_a * b) + m1 * c * z_a;
1534 coeff_eta_eta_xa_x[row] = m2 * a * a + m1 * b;
1535 coeff_eta_w_xa_b[row] = m2 * a;
1536 coeff_eta_w_x_b[row] = m3 * q_a * a + m2 * b * z_a;
1537 coeff_eta_w_x_d1[row] = m2 * (a * z_a + q_a);
1538 coeff_eta_w_xa_d1[row] = m1;
1539 coeff_eta_w_x_d2[row] = m1 * z_a;
1540 coeff_ww_bb[row] = m3 * q_a;
1541 coeff_ww_db[row] = m2 * z_a;
1542 }
1543
1544 let score_w = gam_linalg::faer_ndarray::fast_atv(&geom.basis, &score_w_b)
1545 + gam_linalg::faer_ndarray::fast_atv(&geom.basis_d1, &score_w_d1);
1546
1547 if let Some((action, _)) = implicit_dir {
1548 let score_eta = action.transpose_mul(score_eta_xa.view())
1549 + gam_linalg::faer_ndarray::fast_atv(x_eta.as_ref(), &score_eta_x);
1550 let score_psi = binomial_pack_mean_wiggle_joint_score(&score_eta, &score_w);
1551 let x_eta_arc = shared_dense_arc(x_eta.as_ref());
1552 let basis_arc = Arc::new(geom.basis.clone());
1553 let basis_d1_arc = Arc::new(geom.basis_d1.clone());
1554 let basis_d2_arc = Arc::new(geom.basis_d2.clone());
1555 let zeros = Array1::<f64>::zeros(n);
1556 let operator = CustomFamilyJointPsiOperator::new(
1557 p_eta + pw,
1558 vec![
1559 CustomFamilyJointDesignChannel::new(
1560 0..p_eta,
1561 Arc::clone(&x_eta_arc),
1562 Some(action),
1563 ),
1564 CustomFamilyJointDesignChannel::new(
1565 p_eta..p_eta + pw,
1566 Arc::clone(&basis_arc),
1567 None,
1568 ),
1569 CustomFamilyJointDesignChannel::new(
1570 p_eta..p_eta + pw,
1571 Arc::clone(&basis_d1_arc),
1572 None,
1573 ),
1574 CustomFamilyJointDesignChannel::new(
1575 p_eta..p_eta + pw,
1576 Arc::clone(&basis_d2_arc),
1577 None,
1578 ),
1579 ],
1580 vec![
1581 CustomFamilyJointDesignPairContribution::new(
1582 0,
1583 0,
1584 coeff_eta_eta_xa_x.clone(),
1585 coeff_eta_eta_xx.clone(),
1586 ),
1587 CustomFamilyJointDesignPairContribution::new(
1588 0,
1589 1,
1590 coeff_eta_w_xa_b.clone(),
1591 coeff_eta_w_x_b.clone(),
1592 ),
1593 CustomFamilyJointDesignPairContribution::new(
1594 1,
1595 0,
1596 coeff_eta_w_xa_b.clone(),
1597 coeff_eta_w_x_b.clone(),
1598 ),
1599 CustomFamilyJointDesignPairContribution::new(
1600 0,
1601 2,
1602 coeff_eta_w_xa_d1.clone(),
1603 coeff_eta_w_x_d1.clone(),
1604 ),
1605 CustomFamilyJointDesignPairContribution::new(
1606 2,
1607 0,
1608 coeff_eta_w_xa_d1.clone(),
1609 coeff_eta_w_x_d1.clone(),
1610 ),
1611 CustomFamilyJointDesignPairContribution::new(
1612 0,
1613 3,
1614 zeros.clone(),
1615 coeff_eta_w_x_d2.clone(),
1616 ),
1617 CustomFamilyJointDesignPairContribution::new(
1618 3,
1619 0,
1620 zeros.clone(),
1621 coeff_eta_w_x_d2.clone(),
1622 ),
1623 CustomFamilyJointDesignPairContribution::new(
1624 1,
1625 1,
1626 zeros.clone(),
1627 coeff_ww_bb.clone(),
1628 ),
1629 CustomFamilyJointDesignPairContribution::new(
1630 2,
1631 1,
1632 zeros.clone(),
1633 coeff_ww_db.clone(),
1634 ),
1635 CustomFamilyJointDesignPairContribution::new(1, 2, zeros, coeff_ww_db.clone()),
1636 ],
1637 );
1638 return Ok(Some(gam_problem::ExactNewtonJointPsiTerms {
1639 objective_psi,
1640 score_psi,
1641 hessian_psi: Array2::zeros((0, 0)),
1642 hessian_psi_operator: Some(std::sync::Arc::new(operator)),
1643 }));
1644 }
1645
1646 let dir_a =
1647 dense_dir.expect("dense psi direction should exist when implicit direction is absent");
1648 let x_eta_psi = dir_a
1649 .x_eta_psi
1650 .as_ref()
1651 .expect("dense eta psi design should exist when implicit direction is absent");
1652 let score_psi = binomial_pack_mean_wiggle_joint_score(
1653 &(gam_linalg::faer_ndarray::fast_atv(x_eta_psi, &score_eta_xa)
1654 + gam_linalg::faer_ndarray::fast_atv(x_eta.as_ref(), &score_eta_x)),
1655 &score_w,
1656 );
1657 let a_eta_eta = xt_diag_y_dense(x_eta_psi, &coeff_eta_eta_xa_x, &x_eta)?;
1658 let h_eta_eta = &a_eta_eta + &a_eta_eta.t() + &xt_diag_x_dense(&x_eta, &coeff_eta_eta_xx)?;
1659 let h_eta_w = xt_diag_y_dense(x_eta_psi, &coeff_eta_w_xa_b, &geom.basis)?
1660 + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_b, &geom.basis)?
1661 + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_d1, &geom.basis_d1)?
1662 + &xt_diag_y_dense(x_eta_psi, &coeff_eta_w_xa_d1, &geom.basis_d1)?
1663 + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_d2, &geom.basis_d2)?;
1664 let a_ww = xt_diag_y_dense(&geom.basis_d1, &coeff_ww_db, &geom.basis)?;
1665 let h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww_bb)? + &a_ww + a_ww.t();
1666
1667 Ok(Some(gam_problem::ExactNewtonJointPsiTerms {
1668 objective_psi,
1669 score_psi,
1670 hessian_psi: binomial_pack_mean_wiggle_joint_symmetrichessian(
1671 &h_eta_eta, &h_eta_w, &h_ww,
1672 ),
1673 hessian_psi_operator: None,
1674 }))
1675 }
1676}
1677
1678pub(crate) struct BinomialMeanWiggleHessianWorkspace {
1679 pub(crate) family: BinomialMeanWiggleFamily,
1680 pub(crate) block_states: Vec<ParameterBlockState>,
1681 pub(crate) x_eta: Arc<Array2<f64>>,
1682 pub(crate) hessian_operator: Arc<RowCoeffOperator>,
1683}
1684
1685impl BinomialMeanWiggleHessianWorkspace {
1686 pub(crate) fn new(
1687 family: BinomialMeanWiggleFamily,
1688 block_states: Vec<ParameterBlockState>,
1689 x_eta: Array2<f64>,
1690 ) -> Result<Self, String> {
1691 let x_eta = Arc::new(x_eta);
1692 let hessian_operator = family.bmw_static_hessian_operator(&block_states, x_eta.clone())?;
1693 Ok(Self {
1694 family,
1695 block_states,
1696 x_eta,
1697 hessian_operator,
1698 })
1699 }
1700}
1701
1702impl ExactNewtonJointHessianWorkspace for BinomialMeanWiggleHessianWorkspace {
1703 fn warm_up_outer_caches_for_mode(
1704 &self,
1705 eval_mode: gam_problem::EvalMode,
1706 ) -> Result<(), String> {
1707 match eval_mode {
1708 gam_problem::EvalMode::ValueOnly
1709 | gam_problem::EvalMode::ValueAndGradient
1710 | gam_problem::EvalMode::ValueGradientHessian => Ok(()),
1711 }
1712 }
1713
1714 fn hessian_matvec_available(&self) -> bool {
1715 true
1716 }
1717
1718 fn hessian_matvec(&self, v: &Array1<f64>) -> Result<Option<Array1<f64>>, String> {
1719 Ok(Some(gam_problem::HyperOperator::mul_vec(
1720 self.hessian_operator.as_ref(),
1721 v,
1722 )))
1723 }
1724
1725 fn hessian_diagonal(&self) -> Result<Option<Array1<f64>>, String> {
1726 Ok(Some(self.hessian_operator.diagonal()))
1732 }
1733
1734 fn directional_derivative(
1735 &self,
1736 d_beta_flat: &Array1<f64>,
1737 ) -> Result<Option<Array2<f64>>, String> {
1738 Ok(self
1739 .directional_derivative_operator(d_beta_flat)?
1740 .map(|operator| operator.to_dense()))
1741 }
1742
1743 fn directional_derivative_operator(
1744 &self,
1745 d_beta_flat: &Array1<f64>,
1746 ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
1747 self.family
1748 .bmw_directional_operator(&self.block_states, self.x_eta.clone(), d_beta_flat)
1749 }
1750
1751 fn second_directional_derivative(
1752 &self,
1753 d_beta_u_flat: &Array1<f64>,
1754 d_beta_v_flat: &Array1<f64>,
1755 ) -> Result<Option<Array2<f64>>, String> {
1756 Ok(self
1757 .second_directional_derivative_operator(d_beta_u_flat, d_beta_v_flat)?
1758 .map(|operator| operator.to_dense()))
1759 }
1760
1761 fn second_directional_derivative_operator(
1762 &self,
1763 d_beta_u: &Array1<f64>,
1764 d_beta_v: &Array1<f64>,
1765 ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
1766 self.family.bmw_second_directional_operator(
1767 &self.block_states,
1768 self.x_eta.clone(),
1769 d_beta_u,
1770 d_beta_v,
1771 )
1772 }
1773}
1774
1775impl CustomFamilyGenerative for BinomialMeanWiggleFamily {
1776 fn generativespec(
1777 &self,
1778 block_states: &[ParameterBlockState],
1779 ) -> Result<GenerativeSpec, String> {
1780 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1781 let eta = &block_states[Self::BLOCK_ETA].eta;
1782 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1783 if eta.len() != self.y.len() || etaw.len() != self.y.len() {
1784 return Err(GamlssError::DimensionMismatch {
1785 reason: "BinomialMeanWiggleFamily generative size mismatch".to_string(),
1786 }
1787 .into());
1788 }
1789 let mean = gamlss_rowwise_map_result(self.y.len(), |i| {
1790 let jet = inverse_link_jet_for_inverse_link(&self.link_kind, eta[i] + etaw[i])
1791 .map_err(|e| format!("fixed-link wiggle inverse-link evaluation failed: {e}"))?;
1792 Ok(jet.mu)
1793 })?;
1794 Ok(GenerativeSpec {
1795 mean,
1796 noise: NoiseModel::Bernoulli,
1797 })
1798 }
1799}
1800
1801#[cfg(test)]
1802mod exact_frozen_monotonicity_tests {
1803 use super::*;
1804
1805 fn frozen_family_and_wiggle_spec() -> (BinomialMeanWiggleFamily, ParameterBlockSpec) {
1806 let n = 4;
1807 let p = 3;
1808 let frozen = Array2::<f64>::zeros((n, p));
1809 let family = BinomialMeanWiggleFamily {
1810 y: Array1::zeros(n),
1811 weights: Array1::ones(n),
1812 link_kind: InverseLink::Standard(StandardLink::Logit),
1813 wiggle_knots: Array1::linspace(-1.0, 1.0, 8),
1814 wiggle_degree: 3,
1815 policy: gam_runtime::resource::ResourcePolicy::default_library(),
1816 frozen_warp_design: Some(Arc::new(frozen.clone())),
1817 };
1818 let spec = ParameterBlockSpec {
1819 name: "wiggle".to_string(),
1820 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(frozen)),
1821 offset: Array1::zeros(n),
1822 penalties: vec![],
1823 nullspace_dims: vec![],
1824 initial_log_lambdas: Array1::zeros(0),
1825 initial_beta: None,
1826 gauge_priority: 100,
1827 jacobian_callback: None,
1828 stacked_design: None,
1829 stacked_offset: None,
1830 };
1831 (family, spec)
1832 }
1833
1834 #[test]
1835 fn frozen_warp_keeps_exact_nonnegative_i_spline_cone() {
1836 let (family, spec) = frozen_family_and_wiggle_spec();
1837 let constraints = family
1838 .block_linear_constraints(&[], BinomialMeanWiggleFamily::BLOCK_WIGGLE, &spec)
1839 .expect("frozen constraint construction")
1840 .expect("wiggle block must be constrained")
1841 .to_dense()
1842 .expect("wiggle cone is a small dense system");
1843 assert_eq!(constraints.a, Array2::<f64>::eye(3));
1844 assert_eq!(constraints.b, Array1::<f64>::zeros(3));
1845
1846 let solver_slop = Array1::from_vec(vec![
1847 -0.5 * crate::wiggle::MONOTONE_WIGGLE_ACTIVE_SET_TOL,
1848 0.2,
1849 0.0,
1850 ]);
1851 let projected = family
1852 .post_update_block_beta(
1853 &[],
1854 BinomialMeanWiggleFamily::BLOCK_WIGGLE,
1855 &spec,
1856 solver_slop,
1857 )
1858 .expect("active-set slop projects onto the exact cone");
1859 assert_eq!(projected[0], 0.0);
1860
1861 let material_violation = Array1::from_vec(vec![
1862 -2.0 * crate::wiggle::MONOTONE_WIGGLE_ACTIVE_SET_TOL,
1863 0.2,
1864 0.0,
1865 ]);
1866 assert!(
1867 family
1868 .post_update_block_beta(
1869 &[],
1870 BinomialMeanWiggleFamily::BLOCK_WIGGLE,
1871 &spec,
1872 material_violation,
1873 )
1874 .is_err(),
1875 "a material negative I-spline coefficient must be rejected"
1876 );
1877 }
1878}