Source code for ewoksxas.tasks.tests.test_xftf

"""Tests for the Xftf task, focused on the array-size alignment path.

Like autobk, xftf can emit R-space arrays of differing length when inputs
differ. These tests stub ``larch.xafs.xftf`` to force ragged output and check
the task aligns it onto one grid (no-op when lengths already match).
"""

import numpy as np
import pytest
from ewoksorange.tests.utils import execute_task

from ewoksxas.converters.orange import Converter
from ewoksxas.tasks.xftf import Xftf


def _two_row_kspace_table():
    k = np.arange(200) * 0.05
    chi = np.vstack([np.cos(k), np.cos(k) * 0.5])
    return Converter().add_features(k, chi).to_table()


def _run_xftf(**parameters):
    inputs = {"Data": _two_row_kspace_table(), "parameters": parameters}
    outputs = execute_task(Xftf, inputs=inputs)
    return Converter.from_table(outputs["Data"]).features




@pytest.fixture
def make_stub_xftf(monkeypatch):
    """Factory: stub xftf so successive groups get r arrays of given lengths."""

    def _make(lengths: list[int]):
        state = {"i": 0}

        def fake_xftf(data, group=None, **kwargs):
            n = lengths[state["i"]]
            state["i"] += 1
            r = np.arange(n) * 0.05
            group.r = r
            group.chir_mag = np.abs(np.cos(r))

        monkeypatch.setattr("larch.xafs.xftf", fake_xftf)
        return state

    return _make





def test_matched_lengths_are_not_resampled(make_stub_xftf):
    """Equal-length spectra bypass alignment and stack byte-identically."""
    make_stub_xftf([60, 60])
    r, y = _run_xftf()

    expected_r = np.arange(60) * 0.05
    np.testing.assert_array_equal(r, expected_r)
    np.testing.assert_array_equal(y, np.vstack([np.abs(np.cos(expected_r))] * 2))
    assert not np.isnan(y).any()





def test_interpolation_aligns_ragged_without_nan(make_stub_xftf):
    make_stub_xftf([60, 45])
    r, y = _run_xftf(resample="interpolation")

    assert y.shape == (2, r.size)
    assert not np.isnan(y).any()
    assert r[-1] <= (45 - 1) * 0.05 + 1e-9  # overlap bound (shortest spectrum)





def test_rebinning_aligns_ragged_with_nan_gaps(make_stub_xftf):
    make_stub_xftf([60, 45])
    r, y = _run_xftf(resample="rebinning")

    assert y.shape == (2, r.size)
    assert r[-1] >= (60 - 1) * 0.05 - 1e-9  # union reaches the longest spectrum
    assert not np.isnan(y[0]).any()
    assert np.isnan(y[1]).any()