diff --git a/src/hextools/exporters/tomography.py b/src/hextools/exporters/tomography.py index 3238adb..bd76998 100644 --- a/src/hextools/exporters/tomography.py +++ b/src/hextools/exporters/tomography.py @@ -1,7 +1,9 @@ """Exporters for tomography data.""" import logging +from collections.abc import Mapping from pathlib import Path +from typing import Any import h5py from bluesky_tiled_plugins import CatalogOfBlueskyRuns @@ -19,6 +21,46 @@ logger = logging.getLogger("hextools.exporters.tomography") +def resolve_dark_flat_run( + tiled_client: Mapping[Any, Any], + run_start: Mapping[str, Any], + dark_flat_id: str | int | None = None, +) -> BlueskyRunV3 | None: + """Return the dark/flat run to link into a tomography export, or None. + + Uses ``dark_flat_id`` if given, else the run's ``last_dark_flat_uid``, and + only if that dark/flat was taken in the same proposal and cycle as the run. + """ + uid = ( + dark_flat_id + if dark_flat_id is not None + else run_start.get("last_dark_flat_uid") + ) + if uid is None: + logger.warning("No last_dark_flat_uid in the run start document") + return None + try: + dark_flat_run = tiled_client[uid] + except KeyError: + dark_flat_run = None + if not isinstance(dark_flat_run, BlueskyRunV3): + logger.warning(f"Dark/flat run {uid} not found; omitting dark/flat entries") + return None + # A dark/flat from the same proposal in an earlier cycle is stale too. + dark_flat_from = tuple( + dark_flat_run.start.get(k) for k in ("data_session", "cycle") + ) + scan_from = tuple(run_start.get(k) for k in ("data_session", "cycle")) + if dark_flat_from != scan_from: + logger.warning( + f"Dark/flat run {uid} is from {dark_flat_from}, not {scan_from}; " + "omitting dark/flat entries" + ) + return None + logger.info(f"Using dark/flat run {uid}") + return dark_flat_run + + def export_tomo_flyscan_to_h5( tomo_flyscan: BlueskyRunV3 | str | int, tiled_client: CatalogOfBlueskyRuns | None = None, @@ -44,21 +86,7 @@ def export_tomo_flyscan_to_h5( f" for proposal {proposal}, cycle {cycle}" ) - dark_flat_uid = dark_flat_id or run_start.get("last_dark_flat_uid") - - if dark_flat_uid is None: - logger.warning( - "No last_dark_flat_uid found in the run start document. " - "Omitting dark/flat entries" - ) - dark_flat_run = None - else: - logger.info(f"Using dark/flat run {dark_flat_uid}") - dark_flat_run = tiled_client[dark_flat_uid] - if not isinstance(dark_flat_run, BlueskyRunV3): - raise ValueError( - f"Failed to resolve dark/flat run {dark_flat_uid} to a BlueskyRunV3" - ) + dark_flat_run = resolve_dark_flat_run(tiled_client, run_start, dark_flat_id) scan_id_dir_path = Path( f"/nsls2/data/hex/proposals/{cycle}/{proposal}/tomography/metadata/scan_{scan_id:05}" @@ -101,8 +129,8 @@ def export_tomo_flyscan_to_h5( missing = set(det_names) - found.keys() if missing: logger.warning( - f"Dark/flat run {dark_flat_uid} has no {stream_name} data for " - f"{sorted(missing)}" + f"Dark/flat run {dark_flat_run.start.get('uid')} has no " + f"{stream_name} data for {sorted(missing)}" ) for name in set(det_names) & found.keys(): relpath, dataset, _ = found[name] diff --git a/src/hextools/tomography/utils.py b/src/hextools/tomography/utils.py index 23563a5..cf31a7d 100644 --- a/src/hextools/tomography/utils.py +++ b/src/hextools/tomography/utils.py @@ -3,6 +3,8 @@ from collections.abc import MutableMapping from typing import Any +from hextools.utils import DARK_FLAT_MD_KEYS + class LastDarkFlatUIDTracker: """RunEngine subscription that tracks the last dark-flat scan UID and scan ID.""" @@ -16,7 +18,6 @@ def __call__(self, name: str, doc: dict[str, Any]): and doc.get("uid", None) and doc.get("plan_name", None) == "tomo_dark_flat" ): - last_uid = doc.get("uid", None) - last_scan_id = doc.get("scan_id", None) - self._metadata_dict["last_dark_flat_uid"] = last_uid - self._metadata_dict["last_dark_flat_scan_id"] = last_scan_id + uid_key, scan_id_key = DARK_FLAT_MD_KEYS + self._metadata_dict[uid_key] = doc.get("uid", None) + self._metadata_dict[scan_id_key] = doc.get("scan_id", None) diff --git a/src/hextools/utils/__init__.py b/src/hextools/utils/__init__.py index f64e588..07c6c57 100644 --- a/src/hextools/utils/__init__.py +++ b/src/hextools/utils/__init__.py @@ -588,8 +588,26 @@ def print_device_tree(device: Device, indent: int = 0) -> None: print("\n".join(x)) +# Set by LastDarkFlatUIDTracker. +DARK_FLAT_MD_KEYS = ("last_dark_flat_uid", "last_dark_flat_scan_id") + + +def clear_last_dark_flat(RE: RunEngine | None = None): + """Forget the last dark/flat, so later scans link none until a new one is taken. + + Parameters + ---------- + RE : RunEngine | None, optional + The RunEngine to update. If None, it is taken from the IPython namespace. + """ + RE = ensure_available(RunEngine, RE=RE) + for key in DARK_FLAT_MD_KEYS: + RE.md.pop(key, None) + yield from bps.null() + + def reset_scan_id(new_scan_id: int = 0, RE: RunEngine | None = None): - """Reset the RunEngine's scan ID. + """Reset the RunEngine's scan ID, and forget the last dark/flat. Parameters ---------- @@ -600,4 +618,5 @@ def reset_scan_id(new_scan_id: int = 0, RE: RunEngine | None = None): """ RE = ensure_available(RunEngine, RE=RE) RE.md["scan_id"] = new_scan_id - yield from bps.null() + # Earlier dark/flats belong to the scans being numbered over. + yield from clear_last_dark_flat(RE=RE) diff --git a/tests/test_dark_flat_proposal.py b/tests/test_dark_flat_proposal.py new file mode 100644 index 0000000..68a0003 --- /dev/null +++ b/tests/test_dark_flat_proposal.py @@ -0,0 +1,100 @@ +"""A tomography run's dark/flat must come from its own proposal and cycle. + +RE.md is shared by the whole beamline and keeps ``last_dark_flat_uid`` across a +proposal switch. At HEX on 2026-10-22, tomo scans 249-256 in pass-312143 were +exported with the dark/flat of scan 246 from pass-314022. +""" + +from unittest.mock import Mock + +from bluesky import RunEngine +from bluesky_tiled_plugins.clients.bluesky_run import BlueskyRunV3 + +from hextools.exporters.tomography import resolve_dark_flat_run +from hextools.utils import DARK_FLAT_MD_KEYS, clear_last_dark_flat, reset_scan_id + + +def _run(data_session: str, cycle: str = "2026-3") -> Mock: + run = Mock(spec=BlueskyRunV3) + run.start = { + "uid": f"uid-{data_session}", + "data_session": data_session, + "cycle": cycle, + } + return run + + +def test_clear_last_dark_flat(RE: RunEngine): + RE.md.update({"last_dark_flat_uid": "uid-246", "last_dark_flat_scan_id": 246}) + RE(clear_last_dark_flat(RE=RE)) + assert not set(DARK_FLAT_MD_KEYS) & RE.md.keys() + RE(clear_last_dark_flat(RE=RE)) # nothing left to clear is fine + + +def test_reset_scan_id_forgets_the_last_dark_flat(RE: RunEngine): + RE.md.update({"scan_id": 256, "last_dark_flat_uid": "uid-246"}) + RE(reset_scan_id(RE=RE)) + assert RE.md["scan_id"] == 0 + assert not set(DARK_FLAT_MD_KEYS) & RE.md.keys() + + +def test_dark_flat_from_another_proposal_is_not_used(): + client = {"uid-246": _run("pass-314022")} + start = { + "data_session": "pass-312143", + "cycle": "2026-3", + "last_dark_flat_uid": "uid-246", + } + assert resolve_dark_flat_run(client, start) is None + + +def test_dark_flat_from_an_earlier_cycle_is_not_used(): + client = {"uid-246": _run("pass-312143", cycle="2026-2")} + start = { + "data_session": "pass-312143", + "cycle": "2026-3", + "last_dark_flat_uid": "uid-246", + } + assert resolve_dark_flat_run(client, start) is None + + +def test_dark_flat_from_the_same_proposal_is_used(): + dark_flat = _run("pass-312143") + client = {"uid-260": dark_flat} + start = { + "data_session": "pass-312143", + "cycle": "2026-3", + "last_dark_flat_uid": "uid-260", + } + assert resolve_dark_flat_run(client, start) is dark_flat + + +def test_unresolvable_dark_flat_is_skipped_not_raised(): + start = {"data_session": "pass-312143", "last_dark_flat_uid": "gone"} + assert resolve_dark_flat_run({}, start) is None + + +def test_no_dark_flat_recorded(): + assert resolve_dark_flat_run({}, {"data_session": "pass-312143"}) is None + + +def test_explicit_dark_flat_id_from_another_proposal_is_not_used(): + client = {"uid-246": _run("pass-314022")} + start = {"data_session": "pass-312143", "cycle": "2026-3"} + assert resolve_dark_flat_run(client, start, dark_flat_id="uid-246") is None + + +def test_explicit_dark_flat_id_overrides_the_start_document(): + dark_flat = _run("pass-312143") + client = {"uid-260": dark_flat, "uid-246": _run("pass-314022")} + start = { + "data_session": "pass-312143", + "cycle": "2026-3", + "last_dark_flat_uid": "uid-246", + } + assert resolve_dark_flat_run(client, start, dark_flat_id="uid-260") is dark_flat + + +def test_unresolvable_explicit_dark_flat_id_is_skipped_not_raised(): + start = {"data_session": "pass-312143", "cycle": "2026-3"} + assert resolve_dark_flat_run({}, start, dark_flat_id="gone") is None