1use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Axis};
115use rayon::prelude::*;
116use serde::{Deserialize, Serialize};
117
118use faer::Side;
119
120use gam_linalg::faer_ndarray::{FaerEigh, default_rrqr_rank_alpha, rrqr_nullspace_basis};
121
122use super::{
123 AnisoBasisPsiDerivatives, AnisoPenaltyCrossProvider, BasisBuildResult, BasisError,
124 BasisMetadata, CenterStrategy, PenaltyCandidate, PenaltySource,
125 filter_active_penalty_candidates_with_ops, normalize_penalty,
126 normalize_penalty_cross_psi_derivative, normalize_penaltywith_psi_derivatives,
127 select_centers_by_strategy, trace_of_product,
128};
129
130pub(crate) const MEASURE_JET_PROFILE_CUTOFF: f64 = 3.0;
136
137pub(crate) const MEASURE_JET_PSEUDOINVERSE_RTOL: f64 = 64.0 * f64::EPSILON;
141
142pub(crate) const MEASURE_JET_DEFAULT_ORDER_S: f64 = 1.5;
148
149pub(crate) const MEASURE_JET_MIN_AUTO_SCALES: usize = 3;
153pub(crate) const MEASURE_JET_MAX_AUTO_SCALES: usize = 8;
154
155pub(crate) const MEASURE_JET_AUTO_LENGTH_SCALE_FACTOR: f64 = 1.0;
170
171pub(crate) const MEASURE_JET_PARALLEL_FORM_BUDGET_DOUBLES: usize = 1 << 26;
177
178#[derive(Debug, Clone, Serialize, Deserialize, Default)]
184pub enum MeasureJetIdentifiability {
185 #[default]
190 CenterSumToZero,
191 FrozenTransform { transform: Array2<f64> },
194}
195
196#[derive(Debug, Clone, Serialize, Deserialize)]
201pub struct MeasureJetFrozenQuadrature {
202 pub masses: Array1<f64>,
204 pub eps_band: Vec<f64>,
206 pub support_means: Vec<f64>,
209 pub penalty_normalization_scales: Vec<f64>,
212 pub raw_penalty_normalization_scales: Vec<f64>,
215 pub fused_penalty_normalization_scale: Option<f64>,
218 #[serde(default)]
227 pub sigma_coord: Option<f64>,
228}
229
230fn measure_jet_learn_length_scale_default() -> bool {
234 false
235}
236
237#[derive(Debug, Clone, Serialize, Deserialize)]
245pub struct MeasureJetBasisSpec {
246 pub center_strategy: CenterStrategy,
248 pub order_s: f64,
251 pub alpha: f64,
253 pub tau0: f64,
257 pub num_scales: usize,
259 pub length_scale: f64,
262 pub double_penalty: bool,
265 #[serde(default = "measure_jet_learn_length_scale_default")]
273 pub learn_length_scale: bool,
274 #[serde(default)]
282 pub multiscale: bool,
283 #[serde(default)]
285 pub identifiability: MeasureJetIdentifiability,
286 #[serde(default)]
289 pub frozen_quadrature: Option<MeasureJetFrozenQuadrature>,
290}
291
292impl Default for MeasureJetBasisSpec {
293 fn default() -> Self {
294 Self {
295 center_strategy: CenterStrategy::FarthestPoint { num_centers: 50 },
296 order_s: 0.0,
297 alpha: 1.0,
309 tau0: 1e-3,
310 num_scales: 0,
311 length_scale: 0.0,
312 double_penalty: true,
313 learn_length_scale: false,
314 multiscale: false,
315 identifiability: MeasureJetIdentifiability::CenterSumToZero,
316 frozen_quadrature: None,
317 }
318 }
319}
320
321pub struct MeasureJetBand {
324 pub eps: Vec<f64>,
325 pub log_step: f64,
326}
327
328pub struct MeasureJetEnergyJets {
334 pub q: Array2<f64>,
335 pub dq_ds: Array2<f64>,
336 pub d2q_ds2: Array2<f64>,
337 pub dq_dalpha: Array2<f64>,
338 pub d2q_dalpha2: Array2<f64>,
339 pub d2q_ds_dalpha: Array2<f64>,
340 pub dq_dlogtau: Array2<f64>,
341 pub d2q_dlogtau2: Array2<f64>,
342 pub d2q_ds_dlogtau: Array2<f64>,
343 pub d2q_dalpha_dlogtau: Array2<f64>,
344}
345
346pub(crate) fn householder_sum_to_zero_u(m: usize) -> Array1<f64> {
354 let c = 1.0 / (m as f64).sqrt();
355 let mut u = Array1::<f64>::from_elem(m, c);
356 u[0] -= 1.0;
357 let norm = u.dot(&u).sqrt();
358 u.mapv_inplace(|v| v / norm);
359 u
360}
361
362pub(crate) fn householder_sum_to_zero_z(u: &Array1<f64>) -> Array2<f64> {
366 let m = u.len();
367 let mut z = Array2::<f64>::zeros((m, m - 1));
368 for j in 0..(m - 1) {
369 for i in 0..m {
370 let h = if i == j + 1 { 1.0 } else { 0.0 } - 2.0 * u[i] * u[j + 1];
371 z[(i, j)] = h;
372 }
373 }
374 z
375}
376
377pub(crate) fn symmetric_pseudoinverse(
378 a: &Array2<f64>,
379 label: &str,
380) -> Result<Array2<f64>, BasisError> {
381 let n = a.nrows();
382 if a.ncols() != n {
383 crate::bail_dim_basis!(
384 "measure-jet pseudo-inverse `{label}` needs a square matrix, got {:?}",
385 a.dim()
386 );
387 }
388 let (evals, evecs) = a.eigh(Side::Lower).map_err(|e| {
389 BasisError::InvalidInput(format!(
390 "measure-jet pseudo-inverse `{label}` eigendecomposition failed: {e}"
391 ))
392 })?;
393 let lam_max = evals.iter().fold(0.0_f64, |acc, v| acc.max((*v).max(0.0)));
394 let rank_tol = MEASURE_JET_PSEUDOINVERSE_RTOL * (n.max(1) as f64) * lam_max;
395 let mut scaled = evecs.clone();
396 for (k, mut col) in scaled.axis_iter_mut(Axis(1)).enumerate() {
397 let lam = evals[k].max(0.0);
398 let inv = if lam > rank_tol { 1.0 / lam } else { 0.0 };
399 col.mapv_inplace(|v| v * inv);
400 }
401 Ok(scaled.dot(&evecs.t()))
402}
403
404fn measure_jet_affine_value_basis(
411 centers: ArrayView2<'_, f64>,
412 masses: ArrayView1<'_, f64>,
413) -> Array2<f64> {
414 let m = centers.nrows();
415 let head_transform = measure_jet_affine_head_transform(centers, masses);
416 let head_rank = head_transform.ncols();
417 let mut affine = Array2::<f64>::ones((m, head_rank + 1));
418 if head_rank > 0 {
419 affine
420 .slice_mut(ndarray::s![.., 1..])
421 .assign(¢ers.dot(&head_transform));
422 }
423 affine
424}
425
426fn affine_function_nullspace_form(
434 centers: ArrayView2<'_, f64>,
435 masses: ArrayView1<'_, f64>,
436) -> Result<Array2<f64>, BasisError> {
437 let m = centers.nrows();
438 if masses.len() != m {
439 crate::bail_dim_basis!(
440 "measure-jet affine function-space form shape mismatch: centers {:?}, masses {}",
441 centers.dim(),
442 masses.len()
443 );
444 }
445 let affine = measure_jet_affine_value_basis(centers, masses);
446 let mut weighted_affine = affine.clone();
447 for (i, mut row) in weighted_affine.outer_iter_mut().enumerate() {
448 row.mapv_inplace(|v| v * masses[i]);
449 }
450 let affine_gram = affine.t().dot(&weighted_affine);
451 let affine_gram_pinv = symmetric_pseudoinverse(&affine_gram, "affine function-space Gram")?;
452 let form = weighted_affine
453 .dot(&affine_gram_pinv)
454 .dot(&weighted_affine.t());
455 Ok((&form + &form.t()) * 0.5)
456}
457
458fn pullback_center_form(evaluation: &Array2<f64>, form: &Array2<f64>) -> Array2<f64> {
460 let pulled = evaluation.t().dot(form).dot(evaluation);
461 (&pulled + &pulled.t()) * 0.5
462}
463
464fn pullback_center_form_log_length_jets(
467 evaluation: &Array2<f64>,
468 evaluation_first: &Array2<f64>,
469 evaluation_second: &Array2<f64>,
470 form: &Array2<f64>,
471) -> (Array2<f64>, Array2<f64>) {
472 let h_e = form.dot(evaluation);
473 let h_e_first = form.dot(evaluation_first);
474 let h_e_second = form.dot(evaluation_second);
475 let first_raw = evaluation_first.t().dot(&h_e) + evaluation.t().dot(&h_e_first);
476 let second_raw = evaluation_second.t().dot(&h_e)
477 + evaluation.t().dot(&h_e_second)
478 + evaluation_first.t().dot(&h_e_first) * 2.0;
479 (
480 (&first_raw + &first_raw.t()) * 0.5,
481 (&second_raw + &second_raw.t()) * 0.5,
482 )
483}
484
485fn pullback_center_form_log_length_cross(
487 evaluation: &Array2<f64>,
488 evaluation_first: &Array2<f64>,
489 form_first: &Array2<f64>,
490) -> Array2<f64> {
491 let h_e = form_first.dot(evaluation);
492 let h_e_first = form_first.dot(evaluation_first);
493 let cross_raw = evaluation_first.t().dot(&h_e) + evaluation.t().dot(&h_e_first);
494 (&cross_raw + &cross_raw.t()) * 0.5
495}
496
497pub(crate) fn affine_function_nullspace_penalty(
502 evaluation: &Array2<f64>,
503 centers: ArrayView2<'_, f64>,
504 masses: ArrayView1<'_, f64>,
505) -> Result<Array2<f64>, BasisError> {
506 if evaluation.nrows() != centers.nrows() {
507 crate::bail_dim_basis!(
508 "measure-jet affine function-space penalty shape mismatch: evaluation {:?}, centers {:?}",
509 evaluation.dim(),
510 centers.dim()
511 );
512 }
513 let form = affine_function_nullspace_form(centers, masses)?;
514 Ok(pullback_center_form(evaluation, &form))
515}
516
517pub(crate) fn pairwise_sq_dists(a: ArrayView2<'_, f64>, b: ArrayView2<'_, f64>) -> Array2<f64> {
528 let an: Vec<f64> = a.outer_iter().map(|r| r.dot(&r)).collect();
529 let bn: Vec<f64> = b.outer_iter().map(|r| r.dot(&r)).collect();
530 let mut g = a.dot(&b.t());
531 g.axis_iter_mut(Axis(0))
532 .into_par_iter()
533 .enumerate()
534 .for_each(|(i, mut row)| {
535 for (j, v) in row.iter_mut().enumerate() {
536 *v = (an[i] + bn[j] - 2.0 * *v).max(0.0);
537 }
538 });
539 g
540}
541
542pub(crate) const MEASURE_JET_ASSIGN_BLOCK_ROWS: usize = 65_536;
546
547pub(crate) fn validate_finite_points(
548 points: ArrayView2<'_, f64>,
549 what: &str,
550) -> Result<(), BasisError> {
551 for (i, row) in points.outer_iter().enumerate() {
552 if row.iter().any(|v| !v.is_finite()) {
553 crate::bail_invalid_basis!("measure-jet {what} row {i} has a non-finite coordinate");
554 }
555 }
556 Ok(())
557}
558
559pub(crate) fn median_nearest_center_spacing(dist2: &Array2<f64>) -> Result<f64, BasisError> {
562 let m = dist2.nrows();
563 if m < 2 {
564 return Err(BasisError::InsufficientColumnsForConstraint { found: m });
565 }
566 let mut nearest: Vec<f64> = Vec::with_capacity(m);
567 for i in 0..m {
568 let mut best = f64::INFINITY;
569 for j in 0..m {
570 if j != i && dist2[(i, j)] < best {
571 best = dist2[(i, j)];
572 }
573 }
574 nearest.push(best.sqrt());
575 }
576 nearest.sort_by(|a, b| a.partial_cmp(b).expect("finite center spacings"));
577 let median = nearest[nearest.len() / 2];
578 if !(median.is_finite() && median > 0.0) {
579 crate::bail_invalid_basis!(
580 "measure-jet centers are degenerate (median nearest-center spacing = {median}); \
581 duplicate centers cannot carry a scale band"
582 );
583 }
584 Ok(median)
585}
586
587pub fn measure_jet_band(
595 centers: ArrayView2<'_, f64>,
596 num_scales: usize,
597) -> Result<MeasureJetBand, BasisError> {
598 validate_finite_points(centers, "centers")?;
599 let dist2 = pairwise_sq_dists(centers, centers);
600 let eps_min = median_nearest_center_spacing(&dist2)?;
601 let d = centers.ncols();
603 let mut diag2 = 0.0_f64;
604 for k in 0..d {
605 let col = centers.column(k);
606 let mut lo = f64::INFINITY;
607 let mut hi = f64::NEG_INFINITY;
608 for &v in col.iter() {
609 lo = lo.min(v);
610 hi = hi.max(v);
611 }
612 diag2 += (hi - lo) * (hi - lo);
613 }
614 let eps_max = 0.5 * diag2.sqrt();
615 if !(eps_max.is_finite() && eps_max > eps_min) {
616 return Ok(MeasureJetBand {
617 eps: vec![eps_min],
618 log_step: std::f64::consts::LN_2,
619 });
620 }
621 let auto = ((eps_max / eps_min).log2().ceil() as usize + 1)
622 .clamp(MEASURE_JET_MIN_AUTO_SCALES, MEASURE_JET_MAX_AUTO_SCALES);
623 let count = if num_scales == 0 { auto } else { num_scales };
624 if count == 1 {
625 return Ok(MeasureJetBand {
626 eps: vec![eps_min],
627 log_step: std::f64::consts::LN_2,
628 });
629 }
630 let ratio = (eps_max / eps_min).powf(1.0 / (count as f64 - 1.0));
631 let mut eps = Vec::with_capacity(count);
632 let mut e = eps_min;
633 for _ in 0..count {
634 eps.push(e);
635 e *= ratio;
636 }
637 Ok(MeasureJetBand {
638 eps,
639 log_step: ratio.ln(),
640 })
641}
642
643pub fn measure_jet_quadrature_nodes(
650 data: ArrayView2<'_, f64>,
651 centers: ArrayView2<'_, f64>,
652) -> Result<(Array2<f64>, Array1<f64>), BasisError> {
653 if data.ncols() != centers.ncols() {
654 crate::bail_dim_basis!(
655 "measure-jet mass assignment dimension mismatch: data d={} centers d={}",
656 data.ncols(),
657 centers.ncols()
658 );
659 }
660 validate_finite_points(data, "data")?;
661 validate_finite_points(centers, "centers")?;
662 let n = data.nrows();
663 let m = centers.nrows();
664 let d = centers.ncols();
665 if n == 0 || m == 0 {
666 crate::bail_invalid_basis!("measure-jet mass assignment needs nonempty data and centers");
667 }
668 let cn: Vec<f64> = centers.outer_iter().map(|r| r.dot(&r)).collect();
673 let assignments: Vec<usize> = (0..n)
674 .step_by(MEASURE_JET_ASSIGN_BLOCK_ROWS)
675 .flat_map(|start| {
676 let end = (start + MEASURE_JET_ASSIGN_BLOCK_ROWS).min(n);
677 let g = data.slice(ndarray::s![start..end, ..]).dot(¢ers.t());
678 let block: Vec<usize> = g
679 .axis_iter(Axis(0))
680 .into_par_iter()
681 .map(|row| {
682 let mut best_j = 0usize;
683 let mut best = f64::INFINITY;
684 for (j, &gij) in row.iter().enumerate() {
685 let s = cn[j] - 2.0 * gij;
686 if s < best {
687 best = s;
688 best_j = j;
689 }
690 }
691 best_j
692 })
693 .collect();
694 block
695 })
696 .collect();
697 let mut masses = Array1::<f64>::zeros(m);
698 let mut nodes = centers.to_owned();
699 let mut sums = Array2::<f64>::zeros((m, d));
700 let unit = 1.0 / n as f64;
701 for (i, &j) in assignments.iter().enumerate() {
702 masses[j] += unit;
703 for k in 0..d {
704 sums[(j, k)] += data[(i, k)];
705 }
706 }
707 let mut barycenter = sums;
710 for j in 0..m {
711 let count = masses[j] * n as f64;
712 if count > 0.0 {
713 for k in 0..d {
714 barycenter[(j, k)] /= count;
715 nodes[(j, k)] = barycenter[(j, k)];
716 }
717 }
718 }
719 Ok((nodes, masses))
720}
721
722pub fn measure_jet_center_masses(
725 data: ArrayView2<'_, f64>,
726 centers: ArrayView2<'_, f64>,
727) -> Result<Array1<f64>, BasisError> {
728 measure_jet_quadrature_nodes(data, centers).map(|(_, masses)| masses)
729}
730
731pub(crate) fn assemble_weighted_forms<F>(
754 centers: ArrayView2<'_, f64>,
755 masses: ArrayView1<'_, f64>,
756 band: &MeasureJetBand,
757 order_s: f64,
758 alpha: f64,
759 tau0: f64,
760 n_forms: usize,
761 channels: usize,
762 weights: &F,
763) -> Result<Vec<Array2<f64>>, BasisError>
764where
765 F: Fn(usize, f64, f64, f64, &mut [[f64; 3]]) + Sync,
766{
767 let m = centers.nrows();
768 let d = centers.ncols();
769 if n_forms == 0 || !(1..=3).contains(&channels) {
770 crate::bail_invalid_basis!(
771 "measure-jet assembly needs at least one output form and 1..=3 block channels"
772 );
773 }
774 if masses.len() != m {
775 crate::bail_dim_basis!(
776 "measure-jet energy mass/center mismatch: {} masses for {} centers",
777 masses.len(),
778 m
779 );
780 }
781 if band.eps.is_empty() || band.eps.iter().any(|e| !(e.is_finite() && *e > 0.0)) {
782 crate::bail_invalid_basis!("measure-jet energy needs a nonempty positive scale band");
783 }
784 if !(order_s.is_finite() && order_s > 0.0 && order_s < 2.0) {
785 crate::bail_invalid_basis!(
786 "measure-jet order s must lie in (0, 2) for the affine-jet energy; got {order_s}"
787 );
788 }
789 if !(alpha.is_finite() && tau0.is_finite() && tau0 >= 0.0) {
790 crate::bail_invalid_basis!(
791 "measure-jet energy needs finite alpha and finite tau0 >= 0; got alpha={alpha}, tau0={tau0}"
792 );
793 }
794 if masses.iter().any(|v| !(v.is_finite() && *v >= 0.0)) {
795 crate::bail_invalid_basis!("measure-jet energy needs finite nonnegative center masses");
796 }
797 let dist2 = pairwise_sq_dists(centers, centers);
798
799 let assemble_scale = |scale_idx: usize, eps: f64| -> Result<Vec<Array2<f64>>, BasisError> {
804 let mut out: Vec<Array2<f64>> =
805 (0..n_forms).map(|_| Array2::<f64>::zeros((m, m))).collect();
806 let cutoff2 = (MEASURE_JET_PROFILE_CUTOFF * eps) * (MEASURE_JET_PROFILE_CUTOFF * eps);
807 let inv_two_eps2 = 1.0 / (2.0 * eps * eps);
808 let eta = 2.0 * order_s + (d as f64) * (2.0 - 2.0 * alpha);
809 let scale_weight = band.log_step * eps.powf(-eta);
810 let net_radius2 = 0.25 * eps * eps;
814 let mut outer: Vec<usize> = Vec::new();
815 for i in 0..m {
816 if masses[i] <= 0.0 {
817 continue;
818 }
819 let covered = outer.iter().any(|&o| dist2[(i, o)] <= net_radius2);
820 if !covered {
821 outer.push(i);
822 }
823 }
824 let mut net_mass = vec![0.0_f64; m];
825 for i in 0..m {
826 if masses[i] <= 0.0 {
827 continue;
828 }
829 let mut best = f64::INFINITY;
830 let mut best_o = usize::MAX;
831 for &o in &outer {
832 if dist2[(i, o)] < best {
833 best = dist2[(i, o)];
834 best_o = o;
835 }
836 }
837 if best_o != usize::MAX {
838 net_mass[best_o] += masses[i];
839 }
840 }
841 let mut wbuf = vec![[0.0_f64; 3]; n_forms];
842 for &i in &outer {
843 let mut idx: Vec<usize> = Vec::new();
845 for j in 0..m {
846 if dist2[(i, j)] <= cutoff2 {
847 idx.push(j);
848 }
849 }
850 let ml = idx.len();
851 let mut w = Array1::<f64>::zeros(ml);
853 let mut q = 0.0_f64;
854 for (a, &j) in idx.iter().enumerate() {
855 let wj = masses[j] * (-dist2[(i, j)] * inv_two_eps2).exp();
856 w[a] = wj;
857 q += wj;
858 }
859 if !(q > 0.0) {
860 continue;
861 }
862 let mut phi = Array2::<f64>::zeros((ml, d));
864 for (a, &j) in idx.iter().enumerate() {
865 for k in 0..d {
866 phi[(a, k)] = (centers[(j, k)] - centers[(i, k)]) / eps;
867 }
868 }
869 let a_mean = phi.t().dot(&w) / q;
870 let mut wphi = phi.clone();
872 for (a, mut row) in wphi.outer_iter_mut().enumerate() {
873 row.mapv_inplace(|v| v * w[a]);
874 }
875 let mut b = wphi.clone();
876 for (a, mut row) in b.outer_iter_mut().enumerate() {
877 for k in 0..d {
878 row[k] -= w[a] * a_mean[k];
879 }
880 }
881 let mut g = phi.t().dot(&wphi);
882 g.mapv_inplace(|v| v / q);
883 for r in 0..d {
884 for c in 0..d {
885 g[(r, c)] -= a_mean[r] * a_mean[c];
886 }
887 }
888 let g_pinv = symmetric_pseudoinverse(&g, "local affine Gram")?;
889 let bm = b.dot(&g_pinv);
890 let base = scale_weight * net_mass[i] * q.powf(1.0 - 2.0 * alpha);
891 weights(scale_idx, eps, q, base, &mut wbuf);
892 for (a, &ja) in idx.iter().enumerate() {
895 let bma = bm.row(a);
896 for (c, &jc) in idx.iter().enumerate() {
897 let b_c = b.row(c);
898 let mut val_r = -w[a] * w[c] / q - bma.dot(&b_c) / q;
899 if a == c {
900 val_r += w[a];
901 }
902 for (k, out_k) in out.iter_mut().enumerate() {
903 let wk = wbuf[k];
904 out_k[(ja, jc)] += wk[0] * val_r;
905 }
906 }
907 }
908 }
909 Ok(out)
910 };
911
912 let n_scales = band.eps.len();
913 let parallel_ok = m
914 .saturating_mul(m)
915 .saturating_mul(n_scales)
916 .saturating_mul(n_forms)
917 <= MEASURE_JET_PARALLEL_FORM_BUDGET_DOUBLES;
918 let per_scale: Vec<Vec<Array2<f64>>> = if parallel_ok {
919 band.eps
920 .par_iter()
921 .enumerate()
922 .map(|(scale_idx, &eps)| assemble_scale(scale_idx, eps))
923 .collect::<Result<Vec<_>, BasisError>>()?
924 } else {
925 band.eps
926 .iter()
927 .enumerate()
928 .map(|(scale_idx, &eps)| assemble_scale(scale_idx, eps))
929 .collect::<Result<Vec<_>, BasisError>>()?
930 };
931
932 let mut totals: Vec<Array2<f64>> = (0..n_forms).map(|_| Array2::<f64>::zeros((m, m))).collect();
933 for scale_forms in per_scale {
934 for (total, part) in totals.iter_mut().zip(scale_forms) {
935 *total += ∂
936 }
937 }
938 Ok(totals.into_iter().map(|t| (&t + &t.t()) * 0.5).collect())
940}
941
942pub fn measure_jet_energy_form(
957 centers: ArrayView2<'_, f64>,
958 masses: ArrayView1<'_, f64>,
959 band: &MeasureJetBand,
960 order_s: f64,
961 alpha: f64,
962 tau0: f64,
963) -> Result<Array2<f64>, BasisError> {
964 let mut forms = assemble_weighted_forms(
965 centers,
966 masses,
967 band,
968 order_s,
969 alpha,
970 tau0,
971 1,
972 1,
973 &|_, _, _, base, out: &mut [[f64; 3]]| out[0] = [base, 0.0, 0.0],
974 )?;
975 let q = forms.swap_remove(0);
976 project_symmetric_psd(q, "measure-jet energy form")
984}
985
986pub(crate) fn project_symmetric_psd(
992 a: Array2<f64>,
993 label: &str,
994) -> Result<Array2<f64>, BasisError> {
995 let n = a.nrows();
996 if n == 0 {
997 return Ok(a);
998 }
999 let (evals, evecs) = a.eigh(Side::Lower).map_err(|e| {
1000 BasisError::InvalidInput(format!(
1001 "measure-jet PSD projection `{label}` eigendecomposition failed: {e}"
1002 ))
1003 })?;
1004 if evals.iter().all(|&lam| lam >= 0.0) {
1005 return Ok(a);
1006 }
1007 let mut scaled = evecs.clone();
1008 for (k, mut col) in scaled.axis_iter_mut(Axis(1)).enumerate() {
1009 let lam = evals[k].max(0.0);
1010 col.mapv_inplace(|v| v * lam);
1011 }
1012 let psd = scaled.dot(&evecs.t());
1013 Ok((&psd + &psd.t()) * 0.5)
1014}
1015
1016pub fn measure_jet_energy_form_with_jets(
1031 centers: ArrayView2<'_, f64>,
1032 masses: ArrayView1<'_, f64>,
1033 band: &MeasureJetBand,
1034 order_s: f64,
1035 alpha: f64,
1036 tau0: f64,
1037) -> Result<MeasureJetEnergyJets, BasisError> {
1038 if !(tau0.is_finite() && tau0 > 0.0) {
1039 crate::bail_invalid_basis!(
1040 "measure-jet jets need tau0 > 0 because the retained τ coordinate is ln τ; got {tau0}"
1041 );
1042 }
1043 let mut forms = assemble_weighted_forms(
1044 centers,
1045 masses,
1046 band,
1047 order_s,
1048 alpha,
1049 tau0,
1050 10,
1051 3,
1052 &|_, eps: f64, q: f64, base: f64, out: &mut [[f64; 3]]| {
1053 let gs = -2.0 * eps.ln();
1054 let intrinsic_dim = centers.ncols() as f64;
1055 let ga = 2.0 * intrinsic_dim * eps.ln() - 2.0 * q.max(f64::MIN_POSITIVE).ln();
1056 out[0] = [base, 0.0, 0.0];
1057 out[1] = [gs * base, 0.0, 0.0];
1058 out[2] = [gs * gs * base, 0.0, 0.0];
1059 out[3] = [ga * base, 0.0, 0.0];
1060 out[4] = [ga * ga * base, 0.0, 0.0];
1061 out[5] = [gs * ga * base, 0.0, 0.0];
1062 out[6] = [0.0, 0.0, 0.0];
1063 out[7] = [0.0, 0.0, 0.0];
1064 out[8] = [0.0, 0.0, 0.0];
1065 out[9] = [0.0, 0.0, 0.0];
1066 },
1067 )?;
1068 let d2q_dalpha_dlogtau = forms.pop().expect("ten assembled forms");
1069 let d2q_ds_dlogtau = forms.pop().expect("ten assembled forms");
1070 let d2q_dlogtau2 = forms.pop().expect("ten assembled forms");
1071 let dq_dlogtau = forms.pop().expect("ten assembled forms");
1072 let d2q_ds_dalpha = forms.pop().expect("ten assembled forms");
1073 let d2q_dalpha2 = forms.pop().expect("ten assembled forms");
1074 let dq_dalpha = forms.pop().expect("ten assembled forms");
1075 let d2q_ds2 = forms.pop().expect("ten assembled forms");
1076 let dq_ds = forms.pop().expect("ten assembled forms");
1077 let q = forms.pop().expect("ten assembled forms");
1078 Ok(MeasureJetEnergyJets {
1079 q,
1080 dq_ds,
1081 d2q_ds2,
1082 dq_dalpha,
1083 d2q_dalpha2,
1084 d2q_ds_dalpha,
1085 dq_dlogtau,
1086 d2q_dlogtau2,
1087 d2q_ds_dlogtau,
1088 d2q_dalpha_dlogtau,
1089 })
1090}
1091
1092pub fn measure_jet_scale_spectrum(
1098 centers: ArrayView2<'_, f64>,
1099 masses: ArrayView1<'_, f64>,
1100 band: &MeasureJetBand,
1101 order_s: f64,
1102 alpha: f64,
1103 tau0: f64,
1104 values: ArrayView1<'_, f64>,
1105) -> Result<Vec<f64>, BasisError> {
1106 if values.len() != centers.nrows() {
1107 crate::bail_dim_basis!(
1108 "measure-jet scale spectrum needs one value per center: {} values for {} centers",
1109 values.len(),
1110 centers.nrows()
1111 );
1112 }
1113 let forms = measure_jet_energy_forms_per_scale(centers, masses, band, order_s, alpha, tau0)?;
1114 Ok(forms
1115 .iter()
1116 .map(|q_l| values.dot(&q_l.dot(&values)))
1117 .collect())
1118}
1119
1120pub fn measure_jet_energy_forms_per_scale(
1126 centers: ArrayView2<'_, f64>,
1127 masses: ArrayView1<'_, f64>,
1128 band: &MeasureJetBand,
1129 order_s: f64,
1130 alpha: f64,
1131 tau0: f64,
1132) -> Result<Vec<Array2<f64>>, BasisError> {
1133 let n_scales = band.eps.len();
1134 assemble_weighted_forms(
1135 centers,
1136 masses,
1137 band,
1138 order_s,
1139 alpha,
1140 tau0,
1141 n_scales,
1142 1,
1143 &|scale_idx, _, _, base, out: &mut [[f64; 3]]| {
1144 for (k, slot) in out.iter_mut().enumerate() {
1145 *slot = if k == scale_idx {
1146 [base, 0.0, 0.0]
1147 } else {
1148 [0.0, 0.0, 0.0]
1149 };
1150 }
1151 },
1152 )
1153}
1154
1155pub fn measure_jet_support_curve(
1163 queries: ArrayView2<'_, f64>,
1164 centers: ArrayView2<'_, f64>,
1165 masses: ArrayView1<'_, f64>,
1166 eps_band: &[f64],
1167) -> Result<Array2<f64>, BasisError> {
1168 if queries.ncols() != centers.ncols() {
1169 crate::bail_dim_basis!(
1170 "measure-jet support curve dimension mismatch: queries d={} centers d={}",
1171 queries.ncols(),
1172 centers.ncols()
1173 );
1174 }
1175 if masses.len() != centers.nrows() {
1176 crate::bail_dim_basis!(
1177 "measure-jet support curve mass/center mismatch: {} masses for {} centers",
1178 masses.len(),
1179 centers.nrows()
1180 );
1181 }
1182 if eps_band.is_empty() || eps_band.iter().any(|e| !(e.is_finite() && *e > 0.0)) {
1183 crate::bail_invalid_basis!("measure-jet support curve needs a nonempty positive band");
1184 }
1185 validate_finite_points(queries, "queries")?;
1186 validate_finite_points(centers, "centers")?;
1187 let nq = queries.nrows();
1188 let nl = eps_band.len();
1189 let d2 = pairwise_sq_dists(queries, centers);
1192 let mut out = Array2::<f64>::zeros((nq, nl));
1193 out.axis_iter_mut(Axis(0))
1194 .into_par_iter()
1195 .enumerate()
1196 .for_each(|(qi, mut row)| {
1197 let d2_row = d2.row(qi);
1198 for (li, &eps) in eps_band.iter().enumerate() {
1199 let inv_two_eps2 = 1.0 / (2.0 * eps * eps);
1200 let mut acc = 0.0_f64;
1201 for (j, &dd) in d2_row.iter().enumerate() {
1202 acc += masses[j] * (-dd * inv_two_eps2).exp();
1203 }
1204 row[li] = acc;
1205 }
1206 });
1207 Ok(out)
1208}
1209
1210pub(crate) fn measure_jet_support_means(
1211 centers: ArrayView2<'_, f64>,
1212 masses: ArrayView1<'_, f64>,
1213 eps_band: &[f64],
1214) -> Result<Vec<f64>, BasisError> {
1215 let total_mass = masses.sum();
1216 if !(total_mass.is_finite() && total_mass > 0.0) {
1217 crate::bail_invalid_basis!(
1218 "measure-jet support means need positive finite total mass; got {total_mass}"
1219 );
1220 }
1221 let support = measure_jet_support_curve(centers, centers, masses, eps_band)?;
1222 let mut means = vec![0.0_f64; eps_band.len()];
1223 for (i, row) in support.rows().into_iter().enumerate() {
1224 let mass = masses[i];
1225 for (mean, &q) in means.iter_mut().zip(row.iter()) {
1226 *mean += mass * q;
1227 }
1228 }
1229 for mean in &mut means {
1230 *mean /= total_mass;
1231 if !(*mean).is_finite() || *mean <= 0.0 {
1232 crate::bail_invalid_basis!(
1233 "measure-jet support mean must be positive and finite; got {mean}"
1234 );
1235 }
1236 }
1237 Ok(means)
1238}
1239
1240pub fn measure_jet_design_matrix(
1242 data: ArrayView2<'_, f64>,
1243 centers: ArrayView2<'_, f64>,
1244 length_scale: f64,
1245) -> Result<Array2<f64>, BasisError> {
1246 if data.ncols() != centers.ncols() {
1247 crate::bail_dim_basis!(
1248 "measure-jet design dimension mismatch: data d={} centers d={}",
1249 data.ncols(),
1250 centers.ncols()
1251 );
1252 }
1253 if !(length_scale.is_finite() && length_scale > 0.0) {
1254 crate::bail_invalid_basis!(
1255 "measure-jet design needs a positive finite length_scale; got {length_scale}"
1256 );
1257 }
1258 validate_finite_points(data, "data")?;
1259 validate_finite_points(centers, "centers")?;
1260 let inv_two_l2 = 1.0 / (2.0 * length_scale * length_scale);
1261 let mut out = pairwise_sq_dists(data, centers);
1264 out.axis_iter_mut(Axis(0))
1265 .into_par_iter()
1266 .for_each(|mut row| {
1267 row.mapv_inplace(|d2| (-d2 * inv_two_l2).exp());
1268 });
1269 Ok(out)
1270}
1271
1272fn measure_jet_design_log_length_jets(
1275 data: ArrayView2<'_, f64>,
1276 centers: ArrayView2<'_, f64>,
1277 length_scale: f64,
1278) -> Result<(Array2<f64>, Array2<f64>), BasisError> {
1279 let kernel = measure_jet_design_matrix(data, centers, length_scale)?;
1280 let squared_distances = pairwise_sq_dists(data, centers);
1281 let inv_l2 = 1.0 / (length_scale * length_scale);
1282 let mut first = kernel.clone();
1283 let mut second = kernel;
1284 for ((first_value, second_value), &distance_squared) in first
1285 .iter_mut()
1286 .zip(second.iter_mut())
1287 .zip(squared_distances.iter())
1288 {
1289 let a = distance_squared * inv_l2;
1290 let kernel_value = *first_value;
1291 *first_value = kernel_value * a;
1292 *second_value = kernel_value * (a * a - 2.0 * a);
1293 }
1294 Ok((first, second))
1295}
1296
1297pub fn measure_jet_affine_head_transform(
1328 centers: ArrayView2<'_, f64>,
1329 masses: ArrayView1<'_, f64>,
1330) -> Array2<f64> {
1331 let m = centers.nrows();
1332 let d = centers.ncols();
1333 let total_mass = masses.sum();
1334 let mdot = |u: &Array1<f64>, v: &Array1<f64>| -> f64 {
1336 let mut acc = 0.0;
1337 for i in 0..m {
1338 acc += masses[i] * u[i] * v[i];
1339 }
1340 acc
1341 };
1342 let cols: Vec<Array1<f64>> = (0..d)
1346 .map(|k| {
1347 let col = centers.column(k).to_owned();
1348 let mean = if total_mass > 0.0 {
1349 mdot(&col, &Array1::ones(m)) / total_mass
1350 } else {
1351 0.0
1352 };
1353 col.mapv(|x| x - mean)
1354 })
1355 .collect();
1356 let max_norm = cols
1358 .iter()
1359 .fold(0.0_f64, |acc, c| acc.max(mdot(c, c).sqrt()));
1360 let drop_below =
1361 (MEASURE_JET_PSEUDOINVERSE_RTOL * (d.max(1) as f64) * max_norm).max(f64::MIN_POSITIVE);
1362 let mut q_cols: Vec<Array1<f64>> = Vec::new();
1366 let mut t_cols: Vec<Array1<f64>> = Vec::new();
1367 for k in 0..d {
1368 let mut v = cols[k].clone();
1369 let mut t = Array1::<f64>::zeros(d);
1370 t[k] = 1.0;
1371 for (q, tq) in q_cols.iter().zip(t_cols.iter()) {
1372 let proj = mdot(q, &v);
1373 v.scaled_add(-proj, q);
1374 t.scaled_add(-proj, tq);
1375 }
1376 let norm = mdot(&v, &v).sqrt();
1377 if norm > drop_below {
1378 v.mapv_inplace(|x| x / norm);
1379 t.mapv_inplace(|x| x / norm);
1380 q_cols.push(v);
1381 t_cols.push(t);
1382 }
1383 }
1384 let head_rank = t_cols.len();
1385 let mut t_mat = Array2::<f64>::zeros((d, head_rank));
1386 for (r, t) in t_cols.into_iter().enumerate() {
1387 t_mat.column_mut(r).assign(&t);
1388 }
1389 t_mat
1390}
1391
1392pub fn realized_measure_jet_length_scale(
1397 centers: ArrayView2<'_, f64>,
1398 spec_length_scale: f64,
1399) -> Result<f64, BasisError> {
1400 if spec_length_scale.is_finite() && spec_length_scale > 0.0 {
1401 return Ok(spec_length_scale);
1402 }
1403 if spec_length_scale != 0.0 {
1404 crate::bail_invalid_basis!(
1405 "measure-jet length_scale must be positive (or 0.0 for auto); got {spec_length_scale}"
1406 );
1407 }
1408 let dist2 = pairwise_sq_dists(centers, centers);
1409 let spacing = median_nearest_center_spacing(&dist2)?;
1410 Ok(MEASURE_JET_AUTO_LENGTH_SCALE_FACTOR * spacing)
1411}
1412
1413pub(crate) struct RealizedMeasureJetGeometry {
1418 pub(crate) centers: Array2<f64>,
1419 pub(crate) masses: Array1<f64>,
1420 pub(crate) eps_band: Vec<f64>,
1421 pub(crate) log_step: f64,
1422 pub(crate) length_scale: f64,
1423 pub(crate) order_s_eval: f64,
1427 pub(crate) per_level: bool,
1429 pub(crate) z: Array2<f64>,
1430 pub(crate) coefficient_gauge: gam_problem::Gauge,
1431 pub(crate) kz: Array2<f64>,
1432 pub(crate) head_transform: Array2<f64>,
1438}
1439
1440pub(crate) fn realize_measure_jet_geometry(
1441 data: ArrayView2<'_, f64>,
1442 spec: &MeasureJetBasisSpec,
1443) -> Result<RealizedMeasureJetGeometry, BasisError> {
1444 if data.ncols() == 0 {
1445 crate::bail_invalid_basis!("measure-jet smooth needs at least one feature column");
1446 }
1447 validate_finite_points(data, "data")?;
1448 let seed_centers = select_centers_by_strategy(data, &spec.center_strategy)?;
1449 let m = seed_centers.nrows();
1450 if m < 3 {
1451 return Err(BasisError::InsufficientColumnsForConstraint { found: m });
1452 }
1453 let order_s = if spec.order_s == 0.0 {
1454 MEASURE_JET_DEFAULT_ORDER_S
1455 } else {
1456 spec.order_s
1457 };
1458 let (centers, masses, eps_band, log_step) = match &spec.frozen_quadrature {
1465 Some(frozen) => {
1466 if frozen.masses.len() != m {
1467 crate::bail_dim_basis!(
1468 "frozen measure-jet quadrature mismatch: {} masses for {} centers",
1469 frozen.masses.len(),
1470 m
1471 );
1472 }
1473 if frozen.eps_band.is_empty() {
1474 crate::bail_invalid_basis!("frozen measure-jet quadrature has an empty band");
1475 }
1476 let log_step = if frozen.eps_band.len() >= 2 {
1477 (frozen.eps_band[1] / frozen.eps_band[0]).ln()
1478 } else {
1479 std::f64::consts::LN_2
1480 };
1481 (
1482 seed_centers,
1483 frozen.masses.clone(),
1484 frozen.eps_band.clone(),
1485 log_step,
1486 )
1487 }
1488 None => {
1489 let (nodes, masses) = measure_jet_quadrature_nodes(data, seed_centers.view())?;
1490 let band = measure_jet_band(nodes.view(), spec.num_scales)?;
1491 (nodes, masses, band.eps, band.log_step)
1492 }
1493 };
1494 let length_scale = realized_measure_jet_length_scale(centers.view(), spec.length_scale)?;
1495 let head_transform = if spec.multiscale {
1506 Array2::<f64>::zeros((centers.ncols(), 0))
1507 } else {
1508 measure_jet_affine_head_transform(centers.view(), masses.view())
1509 };
1510 let head_rank = head_transform.ncols();
1511 let m_aug = m + head_rank;
1512 let k_cc = measure_jet_design_matrix(centers.view(), centers.view(), length_scale)?;
1513 let head_cc = centers.dot(&head_transform);
1514 let (z, coefficient_gauge) = match &spec.identifiability {
1529 MeasureJetIdentifiability::FrozenTransform { transform } => {
1530 if transform.nrows() != m_aug {
1531 crate::bail_dim_basis!(
1532 "frozen measure-jet identifiability transform mismatch: {} representers + {} head columns but transform has {} rows",
1533 m,
1534 head_rank,
1535 transform.nrows()
1536 );
1537 }
1538 (
1539 transform.clone(),
1540 gam_problem::Gauge::from_block_transforms(&[transform.clone()]),
1541 )
1542 }
1543 MeasureJetIdentifiability::CenterSumToZero => {
1544 let z_rbf = if head_rank > 0 {
1545 let affine = measure_jet_affine_value_basis(centers.view(), masses.view());
1546 let mut weighted_affine = affine.clone();
1547 for (i, mut row) in weighted_affine.outer_iter_mut().enumerate() {
1548 row.mapv_inplace(|v| v * masses[i]);
1549 }
1550 let constraint_cross = k_cc.t().dot(&weighted_affine);
1554 rrqr_nullspace_basis(&constraint_cross, default_rrqr_rank_alpha())
1555 .map_err(BasisError::LinalgError)?
1556 .0
1557 } else {
1558 let u = householder_sum_to_zero_u(m);
1559 householder_sum_to_zero_z(&u)
1560 };
1561 let rbf_rank = z_rbf.ncols();
1562 let mut z_block = Array2::<f64>::zeros((m_aug, rbf_rank + head_rank));
1563 z_block
1564 .slice_mut(ndarray::s![..m, ..rbf_rank])
1565 .assign(&z_rbf);
1566 for r in 0..head_rank {
1567 z_block[(m + r, rbf_rank + r)] = 1.0;
1568 }
1569 (
1570 z_block.clone(),
1571 gam_problem::Gauge::from_block_transforms(&[z_block]),
1572 )
1573 }
1574 };
1575 let mut k_aug = Array2::<f64>::zeros((m, m_aug));
1580 k_aug.slice_mut(ndarray::s![.., ..m]).assign(&k_cc);
1581 if head_rank > 0 {
1582 k_aug.slice_mut(ndarray::s![.., m..]).assign(&head_cc);
1583 }
1584 let kz = coefficient_gauge.restrict_design(&k_aug);
1585 Ok(RealizedMeasureJetGeometry {
1586 centers,
1587 masses,
1588 eps_band,
1589 log_step,
1590 length_scale,
1591 order_s_eval: order_s,
1592 per_level: spec.multiscale,
1597 z,
1598 coefficient_gauge,
1599 kz,
1600 head_transform,
1601 })
1602}
1603
1604pub fn measure_jet_input_noise_scale(
1627 data: ArrayView2<'_, f64>,
1628 centers: ArrayView2<'_, f64>,
1629) -> Result<Option<f64>, BasisError> {
1630 let d = data.ncols();
1631 let m = centers.nrows();
1632 if d == 0 || m == 0 || data.nrows() == 0 {
1633 return Ok(None);
1634 }
1635 if centers.ncols() != d {
1636 crate::bail_dim_basis!(
1637 "measure-jet input-noise estimate: data d={d} disagrees with centers d={}",
1638 centers.ncols()
1639 );
1640 }
1641 validate_finite_points(data, "data")?;
1642 validate_finite_points(centers, "centers")?;
1643 let sq = pairwise_sq_dists(data, centers);
1646 let mut members: Vec<Vec<usize>> = vec![Vec::new(); m];
1647 for (j, row) in sq.axis_iter(Axis(0)).enumerate() {
1648 let mut best = 0usize;
1649 let mut best_d = f64::INFINITY;
1650 for (i, &dij) in row.iter().enumerate() {
1651 if dij < best_d {
1652 best_d = dij;
1653 best = i;
1654 }
1655 }
1656 members[best].push(j);
1657 }
1658 let mut weighted_sum = 0.0_f64;
1659 let mut weight = 0.0_f64;
1660 for cell in &members {
1661 let n_i = cell.len();
1662 if n_i < d + 1 {
1665 continue;
1666 }
1667 let mut mean = Array1::<f64>::zeros(d);
1669 for &j in cell {
1670 mean += &data.row(j);
1671 }
1672 mean /= n_i as f64;
1673 let mut cov = Array2::<f64>::zeros((d, d));
1674 for &j in cell {
1675 let mut centered = data.row(j).to_owned();
1676 centered -= &mean;
1677 for a in 0..d {
1678 for b in 0..d {
1679 cov[(a, b)] += centered[a] * centered[b];
1680 }
1681 }
1682 }
1683 cov /= n_i as f64;
1684 let cov_sym = (&cov + &cov.t()) * 0.5;
1687 let (evals, _) = cov_sym.eigh(Side::Lower).map_err(|e| {
1688 BasisError::InvalidInput(format!(
1689 "measure-jet input-noise estimate: local covariance eigendecomposition failed: {e}"
1690 ))
1691 })?;
1692 let smallest = evals
1693 .iter()
1694 .copied()
1695 .fold(f64::INFINITY, |acc, v| acc.min(v))
1696 .max(0.0);
1697 if smallest.is_finite() {
1698 weighted_sum += n_i as f64 * smallest;
1699 weight += n_i as f64;
1700 }
1701 }
1702 if weight <= 0.0 {
1703 return Ok(None);
1704 }
1705 let sigma2 = weighted_sum / weight;
1706 if !(sigma2.is_finite() && sigma2 > 0.0) {
1707 return Ok(None);
1708 }
1709 Ok(Some(sigma2.sqrt()))
1710}
1711
1712pub fn measure_jet_multiscale_mode(spec: &MeasureJetBasisSpec) -> bool {
1720 spec.multiscale
1721}
1722
1723pub fn build_measure_jet_basis(
1731 data: ArrayView2<'_, f64>,
1732 spec: &MeasureJetBasisSpec,
1733) -> Result<BasisBuildResult, BasisError> {
1734 let RealizedMeasureJetGeometry {
1735 centers,
1736 masses,
1737 eps_band,
1738 log_step,
1739 length_scale,
1740 order_s_eval: order_s,
1741 per_level,
1742 z,
1743 coefficient_gauge,
1744 kz,
1745 head_transform,
1746 } = realize_measure_jet_geometry(data, spec)?;
1747 let band = MeasureJetBand {
1748 eps: eps_band.clone(),
1749 log_step,
1750 };
1751 let m = centers.nrows();
1752 let head_rank = head_transform.ncols();
1753 let m_aug = m + head_rank;
1754 let kernel_design = measure_jet_design_matrix(data, centers.view(), length_scale)?;
1759 let mut raw_design = Array2::<f64>::zeros((data.nrows(), m_aug));
1760 raw_design
1761 .slice_mut(ndarray::s![.., ..m])
1762 .assign(&kernel_design);
1763 if head_rank > 0 {
1764 let head_design = data.dot(&head_transform);
1765 raw_design
1766 .slice_mut(ndarray::s![.., m..])
1767 .assign(&head_design);
1768 }
1769 let constrained_design = coefficient_gauge.restrict_design(&raw_design);
1770 let design = gam_linalg::matrix::DesignMatrix::Dense(
1771 gam_linalg::matrix::DenseDesignMatrix::from(constrained_design),
1772 );
1773 let support_means = measure_jet_support_means(centers.view(), masses.view(), &eps_band)?;
1774 let mut candidates = Vec::new();
1787 let mut penalty_normalization_scales = Vec::new();
1788 let mut raw_penalty_normalization_scales = Vec::new();
1789 let mut fused_penalty_normalization_scale = None;
1790 if per_level {
1791 let forms = measure_jet_energy_forms_per_scale(
1792 centers.view(),
1793 masses.view(),
1794 &band,
1795 order_s,
1796 spec.alpha,
1797 spec.tau0,
1798 )?;
1799 for (level, q_l) in forms.into_iter().enumerate() {
1800 let s_l = kz.t().dot(&q_l).dot(&kz);
1801 let (s_norm, c_l) = normalize_penalty(&((&s_l + &s_l.t()) * 0.5));
1802 let intrinsic_dim = centers.ncols() as f64;
1803 let eta = 2.0 * order_s + intrinsic_dim * (2.0 - 2.0 * spec.alpha);
1804 let scale_weight = log_step * eps_band[level].powf(-eta);
1805 penalty_normalization_scales.push(c_l);
1806 raw_penalty_normalization_scales.push(c_l / scale_weight);
1807 candidates.push(PenaltyCandidate {
1808 matrix: s_norm,
1809 nullspace_dim_hint: 0,
1810 source: PenaltySource::Other(format!("measure_jet_scale_{level}")),
1811 normalization_scale: c_l,
1812 kronecker_factors: None,
1813 op: None,
1814 });
1815 }
1816 } else {
1817 let q_form = measure_jet_energy_form(
1818 centers.view(),
1819 masses.view(),
1820 &band,
1821 order_s,
1822 spec.alpha,
1823 spec.tau0,
1824 )?;
1825 let penalty = pullback_center_form(&kz, &q_form);
1830 let (penalty_norm, c_primary) = normalize_penalty(&penalty);
1831 fused_penalty_normalization_scale = Some(c_primary);
1832 candidates.push(PenaltyCandidate {
1833 matrix: penalty_norm,
1834 nullspace_dim_hint: 0,
1835 source: PenaltySource::Primary,
1836 normalization_scale: c_primary,
1837 kronecker_factors: None,
1838 op: None,
1839 });
1840 }
1841 if spec.double_penalty {
1847 let null_penalty = affine_function_nullspace_penalty(&kz, centers.view(), masses.view())?;
1848 let (null_penalty_norm, c_null) = normalize_penalty(&null_penalty);
1849 candidates.push(PenaltyCandidate {
1850 matrix: null_penalty_norm,
1851 nullspace_dim_hint: 0,
1852 source: PenaltySource::DoublePenaltyNullspace,
1853 normalization_scale: c_null,
1854 kronecker_factors: None,
1855 op: None,
1856 });
1857 }
1858 let (penalties, nullspace_dims, penaltyinfo, null_eigenvectors, ops) =
1859 filter_active_penalty_candidates_with_ops(candidates)?;
1860 let sigma_coord = measure_jet_input_noise_scale(data, centers.view())?;
1863 Ok(BasisBuildResult {
1864 design,
1865 penalties,
1866 nullspace_dims,
1867 penaltyinfo,
1868 metadata: BasisMetadata::MeasureJet {
1869 centers,
1870 input_scales: None,
1871 length_scale,
1872 eps_band,
1873 order_s: spec.order_s,
1878 alpha: spec.alpha,
1879 tau0: spec.tau0,
1880 masses,
1881 support_means,
1882 penalty_normalization_scales,
1883 raw_penalty_normalization_scales,
1884 fused_penalty_normalization_scale,
1885 constraint_transform: Some(z),
1886 sigma_coord,
1889 },
1890 kronecker_factored: None,
1891 ops,
1892 null_eigenvectors,
1893 joint_null_rotation: None,
1894 })
1895}
1896
1897pub fn build_measure_jet_basis_psi_derivatives(
1920 data: ArrayView2<'_, f64>,
1921 spec: &MeasureJetBasisSpec,
1922) -> Result<AnisoBasisPsiDerivatives, BasisError> {
1923 if !(spec.tau0.is_finite() && spec.tau0 > 0.0) {
1924 crate::bail_invalid_basis!(
1925 "measure-jet ψ derivatives need tau0 > 0 because the retained τ coordinate is ln τ; got {}",
1926 spec.tau0
1927 );
1928 }
1929 let geom = realize_measure_jet_geometry(data, spec)?;
1930 let band = MeasureJetBand {
1931 eps: geom.eps_band.clone(),
1932 log_step: geom.log_step,
1933 };
1934 let n = data.nrows();
1935 let p = geom.kz.ncols();
1936 let m = geom.centers.nrows();
1937 let m_aug = m + geom.head_transform.ncols();
1938
1939 struct LengthScaleJets {
1940 evaluation_first: Array2<f64>,
1941 evaluation_second: Array2<f64>,
1942 design_first: Array2<f64>,
1943 design_second: Array2<f64>,
1944 }
1945
1946 let length_scale_jets = if spec.learn_length_scale {
1953 let (dk_data, d2k_data) =
1954 measure_jet_design_log_length_jets(data, geom.centers.view(), geom.length_scale)?;
1955 let mut dk_data_aug = Array2::<f64>::zeros((n, m_aug));
1956 let mut d2k_data_aug = Array2::<f64>::zeros((n, m_aug));
1957 dk_data_aug.slice_mut(ndarray::s![.., ..m]).assign(&dk_data);
1958 d2k_data_aug
1959 .slice_mut(ndarray::s![.., ..m])
1960 .assign(&d2k_data);
1961
1962 let (dk_centers, d2k_centers) = measure_jet_design_log_length_jets(
1963 geom.centers.view(),
1964 geom.centers.view(),
1965 geom.length_scale,
1966 )?;
1967 let mut dk_centers_aug = Array2::<f64>::zeros((m, m_aug));
1968 let mut d2k_centers_aug = Array2::<f64>::zeros((m, m_aug));
1969 dk_centers_aug
1970 .slice_mut(ndarray::s![.., ..m])
1971 .assign(&dk_centers);
1972 d2k_centers_aug
1973 .slice_mut(ndarray::s![.., ..m])
1974 .assign(&d2k_centers);
1975
1976 Some(LengthScaleJets {
1977 evaluation_first: geom.coefficient_gauge.restrict_design(&dk_centers_aug),
1978 evaluation_second: geom.coefficient_gauge.restrict_design(&d2k_centers_aug),
1979 design_first: geom.coefficient_gauge.restrict_design(&dk_data_aug),
1980 design_second: geom.coefficient_gauge.restrict_design(&d2k_data_aug),
1981 })
1982 } else {
1983 None
1984 };
1985
1986 let coord_offset = usize::from(length_scale_jets.is_some());
1987 let n_coords = coord_offset + if geom.per_level { 2 } else { 0 };
1988 let pairs: Vec<(usize, usize)> = (0..n_coords)
1989 .flat_map(|a| ((a + 1)..n_coords).map(move |b| (a, b)))
1990 .collect();
1991 let zero_p = || Array2::<f64>::zeros((p, p));
1992
1993 struct RawPenaltyJets {
1994 value: Array2<f64>,
1995 first: Vec<Array2<f64>>,
1996 second_diag: Vec<Array2<f64>>,
1997 cross: Vec<Array2<f64>>,
1998 }
1999
2000 let sandwich = |form: &Array2<f64>| pullback_center_form(&geom.kz, form);
2001 let length_diag = |form: &Array2<f64>| {
2002 let jets = length_scale_jets
2003 .as_ref()
2004 .expect("length-scale form jets require an enrolled length coordinate");
2005 pullback_center_form_log_length_jets(
2006 &geom.kz,
2007 &jets.evaluation_first,
2008 &jets.evaluation_second,
2009 form,
2010 )
2011 };
2012 let length_cross = |form_first: &Array2<f64>| {
2013 let jets = length_scale_jets
2014 .as_ref()
2015 .expect("length-scale cross jets require an enrolled length coordinate");
2016 pullback_center_form_log_length_cross(&geom.kz, &jets.evaluation_first, form_first)
2017 };
2018
2019 let mut raw: Vec<RawPenaltyJets> = if geom.per_level {
2025 let l_count = band.eps.len();
2026 let forms = assemble_weighted_forms(
2029 geom.centers.view(),
2030 geom.masses.view(),
2031 &band,
2032 geom.order_s_eval,
2033 spec.alpha,
2034 spec.tau0,
2035 6 * l_count,
2036 3,
2037 &|scale_idx, eps: f64, q: f64, base: f64, out: &mut [[f64; 3]]| {
2038 for slot in out.iter_mut() {
2039 *slot = [0.0, 0.0, 0.0];
2040 }
2041 let intrinsic_dim = geom.centers.ncols() as f64;
2042 let ga = 2.0 * intrinsic_dim * eps.ln() - 2.0 * q.max(f64::MIN_POSITIVE).ln();
2043 let k0 = 6 * scale_idx;
2044 out[k0] = [base, 0.0, 0.0];
2045 out[k0 + 1] = [ga * base, 0.0, 0.0];
2046 out[k0 + 2] = [ga * ga * base, 0.0, 0.0];
2047 out[k0 + 3] = [0.0, 0.0, 0.0];
2048 out[k0 + 4] = [0.0, 0.0, 0.0];
2049 out[k0 + 5] = [0.0, 0.0, 0.0];
2050 },
2051 )?;
2052 let alpha_coord = coord_offset;
2053 let tau_coord = coord_offset + 1;
2054 let mut raw = Vec::with_capacity(l_count + usize::from(spec.double_penalty));
2055 for level in 0..l_count {
2056 let chunk = &forms[6 * level..6 * level + 6];
2057 let mut first: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2058 let mut second_diag: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2059 first[alpha_coord] = sandwich(&chunk[1]);
2060 first[tau_coord] = sandwich(&chunk[3]);
2061 second_diag[alpha_coord] = sandwich(&chunk[2]);
2062 second_diag[tau_coord] = sandwich(&chunk[4]);
2063 if coord_offset == 1 {
2064 let (ell_first, ell_second) = length_diag(&chunk[0]);
2065 first[0] = ell_first;
2066 second_diag[0] = ell_second;
2067 }
2068 let mut cross: Vec<Array2<f64>> = (0..pairs.len()).map(|_| zero_p()).collect();
2069 for (pair_idx, &(a, b)) in pairs.iter().enumerate() {
2070 cross[pair_idx] = if coord_offset == 1 && a == 0 && b == alpha_coord {
2071 length_cross(&chunk[1])
2072 } else if coord_offset == 1 && a == 0 && b == tau_coord {
2073 length_cross(&chunk[3])
2074 } else if a == alpha_coord && b == tau_coord {
2075 sandwich(&chunk[5])
2076 } else {
2077 zero_p()
2078 };
2079 }
2080 raw.push(RawPenaltyJets {
2081 value: sandwich(&chunk[0]),
2082 first,
2083 second_diag,
2084 cross,
2085 });
2086 }
2087 raw
2088 } else {
2089 let q_form = measure_jet_energy_form(
2093 geom.centers.view(),
2094 geom.masses.view(),
2095 &band,
2096 geom.order_s_eval,
2097 spec.alpha,
2098 spec.tau0,
2099 )?;
2100 let mut first: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2101 let mut second_diag: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2102 if coord_offset == 1 {
2103 let (ell_first, ell_second) = length_diag(&q_form);
2104 first[0] = ell_first;
2105 second_diag[0] = ell_second;
2106 }
2107 vec![RawPenaltyJets {
2108 value: sandwich(&q_form),
2109 first,
2110 second_diag,
2111 cross: Vec::new(),
2112 }]
2113 };
2114
2115 if spec.double_penalty {
2116 let null_form = affine_function_nullspace_form(geom.centers.view(), geom.masses.view())?;
2117 let mut first: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2118 let mut second_diag: Vec<Array2<f64>> = (0..n_coords).map(|_| zero_p()).collect();
2119 if coord_offset == 1 {
2120 let (ell_first, ell_second) = length_diag(&null_form);
2121 first[0] = ell_first;
2122 second_diag[0] = ell_second;
2123 }
2124 raw.push(RawPenaltyJets {
2125 value: sandwich(&null_form),
2126 first,
2127 second_diag,
2128 cross: (0..pairs.len()).map(|_| zero_p()).collect(),
2131 });
2132 }
2133
2134 let n_cands = raw.len();
2135 let mut penalties_first: Vec<Vec<Array2<f64>>> =
2136 (0..n_coords).map(|_| Vec::with_capacity(n_cands)).collect();
2137 let mut penalties_second_diag: Vec<Vec<Array2<f64>>> =
2138 (0..n_coords).map(|_| Vec::with_capacity(n_cands)).collect();
2139 let mut crosses: Vec<Vec<Array2<f64>>> = (0..pairs.len()).map(|_| Vec::new()).collect();
2143 for candidate in &raw {
2144 let s_raw = &candidate.value;
2145 let fro = trace_of_product(s_raw, s_raw).sqrt();
2154 let c = if fro.is_finite() && fro > 1e-12 {
2155 fro
2156 } else {
2157 1.0
2158 };
2159 for coord in 0..n_coords {
2160 let (_, s_first, s_second, _) = normalize_penaltywith_psi_derivatives(
2161 s_raw,
2162 &candidate.first[coord],
2163 &candidate.second_diag[coord],
2164 );
2165 penalties_first[coord].push(s_first);
2166 penalties_second_diag[coord].push(s_second);
2167 }
2168 for (pair_idx, &(a, b)) in pairs.iter().enumerate() {
2169 let cross_raw_mat = normalize_penalty_cross_psi_derivative(
2170 s_raw,
2171 &candidate.first[a],
2172 &candidate.first[b],
2173 &candidate.cross[pair_idx],
2174 c,
2175 );
2176 crosses[pair_idx].push(cross_raw_mat);
2177 }
2178 }
2179
2180 let pair_index: Vec<((usize, usize), Vec<Array2<f64>>)> =
2181 pairs.iter().copied().zip(crosses.into_iter()).collect();
2182 let provider = AnisoPenaltyCrossProvider::new(move |a, b| {
2183 pair_index
2184 .iter()
2185 .find(|((pa, pb), _)| (*pa, *pb) == (a, b) || (*pa, *pb) == (b, a))
2186 .map(|(_, mats)| mats.clone())
2187 .ok_or_else(|| {
2188 BasisError::InvalidInput(format!(
2189 "measure-jet ψ cross derivative requested for unknown pair ({a}, {b})"
2190 ))
2191 })
2192 });
2193 let mut design_first: Vec<Array2<f64>> = (0..n_coords)
2194 .map(|_| Array2::<f64>::zeros((n, p)))
2195 .collect();
2196 let mut design_second_diag: Vec<Array2<f64>> = (0..n_coords)
2197 .map(|_| Array2::<f64>::zeros((n, p)))
2198 .collect();
2199 if let Some(jets) = &length_scale_jets {
2200 design_first[0] = jets.design_first.clone();
2201 design_second_diag[0] = jets.design_second.clone();
2202 }
2203 Ok(AnisoBasisPsiDerivatives {
2204 design_first,
2205 design_second_diag,
2206 design_second_cross: Vec::new(),
2207 design_second_cross_pairs: Vec::new(),
2208 penalties_first,
2209 penalties_second_diag,
2210 penalties_cross_pairs: pairs,
2211 penalties_cross_provider: Some(provider),
2212 implicit_operator: None,
2213 })
2214}
2215
2216#[cfg(test)]
2217mod tests {
2218 use super::*;
2219
2220 fn lcg_normal(state: &mut u64) -> f64 {
2223 let mut next = || {
2224 *state = state
2225 .wrapping_mul(6364136223846793005)
2226 .wrapping_add(1442695040888963407);
2227 (((*state >> 11) as f64) + 0.5) / (1u64 << 53) as f64
2229 };
2230 let u1 = next();
2231 let u2 = next();
2232 (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
2233 }
2234
2235 #[test]
2241 pub(crate) fn input_noise_scale_recovers_known_perpendicular_sigma() {
2242 let tang = [1.0 / 5f64.sqrt(), 2.0 / 5f64.sqrt()];
2244 let perp = [2.0 / 5f64.sqrt(), -1.0 / 5f64.sqrt()];
2245 let sigma = 0.05_f64;
2246 let n = 600usize;
2247 let mut state = 0x1234_5678_9abc_def0u64;
2248 let mut data = Array2::<f64>::zeros((n, 2));
2249 for j in 0..n {
2250 let t = 3.0 * (j as f64) / (n as f64 - 1.0);
2252 let noise = sigma * lcg_normal(&mut state);
2253 for a in 0..2 {
2254 data[(j, a)] = t * tang[a] + noise * perp[a];
2255 }
2256 }
2257 let n_centers = 8usize;
2260 let mut centers = Array2::<f64>::zeros((n_centers, 2));
2261 for i in 0..n_centers {
2262 let t = 3.0 * (i as f64 + 0.5) / (n_centers as f64);
2263 for a in 0..2 {
2264 centers[(i, a)] = t * tang[a];
2265 }
2266 }
2267 let est = measure_jet_input_noise_scale(data.view(), centers.view())
2268 .expect("estimate ok")
2269 .expect("noise scale present");
2270 assert!(
2273 (est - sigma).abs() <= 0.4 * sigma,
2274 "estimated σ_coord {est} far from true {sigma}"
2275 );
2276 }
2277
2278 #[test]
2281 pub(crate) fn input_noise_scale_none_when_cells_too_small() {
2282 let data = array![[0.0, 0.0], [1.0, 2.0], [2.0, 4.0]];
2283 let centers = array![[0.0, 0.0], [1.0, 2.0], [2.0, 4.0]];
2284 assert!(
2286 measure_jet_input_noise_scale(data.view(), centers.view())
2287 .expect("estimate ok")
2288 .is_none()
2289 );
2290 }
2291
2292 pub(crate) fn two_cluster_centers() -> (ndarray::Array2<f64>, ndarray::Array1<f64>) {
2293 let centers = array![
2294 [0.00, 0.00],
2295 [0.31, 0.05],
2296 [0.58, -0.07],
2297 [0.93, 0.11],
2298 [1.22, 0.02],
2299 [1.49, -0.04],
2300 [3.10, 2.00],
2301 [3.42, 2.13],
2302 [3.71, 1.91],
2303 [4.05, 2.07],
2304 [4.33, 1.96],
2305 [4.61, 2.12],
2306 ];
2307 let m = centers.nrows();
2308 let masses = ndarray::Array1::<f64>::from_elem(m, 1.0 / m as f64);
2309 (centers, masses)
2310 }
2311 use ndarray::array;
2312
2313 pub(crate) fn band_for(centers: &Array2<f64>) -> MeasureJetBand {
2314 measure_jet_band(centers.view(), 0).expect("band")
2315 }
2316
2317 #[test]
2320 pub(crate) fn energy_form_annihilates_constants_exactly() {
2321 let (centers, masses) = two_cluster_centers();
2322 let band = band_for(¢ers);
2323 let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2324 .expect("energy form");
2325 let m = q.nrows();
2326 let ones = Array1::<f64>::ones(m);
2327 let qv = q.dot(&ones);
2328 let scale = q.iter().fold(0.0_f64, |acc, v| acc.max(v.abs()));
2329 assert!(scale > 0.0, "energy form is identically zero");
2330 for (i, v) in qv.iter().enumerate() {
2331 assert!(
2332 v.abs() <= 1e-12 * scale,
2333 "Q·1 leak at row {i}: {v:.3e} vs scale {scale:.3e}"
2334 );
2335 }
2336 let vqv = ones.dot(&qv);
2337 assert!(
2338 vqv.abs() <= 1e-12 * scale,
2339 "constant carries energy: 1ᵀQ1 = {vqv:.3e}"
2340 );
2341 }
2342
2343 #[test]
2346 pub(crate) fn energy_form_annihilates_affine_at_default_tau() {
2347 let (centers, masses) = two_cluster_centers();
2348 let band = band_for(¢ers);
2349 let m = centers.nrows();
2350 let mut affine = Array1::<f64>::zeros(m);
2352 let mut rough = Array1::<f64>::zeros(m);
2353 for i in 0..m {
2354 affine[i] = 0.7 + 1.3 * centers[(i, 0)] - 0.4 * centers[(i, 1)];
2355 rough[i] = if i % 2 == 0 { 1.0 } else { -1.0 };
2356 }
2357 let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2358 .expect("energy form");
2359 let e_affine = affine.dot(&q.dot(&affine));
2360 let e_rough = rough.dot(&q.dot(&rough));
2361 assert!(e_rough > 0.0, "rough vector must pay energy");
2362 assert!(
2363 e_affine.abs() <= 1e-12 * e_rough,
2364 "default affine energy {e_affine:.3e} vs rough {e_rough:.3e}"
2365 );
2366 }
2367
2368 #[test]
2370 pub(crate) fn energy_form_is_psd() {
2371 let (centers, masses) = two_cluster_centers();
2372 let band = band_for(¢ers);
2373 let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2374 .expect("energy form");
2375 let m = q.nrows();
2376 for trial in 0..5usize {
2377 let v = Array1::<f64>::from_shape_fn(m, |i| {
2378 ((i * 7 + trial * 13) % 11) as f64 / 11.0 - 0.5
2379 });
2380 let e = v.dot(&q.dot(&v));
2381 assert!(e >= -1e-10, "vᵀQv = {e:.3e} < 0 on trial {trial}");
2382 }
2383 }
2384
2385 #[test]
2388 pub(crate) fn rough_vector_pays_more_than_smooth() {
2389 let m = 24usize;
2390 let centers = Array2::<f64>::from_shape_fn((m, 2), |(i, k)| {
2391 let t = i as f64 / (m as f64 - 1.0);
2392 if k == 0 {
2393 t * 4.0
2394 } else {
2395 0.3 * (t * 4.0).sin()
2396 }
2397 });
2398 let masses = Array1::<f64>::from_elem(m, 1.0 / m as f64);
2399 let band = band_for(¢ers);
2400 let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2401 .expect("energy form");
2402 let slow = Array1::<f64>::from_shape_fn(m, |i| (i as f64 / (m as f64 - 1.0)).powi(2));
2403 let fast = Array1::<f64>::from_shape_fn(m, |i| if i % 2 == 0 { 0.5 } else { -0.5 });
2404 let e_slow = slow.dot(&q.dot(&slow));
2405 let e_fast = fast.dot(&q.dot(&fast));
2406 assert!(
2407 e_fast > 10.0 * e_slow,
2408 "alternating values must pay >> a slow trend: fast {e_fast:.3e} vs slow {e_slow:.3e}"
2409 );
2410 }
2411
2412 #[test]
2417 pub(crate) fn energy_jets_match_finite_differences() {
2418 let (centers, masses) = two_cluster_centers();
2419 let band = band_for(¢ers);
2420 let (s0, a0, tau) = (1.3, 0.8, 1e-3);
2421 let jets =
2422 measure_jet_energy_form_with_jets(centers.view(), masses.view(), &band, s0, a0, tau)
2423 .expect("jets");
2424 let q_at = |s: f64, a: f64| {
2425 measure_jet_energy_form(centers.view(), masses.view(), &band, s, a, tau)
2426 .expect("energy form")
2427 };
2428 let q_plain = q_at(s0, a0);
2430 for (a, b) in jets.q.iter().zip(q_plain.iter()) {
2431 assert!(
2432 (a - b).abs() <= 1e-14 * (1.0 + b.abs()),
2433 "Q drift {a} vs {b}"
2434 );
2435 }
2436 let lt0 = tau.ln();
2437 let q_at_lt = |lt: f64| {
2438 measure_jet_energy_form(centers.view(), masses.view(), &band, s0, a0, lt.exp())
2439 .expect("energy form")
2440 };
2441 let h = 1e-4;
2447 let checks: [(&str, &Array2<f64>, Array2<f64>); 9] = [
2448 ("dq_ds", &jets.dq_ds, {
2449 let (p, m_) = (q_at(s0 + h, a0), q_at(s0 - h, a0));
2450 (&p - &m_) / (2.0 * h)
2451 }),
2452 ("d2q_ds2", &jets.d2q_ds2, {
2453 let (p, c, m_) = (q_at(s0 + h, a0), q_at(s0, a0), q_at(s0 - h, a0));
2454 (&(&p + &m_) - &(&c * 2.0)) / (h * h)
2455 }),
2456 ("dq_dalpha", &jets.dq_dalpha, {
2457 let (p, m_) = (q_at(s0, a0 + h), q_at(s0, a0 - h));
2458 (&p - &m_) / (2.0 * h)
2459 }),
2460 ("d2q_dalpha2", &jets.d2q_dalpha2, {
2461 let (p, c, m_) = (q_at(s0, a0 + h), q_at(s0, a0), q_at(s0, a0 - h));
2462 (&(&p + &m_) - &(&c * 2.0)) / (h * h)
2463 }),
2464 ("d2q_ds_dalpha", &jets.d2q_ds_dalpha, {
2465 let pp = q_at(s0 + h, a0 + h);
2466 let pm = q_at(s0 + h, a0 - h);
2467 let mp = q_at(s0 - h, a0 + h);
2468 let mm = q_at(s0 - h, a0 - h);
2469 (&(&pp - &pm) - &(&mp - &mm)) / (4.0 * h * h)
2470 }),
2471 ("dq_dlogtau", &jets.dq_dlogtau, {
2472 let (p, m_) = (q_at_lt(lt0 + h), q_at_lt(lt0 - h));
2473 (&p - &m_) / (2.0 * h)
2474 }),
2475 ("d2q_dlogtau2", &jets.d2q_dlogtau2, {
2476 let (p, c, m_) = (q_at_lt(lt0 + h), q_at_lt(lt0), q_at_lt(lt0 - h));
2477 (&(&p + &m_) - &(&c * 2.0)) / (h * h)
2478 }),
2479 ("d2q_ds_dlogtau", &jets.d2q_ds_dlogtau, {
2480 let f = |s: f64, lt: f64| {
2481 measure_jet_energy_form(centers.view(), masses.view(), &band, s, a0, lt.exp())
2482 .expect("energy form")
2483 };
2484 let pp = f(s0 + h, lt0 + h);
2485 let pm = f(s0 + h, lt0 - h);
2486 let mp = f(s0 - h, lt0 + h);
2487 let mm = f(s0 - h, lt0 - h);
2488 (&(&pp - &pm) - &(&mp - &mm)) / (4.0 * h * h)
2489 }),
2490 ("d2q_dalpha_dlogtau", &jets.d2q_dalpha_dlogtau, {
2491 let f = |a: f64, lt: f64| {
2492 measure_jet_energy_form(centers.view(), masses.view(), &band, s0, a, lt.exp())
2493 .expect("energy form")
2494 };
2495 let pp = f(a0 + h, lt0 + h);
2496 let pm = f(a0 + h, lt0 - h);
2497 let mp = f(a0 - h, lt0 + h);
2498 let mm = f(a0 - h, lt0 - h);
2499 (&(&pp - &pm) - &(&mp - &mm)) / (4.0 * h * h)
2500 }),
2501 ];
2502 for (name, analytic, fd) in checks.iter() {
2503 let scale = fd.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
2504 for (a, b) in analytic.iter().zip(fd.iter()) {
2505 assert!(
2506 (a - b).abs() <= 5e-5 * scale,
2507 "{name} jet mismatch: analytic {a:.6e} vs FD {b:.6e} (scale {scale:.3e})"
2508 );
2509 }
2510 }
2511 }
2512
2513 #[test]
2517 pub(crate) fn scale_spectrum_sums_to_total_and_localizes_roughness() {
2518 let m = 24usize;
2519 let centers = Array2::<f64>::from_shape_fn((m, 2), |(i, k)| {
2520 let t = i as f64 / (m as f64 - 1.0);
2521 if k == 0 { t * 4.0 } else { 0.0 }
2522 });
2523 let masses = Array1::<f64>::from_elem(m, 1.0 / m as f64);
2524 let band = band_for(¢ers);
2525 let q = measure_jet_energy_form(centers.view(), masses.view(), &band, 1.5, 1.0, 1e-3)
2526 .expect("energy form");
2527 let fast = Array1::<f64>::from_shape_fn(m, |i| if i % 2 == 0 { 0.5 } else { -0.5 });
2528 let spec = measure_jet_scale_spectrum(
2529 centers.view(),
2530 masses.view(),
2531 &band,
2532 1.5,
2533 1.0,
2534 1e-3,
2535 fast.view(),
2536 )
2537 .expect("spectrum");
2538 assert_eq!(spec.len(), band.eps.len());
2539 let total = fast.dot(&q.dot(&fast));
2540 let sum: f64 = spec.iter().sum();
2541 assert!(
2542 (sum - total).abs() <= 1e-10 * total.abs().max(1e-30),
2543 "spectrum must sum to vᵀQv: {sum:.6e} vs {total:.6e}"
2544 );
2545 let finest = spec[0];
2547 let coarsest = *spec.last().expect("nonempty spectrum");
2548 assert!(
2549 finest > coarsest,
2550 "alternating values must charge fine scales hardest: fine {finest:.3e} vs coarse {coarsest:.3e}"
2551 );
2552 }
2553
2554 #[test]
2557 pub(crate) fn support_curve_separates_on_web_from_off_web() {
2558 let m = 24usize;
2559 let centers = Array2::<f64>::from_shape_fn((m, 2), |(i, k)| {
2560 let t = i as f64 / (m as f64 - 1.0);
2561 if k == 0 { t * 4.0 } else { 0.0 }
2562 });
2563 let masses = Array1::<f64>::from_elem(m, 1.0 / m as f64);
2564 let band = band_for(¢ers);
2565 let queries = array![[2.0, 0.0], [2.0, 1.5]];
2566 let curves =
2567 measure_jet_support_curve(queries.view(), centers.view(), masses.view(), &band.eps)
2568 .expect("support curve");
2569 assert!(
2571 curves[(0, 0)] > 10.0 * curves[(1, 0)],
2572 "fine-scale support must separate web from void: on {:.3e} vs off {:.3e}",
2573 curves[(0, 0)],
2574 curves[(1, 0)]
2575 );
2576 for qi in 0..2 {
2578 for li in 1..band.eps.len() {
2579 assert!(
2580 curves[(qi, li)] >= curves[(qi, li - 1)] - 1e-15,
2581 "support curve must be monotone in scale (query {qi}, level {li})"
2582 );
2583 }
2584 }
2585 }
2586
2587 #[test]
2595 pub(crate) fn default_stays_single_scale_until_multiscale_opt_in() {
2596 let n = 200usize;
2597 let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2598 let t = i as f64 / (n as f64 - 1.0);
2599 if k == 0 {
2600 t * 3.0
2601 } else {
2602 0.4 * (t * 3.0).sin()
2603 }
2604 });
2605 let single = MeasureJetBasisSpec {
2609 center_strategy: CenterStrategy::FarthestPoint { num_centers: 80 },
2610 ..MeasureJetBasisSpec::default()
2611 };
2612 assert!(
2613 !measure_jet_multiscale_mode(&single),
2614 "default must resolve to single-scale at any center count"
2615 );
2616 let built_single =
2617 build_measure_jet_basis(data.view(), &single).expect("single-scale build");
2618 assert_eq!(
2619 built_single.penalties.len(),
2620 2,
2621 "single-scale double-penalty mode emits Primary + affine/null component"
2622 );
2623 assert!(matches!(
2624 built_single.penaltyinfo[0].source,
2625 PenaltySource::Primary
2626 ));
2627 assert!(matches!(
2628 built_single.penaltyinfo[1].source,
2629 PenaltySource::DoublePenaltyNullspace
2630 ));
2631 let multi = MeasureJetBasisSpec {
2635 center_strategy: CenterStrategy::FarthestPoint { num_centers: 80 },
2636 multiscale: true,
2637 ..MeasureJetBasisSpec::default()
2638 };
2639 assert!(
2640 measure_jet_multiscale_mode(&multi),
2641 "multiscale=true must resolve to multiscale mode"
2642 );
2643 let built_multi = build_measure_jet_basis(data.view(), &multi).expect("multiscale build");
2644 assert!(
2645 built_multi.penalties.len() > built_single.penalties.len(),
2646 "multiscale mode emits the per-scale spectral split plus null selection, got {} (vs single-scale {})",
2647 built_multi.penalties.len(),
2648 built_single.penalties.len()
2649 );
2650 }
2651
2652 #[test]
2656 pub(crate) fn fused_mode_without_double_penalty_emits_single_primary_candidate() {
2657 let n = 40usize;
2658 let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2659 let t = i as f64 / (n as f64 - 1.0);
2660 if k == 0 {
2661 t * 3.0
2662 } else {
2663 0.4 * (t * 3.0).sin()
2664 }
2665 });
2666 let spec = MeasureJetBasisSpec {
2667 center_strategy: CenterStrategy::FarthestPoint { num_centers: 14 },
2668 order_s: 1.3,
2669 double_penalty: false,
2670 ..MeasureJetBasisSpec::default()
2671 };
2672 let built = build_measure_jet_basis(data.view(), &spec).expect("fused build");
2673 assert_eq!(
2674 built.penalties.len(),
2675 1,
2676 "single-scale mode without null recovery emits exactly one Primary"
2677 );
2678 assert!(matches!(
2679 built.penaltyinfo[0].source,
2680 PenaltySource::Primary
2681 ));
2682 let BasisMetadata::MeasureJet { order_s, .. } = &built.metadata else {
2683 panic!("measure-jet build must return MeasureJet metadata");
2684 };
2685 assert_eq!(*order_s, 1.3, "explicit order must persist verbatim");
2686 }
2687
2688 #[test]
2693 pub(crate) fn single_scale_affine_head_gauge_annihilates_center_cross() {
2694 let n = 90usize;
2695 let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2696 let t = i as f64 / (n as f64 - 1.0);
2697 if k == 0 {
2698 3.0 * t
2699 } else {
2700 (2.0 * std::f64::consts::PI * t).sin() + 0.2 * t
2701 }
2702 });
2703 let spec = MeasureJetBasisSpec {
2704 center_strategy: CenterStrategy::FarthestPoint { num_centers: 18 },
2705 double_penalty: false,
2706 multiscale: false,
2707 ..MeasureJetBasisSpec::default()
2708 };
2709 let geom = realize_measure_jet_geometry(data.view(), &spec).expect("realized geometry");
2710 let m = geom.centers.nrows();
2711 let head_rank = geom.head_transform.ncols();
2712 assert!(head_rank > 0, "fixture must realize an affine head");
2713 assert_eq!(
2714 geom.z.ncols(),
2715 m - 1,
2716 "affine gauge replaces duplicated RBF directions without widening the smooth"
2717 );
2718 let rbf_rank = m - (head_rank + 1);
2719 let z_rbf = geom.z.slice(ndarray::s![..m, ..rbf_rank]).to_owned();
2720 let k_cc =
2721 measure_jet_design_matrix(geom.centers.view(), geom.centers.view(), geom.length_scale)
2722 .expect("center kernel");
2723 let affine = measure_jet_affine_value_basis(geom.centers.view(), geom.masses.view());
2724 assert_eq!(affine.ncols(), head_rank + 1);
2725 let mut weighted_affine = affine.clone();
2726 for (i, mut row) in weighted_affine.outer_iter_mut().enumerate() {
2727 row.mapv_inplace(|v| v * geom.masses[i]);
2728 }
2729 let constraint_cross = k_cc.t().dot(&weighted_affine);
2730 let residual = constraint_cross.t().dot(&z_rbf);
2731 let scale = constraint_cross
2732 .iter()
2733 .fold(1.0_f64, |acc, value| acc.max(value.abs()));
2734 assert!(
2735 residual.iter().all(|value| value.abs() <= 1e-10 * scale),
2736 "A^T W Kcc Z_rbf must vanish; max residual {:.3e}",
2737 residual
2738 .iter()
2739 .fold(0.0_f64, |acc, value| acc.max(value.abs()))
2740 );
2741 }
2742
2743 #[test]
2746 pub(crate) fn affine_null_penalty_is_covariant_under_coefficient_reparameterization() {
2747 let centers = array![
2748 [-1.0, 0.2],
2749 [-0.4, -0.3],
2750 [0.1, 0.5],
2751 [0.7, -0.2],
2752 [1.2, 0.4],
2753 [1.8, -0.1],
2754 ];
2755 let masses = array![0.08, 0.12, 0.18, 0.22, 0.17, 0.23];
2756 let evaluation = Array2::<f64>::from_shape_fn((centers.nrows(), 3), |(i, j)| {
2757 ((i + 2 * j + 1) as f64).sin() + 0.15 * (i * (j + 1)) as f64
2758 });
2759 let reparameterization = array![[1.7, 0.2, -0.1], [0.0, 0.6, 0.3], [0.0, 0.0, 1.3]];
2760 let base = affine_function_nullspace_penalty(&evaluation, centers.view(), masses.view())
2761 .expect("base function-space penalty");
2762 let transformed_evaluation = evaluation.dot(&reparameterization);
2763 let transformed = affine_function_nullspace_penalty(
2764 &transformed_evaluation,
2765 centers.view(),
2766 masses.view(),
2767 )
2768 .expect("reparameterized function-space penalty");
2769 let expected = reparameterization.t().dot(&base).dot(&reparameterization);
2770 let scale = expected
2771 .iter()
2772 .fold(1.0_f64, |acc, value| acc.max(value.abs()));
2773 assert!(
2774 transformed
2775 .iter()
2776 .zip(expected.iter())
2777 .all(|(actual, want)| (actual - want).abs() <= 1e-11 * scale),
2778 "S(E R) must equal R^T S(E) R"
2779 );
2780 }
2781
2782 #[test]
2785 pub(crate) fn double_penalty_leaves_primary_matrix_unchanged() {
2786 let n = 64usize;
2787 let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2788 let t = i as f64 / (n as f64 - 1.0);
2789 if k == 0 { 2.5 * t } else { (4.0 * t).cos() }
2790 });
2791 let base = MeasureJetBasisSpec {
2792 center_strategy: CenterStrategy::FarthestPoint { num_centers: 16 },
2793 order_s: 1.25,
2794 double_penalty: false,
2795 ..MeasureJetBasisSpec::default()
2796 };
2797 let without = build_measure_jet_basis(data.view(), &base).expect("primary-only build");
2798 let with = build_measure_jet_basis(
2799 data.view(),
2800 &MeasureJetBasisSpec {
2801 double_penalty: true,
2802 ..base.clone()
2803 },
2804 )
2805 .expect("double-penalty build");
2806 assert_eq!(without.penalties.len(), 1);
2807 assert_eq!(with.penalties.len(), 2);
2808 assert!(matches!(
2809 without.penaltyinfo[0].source,
2810 PenaltySource::Primary
2811 ));
2812 assert!(matches!(with.penaltyinfo[0].source, PenaltySource::Primary));
2813 assert!(matches!(
2814 with.penaltyinfo[1].source,
2815 PenaltySource::DoublePenaltyNullspace
2816 ));
2817 assert!(
2818 without.penalties[0]
2819 .iter()
2820 .zip(with.penalties[0].iter())
2821 .all(|(a, b)| (a - b).abs() <= 1e-13),
2822 "turning on null recovery must not modify Primary"
2823 );
2824 }
2825
2826 #[test]
2828 pub(crate) fn householder_sum_to_zero_basis_is_orthonormal() {
2829 let m = 9usize;
2830 let u = householder_sum_to_zero_u(m);
2831 let z = householder_sum_to_zero_z(&u);
2832 for j in 0..(m - 1) {
2833 let col_j = z.column(j);
2834 assert!(col_j.sum().abs() <= 1e-12, "column {j} must sum to zero");
2835 for j2 in j..(m - 1) {
2836 let dot = col_j.dot(&z.column(j2));
2837 let want = if j == j2 { 1.0 } else { 0.0 };
2838 assert!(
2839 (dot - want).abs() <= 1e-12,
2840 "orthonormality failure at ({j}, {j2}): {dot}"
2841 );
2842 }
2843 }
2844 }
2845
2846 pub(crate) fn frozen_spec_fixture(
2851 order_s: f64,
2852 multiscale: bool,
2853 ) -> (Array2<f64>, MeasureJetBasisSpec) {
2854 let n = 140usize;
2859 let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
2860 let t = i as f64 / (n as f64 - 1.0);
2861 if k == 0 {
2862 t * 3.0
2863 } else {
2864 0.5 * (t * 3.0).cos() + if i % 9 == 0 { 0.8 } else { 0.0 }
2865 }
2866 });
2867 let spec = MeasureJetBasisSpec {
2868 center_strategy: CenterStrategy::FarthestPoint { num_centers: 70 },
2869 order_s,
2870 multiscale,
2871 learn_length_scale: false,
2875 ..MeasureJetBasisSpec::default()
2876 };
2877 let first = build_measure_jet_basis(data.view(), &spec).expect("fixture build");
2878 let BasisMetadata::MeasureJet {
2879 centers,
2880 length_scale,
2881 eps_band,
2882 masses,
2883 support_means,
2884 penalty_normalization_scales,
2885 raw_penalty_normalization_scales,
2886 fused_penalty_normalization_scale,
2887 constraint_transform,
2888 ..
2889 } = &first.metadata
2890 else {
2891 panic!("measure-jet build must return MeasureJet metadata");
2892 };
2893 let frozen = MeasureJetBasisSpec {
2894 center_strategy: CenterStrategy::UserProvided(centers.clone()),
2895 order_s,
2896 alpha: spec.alpha,
2897 tau0: spec.tau0,
2898 num_scales: eps_band.len(),
2899 length_scale: *length_scale,
2900 double_penalty: spec.double_penalty,
2901 learn_length_scale: false,
2902 multiscale,
2903 identifiability: MeasureJetIdentifiability::FrozenTransform {
2904 transform: constraint_transform.clone().expect("fit-time z"),
2905 },
2906 frozen_quadrature: Some(MeasureJetFrozenQuadrature {
2907 masses: masses.clone(),
2908 eps_band: eps_band.clone(),
2909 support_means: support_means.clone(),
2910 penalty_normalization_scales: penalty_normalization_scales.clone(),
2911 raw_penalty_normalization_scales: raw_penalty_normalization_scales.clone(),
2912 fused_penalty_normalization_scale: *fused_penalty_normalization_scale,
2913 sigma_coord: None,
2914 }),
2915 };
2916 (data, frozen)
2917 }
2918
2919 #[test]
2924 pub(crate) fn psi_producer_matches_fd_per_level_mode() {
2925 let (data, frozen) = frozen_spec_fixture(0.0, true);
2926 let derivs =
2927 build_measure_jet_basis_psi_derivatives(data.view(), &frozen).expect("psi derivatives");
2928 let l_count = frozen
2929 .frozen_quadrature
2930 .as_ref()
2931 .expect("frozen quadrature")
2932 .eps_band
2933 .len();
2934 assert_eq!(
2935 derivs.penalties_first.len(),
2936 2,
2937 "per-level coords are (α, lnτ)"
2938 );
2939 assert_eq!(derivs.penalties_first[0].len(), l_count + 1);
2940 assert_eq!(derivs.penalties_cross_pairs, vec![(0, 1)]);
2941 let pen_at = |alpha: f64, tau0: f64| {
2942 let trial = MeasureJetBasisSpec {
2943 alpha,
2944 tau0,
2945 ..frozen.clone()
2946 };
2947 build_measure_jet_basis(data.view(), &trial)
2948 .expect("trial build")
2949 .penalties
2950 };
2951 let h = 1e-4;
2954 let (a0, t0) = (frozen.alpha, frozen.tau0);
2955 let ap = pen_at(a0 + h, t0);
2956 let am = pen_at(a0 - h, t0);
2957 let tp = pen_at(a0, t0 * h.exp());
2958 let tm = pen_at(a0, t0 * (-h).exp());
2959 assert_eq!(
2960 ap.len(),
2961 l_count + 1,
2962 "fixture must keep every scale active"
2963 );
2964 for level in 0..l_count {
2965 let fd_alpha = (&ap[level] - &am[level]) / (2.0 * h);
2966 let fd_tau = (&tp[level] - &tm[level]) / (2.0 * h);
2967 for (name, analytic, fd) in [
2968 ("alpha", &derivs.penalties_first[0][level], fd_alpha),
2969 ("ln_tau", &derivs.penalties_first[1][level], fd_tau),
2970 ] {
2971 let scale = fd.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
2972 for (x, y) in analytic.iter().zip(fd.iter()) {
2973 assert!(
2974 (x - y).abs() <= 5e-5 * scale,
2975 "{name} jet of scale-candidate {level}: analytic {x:.6e} vs FD {y:.6e}"
2976 );
2977 }
2978 }
2979 }
2980 for coord in 0..2 {
2982 assert!(
2983 derivs.penalties_first[coord][l_count]
2984 .iter()
2985 .all(|v| *v == 0.0),
2986 "null-component candidate must have zero (α, lnτ) drift"
2987 );
2988 }
2989 let provider = derivs
2991 .penalties_cross_provider
2992 .as_ref()
2993 .expect("cross provider");
2994 let cross = provider.evaluate(0, 1).expect("cross pair (α, lnτ)");
2995 let pp = pen_at(a0 + h, t0 * h.exp());
2996 let pm = pen_at(a0 + h, t0 * (-h).exp());
2997 let mp = pen_at(a0 - h, t0 * h.exp());
2998 let mm = pen_at(a0 - h, t0 * (-h).exp());
2999 for level in 0..l_count {
3000 let fd = (&(&pp[level] - &pm[level]) - &(&mp[level] - &mm[level])) / (4.0 * h * h);
3001 let scale = fd.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
3002 for (x, y) in cross[level].iter().zip(fd.iter()) {
3003 assert!(
3004 (x - y).abs() <= 5e-4 * scale,
3005 "cross (α, lnτ) jet of scale-candidate {level}: analytic {x:.6e} vs FD {y:.6e}"
3006 );
3007 }
3008 }
3009 }
3010
3011 #[test]
3017 pub(crate) fn psi_producer_matches_fd_length_scale() {
3018 let (data, mut frozen) = frozen_spec_fixture(0.0, false);
3021 frozen.learn_length_scale = true;
3022 let derivs =
3023 build_measure_jet_basis_psi_derivatives(data.view(), &frozen).expect("psi derivatives");
3024 assert_eq!(
3026 derivs.design_first.len(),
3027 1,
3028 "single-scale + learn_length_scale enrolls exactly the ℓ coordinate"
3029 );
3030 assert_eq!(
3031 derivs.penalties_first[0].len(),
3032 2,
3033 "single-scale double penalty carries Primary + affine/null component"
3034 );
3035 let ell0 = frozen.length_scale;
3039 let build_at = |ell: f64| {
3040 let trial = MeasureJetBasisSpec {
3041 length_scale: ell,
3042 ..frozen.clone()
3043 };
3044 build_measure_jet_basis(data.view(), &trial).expect("trial build")
3045 };
3046 let h: f64 = 1e-4;
3047 let plus = build_at(ell0 * h.exp());
3048 let minus = build_at(ell0 * (-h).exp());
3049 let at = build_at(ell0);
3050 assert_eq!(
3051 plus.penalties.len(),
3052 2,
3053 "fixture must keep both candidates active"
3054 );
3055 assert_eq!(
3056 minus.penalties.len(),
3057 2,
3058 "fixture must keep both candidates active"
3059 );
3060 assert_eq!(
3061 at.penalties.len(),
3062 2,
3063 "fixture must keep both candidates active"
3064 );
3065
3066 let x_plus = plus.design.to_dense();
3067 let x_minus = minus.design.to_dense();
3068 let x_0 = at.design.to_dense();
3069 let fd_first = (&x_plus - &x_minus) / (2.0 * h);
3070 let fd_second = (&x_plus - &(&x_0 * 2.0) + &x_minus) / (h * h);
3071 let scale1 = fd_first.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
3072 for (x, y) in derivs.design_first[0].iter().zip(fd_first.iter()) {
3073 assert!(
3074 (x - y).abs() <= 5e-5 * scale1,
3075 "∂X/∂lnℓ: analytic {x:.6e} vs FD {y:.6e}"
3076 );
3077 }
3078 let scale2 = fd_second.iter().fold(1e-30_f64, |acc, v| acc.max(v.abs()));
3079 for (x, y) in derivs.design_second_diag[0].iter().zip(fd_second.iter()) {
3080 assert!(
3081 (x - y).abs() <= 1e-3 * scale2,
3082 "∂²X/∂lnℓ²: analytic {x:.6e} vs FD {y:.6e}"
3083 );
3084 }
3085
3086 for candidate in 0..2 {
3087 let fd_penalty_first =
3088 (&plus.penalties[candidate] - &minus.penalties[candidate]) / (2.0 * h);
3089 let fd_penalty_second = (&plus.penalties[candidate]
3090 - &(&at.penalties[candidate] * 2.0)
3091 + &minus.penalties[candidate])
3092 / (h * h);
3093 let first_scale = fd_penalty_first
3094 .iter()
3095 .fold(1e-12_f64, |acc, value| acc.max(value.abs()));
3096 let second_scale = fd_penalty_second
3097 .iter()
3098 .fold(1e-10_f64, |acc, value| acc.max(value.abs()));
3099 for (analytic, finite_difference) in derivs.penalties_first[0][candidate]
3100 .iter()
3101 .zip(fd_penalty_first.iter())
3102 {
3103 assert!(
3104 (analytic - finite_difference).abs() <= 1e-4 * first_scale,
3105 "candidate {candidate} ∂S~/∂lnℓ: analytic {analytic:.6e} vs FD {finite_difference:.6e}"
3106 );
3107 }
3108 for (analytic, finite_difference) in derivs.penalties_second_diag[0][candidate]
3109 .iter()
3110 .zip(fd_penalty_second.iter())
3111 {
3112 assert!(
3113 (analytic - finite_difference).abs() <= 5e-3 * second_scale,
3114 "candidate {candidate} ∂²S~/∂lnℓ²: analytic {analytic:.6e} vs FD {finite_difference:.6e}"
3115 );
3116 }
3117 }
3118 }
3119
3120 #[test]
3124 pub(crate) fn quadrature_nodes_are_cell_barycenters() {
3125 let data = array![
3128 [0.0, 0.2],
3129 [0.4, -0.2],
3130 [0.2, 0.0],
3131 [9.8, 10.1],
3132 [10.2, 9.9],
3133 ];
3134 let seeds = array![[0.1, 0.1], [10.0, 10.0], [-50.0, -50.0]];
3135 let (nodes, masses) =
3136 measure_jet_quadrature_nodes(data.view(), seeds.view()).expect("quadrature nodes");
3137 assert!((masses.sum() - 1.0).abs() <= 1e-15, "masses must sum to 1");
3138 assert!((masses[0] - 0.6).abs() <= 1e-15);
3139 assert!((masses[1] - 0.4).abs() <= 1e-15);
3140 assert_eq!(masses[2], 0.0);
3141 assert_eq!(nodes[(0, 0)], 0.2);
3143 assert_eq!(nodes[(0, 1)], 0.0);
3144 assert_eq!(nodes[(1, 0)], 10.0);
3146 assert_eq!(nodes[(1, 1)], 10.0);
3147 assert_eq!(nodes[(2, 0)], -50.0);
3149 assert_eq!(nodes[(2, 1)], -50.0);
3150 }
3151
3152 #[test]
3156 pub(crate) fn build_replay_roundtrip_reproduces_design_and_penalty() {
3157 let n = 140usize;
3160 let data = Array2::<f64>::from_shape_fn((n, 2), |(i, k)| {
3161 let t = i as f64 / (n as f64 - 1.0);
3162 if k == 0 {
3163 t * 3.0
3164 } else {
3165 0.5 * (t * 3.0).cos() + if i % 9 == 0 { 0.8 } else { 0.0 }
3166 }
3167 });
3168 let spec = MeasureJetBasisSpec {
3169 center_strategy: CenterStrategy::FarthestPoint { num_centers: 70 },
3170 multiscale: true,
3171 ..MeasureJetBasisSpec::default()
3172 };
3173 let first = build_measure_jet_basis(data.view(), &spec).expect("first build");
3174 let BasisMetadata::MeasureJet {
3175 centers,
3176 length_scale,
3177 eps_band,
3178 order_s,
3179 alpha,
3180 tau0,
3181 masses,
3182 support_means,
3183 penalty_normalization_scales,
3184 raw_penalty_normalization_scales,
3185 fused_penalty_normalization_scale,
3186 constraint_transform,
3187 ..
3188 } = &first.metadata
3189 else {
3190 panic!("measure-jet build must return MeasureJet metadata");
3191 };
3192 let replay_spec = MeasureJetBasisSpec {
3193 center_strategy: CenterStrategy::UserProvided(centers.clone()),
3194 order_s: *order_s,
3195 alpha: *alpha,
3196 tau0: *tau0,
3197 num_scales: eps_band.len(),
3198 length_scale: *length_scale,
3199 double_penalty: spec.double_penalty,
3200 learn_length_scale: spec.learn_length_scale,
3201 multiscale: spec.multiscale,
3202 identifiability: MeasureJetIdentifiability::FrozenTransform {
3203 transform: constraint_transform.clone().expect("fit-time z"),
3204 },
3205 frozen_quadrature: Some(MeasureJetFrozenQuadrature {
3206 masses: masses.clone(),
3207 eps_band: eps_band.clone(),
3208 support_means: support_means.clone(),
3209 penalty_normalization_scales: penalty_normalization_scales.clone(),
3210 raw_penalty_normalization_scales: raw_penalty_normalization_scales.clone(),
3211 fused_penalty_normalization_scale: *fused_penalty_normalization_scale,
3212 sigma_coord: None,
3213 }),
3214 };
3215 assert_eq!(
3218 first.penalties.len(),
3219 eps_band.len() + 1,
3220 "per-level mode must emit one candidate per scale + null component"
3221 );
3222 let second = build_measure_jet_basis(data.view(), &replay_spec).expect("replay build");
3223 let x1 = first.design.to_dense();
3224 let x2 = second.design.to_dense();
3225 assert_eq!(x1.shape(), x2.shape());
3226 for (a, b) in x1.iter().zip(x2.iter()) {
3227 assert!((a - b).abs() <= 1e-12, "design replay drift: {a} vs {b}");
3228 }
3229 assert_eq!(first.penalties.len(), second.penalties.len());
3230 for (p1, p2) in first.penalties.iter().zip(second.penalties.iter()) {
3231 for (a, b) in p1.iter().zip(p2.iter()) {
3232 assert!((a - b).abs() <= 1e-12, "penalty replay drift: {a} vs {b}");
3233 }
3234 }
3235 }
3236}