use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
use faer::Side;
use super::BasisError;
use crate::faer_ndarray::FaerEigh;
pub const MEASURE_JET_FRAME_DEFAULT_ORDER_R: usize = 2;
const MEASURE_JET_FRAME_NET_RADIUS_FACTOR: f64 = 0.5;
const MEASURE_JET_FRAME_RANK_RTOL: f64 = 1e-10;
const MEASURE_JET_FRAME_POWER_ITERS: usize = 64;
pub struct MeasureJetFrame {
n_nodes: usize,
head_dim: usize,
level_dims: Vec<usize>,
level_precisions: Vec<f64>,
synthesis: Array2<f64>,
eps_band: Vec<f64>,
order_s: f64,
}
#[derive(Clone, Copy, Debug)]
pub struct MeasureJetFrameRatioCertificate {
pub lower_a: f64,
pub upper_b: f64,
pub ratio: f64,
pub iterations: usize,
}
impl MeasureJetFrame {
pub fn total_dim(&self) -> usize {
self.head_dim + self.level_dims.iter().sum::<usize>()
}
pub fn n_nodes(&self) -> usize {
self.n_nodes
}
pub fn head_dim(&self) -> usize {
self.head_dim
}
pub fn level_dims(&self) -> &[usize] {
&self.level_dims
}
pub fn level_precisions(&self) -> &[f64] {
&self.level_precisions
}
pub fn eps_band(&self) -> &[f64] {
&self.eps_band
}
pub fn order_s(&self) -> f64 {
self.order_s
}
pub fn synthesis_matrix(&self) -> ArrayView2<'_, f64> {
self.synthesis.view()
}
pub fn synthesize(&self, coef: ArrayView1<'_, f64>) -> Result<Array1<f64>, BasisError> {
if coef.len() != self.total_dim() {
crate::bail_dim_basis!(
"measure-jet frame synthesis expects {} coefficients, got {}",
self.total_dim(),
coef.len()
);
}
Ok(self.synthesis.dot(&coef))
}
pub fn analyze(&self, values: ArrayView1<'_, f64>) -> Result<Array1<f64>, BasisError> {
if values.len() != self.n_nodes {
crate::bail_dim_basis!(
"measure-jet frame analysis expects {} nodal values, got {}",
self.n_nodes,
values.len()
);
}
Ok(self.synthesis.t().dot(&values))
}
pub fn whitening_sqrt(&self) -> Array1<f64> {
let mut w = Array1::<f64>::ones(self.total_dim());
let mut off = self.head_dim;
for (&dim, &lambda) in self.level_dims.iter().zip(self.level_precisions.iter()) {
let inv_sqrt = 1.0 / lambda.sqrt();
for j in off..off + dim {
w[j] = inv_sqrt;
}
off += dim;
}
w
}
pub fn frame_ratio_certificate(&self) -> MeasureJetFrameRatioCertificate {
let p = self.total_dim();
let n_innov = p - self.head_dim;
if n_innov == 0 {
return MeasureJetFrameRatioCertificate {
lower_a: 1.0,
upper_b: 1.0,
ratio: 1.0,
iterations: 0,
};
}
let w = self.whitening_sqrt();
let mut s_white = Array2::<f64>::zeros((self.n_nodes, n_innov));
for j in 0..n_innov {
let col = self.head_dim + j;
let scale = w[col];
for i in 0..self.n_nodes {
s_white[(i, j)] = self.synthesis[(i, col)] * scale;
}
}
let m = s_white.t().dot(&s_white);
let top = power_top_eigenvalue(&m, MEASURE_JET_FRAME_POWER_ITERS);
let bottom = power_bottom_eigenvalue(&m, top, MEASURE_JET_FRAME_POWER_ITERS);
let lower_a = bottom.max(0.0);
let upper_b = top.max(lower_a);
let ratio = if lower_a > 0.0 {
upper_b / lower_a
} else {
f64::INFINITY
};
MeasureJetFrameRatioCertificate {
lower_a,
upper_b,
ratio,
iterations: MEASURE_JET_FRAME_POWER_ITERS,
}
}
}
pub fn build_measure_jet_frame(
centers: ArrayView2<'_, f64>,
masses: ArrayView1<'_, f64>,
eps_band: &[f64],
order_s: f64,
order_r: usize,
) -> Result<MeasureJetFrame, BasisError> {
let n = centers.nrows();
let d = centers.ncols();
if n == 0 || d == 0 {
crate::bail_invalid_basis!("measure-jet frame needs nonempty centers");
}
if masses.len() != n {
crate::bail_dim_basis!(
"measure-jet frame needs one mass per center: {} masses, {} centers",
masses.len(),
n
);
}
if eps_band.is_empty() {
crate::bail_invalid_basis!("measure-jet frame needs a nonempty scale band");
}
for (l, pair) in eps_band.windows(2).enumerate() {
if pair[1] <= pair[0] {
crate::bail_invalid_basis!(
"measure-jet frame 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 frame band scales must be finite and positive");
}
if !(order_s.is_finite() && order_s > 0.0) {
crate::bail_invalid_basis!("measure-jet frame order_s must be finite and positive");
}
if order_r == 0 {
crate::bail_invalid_basis!("measure-jet frame order_r must be at least 1");
}
if centers.iter().any(|v| !v.is_finite()) {
crate::bail_invalid_basis!("measure-jet frame centers must be finite");
}
if masses.iter().any(|v| !(v.is_finite() && *v >= 0.0)) {
crate::bail_invalid_basis!("measure-jet frame masses must be finite and nonnegative");
}
let head_exponents = enumerate_monomials_below_degree(d, order_r);
let head_raw = monomial_design(centers, &head_exponents);
let head = orthonormalize_columns(&head_raw, MEASURE_JET_FRAME_RANK_RTOL)?;
let head_dim = head.ncols();
let mut ortho_pool: Vec<Array1<f64>> = (0..head_dim).map(|j| head.column(j).to_owned()).collect();
let mut innovation_cols: Vec<Array1<f64>> = Vec::new();
let mut level_dims: Vec<usize> = Vec::with_capacity(eps_band.len());
let mut level_precisions: Vec<f64> = Vec::with_capacity(eps_band.len());
for (rev_idx, &eps) in eps_band.iter().enumerate().rev() {
let net = build_eps_half_net(centers, eps);
let mut level_count = 0usize;
for &node in net.iter() {
let atom_block = jet_atom_columns(centers, masses, node, eps, d);
for raw in atom_block.into_iter() {
let mut resid = raw;
for basis_col in ortho_pool.iter() {
let proj = basis_col.dot(&resid);
resid.scaled_add(-proj, basis_col);
}
let norm = resid.dot(&resid).sqrt();
if norm > MEASURE_JET_FRAME_RANK_RTOL {
resid /= norm;
ortho_pool.push(resid.clone());
innovation_cols.push(resid);
level_count += 1;
}
}
}
level_dims.push(level_count);
level_precisions.push(eps.powf(2.0 * order_s));
}
level_dims.reverse();
level_precisions.reverse();
let mut total_dim = head_dim;
for &dim in &level_dims {
total_dim += dim;
}
let mut synthesis = Array2::<f64>::zeros((n, total_dim));
for j in 0..head_dim {
let mut col = synthesis.column_mut(j);
col.assign(&head.column(j));
}
let mut coarse_to_fine_dims: Vec<usize> = level_dims.clone();
coarse_to_fine_dims.reverse(); let mut ascending_offsets = vec![head_dim; level_dims.len()];
{
let mut acc = head_dim;
for (l, &dim) in level_dims.iter().enumerate() {
ascending_offsets[l] = acc;
acc += dim;
}
}
{
let mut k = 0usize;
for (c, &dim) in coarse_to_fine_dims.iter().enumerate() {
let ascending_level = level_dims.len() - 1 - c;
let mut write = ascending_offsets[ascending_level];
for _ in 0..dim {
let mut col = synthesis.column_mut(write);
col.assign(&innovation_cols[k]);
k += 1;
write += 1;
}
}
}
Ok(MeasureJetFrame {
n_nodes: n,
head_dim,
level_dims,
level_precisions,
synthesis,
eps_band: eps_band.to_vec(),
order_s,
})
}
fn enumerate_monomials_below_degree(d: usize, order_r: usize) -> Vec<Vec<usize>> {
let mut out: Vec<Vec<usize>> = Vec::new();
let max_total = order_r - 1;
for total in 0..=max_total {
let mut current = vec![0usize; d];
enumerate_fixed_total(d, total, 0, &mut current, &mut out);
}
out
}
fn enumerate_fixed_total(
d: usize,
remaining: usize,
pos: usize,
current: &mut Vec<usize>,
out: &mut Vec<Vec<usize>>,
) {
if pos == d - 1 {
current[pos] = remaining;
out.push(current.clone());
current[pos] = 0;
return;
}
for e in 0..=remaining {
current[pos] = e;
enumerate_fixed_total(d, remaining - e, pos + 1, current, out);
}
current[pos] = 0;
}
fn monomial_design(centers: ArrayView2<'_, f64>, exponents: &[Vec<usize>]) -> Array2<f64> {
let n = centers.nrows();
let mut out = Array2::<f64>::zeros((n, exponents.len()));
for (j, alpha) in exponents.iter().enumerate() {
for i in 0..n {
let mut v = 1.0;
for (k, &e) in alpha.iter().enumerate() {
if e > 0 {
v *= centers[(i, k)].powi(e as i32);
}
}
out[(i, j)] = v;
}
}
out
}
fn build_eps_half_net(centers: ArrayView2<'_, f64>, eps: f64) -> Vec<usize> {
let n = centers.nrows();
let radius = MEASURE_JET_FRAME_NET_RADIUS_FACTOR * eps;
let radius2 = radius * radius;
let mut claimed = vec![false; n];
let mut nodes: Vec<usize> = Vec::new();
for i in 0..n {
if claimed[i] {
continue;
}
nodes.push(i);
let ci = centers.row(i);
for j in i..n {
if claimed[j] {
continue;
}
let cj = centers.row(j);
let mut dd = 0.0;
for k in 0..centers.ncols() {
let diff = ci[k] - cj[k];
dd += diff * diff;
}
if dd <= radius2 {
claimed[j] = true;
}
}
}
nodes
}
fn jet_atom_columns(
centers: ArrayView2<'_, f64>,
masses: ArrayView1<'_, f64>,
node: usize,
eps: f64,
d: usize,
) -> Vec<Array1<f64>> {
let n = centers.nrows();
let inv_two_eps2 = 1.0 / (2.0 * eps * eps);
let inv_eps = 1.0 / eps;
let c = centers.row(node);
let mut bump = Array1::<f64>::zeros(n);
for i in 0..n {
let xi = centers.row(i);
let mut dd = 0.0;
for k in 0..d {
let diff = xi[k] - c[k];
dd += diff * diff;
}
bump[i] = masses[i] * (-dd * inv_two_eps2).exp();
}
let mut cols: Vec<Array1<f64>> = Vec::with_capacity(1 + d);
cols.push(bump.clone());
for k in 0..d {
let mut col = Array1::<f64>::zeros(n);
for i in 0..n {
let frame = (centers[(i, k)] - c[k]) * inv_eps;
col[i] = bump[i] * frame;
}
cols.push(col);
}
cols
}
fn orthonormalize_columns(a: &Array2<f64>, rtol: f64) -> Result<Array2<f64>, BasisError> {
let n = a.nrows();
let p = a.ncols();
if p == 0 {
crate::bail_invalid_basis!("measure-jet frame head must have at least one column");
}
let mut max_norm = 0.0_f64;
for j in 0..p {
let nj = a.column(j).dot(&a.column(j)).sqrt();
max_norm = max_norm.max(nj);
}
let drop_below = (rtol * max_norm).max(f64::MIN_POSITIVE);
let mut basis: Vec<Array1<f64>> = Vec::new();
for j in 0..p {
let mut v = a.column(j).to_owned();
for q in basis.iter() {
let proj = q.dot(&v);
v.scaled_add(-proj, q);
}
let norm = v.dot(&v).sqrt();
if norm > drop_below {
v /= norm;
basis.push(v);
}
}
if basis.is_empty() {
crate::bail_invalid_basis!("measure-jet frame head collapsed to rank zero");
}
let rank = basis.len();
let mut out = Array2::<f64>::zeros((n, rank));
for (j, col) in basis.into_iter().enumerate() {
out.column_mut(j).assign(&col);
}
Ok(out)
}
fn power_top_eigenvalue(m: &Array2<f64>, iters: usize) -> f64 {
let p = m.nrows();
if p == 0 {
return 0.0;
}
let mut v = Array1::<f64>::ones(p);
let mut norm = v.dot(&v).sqrt();
if norm == 0.0 {
return 0.0;
}
v /= norm;
let mut lambda = 0.0;
for _ in 0..iters {
let w = m.dot(&v);
lambda = v.dot(&w);
norm = w.dot(&w).sqrt();
if !(norm > 0.0) {
break;
}
v = w / norm;
}
if let Ok((eigs, _)) = m.eigh(Side::Lower) {
if let Some(max) = eigs.iter().cloned().fold(None, |acc: Option<f64>, e| {
Some(acc.map_or(e, |a| a.max(e)))
}) {
return max.max(lambda).max(0.0);
}
}
lambda.max(0.0)
}
fn power_bottom_eigenvalue(m: &Array2<f64>, top: f64, iters: usize) -> f64 {
let p = m.nrows();
if p == 0 {
return 0.0;
}
if let Ok((eigs, _)) = m.eigh(Side::Lower) {
if let Some(min) = eigs.iter().cloned().fold(None, |acc: Option<f64>, e| {
Some(acc.map_or(e, |a| a.min(e)))
}) {
return min.max(0.0);
}
}
let mut shifted = Array2::<f64>::zeros((p, p));
for i in 0..p {
for j in 0..p {
shifted[(i, j)] = if i == j { top - m[(i, j)] } else { -m[(i, j)] };
}
}
let shifted_top = power_top_eigenvalue(&shifted, iters);
(top - shifted_top).max(0.0)
}
#[cfg(test)]
mod tests {
use super::*;
use ndarray::Array2;
fn grid_centers(side: usize) -> (Array2<f64>, Array1<f64>) {
let n = side * side;
let mut centers = Array2::<f64>::zeros((n, 2));
let step = 1.0 / (side as f64 - 1.0).max(1.0);
let mut idx = 0;
for i in 0..side {
for j in 0..side {
centers[(idx, 0)] = i as f64 * step;
centers[(idx, 1)] = j as f64 * step;
idx += 1;
}
}
let masses = Array1::<f64>::from_elem(n, 1.0 / n as f64);
(centers, masses)
}
fn grid_band() -> Vec<f64> {
vec![0.15, 0.30, 0.60]
}
fn affine_values(centers: ArrayView2<'_, f64>, a: f64, b: f64, c: f64) -> Array1<f64> {
let n = centers.nrows();
let mut v = Array1::<f64>::zeros(n);
for i in 0..n {
v[i] = a + b * centers[(i, 0)] + c * centers[(i, 1)];
}
v
}
#[test]
fn affine_functions_in_exact_null_space() {
let (centers, masses) = grid_centers(7);
let band = grid_band();
let frame = build_measure_jet_frame(
centers.view(),
masses.view(),
&band,
1.5,
MEASURE_JET_FRAME_DEFAULT_ORDER_R,
)
.expect("frame builds");
let f = affine_values(centers.view(), 0.4, -1.3, 2.1);
let rough = f.dot(&f);
let dual = frame.analyze(f.view()).expect("analyze");
let head_dim = frame.head_dim();
let mut innov_energy = 0.0;
for j in head_dim..frame.total_dim() {
innov_energy += dual[j] * dual[j];
}
assert!(
innov_energy <= 1e-10 * rough.max(1.0),
"affine field must leave innovations empty: innov_energy={innov_energy:.3e}, rough={rough:.3e}"
);
let mut head_coef = Array1::<f64>::zeros(frame.total_dim());
for j in 0..head_dim {
head_coef[j] = dual[j];
}
let recon = frame.synthesize(head_coef.view()).expect("synthesize");
let mut resid = 0.0;
for i in 0..centers.nrows() {
let e = recon[i] - f[i];
resid += e * e;
}
assert!(
resid <= 1e-10 * rough.max(1.0),
"head must reproduce the affine field exactly: resid={resid:.3e}"
);
}
#[test]
fn innovations_have_vanishing_moments() {
let (centers, masses) = grid_centers(7);
let band = grid_band();
let frame = build_measure_jet_frame(
centers.view(),
masses.view(),
&band,
1.5,
MEASURE_JET_FRAME_DEFAULT_ORDER_R,
)
.expect("frame builds");
let s = frame.synthesis_matrix();
let head_dim = frame.head_dim();
let affine_basis = [
affine_values(centers.view(), 1.0, 0.0, 0.0),
affine_values(centers.view(), 0.0, 1.0, 0.0),
affine_values(centers.view(), 0.0, 0.0, 1.0),
];
let mut col = head_dim;
for (l, &dim) in frame.level_dims().iter().enumerate() {
for _ in 0..dim {
let column = s.column(col);
for (bidx, basis_field) in affine_basis.iter().enumerate() {
let ip = column.dot(basis_field);
assert!(
ip.abs() <= 1e-9,
"level {l} innovation column {col} couples to affine basis {bidx}: <c,p>={ip:.3e}"
);
}
col += 1;
}
}
}
#[test]
fn synthesis_and_analysis_are_exact_adjoints() {
let (centers, masses) = grid_centers(6);
let band = grid_band();
let frame = build_measure_jet_frame(
centers.view(),
masses.view(),
&band,
1.5,
MEASURE_JET_FRAME_DEFAULT_ORDER_R,
)
.expect("frame builds");
let p = frame.total_dim();
let n = frame.n_nodes();
let x = Array1::<f64>::from_shape_fn(p, |i| ((i % 7) as f64) - 3.0 + 0.25 * i as f64);
let y = Array1::<f64>::from_shape_fn(n, |i| 1.0 - 0.5 * (i as f64) + ((i % 5) as f64));
let sx = frame.synthesize(x.view()).expect("synthesize");
let sty = frame.analyze(y.view()).expect("analyze");
let lhs = sx.dot(&y);
let rhs = x.dot(&sty);
let scale = lhs.abs().max(rhs.abs()).max(1.0);
assert!(
(lhs - rhs).abs() <= 1e-10 * scale,
"S and Sᵀ must be exact adjoints: <Sx,y>={lhs:.6e}, <x,Sᵀy>={rhs:.6e}"
);
}
#[test]
fn frame_ratio_is_finite_and_at_least_one() {
let (centers, masses) = grid_centers(7);
let band = grid_band();
let frame = build_measure_jet_frame(
centers.view(),
masses.view(),
&band,
1.5,
MEASURE_JET_FRAME_DEFAULT_ORDER_R,
)
.expect("frame builds");
let cert = frame.frame_ratio_certificate();
assert!(
cert.lower_a > 0.0,
"lower frame constant A must be positive on a quasi-uniform net: A={:.3e}",
cert.lower_a
);
assert!(
cert.upper_b.is_finite() && cert.upper_b >= cert.lower_a,
"upper frame constant B must be finite and ≥ A: B={:.3e}, A={:.3e}",
cert.upper_b,
cert.lower_a
);
assert!(
cert.ratio.is_finite() && cert.ratio >= 1.0 - 1e-12,
"frame ratio B/A must be finite and ≥ 1: ratio={:.3e}",
cert.ratio
);
}
#[test]
fn affine_null_space_holds_at_default_order() {
let (centers, masses) = grid_centers(8);
let band = grid_band();
let frame = build_measure_jet_frame(
centers.view(),
masses.view(),
&band,
1.5,
MEASURE_JET_FRAME_DEFAULT_ORDER_R,
)
.expect("frame builds");
let f = affine_values(centers.view(), -2.0, 3.5, -0.75);
let rough = f.dot(&f);
let dual = frame.analyze(f.view()).expect("analyze");
let mut innov_energy = 0.0;
for j in frame.head_dim()..frame.total_dim() {
innov_energy += dual[j] * dual[j];
}
assert!(
innov_energy <= 1e-10 * rough.max(1.0),
"default-band affine null space failed: innov_energy={innov_energy:.3e}"
);
}
}