From abc2a3d0415071a805e7d60cf66387762abd68ec Mon Sep 17 00:00:00 2001 From: x01dc Date: Mon, 20 Jul 2026 22:00:22 +0200 Subject: [PATCH] feat/tomo queue reacquire from an earlier projection Add tomo_queue_reacquire(job_index, projection_number) to the shared TomoQueueMixin: reopens any queue job (even an already-"done" one) at an earlier projection than its own recorded progress, and resets every job queued after it -- tomo and command jobs alike -- to "pending", so a following tomo_queue_resume() (a new alias for tomo_queue_execute(), added for naming symmetry with tomo_scan_resume()) re-runs everything from there forward in order. Reuses the existing tomo_scan_resume() machinery by writing the requested resume point into the shared progress global var, rather than adding new parameters to tomo_scan()/tomo_queue_execute(). Guards against clobbering another job's progress if one is already "incomplete"/"running" elsewhere in the queue. A flat projection number is used as the resume-point unit for all three tomo_types, including type 1 (8 equally-spaced sub-tomograms), via a new per-setup _resolve_type1_projection() that maps it to (subtomo_start, start_angle). flomni's version reuses its existing _subtomo_angle_plan(); lamni's inline angle math in sub_tomo_scan() was extracted into an analogous _subtomo_angle_plan() static method first (behavior-preserving -- verified against the existing angle-math test suite) so both setups share the same pattern. Co-Authored-By: Claude Sonnet 5 --- .../bec_ipython_client/plugins/LamNI/lamni.py | 51 +++- .../plugins/OMNY_shared/tomo_queue_mixin.py | 145 +++++++++++ .../plugins/flomni/flomni.py | 24 ++ .../test_tomo_queue_reacquire.py | 244 ++++++++++++++++++ 4 files changed, 460 insertions(+), 4 deletions(-) create mode 100644 tests/tests_bec_ipython_client/test_tomo_queue_reacquire.py diff --git a/csaxs_bec/bec_ipython_client/plugins/LamNI/lamni.py b/csaxs_bec/bec_ipython_client/plugins/LamNI/lamni.py index fcb659e..0e4b28e 100644 --- a/csaxs_bec/bec_ipython_client/plugins/LamNI/lamni.py +++ b/csaxs_bec/bec_ipython_client/plugins/LamNI/lamni.py @@ -859,11 +859,22 @@ class LamNI(TomoQueueMixin, LamNIAlignmentMixin, LamNIOpticsMixin, LamniGuiTools step = 360.0 / N return N, step, N * 8 - def sub_tomo_scan(self, subtomo_number, start_angle=None): - """Perform one sub-tomogram (tomo_type 1 only).""" - self._write_subtomo_to_scilog(subtomo_number) + @staticmethod + def _subtomo_angle_plan(subtomo_number, tomo_angle_stepsize, start_angle=None): + """Pure angle-generation logic for one sub-tomogram (tomo_type 1 + only, 360-degree span). No device I/O / progress side effects -- + kept separate from sub_tomo_scan() so the angle math is directly + unit-testable and reusable (e.g. by tomo_queue_reacquire()'s + projection-number resolver). Mirrors Flomni._subtomo_angle_plan() + (no tomo_angle_range/forward-reverse split -- lamni always scans + the full 360 degrees). - N, achievable_step, total_projections = self._tomo_type1_actual_grid() + Returns: + angles (np.ndarray): the N angles (degrees) for this sub-tomogram. + N (int): number of projections in this (and every) sub-tomogram. + """ + N = int(360.0 / tomo_angle_stepsize) + achievable_step = 360.0 / N if start_angle is None: # Phase offset must be a fraction of the ACHIEVABLE step (after @@ -886,6 +897,38 @@ class LamNI(TomoQueueMixin, LamNIAlignmentMixin, LamNIOpticsMixin, LamniGuiTools if not (subtomo_number % 2): angles = np.flip(angles) + return angles, N + + def _resolve_type1_projection(self, projection_number: int) -> tuple[int, float]: + """Map a flat 0-indexed tomo_type-1 projection number to + (subtomo_start, start_angle), for tomo_queue_reacquire(). + + The mapping mirrors exactly how sub_tomo_scan() itself numbers + projections (self.progress["projection"] = (subtomo_number - 1) * N + + subtomo_projection), so a projection_number read off e.g. the + tomography_scannumbers.txt log resolves to the same angle + sub_tomo_scan() would have produced at that position. + """ + _, N = self._subtomo_angle_plan(1, self.tomo_angle_stepsize) + total = N * 8 + if not 0 <= projection_number < total: + raise ValueError( + f"projection_number must be in 0..{total - 1} for this job's grid " + f"(tomo_angle_stepsize={self.tomo_angle_stepsize})." + ) + subtomo_number = projection_number // N + 1 + angles, _ = self._subtomo_angle_plan(subtomo_number, self.tomo_angle_stepsize) + return subtomo_number, float(angles[projection_number % N]) + + def sub_tomo_scan(self, subtomo_number, start_angle=None): + """Perform one sub-tomogram (tomo_type 1 only).""" + self._write_subtomo_to_scilog(subtomo_number) + + angles, N = self._subtomo_angle_plan( + subtomo_number, self.tomo_angle_stepsize, start_angle=start_angle + ) + total_projections = N * 8 + for angle in angles: self.progress["tomo_type"] = "Equally spaced sub-tomograms" self.progress["subtomo"] = subtomo_number diff --git a/csaxs_bec/bec_ipython_client/plugins/OMNY_shared/tomo_queue_mixin.py b/csaxs_bec/bec_ipython_client/plugins/OMNY_shared/tomo_queue_mixin.py index 6465b6a..2740e56 100644 --- a/csaxs_bec/bec_ipython_client/plugins/OMNY_shared/tomo_queue_mixin.py +++ b/csaxs_bec/bec_ipython_client/plugins/OMNY_shared/tomo_queue_mixin.py @@ -26,6 +26,11 @@ Consumers (e.g. ``Flomni``, ``LamNI``) must: the same way, but is setup-specific enough -- e.g. flomni's legacy ``flomni_at_each_angle`` builtins hook, single-point acquisition -- that it stays out of this mixin) + - implement their own ``_resolve_type1_projection()`` (dispatched + generically by ``tomo_queue_reacquire()``), mapping a flat tomo_type-1 + projection number to a ``(subtomo_start, start_angle)`` pair -- the + angle math is setup-specific (see each setup's own + ``_subtomo_angle_plan()``), same reasoning as ``tomo_scan()`` above """ from __future__ import annotations @@ -756,3 +761,143 @@ class TomoQueueMixin: spec = self._TOMO_QUEUE_ACTIONS[step["action"]] method = getattr(self, spec["func"]) method(**step["kwargs"]) + + # ── reacquire-from-an-earlier-projection ──────────────────────────────── + + def _resolve_type1_projection(self, projection_number: int) -> tuple[int, float]: + """Map a flat 0-indexed tomo_type-1 projection number to + (subtomo_start, start_angle), for tomo_queue_reacquire(). + + Subclasses must override this -- the angle math is setup-specific, + see each setup's own ``_subtomo_angle_plan()``. Not needed for + tomo_types 2/3, where a projection number is already the native + resume unit (see ``tomo_scan_resume()``). + """ + raise NotImplementedError( + f"{type(self).__name__} must implement _resolve_type1_projection() to use " + "tomo_queue_reacquire() on a tomo_type-1 job." + ) + + def tomo_queue_resume(self, start_index: int = 0) -> None: + """Continue running the tomo queue from wherever it currently + stands -- an alias for ``tomo_queue_execute()``, added purely so + the queue has a command name that mirrors ``tomo_scan_resume()``'s. + The underlying pick-next mechanism (resume any incomplete/running + job first, else the next pending one) is identical either way. + """ + self.tomo_queue_execute(start_index=start_index) + + def tomo_queue_reacquire(self, job_index: int, projection_number: int) -> None: + """Reopen tomo queue job #``job_index`` at an earlier projection + than wherever it currently stands -- including a job that's + already "done" -- and reset every job queued after it (tomo and + command jobs alike, regardless of their own current status) to + "pending", so a following ``tomo_queue_resume()``/ + ``tomo_queue_execute()`` call re-runs everything from here + forward, in order. + + This does NOT run anything itself. It only: + - restores ``job_index``'s own snapshotted params onto the live + properties (so ``tomo_type`` etc. reflect the job being + reopened, not whatever ran most recently) + - writes the requested resume point into the shared ``progress`` + global var exactly as if a crash had just occurred there, and + resets its timing bookkeeping (``tomo_start_time``/ + ``accumulated_idle_time``/``heartbeat``/ETA) the same way + ``tomo_scan()``'s own "new scan" branch does -- this is a + fresh acquisition attempt, not a continuation of whatever job + ran most recently + - flips ``job_index``'s own status to "incomplete" + + Call ``tomo_queue_resume()`` (or ``tomo_queue_execute()``) + afterward to actually re-acquire. + + Known limitation: does not re-register a new tomo_id/sample- + database entry for the reopened job -- it keeps whatever tomo_id + is currently live in the session, exactly like a normal + ``tomo_scan_resume()`` already does. + + Args: + job_index: Position of the job to reopen, in the *current* + queue order (same indexing as ``tomo_queue_show()``/ + ``tomo_queue_delete()``). + projection_number: Flat 0-indexed projection to resume from, + for any tomo_type (1, 2, or 3) -- see + ``_resolve_type1_projection()`` for how this maps onto a + tomo_type-1 job's (subtomo, angle). + + Raises: + TomoQueueError: ``job_index`` is out of range, the job at that + index is a command job (not a tomo scan), or some *other* + job in the queue is already "incomplete"/"running" (the + queue's single-resumable-job invariant -- resolve that one + first, or pass its own index). + ValueError: ``projection_number`` is out of range for this + job's own tomo scan parameters. + """ + jobs = self._tomo_queue_proxy.ensure_ids() + if not 0 <= job_index < len(jobs): + raise TomoQueueError(f"tomo_queue_reacquire: no job at index {job_index}.") + + job = jobs[job_index] + if job.get("kind", "tomo") != "tomo": + raise TomoQueueError( + f"tomo_queue_reacquire: job #{job_index} ('{job['label']}') is a " + "command job, not a tomo scan -- nothing to reacquire." + ) + + for idx, other in enumerate(jobs): + if idx != job_index and other.get("status") in ("incomplete", "running"): + raise TomoQueueError( + f"tomo_queue_reacquire: job #{idx} ('{other['label']}') is already " + f"'{other['status']}' -- resolve it first (or pass index {idx} " + "instead), since only one job's progress can be tracked at a time." + ) + + for name, value in job["params"].items(): + setattr(self, name, value) + + if self.tomo_type == 1: + subtomo_start, start_angle = self._resolve_type1_projection(projection_number) + self.progress["subtomo"] = subtomo_start + self.progress["angle"] = start_angle + elif self.tomo_type in (2, 3): + if projection_number < 0: + raise ValueError("tomo_queue_reacquire: projection_number must be >= 0.") + max_prj = job["params"].get("golden_max_number_of_projections", 0) or 0 + if max_prj > 0 and projection_number >= max_prj: + raise ValueError( + f"projection_number must be < {max_prj} " + "(golden_max_number_of_projections) for this job." + ) + self.progress["projection"] = projection_number + else: + raise TomoQueueError( + f"tomo_queue_reacquire: unknown tomo_type {self.tomo_type} for job " + f"#{job_index} ('{job['label']}')." + ) + + # Fresh acquisition attempt -- reset timing bookkeeping the same way + # tomo_scan()'s own "new scan" branch does, so the eventual + # end-of-scan summary reflects this attempt, not whatever job ran + # most recently. + self.progress["tomo_start_time"] = datetime.datetime.now().isoformat() + self.progress["estimated_remaining_time"] = None + self.progress["estimated_finish_time"] = None + self.progress["accumulated_idle_time"] = 0.0 + self.progress["heartbeat"] = None + + job_id = job["id"] + self._tomo_queue_proxy.update_by_id(job_id, status="incomplete") + + reset_count = 0 + for later_job in jobs[job_index + 1 :]: + if later_job.get("status") != "pending": + self._tomo_queue_proxy.update_by_id(later_job["id"], status="pending") + reset_count += 1 + + print( + f"Job #{job_index} ('{job['label']}') reopened at projection " + f"{projection_number}. {reset_count} later job(s) reset to pending. " + "Call tomo_queue_resume() to run." + ) diff --git a/csaxs_bec/bec_ipython_client/plugins/flomni/flomni.py b/csaxs_bec/bec_ipython_client/plugins/flomni/flomni.py index ed26e01..d5685c4 100644 --- a/csaxs_bec/bec_ipython_client/plugins/flomni/flomni.py +++ b/csaxs_bec/bec_ipython_client/plugins/flomni/flomni.py @@ -2272,6 +2272,30 @@ class Flomni( return angles, subtomo_offset, N, step + def _resolve_type1_projection(self, projection_number: int) -> tuple[int, float]: + """Map a flat 0-indexed tomo_type-1 projection number to + (subtomo_start, start_angle), for tomo_queue_reacquire(). + + The mapping mirrors exactly how sub_tomo_scan() itself numbers + projections (self.progress["projection"] = (subtomo_number - 1) * N + + subtomo_projection), so a projection_number read off e.g. the + tomography_scannumbers.txt log resolves to the same angle + sub_tomo_scan() would have produced at that position. + """ + _, _, N, _ = self._subtomo_angle_plan(1, self.tomo_angle_range, self.tomo_angle_stepsize) + total = N * 8 + if not 0 <= projection_number < total: + raise ValueError( + f"projection_number must be in 0..{total - 1} for this job's grid " + f"(tomo_angle_range={self.tomo_angle_range}, " + f"tomo_angle_stepsize={self.tomo_angle_stepsize})." + ) + subtomo_number = projection_number // N + 1 + angles, _, _, _ = self._subtomo_angle_plan( + subtomo_number, self.tomo_angle_range, self.tomo_angle_stepsize + ) + return subtomo_number, float(angles[projection_number % N]) + def sub_tomo_scan(self, subtomo_number, start_angle=None): """ Performs a sub tomogram scan. diff --git a/tests/tests_bec_ipython_client/test_tomo_queue_reacquire.py b/tests/tests_bec_ipython_client/test_tomo_queue_reacquire.py new file mode 100644 index 0000000..7fcaed6 --- /dev/null +++ b/tests/tests_bec_ipython_client/test_tomo_queue_reacquire.py @@ -0,0 +1,244 @@ +"""Tests for TomoQueueMixin.tomo_queue_reacquire()/tomo_queue_resume() and +the per-setup _resolve_type1_projection() resolvers it depends on (see +csaxs_bec/bec_ipython_client/plugins/OMNY_shared/tomo_queue_mixin.py). + +Round-trip tests build a bare LamNI/Flomni instance (bypassing __init__'s +heavy side effects, same pattern as test_lamni_tomo_queue.py) and confirm +_resolve_type1_projection() agrees with what sub_tomo_scan() itself would +actually produce at a given flat projection number -- not a re-derivation +of the same formula, but a comparison against the real scan-execution path. +""" + +import types + +import pytest + +import csaxs_bec.bec_ipython_client.plugins.LamNI.lamni as lamni_module +from csaxs_bec.bec_ipython_client.plugins.flomni.flomni import Flomni +from csaxs_bec.bec_ipython_client.plugins.flomni.flomni import _ProgressProxy as FlomniProgressProxy +from csaxs_bec.bec_ipython_client.plugins.LamNI.lamni import LamNI, _ProgressProxy +from csaxs_bec.bec_ipython_client.plugins.OMNY_shared.tomo_queue_mixin import TomoQueueError + +STEPSIZES = [10.0, 7.0, 25.0, 12.5] + + +class FakeClient: + """Minimal in-memory stand-in for BEC's global-var store.""" + + def __init__(self): + self._vars = {} + + def get_global_var(self, key): + return self._vars.get(key) + + def set_global_var(self, key, value): + self._vars[key] = value + + +class FakeOMNYTools: + def yesno(self, *args, **kwargs): + return True + + +class _FakeRtx: + user_parameter = {"large_range_scan": True} + + +class _FakeDev: + rtx = _FakeRtx() + + def __contains__(self, name): + return hasattr(self, name) + + def __getitem__(self, name): + return getattr(self, name) + + +def make_lamni(): + lamni_module.dev = _FakeDev() + obj = object.__new__(LamNI) + obj.client = FakeClient() + obj.OMNYTools = FakeOMNYTools() + obj.reconstructor = types.SimpleNamespace(folder_name=None) + obj.tomo_id = -1 + obj._progress_proxy = _ProgressProxy(obj.client) + obj._init_tomo_queue() + return obj + + +def make_flomni(): + obj = object.__new__(Flomni) + obj.client = FakeClient() + obj.OMNYTools = FakeOMNYTools() + obj._progress_proxy = FlomniProgressProxy(obj.client) + obj._init_tomo_queue() + return obj + + +# ── _resolve_type1_projection round-trip: lamni ───────────────────────────── + + +@pytest.mark.parametrize("stepsize", STEPSIZES) +def test_lamni_resolve_type1_projection_matches_sub_tomo_scan(stepsize): + lamni = make_lamni() + lamni.tomo_angle_stepsize = stepsize + lamni._write_subtomo_to_scilog = lambda subtomo_number: None + recorded = {} + lamni._tomo_scan_at_angle = lambda angle, subtomo_number: recorded.__setitem__( + lamni.progress["projection"], (subtomo_number, angle) + ) + for subtomo in range(1, 9): + lamni.sub_tomo_scan(subtomo) + + assert len(recorded) > 0 + for projection_number, (expected_subtomo, expected_angle) in recorded.items(): + subtomo_start, start_angle = lamni._resolve_type1_projection(projection_number) + assert subtomo_start == expected_subtomo + assert start_angle == pytest.approx(expected_angle) + + +def test_lamni_resolve_type1_projection_out_of_range(): + lamni = make_lamni() + lamni.tomo_angle_stepsize = 10.0 + with pytest.raises(ValueError): + lamni._resolve_type1_projection(-1) + with pytest.raises(ValueError): + lamni._resolve_type1_projection(10_000) + + +# ── _resolve_type1_projection round-trip: flomni ──────────────────────────── + + +@pytest.mark.parametrize("stepsize", STEPSIZES) +@pytest.mark.parametrize("angle_range", [180, 360]) +def test_flomni_resolve_type1_projection_matches_sub_tomo_scan(stepsize, angle_range): + flomni = make_flomni() + flomni.tomo_angle_range = angle_range + flomni.tomo_angle_stepsize = stepsize + recorded = {} + flomni._tomo_scan_at_angle = lambda angle, subtomo_number: recorded.__setitem__( + flomni.progress["projection"], (subtomo_number, angle) + ) + for subtomo in range(1, 9): + flomni.sub_tomo_scan(subtomo) + + assert len(recorded) > 0 + for projection_number, (expected_subtomo, expected_angle) in recorded.items(): + subtomo_start, start_angle = flomni._resolve_type1_projection(projection_number) + assert subtomo_start == expected_subtomo + assert start_angle == pytest.approx(expected_angle) + + +def test_flomni_resolve_type1_projection_out_of_range(): + flomni = make_flomni() + flomni.tomo_angle_range = 180 + flomni.tomo_angle_stepsize = 10.0 + with pytest.raises(ValueError): + flomni._resolve_type1_projection(-1) + with pytest.raises(ValueError): + flomni._resolve_type1_projection(10_000) + + +# ── tomo_queue_reacquire() / tomo_queue_resume() mechanics (lamni) ────────── + + +def test_tomo_queue_resume_is_an_alias_for_execute(): + lamni = make_lamni() + calls = [] + lamni.tomo_queue_execute = lambda start_index=0: calls.append(start_index) + + lamni.tomo_queue_resume() + lamni.tomo_queue_resume(start_index=2) + + assert calls == [0, 2] + + +def test_tomo_queue_reacquire_type1_sets_progress_and_status(): + lamni = make_lamni() + lamni.tomo_type = 1 + lamni.tomo_angle_stepsize = 10.0 + index = lamni.tomo_queue_add(label="job0") + lamni._tomo_queue_proxy.update(index, status="done") + + lamni.tomo_queue_reacquire(index, projection_number=5) + + job = lamni._tomo_queue_proxy.as_list()[index] + assert job["status"] == "incomplete" + subtomo_start, start_angle = lamni._resolve_type1_projection(5) + assert lamni.progress["subtomo"] == subtomo_start + assert lamni.progress["angle"] == pytest.approx(start_angle) + assert lamni.progress["tomo_start_time"] is not None + assert lamni.progress["accumulated_idle_time"] == 0.0 + + +def test_tomo_queue_reacquire_type23_sets_projection_directly(): + lamni = make_lamni() + lamni.tomo_type = 2 + lamni.golden_max_number_of_projections = 0 + index = lamni.tomo_queue_add(label="job0") + lamni._tomo_queue_proxy.update(index, status="done") + + lamni.tomo_queue_reacquire(index, projection_number=17) + + assert lamni.progress["projection"] == 17 + job = lamni._tomo_queue_proxy.as_list()[index] + assert job["status"] == "incomplete" + + +def test_tomo_queue_reacquire_resets_later_jobs_to_pending_tomo_and_command(): + lamni = make_lamni() + lamni.tomo_type = 1 + lamni.tomo_angle_stepsize = 10.0 + + i0 = lamni.tomo_queue_add(label="job0") + i1 = lamni.tomo_queue_add(label="job1") + lamni.tomo_queue_add_command({"action": "move", "kwargs": {"positions": {"mokev": 6.2}}}) + i3 = lamni.tomo_queue_add(label="job3") + + for i in (i0, i1, i3): + lamni._tomo_queue_proxy.update(i, status="done") + + lamni.tomo_queue_reacquire(i0, projection_number=0) + + jobs = lamni._tomo_queue_proxy.as_list() + assert jobs[i0]["status"] == "incomplete" + for later in jobs[i0 + 1 :]: + assert later["status"] == "pending" + + +def test_tomo_queue_reacquire_rejects_out_of_range_job_index(): + lamni = make_lamni() + with pytest.raises(TomoQueueError): + lamni.tomo_queue_reacquire(0, projection_number=0) + + +def test_tomo_queue_reacquire_rejects_command_job(): + lamni = make_lamni() + lamni.tomo_queue_add_command({"action": "move", "kwargs": {"positions": {"mokev": 6.2}}}) + with pytest.raises(TomoQueueError): + lamni.tomo_queue_reacquire(0, projection_number=0) + + +def test_tomo_queue_reacquire_rejects_when_another_job_is_incomplete(): + lamni = make_lamni() + lamni.tomo_type = 1 + lamni.tomo_angle_stepsize = 10.0 + i0 = lamni.tomo_queue_add(label="job0") + i1 = lamni.tomo_queue_add(label="job1") + lamni._tomo_queue_proxy.update(i1, status="incomplete") + + with pytest.raises(TomoQueueError): + lamni.tomo_queue_reacquire(i0, projection_number=0) + + +def test_tomo_queue_reacquire_allows_targeting_the_incomplete_job_itself(): + lamni = make_lamni() + lamni.tomo_type = 1 + lamni.tomo_angle_stepsize = 10.0 + index = lamni.tomo_queue_add(label="job0") + lamni._tomo_queue_proxy.update(index, status="incomplete") + + lamni.tomo_queue_reacquire(index, projection_number=0) + + job = lamni._tomo_queue_proxy.as_list()[index] + assert job["status"] == "incomplete"