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 = 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 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 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 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);
1246 let b2u = geom.basis_d2.dot(&uw); let b2v = geom.basis_d2.dot(&vw);
1248 let b3u = geom.basis_d3.dot(&uw); let b3v = geom.basis_d3.dot(&vw);
1250
1251 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);
1268
1269 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 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 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 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 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 coeff_eta[row] = second_directionalhessian_coeff_fromobjective_q_terms(
1367 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, );
1374
1375 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 let dc_b_u = m3 * q_u * a + m2 * a_u;
1383 let dc_b_v = m3 * q_v * a + m2 * a_v;
1384 let c_b_static = m2 * a;
1386 let d2_c_b1 = m3 * q_u * q_v + m2 * q_uv;
1388 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 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 let d2_h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
1408
1409 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 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 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 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 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}