Source code for ewoksxas.fit.tests.test_peak_fit

"""Tests for the peak-fitting module ewoksxas.fit.peak_fit.

Peaks are synthesised from silx's own profile functions on a noiseless grid (see
conftest), so the fitted center, area, width(s) and mixing are known exactly.
The Strip/Snip backgrounds are silx's data-derived peak-stripping filters: with
a filter width suited to the peak they recover the truth, so their tests pass
tuned widths; at unsuited widths the filter absorbs part of a broad peak, which
is inherent to those backgrounds (the widths are user-tunable in the GUI).
"""

import numpy as np
import pytest

from ewoksxas.fit.data import XYData
from ewoksxas.fit.peak_fit import Baseline, PeakModel, PeakShape
from ewoksxas.fit.tests.conftest import (
    TRUE_ETA,
    TRUE_FWHM,
    TRUE_HIGH_FWHM,
    TRUE_LOW_FWHM,
)

SHAPES = [
    "gaussian",
    "skewed gaussian",
    "lorentzian",
    "split lorentzian",
    "voigt",
    "skewed voigt",
]
# Baselines with strict recovery guarantees (the parametric silx backgrounds).
BASELINES = ["none", "offset", "linear"]


def _add_baseline(x, peak, baseline):
    """Return ``peak`` plus a synthetic baseline of the requested kind."""
    if baseline == "offset":
        background = np.full_like(x, 5.0)
    elif baseline == "linear":
        background = 0.3 * x + 4.0
    else:
        background = np.zeros_like(x)
    return peak + background




def test_peakshape_metadata_and_from_label():
    # Six distinct members carrying their silx theory name, family and symmetry.
    assert len(list(PeakShape)) == 6
    assert not PeakShape.GAUSSIAN.asymmetric
    assert PeakShape.SKEWED_GAUSSIAN.asymmetric
    assert PeakShape.SPLIT_LORENTZIAN.theory_name == "Split Lorentz"
    assert PeakShape.VOIGT.family == "voigt"
    assert PeakShape.from_label("skewed voigt") is PeakShape.SKEWED_VOIGT
    assert PeakShape.from_label(PeakShape.VOIGT) is PeakShape.VOIGT
    assert PeakShape.from_theory_name("Gaussians") is PeakShape.GAUSSIAN
    with pytest.raises(ValueError, match="Unknown shape"):
        PeakShape.from_label("pearson")
    with pytest.raises(ValueError, match="Unknown silx theory"):
        PeakShape.from_theory_name("Hypermet")





def test_baseline_metadata_and_from_label():
    # Five members mapping straight to silx background theories.
    assert len(list(Baseline)) == 5
    assert Baseline.NONE.theory_name == "No Background"
    assert Baseline.OFFSET.theory_name == "Constant"
    assert Baseline.LINEAR.theory_name == "Linear"
    assert Baseline.STRIP.theory_name == "Strip"
    assert Baseline.SNIP.theory_name == "Snip"
    assert Baseline.from_label("strip") is Baseline.STRIP
    assert Baseline.from_theory_name("Snip") is Baseline.SNIP
    with pytest.raises(ValueError, match="Unknown baseline"):
        Baseline.from_label("polynomial")
    with pytest.raises(ValueError, match="Unknown silx background"):
        Baseline.from_theory_name("Degree 2 Polynomial")





@pytest.mark.filterwarnings("ignore::RuntimeWarning")
@pytest.mark.parametrize("shape", SHAPES)
@pytest.mark.parametrize("baseline", BASELINES)
def test_fit_recovers_peak(x, peak_factory, shape, baseline):
    peak, truth = peak_factory(shape)
    data = _add_baseline(x, peak, baseline)

    result = PeakModel(XYData(x=x, y=data), shape=shape, baseline=baseline).fit()

    assert result.r2 == pytest.approx(1.0, abs=1e-3)
    assert result.center == pytest.approx(truth["center"], abs=5e-2)
    assert result.area == pytest.approx(truth["area"], rel=3e-2)
    assert result.height > 0
    # Asymmetric shapes report per-side widths; symmetric ones a single FWHM.
    if result.shape.asymmetric:
        assert result.fwhm_low > 0
        assert result.fwhm_high > 0
    else:
        assert result.fwhm > 0
    # The total model (background + peaks) reconstructs the data.
    np.testing.assert_allclose(result.model.y, data, atol=0.15)



# The conftest grid samples the FWHM-4 peak with 160 points; these stripping
# widths exceed the peak width so the filters leave the peak intact.
_STRIPPING = [("strip", {"strip_width": 200}), ("snip", {"snip_width": 300})]




@pytest.mark.parametrize(("baseline", "settings"), _STRIPPING)
def test_stripping_backgrounds_recover_peak(x, peak_factory, baseline, settings):
    peak, truth = peak_factory("gaussian")
    data = peak + 5.0 + 0.1 * x

    result = PeakModel(
        XYData(x=x, y=data), shape="gaussian", baseline=baseline, **settings
    ).fit()

    assert result.center == pytest.approx(truth["center"], abs=5e-2)
    assert result.area == pytest.approx(truth["area"], rel=3e-2)
    # The stripped background follows the synthetic one.
    np.testing.assert_allclose(result.baseline.y, 5.0 + 0.1 * x, atol=0.5)





@pytest.mark.parametrize(("baseline", "settings"), _STRIPPING)
def test_stripping_baseline_with_roi_after_full_range(
    x, peak_factory, baseline, settings
):
    # Regression: re-fitting with a different data length (a region-of-interest
    # change) must not crash. silx's stock Snip background cached the previous
    # y-array process-wide and raised a broadcast ValueError on the second fit
    # (replaced by the cached background theories in ewoksxas.fit.silx_engine).
    peak, truth = peak_factory("gaussian")
    data = peak + 5.0 + 0.1 * x

    full = PeakModel(
        XYData(x=x, y=data), shape="gaussian", baseline=baseline, **settings
    ).fit()
    restricted = PeakModel(
        XYData(x=x, y=data),
        shape="gaussian",
        baseline=baseline,
        roi=(-5.0, 9.0),
        **settings,
    ).fit()

    assert full.center == pytest.approx(truth["center"], abs=5e-2)
    assert restricted.center == pytest.approx(truth["center"], abs=5e-2)





@pytest.mark.parametrize("baseline", ["offset", "linear", "strip", "snip"])
@pytest.mark.parametrize("roi", [None, (-5.0, 9.0)])
def test_baseline_roi_matrix(x, peak_factory, baseline, roi):
    # Every baseline works with and without a region of interest, and the
    # fitted background is evaluated on the (possibly restricted) model grid.
    peak, truth = peak_factory("gaussian")
    background = np.full_like(x, 5.0) if baseline == "offset" else 5.0 + 0.1 * x
    settings = dict(_STRIPPING).get(baseline, {})

    result = PeakModel(
        XYData(x=x, y=peak + background),
        shape="gaussian",
        baseline=baseline,
        roi=roi,
        **settings,
    ).fit()

    assert result.center == pytest.approx(truth["center"], abs=5e-2)
    assert result.baseline.y.shape == result.model.x.shape





def test_accepts_enum_members(x, peak_factory):
    # The constructor accepts PeakShape/Baseline members, not only labels.
    peak, _ = peak_factory("skewed voigt")
    result = PeakModel(
        XYData(x=x, y=peak + 5.0),
        shape=PeakShape.SKEWED_VOIGT,
        baseline=Baseline.OFFSET,
    ).fit()
    assert result.shape is PeakShape.SKEWED_VOIGT
    assert result.center == pytest.approx(2.0, abs=5e-2)





def test_lmfit_gaussian_cross_check(x, lmfit_gaussian):
    # Independent oracle: lmfit's Gaussian is identical maths to silx's.
    y, center, area = lmfit_gaussian
    result = PeakModel(XYData(x=x, y=y), shape="gaussian").fit()
    assert result.center == pytest.approx(center, abs=1e-3)
    assert result.area == pytest.approx(area, rel=1e-3)





def test_gaussian_descriptors(x, peak_factory):
    peak, truth = peak_factory("gaussian")
    result = PeakModel(XYData(x=x, y=peak), shape="gaussian").fit()
    assert result.center == pytest.approx(truth["center"], abs=1e-3)
    assert result.area == pytest.approx(truth["area"], rel=1e-3)
    assert result.fwhm == pytest.approx(TRUE_FWHM, abs=1e-2)
    # A symmetric, non-Voigt shape reports only a single fwhm.
    descriptors = result.as_dict()
    assert "fwhm" in descriptors
    assert {"fwhm_low", "fwhm_high", "eta"}.isdisjoint(descriptors)





def test_lorentzian_descriptors(x, peak_factory):
    peak, _ = peak_factory("lorentzian")
    result = PeakModel(XYData(x=x, y=peak), shape="lorentzian").fit()
    assert result.fwhm == pytest.approx(TRUE_FWHM, abs=1e-2)
    assert {"fwhm_low", "fwhm_high", "eta"}.isdisjoint(result.as_dict())





def test_split_lorentzian_is_asymmetric(x, peak_factory):
    peak, _ = peak_factory("split lorentzian")
    result = PeakModel(XYData(x=x, y=peak), shape="split lorentzian").fit()
    # Low-side FWHM 4, high-side FWHM 6.
    assert result.fwhm_low == pytest.approx(TRUE_LOW_FWHM, abs=1e-2)
    assert result.fwhm_high == pytest.approx(TRUE_HIGH_FWHM, abs=1e-2)
    assert {"fwhm", "eta"}.isdisjoint(result.as_dict())





def test_voigt_reports_eta(x, peak_factory):
    peak, _ = peak_factory("voigt")
    result = PeakModel(XYData(x=x, y=peak), shape="voigt").fit()
    # Voigt is symmetric -> a single fwhm, plus the pseudo-Voigt mixing eta.
    assert result.fwhm > 0.0
    # silx fits eta directly; it is recovered from the synthesised mixing.
    assert result.eta == pytest.approx(TRUE_ETA, abs=2e-2)
    assert {"fwhm_low", "fwhm_high"}.isdisjoint(result.as_dict())





def test_roi_restricts_fit(x, peak_factory):
    peak, _ = peak_factory("gaussian")
    result = PeakModel(XYData(x=x, y=peak), roi=(-5.0, 9.0)).fit()
    assert result.peak.x.min() >= -5.0
    assert result.peak.x.max() <= 9.0
    assert result.peak.x.size < x.size
    assert result.center == pytest.approx(2.0, abs=1e-3)





def test_from_dict(x, peak_factory):
    peak, _ = peak_factory("voigt")
    model = PeakModel.from_dict(
        {"x": x, "y": peak + 5.0},
        {
            "shape": "voigt",
            "baseline": "offset",
            "xmin": -10.0,
            "xmax": 12.0,
        },
    )
    assert model.shape is PeakShape.VOIGT
    assert model.baseline is Baseline.OFFSET
    assert model.xmin == -10.0
    assert model.xmax == 12.0
    assert model.fit().center == pytest.approx(2.0, abs=5e-2)





def test_from_dict_stripping_settings(x, peak_factory):
    peak, _ = peak_factory("gaussian")
    model = PeakModel.from_dict(
        {"x": x, "y": peak + 5.0},
        {"baseline": "strip", "strip_width": 200, "strip_iterations": 4000},
    )
    assert model.baseline is Baseline.STRIP
    assert model.strip_width == 200
    assert model.strip_iterations == 4000
    assert model.fit().center == pytest.approx(2.0, abs=5e-2)





def test_invalid_shape(x):
    with pytest.raises(ValueError, match="Unknown shape"):
        PeakModel(XYData(x=x, y=x), shape="step")





def test_invalid_baseline(x):
    with pytest.raises(ValueError, match="Unknown baseline"):
        PeakModel(XYData(x=x, y=x), baseline="polynomial")





def test_too_few_points():
    with pytest.raises(ValueError, match="too few points"):
        PeakModel(XYData(x=np.array([0.0, 1.0]), y=np.array([0.0, 1.0])))





def test_result_as_dict_and_reports(x, peak_factory):
    peak, _ = peak_factory("voigt")
    result = PeakModel(XYData(x=x, y=peak), shape="voigt").fit()

    descriptors = result.as_dict()
    assert {
        "center",
        "height",
        "area",
        "fwhm",
        "eta",
        "r2",
        "rmse",
    } <= set(descriptors)
    assert all(isinstance(value, float) for value in descriptors.values())

    assert isinstance(result.summary(), str)
    assert "Peak Fit Results" in repr(result)
    assert "voigt" in repr(result)  # the shape label appears in the report
    assert result.dof > 0



# (center, height) of three well-separated Gaussians used by the multi-peak
# tests; widely spaced relative to the FWHM so each peak is cleanly resolvable.
_MULTI_PEAKS = [(-8.0, 20.0), (0.0, 15.0), (8.0, 25.0)]




def test_auto_n_selects_true_peak_count(x, multi_peak_factory):
    y = multi_peak_factory(_MULTI_PEAKS)
    result = PeakModel(XYData(x=x, y=y), shape="gaussian").fit()

    assert result.n_peaks == 3
    # Peaks are reported left to right.
    centers = [peak["center"] for peak in result.peaks]
    assert centers == sorted(centers)
    np.testing.assert_allclose(centers, [-8.0, 0.0, 8.0], atol=1e-2)
    assert result.r2 == pytest.approx(1.0, abs=1e-3)
    # The bare scalar properties alias the leftmost peak.
    assert result.center == pytest.approx(-8.0, abs=1e-2)





def test_auto_n_single_peak_selects_one(x, peak_factory):
    peak, _ = peak_factory("gaussian")
    result = PeakModel(XYData(x=x, y=peak), shape="gaussian").fit()
    assert result.n_peaks == 1
    assert result.center == pytest.approx(2.0, abs=1e-3)





def test_auto_n_two_peaks(x, multi_peak_factory):
    y = multi_peak_factory([(-6.0, 20.0), (6.0, 15.0)])
    result = PeakModel(XYData(x=x, y=y), shape="gaussian").fit()
    assert result.n_peaks == 2





def test_multi_peak_descriptors_have_all_keys(x, multi_peak_factory):
    y = multi_peak_factory(_MULTI_PEAKS[:2])
    result = PeakModel(XYData(x=x, y=y), shape="gaussian").fit()
    expected = {
        "center",
        "center_stderr",
        "height",
        "area",
        "area_stderr",
        "fwhm",
    }
    for peak in result.peaks:
        assert expected <= set(peak)
        assert all(isinstance(value, float) for value in peak.values())