1use ndarray::{Array1, ArrayView1, ArrayView2};
76
77use super::BasisError;
78
79pub fn measure_jet_ambient_gradient(
111 query: ArrayView1<'_, f64>,
112 centers: ArrayView2<'_, f64>,
113 representer_coeffs: ArrayView1<'_, f64>,
114 length_scale: f64,
115 head_transform: Option<ArrayView2<'_, f64>>,
116 head_coeffs: ArrayView1<'_, f64>,
117) -> Result<Array1<f64>, BasisError> {
118 let d = query.len();
119 let m = centers.nrows();
120 if centers.ncols() != d {
121 crate::bail_dim_basis!(
122 "measure-jet ambient gradient: query dimension {d} disagrees with centers ({} × {})",
123 m,
124 centers.ncols()
125 );
126 }
127 if representer_coeffs.len() != m {
128 crate::bail_dim_basis!(
129 "measure-jet ambient gradient: {} representer coefficients for {m} centers",
130 representer_coeffs.len()
131 );
132 }
133 if !(length_scale.is_finite() && length_scale > 0.0) {
134 crate::bail_invalid_basis!(
135 "measure-jet ambient gradient needs a positive finite length_scale; got {length_scale}"
136 );
137 }
138 let inv_l2 = 1.0 / (length_scale * length_scale);
139 let mut grad = Array1::<f64>::zeros(d);
140 for i in 0..m {
141 let center = centers.row(i);
142 let mut sq = 0.0_f64;
143 for a in 0..d {
144 let delta = query[a] - center[a];
145 sq += delta * delta;
146 }
147 let k = (-0.5 * sq * inv_l2).exp();
148 let coeff = representer_coeffs[i] * k * inv_l2;
149 for a in 0..d {
150 grad[a] += coeff * (center[a] - query[a]);
152 }
153 }
154 if let Some(t) = head_transform {
155 let head_rank = t.ncols();
156 if head_coeffs.len() != head_rank {
157 crate::bail_dim_basis!(
158 "measure-jet ambient gradient: {} head coefficients for a head lift with \
159 {head_rank} columns",
160 head_coeffs.len()
161 );
162 }
163 if t.nrows() != d {
164 crate::bail_dim_basis!(
165 "measure-jet ambient gradient: head lift has {} rows but ambient dimension is {d}",
166 t.nrows()
167 );
168 }
169 for a in 0..d {
170 let mut acc = 0.0_f64;
171 for h in 0..head_rank {
172 acc += head_coeffs[h] * t[(a, h)];
173 }
174 grad[a] += acc;
175 }
176 } else if !head_coeffs.is_empty() {
177 crate::bail_dim_basis!(
178 "measure-jet ambient gradient: {} head coefficients supplied without a head lift",
179 head_coeffs.len()
180 );
181 }
182 Ok(grad)
183}
184
185#[derive(Clone, Copy)]
186pub enum MeasureJetExtrapolationSpectrum<'a> {
187 PerLevel(&'a [f64]),
189 Fused(f64),
192}
193
194pub fn measure_jet_extrapolation_variance(
212 support_row: ArrayView1<'_, f64>,
213 eps_band: &[f64],
214 support_means: &[f64],
215 spectrum: MeasureJetExtrapolationSpectrum<'_>,
216 coverage_floor: f64,
217) -> Result<f64, BasisError> {
218 let n_levels = eps_band.len();
219 if n_levels == 0 {
220 crate::bail_invalid_basis!("measure-jet extrapolation variance needs a nonempty band");
221 }
222 if support_row.len() != n_levels || support_means.len() != n_levels {
223 crate::bail_dim_basis!(
224 "measure-jet extrapolation variance needs one support value and one support mean per \
225 band scale: {} support values, {} support means, {} scales",
226 support_row.len(),
227 support_means.len(),
228 n_levels
229 );
230 }
231 for (l, pair) in eps_band.windows(2).enumerate() {
232 if pair[1] <= pair[0] {
233 crate::bail_invalid_basis!(
234 "measure-jet band must be strictly ascending: eps[{l}] = {} vs eps[{}] = {}",
235 pair[0],
236 l + 1,
237 pair[1]
238 );
239 }
240 }
241 if eps_band.iter().any(|e| !(e.is_finite() && *e > 0.0)) {
242 crate::bail_invalid_basis!("measure-jet band scales must be finite and positive");
243 }
244 if support_row.iter().any(|q| !(q.is_finite() && *q >= 0.0)) {
245 crate::bail_invalid_basis!(
246 "measure-jet support row must be finite and nonnegative (kernel masses)"
247 );
248 }
249 if support_means.iter().any(|q| !(q.is_finite() && *q > 0.0)) {
250 crate::bail_invalid_basis!("measure-jet support means must be finite and positive");
251 }
252 if !(coverage_floor.is_finite() && coverage_floor > 0.0 && coverage_floor < 1.0) {
253 crate::bail_invalid_basis!(
254 "measure-jet coverage floor must lie strictly in (0, 1); got {coverage_floor}"
255 );
256 }
257 match spectrum {
258 MeasureJetExtrapolationSpectrum::PerLevel(lambda_hat) => {
259 if lambda_hat.len() != n_levels {
260 crate::bail_dim_basis!(
261 "measure-jet per-level extrapolation variance needs one physical precision per \
262 band scale: {} precisions, {} scales",
263 lambda_hat.len(),
264 n_levels
265 );
266 }
267 if lambda_hat.iter().any(|l| !(l.is_finite() && *l > 0.0)) {
268 crate::bail_invalid_basis!(
269 "measure-jet per-scale amplitudes must be finite and positive (physical precisions)"
270 );
271 }
272 let first_covering = support_row
273 .iter()
274 .zip(support_means.iter())
275 .position(|(q, q_bar)| *q >= coverage_floor * *q_bar)
276 .unwrap_or(n_levels);
277 let mut variance = 0.0_f64;
278 for (l, ((&q, &q_bar), &lam)) in support_row
279 .iter()
280 .zip(support_means.iter())
281 .zip(lambda_hat.iter())
282 .enumerate()
283 {
284 let weight = if l < first_covering {
285 1.0
286 } else {
287 1.0 - (q / q_bar).min(1.0)
288 };
289 variance += weight / lam;
290 }
291 Ok(variance)
292 }
293 MeasureJetExtrapolationSpectrum::Fused(lambda_hat) => {
294 if !(lambda_hat.is_finite() && lambda_hat > 0.0) {
295 crate::bail_invalid_basis!(
296 "measure-jet fused amplitude must be finite and positive (physical precision)"
297 );
298 }
299 let mut best_coverage = 0.0_f64;
300 let mut covered = false;
301 for (&q, &q_bar) in support_row.iter().zip(support_means.iter()) {
302 let coverage = (q / q_bar).min(1.0);
303 best_coverage = best_coverage.max(coverage);
304 if q >= coverage_floor * q_bar {
305 covered = true;
306 }
307 }
308 let weight = if covered { 1.0 - best_coverage } else { 1.0 };
309 Ok(weight / lambda_hat)
310 }
311 }
312}
313
314#[cfg(test)]
315mod tests {
316 use super::*;
317 use ndarray::{Array1, arr1};
318
319 pub(crate) fn band() -> Vec<f64> {
322 vec![0.05, 0.1, 0.2, 0.4, 0.8]
323 }
324
325 pub(crate) fn lambdas() -> Vec<f64> {
326 vec![40.0, 11.0, 3.5, 1.25, 0.6]
327 }
328
329 pub(crate) fn support_means(eps: &[f64]) -> Vec<f64> {
330 vec![TOTAL; eps.len()]
331 }
332
333 pub(crate) const FLOOR: f64 = 0.05;
334 pub(crate) const TOTAL: f64 = 1.0;
335
336 pub(crate) fn total_ignorance(lams: &[f64]) -> f64 {
337 lams.iter().map(|l| 1.0 / l).sum()
338 }
339
340 pub(crate) fn support_at_distance(d: f64, eps: &[f64]) -> Array1<f64> {
344 Array1::from_iter(eps.iter().map(|e| TOTAL * (-d * d / (2.0 * e * e)).exp()))
345 }
346
347 #[test]
352 pub(crate) fn extrapolation_variance_is_monotone_in_distance() {
353 let eps = band();
354 let lams = lambdas();
355 let q_bar = support_means(&eps);
356 let mut prev = -1.0_f64;
357 for step in 0..400 {
360 let d = 0.015 * step as f64;
361 let row = support_at_distance(d, &eps);
362 let v = measure_jet_extrapolation_variance(
363 row.view(),
364 &eps,
365 &q_bar,
366 MeasureJetExtrapolationSpectrum::PerLevel(&lams),
367 FLOOR,
368 )
369 .expect("valid inputs");
370 assert!(
371 v >= prev,
372 "variance decreased with distance: variance({d:.3}) = {v:.12} < {prev:.12}"
373 );
374 prev = v;
375 }
376 assert!(
379 (prev - total_ignorance(&lams)).abs() <= 1e-12,
380 "far-field variance must saturate at Σ 1/λ̂: got {prev}"
381 );
382 }
383
384 #[test]
388 pub(crate) fn extrapolation_variance_is_monotone_under_pointwise_domination() {
389 let eps = band();
390 let lams = lambdas();
391 let q_bar = support_means(&eps);
392 let rows = [
393 arr1(&[0.9, 0.95, 0.99, 1.0, 1.0]),
394 arr1(&[0.02, 0.3, 0.06, 0.8, 0.97]),
395 arr1(&[0.0, 0.0, 0.04, 0.2, 0.6]),
396 arr1(&[0.04, 0.04, 0.04, 0.04, 0.049]),
397 ];
398 for row in &rows {
399 for shrink in [1.0, 0.9, 0.7, 0.3, 0.0] {
400 let smaller = row.mapv(|q| shrink * q);
401 let v_big = measure_jet_extrapolation_variance(
402 row.view(),
403 &eps,
404 &q_bar,
405 MeasureJetExtrapolationSpectrum::PerLevel(&lams),
406 FLOOR,
407 )
408 .expect("valid inputs");
409 let v_small = measure_jet_extrapolation_variance(
410 smaller.view(),
411 &eps,
412 &q_bar,
413 MeasureJetExtrapolationSpectrum::PerLevel(&lams),
414 FLOOR,
415 )
416 .expect("valid inputs");
417 assert!(
418 v_small >= v_big,
419 "pointwise-smaller support gave smaller variance: {v_small} < {v_big} \
420 (row {row:?}, shrink {shrink})"
421 );
422 }
423 }
424 }
425
426 #[test]
430 pub(crate) fn extrapolation_variance_vanishes_on_web() {
431 let eps = band();
432 let lams = lambdas();
433 let q_bar = support_means(&eps);
434 let full = Array1::from_elem(eps.len(), TOTAL);
435 let v_full = measure_jet_extrapolation_variance(
436 full.view(),
437 &eps,
438 &q_bar,
439 MeasureJetExtrapolationSpectrum::PerLevel(&lams),
440 FLOOR,
441 )
442 .expect("valid inputs");
443 assert_eq!(v_full, 0.0, "full coverage must price zero extra variance");
444
445 let near = Array1::from_elem(eps.len(), 0.97 * TOTAL);
446 let v_near = measure_jet_extrapolation_variance(
447 near.view(),
448 &eps,
449 &q_bar,
450 MeasureJetExtrapolationSpectrum::PerLevel(&lams),
451 FLOOR,
452 )
453 .expect("valid inputs");
454 let budget = total_ignorance(&lams);
455 assert!(
456 v_near <= 0.05 * budget,
457 "near-full coverage must price a small fraction of Σ 1/λ̂: {v_near} vs budget {budget}"
458 );
459 }
460
461 #[test]
464 pub(crate) fn extrapolation_variance_saturates_off_web() {
465 let eps = band();
466 let lams = lambdas();
467 let q_bar = support_means(&eps);
468 let zero = Array1::<f64>::zeros(eps.len());
469 let v = measure_jet_extrapolation_variance(
470 zero.view(),
471 &eps,
472 &q_bar,
473 MeasureJetExtrapolationSpectrum::PerLevel(&lams),
474 FLOOR,
475 )
476 .expect("valid inputs");
477 assert_eq!(
478 v,
479 total_ignorance(&lams),
480 "never-covered query must pay Σ 1/λ̂ exactly"
481 );
482 }
483
484 #[test]
487 pub(crate) fn extrapolation_variance_halves_when_amplitudes_double() {
488 let eps = band();
489 let lams = lambdas();
490 let q_bar = support_means(&eps);
491 let doubled: Vec<f64> = lams.iter().map(|l| 2.0 * l).collect();
492 let rows = [
495 support_at_distance(0.35, &eps),
496 Array1::<f64>::zeros(eps.len()),
497 Array1::from_elem(eps.len(), 0.5),
498 ];
499 for row in &rows {
500 let v1 = measure_jet_extrapolation_variance(
501 row.view(),
502 &eps,
503 &q_bar,
504 MeasureJetExtrapolationSpectrum::PerLevel(&lams),
505 FLOOR,
506 )
507 .expect("valid inputs");
508 let v2 = measure_jet_extrapolation_variance(
509 row.view(),
510 &eps,
511 &q_bar,
512 MeasureJetExtrapolationSpectrum::PerLevel(&doubled),
513 FLOOR,
514 )
515 .expect("valid inputs");
516 assert!(
517 (2.0 * v2 - v1).abs() <= 1e-15 * v1.max(1.0),
518 "doubling λ̂ must halve the variance: {v1} vs 2×{v2}"
519 );
520 }
521 }
522
523 #[test]
529 pub(crate) fn extrapolation_variance_gate_convention() {
530 let eps = band();
531 let lams = lambdas();
532 let q_bar = support_means(&eps);
533 let sub_floor = Array1::from_elem(eps.len(), 0.049 * TOTAL);
534 let v_sub = measure_jet_extrapolation_variance(
535 sub_floor.view(),
536 &eps,
537 &q_bar,
538 MeasureJetExtrapolationSpectrum::PerLevel(&lams),
539 FLOOR,
540 )
541 .expect("valid inputs");
542 assert_eq!(
543 v_sub,
544 total_ignorance(&lams),
545 "sub-floor mass earns no credit: full Σ 1/λ̂"
546 );
547
548 let mut at_floor = sub_floor.clone();
550 at_floor[eps.len() - 1] = FLOOR * TOTAL;
551 let v_floor = measure_jet_extrapolation_variance(
552 at_floor.view(),
553 &eps,
554 &q_bar,
555 MeasureJetExtrapolationSpectrum::PerLevel(&lams),
556 FLOOR,
557 )
558 .expect("valid inputs");
559 let expected: f64 = lams[..eps.len() - 1].iter().map(|l| 1.0 / l).sum::<f64>()
560 + (1.0 - FLOOR) / lams[eps.len() - 1];
561 assert!(
562 (v_floor - expected).abs() <= 1e-15,
563 "floor-clearing coarsest level must take weight 1 − a: {v_floor} vs {expected}"
564 );
565 assert!(
568 v_sub - v_floor <= FLOOR * total_ignorance(&lams) + 1e-15,
569 "gate jump exceeds the documented coverage_floor bound"
570 );
571 }
572
573 #[test]
574 pub(crate) fn fused_extrapolation_charges_single_band_amplitude_once() {
575 let eps = band();
576 let q_bar = support_means(&eps);
577 let lam = 2.5;
578 let zero = Array1::<f64>::zeros(eps.len());
579 let v_zero = measure_jet_extrapolation_variance(
580 zero.view(),
581 &eps,
582 &q_bar,
583 MeasureJetExtrapolationSpectrum::Fused(lam),
584 FLOOR,
585 )
586 .expect("valid inputs");
587 assert_eq!(
588 v_zero,
589 1.0 / lam,
590 "never-covered fused band must pay one amplitude, not one per level"
591 );
592
593 let covered = arr1(&[0.01, 0.2, 0.4, 0.75, 0.5]);
594 let v_covered = measure_jet_extrapolation_variance(
595 covered.view(),
596 &eps,
597 &q_bar,
598 MeasureJetExtrapolationSpectrum::Fused(lam),
599 FLOOR,
600 )
601 .expect("valid inputs");
602 let expected = (1.0 - 0.75) / lam;
603 assert!(
604 (v_covered - expected).abs() <= 1e-15,
605 "fused band must use the best covered level once: {v_covered} vs {expected}"
606 );
607 }
608
609 fn eval_fitted(
612 query: ArrayView1<'_, f64>,
613 centers: ArrayView2<'_, f64>,
614 z: ArrayView1<'_, f64>,
615 length_scale: f64,
616 head: Option<ArrayView2<'_, f64>>,
617 head_coeffs: ArrayView1<'_, f64>,
618 ) -> f64 {
619 let inv_two_l2 = 1.0 / (2.0 * length_scale * length_scale);
620 let mut val = 0.0_f64;
621 for i in 0..centers.nrows() {
622 let mut sq = 0.0_f64;
623 for a in 0..query.len() {
624 let dlt = query[a] - centers[(i, a)];
625 sq += dlt * dlt;
626 }
627 val += z[i] * (-sq * inv_two_l2).exp();
628 }
629 if let Some(t) = head {
630 for h in 0..t.ncols() {
631 let mut proj = 0.0_f64;
632 for a in 0..query.len() {
633 proj += query[a] * t[(a, h)];
634 }
635 val += head_coeffs[h] * proj;
636 }
637 }
638 val
639 }
640
641 #[test]
645 pub(crate) fn ambient_gradient_matches_central_difference() {
646 use ndarray::{arr1, arr2};
647 let centers = arr2(&[[0.0, 0.0], [1.0, 0.5], [-0.7, 0.9], [0.4, -1.1]]);
648 let z = arr1(&[0.8, -1.3, 0.5, 2.0]);
649 let length_scale = 0.6;
650 let t = arr2(&[[1.0, 0.2], [-0.3, 0.9]]);
652 let head_coeffs = arr1(&[0.7, -0.4]);
653 let query = arr1(&[0.15, -0.2]);
654
655 let grad = measure_jet_ambient_gradient(
656 query.view(),
657 centers.view(),
658 z.view(),
659 length_scale,
660 Some(t.view()),
661 head_coeffs.view(),
662 )
663 .expect("valid gradient");
664
665 let h = 1e-6;
666 for a in 0..query.len() {
667 let mut qp = query.clone();
668 let mut qm = query.clone();
669 qp[a] += h;
670 qm[a] -= h;
671 let fp = eval_fitted(
672 qp.view(),
673 centers.view(),
674 z.view(),
675 length_scale,
676 Some(t.view()),
677 head_coeffs.view(),
678 );
679 let fm = eval_fitted(
680 qm.view(),
681 centers.view(),
682 z.view(),
683 length_scale,
684 Some(t.view()),
685 head_coeffs.view(),
686 );
687 let fd = (fp - fm) / (2.0 * h);
688 assert!(
689 (grad[a] - fd).abs() <= 1e-6 * (1.0 + fd.abs()),
690 "axis {a}: analytic {} vs central FD {fd}",
691 grad[a]
692 );
693 }
694 }
695
696 #[test]
699 pub(crate) fn ambient_gradient_representer_only_and_head_guard() {
700 use ndarray::{arr1, arr2};
701 let centers = arr2(&[[0.0], [0.5], [-0.4]]);
702 let z = arr1(&[1.0, -2.0, 0.5]);
703 let length_scale = 0.3;
704 let query = arr1(&[0.1]);
705 let empty = Array1::<f64>::zeros(0);
706 let grad = measure_jet_ambient_gradient(
707 query.view(),
708 centers.view(),
709 z.view(),
710 length_scale,
711 None,
712 empty.view(),
713 )
714 .expect("valid gradient");
715 let h = 1e-6;
716 let mut qp = query.clone();
717 let mut qm = query.clone();
718 qp[0] += h;
719 qm[0] -= h;
720 let fd = (eval_fitted(
721 qp.view(),
722 centers.view(),
723 z.view(),
724 length_scale,
725 None,
726 empty.view(),
727 ) - eval_fitted(
728 qm.view(),
729 centers.view(),
730 z.view(),
731 length_scale,
732 None,
733 empty.view(),
734 )) / (2.0 * h);
735 assert!((grad[0] - fd).abs() <= 1e-6 * (1.0 + fd.abs()));
736
737 let stray = arr1(&[1.0]);
738 assert!(
739 measure_jet_ambient_gradient(
740 query.view(),
741 centers.view(),
742 z.view(),
743 length_scale,
744 None,
745 stray.view(),
746 )
747 .is_err(),
748 "head coefficients without a head lift must be rejected"
749 );
750 }
751}