use super::block::BlockSparseFit;
use crate::dual_certificate::harmonic_dual_birth_eta;
use crate::super_resolution::{recover_spikes, separation_limit};
use ndarray::{Array2, ArrayView2, ArrayView3};
use std::f64::consts::TAU;
#[derive(Clone, Copy, Debug)]
pub struct FiringCoordinate {
pub block: usize,
pub row: usize,
pub t: f64,
pub amplitude: f64,
pub t_se: f64,
pub amplitude_se: f64,
pub t_se_clamped: bool,
}
#[derive(Clone, Debug)]
pub struct BlockCoordinateReport {
pub sigma_hat: f64,
pub mean_radius: f64,
pub n_firings: usize,
pub firings: Vec<FiringCoordinate>,
}
#[derive(Clone, Copy, Debug)]
pub struct MeasureSpikeCoordinate {
pub amplitude: f64,
pub coordinate: f64,
pub coordinate_se: f64,
}
#[derive(Clone, Debug)]
pub struct MeasureValuedCode {
pub block: usize,
pub row: usize,
pub spikes: Vec<MeasureSpikeCoordinate>,
pub dual_eta: f64,
pub used_super_resolution: bool,
}
#[derive(Clone, Debug)]
pub struct BlockMeasureCoordinateReport {
pub sigma_hat: f64,
pub mean_radius: f64,
pub n_firings: usize,
pub firings: Vec<MeasureValuedCode>,
}
fn uniform_phase_sd() -> f64 {
(1.0f64 / 12.0).sqrt()
}
fn collect_route_firings(
blocks: ArrayView2<'_, u32>,
codes: ArrayView3<'_, f32>,
n_blocks: usize,
block: usize,
block_size: usize,
) -> Result<Vec<(usize, Vec<f64>)>, String> {
if block_size == 0 {
return Err("coordinate route block_size must be >= 1".to_string());
}
let (n, width) = blocks.dim();
if codes.shape() != [n, width, block_size] {
return Err(format!(
"coordinate route codes shape {:?} does not match blocks {:?} and block_size {block_size}",
codes.shape(),
blocks.dim()
));
}
if block >= n_blocks {
return Err(format!(
"coordinate route block {block} out of range 0..{n_blocks}"
));
}
let mut out = Vec::new();
for i in 0..n {
let mut found = false;
for j in 0..width {
let routed_block = blocks[[i, j]] as usize;
if routed_block >= n_blocks {
return Err(format!(
"coordinate route block index {routed_block} at row {i}, slot {j} is outside 0..{n_blocks}"
));
}
let mut norm2 = 0.0_f64;
let mut z = Vec::with_capacity(block_size);
for r in 0..block_size {
let value = codes[[i, j, r]] as f64;
if !value.is_finite() {
return Err(format!(
"coordinate route code at row {i}, slot {j}, offset {r} is not finite"
));
}
norm2 += value * value;
z.push(value);
}
if routed_block != block || norm2 == 0.0 {
continue;
}
if found {
return Err(format!(
"coordinate route repeats live block {block} in row {i}"
));
}
out.push((i, z));
found = true;
}
}
Ok(out)
}
fn collect_firings(
fit: &BlockSparseFit,
block: usize,
block_size: usize,
) -> Result<Vec<(usize, Vec<f64>)>, String> {
if block_size == 0 || fit.decoder.nrows() % block_size != 0 {
return Err(format!(
"coordinate fit decoder rows {} not divisible by block_size {block_size}",
fit.decoder.nrows()
));
}
collect_route_firings(
fit.blocks.view(),
fit.codes.view(),
fit.decoder.nrows() / block_size,
block,
block_size,
)
}
fn radius_and_sigma(firings: &[(usize, Vec<f64>)]) -> (f64, f64) {
let n = firings.len();
if n == 0 {
return (0.0, 0.0);
}
let mut norms = Vec::with_capacity(n);
let mut sum = 0.0f64;
for (_, z) in firings {
let nrm = z.iter().map(|v| v * v).sum::<f64>().sqrt();
sum += nrm;
norms.push(nrm);
}
let mean = sum / n as f64;
if n < 2 {
return (mean, 0.0);
}
let ss: f64 = norms.iter().map(|&r| (r - mean) * (r - mean)).sum();
let sigma = (ss / (n - 1) as f64).sqrt();
(mean, sigma)
}
fn phase_se_b2(sigma: f64, norm: f64) -> (f64, bool) {
let ceiling = uniform_phase_sd();
let amp_sq = norm * norm - 2.0 * sigma * sigma;
let raw = if amp_sq > 0.0 {
sigma / (TAU * amp_sq.sqrt())
} else {
f64::INFINITY
};
if raw >= ceiling {
(ceiling, true)
} else {
(raw, false)
}
}
pub(crate) fn phase_coordinate_se(sigma: f64, norm: f64) -> f64 {
phase_se_b2(sigma, norm).0
}
fn coeffs_from_code(z: &[f64]) -> Vec<(f64, f64)> {
z.chunks_exact(2).map(|pair| (pair[0], pair[1])).collect()
}
fn code_from_spikes(spikes: &[MeasureSpikeCoordinate], h_count: usize) -> Vec<f64> {
let mut z = vec![0.0; 2 * h_count];
for spike in spikes {
for h in 1..=h_count {
let phase = TAU * h as f64 * spike.coordinate;
let (s, c) = phase.sin_cos();
z[2 * (h - 1)] += spike.amplitude * c;
z[2 * (h - 1) + 1] += spike.amplitude * s;
}
}
z
}
fn single_harmonic_spike(z: &[f64], sigma: f64) -> (MeasureSpikeCoordinate, Vec<f64>, f64) {
let h_count = z.len() / 2;
let (coordinate, _curvature) = harmonic_argmax(z);
let matched = harmonic_f(z, coordinate);
let amplitude = (matched / h_count.max(1) as f64).max(0.0);
let spike = MeasureSpikeCoordinate {
amplitude,
coordinate,
coordinate_se: spike_coordinate_se(sigma, amplitude, h_count),
};
let fitted = code_from_spikes(&[spike], h_count);
let residual: Vec<f64> = z
.iter()
.zip(fitted.iter())
.map(|(&observed, &pred)| observed - pred)
.collect();
let residual_norm = residual.iter().map(|v| v * v).sum::<f64>().sqrt();
(spike, residual, residual_norm)
}
fn spike_coordinate_se(sigma: f64, amplitude: f64, h_count: usize) -> f64 {
if sigma <= 0.0 {
return 0.0;
}
let ceiling = uniform_phase_sd();
let h_sq_sum = (1..=h_count).map(|h| (h * h) as f64).sum::<f64>();
let slope = amplitude * TAU * h_sq_sum.sqrt();
if slope <= 0.0 {
ceiling
} else {
(sigma / slope).min(ceiling)
}
}
fn circle_dist(a: f64, b: f64) -> f64 {
let d = (a - b).abs();
d.min(1.0 - d)
}
fn separated_from_all(t: f64, accepted: &[MeasureSpikeCoordinate], min_sep: f64) -> bool {
accepted
.iter()
.all(|spike| circle_dist(t, spike.coordinate) + f64::EPSILON >= min_sep)
}
fn count_separated_positive_modes(z: &[f64], min_sep: f64) -> usize {
let stationary = harmonic_stationary_points(z);
let derivative_scale = harmonic_derivative_scale(z);
let sign_tol = f64::EPSILON.sqrt() * derivative_scale;
let mut candidates = Vec::new();
for (idx, &t) in stationary.iter().enumerate() {
let previous = stationary[(idx + stationary.len() - 1) % stationary.len()];
let next = stationary[(idx + 1) % stationary.len()];
let left_span = (t - previous).rem_euclid(1.0);
let right_span = (next - t).rem_euclid(1.0);
let left = (t - 0.5 * left_span).rem_euclid(1.0);
let right = (t + 0.5 * right_span).rem_euclid(1.0);
let val = harmonic_f(z, t);
if val > 0.0 && harmonic_fp(z, left) > sign_tol && harmonic_fp(z, right) < -sign_tol {
candidates.push((t, val));
}
}
candidates.sort_by(|a, b| b.1.total_cmp(&a.1));
let mut accepted: Vec<MeasureSpikeCoordinate> = Vec::new();
for (t, val) in candidates {
if separated_from_all(t, &accepted, min_sep) {
accepted.push(MeasureSpikeCoordinate {
amplitude: val,
coordinate: t,
coordinate_se: 0.0,
});
}
}
accepted.len()
}
fn maybe_super_resolve(z: &[f64], sigma: f64) -> (Vec<MeasureSpikeCoordinate>, f64, bool) {
let h_count = z.len() / 2;
let (single, single_residual, single_residual_norm) = single_harmonic_spike(z, sigma);
if h_count < 2 {
return (vec![single], 0.0, false);
}
let min_sep = separation_limit(h_count);
let single_residual_coeffs = coeffs_from_code(&single_residual);
let eta = harmonic_dual_birth_eta(&single_residual_coeffs, single.amplitude);
let residual_is_multimodal = count_separated_positive_modes(&single_residual, min_sep) > 1;
let code_is_multimodal = count_separated_positive_modes(z, min_sep) > 1;
if eta <= 1.0 && !residual_is_multimodal && !code_is_multimodal {
return (vec![single], eta, false);
}
let coeffs = coeffs_from_code(z);
let recovery = match recover_spikes(&coeffs, sigma) {
Ok(recovery) => recovery,
Err(_err) => return (vec![single], eta, false),
};
if recovery.spikes.len() <= 1 || recovery.residual >= single_residual_norm {
return (vec![single], eta, false);
}
let max_by_separation = if min_sep.is_finite() && min_sep > 0.0 {
(1.0 / min_sep).floor().max(1.0) as usize
} else {
recovery.spikes.len()
};
let mut sorted = recovery.spikes;
sorted.sort_by(|a, b| b.amplitude.total_cmp(&a.amplitude));
let mut accepted = Vec::new();
for spike in sorted {
if spike.amplitude <= 0.0 {
continue;
}
if accepted.len() >= max_by_separation {
break;
}
if !separated_from_all(spike.t, &accepted, min_sep) {
continue;
}
accepted.push(MeasureSpikeCoordinate {
amplitude: spike.amplitude,
coordinate: spike.t,
coordinate_se: spike_coordinate_se(sigma, spike.amplitude, h_count),
});
}
accepted.sort_by(|a, b| a.coordinate.total_cmp(&b.coordinate));
if accepted.len() <= 1 {
(vec![single], eta, false)
} else {
(accepted, eta, true)
}
}
pub fn recover_measure_from_code(
z: &[f64],
sigma: f64,
) -> (Vec<MeasureSpikeCoordinate>, f64, bool) {
maybe_super_resolve(z, sigma)
}
pub fn block_firing_coordinates(
fit: &BlockSparseFit,
block: usize,
) -> Result<BlockCoordinateReport, String> {
let block_size = fit.block_size;
if block_size == 0 || fit.decoder.nrows() % block_size != 0 {
return Err(format!(
"block_firing_coordinates: decoder rows {} not divisible by block_size {block_size}",
fit.decoder.nrows()
));
}
block_route_firing_coordinates(
fit.blocks.view(),
fit.codes.view(),
fit.decoder.nrows() / block_size,
block,
)
}
pub fn block_route_firing_coordinates(
blocks: ArrayView2<'_, u32>,
codes: ArrayView3<'_, f32>,
n_blocks: usize,
block: usize,
) -> Result<BlockCoordinateReport, String> {
let b = codes.shape()[2];
if b != 2 {
return Err(format!(
"block_route_firing_coordinates: circle readout requires block_size b = 2, got b = {b}; \
use harmonic_firing_coordinates for b = 2H"
));
}
let firings = collect_route_firings(blocks, codes, n_blocks, block, b)?;
let (mean_radius, sigma_hat) = radius_and_sigma(&firings);
let mut coords = Vec::with_capacity(firings.len());
for (row, z) in &firings {
let norm = (z[0] * z[0] + z[1] * z[1]).sqrt();
let t = (z[1].atan2(z[0]) / TAU).rem_euclid(1.0);
let (t_se, t_se_clamped) = phase_se_b2(sigma_hat, norm);
coords.push(FiringCoordinate {
block,
row: *row,
t,
amplitude: norm,
t_se,
amplitude_se: sigma_hat,
t_se_clamped,
});
}
Ok(BlockCoordinateReport {
sigma_hat,
mean_radius,
n_firings: firings.len(),
firings: coords,
})
}
fn harmonic_f(rho: &[f64], t: f64) -> f64 {
let h_count = rho.len() / 2;
let mut acc = 0.0;
for h in 0..h_count {
let w = TAU * (h + 1) as f64;
let (s, c) = (w * t).sin_cos();
acc += rho[2 * h] * c + rho[2 * h + 1] * s;
}
acc
}
fn harmonic_fp(rho: &[f64], t: f64) -> f64 {
let h_count = rho.len() / 2;
let mut acc = 0.0;
for h in 0..h_count {
let w = TAU * (h + 1) as f64;
let (s, c) = (w * t).sin_cos();
acc += w * (-rho[2 * h] * s + rho[2 * h + 1] * c);
}
acc
}
fn harmonic_fpp(rho: &[f64], t: f64) -> f64 {
let h_count = rho.len() / 2;
let mut acc = 0.0;
for h in 0..h_count {
let w = TAU * (h + 1) as f64;
let (s, c) = (w * t).sin_cos();
acc += w * w * (-rho[2 * h] * c - rho[2 * h + 1] * s);
}
acc
}
fn harmonic_derivative_scale(rho: &[f64]) -> f64 {
rho.chunks_exact(2)
.enumerate()
.map(|(h, pair)| TAU * (h + 1) as f64 * (pair[0].abs() + pair[1].abs()))
.sum()
}
fn stationary_eliminant(rho: &[f64]) -> Vec<f64> {
let h_count = rho.len() / 2;
let rho_scale = rho
.iter()
.fold(0.0_f64, |scale, value| scale.max(value.abs()));
if h_count == 0 || rho_scale == 0.0 || !rho_scale.is_finite() {
return vec![0.0];
}
let mut alpha = vec![0.0; h_count + 1];
let mut beta = vec![0.0; h_count + 1];
for h in 1..=h_count {
alpha[h] = h as f64 * (rho[2 * (h - 1)] / rho_scale);
beta[h] = h as f64 * (rho[2 * (h - 1) + 1] / rho_scale);
}
let mut coefficients = vec![0.0; 2 * h_count + 1];
for h in 1..=h_count {
for k in 1..=h_count {
let aa = alpha[h] * alpha[k];
let bb = beta[h] * beta[k];
coefficients[h + k] += 0.5 * (bb + aa);
coefficients[h.abs_diff(k)] += 0.5 * (bb - aa);
}
}
coefficients
}
fn normalize_chebyshev(mut coefficients: Vec<f64>) -> Vec<f64> {
let scale = coefficients
.iter()
.fold(0.0_f64, |largest, value| largest.max(value.abs()));
if scale == 0.0 || !scale.is_finite() {
return vec![0.0];
}
let rounding_floor = f64::EPSILON * (coefficients.len() * coefficients.len()) as f64 * scale;
while coefficients.len() > 1
&& coefficients
.last()
.is_some_and(|value| value.abs() <= rounding_floor)
{
coefficients.pop();
}
for value in &mut coefficients {
*value /= scale;
}
coefficients
}
fn evaluate_chebyshev(coefficients: &[f64], x: f64) -> f64 {
let mut next = 0.0;
let mut next_next = 0.0;
for &coefficient in coefficients.iter().skip(1).rev() {
let current = coefficient + 2.0 * x * next - next_next;
next_next = next;
next = current;
}
coefficients[0] + x * next - next_next
}
fn differentiate_chebyshev(coefficients: &[f64]) -> Vec<f64> {
let degree = coefficients.len().saturating_sub(1);
if degree == 0 {
return vec![0.0];
}
let mut derivative = vec![0.0; degree];
derivative[degree - 1] = 2.0 * degree as f64 * coefficients[degree];
if degree >= 2 {
derivative[degree - 2] = 2.0 * (degree - 1) as f64 * coefficients[degree - 1];
for k in (0..degree - 2).rev() {
derivative[k] = derivative[k + 2] + 2.0 * (k + 1) as f64 * coefficients[k + 1];
}
}
derivative[0] *= 0.5;
normalize_chebyshev(derivative)
}
fn chebyshev_zero_tolerance(coefficients: &[f64]) -> f64 {
f64::EPSILON.sqrt() * coefficients.iter().map(|value| value.abs()).sum::<f64>()
}
fn push_distinct_root(roots: &mut Vec<f64>, root: f64) {
let merge_tol = f64::EPSILON.sqrt();
if !roots
.iter()
.any(|existing| (existing - root).abs() <= merge_tol)
{
roots.push(root.clamp(-1.0, 1.0));
}
}
fn bisect_chebyshev_root(coefficients: &[f64], mut left: f64, mut right: f64) -> f64 {
let mut left_value = evaluate_chebyshev(coefficients, left);
loop {
let middle = left + 0.5 * (right - left);
if middle == left || middle == right {
break;
}
let middle_value = evaluate_chebyshev(coefficients, middle);
if middle_value == 0.0 {
return middle;
}
if left_value.is_sign_negative() != middle_value.is_sign_negative() {
right = middle;
} else {
left = middle;
left_value = middle_value;
}
}
if evaluate_chebyshev(coefficients, left).abs() <= evaluate_chebyshev(coefficients, right).abs()
{
left
} else {
right
}
}
fn chebyshev_roots_unit_interval(coefficients: Vec<f64>) -> Vec<f64> {
let coefficients = normalize_chebyshev(coefficients);
let degree = coefficients.len() - 1;
if degree == 0 {
return Vec::new();
}
if degree == 1 {
let root = -coefficients[0] / coefficients[1];
return if (-1.0..=1.0).contains(&root) {
vec![root]
} else {
Vec::new()
};
}
let critical = chebyshev_roots_unit_interval(differentiate_chebyshev(&coefficients));
let tolerance = chebyshev_zero_tolerance(&coefficients);
let mut roots = Vec::new();
for &candidate in critical.iter().chain([-1.0, 1.0].iter()) {
if evaluate_chebyshev(&coefficients, candidate).abs() <= tolerance {
push_distinct_root(&mut roots, candidate);
}
}
let mut boundaries = Vec::with_capacity(critical.len() + 2);
boundaries.push(-1.0);
boundaries.extend(critical.iter().copied());
boundaries.push(1.0);
boundaries.sort_by(f64::total_cmp);
for interval in boundaries.windows(2) {
let left = interval[0];
let right = interval[1];
let left_value = evaluate_chebyshev(&coefficients, left);
let right_value = evaluate_chebyshev(&coefficients, right);
if left_value.abs() > tolerance
&& right_value.abs() > tolerance
&& left_value.is_sign_negative() != right_value.is_sign_negative()
{
push_distinct_root(
&mut roots,
bisect_chebyshev_root(&coefficients, left, right),
);
}
}
roots.sort_by(f64::total_cmp);
roots
}
fn push_distinct_phase(phases: &mut Vec<f64>, phase: f64) {
let phase = phase.rem_euclid(1.0);
let merge_tol = f64::EPSILON.sqrt();
if !phases
.iter()
.any(|existing| circle_dist(*existing, phase) <= merge_tol)
{
phases.push(phase);
}
}
fn harmonic_stationary_points(rho: &[f64]) -> Vec<f64> {
let derivative_scale = harmonic_derivative_scale(rho);
if derivative_scale == 0.0 || !derivative_scale.is_finite() {
return Vec::new();
}
let residual_tol = f64::EPSILON.sqrt() * derivative_scale;
let projected = chebyshev_roots_unit_interval(stationary_eliminant(rho));
let mut phases = Vec::with_capacity(2 * projected.len());
for x in projected {
let theta = x.clamp(-1.0, 1.0).acos();
for lifted in [theta / TAU, (-theta) / TAU] {
let phase = lifted.rem_euclid(1.0);
if harmonic_fp(rho, phase).abs() <= residual_tol {
push_distinct_phase(&mut phases, phase);
}
}
}
phases.sort_by(f64::total_cmp);
phases
}
fn harmonic_argmax(rho: &[f64]) -> (f64, f64) {
let stationary = harmonic_stationary_points(rho);
let Some(&first) = stationary.first() else {
return (0.0, harmonic_fpp(rho, 0.0));
};
let mut best_t = first;
let mut best_value = harmonic_f(rho, first);
for &candidate in stationary.iter().skip(1) {
let value = harmonic_f(rho, candidate);
if value > best_value || (value == best_value && candidate < best_t) {
best_t = candidate;
best_value = value;
}
}
(best_t, harmonic_fpp(rho, best_t))
}
pub fn harmonic_firing_coordinates(
fit: &BlockSparseFit,
block: usize,
) -> Result<BlockCoordinateReport, String> {
let block_size = fit.block_size;
if block_size == 0 || fit.decoder.nrows() % block_size != 0 {
return Err(format!(
"harmonic_firing_coordinates: decoder rows {} not divisible by block_size {block_size}",
fit.decoder.nrows()
));
}
harmonic_route_firing_coordinates(
fit.blocks.view(),
fit.codes.view(),
fit.decoder.nrows() / block_size,
block,
)
}
pub fn harmonic_route_firing_coordinates(
blocks: ArrayView2<'_, u32>,
codes: ArrayView3<'_, f32>,
n_blocks: usize,
block: usize,
) -> Result<BlockCoordinateReport, String> {
let b = codes.shape()[2];
if b < 2 || b % 2 != 0 {
return Err(format!(
"harmonic_route_firing_coordinates: harmonic readout requires block_size b = 2H (even, \
≥ 2), got b = {b}"
));
}
let h_count = b / 2;
let omega_sq_sum: f64 = (1..=h_count).map(|h| TAU * h as f64).map(|w| w * w).sum();
let firings = collect_route_firings(blocks, codes, n_blocks, block, b)?;
let (mean_radius, sigma_hat) = radius_and_sigma(&firings);
let ceiling = uniform_phase_sd();
let mut coords = Vec::with_capacity(firings.len());
for (row, z) in &firings {
let norm = z.iter().map(|v| v * v).sum::<f64>().sqrt();
let (t_hat, fpp) = harmonic_argmax(z);
let raw = if fpp < 0.0 {
(sigma_hat * sigma_hat * omega_sq_sum).sqrt() / (-fpp)
} else {
f64::INFINITY
};
let (t_se, t_se_clamped) = if raw >= ceiling {
(ceiling, true)
} else {
(raw, false)
};
coords.push(FiringCoordinate {
block,
row: *row,
t: t_hat,
amplitude: norm,
t_se,
amplitude_se: sigma_hat,
t_se_clamped,
});
}
Ok(BlockCoordinateReport {
sigma_hat,
mean_radius,
n_firings: firings.len(),
firings: coords,
})
}
pub fn harmonic_measure_coordinates(
fit: &BlockSparseFit,
block: usize,
) -> Result<BlockMeasureCoordinateReport, String> {
let b = fit.block_size;
if b < 2 || b % 2 != 0 {
return Err(format!(
"harmonic_measure_coordinates: harmonic readout requires block_size b = 2H (even, \
>= 2), got b = {b}"
));
}
let g_total = fit.decoder.nrows() / b;
if block >= g_total {
return Err(format!(
"harmonic_measure_coordinates: block {block} out of range 0..{g_total}"
));
}
let firings = collect_firings(fit, block, b)?;
let (mean_radius, sigma_hat) = radius_and_sigma(&firings);
let mut measures = Vec::with_capacity(firings.len());
for (row, z) in &firings {
let (spikes, dual_eta, used_super_resolution) = maybe_super_resolve(z, sigma_hat);
measures.push(MeasureValuedCode {
block,
row: *row,
spikes,
dual_eta,
used_super_resolution,
});
}
Ok(BlockMeasureCoordinateReport {
sigma_hat,
mean_radius,
n_firings: firings.len(),
firings: measures,
})
}
pub fn block_measure_valued_codes(fit: &BlockSparseFit) -> Result<Vec<MeasureValuedCode>, String> {
let b = fit.block_size;
if b < 2 || b % 2 != 0 {
return Err(format!(
"block_measure_valued_codes: harmonic readout requires block_size b = 2H (even, \
>= 2), got b = {b}"
));
}
let g_total = fit.decoder.nrows() / b;
let mut all = Vec::new();
for block in 0..g_total {
let mut report = harmonic_measure_coordinates(fit, block)?;
all.append(&mut report.firings);
}
all.sort_by(|a, b| a.row.cmp(&b.row).then(a.block.cmp(&b.block)));
Ok(all)
}
pub fn reconstruct_measure_valued_rows(
decoder: ArrayView2<'_, f32>,
measures: &[MeasureValuedCode],
n_rows: usize,
block_size: usize,
) -> Result<Array2<f32>, String> {
let b = block_size;
if b < 2 || b % 2 != 0 {
return Err(format!(
"reconstruct_measure_valued_rows: block_size must be even and >= 2, got {b}"
));
}
if decoder.nrows() % b != 0 {
return Err(format!(
"reconstruct_measure_valued_rows: decoder rows {} not divisible by block_size {b}",
decoder.nrows()
));
}
let g_total = decoder.nrows() / b;
let h_count = b / 2;
let p = decoder.ncols();
let mut out = Array2::<f32>::zeros((n_rows, p));
for measure in measures {
if measure.row >= n_rows {
return Err(format!(
"reconstruct_measure_valued_rows: row {} out of range 0..{n_rows}",
measure.row
));
}
if measure.block >= g_total {
return Err(format!(
"reconstruct_measure_valued_rows: block {} out of range 0..{g_total}",
measure.block
));
}
let code = code_from_spikes(&measure.spikes, h_count);
for r in 0..b {
let coeff = code[r] as f32;
if coeff == 0.0 {
continue;
}
let atom = decoder.row(measure.block * b + r);
for c in 0..p {
out[[measure.row, c]] += coeff * atom[c];
}
}
}
Ok(out)
}
pub fn reconstruct_single_coordinate_rows(fit: &BlockSparseFit) -> Result<Array2<f32>, String> {
let b = fit.block_size;
if b < 2 || b % 2 != 0 {
return Err(format!(
"reconstruct_single_coordinate_rows: harmonic readout requires block_size b = 2H \
(even, >= 2), got b = {b}"
));
}
let mut measures = Vec::new();
for block in 0..fit.decoder.nrows() / b {
for (row, z) in collect_firings(fit, block, b)? {
let (spike, residual, residual_norm) = single_harmonic_spike(&z, 0.0);
assert_eq!(residual.len(), b);
assert!(residual_norm.is_finite());
measures.push(MeasureValuedCode {
block,
row,
spikes: vec![spike],
dual_eta: harmonic_dual_birth_eta(&coeffs_from_code(&residual), spike.amplitude),
used_super_resolution: false,
});
}
}
reconstruct_measure_valued_rows(fit.decoder.view(), &measures, fit.blocks.nrows(), b)
}
pub fn explained_variance_from_reconstruction(
x: ArrayView2<'_, f32>,
reconstruction: ArrayView2<'_, f32>,
) -> Result<f64, String> {
if x.dim() != reconstruction.dim() {
return Err(format!(
"explained_variance_from_reconstruction: X shape {:?} != reconstruction shape {:?}",
x.dim(),
reconstruction.dim()
));
}
let (n, p) = x.dim();
let mut means = vec![0.0; p];
for i in 0..n {
for c in 0..p {
means[c] += x[[i, c]] as f64;
}
}
for mean in &mut means {
*mean /= n.max(1) as f64;
}
let mut rss = 0.0;
let mut tss = 0.0;
for i in 0..n {
for c in 0..p {
let r = x[[i, c]] as f64 - reconstruction[[i, c]] as f64;
rss += r * r;
let centered = x[[i, c]] as f64 - means[c];
tss += centered * centered;
}
}
if tss <= f64::MIN_POSITIVE {
Ok(if rss <= f64::MIN_POSITIVE { 1.0 } else { 0.0 })
} else {
Ok(1.0 - rss / tss)
}
}
#[cfg(test)]
mod tests {
use super::*;
use ndarray::{Array2, Array3};
fn u01(state: &mut u64) -> f64 {
*state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
(*state >> 11) as f64 / (1u64 << 53) as f64
}
fn gauss(state: &mut u64) -> f64 {
let u1 = u01(state).max(f64::MIN_POSITIVE);
let u2 = u01(state);
(-2.0 * u1.ln()).sqrt() * (TAU * u2).cos()
}
fn circ_err(a: f64, b: f64) -> f64 {
let d = (a - b).rem_euclid(1.0);
d.min(1.0 - d)
}
#[test]
fn stationary_root_argmax_beats_the_old_lattice_basin() {
let rho = [
-0.2203432413122001,
1.7750647300698044,
1.4778082357960907,
0.4751086996188798,
];
let expected = 0.07077957313295;
let (phase, curvature) = harmonic_argmax(&rho);
assert!(
circ_err(phase, expected) <= f64::EPSILON.sqrt(),
"stationary-root global phase {phase} missed {expected}"
);
assert!(curvature < 0.0);
assert!(
harmonic_f(&rho, phase) > 1.86,
"selected local rather than global maximum"
);
}
#[test]
fn stationary_roots_include_tangencies_and_circle_seam() {
let tangent = [0.0, 1.0, 0.0, -0.5];
let roots = harmonic_stationary_points(&tangent);
for expected in [0.0, 1.0 / 3.0, 2.0 / 3.0] {
assert!(
roots
.iter()
.any(|&root| circ_err(root, expected) <= f64::EPSILON.sqrt()),
"missing stationary root {expected}; got {roots:?}"
);
}
let seam = [-1.0, 0.0];
let (phase, curvature) = harmonic_argmax(&seam);
assert!(circ_err(phase, 0.5) <= f64::EPSILON.sqrt());
assert!(curvature < 0.0);
let zero = [0.0, 0.0, 0.0, 0.0];
assert_eq!(harmonic_argmax(&zero), (0.0, 0.0));
}
#[test]
fn exact_mode_positions_prevent_grid_snapping_separation_error() {
let rho = [
-0.19136859058260647,
0.5744053053825253,
-0.8795719735057727,
-0.16915833401458946,
0.1636929927729193,
-0.03509594833820036,
0.05717764914618254,
-0.1065206988264421,
];
assert_eq!(count_separated_positive_modes(&rho, separation_limit(4)), 1);
}
fn fit_from_codes(codes_rows: &[Vec<f32>], b: usize) -> BlockSparseFit {
let n = codes_rows.len();
let mut blocks = Array2::<u32>::zeros((n, 1));
let mut gates = Array2::<f32>::zeros((n, 1));
let mut codes = Array3::<f32>::zeros((n, 1, b));
for (i, z) in codes_rows.iter().enumerate() {
blocks[[i, 0]] = 0;
let mut nrm = 0.0f32;
for r in 0..b {
codes[[i, 0, r]] = z[r];
nrm += z[r] * z[r];
}
gates[[i, 0]] = nrm.sqrt();
}
let mut decoder = Array2::<f32>::zeros((b, b));
for i in 0..b {
decoder[[i, i]] = 1.0;
}
BlockSparseFit {
decoder,
blocks,
gates,
codes,
gamma: 1.0,
block_utilization: vec![1.0],
block_stable_rank: vec![1.0],
matryoshka_prefix_losses: Vec::new(),
explained_variance: 1.0,
epochs: 1,
convergence: crate::sparse_dict::BlockSparseConvergence::trivially_converged(),
block_topk: 1,
block_size: b,
}
}
fn plant_circle(
n: usize,
sigma: f64,
amp_lo: f64,
amp_hi: f64,
seed: u64,
) -> (Vec<Vec<f32>>, Vec<f64>, Vec<f64>) {
let mut s = seed;
let mut rows = Vec::with_capacity(n);
let mut phase = Vec::with_capacity(n);
let mut amp = Vec::with_capacity(n);
for _ in 0..n {
let t = u01(&mut s);
let a = amp_lo + (amp_hi - amp_lo) * u01(&mut s);
let z0 = a * (TAU * t).cos() + sigma * gauss(&mut s);
let z1 = a * (TAU * t).sin() + sigma * gauss(&mut s);
rows.push(vec![z0 as f32, z1 as f32]);
phase.push(t);
amp.push(a);
}
(rows, phase, amp)
}
fn median(mut v: Vec<f64>) -> f64 {
v.sort_by(|a, b| a.partial_cmp(b).unwrap());
let n = v.len();
if n == 0 {
return 0.0;
}
if n % 2 == 1 {
v[n / 2]
} else {
0.5 * (v[n / 2 - 1] + v[n / 2])
}
}
#[test]
fn circle_phase_se_is_calibrated_and_covers() {
let sigma = 0.05;
let n = 4000;
let (rows, phase, _amp) = plant_circle(n, sigma, 1.0, 1.0, 0x9E3779B97F4A7C15);
let fit = fit_from_codes(&rows, 2);
let report = block_firing_coordinates(&fit, 0).expect("circle readout");
assert_eq!(report.n_firings, n);
let sig_err = (report.sigma_hat - sigma).abs() / sigma;
assert!(
sig_err < 0.30,
"σ̂ = {} should recover planted σ = {sigma} within 30% (rel err {sig_err})",
report.sigma_hat
);
let mut errs = Vec::with_capacity(n);
let mut ses = Vec::with_capacity(n);
let mut covered = 0usize;
for fc in &report.firings {
let e = circ_err(fc.t, phase[fc.row]);
errs.push(e);
ses.push(fc.t_se);
if e < 2.0 * fc.t_se {
covered += 1;
}
}
let normal_mad = 0.6744897501960817; let predicted = normal_mad * median(ses.clone());
let measured = median(errs);
let ratio = measured / predicted;
assert!(
(1.0 / 1.5..=1.5).contains(&ratio),
"median error {measured} vs predicted median SE {predicted} (ratio {ratio}) \
must agree within ~1.5×"
);
let coverage = covered as f64 / n as f64;
assert!(
coverage >= 0.90,
"|t̂−t| < 2·t_se coverage {coverage} must be ≥ 0.90"
);
}
#[test]
fn amplitude_readout_and_conservative_sigma() {
let sigma = 0.05;
let n = 4000;
let (rows, _phase, amp) = plant_circle(n, sigma, 0.5, 2.0, 0x243F6A8885A308D3);
let fit = fit_from_codes(&rows, 2);
let report = block_firing_coordinates(&fit, 0).expect("circle readout");
let mut amp_err = Vec::with_capacity(n);
for fc in &report.firings {
amp_err.push((fc.amplitude - amp[fc.row]).abs());
}
let med_amp_err = median(amp_err);
assert!(
med_amp_err < 3.0 * sigma,
"median |â − a| = {med_amp_err} should track the planted amplitude (~σ = {sigma})"
);
assert!(
report.sigma_hat > sigma,
"radial-scatter σ̂ = {} must be conservative (> true σ = {sigma}) under \
amplitude variation",
report.sigma_hat
);
}
#[test]
fn weak_signal_phase_se_clamps_to_uniform_never_nan() {
let mut rows = vec![vec![1.0f32, 0.0]; 8];
rows.push(vec![1.0e-6f32, 1.0e-6]);
let fit = fit_from_codes(&rows, 2);
let report = block_firing_coordinates(&fit, 0).expect("readout");
let weak = report.firings.last().unwrap();
assert!(weak.t_se.is_finite(), "weak-signal t_se must be finite");
assert!(weak.t_se_clamped, "weak-signal firing must flag the clamp");
assert!(
(weak.t_se - uniform_phase_sd()).abs() < 1.0e-12,
"clamped t_se must equal the uniform-phase SD"
);
assert!(!report.firings[0].t_se_clamped);
}
#[test]
fn harmonic_h2_recovers_phase() {
let sigma = 0.05;
let n = 2000;
let mut s = 0xD1B54A32D192ED03u64;
let mut rows = Vec::with_capacity(n);
let mut truth = Vec::with_capacity(n);
for _ in 0..n {
let t = u01(&mut s);
let a = 0.5 + 1.5 * u01(&mut s);
let mut z = [0.0f64; 4];
for h in 0..2 {
let w = TAU * (h + 1) as f64;
z[2 * h] = a * (w * t).cos();
z[2 * h + 1] = a * (w * t).sin();
}
for zc in z.iter_mut() {
*zc += sigma * gauss(&mut s);
}
rows.push(z.iter().map(|&v| v as f32).collect::<Vec<_>>());
truth.push(t);
}
let fit = fit_from_codes(&rows, 4);
let report = harmonic_firing_coordinates(&fit, 0).expect("harmonic readout");
assert_eq!(report.n_firings, n);
let mut errs = Vec::with_capacity(n);
let mut covered = 0usize;
for fc in &report.firings {
let e = circ_err(fc.t, truth[fc.row]);
errs.push(e);
if e < 2.0 * fc.t_se {
covered += 1;
}
}
let med = median(errs);
assert!(
med < 0.02,
"H=2 harmonic phase median error {med} must be well under 0.02"
);
let coverage = covered as f64 / n as f64;
assert!(
coverage >= 0.85,
"harmonic 2·SE coverage {coverage} must be ≥ 0.85"
);
}
fn harmonic_code(spikes: &[(f64, f64)], h_count: usize) -> Vec<f32> {
let measure_spikes: Vec<MeasureSpikeCoordinate> = spikes
.iter()
.map(|&(coordinate, amplitude)| MeasureSpikeCoordinate {
amplitude,
coordinate,
coordinate_se: 0.0,
})
.collect();
code_from_spikes(&measure_spikes, h_count)
.into_iter()
.map(|v| v as f32)
.collect()
}
fn row_matrix(rows: &[Vec<f32>]) -> Array2<f32> {
let n = rows.len();
let p = rows[0].len();
let mut out = Array2::<f32>::zeros((n, p));
for i in 0..n {
for c in 0..p {
out[[i, c]] = rows[i][c];
}
}
out
}
#[test]
fn measure_readout_recovers_two_separated_spikes_on_one_circle() {
let h_count = 8;
let t1 = 0.20;
let t2 = 0.55;
assert!(
circ_err(t1, t2) > separation_limit(h_count),
"planted spikes must be beyond the super-resolution separation guarantee"
);
let rows = vec![
harmonic_code(&[(t1, 1.0), (t2, 0.7)], h_count),
harmonic_code(&[(0.05, 0.9)], h_count),
];
let fit = fit_from_codes(&rows, 2 * h_count);
let single = harmonic_firing_coordinates(&fit, 0).expect("single coordinate readout");
assert_eq!(
single.firings.iter().filter(|f| f.row == 0).count(),
1,
"legacy harmonic readout emits exactly one coordinate for the two-spike row"
);
let report = harmonic_measure_coordinates(&fit, 0).expect("measure readout");
let row0 = report
.firings
.iter()
.find(|firing| firing.row == 0)
.expect("row 0 firing");
assert!(
row0.used_super_resolution,
"two separated modes should invoke matrix-pencil recovery"
);
assert_eq!(
row0.spikes.len(),
2,
"measure readout must recover both spikes"
);
assert!(
row0.spikes
.iter()
.any(|spike| circ_err(spike.coordinate, t1) < 1.0e-8),
"missing spike at t1={t1}"
);
assert!(
row0.spikes
.iter()
.any(|spike| circ_err(spike.coordinate, t2) < 1.0e-8),
"missing spike at t2={t2}"
);
}
#[test]
fn measure_readout_keeps_single_spike_single() {
let h_count = 8;
let t = 0.37;
let rows = vec![
harmonic_code(&[(t, 1.25)], h_count),
harmonic_code(&[(0.62, 0.8)], h_count),
];
let fit = fit_from_codes(&rows, 2 * h_count);
let report = harmonic_measure_coordinates(&fit, 0).expect("measure readout");
for firing in &report.firings {
assert_eq!(
firing.spikes.len(),
1,
"single-spike harmonic code must not receive spurious multiplicity"
);
assert!(
!firing.used_super_resolution,
"exact one-spike code should stay on the single-coordinate path"
);
}
}
#[test]
fn measure_reconstruction_improves_two_spike_ev_and_preserves_one_spike_ev() {
let h_count = 8;
let rows = vec![
harmonic_code(&[(0.20, 1.0), (0.55, 0.7)], h_count),
harmonic_code(&[(0.05, 0.9)], h_count),
];
let x = row_matrix(&rows);
let fit = fit_from_codes(&rows, 2 * h_count);
let single_recon = reconstruct_single_coordinate_rows(&fit).expect("single recon");
let measures = block_measure_valued_codes(&fit).expect("measure codes");
let measure_recon =
reconstruct_measure_valued_rows(fit.decoder.view(), &measures, rows.len(), 2 * h_count)
.expect("measure recon");
let single_ev =
explained_variance_from_reconstruction(x.view(), single_recon.view()).expect("ev");
let measure_ev =
explained_variance_from_reconstruction(x.view(), measure_recon.view()).expect("ev");
assert!(
measure_ev > single_ev,
"measure reconstruction EV {measure_ev} must improve over single-coordinate EV {single_ev}"
);
assert!(
(1.0 - measure_ev).abs() < 1.0e-10,
"exact two-spike measure reconstruction should recover the harmonic code"
);
let one_rows = vec![
harmonic_code(&[(0.20, 1.0)], h_count),
harmonic_code(&[(0.55, 0.7)], h_count),
];
let one_x = row_matrix(&one_rows);
let one_fit = fit_from_codes(&one_rows, 2 * h_count);
let one_single = reconstruct_single_coordinate_rows(&one_fit).expect("single recon");
let one_measures = block_measure_valued_codes(&one_fit).expect("measure codes");
let one_measure = reconstruct_measure_valued_rows(
one_fit.decoder.view(),
&one_measures,
one_rows.len(),
2 * h_count,
)
.expect("measure recon");
let one_single_ev =
explained_variance_from_reconstruction(one_x.view(), one_single.view()).expect("ev");
let one_measure_ev =
explained_variance_from_reconstruction(one_x.view(), one_measure.view()).expect("ev");
assert!(
(one_measure_ev - one_single_ev).abs() < 1.0e-10,
"one-spike EV must be unchanged: measure {one_measure_ev}, single {one_single_ev}"
);
}
}