import h5py
import numpy as np
import pytest
from pathlib import Path
from unittest.mock import patch
import shutil
from ewokstomo.tasks.online.reducedarkflat import OnlineReduceDarkFlat
from ewokstomo.tasks.reducedarkflat import ReduceDarkFlat
from ewokstomo.tests.test_reducedarkflat import get_data_dir, get_raw_data_dir
from ewokstomo.tests.online.mock import FakeScan
[docs]
@pytest.fixture
def TestEwoksTomo_0010_dataset(tmp_path) -> Path:
scan = "TestEwoksTomo_0010"
processed_dir = get_data_dir(scan)
raw_dir = get_raw_data_dir(scan)
dst_dir = tmp_path / scan
shutil.copytree(processed_dir, dst_dir)
shutil.copy(raw_dir / f"{scan}.h5", dst_dir / f"{scan}.h5")
for pattern in ("*_darks.hdf5", "*_flats.hdf5"):
for f in dst_dir.glob(pattern):
f.unlink()
return dst_dir
[docs]
def test_reducedarkflat_online(TestEwoksTomo_0010_dataset):
nx_file = TestEwoksTomo_0010_dataset / "projections" / "TestEwoksTomo_0010.nx"
h5py_file = TestEwoksTomo_0010_dataset / "TestEwoksTomo_0010.h5"
# --- Run offline task to produce ground-truth ---
offline_task = ReduceDarkFlat(
inputs={
"nx_path": str(nx_file),
"dark_reduction_method": "mean",
"flat_reduction_method": "median",
"overwrite": True,
"return_info": False,
},
)
offline_task.execute()
darks_offline = Path(offline_task.outputs.reduced_darks_path)
flats_offline = Path(offline_task.outputs.reduced_flats_path)
# Load offline reduced frames
with h5py.File(darks_offline, "r") as f:
offline_dark = f["entry0000/darks/0"][()]
with h5py.File(flats_offline, "r") as f:
offline_flat = f["entry0000/flats/2"][()]
# --- Prepare fake online scan using raw frames ---
with h5py.File(h5py_file, "r") as f:
dark_frames = [np.array(frame) for frame in f["2.1/measurement/edgetwinmic"]]
flat_frames = [np.array(frame) for frame in f["3.1/measurement/edgetwinmic"]]
# Test both darks and flats
for frames, title, method, expected in [
(dark_frames, "darks", "mean", offline_dark),
(flat_frames, "flats", "median", offline_flat),
]:
fake_scan = FakeScan(frames, title=title)
with (
patch("ewokstomo.tasks.online.reducedarkflat.BeaconData") as MockBeacon,
patch("ewokstomo.tasks.online.reducedarkflat.DataStore") as MockStore,
):
MockBeacon.return_value.get_redis_data_db.return_value = "redis://fake"
MockStore.return_value.load_scan.return_value = fake_scan
output_file = TestEwoksTomo_0010_dataset / f"online_{title}.h5"
task = OnlineReduceDarkFlat(
inputs={
"scan_key": f"scan_{title}",
"index": 0,
"reduction_method": method,
"output_file_path": str(output_file),
}
)
task.execute()
assert output_file.is_file()
# Load online reduced frames
with h5py.File(output_file, "r") as f:
reduced_online = f[f"entry0000/{title}s"][()]
# --- Compare online vs offline results ---
np.testing.assert_allclose(reduced_online, expected, rtol=1e-6)
[docs]
def test_reducedarkflat_offline(TestEwoksTomo_0010_dataset):
nx_file = TestEwoksTomo_0010_dataset / "projections" / "TestEwoksTomo_0010.nx"
h5py_file = TestEwoksTomo_0010_dataset / "TestEwoksTomo_0010.h5"
# --- Run offline task to produce ground-truth ---
offline_task = ReduceDarkFlat(
inputs={
"nx_path": str(nx_file),
"dark_reduction_method": "mean",
"flat_reduction_method": "median",
"overwrite": True,
"return_info": False,
},
)
offline_task.execute()
darks_offline = Path(offline_task.outputs.reduced_darks_path)
flats_offline = Path(offline_task.outputs.reduced_flats_path)
# Load offline reduced frames
with h5py.File(darks_offline, "r") as f:
offline_dark = f["entry0000/darks/0"][()]
with h5py.File(flats_offline, "r") as f:
offline_flat = f["entry0000/flats/2"][()]
# Test both darks and flats
for title, method, expected in [
("dark", "mean", offline_dark),
("flat", "median", offline_flat),
]:
output_file = TestEwoksTomo_0010_dataset / f"offline_{title}.h5"
task = OnlineReduceDarkFlat(
inputs={
"scan_file_path": str(h5py_file),
"scan_number": 2 if "dark" in title else 3,
"index": 0,
"reduction_method": method,
"output_file_path": str(output_file),
}
)
task.execute()
assert output_file.is_file()
with h5py.File(output_file, "r") as f:
reduced_online = f[f"entry0000/{title}s"][()]
np.testing.assert_allclose(reduced_online, expected, rtol=1e-6)
