powder.rs

1use mrlycore::Tensor;2use mrlynum::fft::fft2;34pub struct Powder {5    pub slope: f64,6    pub low: f64,7    pub high: f64,8    pub swing: f64,9}1011fn rings(grid: &Tensor, pad: usize, bins: usize) -> Vec<(f64, f64, f64)> {12    let side = grid.shape[0];13    let bytes = grid.bytes();14    let mut re = vec![0.0f64; pad * pad];15    let mut im = vec![0.0f64; pad * pad];16    for row in 0..side {17        for col in 0..side {18            re[row * pad + col] = bytes[row * side + col] as f64;19        }20    }21    fft2(&mut re, &mut im, pad, false);22    let half = (pad / 2) as i64;23    let top = (half as f64 * 2f64.sqrt()).ln();24    let mut sums = vec![0.0f64; bins];25    let mut logs = vec![0.0f64; bins];26    let mut hits = vec![0.0f64; bins];27    for row in 0..pad {28        for col in 0..pad {29            let u = if (row as i64) <= half {30                row as i6431            } else {32                row as i64 - pad as i6433            };34            let v = if (col as i64) <= half {35                col as i6436            } else {37                col as i64 - pad as i6438            };39            let norm = ((u * u + v * v) as f64).sqrt();40            if norm < 1.0 {41                continue;42            }43            let bin = ((norm.ln() / top) * bins as f64) as usize;44            if bin >= bins {45                continue;46            }47            let flat = row * pad + col;48            let power = re[flat] * re[flat] + im[flat] * im[flat];49            sums[bin] += power;50            logs[bin] += power.max(1e-300).ln();51            hits[bin] += 1.0;52        }53    }54    let mut out = Vec::new();55    for bin in 0..bins {56        if hits[bin] == 0.0 {57            continue;58        }59        let centre = ((bin as f64 + 0.5) / bins as f64 * top).exp();60        out.push((centre, sums[bin] / hits[bin], (logs[bin] / hits[bin]).exp()));61    }62    out63}6465fn fit(points: &[(f64, f64)]) -> (f64, f64) {66    let n = points.len() as f64;67    let sx: f64 = points.iter().map(|p| p.0.ln()).sum();68    let sy: f64 = points.iter().map(|p| p.1.ln()).sum();69    let sxx: f64 = points.iter().map(|p| p.0.ln() * p.0.ln()).sum();70    let sxy: f64 = points.iter().map(|p| p.0.ln() * p.1.ln()).sum();71    let slope = (n * sxy - sx * sy) / (n * sxx - sx * sx);72    (slope, (sy - slope * sx) / n)73}7475fn slide(points: &[(f64, f64)], width: f64, step: f64) -> (f64, f64) {76    let first = points.first().map(|p| p.0.ln()).unwrap_or(0.0);77    let last = points.last().map(|p| p.0.ln()).unwrap_or(0.0);78    let mut low = f64::INFINITY;79    let mut high = f64::NEG_INFINITY;80    let mut start = first;81    while start + width <= last + 1e-12 {82        let window: Vec<(f64, f64)> = points83            .iter()84            .copied()85            .filter(|(k, _)| k.ln() >= start && k.ln() <= start + width)86            .collect();87        if window.len() >= 8 {88            let (value, _) = fit(&window);89            low = low.min(value);90            high = high.max(value);91        }92        start += step;93    }94    (low, high)95}9697pub fn powder(98    grid: &Tensor,99    pad: usize,100    low: f64,101    high: f64,102    bins: usize,103    phases: usize,104) -> Powder {105    let all = rings(grid, pad, bins);106    let band: Vec<(f64, f64, f64)> = all107        .iter()108        .copied()109        .filter(|(k, p, g)| *k >= low && *k <= high && *p > 0.0 && *g > 0.0)110        .collect();111    let arithmetic: Vec<(f64, f64)> = band.iter().map(|(k, p, _)| (*k, *p)).collect();112    let (slope, intercept) = fit(&arithmetic);113    let period = (crate::design::BASE as f64).ln();114    let (least, most) = slide(&arithmetic, 3.0 * period, period / 4.0);115    let mut sums = vec![0.0f64; phases];116    let mut hits = vec![0.0f64; phases];117    for (k, power) in &arithmetic {118        let residual = power.ln() - slope * k.ln() - intercept;119        let phase = (k.ln() / period).rem_euclid(1.0);120        let bin = ((phase * phases as f64) as usize).min(phases - 1);121        sums[bin] += residual;122        hits[bin] += 1.0;123    }124    let curve: Vec<f64> = sums125        .iter()126        .zip(&hits)127        .map(|(sum, hit)| if *hit == 0.0 { f64::NAN } else { sum / hit })128        .collect();129    let swing = crate::mass::spread(&curve);130    Powder {131        slope,132        low: least,133        high: most,134        swing,135    }136}