feffit 0.1.0

//! Parity tests for `feffit` against larch's `FeffitDataSet._residual`, via the
//! labeled-block references in `tests/data/` (generated by `scripts/ref_feffit.py`,
//! which drives the real `larch.xafs.feffit`).
//!
//! Tolerances absorb FFT round-off (rustfft vs scipy.fftpack) and the cubic
//! spline difference (~5e-14) but catch any real divergence.

use std::collections::HashMap;
use std::path::PathBuf;

use feffit::feffdat::FeffPath;
use feffit::xafsft::Window;
use feffit::{DataSet, FitSpace, Transform};

fn data_dir() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/data")
}

/// A parsed reference: scalar params, path specs, and named float blocks.
struct Ref {
    params: HashMap<String, String>,
    paths: Vec<(String, [f64; 7])>,
    blocks: HashMap<String, Vec<f64>>,
}

impl Ref {
    fn load(name: &str) -> Ref {
        let text = std::fs::read_to_string(data_dir().join(name)).unwrap();
        let mut params = HashMap::new();
        let mut paths = Vec::new();
        let mut blocks: HashMap<String, Vec<f64>> = HashMap::new();
        let mut cur: Option<String> = None;
        for line in text.lines() {
            if let Some(rest) = line.strip_prefix("#param ") {
                let mut it = rest.split_whitespace();
                params.insert(
                    it.next().unwrap().to_string(),
                    it.next().unwrap().to_string(),
                );
            } else if let Some(rest) = line.strip_prefix("#path ") {
                let mut it = rest.split_whitespace();
                let fname = it.next().unwrap().to_string();
                let v: Vec<f64> = it.map(|s| s.parse().unwrap()).collect();
                assert_eq!(v.len(), 7, "expected 7 path params");
                paths.push((fname, [v[0], v[1], v[2], v[3], v[4], v[5], v[6]]));
            } else if let Some(name) = line.strip_prefix("#begin ") {
                cur = Some(name.to_string());
                blocks.insert(name.to_string(), Vec::new());
            } else if line == "#end" {
                cur = None;
            } else if let Some(ref name) = cur {
                blocks.get_mut(name).unwrap().push(line.parse().unwrap());
            }
        }
        Ref {
            params,
            paths,
            blocks,
        }
    }

    fn f(&self, k: &str) -> f64 {
        self.params[k].parse().unwrap()
    }
    fn i(&self, k: &str) -> i32 {
        self.params[k].parse().unwrap()
    }
    fn s(&self, k: &str) -> &str {
        &self.params[k]
    }
    fn blk(&self, k: &str) -> &[f64] {
        &self.blocks[k]
    }

    /// Build the FeffPath list, applying the 7 stored path parameters
    /// (degen stays at the file default, matching the reference generator).
    fn build_paths(&self) -> Vec<FeffPath> {
        self.paths
            .iter()
            .map(|(fname, p)| {
                let mut path = FeffPath::from_path(data_dir().join(fname)).unwrap();
                let pp = &mut path.params;
                pp.s02 = p[0];
                pp.e0 = p[1];
                pp.deltar = p[2];
                pp.sigma2 = p[3];
                pp.third = p[4];
                pp.fourth = p[5];
                pp.ei = p[6];
                path
            })
            .collect()
    }

    fn transform(&self) -> Transform {
        let fitspace = match self.s("fitspace") {
            "r" => FitSpace::R,
            "k" => FitSpace::K,
            "q" => FitSpace::Q,
            "w" => FitSpace::W,
            other => panic!("unknown fitspace {other}"),
        };
        Transform::new(
            self.f("kmin"),
            self.f("kmax"),
            vec![self.i("kweight")],
            self.f("dk"),
            None,
            Window::from_name(self.s("window")).unwrap(),
            self.f("nfft") as usize,
            self.f("kstep"),
            self.f("rmin"),
            self.f("rmax"),
            self.f("dr"),
            None,
            Window::from_name(self.s("rwindow")).unwrap(),
            0.0,
            fitspace,
        )
    }

    fn dataset(&self) -> DataSet {
        DataSet::new(
            self.blk("data_k").to_vec(),
            self.blk("data_chi").to_vec(),
            self.build_paths(),
            self.transform(),
        )
    }
}

fn max_abs_diff(a: &[f64], b: &[f64]) -> f64 {
    assert_eq!(
        a.len(),
        b.len(),
        "length mismatch {} vs {}",
        a.len(),
        b.len()
    );
    a.iter()
        .zip(b)
        .map(|(x, y)| (x - y).abs())
        .fold(0.0, f64::max)
}

fn check_close(label: &str, got: &[f64], want: &[f64], tol: f64) {
    let d = max_abs_diff(got, want);
    println!("{label}: max|Δ| = {d:.3e}  (n={})", got.len());
    assert!(d < tol, "{label}: diff {d:.3e} exceeds {tol:.0e}");
}

fn check_rel(label: &str, got: f64, want: f64, rtol: f64) {
    let d = (got - want).abs();
    let r = d / want.abs().max(1e-300);
    println!("{label}: got={got:.8e} want={want:.8e} rel={r:.3e}");
    assert!(r < rtol, "{label}: rel {r:.3e} exceeds {rtol:.0e}");
}

fn run_residual_case(file: &str) {
    let r = Ref::load(file);
    let mut ds = r.dataset();
    ds.prepare_fit(Some(r.f("epsilon_k")));

    // model chi(k) from ff2chi
    let resid = ds.residual(false);

    // epsilon_r and n_idp (pure arithmetic) should match closely
    check_rel(
        &format!("{file} epsilon_r"),
        ds.epsilon_r()[0],
        r.f("epsilon_r"),
        1e-12,
    );
    check_rel(&format!("{file} n_idp"), ds.n_idp(), r.f("n_idp"), 1e-12);

    check_close(&format!("{file} residual"), &resid, r.blk("residual"), 1e-8);
}

#[test]
fn residual_r_matches_larch() {
    run_residual_case("ref_feffit_r.txt");
}

#[test]
fn residual_k_matches_larch() {
    run_residual_case("ref_feffit_k.txt");
}

#[test]
fn residual_q_matches_larch() {
    run_residual_case("ref_feffit_q.txt");
}

#[test]
fn residual_w_matches_larch() {
    // Cauchy-wavelet residual: realimag(cwt(diff/eps_k)).ravel() over the
    // [rmin,rmax)x[kmin,kmax) mask. It has a large dynamic range (peak ~36 here,
    // since eps_k=1e-3 scales diff up ~1000x), so compare the max abs error to
    // the peak magnitude rather than an absolute floor.
    let r = Ref::load("ref_feffit_w.txt");
    let mut ds = r.dataset();
    ds.prepare_fit(Some(r.f("epsilon_k")));
    let resid = ds.residual(false);

    check_rel("ref_feffit_w.txt n_idp", ds.n_idp(), r.f("n_idp"), 1e-12);

    let want = r.blk("residual");
    assert_eq!(resid.len(), want.len(), "w residual length");
    let peak = want.iter().fold(0.0f64, |m, &v| m.max(v.abs()));
    let maxd = max_abs_diff(&resid, want);
    println!(
        "w residual: max|Δ|={maxd:.3e} peak={peak:.3e} rel={:.3e} (n={})",
        maxd / peak,
        resid.len()
    );
    // rustfft vs numpy.fft round-off across the cwt's two FFT stages.
    assert!(
        maxd < 1e-9 * peak,
        "w residual max|Δ| {maxd:.3e} exceeds {:.0e}",
        1e-9 * peak
    );
}

#[test]
fn model_chi_matches_larch() {
    // verify the summed model chi(k) itself (ff2chi over the model paths)
    let r = Ref::load("ref_feffit_r.txt");
    let mut ds = r.dataset();
    ds.prepare_fit(Some(r.f("epsilon_k")));
    let model_chi = ds.model_chi_sum();
    check_close("model_chi", &model_chi, r.blk("model_chi"), 1e-9);
}

#[test]
fn estimate_noise_matches_larch() {
    let r = Ref::load("ref_feffit_noise.txt");
    let mut ds = r.dataset();
    ds.prepare_fit(Some(r.f("epsilon_k"))); // sets chi_interp
    let chi = ds.chi_interp().to_vec();
    ds.estimate_noise(&chi, r.f("rmin_noise"), r.f("rmax_noise"));
    check_rel(
        "est_epsilon_k",
        ds.epsilon_k()[0],
        r.f("est_epsilon_k"),
        1e-9,
    );
    check_rel(
        "est_epsilon_r",
        ds.epsilon_r()[0],
        r.f("est_epsilon_r"),
        1e-9,
    );
}