use std::f32::consts::PI;
#[inline]
pub fn wrap_pi(theta: f32) -> f32 {
let mut t = theta % PI;
if t < 0.0 {
t += PI;
}
if t >= PI {
t -= PI;
}
t
}
#[inline]
pub fn angular_dist_pi(a: f32, b: f32) -> f32 {
let diff = ((a - b) % PI + PI) % PI;
diff.min(PI - diff)
}
#[inline]
pub fn angle_to_bin(theta: f32, n: usize) -> usize {
let t = wrap_pi(theta);
let x = t / PI * n as f32;
let mut idx = x.floor() as isize;
if idx < 0 {
idx = 0;
}
if idx as usize >= n {
idx = (n - 1) as isize;
}
idx as usize
}
#[inline]
pub fn bin_to_angle(bin: usize, n: usize) -> f32 {
let step = PI / n as f32;
(bin as f32 + 0.5) * step
}
pub fn smooth_circular_5(hist: &[f32]) -> Vec<f32> {
let n = hist.len();
if n == 0 {
return Vec::new();
}
const K: [f32; 5] = [1.0, 4.0, 6.0, 4.0, 1.0];
const K_SUM: f32 = 16.0;
let mut out = vec![0.0_f32; n];
for (i, bin) in out.iter_mut().enumerate() {
let mut acc = 0.0_f32;
for (k, &w) in K.iter().enumerate() {
let offset = k as isize - 2;
let j = ((i as isize + offset).rem_euclid(n as isize)) as usize;
acc += w * hist[j];
}
*bin = acc / K_SUM;
}
out
}
#[non_exhaustive]
#[derive(Clone, Copy, Debug)]
pub struct PeakPickOptions {
pub min_peak_weight_fraction: f32,
pub min_separation: f32,
}
impl PeakPickOptions {
pub fn new(min_peak_weight_fraction: f32, min_separation: f32) -> Self {
Self {
min_peak_weight_fraction,
min_separation,
}
}
}
pub fn pick_two_peaks(
smoothed: &[f32],
total_weight: f32,
opts: &PeakPickOptions,
) -> Option<(f32, f32)> {
let n = smoothed.len();
if n == 0 {
return None;
}
let min_w = total_weight * opts.min_peak_weight_fraction;
let mut peaks: Vec<(usize, f32)> = Vec::new();
let mut visited = vec![false; n];
for start in 0..n {
if visited[start] {
continue;
}
let here = smoothed[start];
if here < min_w {
visited[start] = true;
continue;
}
let mut len = 1usize;
while len < n {
let next_idx = (start + len) % n;
if smoothed[next_idx] != here {
break;
}
len += 1;
}
for k in 0..len {
visited[(start + k) % n] = true;
}
if len == n {
continue;
}
let left = smoothed[(start + n - 1) % n];
let right = smoothed[(start + len) % n];
if here > left && here > right {
let mid = (start + len / 2) % n;
peaks.push((mid, here));
}
}
peaks.sort_by(|a, b| b.1.total_cmp(&a.1));
if peaks.is_empty() {
return None;
}
let theta_of = |bin: usize| bin_to_angle(bin, n);
let first = theta_of(peaks[0].0);
for (bin, _w) in peaks.iter().skip(1) {
let cand = theta_of(*bin);
if angular_dist_pi(first, cand) >= opts.min_separation {
return Some((first, cand));
}
}
None
}
#[derive(Clone, Copy, Debug)]
pub struct AngleVote {
pub angle: f32,
pub weight: f32,
}
pub fn refine_2means_double_angle(
votes: &[AngleVote],
seed: [f32; 2],
max_iters: usize,
) -> (f32, f32) {
if votes.is_empty() {
return (seed[0], seed[1]);
}
let mut centers = seed;
for _ in 0..max_iters {
let mut sum_2cos = [0.0_f32; 2];
let mut sum_2sin = [0.0_f32; 2];
let mut sum_w = [0.0_f32; 2];
for v in votes {
let d0 = angular_dist_pi(v.angle, centers[0]);
let d1 = angular_dist_pi(v.angle, centers[1]);
let k = if d0 <= d1 { 0 } else { 1 };
let two_theta = 2.0 * v.angle;
sum_2cos[k] += v.weight * two_theta.cos();
sum_2sin[k] += v.weight * two_theta.sin();
sum_w[k] += v.weight;
}
let mut new_centers = centers;
for k in 0..2 {
if sum_w[k] > 0.0 {
let two_theta = sum_2sin[k].atan2(sum_2cos[k]);
new_centers[k] = wrap_pi(two_theta * 0.5);
}
}
if (new_centers[0] - centers[0]).abs() < 1e-5 && (new_centers[1] - centers[1]).abs() < 1e-5
{
return (new_centers[0], new_centers[1]);
}
centers = new_centers;
}
(centers[0], centers[1])
}
#[cfg(test)]
mod tests {
use super::*;
use std::f32::consts::FRAC_PI_2;
#[test]
fn wrap_pi_handles_boundary() {
assert!((wrap_pi(0.0) - 0.0).abs() < 1e-6);
assert!((wrap_pi(PI) - 0.0).abs() < 1e-6);
assert!((wrap_pi(PI + 0.1) - 0.1).abs() < 1e-5);
assert!((wrap_pi(-0.1) - (PI - 0.1)).abs() < 1e-5);
}
#[test]
fn angular_dist_pi_wraps() {
assert!((angular_dist_pi(0.1, PI - 0.1) - 0.2).abs() < 1e-5);
assert!((angular_dist_pi(0.0, FRAC_PI_2) - FRAC_PI_2).abs() < 1e-6);
}
#[test]
fn smooth_5_preserves_total() {
let hist = vec![0.0, 0.0, 16.0, 0.0, 0.0];
let out = smooth_circular_5(&hist);
let sum: f32 = out.iter().sum();
assert!((sum - 16.0).abs() < 1e-4, "got {sum}");
}
#[test]
fn pick_two_peaks_separates_orthogonal_peaks() {
let mut hist = vec![0.0_f32; 18];
hist[0] = 100.0;
hist[9] = 100.0;
let smoothed = smooth_circular_5(&hist);
let peaks = pick_two_peaks(
&smoothed,
200.0,
&PeakPickOptions::new(0.02, 60.0_f32.to_radians()),
)
.expect("two peaks");
let lo = peaks.0.min(peaks.1);
let hi = peaks.0.max(peaks.1);
assert!((lo).abs() < 0.1, "lo too far from 0: {lo}");
assert!((hi - FRAC_PI_2).abs() < 0.1, "hi too far from π/2: {hi}");
}
#[test]
fn pick_two_peaks_handles_plateau_at_boundary() {
let n = 18;
let mut hist = vec![0.0_f32; n];
hist[0] = 50.0;
hist[n - 1] = 50.0;
hist[9] = 100.0;
let smoothed = smooth_circular_5(&hist);
let peaks = pick_two_peaks(
&smoothed,
200.0,
&PeakPickOptions::new(0.02, 60.0_f32.to_radians()),
)
.expect("should recover two peaks despite the boundary plateau");
let lo = peaks.0.min(peaks.1);
let hi = peaks.0.max(peaks.1);
assert!(
lo.abs() < 0.2 || (PI - lo).abs() < 0.2,
"low peak at wrong angle: {lo}"
);
assert!((hi - FRAC_PI_2).abs() < 0.2, "hi at wrong angle: {hi}");
}
#[test]
fn two_means_converges_on_orthogonal_votes() {
let votes: Vec<AngleVote> = (0..50)
.flat_map(|_| {
[
AngleVote {
angle: 0.1,
weight: 1.0,
},
AngleVote {
angle: FRAC_PI_2 - 0.1,
weight: 1.0,
},
]
})
.collect();
let (c0, c1) = refine_2means_double_angle(&votes, [0.2, FRAC_PI_2 - 0.2], 10);
assert!((c0 - 0.1).abs() < 0.05 || (c1 - 0.1).abs() < 0.05);
assert!((c0 - (FRAC_PI_2 - 0.1)).abs() < 0.05 || (c1 - (FRAC_PI_2 - 0.1)).abs() < 0.05);
}
#[test]
fn two_means_empty_returns_seed() {
let seed = [0.1_f32, 1.2];
let (c0, c1) = refine_2means_double_angle(&[], seed, 5);
assert_eq!((c0, c1), (seed[0], seed[1]));
}
}