Source code for ewoksxas.resources.data.regenerate_reference

"""Regenerate the reference snapshots used by the test suite.

The tests compare computed arrays against stored ``.npz`` snapshots (in the
``snapshots`` directory next to this module) with a tolerance, which is portable
across operating systems where byte-level hashing is not. Run this only after a
dependency change (e.g. a larch upgrade) has been reviewed and accepted, to
refresh those snapshots::

    python -m ewoksxas.resources.data.regenerate_reference              # all
    python -m ewoksxas.resources.data.regenerate_reference autobk_chi   # one

Adding a snapshot for a new test takes two steps:

1. Write a builder function returning a ``dict`` of arrays and register it in
   ``SNAPSHOTS`` under the snapshot name.
2. In the test, load it with
   ``np.load(resource_filename("ewoksxas:data/snapshots/<name>.npz"))`` and
   compare with ``numpy.testing.assert_allclose``.

The data-loading helpers mirror the fixtures in ``tasks/tests/conftest.py``.
"""

from __future__ import annotations

import sys
from pathlib import Path
from typing import TYPE_CHECKING

import larch
import numpy as np
from ewoksorange.tests.utils import execute_task
from silx.resources import resource_filename

from ewoksxas.converters.orange import Converter
from ewoksxas.tasks.autobk import AutoBK
from ewoksxas.tasks.normalize import Normalize
from ewoksxas.tasks.read_scans import ReadScans
from ewoksxas.tasks.read_sources import ReadSources

if TYPE_CHECKING:
    from collections.abc import Callable

SNAPSHOT_DIR = Path(__file__).parent / "snapshots"


# --- shared data loaders (mirror tasks/tests/conftest.py) ------------------




def raw_fe_foil():
    """Return the raw Fe-foil scans, as the fe_foil_data fixture does."""
    data_path = resource_filename("ewoksxas:data/Fe_foil_0001.h5")
    sources = execute_task(
        ReadSources, inputs={"file_paths": [data_path], "filters": []}
    )
    return execute_task(
        ReadScans,
        inputs={
            "Data": sources["Data"],
            "x": "instrument/energy_enc/data",
            "counters": [{"name": "mu_trans", "path": "measurement/mu_trans"}],
            "metadata": [],
        },
    )["Data"]





def normalized_fe_foil_one():
    """Return a single normalized Fe-foil spectrum, as the fe_foil_one fixture."""
    return execute_task(Normalize, inputs={"Data": raw_fe_foil()})["Data"][:1]



# --- snapshot builders: name -> dict of arrays to store --------------------


def _build_fe_foil_raw() -> dict:
    energy, mu = Converter.from_table(raw_fe_foil()).features
    return {"energy": energy, "mu": mu}


def _build_autobk_chi() -> dict:
    outputs = execute_task(
        AutoBK,
        inputs={"Data": normalized_fe_foil_one(), "parameters": {"output": "chi"}},
    )
    k, chi = Converter.from_table(outputs["Data"]).features
    return {
        "k": k,
        "chi": chi,
        "larch_version": getattr(larch, "__version__", "unknown"),
    }


# Registry of available snapshots. Add a builder here to support a new test.
SNAPSHOTS: dict[str, Callable[[], dict]] = {
    "fe_foil_raw_reference": _build_fe_foil_raw,
    "autobk_chi_reference": _build_autobk_chi,
}




def regenerate(name: str) -> Path:
    """Rebuild one snapshot and write it into the snapshots directory."""
    try:
        builder = SNAPSHOTS[name]
    except KeyError:
        known = ", ".join(sorted(SNAPSHOTS))
        raise KeyError(f"Unknown snapshot '{name}'. Known snapshots: {known}") from None
    SNAPSHOT_DIR.mkdir(parents=True, exist_ok=True)
    path = SNAPSHOT_DIR / f"{name}.npz"
    np.savez(path, **builder())
    print(f"Wrote {path}")
    return path





def main(argv: list[str] | None = None) -> None:
    """Regenerate the snapshots named on the command line, or all of them."""
    names = argv if argv else list(SNAPSHOTS)
    for name in names:
        regenerate(name)



if __name__ == "__main__":
    main(sys.argv[1:])