use ndarray::{Array1, ArrayView1, ArrayView2};
use super::BasisError;
pub fn measure_jet_ambient_gradient(
query: ArrayView1<'_, f64>,
centers: ArrayView2<'_, f64>,
representer_coeffs: ArrayView1<'_, f64>,
length_scale: f64,
head_transform: Option<ArrayView2<'_, f64>>,
head_coeffs: ArrayView1<'_, f64>,
) -> Result<Array1<f64>, BasisError> {
let d = query.len();
let m = centers.nrows();
if centers.ncols() != d {
crate::bail_dim_basis!(
"measure-jet ambient gradient: query dimension {d} disagrees with centers ({} × {})",
m,
centers.ncols()
);
}
if representer_coeffs.len() != m {
crate::bail_dim_basis!(
"measure-jet ambient gradient: {} representer coefficients for {m} centers",
representer_coeffs.len()
);
}
if !(length_scale.is_finite() && length_scale > 0.0) {
crate::bail_invalid_basis!(
"measure-jet ambient gradient needs a positive finite length_scale; got {length_scale}"
);
}
let inv_l2 = 1.0 / (length_scale * length_scale);
let mut grad = Array1::<f64>::zeros(d);
for i in 0..m {
let center = centers.row(i);
let mut sq = 0.0_f64;
for a in 0..d {
let delta = query[a] - center[a];
sq += delta * delta;
}
let k = (-0.5 * sq * inv_l2).exp();
let coeff = representer_coeffs[i] * k * inv_l2;
for a in 0..d {
grad[a] += coeff * (center[a] - query[a]);
}
}
if let Some(t) = head_transform {
let head_width = t.ncols();
if head_coeffs.len() != head_width {
crate::bail_dim_basis!(
"measure-jet ambient gradient: {} head coefficients for a head lift with \
{head_width} columns",
head_coeffs.len()
);
}
if t.nrows() != d + 1 {
crate::bail_dim_basis!(
"measure-jet ambient gradient: affine head lift has {} rows but ambient \
dimension is {d} (expected d+1)",
t.nrows()
);
}
for a in 0..d {
let mut acc = 0.0_f64;
for h in 0..head_width {
acc += head_coeffs[h] * t[(a + 1, h)];
}
grad[a] += acc;
}
} else if !head_coeffs.is_empty() {
crate::bail_dim_basis!(
"measure-jet ambient gradient: {} head coefficients supplied without a head lift",
head_coeffs.len()
);
}
Ok(grad)
}
#[derive(Clone, Copy)]
pub enum MeasureJetExtrapolationSpectrum<'a> {
PerLevel(&'a [f64]),
Fused(f64),
}
pub fn measure_jet_extrapolation_variance(
support_row: ArrayView1<'_, f64>,
eps_band: &[f64],
support_means: &[f64],
spectrum: MeasureJetExtrapolationSpectrum<'_>,
coverage_floor: f64,
) -> Result<f64, BasisError> {
let n_levels = eps_band.len();
if n_levels == 0 {
crate::bail_invalid_basis!("measure-jet extrapolation variance needs a nonempty band");
}
if support_row.len() != n_levels || support_means.len() != n_levels {
crate::bail_dim_basis!(
"measure-jet extrapolation variance needs one support value and one support mean per \
band scale: {} support values, {} support means, {} scales",
support_row.len(),
support_means.len(),
n_levels
);
}
for (l, pair) in eps_band.windows(2).enumerate() {
if pair[1] <= pair[0] {
crate::bail_invalid_basis!(
"measure-jet band must be strictly ascending: eps[{l}] = {} vs eps[{}] = {}",
pair[0],
l + 1,
pair[1]
);
}
}
if eps_band.iter().any(|e| !(e.is_finite() && *e > 0.0)) {
crate::bail_invalid_basis!("measure-jet band scales must be finite and positive");
}
if support_row.iter().any(|q| !(q.is_finite() && *q >= 0.0)) {
crate::bail_invalid_basis!(
"measure-jet support row must be finite and nonnegative (kernel masses)"
);
}
if support_means.iter().any(|q| !(q.is_finite() && *q > 0.0)) {
crate::bail_invalid_basis!("measure-jet support means must be finite and positive");
}
if !(coverage_floor.is_finite() && coverage_floor > 0.0 && coverage_floor < 1.0) {
crate::bail_invalid_basis!(
"measure-jet coverage floor must lie strictly in (0, 1); got {coverage_floor}"
);
}
match spectrum {
MeasureJetExtrapolationSpectrum::PerLevel(lambda_hat) => {
if lambda_hat.len() != n_levels {
crate::bail_dim_basis!(
"measure-jet per-level extrapolation variance needs one physical precision per \
band scale: {} precisions, {} scales",
lambda_hat.len(),
n_levels
);
}
if lambda_hat.iter().any(|l| !(l.is_finite() && *l > 0.0)) {
crate::bail_invalid_basis!(
"measure-jet per-scale amplitudes must be finite and positive (physical precisions)"
);
}
let first_covering = support_row
.iter()
.zip(support_means.iter())
.position(|(q, q_bar)| *q >= coverage_floor * *q_bar)
.unwrap_or(n_levels);
let mut variance = 0.0_f64;
for (l, ((&q, &q_bar), &lam)) in support_row
.iter()
.zip(support_means.iter())
.zip(lambda_hat.iter())
.enumerate()
{
let weight = if l < first_covering {
1.0
} else {
1.0 - (q / q_bar).min(1.0)
};
variance += weight / lam;
}
Ok(variance)
}
MeasureJetExtrapolationSpectrum::Fused(lambda_hat) => {
if !(lambda_hat.is_finite() && lambda_hat > 0.0) {
crate::bail_invalid_basis!(
"measure-jet fused amplitude must be finite and positive (physical precision)"
);
}
let mut best_coverage = 0.0_f64;
let mut covered = false;
for (&q, &q_bar) in support_row.iter().zip(support_means.iter()) {
let coverage = (q / q_bar).min(1.0);
best_coverage = best_coverage.max(coverage);
if q >= coverage_floor * q_bar {
covered = true;
}
}
let weight = if covered { 1.0 - best_coverage } else { 1.0 };
Ok(weight / lambda_hat)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use ndarray::{Array1, arr1};
pub(crate) fn band() -> Vec<f64> {
vec![0.05, 0.1, 0.2, 0.4, 0.8]
}
pub(crate) fn lambdas() -> Vec<f64> {
vec![40.0, 11.0, 3.5, 1.25, 0.6]
}
pub(crate) fn support_means(eps: &[f64]) -> Vec<f64> {
vec![TOTAL; eps.len()]
}
pub(crate) const FLOOR: f64 = 0.05;
pub(crate) const TOTAL: f64 = 1.0;
pub(crate) fn total_ignorance(lams: &[f64]) -> f64 {
lams.iter().map(|l| 1.0 / l).sum()
}
pub(crate) fn support_at_distance(d: f64, eps: &[f64]) -> Array1<f64> {
Array1::from_iter(eps.iter().map(|e| TOTAL * (-d * d / (2.0 * e * e)).exp()))
}
#[test]
pub(crate) fn extrapolation_variance_is_monotone_in_distance() {
let eps = band();
let lams = lambdas();
let q_bar = support_means(&eps);
let mut prev = -1.0_f64;
for step in 0..400 {
let d = 0.015 * step as f64;
let row = support_at_distance(d, &eps);
let v = measure_jet_extrapolation_variance(
row.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::PerLevel(&lams),
FLOOR,
)
.expect("valid inputs");
assert!(
v >= prev,
"variance decreased with distance: variance({d:.3}) = {v:.12} < {prev:.12}"
);
prev = v;
}
assert!(
(prev - total_ignorance(&lams)).abs() <= 1e-12,
"far-field variance must saturate at Σ 1/λ̂: got {prev}"
);
}
#[test]
pub(crate) fn extrapolation_variance_is_monotone_under_pointwise_domination() {
let eps = band();
let lams = lambdas();
let q_bar = support_means(&eps);
let rows = [
arr1(&[0.9, 0.95, 0.99, 1.0, 1.0]),
arr1(&[0.02, 0.3, 0.06, 0.8, 0.97]),
arr1(&[0.0, 0.0, 0.04, 0.2, 0.6]),
arr1(&[0.04, 0.04, 0.04, 0.04, 0.049]),
];
for row in &rows {
for shrink in [1.0, 0.9, 0.7, 0.3, 0.0] {
let smaller = row.mapv(|q| shrink * q);
let v_big = measure_jet_extrapolation_variance(
row.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::PerLevel(&lams),
FLOOR,
)
.expect("valid inputs");
let v_small = measure_jet_extrapolation_variance(
smaller.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::PerLevel(&lams),
FLOOR,
)
.expect("valid inputs");
assert!(
v_small >= v_big,
"pointwise-smaller support gave smaller variance: {v_small} < {v_big} \
(row {row:?}, shrink {shrink})"
);
}
}
}
#[test]
pub(crate) fn extrapolation_variance_vanishes_on_web() {
let eps = band();
let lams = lambdas();
let q_bar = support_means(&eps);
let full = Array1::from_elem(eps.len(), TOTAL);
let v_full = measure_jet_extrapolation_variance(
full.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::PerLevel(&lams),
FLOOR,
)
.expect("valid inputs");
assert_eq!(v_full, 0.0, "full coverage must price zero extra variance");
let near = Array1::from_elem(eps.len(), 0.97 * TOTAL);
let v_near = measure_jet_extrapolation_variance(
near.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::PerLevel(&lams),
FLOOR,
)
.expect("valid inputs");
let budget = total_ignorance(&lams);
assert!(
v_near <= 0.05 * budget,
"near-full coverage must price a small fraction of Σ 1/λ̂: {v_near} vs budget {budget}"
);
}
#[test]
pub(crate) fn extrapolation_variance_saturates_off_web() {
let eps = band();
let lams = lambdas();
let q_bar = support_means(&eps);
let zero = Array1::<f64>::zeros(eps.len());
let v = measure_jet_extrapolation_variance(
zero.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::PerLevel(&lams),
FLOOR,
)
.expect("valid inputs");
assert_eq!(
v,
total_ignorance(&lams),
"never-covered query must pay Σ 1/λ̂ exactly"
);
}
#[test]
pub(crate) fn extrapolation_variance_halves_when_amplitudes_double() {
let eps = band();
let lams = lambdas();
let q_bar = support_means(&eps);
let doubled: Vec<f64> = lams.iter().map(|l| 2.0 * l).collect();
let rows = [
support_at_distance(0.35, &eps),
Array1::<f64>::zeros(eps.len()),
Array1::from_elem(eps.len(), 0.5),
];
for row in &rows {
let v1 = measure_jet_extrapolation_variance(
row.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::PerLevel(&lams),
FLOOR,
)
.expect("valid inputs");
let v2 = measure_jet_extrapolation_variance(
row.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::PerLevel(&doubled),
FLOOR,
)
.expect("valid inputs");
assert!(
(2.0 * v2 - v1).abs() <= 1e-15 * v1.max(1.0),
"doubling λ̂ must halve the variance: {v1} vs 2×{v2}"
);
}
}
#[test]
pub(crate) fn extrapolation_variance_gate_convention() {
let eps = band();
let lams = lambdas();
let q_bar = support_means(&eps);
let sub_floor = Array1::from_elem(eps.len(), 0.049 * TOTAL);
let v_sub = measure_jet_extrapolation_variance(
sub_floor.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::PerLevel(&lams),
FLOOR,
)
.expect("valid inputs");
assert_eq!(
v_sub,
total_ignorance(&lams),
"sub-floor mass earns no credit: full Σ 1/λ̂"
);
let mut at_floor = sub_floor.clone();
at_floor[eps.len() - 1] = FLOOR * TOTAL;
let v_floor = measure_jet_extrapolation_variance(
at_floor.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::PerLevel(&lams),
FLOOR,
)
.expect("valid inputs");
let expected: f64 = lams[..eps.len() - 1].iter().map(|l| 1.0 / l).sum::<f64>()
+ (1.0 - FLOOR) / lams[eps.len() - 1];
assert!(
(v_floor - expected).abs() <= 1e-15,
"floor-clearing coarsest level must take weight 1 − a: {v_floor} vs {expected}"
);
assert!(
v_sub - v_floor <= FLOOR * total_ignorance(&lams) + 1e-15,
"gate jump exceeds the documented coverage_floor bound"
);
}
#[test]
pub(crate) fn fused_extrapolation_charges_single_band_amplitude_once() {
let eps = band();
let q_bar = support_means(&eps);
let lam = 2.5;
let zero = Array1::<f64>::zeros(eps.len());
let v_zero = measure_jet_extrapolation_variance(
zero.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::Fused(lam),
FLOOR,
)
.expect("valid inputs");
assert_eq!(
v_zero,
1.0 / lam,
"never-covered fused band must pay one amplitude, not one per level"
);
let covered = arr1(&[0.01, 0.2, 0.4, 0.75, 0.5]);
let v_covered = measure_jet_extrapolation_variance(
covered.view(),
&eps,
&q_bar,
MeasureJetExtrapolationSpectrum::Fused(lam),
FLOOR,
)
.expect("valid inputs");
let expected = (1.0 - 0.75) / lam;
assert!(
(v_covered - expected).abs() <= 1e-15,
"fused band must use the best covered level once: {v_covered} vs {expected}"
);
}
fn eval_fitted(
query: ArrayView1<'_, f64>,
centers: ArrayView2<'_, f64>,
z: ArrayView1<'_, f64>,
length_scale: f64,
head: Option<ArrayView2<'_, f64>>,
head_coeffs: ArrayView1<'_, f64>,
) -> f64 {
let inv_two_l2 = 1.0 / (2.0 * length_scale * length_scale);
let mut val = 0.0_f64;
for i in 0..centers.nrows() {
let mut sq = 0.0_f64;
for a in 0..query.len() {
let dlt = query[a] - centers[(i, a)];
sq += dlt * dlt;
}
val += z[i] * (-sq * inv_two_l2).exp();
}
if let Some(t) = head {
for h in 0..t.ncols() {
let mut proj = t[(0, h)];
for a in 0..query.len() {
proj += query[a] * t[(a + 1, h)];
}
val += head_coeffs[h] * proj;
}
}
val
}
#[test]
pub(crate) fn ambient_gradient_matches_central_difference() {
use ndarray::{arr1, arr2};
let centers = arr2(&[[0.0, 0.0], [1.0, 0.5], [-0.7, 0.9], [0.4, -1.1]]);
let z = arr1(&[0.8, -1.3, 0.5, 2.0]);
let length_scale = 0.6;
let t = arr2(&[[1.0, 0.0, 0.0], [0.0, 1.0, 0.2], [0.0, -0.3, 0.9]]);
let head_coeffs = arr1(&[1.9, 0.7, -0.4]);
let query = arr1(&[0.15, -0.2]);
let grad = measure_jet_ambient_gradient(
query.view(),
centers.view(),
z.view(),
length_scale,
Some(t.view()),
head_coeffs.view(),
)
.expect("valid gradient");
let h = 1e-6;
for a in 0..query.len() {
let mut qp = query.clone();
let mut qm = query.clone();
qp[a] += h;
qm[a] -= h;
let fp = eval_fitted(
qp.view(),
centers.view(),
z.view(),
length_scale,
Some(t.view()),
head_coeffs.view(),
);
let fm = eval_fitted(
qm.view(),
centers.view(),
z.view(),
length_scale,
Some(t.view()),
head_coeffs.view(),
);
let fd = (fp - fm) / (2.0 * h);
assert!(
(grad[a] - fd).abs() <= 1e-6 * (1.0 + fd.abs()),
"axis {a}: analytic {} vs central FD {fd}",
grad[a]
);
}
}
#[test]
pub(crate) fn ambient_gradient_representer_only_and_head_guard() {
use ndarray::{arr1, arr2};
let centers = arr2(&[[0.0], [0.5], [-0.4]]);
let z = arr1(&[1.0, -2.0, 0.5]);
let length_scale = 0.3;
let query = arr1(&[0.1]);
let empty = Array1::<f64>::zeros(0);
let grad = measure_jet_ambient_gradient(
query.view(),
centers.view(),
z.view(),
length_scale,
None,
empty.view(),
)
.expect("valid gradient");
let h = 1e-6;
let mut qp = query.clone();
let mut qm = query.clone();
qp[0] += h;
qm[0] -= h;
let fd = (eval_fitted(
qp.view(),
centers.view(),
z.view(),
length_scale,
None,
empty.view(),
) - eval_fitted(
qm.view(),
centers.view(),
z.view(),
length_scale,
None,
empty.view(),
)) / (2.0 * h);
assert!((grad[0] - fd).abs() <= 1e-6 * (1.0 + fd.abs()));
let stray = arr1(&[1.0]);
assert!(
measure_jet_ambient_gradient(
query.view(),
centers.view(),
z.view(),
length_scale,
None,
stray.view(),
)
.is_err(),
"head coefficients without a head lift must be rejected"
);
}
}