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>>>,
49}
50
51pub(crate) struct BinomialMeanWiggleGeometry {
52 pub(crate) basis: Array2<f64>,
53 pub(crate) basis_d1: Array2<f64>,
54 pub(crate) basis_d2: Array2<f64>,
55 pub(crate) basis_d3: Array2<f64>,
56 pub(crate) dq_dq0: Array1<f64>,
57 pub(crate) d2q_dq02: Array1<f64>,
58 pub(crate) d3q_dq03: Array1<f64>,
59 pub(crate) d4q_dq04: Array1<f64>,
60}
61
62pub(crate) struct BinomialMeanWiggleJointPsiDirection {
63 pub(crate) x_eta_psi: Option<Array2<f64>>,
64 pub(crate) z_eta_psi: Array1<f64>,
65}
66
67impl BinomialMeanWiggleFamily {
68 pub const BLOCK_ETA: usize = 0;
69 pub const BLOCK_WIGGLE: usize = 1;
70
71 pub(crate) fn wiggle_basiswith_options(
72 &self,
73 q0: ArrayView1<'_, f64>,
74 options: BasisOptions,
75 ) -> Result<Array2<f64>, String> {
76 monotone_wiggle_basis_with_derivative_order(
77 q0,
78 &self.wiggle_knots,
79 self.wiggle_degree,
80 options.derivative_order,
81 )
82 }
83
84 pub(crate) fn wiggle_design(&self, q0: ArrayView1<'_, f64>) -> Result<Array2<f64>, String> {
85 self.wiggle_basiswith_options(q0, BasisOptions::value())
86 }
87
88 pub(crate) fn wiggle_dq_dq0(
89 &self,
90 q0: ArrayView1<'_, f64>,
91 beta_link_wiggle: ArrayView1<'_, f64>,
92 ) -> Result<Array1<f64>, String> {
93 let d_constrained = self.wiggle_basiswith_options(q0, BasisOptions::first_derivative())?;
94 if d_constrained.ncols() != beta_link_wiggle.len() {
95 return Err(GamlssError::DimensionMismatch { reason: format!(
96 "wiggle derivative/beta mismatch: basis has {} columns but beta_link_wiggle has {} coefficients",
97 d_constrained.ncols(),
98 beta_link_wiggle.len()
99 ) }.into());
100 }
101 Ok(d_constrained.dot(&beta_link_wiggle) + 1.0)
102 }
103
104 pub(crate) fn wiggle_d2q_dq02(
105 &self,
106 q0: ArrayView1<'_, f64>,
107 beta_link_wiggle: ArrayView1<'_, f64>,
108 ) -> Result<Array1<f64>, String> {
109 let d2 = self.wiggle_basiswith_options(q0, BasisOptions::second_derivative())?;
110 if d2.ncols() != beta_link_wiggle.len() {
111 return Err(GamlssError::DimensionMismatch { reason: format!(
112 "wiggle second-derivative/beta mismatch: basis has {} columns but beta_link_wiggle has {} coefficients",
113 d2.ncols(),
114 beta_link_wiggle.len()
115 ) }.into());
116 }
117 Ok(d2.dot(&beta_link_wiggle))
118 }
119
120 pub(crate) fn wiggle_d3basis_constrained(
121 &self,
122 q0: ArrayView1<'_, f64>,
123 ) -> Result<Array2<f64>, String> {
124 monotone_wiggle_basis_with_derivative_order(q0, &self.wiggle_knots, self.wiggle_degree, 3)
125 }
126
127 pub(crate) fn wiggle_d3q_dq03(
128 &self,
129 q0: ArrayView1<'_, f64>,
130 beta_link_wiggle: ArrayView1<'_, f64>,
131 ) -> Result<Array1<f64>, String> {
132 let d3 = self.wiggle_d3basis_constrained(q0)?;
133 if d3.ncols() != beta_link_wiggle.len() {
134 return Err(GamlssError::DimensionMismatch { reason: format!(
135 "wiggle third-derivative/beta mismatch: basis has {} columns but beta_link_wiggle has {} coefficients",
136 d3.ncols(),
137 beta_link_wiggle.len()
138 ) }.into());
139 }
140 Ok(d3.dot(&beta_link_wiggle))
141 }
142
143 pub(crate) fn wiggle_d4q_dq04(
144 &self,
145 q0: ArrayView1<'_, f64>,
146 beta_link_wiggle: ArrayView1<'_, f64>,
147 ) -> Result<Array1<f64>, String> {
148 let d4 = monotone_wiggle_basis_with_derivative_order(
149 q0,
150 &self.wiggle_knots,
151 self.wiggle_degree,
152 4,
153 )?;
154 if d4.ncols() != beta_link_wiggle.len() {
155 return Err(GamlssError::DimensionMismatch { reason: format!(
156 "wiggle fourth-derivative/beta mismatch: basis has {} columns but beta_link_wiggle has {} coefficients",
157 d4.ncols(),
158 beta_link_wiggle.len()
159 ) }.into());
160 }
161 Ok(d4.dot(&beta_link_wiggle))
162 }
163
164 pub(crate) fn wiggle_geometry(
165 &self,
166 q0: ArrayView1<'_, f64>,
167 beta_link_wiggle: ArrayView1<'_, f64>,
168 ) -> Result<BinomialMeanWiggleGeometry, String> {
169 if let Some(frozen) = self.frozen_warp_design.as_ref() {
177 let n = frozen.nrows();
178 let pw = frozen.ncols();
179 return Ok(BinomialMeanWiggleGeometry {
180 basis: frozen.as_ref().clone(),
181 basis_d1: Array2::zeros((n, pw)),
182 basis_d2: Array2::zeros((n, pw)),
183 basis_d3: Array2::zeros((n, pw)),
184 dq_dq0: Array1::ones(n),
185 d2q_dq02: Array1::zeros(n),
186 d3q_dq03: Array1::zeros(n),
187 d4q_dq04: Array1::zeros(n),
188 });
189 }
190 let basis = self.wiggle_design(q0)?;
191 let basis_d1 = self.wiggle_basiswith_options(q0, BasisOptions::first_derivative())?;
192 let basis_d2 = self.wiggle_basiswith_options(q0, BasisOptions::second_derivative())?;
193 let basis_d3 = self.wiggle_d3basis_constrained(q0)?;
194 let dq_dq0 = self.wiggle_dq_dq0(q0, beta_link_wiggle)?;
195 let d2q_dq02 = self.wiggle_d2q_dq02(q0, beta_link_wiggle)?;
196 let d3q_dq03 = self.wiggle_d3q_dq03(q0, beta_link_wiggle)?;
197 let d4q_dq04 = self.wiggle_d4q_dq04(q0, beta_link_wiggle)?;
198 Ok(BinomialMeanWiggleGeometry {
199 basis,
200 basis_d1,
201 basis_d2,
202 basis_d3,
203 dq_dq0,
204 d2q_dq02,
205 d3q_dq03,
206 d4q_dq04,
207 })
208 }
209
210 pub(crate) fn neglog_q_derivatives(
211 &self,
212 y: f64,
213 weight: f64,
214 q: f64,
215 ) -> Result<(f64, f64, f64), String> {
216 let jet = inverse_link_jet_for_inverse_link(&self.link_kind, q)
217 .map_err(|e| format!("fixed-link wiggle inverse-link evaluation failed: {e}"))?;
218 Ok(binomial_neglog_q_derivatives_dispatch(
222 y,
223 weight,
224 q,
225 jet.mu,
226 jet.d1,
227 jet.d2,
228 jet.d3,
229 &self.link_kind,
230 ))
231 }
232
233 pub(crate) fn neglog_q_fourth_derivative(
234 &self,
235 y: f64,
236 weight: f64,
237 q: f64,
238 ) -> Result<f64, String> {
239 let jet = inverse_link_jet_for_inverse_link(&self.link_kind, q)
240 .map_err(|e| format!("fixed-link wiggle inverse-link evaluation failed: {e}"))?;
241 binomial_neglog_q_fourth_derivative_dispatch(
243 y,
244 weight,
245 q,
246 jet.mu,
247 jet.d1,
248 jet.d2,
249 jet.d3,
250 &self.link_kind,
251 )
252 }
253
254 pub(crate) fn dense_eta_design_fromspecs<'a>(
255 &self,
256 specs: &'a [ParameterBlockSpec],
257 ) -> Result<Cow<'a, Array2<f64>>, String> {
258 if specs.len() != 2 {
259 return Err(GamlssError::DimensionMismatch {
260 reason: format!(
261 "BinomialMeanWiggleFamily expects 2 specs, got {}",
262 specs.len()
263 ),
264 }
265 .into());
266 }
267 Ok(match specs[Self::BLOCK_ETA].design.as_dense_ref() {
268 Some(d) => Cow::Borrowed(d),
269 None => Cow::Owned(
270 specs[Self::BLOCK_ETA]
271 .design
272 .try_to_dense_with_policy(
273 &self.policy.material_policy(),
274 "BinomialMeanWiggle dense_eta_design_fromspecs eta",
275 )
276 .map_err(|e| e.to_string())?
277 .as_ref()
278 .clone(),
279 ),
280 })
281 }
282
283 pub(crate) fn exact_newton_joint_psi_direction(
284 &self,
285 block_states: &[ParameterBlockState],
286 derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
287 psi_index: usize,
288 x_eta: &Array2<f64>,
289 ) -> Result<Option<BinomialMeanWiggleJointPsiDirection>, String> {
290 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
291 if derivative_blocks.len() != 2 {
292 return Err(GamlssError::DimensionMismatch { reason: format!(
293 "BinomialMeanWiggleFamily joint psi direction expects 2 derivative block lists, got {}",
294 derivative_blocks.len()
295 ) }.into());
296 }
297 let n = self.y.len();
298 let p_eta = x_eta.ncols();
299 let beta_eta = &block_states[Self::BLOCK_ETA].beta;
300 let mut global = 0usize;
301 for (block_idx, block_derivs) in derivative_blocks.iter().enumerate() {
302 for deriv in block_derivs {
303 if global == psi_index {
304 if block_idx != Self::BLOCK_ETA {
305 return Ok(None);
306 }
307 let x_eta_psi_map = resolve_custom_family_x_psi_map(
308 deriv,
309 n,
310 p_eta,
311 0..n,
312 "BinomialMeanWiggleFamily eta",
313 &self.policy,
314 )?;
315 let x_eta_psi = x_eta_psi_map.row_chunk(0..n)?;
316 let z_eta_psi = x_eta_psi.dot(beta_eta);
317 return Ok(Some(BinomialMeanWiggleJointPsiDirection {
318 x_eta_psi: Some(x_eta_psi),
319 z_eta_psi,
320 }));
321 }
322 global += 1;
323 }
324 }
325 Ok(None)
326 }
327
328 pub(crate) fn exact_newton_joint_psi_action(
329 &self,
330 block_states: &[ParameterBlockState],
331 derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
332 psi_index: usize,
333 p_eta: usize,
334 ) -> Result<Option<(CustomFamilyPsiDesignAction, Array1<f64>)>, String> {
335 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
336 if derivative_blocks.len() != 2 {
337 return Err(GamlssError::DimensionMismatch { reason: format!(
338 "BinomialMeanWiggleFamily joint psi action expects 2 derivative block lists, got {}",
339 derivative_blocks.len()
340 ) }.into());
341 }
342 let n = self.y.len();
343 let beta_eta = &block_states[Self::BLOCK_ETA].beta;
344 let mut global = 0usize;
345 for (block_idx, block_derivs) in derivative_blocks.iter().enumerate() {
346 for deriv in block_derivs {
347 if global == psi_index {
348 if block_idx != Self::BLOCK_ETA {
349 return Ok(None);
350 }
351 let action = match CustomFamilyPsiDesignAction::from_first_derivative(
352 deriv,
353 n,
354 p_eta,
355 0..n,
356 "BinomialMeanWiggleFamily eta",
357 ) {
358 Ok(action) => action,
359 Err(_) => return Ok(None),
360 };
361 let z_eta_psi = action.forward_mul(beta_eta.view());
362 return Ok(Some((action, z_eta_psi)));
363 }
364 global += 1;
365 }
366 }
367 Ok(None)
368 }
369
370 pub(crate) fn bmw_static_hessian_operator(
371 &self,
372 block_states: &[ParameterBlockState],
373 x_eta_arc: Arc<Array2<f64>>,
374 ) -> Result<Arc<RowCoeffOperator>, String> {
375 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
376 let eta = &block_states[Self::BLOCK_ETA].eta;
377 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
378 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
379 let n = self.y.len();
380 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
381 return Err(GamlssError::DimensionMismatch {
382 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
383 }
384 .into());
385 }
386 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
387 let p_eta = x_eta_arc.ncols();
388 let pw = geom.basis.ncols();
389 let mut coeff_eta = Array1::<f64>::zeros(n);
390 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
391 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
392 let mut coeff_ww = Array1::<f64>::zeros(n);
393 for row in 0..n {
394 let q = eta[row] + etaw[row];
395 let (m1, m2, _) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
396 let a = geom.dq_dq0[row];
397 let b = geom.d2q_dq02[row];
398 coeff_eta[row] = hessian_coeff_fromobjective_q_terms(m1, m2, a, a, b);
399 coeff_etaw_b[row] = m2 * a;
400 coeff_etaw_d1[row] = m1;
401 coeff_ww[row] = m2;
402 }
403 Ok(Arc::new(RowCoeffOperator::from_directions(
404 vec![p_eta, pw],
405 vec![
406 (0, x_eta_arc),
407 (1, Arc::new(geom.basis)),
408 (1, Arc::new(geom.basis_d1)),
409 ],
410 vec![
411 (0, 0, coeff_eta),
412 (0, 1, coeff_etaw_b),
413 (0, 2, coeff_etaw_d1),
414 (1, 1, coeff_ww),
415 ],
416 n,
417 )))
418 }
419
420 pub(crate) fn bmw_directional_operator(
421 &self,
422 block_states: &[ParameterBlockState],
423 x_eta_arc: Arc<Array2<f64>>,
424 d_beta_flat: &Array1<f64>,
425 ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
426 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
427 let eta = &block_states[Self::BLOCK_ETA].eta;
428 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
429 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
430 let n = self.y.len();
431 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
432 return Err(GamlssError::DimensionMismatch {
433 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
434 }
435 .into());
436 }
437 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
438 let p_eta = x_eta_arc.ncols();
439 let pw = geom.basis.ncols();
440 let total = p_eta + pw;
441 if d_beta_flat.len() != total {
442 return Err(GamlssError::DimensionMismatch {
443 reason: format!(
444 "BinomialMeanWiggleFamily joint d_beta length mismatch: got {}, expected {}",
445 d_beta_flat.len(),
446 total
447 ),
448 }
449 .into());
450 }
451 let u_eta = d_beta_flat.slice(s![0..p_eta]).to_owned();
452 let uw = d_beta_flat.slice(s![p_eta..total]).to_owned();
453 let xi = fast_av(x_eta_arc.as_ref(), &u_eta);
454 let phi = fast_av(&geom.basis, &uw);
455 let basis1_u = fast_av(&geom.basis_d1, &uw);
456 let basis2_u = fast_av(&geom.basis_d2, &uw);
457
458 let mut coeff_eta = Array1::<f64>::zeros(n);
459 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
460 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
461 let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
462 let mut coeff_ww_bb = Array1::<f64>::zeros(n);
463 let mut coeff_ww_db = Array1::<f64>::zeros(n);
464 for row in 0..n {
465 let q = eta[row] + etaw[row];
466 let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
467 let a = geom.dq_dq0[row];
468 let b = geom.d2q_dq02[row];
469 let c = geom.d3q_dq03[row];
470 let q_u = a * xi[row] + phi[row];
471 let a_u = b * xi[row] + basis1_u[row];
472 let b_u = c * xi[row] + basis2_u[row];
473 coeff_eta[row] = directionalhessian_coeff_fromobjective_q_terms(
474 m1, m2, m3, q_u, a, a, b, a_u, a_u, b_u,
475 );
476 coeff_etaw_b[row] = m3 * q_u * a + m2 * a_u;
477 coeff_etaw_d1[row] = m2 * (a * xi[row] + q_u);
478 coeff_etaw_d2[row] = m1 * xi[row];
479 coeff_ww_bb[row] = m3 * q_u;
480 coeff_ww_db[row] = m2 * xi[row];
481 }
482 Ok(Some(Arc::new(RowCoeffOperator::from_directions(
483 vec![p_eta, pw],
484 vec![
485 (0, x_eta_arc),
486 (1, Arc::new(geom.basis)),
487 (1, Arc::new(geom.basis_d1)),
488 (1, Arc::new(geom.basis_d2)),
489 ],
490 vec![
491 (0, 0, coeff_eta),
492 (0, 1, coeff_etaw_b),
493 (0, 2, coeff_etaw_d1),
494 (0, 3, coeff_etaw_d2),
495 (1, 1, coeff_ww_bb),
496 (1, 2, coeff_ww_db),
497 ],
498 n,
499 ))))
500 }
501
502 pub(crate) fn bmw_second_directional_operator(
503 &self,
504 block_states: &[ParameterBlockState],
505 x_eta_arc: Arc<Array2<f64>>,
506 d_beta_u_flat: &Array1<f64>,
507 d_beta_v_flat: &Array1<f64>,
508 ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
509 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
510 let eta = &block_states[Self::BLOCK_ETA].eta;
511 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
512 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
513 let n = self.y.len();
514 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
515 return Err(GamlssError::DimensionMismatch {
516 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
517 }
518 .into());
519 }
520 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
521 let p_eta = x_eta_arc.ncols();
522 let pw = geom.basis.ncols();
523 let total = p_eta + pw;
524 if d_beta_u_flat.len() != total || d_beta_v_flat.len() != total {
525 return Err(GamlssError::DimensionMismatch { reason: format!(
526 "BinomialMeanWiggleFamily joint second d_beta length mismatch: got {} and {}, expected {}",
527 d_beta_u_flat.len(),
528 d_beta_v_flat.len(),
529 total
530 ) }.into());
531 }
532 let u_eta = d_beta_u_flat.slice(s![0..p_eta]).to_owned();
533 let v_eta = d_beta_v_flat.slice(s![0..p_eta]).to_owned();
534 let uw = d_beta_u_flat.slice(s![p_eta..total]).to_owned();
535 let vw = d_beta_v_flat.slice(s![p_eta..total]).to_owned();
536
537 let xi_u = fast_av(x_eta_arc.as_ref(), &u_eta);
538 let xi_v = fast_av(x_eta_arc.as_ref(), &v_eta);
539 let phi_u = fast_av(&geom.basis, &uw);
540 let phi_v = fast_av(&geom.basis, &vw);
541 let b1u = fast_av(&geom.basis_d1, &uw);
542 let b1v = fast_av(&geom.basis_d1, &vw);
543 let b2u = fast_av(&geom.basis_d2, &uw);
544 let b2v = fast_av(&geom.basis_d2, &vw);
545 let b3u = fast_av(&geom.basis_d3, &uw);
546 let b3v = fast_av(&geom.basis_d3, &vw);
547
548 let mut coeff_eta = Array1::<f64>::zeros(n);
549 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
550 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
551 let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
552 let mut coeff_etaw_d3 = Array1::<f64>::zeros(n);
553 let mut coeff_ww_bb = Array1::<f64>::zeros(n);
554 let mut coeff_ww_db = Array1::<f64>::zeros(n);
555 let mut coeff_ww_ddb = Array1::<f64>::zeros(n);
556 let mut coeff_ww_dd = Array1::<f64>::zeros(n);
557
558 for row in 0..n {
559 let q = eta[row] + etaw[row];
560 let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
561 let m4 = self.neglog_q_fourth_derivative(self.y[row], self.weights[row], q)?;
562 let a = geom.dq_dq0[row];
563 let b = geom.d2q_dq02[row];
564 let c = geom.d3q_dq03[row];
565 let d = geom.d4q_dq04[row];
566
567 let q_u = a * xi_u[row] + phi_u[row];
568 let a_u = b * xi_u[row] + b1u[row];
569 let b_u = c * xi_u[row] + b2u[row];
570 let q_v = a * xi_v[row] + phi_v[row];
571 let a_v = b * xi_v[row] + b1v[row];
572 let b_v = c * xi_v[row] + b2v[row];
573 let q_uv = b * xi_u[row] * xi_v[row] + b1u[row] * xi_v[row] + b1v[row] * xi_u[row];
574 let a_uv = c * xi_u[row] * xi_v[row] + b2u[row] * xi_v[row] + b2v[row] * xi_u[row];
575 let b_uv = d * xi_u[row] * xi_v[row] + b3u[row] * xi_v[row] + b3v[row] * xi_u[row];
576
577 coeff_eta[row] = second_directionalhessian_coeff_fromobjective_q_terms(
578 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,
579 b_uv,
580 );
581 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;
582 let dc_b_u = m3 * q_u * a + m2 * a_u;
583 let dc_b_v = m3 * q_v * a + m2 * a_v;
584 let c_b_static = m2 * a;
585 let d2_c_b1 = m3 * q_u * q_v + m2 * q_uv;
586 let dc_b1_u = m2 * q_u;
587 let dc_b1_v = m2 * q_v;
588
589 coeff_etaw_b[row] = d2_c_b;
590 coeff_etaw_d1[row] = dc_b_u * xi_v[row] + dc_b_v * xi_u[row] + d2_c_b1;
591 coeff_etaw_d2[row] =
592 c_b_static * xi_u[row] * xi_v[row] + dc_b1_u * xi_v[row] + dc_b1_v * xi_u[row];
593 coeff_etaw_d3[row] = m1 * xi_u[row] * xi_v[row];
594
595 let dw = m2;
596 let dw_u = m3 * q_u;
597 let dw_v = m3 * q_v;
598 let dw_uv = m4 * q_u * q_v + m3 * q_uv;
599 let xixj = xi_u[row] * xi_v[row];
600 coeff_ww_bb[row] = dw_uv;
601 coeff_ww_db[row] = dw_v * xi_u[row] + dw_u * xi_v[row];
602 coeff_ww_ddb[row] = dw * xixj;
603 coeff_ww_dd[row] = 2.0 * dw * xixj;
604 }
605
606 Ok(Some(Arc::new(RowCoeffOperator::from_directions(
607 vec![p_eta, pw],
608 vec![
609 (0, x_eta_arc),
610 (1, Arc::new(geom.basis)),
611 (1, Arc::new(geom.basis_d1)),
612 (1, Arc::new(geom.basis_d2)),
613 (1, Arc::new(geom.basis_d3)),
614 ],
615 vec![
616 (0, 0, coeff_eta),
617 (0, 1, coeff_etaw_b),
618 (0, 2, coeff_etaw_d1),
619 (0, 3, coeff_etaw_d2),
620 (0, 4, coeff_etaw_d3),
621 (1, 1, coeff_ww_bb),
622 (1, 2, coeff_ww_db),
623 (1, 3, coeff_ww_ddb),
624 (2, 2, coeff_ww_dd),
625 ],
626 n,
627 ))))
628 }
629
630 pub fn block_effective_jacobian(
636 specs: &[ParameterBlockSpec],
637 block_idx: usize,
638 ) -> Result<Box<dyn BlockEffectiveJacobian>, String> {
639 crate::block_layout::block_jacobian::AdditiveWiggleBlockLayout {
640 family: "BinomialMeanWiggleFamily",
641 n_outputs: 1,
642 additive_blocks: &[Self::BLOCK_ETA],
643 wiggle_block: Some(Self::BLOCK_WIGGLE),
644 }
645 .block_effective_jacobian(specs, block_idx)
646 }
647}
648
649impl CustomFamily for BinomialMeanWiggleFamily {
650 fn exact_newton_joint_hessian_beta_dependent(&self) -> bool {
651 true
652 }
653
654 fn joint_jeffreys_term_required(&self) -> bool {
692 false
693 }
694
695 fn coefficient_hessian_cost(&self, specs: &[ParameterBlockSpec]) -> u64 {
696 let p_total = specs
700 .iter()
701 .map(|s| s.design.ncols() as u64)
702 .fold(0u64, |acc, p| acc.saturating_add(p));
703 (self.y.len() as u64).saturating_mul(p_total.max(1))
704 }
705
706 fn block_linear_constraints(
707 &self,
708 _: &[ParameterBlockState],
709 block_idx: usize,
710 spec: &ParameterBlockSpec,
711 ) -> Result<Option<LinearInequalityConstraints>, String> {
712 if block_idx != Self::BLOCK_WIGGLE {
713 return Ok(None);
714 }
715 if self.frozen_warp_design.is_some() {
723 return Ok(None);
724 }
725 Ok(monotone_wiggle_nonnegative_constraints(spec.design.ncols()))
726 }
727
728 fn post_update_block_beta(
729 &self,
730 _: &[ParameterBlockState],
731 block_idx: usize,
732 block_spec: &ParameterBlockSpec,
733 beta: Array1<f64>,
734 ) -> Result<Array1<f64>, String> {
735 assert!(!block_spec.name.is_empty());
736 if block_idx != Self::BLOCK_WIGGLE {
737 return Ok(beta);
738 }
739 if self.frozen_warp_design.is_some() {
741 return Ok(beta);
742 }
743 let beta = project_monotone_wiggle_beta_nonnegative(beta);
744 validate_monotone_wiggle_beta_nonnegative(&beta, "BinomialMeanWiggleFamily post-update")?;
745 Ok(beta)
746 }
747
748 fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
749 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
750 let eta = &block_states[Self::BLOCK_ETA].eta;
751 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
752 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
753 let n = self.y.len();
754 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
755 return Err(GamlssError::DimensionMismatch {
756 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
757 }
758 .into());
759 }
760 let dq_dq0 = if self.frozen_warp_design.is_some() {
764 Array1::<f64>::ones(n)
765 } else {
766 self.wiggle_dq_dq0(eta.view(), betaw.view())?
767 };
768 if dq_dq0.len() != n {
769 return Err(GamlssError::DimensionMismatch {
770 reason: format!(
771 "BinomialMeanWiggleFamily dq/dq0 length mismatch: got {}, expected {}",
772 dq_dq0.len(),
773 n
774 ),
775 }
776 .into());
777 }
778
779 let mut ll = 0.0;
780 let mut z_eta = Array1::<f64>::zeros(n);
781 let mut w_eta = Array1::<f64>::zeros(n);
782 let mut z_wiggle = Array1::<f64>::zeros(n);
783 let mut w_wiggle = Array1::<f64>::zeros(n);
784 for i in 0..n {
785 let q = eta[i] + etaw[i];
786 let (mu_q, d1_q) = inverse_link_mu_d1_for_inverse_link(&self.link_kind, q)
787 .map_err(|e| format!("fixed-link wiggle inverse-link evaluation failed: {e}"))?;
788 let yi = self.y[i];
789 let wi = self.weights[i];
790 ll += binomial_location_scale_log_likelihood(yi, wi, q, &self.link_kind, mu_q)?;
791
792 let mu = mu_q.clamp(1e-12, 1.0 - 1e-12);
793 let var = (mu * (1.0 - mu)).max(MIN_PROB);
794 let dmu_deta = d1_q * dq_dq0[i];
795 let dmu_dw = d1_q;
796 if wi == 0.0 || !var.is_finite() {
797 z_eta[i] = eta[i];
798 z_wiggle[i] = etaw[i];
799 continue;
800 }
801
802 if dmu_deta.is_finite() {
803 w_eta[i] = floor_positiveweight(wi * (dmu_deta * dmu_deta / var), MIN_WEIGHT);
804 z_eta[i] = eta[i] + (yi - mu) / signedwith_floor(dmu_deta, MIN_DERIV);
805 } else {
806 z_eta[i] = eta[i];
807 }
808
809 if dmu_dw.is_finite() {
810 w_wiggle[i] = floor_positiveweight(wi * (dmu_dw * dmu_dw / var), MIN_WEIGHT);
811 z_wiggle[i] = etaw[i] + (yi - mu) / signedwith_floor(dmu_dw, MIN_DERIV);
812 } else {
813 z_wiggle[i] = etaw[i];
814 }
815 }
816
817 Ok(FamilyEvaluation {
818 log_likelihood: ll,
819 blockworking_sets: vec![
820 BlockWorkingSet::diagonal_checked(z_eta, w_eta)?,
821 BlockWorkingSet::diagonal_checked(z_wiggle, w_wiggle)?,
822 ],
823 })
824 }
825
826 fn exact_newton_joint_gradient_evaluation(
827 &self,
828 block_states: &[ParameterBlockState],
829 specs: &[ParameterBlockSpec],
830 ) -> Result<Option<ExactNewtonJointGradientEvaluation>, String> {
831 let eval = self.evaluate(block_states)?;
836 gamlss_joint_gradient_from_working_sets(&eval, specs, block_states).map(Some)
837 }
838
839 fn block_geometry(
840 &self,
841 block_states: &[ParameterBlockState],
842 spec: &ParameterBlockSpec,
843 ) -> Result<(DesignMatrix, Array1<f64>), String> {
844 if spec.name != "wiggle" {
845 return Ok((spec.design.clone(), spec.offset.clone()));
846 }
847 if block_states.is_empty() {
848 return Err(GamlssError::UnsupportedConfiguration {
849 reason: "wiggle geometry requires eta block".to_string(),
850 }
851 .into());
852 }
853 let eta = &block_states[Self::BLOCK_ETA].eta;
854 if eta.len() != self.y.len() {
855 return Err(GamlssError::DimensionMismatch {
856 reason: "BinomialMeanWiggleFamily eta size mismatch".to_string(),
857 }
858 .into());
859 }
860 let x = match self.frozen_warp_design.as_ref() {
866 Some(frozen) => frozen.as_ref().clone(),
867 None => self.wiggle_design(eta.view())?,
868 };
869 if x.ncols() != spec.design.ncols() {
870 return Err(GamlssError::DimensionMismatch {
871 reason: format!(
872 "dynamic wiggle design col mismatch: got {}, expected {}",
873 x.ncols(),
874 spec.design.ncols()
875 ),
876 }
877 .into());
878 }
879 let nrows = x.nrows();
880 Ok((
881 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(x)),
882 Array1::zeros(nrows),
883 ))
884 }
885
886 fn block_geometry_is_dynamic(&self) -> bool {
887 true
888 }
889
890 fn exact_newton_joint_hessian_workspace(
891 &self,
892 block_states: &[ParameterBlockState],
893 specs: &[ParameterBlockSpec],
894 ) -> Result<Option<Arc<dyn ExactNewtonJointHessianWorkspace>>, String> {
895 let x_eta = self.dense_eta_design_fromspecs(specs)?.into_owned();
896 let workspace =
897 BinomialMeanWiggleHessianWorkspace::new(self.clone(), block_states.to_vec(), x_eta)?;
898 Ok(Some(Arc::new(workspace)))
899 }
900
901 fn inner_coefficient_hessian_hvp_available(&self, specs: &[ParameterBlockSpec]) -> bool {
902 self.dense_eta_design_fromspecs(specs).is_ok()
903 }
904
905 fn exact_newton_joint_hessian_with_specs(
906 &self,
907 block_states: &[ParameterBlockState],
908 specs: &[ParameterBlockSpec],
909 ) -> Result<Option<Array2<f64>>, String> {
910 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
911 let x_eta = self.dense_eta_design_fromspecs(specs)?;
912 let eta = &block_states[Self::BLOCK_ETA].eta;
913 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
914 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
915 let n = self.y.len();
916 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
917 return Err(GamlssError::DimensionMismatch {
918 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
919 }
920 .into());
921 }
922 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
923 let p_eta = x_eta.ncols();
924 let pw = geom.basis.ncols();
925 let mut coeff_eta = Array1::<f64>::zeros(n);
926 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
927 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
928 let mut coeff_ww = Array1::<f64>::zeros(n);
929 for row in 0..n {
930 let q = eta[row] + etaw[row];
931 let (m1, m2, _) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
932 let a = geom.dq_dq0[row];
933 let b = geom.d2q_dq02[row];
934 coeff_eta[row] = hessian_coeff_fromobjective_q_terms(m1, m2, a, a, b);
935 coeff_etaw_b[row] = m2 * a;
936 coeff_etaw_d1[row] = m1;
937 coeff_ww[row] = m2;
938 }
939 let h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
940 let h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
941 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?;
942 let h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww)?;
943 assert_eq!(h_eta_eta.nrows(), p_eta);
944 assert_eq!(h_ww.nrows(), pw);
945 Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
946 &h_eta_eta, &h_eta_w, &h_ww,
947 )))
948 }
949
950 fn exact_newton_joint_hessian_directional_derivative_with_specs(
951 &self,
952 block_states: &[ParameterBlockState],
953 specs: &[ParameterBlockSpec],
954 d_beta_flat: &Array1<f64>,
955 ) -> Result<Option<Array2<f64>>, String> {
956 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
957 let x_eta = self.dense_eta_design_fromspecs(specs)?;
958 let eta = &block_states[Self::BLOCK_ETA].eta;
959 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
960 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
961 let n = self.y.len();
962 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
963 return Err(GamlssError::DimensionMismatch {
964 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
965 }
966 .into());
967 }
968 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
969 let p_eta = x_eta.ncols();
970 let pw = geom.basis.ncols();
971 if d_beta_flat.len() != p_eta + pw {
972 return Err(GamlssError::DimensionMismatch {
973 reason: format!(
974 "BinomialMeanWiggleFamily joint d_beta length mismatch: got {}, expected {}",
975 d_beta_flat.len(),
976 p_eta + pw
977 ),
978 }
979 .into());
980 }
981 let u_eta = d_beta_flat.slice(s![0..p_eta]).to_owned();
982 let uw = d_beta_flat.slice(s![p_eta..p_eta + pw]).to_owned();
983 let xi = x_eta.dot(&u_eta);
984 let phi = geom.basis.dot(&uw);
985 let basis1_u = geom.basis_d1.dot(&uw);
986 let basis2_u = geom.basis_d2.dot(&uw);
987
988 let mut coeff_eta = Array1::<f64>::zeros(n);
989 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
990 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
991 let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
992 let mut coeff_ww_bb = Array1::<f64>::zeros(n);
993 let mut coeff_ww_db = Array1::<f64>::zeros(n);
994 for row in 0..n {
995 let q = eta[row] + etaw[row];
996 let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
997 let a = geom.dq_dq0[row];
998 let b = geom.d2q_dq02[row];
999 let c = geom.d3q_dq03[row];
1000 let q_u = a * xi[row] + phi[row];
1001 let a_u = b * xi[row] + basis1_u[row];
1002 let b_u = c * xi[row] + basis2_u[row];
1003 coeff_eta[row] = directionalhessian_coeff_fromobjective_q_terms(
1004 m1, m2, m3, q_u, a, a, b, a_u, a_u, b_u,
1005 );
1006 coeff_etaw_b[row] = m3 * q_u * a + m2 * a_u;
1007 coeff_etaw_d1[row] = m2 * (a * xi[row] + q_u);
1008 coeff_etaw_d2[row] = m1 * xi[row];
1009 coeff_ww_bb[row] = m3 * q_u;
1010 coeff_ww_db[row] = m2 * xi[row];
1011 }
1012
1013 let d_h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
1014 let d_h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
1015 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?
1016 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d2, &geom.basis_d2)?;
1017 let a_ww = xt_diag_y_dense(&geom.basis_d1, &coeff_ww_db, &geom.basis)?;
1018 let d_h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww_bb)? + &a_ww + a_ww.t();
1019 Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
1020 &d_h_eta_eta,
1021 &d_h_eta_w,
1022 &d_h_ww,
1023 )))
1024 }
1025
1026 fn exact_newton_joint_hessian_second_directional_derivative_with_specs(
1084 &self,
1085 block_states: &[ParameterBlockState],
1086 specs: &[ParameterBlockSpec],
1087 d_beta_u_flat: &Array1<f64>,
1088 d_beta_v_flat: &Array1<f64>,
1089 ) -> Result<Option<Array2<f64>>, String> {
1090 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1091 let x_eta = self.dense_eta_design_fromspecs(specs)?;
1092 let eta = &block_states[Self::BLOCK_ETA].eta;
1093 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1094 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1095 let n = self.y.len();
1096 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1097 return Err(GamlssError::DimensionMismatch {
1098 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1099 }
1100 .into());
1101 }
1102 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1103 let p_eta = x_eta.ncols();
1104 let pw = geom.basis.ncols();
1105 let total = p_eta + pw;
1106 if d_beta_u_flat.len() != total || d_beta_v_flat.len() != total {
1107 return Err(GamlssError::DimensionMismatch { reason: format!(
1108 "BinomialMeanWiggleFamily joint second d_beta length mismatch: got {} and {}, expected {}",
1109 d_beta_u_flat.len(),
1110 d_beta_v_flat.len(),
1111 total
1112 ) }.into());
1113 }
1114
1115 let u_eta = d_beta_u_flat.slice(s![0..p_eta]).to_owned();
1117 let v_eta = d_beta_v_flat.slice(s![0..p_eta]).to_owned();
1118 let uw = d_beta_u_flat.slice(s![p_eta..total]).to_owned();
1119 let vw = d_beta_v_flat.slice(s![p_eta..total]).to_owned();
1120
1121 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);
1128 let b2u = geom.basis_d2.dot(&uw); let b2v = geom.basis_d2.dot(&vw);
1130 let b3u = geom.basis_d3.dot(&uw); let b3v = geom.basis_d3.dot(&vw);
1132
1133 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);
1150
1151 let mut coeff_etaw_b = Array1::<f64>::zeros(n);
1186 let mut coeff_etaw_d1 = Array1::<f64>::zeros(n);
1187 let mut coeff_etaw_d2 = Array1::<f64>::zeros(n);
1188 let mut coeff_etaw_d3 = Array1::<f64>::zeros(n);
1189
1190 let mut dw = Array1::<f64>::zeros(n);
1215 let mut dw_u = Array1::<f64>::zeros(n);
1216 let mut dw_v = Array1::<f64>::zeros(n);
1217 let mut dw_uv = Array1::<f64>::zeros(n);
1218
1219 for row in 0..n {
1220 let q = eta[row] + etaw[row];
1221 let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1222 let m4 = self.neglog_q_fourth_derivative(self.y[row], self.weights[row], q)?;
1223 let a = geom.dq_dq0[row];
1224 let b = geom.d2q_dq02[row];
1225 let c = geom.d3q_dq03[row];
1226 let d = geom.d4q_dq04[row];
1227
1228 let q_u = a * xi_u[row] + phi_u[row];
1230 let a_u = b * xi_u[row] + b1u[row];
1231 let b_u = c * xi_u[row] + b2u[row];
1232
1233 let q_v = a * xi_v[row] + phi_v[row];
1235 let a_v = b * xi_v[row] + b1v[row];
1236 let b_v = c * xi_v[row] + b2v[row];
1237
1238 let q_uv = b * xi_u[row] * xi_v[row] + b1u[row] * xi_v[row] + b1v[row] * xi_u[row];
1240 let a_uv = c * xi_u[row] * xi_v[row] + b2u[row] * xi_v[row] + b2v[row] * xi_u[row];
1241 let b_uv = d * xi_u[row] * xi_v[row] + b3u[row] * xi_v[row] + b3v[row] * xi_u[row];
1242
1243 coeff_eta[row] = second_directionalhessian_coeff_fromobjective_q_terms(
1249 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, );
1256
1257 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;
1263 let dc_b_u = m3 * q_u * a + m2 * a_u;
1265 let dc_b_v = m3 * q_v * a + m2 * a_v;
1266 let c_b_static = m2 * a;
1268 let d2_c_b1 = m3 * q_u * q_v + m2 * q_uv;
1270 let dc_b1_u = m2 * q_u;
1272 let dc_b1_v = m2 * q_v;
1273
1274 coeff_etaw_b[row] = d2_c_b;
1275 coeff_etaw_d1[row] = dc_b_u * xi_v[row] + dc_b_v * xi_u[row] + d2_c_b1;
1276 coeff_etaw_d2[row] =
1277 c_b_static * xi_u[row] * xi_v[row] + dc_b1_u * xi_v[row] + dc_b1_v * xi_u[row];
1278 coeff_etaw_d3[row] = m1 * xi_u[row] * xi_v[row];
1279
1280 dw[row] = m2;
1283 dw_u[row] = m3 * q_u;
1284 dw_v[row] = m3 * q_v;
1285 dw_uv[row] = m4 * q_u * q_v + m3 * q_uv;
1286 }
1287
1288 let d2_h_eta_eta = xt_diag_x_dense(&x_eta, &coeff_eta)?;
1290
1291 let d2_h_eta_w = xt_diag_y_dense(&x_eta, &coeff_etaw_b, &geom.basis)?
1298 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d1, &geom.basis_d1)?
1299 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d2, &geom.basis_d2)?
1300 + &xt_diag_y_dense(&x_eta, &coeff_etaw_d3, &geom.basis_d3)?;
1301
1302 let a_ab = xt_diag_y_dense(&basis_uv, &dw, &geom.basis)?;
1313 let a_ij = xt_diag_y_dense(&basis_u, &dw, &basis_v)?;
1314 let a_iwj = xt_diag_y_dense(&basis_u, &dw_v, &geom.basis)?;
1315 let a_jwi = xt_diag_y_dense(&basis_v, &dw_u, &geom.basis)?;
1316 let d2_h_ww = &a_ab
1317 + &a_ab.t()
1318 + &a_ij
1319 + a_ij.t()
1320 + &a_iwj
1321 + a_iwj.t()
1322 + &a_jwi
1323 + a_jwi.t()
1324 + &xt_diag_x_dense(&geom.basis, &dw_uv)?;
1325
1326 Ok(Some(binomial_pack_mean_wiggle_joint_symmetrichessian(
1327 &d2_h_eta_eta,
1328 &d2_h_eta_w,
1329 &d2_h_ww,
1330 )))
1331 }
1332
1333 fn exact_newton_joint_psi_terms(
1334 &self,
1335 block_states: &[ParameterBlockState],
1336 specs: &[ParameterBlockSpec],
1337 derivative_blocks: &[Vec<CustomFamilyBlockPsiDerivative>],
1338 psi_index: usize,
1339 ) -> Result<Option<crate::custom_family::ExactNewtonJointPsiTerms>, String> {
1340 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1341 if derivative_blocks.len() != 2 {
1342 return Err(GamlssError::DimensionMismatch { reason: format!(
1343 "BinomialMeanWiggleFamily joint psi terms expect 2 derivative block lists, got {}",
1344 derivative_blocks.len()
1345 ) }.into());
1346 }
1347 let x_eta = self.dense_eta_design_fromspecs(specs)?;
1348 let eta = &block_states[Self::BLOCK_ETA].eta;
1349 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1350 let betaw = &block_states[Self::BLOCK_WIGGLE].beta;
1351 let n = self.y.len();
1352 if eta.len() != n || etaw.len() != n || self.weights.len() != n {
1353 return Err(GamlssError::DimensionMismatch {
1354 reason: "BinomialMeanWiggleFamily input size mismatch".to_string(),
1355 }
1356 .into());
1357 }
1358 let geom = self.wiggle_geometry(eta.view(), betaw.view())?;
1359 let p_eta = x_eta.ncols();
1360 let pw = geom.basis.ncols();
1361 let implicit_dir =
1362 self.exact_newton_joint_psi_action(block_states, derivative_blocks, psi_index, p_eta)?;
1363 let dense_dir = if implicit_dir.is_none() {
1364 self.exact_newton_joint_psi_direction(
1365 block_states,
1366 derivative_blocks,
1367 psi_index,
1368 &x_eta,
1369 )?
1370 } else {
1371 None
1372 };
1373 let z_eta_psi = if let Some((_, ref z_eta_psi)) = implicit_dir {
1374 z_eta_psi
1375 } else if let Some(ref dir_a) = dense_dir {
1376 &dir_a.z_eta_psi
1377 } else {
1378 return Ok(None);
1379 };
1380
1381 let mut objective_psi = 0.0;
1382 let mut score_eta_xa = Array1::<f64>::zeros(n);
1383 let mut score_eta_x = Array1::<f64>::zeros(n);
1384 let mut score_w_b = Array1::<f64>::zeros(n);
1385 let mut score_w_d1 = Array1::<f64>::zeros(n);
1386
1387 let mut coeff_eta_eta_xx = Array1::<f64>::zeros(n);
1388 let mut coeff_eta_eta_xa_x = Array1::<f64>::zeros(n);
1389 let mut coeff_eta_w_xa_b = Array1::<f64>::zeros(n);
1390 let mut coeff_eta_w_x_b = Array1::<f64>::zeros(n);
1391 let mut coeff_eta_w_x_d1 = Array1::<f64>::zeros(n);
1392 let mut coeff_eta_w_xa_d1 = Array1::<f64>::zeros(n);
1393 let mut coeff_eta_w_x_d2 = Array1::<f64>::zeros(n);
1394 let mut coeff_ww_bb = Array1::<f64>::zeros(n);
1395 let mut coeff_ww_db = Array1::<f64>::zeros(n);
1396
1397 for row in 0..n {
1398 let q = eta[row] + etaw[row];
1399 let (m1, m2, m3) = self.neglog_q_derivatives(self.y[row], self.weights[row], q)?;
1400 let z_a = z_eta_psi[row];
1401 let a = geom.dq_dq0[row];
1402 let b = geom.d2q_dq02[row];
1403 let c = geom.d3q_dq03[row];
1404 let q_a = a * z_a;
1405
1406 objective_psi += m1 * q_a;
1407
1408 score_eta_xa[row] = m1 * a;
1409 score_eta_x[row] = m2 * q_a * a + m1 * b * z_a;
1410 score_w_b[row] = m2 * q_a;
1411 score_w_d1[row] = m1 * z_a;
1412
1413 coeff_eta_eta_xx[row] =
1414 m3 * q_a * a * a + m2 * (2.0 * a * b * z_a + q_a * b) + m1 * c * z_a;
1415 coeff_eta_eta_xa_x[row] = m2 * a * a + m1 * b;
1416 coeff_eta_w_xa_b[row] = m2 * a;
1417 coeff_eta_w_x_b[row] = m3 * q_a * a + m2 * b * z_a;
1418 coeff_eta_w_x_d1[row] = m2 * (a * z_a + q_a);
1419 coeff_eta_w_xa_d1[row] = m1;
1420 coeff_eta_w_x_d2[row] = m1 * z_a;
1421 coeff_ww_bb[row] = m3 * q_a;
1422 coeff_ww_db[row] = m2 * z_a;
1423 }
1424
1425 let score_w = gam_linalg::faer_ndarray::fast_atv(&geom.basis, &score_w_b)
1426 + gam_linalg::faer_ndarray::fast_atv(&geom.basis_d1, &score_w_d1);
1427
1428 if let Some((action, _)) = implicit_dir {
1429 let score_eta = action.transpose_mul(score_eta_xa.view())
1430 + gam_linalg::faer_ndarray::fast_atv(x_eta.as_ref(), &score_eta_x);
1431 let score_psi = binomial_pack_mean_wiggle_joint_score(&score_eta, &score_w);
1432 let x_eta_arc = shared_dense_arc(x_eta.as_ref());
1433 let basis_arc = Arc::new(geom.basis.clone());
1434 let basis_d1_arc = Arc::new(geom.basis_d1.clone());
1435 let basis_d2_arc = Arc::new(geom.basis_d2.clone());
1436 let zeros = Array1::<f64>::zeros(n);
1437 let operator = CustomFamilyJointPsiOperator::new(
1438 p_eta + pw,
1439 vec![
1440 CustomFamilyJointDesignChannel::new(
1441 0..p_eta,
1442 Arc::clone(&x_eta_arc),
1443 Some(action),
1444 ),
1445 CustomFamilyJointDesignChannel::new(
1446 p_eta..p_eta + pw,
1447 Arc::clone(&basis_arc),
1448 None,
1449 ),
1450 CustomFamilyJointDesignChannel::new(
1451 p_eta..p_eta + pw,
1452 Arc::clone(&basis_d1_arc),
1453 None,
1454 ),
1455 CustomFamilyJointDesignChannel::new(
1456 p_eta..p_eta + pw,
1457 Arc::clone(&basis_d2_arc),
1458 None,
1459 ),
1460 ],
1461 vec![
1462 CustomFamilyJointDesignPairContribution::new(
1463 0,
1464 0,
1465 coeff_eta_eta_xa_x.clone(),
1466 coeff_eta_eta_xx.clone(),
1467 ),
1468 CustomFamilyJointDesignPairContribution::new(
1469 0,
1470 1,
1471 coeff_eta_w_xa_b.clone(),
1472 coeff_eta_w_x_b.clone(),
1473 ),
1474 CustomFamilyJointDesignPairContribution::new(
1475 1,
1476 0,
1477 coeff_eta_w_xa_b.clone(),
1478 coeff_eta_w_x_b.clone(),
1479 ),
1480 CustomFamilyJointDesignPairContribution::new(
1481 0,
1482 2,
1483 coeff_eta_w_xa_d1.clone(),
1484 coeff_eta_w_x_d1.clone(),
1485 ),
1486 CustomFamilyJointDesignPairContribution::new(
1487 2,
1488 0,
1489 coeff_eta_w_xa_d1.clone(),
1490 coeff_eta_w_x_d1.clone(),
1491 ),
1492 CustomFamilyJointDesignPairContribution::new(
1493 0,
1494 3,
1495 zeros.clone(),
1496 coeff_eta_w_x_d2.clone(),
1497 ),
1498 CustomFamilyJointDesignPairContribution::new(
1499 3,
1500 0,
1501 zeros.clone(),
1502 coeff_eta_w_x_d2.clone(),
1503 ),
1504 CustomFamilyJointDesignPairContribution::new(
1505 1,
1506 1,
1507 zeros.clone(),
1508 coeff_ww_bb.clone(),
1509 ),
1510 CustomFamilyJointDesignPairContribution::new(
1511 2,
1512 1,
1513 zeros.clone(),
1514 coeff_ww_db.clone(),
1515 ),
1516 CustomFamilyJointDesignPairContribution::new(1, 2, zeros, coeff_ww_db.clone()),
1517 ],
1518 );
1519 return Ok(Some(crate::custom_family::ExactNewtonJointPsiTerms {
1520 objective_psi,
1521 score_psi,
1522 hessian_psi: Array2::zeros((0, 0)),
1523 hessian_psi_operator: Some(std::sync::Arc::new(operator)),
1524 }));
1525 }
1526
1527 let dir_a =
1528 dense_dir.expect("dense psi direction should exist when implicit direction is absent");
1529 let x_eta_psi = dir_a
1530 .x_eta_psi
1531 .as_ref()
1532 .expect("dense eta psi design should exist when implicit direction is absent");
1533 let score_psi = binomial_pack_mean_wiggle_joint_score(
1534 &(gam_linalg::faer_ndarray::fast_atv(x_eta_psi, &score_eta_xa)
1535 + gam_linalg::faer_ndarray::fast_atv(x_eta.as_ref(), &score_eta_x)),
1536 &score_w,
1537 );
1538 let a_eta_eta = xt_diag_y_dense(x_eta_psi, &coeff_eta_eta_xa_x, &x_eta)?;
1539 let h_eta_eta = &a_eta_eta + &a_eta_eta.t() + &xt_diag_x_dense(&x_eta, &coeff_eta_eta_xx)?;
1540 let h_eta_w = xt_diag_y_dense(x_eta_psi, &coeff_eta_w_xa_b, &geom.basis)?
1541 + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_b, &geom.basis)?
1542 + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_d1, &geom.basis_d1)?
1543 + &xt_diag_y_dense(x_eta_psi, &coeff_eta_w_xa_d1, &geom.basis_d1)?
1544 + &xt_diag_y_dense(&x_eta, &coeff_eta_w_x_d2, &geom.basis_d2)?;
1545 let a_ww = xt_diag_y_dense(&geom.basis_d1, &coeff_ww_db, &geom.basis)?;
1546 let h_ww = xt_diag_x_dense(&geom.basis, &coeff_ww_bb)? + &a_ww + a_ww.t();
1547
1548 Ok(Some(crate::custom_family::ExactNewtonJointPsiTerms {
1549 objective_psi,
1550 score_psi,
1551 hessian_psi: binomial_pack_mean_wiggle_joint_symmetrichessian(
1552 &h_eta_eta, &h_eta_w, &h_ww,
1553 ),
1554 hessian_psi_operator: None,
1555 }))
1556 }
1557}
1558
1559pub(crate) struct BinomialMeanWiggleHessianWorkspace {
1560 pub(crate) family: BinomialMeanWiggleFamily,
1561 pub(crate) block_states: Vec<ParameterBlockState>,
1562 pub(crate) x_eta: Arc<Array2<f64>>,
1563 pub(crate) hessian_operator: Arc<RowCoeffOperator>,
1564}
1565
1566impl BinomialMeanWiggleHessianWorkspace {
1567 pub(crate) fn new(
1568 family: BinomialMeanWiggleFamily,
1569 block_states: Vec<ParameterBlockState>,
1570 x_eta: Array2<f64>,
1571 ) -> Result<Self, String> {
1572 let x_eta = Arc::new(x_eta);
1573 let hessian_operator = family.bmw_static_hessian_operator(&block_states, x_eta.clone())?;
1574 Ok(Self {
1575 family,
1576 block_states,
1577 x_eta,
1578 hessian_operator,
1579 })
1580 }
1581}
1582
1583impl ExactNewtonJointHessianWorkspace for BinomialMeanWiggleHessianWorkspace {
1584 fn hessian_matvec_available(&self) -> bool {
1585 true
1586 }
1587
1588 fn hessian_matvec(&self, v: &Array1<f64>) -> Result<Option<Array1<f64>>, String> {
1589 Ok(Some(gam_problem::HyperOperator::mul_vec(
1590 self.hessian_operator.as_ref(),
1591 v,
1592 )))
1593 }
1594
1595 fn hessian_diagonal(&self) -> Result<Option<Array1<f64>>, String> {
1596 Ok(None)
1597 }
1598
1599 fn directional_derivative(
1600 &self,
1601 d_beta_flat: &Array1<f64>,
1602 ) -> Result<Option<Array2<f64>>, String> {
1603 Ok(self
1604 .directional_derivative_operator(d_beta_flat)?
1605 .map(|operator| operator.to_dense()))
1606 }
1607
1608 fn directional_derivative_operator(
1609 &self,
1610 d_beta_flat: &Array1<f64>,
1611 ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
1612 self.family
1613 .bmw_directional_operator(&self.block_states, self.x_eta.clone(), d_beta_flat)
1614 }
1615
1616 fn second_directional_derivative(
1617 &self,
1618 d_beta_u_flat: &Array1<f64>,
1619 d_beta_v_flat: &Array1<f64>,
1620 ) -> Result<Option<Array2<f64>>, String> {
1621 Ok(self
1622 .second_directional_derivative_operator(d_beta_u_flat, d_beta_v_flat)?
1623 .map(|operator| operator.to_dense()))
1624 }
1625
1626 fn second_directional_derivative_operator(
1627 &self,
1628 d_beta_u: &Array1<f64>,
1629 d_beta_v: &Array1<f64>,
1630 ) -> Result<Option<Arc<dyn gam_problem::HyperOperator>>, String> {
1631 self.family.bmw_second_directional_operator(
1632 &self.block_states,
1633 self.x_eta.clone(),
1634 d_beta_u,
1635 d_beta_v,
1636 )
1637 }
1638}
1639
1640impl CustomFamilyGenerative for BinomialMeanWiggleFamily {
1641 fn generativespec(
1642 &self,
1643 block_states: &[ParameterBlockState],
1644 ) -> Result<GenerativeSpec, String> {
1645 validate_block_count::<GamlssError>("BinomialMeanWiggleFamily", 2, block_states.len())?;
1646 let eta = &block_states[Self::BLOCK_ETA].eta;
1647 let etaw = &block_states[Self::BLOCK_WIGGLE].eta;
1648 if eta.len() != self.y.len() || etaw.len() != self.y.len() {
1649 return Err(GamlssError::DimensionMismatch {
1650 reason: "BinomialMeanWiggleFamily generative size mismatch".to_string(),
1651 }
1652 .into());
1653 }
1654 let mean = gamlss_rowwise_map_result(self.y.len(), |i| {
1655 let jet = inverse_link_jet_for_inverse_link(&self.link_kind, eta[i] + etaw[i])
1656 .map_err(|e| format!("fixed-link wiggle inverse-link evaluation failed: {e}"))?;
1657 Ok(jet.mu)
1658 })?;
1659 Ok(GenerativeSpec {
1660 mean,
1661 noise: NoiseModel::Bernoulli,
1662 })
1663 }
1664}