Source code for est.tests.pymca.test_io_with_pymca

import numpy

from ...core.io.read_xas import build_ascii_data_url
from ...core.io.read_xas import read_from_url
from ...core.io.write_xas import write_xas
from ...core.types.xasobject import XASObject
from ..data import example_spectra
from ..hdf5 import assert_hdf5_content




def test_read_spectrum(filename_cu_from_pymca):
    """Test read function for spectra and configuration"""
    res = read_from_url(
        spectra_url=build_ascii_data_url(
            file_path=filename_cu_from_pymca,
            col_name="Column 1",
            scan_title=None,
            data_slice=None,
        ),
        channel_url=build_ascii_data_url(
            file_path=filename_cu_from_pymca,
            col_name="Column 2",
            scan_title=None,
            data_slice=None,
        ),
    )
    assert isinstance(res, XASObject)
    assert res.n_spectrum == 1
    spectrum_0 = res.spectra.data.flat[0]
    assert spectrum_0.mu is not None
    assert spectrum_0.energy is not None





def test_nx_writing(tmp_path):
    """Test that the nx process is correctly store ad the output data"""
    dim1 = 2
    dim2 = 1
    energy, spectra = example_spectra(shape=(256, dim1, dim2))
    xas_obj = XASObject(spectra=spectra, energy=energy, dim1=dim1, dim2=dim2)
    h5_file = str(tmp_path / "output_file.h5")

    write_xas(h5_file, xas_obj)

    expected = {
        "data": {
            "data@NX_class": "NXdata",
            "data@axes": ["energy"],
            "data@signal": "mu",
            "energy": {"energy@long_name": "Energy (eV)", "energy@units": "eV"},
            "mu": {"mu@long_name": "μ"},
        },
    }

    assert_hdf5_content(h5_file, expected, dim1 * dim2)

    original_energy = xas_obj.energy
    original_mu = xas_obj.absorbed_beam()

    shape = xas_obj.spectra.data.shape
    for i in range(dim1 * dim2):
        loaded_xas_obj = XASObject.from_file(h5_file, entry=f"1.{i+1}")

        loaded_energy = loaded_xas_obj.energy
        numpy.testing.assert_allclose(loaded_energy, original_energy)

        loaded_mu = loaded_xas_obj.absorbed_beam().ravel()
        idx = (slice(None),) + numpy.unravel_index(i, shape)
        original_mu_i = original_mu[idx]
        numpy.testing.assert_allclose(loaded_mu, original_mu_i)