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<LinearInequalityConstraints>, 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 ll = 0.0;
768 let mut z_eta = Array1::<f64>::zeros(n);
769 let mut w_eta = Array1::<f64>::zeros(n);
770 let mut z_wiggle = Array1::<f64>::zeros(n);
771 let mut w_wiggle = Array1::<f64>::zeros(n);
772 for i in 0..n {
773 let q = eta[i] + etaw[i];
774 let (mu_q, d1_q) = inverse_link_mu_d1_for_inverse_link(&self.link_kind, q)
775 .map_err(|e| format!("fixed-link wiggle inverse-link evaluation failed: {e}"))?;
776 let yi = self.y[i];
777 let wi = self.weights[i];
778 ll += binomial_location_scale_log_likelihood(yi, wi, q, &self.link_kind, mu_q)?;
779
780 let mu = mu_q.clamp(1e-12, 1.0 - 1e-12);
781 let var = (mu * (1.0 - mu)).max(MIN_PROB);
782 let dmu_deta = d1_q * dq_dq0[i];
783 let dmu_dw = d1_q;
784 if wi == 0.0 || !var.is_finite() {
785 z_eta[i] = eta[i];
786 z_wiggle[i] = etaw[i];
787 continue;
788 }
789
790 if dmu_deta.is_finite() {
791 w_eta[i] = floor_positiveweight(wi * (dmu_deta * dmu_deta / var), MIN_WEIGHT);
792 z_eta[i] = eta[i] + (yi - mu) / signedwith_floor(dmu_deta, MIN_DERIV);
793 } else {
794 z_eta[i] = eta[i];
795 }
796
797 if dmu_dw.is_finite() {
798 w_wiggle[i] = floor_positiveweight(wi * (dmu_dw * dmu_dw / var), MIN_WEIGHT);
799 z_wiggle[i] = etaw[i] + (yi - mu) / signedwith_floor(dmu_dw, MIN_DERIV);
800 } else {
801 z_wiggle[i] = etaw[i];
802 }
803 }
804
805 Ok(FamilyEvaluation {
806 log_likelihood: ll,
807 blockworking_sets: vec![
808 BlockWorkingSet::diagonal_checked(z_eta, w_eta)?,
809 BlockWorkingSet::diagonal_checked(z_wiggle, w_wiggle)?,
810 ],
811 })
812 }
813
814 fn exact_newton_joint_gradient_evaluation(
815 &self,
816 block_states: &[ParameterBlockState],
817 specs: &[ParameterBlockSpec],
818 ) -> Result<Option<ExactNewtonJointGradientEvaluation>, String> {
819 let eval = self.evaluate(block_states)?;
824 gamlss_joint_gradient_from_working_sets(&eval, specs, block_states).map(Some)
825 }
826
827 fn block_geometry(
828 &self,
829 block_states: &[ParameterBlockState],
830 spec: &ParameterBlockSpec,
831 ) -> Result<(DesignMatrix, Array1<f64>), String> {
832 if spec.name != "wiggle" {
833 return Ok((spec.design.clone(), spec.offset.clone()));
834 }
835 if block_states.is_empty() {
836 return Err(GamlssError::UnsupportedConfiguration {
837 reason: "wiggle geometry requires eta block".to_string(),
838 }
839 .into());
840 }
841 let eta = &block_states[Self::BLOCK_ETA].eta;
842 if eta.len() != self.y.len() {
843 return Err(GamlssError::DimensionMismatch {
844 reason: "BinomialMeanWiggleFamily eta size mismatch".to_string(),
845 }
846 .into());
847 }
848 let x = match self.frozen_warp_design.as_ref() {
854 Some(frozen) => frozen.as_ref().clone(),
855 None => self.wiggle_design(eta.view())?,
856 };
857 if x.ncols() != spec.design.ncols() {
858 return Err(GamlssError::DimensionMismatch {
859 reason: format!(
860 "dynamic wiggle design col mismatch: got {}, expected {}",
861 x.ncols(),
862 spec.design.ncols()
863 ),
864 }
865 .into());
866 }
867 let nrows = x.nrows();
868 Ok((
869 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(x)),
870 Array1::zeros(nrows),
871 ))
872 }
873
874 fn block_geometry_is_dynamic(&self) -> bool {
875 true
876 }
877
878 fn exact_newton_joint_hessian_workspace(
879 &self,
880 block_states: &[ParameterBlockState],
881 specs: &[ParameterBlockSpec],
882 ) -> Result<Option<Arc<dyn ExactNewtonJointHessianWorkspace>>, String> {
883 let x_eta = self.dense_eta_design_fromspecs(specs)?.into_owned();
884 let workspace =
885 BinomialMeanWiggleHessianWorkspace::new(self.clone(), block_states.to_vec(), x_eta)?;
886 Ok(Some(Arc::new(workspace)))
887 }
888
889 fn inner_coefficient_hessian_hvp_available(&self, specs: &[ParameterBlockSpec]) -> bool {
890 self.dense_eta_design_fromspecs(specs).is_ok()
891 }
892
893 fn exact_newton_joint_hessian_with_specs(
894 &self,
895 block_states: &[ParameterBlockState],
896 specs: &[ParameterBlockSpec],
897 ) -> Result<Option<Array2<f64>>, String> {
898 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
899 let x_eta = self.dense_eta_design_fromspecs(specs)?;
900 let eta = &block_states[Self::BLOCK_ETA].eta;
901 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
902 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
903 let n = self.y.len();
904 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
905 return Err(GamlssError::DimensionMismatch {
906 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
907 }
908 .into());
909 }
910 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
911 let p_eta = x_eta.ncols();
912 let pw = geom.basis.ncols();
913 let mut coeff_eta = Array1::<f64>::zeros(n);
914 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
915 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
916 let mut coeff_ww = Array1::<f64>::zeros(n);
917 for row in 0..n {
918 let q = eta[row] + etaw[row];
919 let (m1, m2, _) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
920 let a = geom.dq_dq0[row];
921 let b = geom.d2q_dq02[row];
922 coeff_eta[row] = hessian_coeff_fromobjective_q_terms(m1, m2, a, a, b);
923 coeff_etaw_b[row] = m2 * a;
924 coeff_etaw_d1[row] = m1;
925 coeff_ww[row] = m2;
926 }
927 let h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
928 let h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
929 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?;
930 let h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww)?;
931 assert_eq!(h_eta_eta.nrows(), p_eta);
932 assert_eq!(h_ww.nrows(), pw);
933 Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
934 &h_eta_eta, &h_eta_w, &h_ww,
935 )))
936 }
937
938 fn exact_newton_joint_hessian_directional_derivative_with_specs(
939 &self,
940 block_states: &[ParameterBlockState],
941 specs: &[ParameterBlockSpec],
942 d_beta_flat: &Array1<f64>,
943 ) -> Result<Option<Array2<f64>>, String> {
944 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
945 let x_eta = self.dense_eta_design_fromspecs(specs)?;
946 let eta = &block_states[Self::BLOCK_ETA].eta;
947 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
948 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
949 let n = self.y.len();
950 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
951 return Err(GamlssError::DimensionMismatch {
952 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
953 }
954 .into());
955 }
956 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
957 let p_eta = x_eta.ncols();
958 let pw = geom.basis.ncols();
959 if d_beta_flat.len() != p_eta + pw {
960 return Err(GamlssError::DimensionMismatch {
961 reason: format!(
962 "BinomialMeanWiggleFamily joint d_beta length mismatch: got {}, expected {}",
963 d_beta_flat.len(),
964 p_eta + pw
965 ),
966 }
967 .into());
968 }
969 let u_eta = d_beta_flat.slice(s![0..p_eta]).to_owned();
970 let uw = d_beta_flat.slice(s![p_eta..p_eta + pw]).to_owned();
971 let xi = x_eta.dot(&u_eta);
972 let phi = geom.basis.dot(&uw);
973 let basis1_u = geom.basis_d1.dot(&uw);
974 let basis2_u = geom.basis_d2.dot(&uw);
975
976 let mut coeff_eta = Array1::<f64>::zeros(n);
977 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
978 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
979 let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
980 let mut coeff_ww_bb = Array1::<f64>::zeros(n);
981 let mut coeff_ww_db = Array1::<f64>::zeros(n);
982 for row in 0..n {
983 let q = eta[row] + etaw[row];
984 let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
985 let a = geom.dq_dq0[row];
986 let b = geom.d2q_dq02[row];
987 let c = geom.d3q_dq03[row];
988 let q_u = a * xi[row] + phi[row];
989 let a_u = b * xi[row] + basis1_u[row];
990 let b_u = c * xi[row] + basis2_u[row];
991 coeff_eta[row] = directionalhessian_coeff_fromobjective_q_terms(
992 m1, m2, m3, q_u, a, a, b, a_u, a_u, b_u,
993 );
994 coeff_etaw_b[row] = m3 * q_u * a + m2 * a_u;
995 coeff_etaw_d1[row] = m2 * (a * xi[row] + q_u);
996 coeff_etaw_d2[row] = m1 * xi[row];
997 coeff_ww_bb[row] = m3 * q_u;
998 coeff_ww_db[row] = m2 * xi[row];
999 }
1000
1001 let d_h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
1002 let d_h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
1003 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?
1004 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d2, &geom.basis_d2)?;
1005 let a_ww = xt_diag_y_dense(&geom.basis_d1, &coeff_ww_db, &geom.basis)?;
1006 let d_h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww_bb)? + &a_ww + a_ww.t();
1007 Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
1008 &d_h_eta_eta,
1009 &d_h_eta_w,
1010 &d_h_ww,
1011 )))
1012 }
1013
1014 fn exact_newton_joint_hessian_second_directional_derivative_with_specs(
1072 &self,
1073 block_states: &[ParameterBlockState],
1074 specs: &[ParameterBlockSpec],
1075 d_beta_u_flat: &Array1<f64>,
1076 d_beta_v_flat: &Array1<f64>,
1077 ) -> Result<Option<Array2<f64>>, String> {
1078 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1079 let x_eta = self.dense_eta_design_fromspecs(specs)?;
1080 let eta = &block_states[Self::BLOCK_ETA].eta;
1081 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1082 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1083 let n = self.y.len();
1084 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1085 return Err(GamlssError::DimensionMismatch {
1086 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1087 }
1088 .into());
1089 }
1090 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1091 let p_eta = x_eta.ncols();
1092 let pw = geom.basis.ncols();
1093 let total = p_eta + pw;
1094 if d_beta_u_flat.len() != total || d_beta_v_flat.len() != total {
1095 return Err(GamlssError::DimensionMismatch { reason: format!(
1096 "BinomialMeanWiggleFamily joint second d_beta length mismatch: got {} and {}, expected {}",
1097 d_beta_u_flat.len(),
1098 d_beta_v_flat.len(),
1099 total
1100 ) }.into());
1101 }
1102
1103 let u_eta = d_beta_u_flat.slice(s![0..p_eta]).to_owned();
1105 let v_eta = d_beta_v_flat.slice(s![0..p_eta]).to_owned();
1106 let uw = d_beta_u_flat.slice(s![p_eta..total]).to_owned();
1107 let vw = d_beta_v_flat.slice(s![p_eta..total]).to_owned();
1108
1109 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);
1116 let b2u = geom.basis_d2.dot(&uw); let b2v = geom.basis_d2.dot(&vw);
1118 let b3u = geom.basis_d3.dot(&uw); let b3v = geom.basis_d3.dot(&vw);
1120
1121 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);
1138
1139 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
1174 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
1175 let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
1176 let mut coeff_etaw_d3 = Array1::<f64>::zeros(n);
1177
1178 let mut dw = Array1::<f64>::zeros(n);
1203 let mut dw_u = Array1::<f64>::zeros(n);
1204 let mut dw_v = Array1::<f64>::zeros(n);
1205 let mut dw_uv = Array1::<f64>::zeros(n);
1206
1207 for row in 0..n {
1208 let q = eta[row] + etaw[row];
1209 let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1210 let m4 = self.neglog_q_fourth_derivative(self.y[row], self.weights[row], q)?;
1211 let a = geom.dq_dq0[row];
1212 let b = geom.d2q_dq02[row];
1213 let c = geom.d3q_dq03[row];
1214 let d = geom.d4q_dq04[row];
1215
1216 let q_u = a * xi_u[row] + phi_u[row];
1218 let a_u = b * xi_u[row] + b1u[row];
1219 let b_u = c * xi_u[row] + b2u[row];
1220
1221 let q_v = a * xi_v[row] + phi_v[row];
1223 let a_v = b * xi_v[row] + b1v[row];
1224 let b_v = c * xi_v[row] + b2v[row];
1225
1226 let q_uv = b * xi_u[row] * xi_v[row] + b1u[row] * xi_v[row] + b1v[row] * xi_u[row];
1228 let a_uv = c * xi_u[row] * xi_v[row] + b2u[row] * xi_v[row] + b2v[row] * xi_u[row];
1229 let b_uv = d * xi_u[row] * xi_v[row] + b3u[row] * xi_v[row] + b3v[row] * xi_u[row];
1230
1231 coeff_eta[row] = second_directionalhessian_coeff_fromobjective_q_terms(
1237 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, );
1244
1245 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;
1251 let dc_b_u = m3 * q_u * a + m2 * a_u;
1253 let dc_b_v = m3 * q_v * a + m2 * a_v;
1254 let c_b_static = m2 * a;
1256 let d2_c_b1 = m3 * q_u * q_v + m2 * q_uv;
1258 let dc_b1_u = m2 * q_u;
1260 let dc_b1_v = m2 * q_v;
1261
1262 coeff_etaw_b[row] = d2_c_b;
1263 coeff_etaw_d1[row] = dc_b_u * xi_v[row] + dc_b_v * xi_u[row] + d2_c_b1;
1264 coeff_etaw_d2[row] =
1265 c_b_static * xi_u[row] * xi_v[row] + dc_b1_u * xi_v[row] + dc_b1_v * xi_u[row];
1266 coeff_etaw_d3[row] = m1 * xi_u[row] * xi_v[row];
1267
1268 dw[row] = m2;
1271 dw_u[row] = m3 * q_u;
1272 dw_v[row] = m3 * q_v;
1273 dw_uv[row] = m4 * q_u * q_v + m3 * q_uv;
1274 }
1275
1276 let d2_h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
1278
1279 let d2_h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
1286 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?
1287 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d2, &geom.basis_d2)?
1288 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d3, &geom.basis_d3)?;
1289
1290 let a_ab = xt_diag_y_dense(&basis_uv, &dw, &geom.basis)?;
1301 let a_ij = xt_diag_y_dense(&basis_u, &dw, &basis_v)?;
1302 let a_iwj = xt_diag_y_dense(&basis_u, &dw_v, &geom.basis)?;
1303 let a_jwi = xt_diag_y_dense(&basis_v, &dw_u, &geom.basis)?;
1304 let d2_h_ww = &a_ab
1305 + &a_ab.t()
1306 + &a_ij
1307 + a_ij.t()
1308 + &a_iwj
1309 + a_iwj.t()
1310 + &a_jwi
1311 + a_jwi.t()
1312 + &xt_diag_x_dense(&geom.basis, &dw_uv)?;
1313
1314 Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
1315 &d2_h_eta_eta,
1316 &d2_h_eta_w,
1317 &d2_h_ww,
1318 )))
1319 }
1320
1321 fn exact_newton_joint_psi_terms(
1322 &self,
1323 block_states: &[ParameterBlockState],
1324 specs: &[ParameterBlockSpec],
1325 derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
1326 psi_index: usize,
1327 ) -> Result<Option<gam_problem::ExactNewtonJointPsiTerms>, String> {
1328 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1329 if derivative_blocks.len() != 2 {
1330 return Err(GamlssError::DimensionMismatch { reason: format!(
1331 "BinomialMeanWiggleFamily joint psi terms expect 2 derivative block lists, got {}",
1332 derivative_blocks.len()
1333 ) }.into());
1334 }
1335 let x_eta = self.dense_eta_design_fromspecs(specs)?;
1336 let eta = &block_states[Self::BLOCK_ETA].eta;
1337 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1338 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1339 let n = self.y.len();
1340 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1341 return Err(GamlssError::DimensionMismatch {
1342 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1343 }
1344 .into());
1345 }
1346 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1347 let p_eta = x_eta.ncols();
1348 let pw = geom.basis.ncols();
1349 let implicit_dir =
1350 self.exact_newton_joint_psi_action(block_states, derivative_blocks, psi_index, p_eta)?;
1351 let dense_dir = if implicit_dir.is_none() {
1352 self.exact_newton_joint_psi_direction(
1353 block_states,
1354 derivative_blocks,
1355 psi_index,
1356 &x_eta,
1357 )?
1358 } else {
1359 None
1360 };
1361 let z_eta_psi = if let Some((_, ref z_eta_psi)) = implicit_dir {
1362 z_eta_psi
1363 } else if let Some(ref dir_a) = dense_dir {
1364 &dir_a.z_eta_psi
1365 } else {
1366 return Ok(None);
1367 };
1368
1369 let mut objective_psi = 0.0;
1370 let mut score_eta_xa = Array1::<f64>::zeros(n);
1371 let mut score_eta_x = Array1::<f64>::zeros(n);
1372 let mut score_w_b = Array1::<f64>::zeros(n);
1373 let mut score_w_d1 = Array1::<f64>::zeros(n);
1374
1375 let mut coeff_eta_eta_xx = Array1::<f64>::zeros(n);
1376 let mut coeff_eta_eta_xa_x = Array1::<f64>::zeros(n);
1377 let mut coeff_eta_w_xa_b = Array1::<f64>::zeros(n);
1378 let mut coeff_eta_w_x_b = Array1::<f64>::zeros(n);
1379 let mut coeff_eta_w_x_d1 = Array1::<f64>::zeros(n);
1380 let mut coeff_eta_w_xa_d1 = Array1::<f64>::zeros(n);
1381 let mut coeff_eta_w_x_d2 = Array1::<f64>::zeros(n);
1382 let mut coeff_ww_bb = Array1::<f64>::zeros(n);
1383 let mut coeff_ww_db = Array1::<f64>::zeros(n);
1384
1385 for row in 0..n {
1386 let q = eta[row] + etaw[row];
1387 let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1388 let z_a = z_eta_psi[row];
1389 let a = geom.dq_dq0[row];
1390 let b = geom.d2q_dq02[row];
1391 let c = geom.d3q_dq03[row];
1392 let q_a = a * z_a;
1393
1394 objective_psi += m1 * q_a;
1395
1396 score_eta_xa[row] = m1 * a;
1397 score_eta_x[row] = m2 * q_a * a + m1 * b * z_a;
1398 score_w_b[row] = m2 * q_a;
1399 score_w_d1[row] = m1 * z_a;
1400
1401 coeff_eta_eta_xx[row] =
1402 m3 * q_a * a * a + m2 * (2.0 * a * b * z_a + q_a * b) + m1 * c * z_a;
1403 coeff_eta_eta_xa_x[row] = m2 * a * a + m1 * b;
1404 coeff_eta_w_xa_b[row] = m2 * a;
1405 coeff_eta_w_x_b[row] = m3 * q_a * a + m2 * b * z_a;
1406 coeff_eta_w_x_d1[row] = m2 * (a * z_a + q_a);
1407 coeff_eta_w_xa_d1[row] = m1;
1408 coeff_eta_w_x_d2[row] = m1 * z_a;
1409 coeff_ww_bb[row] = m3 * q_a;
1410 coeff_ww_db[row] = m2 * z_a;
1411 }
1412
1413 let score_w = gam_linalg::faer_ndarray::fast_atv(&geom.basis, &score_w_b)
1414 + gam_linalg::faer_ndarray::fast_atv(&geom.basis_d1, &score_w_d1);
1415
1416 if let Some((action, _)) = implicit_dir {
1417 let score_eta = action.transpose_mul(score_eta_xa.view())
1418 + gam_linalg::faer_ndarray::fast_atv(x_eta.as_ref(), &score_eta_x);
1419 let score_psi = binomial_pack_mean_wiggle_joint_score(&score_eta, &score_w);
1420 let x_eta_arc = shared_dense_arc(x_eta.as_ref());
1421 let basis_arc = Arc::new(geom.basis.clone());
1422 let basis_d1_arc = Arc::new(geom.basis_d1.clone());
1423 let basis_d2_arc = Arc::new(geom.basis_d2.clone());
1424 let zeros = Array1::<f64>::zeros(n);
1425 let operator = CustomFamilyJointPsiOperator::new(
1426 p_eta + pw,
1427 vec![
1428 CustomFamilyJointDesignChannel::new(
1429 0..p_eta,
1430 Arc::clone(&x_eta_arc),
1431 Some(action),
1432 ),
1433 CustomFamilyJointDesignChannel::new(
1434 p_eta..p_eta + pw,
1435 Arc::clone(&basis_arc),
1436 None,
1437 ),
1438 CustomFamilyJointDesignChannel::new(
1439 p_eta..p_eta + pw,
1440 Arc::clone(&basis_d1_arc),
1441 None,
1442 ),
1443 CustomFamilyJointDesignChannel::new(
1444 p_eta..p_eta + pw,
1445 Arc::clone(&basis_d2_arc),
1446 None,
1447 ),
1448 ],
1449 vec![
1450 CustomFamilyJointDesignPairContribution::new(
1451 0,
1452 0,
1453 coeff_eta_eta_xa_x.clone(),
1454 coeff_eta_eta_xx.clone(),
1455 ),
1456 CustomFamilyJointDesignPairContribution::new(
1457 0,
1458 1,
1459 coeff_eta_w_xa_b.clone(),
1460 coeff_eta_w_x_b.clone(),
1461 ),
1462 CustomFamilyJointDesignPairContribution::new(
1463 1,
1464 0,
1465 coeff_eta_w_xa_b.clone(),
1466 coeff_eta_w_x_b.clone(),
1467 ),
1468 CustomFamilyJointDesignPairContribution::new(
1469 0,
1470 2,
1471 coeff_eta_w_xa_d1.clone(),
1472 coeff_eta_w_x_d1.clone(),
1473 ),
1474 CustomFamilyJointDesignPairContribution::new(
1475 2,
1476 0,
1477 coeff_eta_w_xa_d1.clone(),
1478 coeff_eta_w_x_d1.clone(),
1479 ),
1480 CustomFamilyJointDesignPairContribution::new(
1481 0,
1482 3,
1483 zeros.clone(),
1484 coeff_eta_w_x_d2.clone(),
1485 ),
1486 CustomFamilyJointDesignPairContribution::new(
1487 3,
1488 0,
1489 zeros.clone(),
1490 coeff_eta_w_x_d2.clone(),
1491 ),
1492 CustomFamilyJointDesignPairContribution::new(
1493 1,
1494 1,
1495 zeros.clone(),
1496 coeff_ww_bb.clone(),
1497 ),
1498 CustomFamilyJointDesignPairContribution::new(
1499 2,
1500 1,
1501 zeros.clone(),
1502 coeff_ww_db.clone(),
1503 ),
1504 CustomFamilyJointDesignPairContribution::new(1, 2, zeros, coeff_ww_db.clone()),
1505 ],
1506 );
1507 return Ok(Some(gam_problem::ExactNewtonJointPsiTerms {
1508 objective_psi,
1509 score_psi,
1510 hessian_psi: Array2::zeros((0, 0)),
1511 hessian_psi_operator: Some(std::sync::Arc::new(operator)),
1512 }));
1513 }
1514
1515 let dir_a =
1516 dense_dir.expect("dense psi direction should exist when implicit direction is absent");
1517 let x_eta_psi = dir_a
1518 .x_eta_psi
1519 .as_ref()
1520 .expect("dense eta psi design should exist when implicit direction is absent");
1521 let score_psi = binomial_pack_mean_wiggle_joint_score(
1522 &(gam_linalg::faer_ndarray::fast_atv(x_eta_psi, &score_eta_xa)
1523 + gam_linalg::faer_ndarray::fast_atv(x_eta.as_ref(), &score_eta_x)),
1524 &score_w,
1525 );
1526 let a_eta_eta = xt_diag_y_dense(x_eta_psi, &coeff_eta_eta_xa_x, &x_eta)?;
1527 let h_eta_eta = &a_eta_eta + &a_eta_eta.t() + &xt_diag_x_dense(&x_eta, &coeff_eta_eta_xx)?;
1528 let h_eta_w = xt_diag_y_dense(x_eta_psi, &coeff_eta_w_xa_b, &geom.basis)?
1529 + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_b, &geom.basis)?
1530 + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_d1, &geom.basis_d1)?
1531 + &xt_diag_y_dense(x_eta_psi, &coeff_eta_w_xa_d1, &geom.basis_d1)?
1532 + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_d2, &geom.basis_d2)?;
1533 let a_ww = xt_diag_y_dense(&geom.basis_d1, &coeff_ww_db, &geom.basis)?;
1534 let h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww_bb)? + &a_ww + a_ww.t();
1535
1536 Ok(Some(gam_problem::ExactNewtonJointPsiTerms {
1537 objective_psi,
1538 score_psi,
1539 hessian_psi: binomial_pack_mean_wiggle_joint_symmetrichessian(
1540 &h_eta_eta, &h_eta_w, &h_ww,
1541 ),
1542 hessian_psi_operator: None,
1543 }))
1544 }
1545}
1546
1547pub(crate) struct BinomialMeanWiggleHessianWorkspace {
1548 pub(crate) family: BinomialMeanWiggleFamily,
1549 pub(crate) block_states: Vec<ParameterBlockState>,
1550 pub(crate) x_eta: Arc<Array2<f64>>,
1551 pub(crate) hessian_operator: Arc<RowCoeffOperator>,
1552}
1553
1554impl BinomialMeanWiggleHessianWorkspace {
1555 pub(crate) fn new(
1556 family: BinomialMeanWiggleFamily,
1557 block_states: Vec<ParameterBlockState>,
1558 x_eta: Array2<f64>,
1559 ) -> Result<Self, String> {
1560 let x_eta = Arc::new(x_eta);
1561 let hessian_operator = family.bmw_static_hessian_operator(&block_states, x_eta.clone())?;
1562 Ok(Self {
1563 family,
1564 block_states,
1565 x_eta,
1566 hessian_operator,
1567 })
1568 }
1569}
1570
1571impl ExactNewtonJointHessianWorkspace for BinomialMeanWiggleHessianWorkspace {
1572 fn warm_up_outer_caches_for_mode(
1573 &self,
1574 eval_mode: gam_problem::EvalMode,
1575 ) -> Result<(), String> {
1576 match eval_mode {
1577 gam_problem::EvalMode::ValueOnly
1578 | gam_problem::EvalMode::ValueAndGradient
1579 | gam_problem::EvalMode::ValueGradientHessian => Ok(()),
1580 }
1581 }
1582
1583 fn hessian_matvec_available(&self) -> bool {
1584 true
1585 }
1586
1587 fn hessian_matvec(&self, v: &Array1<f64>) -> Result<Option<Array1<f64>>, String> {
1588 Ok(Some(gam_problem::HyperOperator::mul_vec(
1589 self.hessian_operator.as_ref(),
1590 v,
1591 )))
1592 }
1593
1594 fn hessian_diagonal(&self) -> Result<Option<Array1<f64>>, String> {
1595 Ok(None)
1596 }
1597
1598 fn directional_derivative(
1599 &self,
1600 d_beta_flat: &Array1<f64>,
1601 ) -> Result<Option<Array2<f64>>, String> {
1602 Ok(self
1603 .directional_derivative_operator(d_beta_flat)?
1604 .map(|operator| operator.to_dense()))
1605 }
1606
1607 fn directional_derivative_operator(
1608 &self,
1609 d_beta_flat: &Array1<f64>,
1610 ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
1611 self.family
1612 .bmw_directional_operator(&self.block_states, self.x_eta.clone(), d_beta_flat)
1613 }
1614
1615 fn second_directional_derivative(
1616 &self,
1617 d_beta_u_flat: &Array1<f64>,
1618 d_beta_v_flat: &Array1<f64>,
1619 ) -> Result<Option<Array2<f64>>, String> {
1620 Ok(self
1621 .second_directional_derivative_operator(d_beta_u_flat, d_beta_v_flat)?
1622 .map(|operator| operator.to_dense()))
1623 }
1624
1625 fn second_directional_derivative_operator(
1626 &self,
1627 d_beta_u: &Array1<f64>,
1628 d_beta_v: &Array1<f64>,
1629 ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
1630 self.family.bmw_second_directional_operator(
1631 &self.block_states,
1632 self.x_eta.clone(),
1633 d_beta_u,
1634 d_beta_v,
1635 )
1636 }
1637}
1638
1639impl CustomFamilyGenerative for BinomialMeanWiggleFamily {
1640 fn generativespec(
1641 &self,
1642 block_states: &[ParameterBlockState],
1643 ) -> Result<GenerativeSpec, String> {
1644 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1645 let eta = &block_states[Self::BLOCK_ETA].eta;
1646 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1647 if eta.len() != self.y.len() || etaw.len() != self.y.len() {
1648 return Err(GamlssError::DimensionMismatch {
1649 reason: "BinomialMeanWiggleFamily generative size mismatch".to_string(),
1650 }
1651 .into());
1652 }
1653 let mean = gamlss_rowwise_map_result(self.y.len(), |i| {
1654 let jet = inverse_link_jet_for_inverse_link(&self.link_kind, eta[i] + etaw[i])
1655 .map_err(|e| format!("fixed-link wiggle inverse-link evaluation failed: {e}"))?;
1656 Ok(jet.mu)
1657 })?;
1658 Ok(GenerativeSpec {
1659 mean,
1660 noise: NoiseModel::Bernoulli,
1661 })
1662 }
1663}
1664
1665#[cfg(test)]
1666mod exact_frozen_monotonicity_tests {
1667 use super::*;
1668
1669 fn frozen_family_and_wiggle_spec() -> (BinomialMeanWiggleFamily, ParameterBlockSpec) {
1670 let n = 4;
1671 let p = 3;
1672 let frozen = Array2::<f64>::zeros((n, p));
1673 let family = BinomialMeanWiggleFamily {
1674 y: Array1::zeros(n),
1675 weights: Array1::ones(n),
1676 link_kind: InverseLink::Standard(StandardLink::Logit),
1677 wiggle_knots: Array1::linspace(-1.0, 1.0, 8),
1678 wiggle_degree: 3,
1679 policy: gam_runtime::resource::ResourcePolicy::default_library(),
1680 frozen_warp_design: Some(Arc::new(frozen.clone())),
1681 };
1682 let spec = ParameterBlockSpec {
1683 name: "wiggle".to_string(),
1684 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(frozen)),
1685 offset: Array1::zeros(n),
1686 penalties: vec![],
1687 nullspace_dims: vec![],
1688 initial_log_lambdas: Array1::zeros(0),
1689 initial_beta: None,
1690 gauge_priority: 100,
1691 jacobian_callback: None,
1692 stacked_design: None,
1693 stacked_offset: None,
1694 };
1695 (family, spec)
1696 }
1697
1698 #[test]
1699 fn frozen_warp_keeps_exact_nonnegative_i_spline_cone() {
1700 let (family, spec) = frozen_family_and_wiggle_spec();
1701 let constraints = family
1702 .block_linear_constraints(&[], BinomialMeanWiggleFamily::BLOCK_WIGGLE, &spec)
1703 .expect("frozen constraint construction")
1704 .expect("wiggle block must be constrained");
1705 assert_eq!(constraints.a, Array2::<f64>::eye(3));
1706 assert_eq!(constraints.b, Array1::<f64>::zeros(3));
1707
1708 let solver_slop = Array1::from_vec(vec![
1709 -0.5 * crate::wiggle::MONOTONE_WIGGLE_ACTIVE_SET_TOL,
1710 0.2,
1711 0.0,
1712 ]);
1713 let projected = family
1714 .post_update_block_beta(
1715 &[],
1716 BinomialMeanWiggleFamily::BLOCK_WIGGLE,
1717 &spec,
1718 solver_slop,
1719 )
1720 .expect("active-set slop projects onto the exact cone");
1721 assert_eq!(projected[0], 0.0);
1722
1723 let material_violation = Array1::from_vec(vec![
1724 -2.0 * crate::wiggle::MONOTONE_WIGGLE_ACTIVE_SET_TOL,
1725 0.2,
1726 0.0,
1727 ]);
1728 assert!(
1729 family
1730 .post_update_block_beta(
1731 &[],
1732 BinomialMeanWiggleFamily::BLOCK_WIGGLE,
1733 &spec,
1734 material_violation,
1735 )
1736 .is_err(),
1737 "a material negative I-spline coefficient must be rejected"
1738 );
1739 }
1740}