Source code for ewoksxrdct.tests.test_reshape_aeroystepscan

import shutil
from pathlib import Path

import h5py
import numpy as np
import pytest
from ewoksutils.exceptions import TaskExecutionError

from ..tasks.reshape_aeroystepscan import ReshapeAero
from .conftest import write_integration_output as _write_integration_output

BLISS_MASTER = Path(__file__).parent / "data" / "test_aeroystepscan.h5"

# BLISS_MASTER records an aeroystepscan with 36 rotation points (x_range=360 in
# steps of 10 from 0) at each of 3 translation points. The rotation motor keeps
# turning between rows, so rows 1 and 2 wrap through 360 degrees and are not
# recorded in ascending order.
N_TRANS = 3
N_ROT = 36
N_POINTS = N_TRANS * N_ROT

# Parameters to use to generate test ai file
N_RADIAL = 1500
RADIAL_AXIS = np.linspace(0.003, 5.477, N_RADIAL, dtype=np.float32)
INTENSITY = np.arange(N_POINTS * N_RADIAL, dtype=np.float32).reshape(N_POINTS, N_RADIAL)




@pytest.fixture(scope="module")
def ai_file(tmp_path_factory):
    """An azimuthal-integration output with one frame per scan point, built
    once and shared across tests since they only ever read it."""
    path = tmp_path_factory.mktemp("ai") / "ai.h5"
    _write_integration_output(path, radial_axis=RADIAL_AXIS, intensity=INTENSITY)
    return path



def _recorded_positions(motor_name):
    """The positions recorded for `motor_name` in BLISS_MASTER."""
    with h5py.File(BLISS_MASTER, "r") as f:
        return f[f"1.1/measurement/{motor_name}"][()]


def _copy_master_with_overrides(tmp_path, overrides):
    """Copy BLISS_MASTER, replacing or adding the datasets named by the HDF5
    paths in `overrides`, to build a master that differs from the recorded
    scan."""
    master = tmp_path / "master.h5"
    shutil.copy(BLISS_MASTER, master)
    with h5py.File(master, "r+") as f:
        for path, value in overrides.items():
            if path in f:
                del f[path]
            f[path] = value
    return master


def _make_task(ai_file, master=BLISS_MASTER, **input_overrides):
    inputs = {
        "bliss_master_path": str(master),
        "integration_output_path": str(ai_file),
        "scan_number": 1,
        "translation_motor_name": "hry",
        "rotation_motor_name": "hrrz",
        "integration_intensity_path": "/intensities",
        "integration_radial_axis_path": "/2th_deg",
    }
    inputs.update(input_overrides)
    return ReshapeAero(inputs=inputs)




def test_run_reshapes_input_data(ai_file):
    """Reshapes into (N_trans, N_rot) and sorts each row by rotation, moving
    the frames and translations by the same permutation."""
    task = _make_task(ai_file)
    task.execute()

    rotations = task.outputs.rotation_angles
    assert rotations.shape == (N_POINTS,)
    assert task.outputs.integration_intensity_values.shape == (N_POINTS, N_RADIAL)

    # Every row now ascends, including the two that wrapped through 360.
    for row in rotations.reshape(N_TRANS, N_ROT):
        assert np.all(np.diff(row) > 0)

    # Each frame of INTENSITY starts at its own scan point index times N_RADIAL,
    # so the first radial bin says where each output frame came from.
    origin = np.rint(task.outputs.integration_intensity_values[:, 0] / N_RADIAL)
    origin = origin.astype(int)
    np.testing.assert_array_equal(rotations, _recorded_positions("hrrz")[origin])
    np.testing.assert_array_equal(
        task.outputs.translation_values, _recorded_positions("hry")[origin]
    )

    assert task.outputs.integration_radial_axis == "2th_deg"
    np.testing.assert_array_equal(
        task.outputs.integration_radial_axis_values, RADIAL_AXIS
    )

    # Grid values derived from the fscan_parameters of BLISS_MASTER.
    assert task.outputs.rotation_grid_params == pytest.approx((5, 355, N_ROT))





def test_run_wrong_motor_name_raises(ai_file):
    """A motor name that doesn't exist in the bliss master file raises ValueError."""
    task = _make_task(ai_file, rotation_motor_name="does_not_exist")
    with pytest.raises(TaskExecutionError, match="does_not_exist") as excinfo:
        task.execute()
    assert isinstance(excinfo.value.__cause__, ValueError)





def test_run_wrong_integration_path_raises(ai_file):
    """A radial axis or intensity HDF5 path that does not exist raises ValueError."""
    task = _make_task(ai_file, integration_intensity_path="/does/not/exist")
    with pytest.raises(TaskExecutionError, match="/does/not/exist") as excinfo:
        task.execute()
    assert isinstance(excinfo.value.__cause__, ValueError)





def test_run_warns_on_rotation_count_mismatch(tmp_path, ai_file):
    """A mismatch between the computed and the fscan-derived rotation count warns."""
    # x_range=400 with x_step_size=10 implies 40 rotation points, against the 36
    # the recorded array length gives.
    master = _copy_master_with_overrides(
        tmp_path, {"1.1/instrument/fscan_parameters/x_range": 400}
    )

    task = _make_task(ai_file, master=master)
    with pytest.warns(UserWarning, match="reshaped incorrectly"):
        task.execute()