From e6928f295e7cc3be67c14f77cbbef0eb3a7e9728 Mon Sep 17 00:00:00 2001 From: x01dc Date: Tue, 1 Sep 2026 14:26:10 +0200 Subject: [PATCH] feat(omny): add 4/8 sub-tomogram tomo modes, hide advanced modes from regular users - New shared static helper _equally_spaced_subtomo_angles(), ported from flomni's _subtomo_angle_plan() with the 180/360-degree-range branch removed (OMNY only ever scans 180 degrees). Generalizes the bit-reversal sub-tomogram interleaving to any power-of-two sub-tomogram count; for n_subtomos=2 it reduces to OMNY's original type-1 scheme. - sub_tomo_scan() refactored onto the new helper (gains an n_subtomos param); tomo_type 1 (2 sub-tomograms) is otherwise unchanged. - Adds tomo_type 4 (4 sub-tomograms) and 5 (8 sub-tomograms), gated by the same "wear in OMNY" authorization-code prompt already used for tomo_type 2/3. - Advanced tomo modes (2, 3, 4, 5) are now hidden from tomo_parameters()'s menu unless bec.active_account is in the new _SUPERUSER_ACCOUNTS allowlist; a non-superuser landing on an advanced type is forced back to type 1. _SUPERUSER_ACCOUNTS is left empty -- needs real account codes before anyone can use the advanced modes. - As a side effect of reusing flomni's formula, also fixes a latent bug: the old code derived the inter-sub-tomogram phase offset from the raw configured tomo_angle_stepsize rather than the achievable (int-truncated) step, producing an unevenly-spaced combined grid whenever stepsize didn't divide 180 evenly. - New tests/tests_bec_ipython_client/test_omny_tomo_angles.py: angle-count and no-duplicate-angle checks for n_subtomos in (2, 4, 8), a regression check against the old type-1 formula, an explicit test of the achievable-step fix, a cross-check against flomni's hardcoded N=8 bit-reversal table, and superuser-gating tests (including the full tomo_parameters() interactive flow forcing a non-superuser back to type 1, and a superuser successfully unlocking type 5). --- .../bec_ipython_client/plugins/omny/omny.py | 165 ++++++++--- .../test_omny_tomo_angles.py | 259 ++++++++++++++++++ 2 files changed, 390 insertions(+), 34 deletions(-) create mode 100644 tests/tests_bec_ipython_client/test_omny_tomo_angles.py diff --git a/csaxs_bec/bec_ipython_client/plugins/omny/omny.py b/csaxs_bec/bec_ipython_client/plugins/omny/omny.py index 2a987d75..a11c7bce 100644 --- a/csaxs_bec/bec_ipython_client/plugins/omny/omny.py +++ b/csaxs_bec/bec_ipython_client/plugins/omny/omny.py @@ -29,9 +29,11 @@ if builtins.__dict__.get("bec") is not None: dev = builtins.__dict__.get("dev") scans = builtins.__dict__.get("scans") + def umv(*args): return scans.umv(*args, relative=False) + class OMNYInitError(Exception): pass @@ -40,6 +42,13 @@ class OMNYError(Exception): pass +# BEC accounts (bec.active_account, e.g. "e12345") allowed to select advanced tomo +# modes (golden ratio, 4/8 sub-tomogram) in OMNY.tomo_parameters(). Regular users only +# ever see "1: 2 equally spaced sub-tomograms". Populate with real beamline-staff +# account codes -- left empty until populated, so nobody sees the advanced options yet. +_SUPERUSER_ACCOUNTS: frozenset[str] = frozenset({}) + + class OMNYInitStagesMixin: def enable_all_devices(self): @@ -744,6 +753,12 @@ class OMNY( elif val == 3: # equally spaced tomography with starting angles shifted by golden ratio self.client.set_global_var("tomo_type", val) + elif val == 4: + # equally spaced tomography with 4 sub tomograms + self.client.set_global_var("tomo_type", val) + elif val == 5: + # equally spaced tomography with 8 sub tomograms + self.client.set_global_var("tomo_type", val) else: raise ValueError("Unknown tomo_type.") @@ -925,41 +940,71 @@ class OMNY( tags, ) - def sub_tomo_scan(self, subtomo_number, start_angle=None): + @staticmethod + def _equally_spaced_subtomo_angles( + subtomo_number, n_subtomos, tomo_angle_stepsize, start_angle=None + ): + """ + Compute the projection angles for one sub-tomogram of an N-way, bit-reversal + interleaved, equally spaced 180 degree tomogram (n_subtomos a power of two: + 2, 4, or 8). Ported from flomni's ``_subtomo_angle_plan()`` + (flomni/flomni.py), with the 180-vs-360-degree-range branch removed since + OMNY only ever scans 180 degrees. + + Odd sub-tomograms scan forward (ascending), even sub-tomograms scan in + reverse (descending, to minimize travel between sub-tomograms). Each + sub-tomogram's phase offset within a step is picked via a bit-reversal + permutation of 0..n_subtomos-1, so that successive sub-tomograms maximally + subdivide the angular gaps left by the previous ones. For n_subtomos=2 this + reduces to phases {0, step/2} -- the original OMNY scheme. + + Always returns exactly ``int(180 / tomo_angle_stepsize)`` angles (no + redundant 0/180 duplicate point), matching the convention already used by + flomni's and lamni's own sub_tomo_scan(). + + Args: + subtomo_number (int): 1-indexed sub-tomogram number. + n_subtomos (int): Total number of sub-tomograms (2, 4, or 8). + tomo_angle_stepsize (float): Angular step size in degrees. + start_angle (float, optional): Explicit start angle override. Defaults + to the bit-reversal phase offset described above. + """ + n_points = int(180.0 / tomo_angle_stepsize) + step = 180.0 / n_points + + bits = n_subtomos.bit_length() - 1 + subtomo_index = subtomo_number - 1 + phase_index = int(format(subtomo_index, f"0{bits}b")[::-1], 2) if bits else 0 + phase = step / n_subtomos * phase_index + + forward = bool(subtomo_number % 2) + if start_angle is None: + start_angle = phase if forward else (180.0 - step + phase) + + if forward: + return start_angle + np.arange(n_points) * step + return start_angle - np.arange(n_points) * step + + def sub_tomo_scan(self, subtomo_number, n_subtomos=2, start_angle=None): """ Performs a sub tomogram scan. Args: subtomo_number (int): The sub tomogram number. + n_subtomos (int): Total number of sub-tomograms in this scan (2, 4, or 8). start_angle (float, optional): The start angle of the scan. Defaults to None. """ self._write_subtomo_to_scilog(subtomo_number) - if start_angle is None: - if subtomo_number == 1: - start_angle = 0 - elif subtomo_number == 2: - start_angle = self.tomo_angle_stepsize / 2.0 - - # _tomo_shift_angles (potential global variable) - _tomo_shift_angles = 0 - angle_end = start_angle + 180 - angles = np.linspace( - start_angle + _tomo_shift_angles, - angle_end, - num=int(180 / self.tomo_angle_stepsize) + 1, - endpoint=True, + angles = self._equally_spaced_subtomo_angles( + subtomo_number, n_subtomos, self.tomo_angle_stepsize, start_angle ) - # reverse even sub-tomogram - if not (subtomo_number % 2): - angles = np.flip(angles) - for angle in angles: + n_points = len(angles) + for idx, angle in enumerate(angles): self.progress["subtomo"] = subtomo_number - self.progress["subtomo_projection"] = np.where(angles == angle)[0][0] - self.progress["subtomo_total_projections"] = 180 / self.tomo_angle_stepsize - self.progress["projection"] = (subtomo_number - 1) * self.progress[ - "subtomo_total_projections" - ] + self.progress["subtomo_projection"] - self.progress["total_projections"] = 180 / self.tomo_angle_stepsize * 2 + self.progress["subtomo_projection"] = idx + self.progress["subtomo_total_projections"] = n_points + self.progress["projection"] = (subtomo_number - 1) * n_points + idx + self.progress["total_projections"] = n_points * n_subtomos self.progress["angle"] = angle self._tomo_scan_at_angle(angle, subtomo_number) @@ -1009,7 +1054,7 @@ class OMNY( self._current_special_angles = self.special_angles.copy() # a new tomo scan was started if ( - (self.tomo_type == 1 and subtomo_start == 1 and start_angle is None) + (self.tomo_type in (1, 4, 5) and subtomo_start == 1 and start_angle is None) or (self.tomo_type == 2 and projection_number == None) or (self.tomo_type == 3 and projection_number == None) ): @@ -1030,11 +1075,12 @@ class OMNY( self.tomo_id = 0 with scans.dataset_id_on_hold: - if self.tomo_type == 1: - # 2 equally spaced sub-tomograms - self.progress["tomo_type"] = "Equally spaced sub-tomograms" - for ii in range(subtomo_start, 3): - self.sub_tomo_scan(ii, start_angle=start_angle) + if self.tomo_type in (1, 4, 5): + # equally spaced sub-tomograms (2, 4, or 8) + n_subtomos = {1: 2, 4: 4, 5: 8}[self.tomo_type] + self.progress["tomo_type"] = f"Equally spaced sub-tomograms ({n_subtomos})" + for ii in range(subtomo_start, n_subtomos + 1): + self.sub_tomo_scan(ii, n_subtomos=n_subtomos, start_angle=start_angle) start_angle = None elif self.tomo_type == 2: @@ -1289,6 +1335,17 @@ class OMNY( corridor_size=corridor_size, ) + def _is_superuser_account(self) -> bool: + """Whether the active BEC account is allowed to select advanced tomo modes.""" + bec = builtins.__dict__.get("bec") + try: + account = bec.active_account + except Exception: + return False + if isinstance(account, bytes): + account = account.decode() + return bool(account) and account in _SUPERUSER_ACCOUNTS + def tomo_parameters(self): """print and update the tomo parameters""" print("Current settings:") @@ -1329,6 +1386,13 @@ class OMNY( print( "repeating prjections at 0 degrees at the beginning of every second subtomogram" ) + if self.tomo_type in (4, 5): + n_subtomos = {4: 4, 5: 8}[self.tomo_type] + print( + f"\x1b[1mTomo type {self.tomo_type}:\x1b[0m {n_subtomos} equally spaced sub-tomograms" + ) + print(f"Total number of projections: {180/self.tomo_angle_stepsize*n_subtomos}") + print(f"Angular step within sub-tomogram: {self.tomo_angle_stepsize} degrees") print(f"\nSample name: {self.sample_name}\n") if self.OMNYTools.yesno("Are these parameters correctly set for your scan?", "y"): @@ -1344,12 +1408,23 @@ class OMNY( "Reconstruction queue ", self.ptycho_reconstruct_foldername, str ) + is_superuser = self._is_superuser_account() print("Tomography type:") print(" 1: 2 equally spaced sub-tomograms") - print(" 2: Golden ratio tomography") - print(" 3: Equally spaced tomography, golden ratio starting angle") + if is_superuser: + print(" 2: Golden ratio tomography") + print(" 3: Equally spaced tomography, golden ratio starting angle") + print(" 4: 4 equally spaced sub-tomograms") + print(" 5: 8 equally spaced sub-tomograms") self.tomo_type = self._get_val("Tomography type", self.tomo_type, int) + if self.tomo_type != 1 and not is_superuser: + print( + "Advanced tomo modes are only available to superuser accounts. Tomo type 1" + " selected: 2 sub-tomograms selected." + ) + self.tomo_type = 1 + if self.tomo_type == 1: tomo_numberofprojections = self._get_val( "Total number of projections", 180 / self.tomo_angle_stepsize * 2, int @@ -1417,6 +1492,28 @@ class OMNY( ) self.tomo_type = 1 + if self.tomo_type in (4, 5): + n_subtomos = {4: 4, 5: 8}[self.tomo_type] + code = self._get_val( + "This mode causes significant wear in OMNY. Enter authorization code.", 0, str + ) + if code == "x12sa": + tomo_numberofprojections = self._get_val( + "Total number of projections", + 180 / self.tomo_angle_stepsize * n_subtomos, + int, + ) + print(f"The angular step will be {180/tomo_numberofprojections}") + self.tomo_angle_stepsize = 180 / tomo_numberofprojections * n_subtomos + print( + f"The angular step in a subtomogram it will be {self.tomo_angle_stepsize}" + ) + else: + print( + "Wrong authorization code. Tomo type 1 selected: 2 sub-tomograms selected." + ) + self.tomo_type = 1 + @staticmethod def _get_val(msg: str, default_value, data_type): return data_type(input(f"{msg} ({default_value}): ") or default_value) @@ -1507,7 +1604,7 @@ class OMNY( def dewar_show_all(): dev.omny_dewar.show_all() - + if __name__ == "__main__": import builtins diff --git a/tests/tests_bec_ipython_client/test_omny_tomo_angles.py b/tests/tests_bec_ipython_client/test_omny_tomo_angles.py new file mode 100644 index 00000000..c1c25255 --- /dev/null +++ b/tests/tests_bec_ipython_client/test_omny_tomo_angles.py @@ -0,0 +1,259 @@ +"""Tests for OMNY's equally-spaced sub-tomogram angle scheme and the +account-based superuser gating in tomo_parameters(). + +OMNY._equally_spaced_subtomo_angles() is a static method ported from flomni's +_subtomo_angle_plan() (flomni/flomni.py), with the 180-vs-360-degree-range +branch removed since OMNY only ever scans 180 degrees. It backs tomo_type 1 +(2 sub-tomograms, unchanged/pre-existing), 4 (4 sub-tomograms), and 5 (8 +sub-tomograms, new). +""" + +import builtins +import types + +import numpy as np +import pytest + +from csaxs_bec.bec_ipython_client.plugins.omny import omny as omny_module +from csaxs_bec.bec_ipython_client.plugins.omny.omny import OMNY + +STEPSIZES = [10.0, 7.0, 25.0, 5.0] + + +def _angles(subtomo_number, n_subtomos, stepsize): + return np.asarray(OMNY._equally_spaced_subtomo_angles(subtomo_number, n_subtomos, stepsize)) + + +def _old_omny_type1_angles(subtomo_number, stepsize): + """Reimplementation of OMNY's original (pre-refactor) type-1 angle logic, + for regression comparison against _equally_spaced_subtomo_angles(). Mirrors + the exact math that used to live inline in sub_tomo_scan(), including the + silent `_tomo_scan_at_angle()` filtering of angles outside [0, 180.05). + """ + if subtomo_number == 1: + start_angle = 0 + elif subtomo_number == 2: + start_angle = stepsize / 2.0 + else: + raise ValueError("old scheme only supports 2 sub-tomograms") + angle_end = start_angle + 180 + angles = np.linspace(start_angle, angle_end, num=int(180 / stepsize) + 1, endpoint=True) + if not (subtomo_number % 2): + angles = np.flip(angles) + return angles[(angles >= 0) & (angles < 180.05)] + + +@pytest.mark.parametrize("stepsize", STEPSIZES) +@pytest.mark.parametrize("n_subtomos", [2, 4, 8]) +def test_each_subtomo_generates_exactly_n_angles(stepsize, n_subtomos): + n_points = int(180.0 / stepsize) + for subtomo_number in range(1, n_subtomos + 1): + assert len(_angles(subtomo_number, n_subtomos, stepsize)) == n_points + + +@pytest.mark.parametrize("stepsize", STEPSIZES) +@pytest.mark.parametrize("n_subtomos", [2, 4, 8]) +def test_combined_subtomos_cover_180_degrees_without_duplicates(stepsize, n_subtomos): + n_points = int(180.0 / stepsize) + step = 180.0 / n_points + all_angles = [] + for subtomo_number in range(1, n_subtomos + 1): + all_angles.extend(_angles(subtomo_number, n_subtomos, stepsize).tolist()) + + assert len(all_angles) == n_points * n_subtomos + unique = {round(a, 6) for a in all_angles} + assert len(unique) == n_points * n_subtomos, "combined sub-tomograms must not repeat an angle" + + # each additional sub-tomogram doubles the resolution of the combined grid + combined_step = step / n_subtomos + sorted_angles = np.sort(np.asarray(all_angles)) + diffs = np.diff(sorted_angles) + assert np.allclose( + diffs, combined_step, atol=1e-6 + ), f"combined angles are not evenly spaced by {combined_step}: got {diffs}" + + +# stepsizes that divide 180 evenly, where the old code's raw-stepsize-based +# phase offset happens to coincide with the achievable-step-based one below +EXACT_STEPSIZES = [10.0, 5.0, 4.0, 20.0] + + +@pytest.mark.parametrize("subtomo_number", [1, 2]) +@pytest.mark.parametrize("stepsize", EXACT_STEPSIZES) +def test_n_subtomos_2_matches_omnys_original_scheme(subtomo_number, stepsize): + old = _old_omny_type1_angles(subtomo_number, stepsize) + new = _angles(subtomo_number, 2, stepsize) + if subtomo_number == 1: + # the old code additionally scanned a redundant point at exactly 180 + # deg (the same physical projection as 0 deg); the new implementation + # intentionally drops it, matching the "exactly N points, no + # redundant endpoint" convention already used by flomni's and lamni's + # own sub-tomogram scans. + np.testing.assert_allclose(new, old[:-1], atol=1e-9) + else: + np.testing.assert_allclose(new, old, atol=1e-9) + + +@pytest.mark.parametrize("stepsize", [7.0, 25.0]) +def test_phase_offset_uses_achievable_step_not_raw_stepsize(stepsize): + """For stepsizes that don't divide 180 evenly, the old code derived + subtomo 2's phase offset from the raw configured tomo_angle_stepsize + (stepsize/2), producing an unevenly-spaced combined grid when overlaid + with subtomo 1. flomni's _subtomo_angle_plan() fixes this by deriving the + phase offset from the ACHIEVABLE step (180/int(180/stepsize)) instead -- + this is the same fix, ported to OMNY. Regression-guards against + reintroducing the raw-stepsize version. + """ + n_points = int(180.0 / stepsize) + achievable_step = 180.0 / n_points + new = _angles(2, 2, stepsize) + + raw_stepsize_start = 180.0 - achievable_step + stepsize / 2.0 + achievable_step_start = 180.0 - achievable_step + achievable_step / 2.0 + + assert not np.isclose( + new[0], raw_stepsize_start, atol=1e-9 + ), "phase offset must not be derived from the raw configured tomo_angle_stepsize" + assert np.isclose(new[0], achievable_step_start, atol=1e-9) + + +def test_bit_reversal_matches_flomnis_hardcoded_phase_eighths_table(): + """flomni.py hardcodes the N=8 bit-reversal permutation as a literal dict + (phase_eighths = {1:0, 2:4, 3:2, 4:6, 5:1, 6:5, 7:3, 8:7}). The generalized + formula here must reproduce the same first-angle phase for n_subtomos=8.""" + phase_eighths = {1: 0, 2: 4, 3: 2, 4: 6, 5: 1, 6: 5, 7: 3, 8: 7} + stepsize = 10.0 + n_points = int(180.0 / stepsize) + step = 180.0 / n_points + for subtomo_number, expected_phase_index in phase_eighths.items(): + angles = _angles(subtomo_number, 8, stepsize) + forward = bool(subtomo_number % 2) + expected_phase = step / 8.0 * expected_phase_index + expected_start = expected_phase if forward else (180.0 - step + expected_phase) + assert np.isclose(angles[0], expected_start, atol=1e-9) + + +# --- superuser gating --------------------------------------------------------- + + +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 + + +def make_omny() -> OMNY: + """Bare OMNY instance with only what tomo_type/_is_superuser_account() need, + bypassing __init__'s heavy side effects.""" + obj = object.__new__(OMNY) + obj.client = FakeClient() + return obj + + +@pytest.fixture +def fake_bec(monkeypatch): + fake = types.SimpleNamespace(active_account="") + monkeypatch.setattr(builtins, "bec", fake, raising=False) + return fake + + +def test_is_superuser_account_true_for_allowlisted_account(monkeypatch, fake_bec): + monkeypatch.setattr(omny_module, "_SUPERUSER_ACCOUNTS", frozenset({"e99999"})) + fake_bec.active_account = "e99999" + omny = make_omny() + assert omny._is_superuser_account() is True + + +def test_is_superuser_account_false_for_other_account(monkeypatch, fake_bec): + monkeypatch.setattr(omny_module, "_SUPERUSER_ACCOUNTS", frozenset({"e99999"})) + fake_bec.active_account = "e00001" + omny = make_omny() + assert omny._is_superuser_account() is False + + +def test_is_superuser_account_handles_bytes_account(monkeypatch, fake_bec): + monkeypatch.setattr(omny_module, "_SUPERUSER_ACCOUNTS", frozenset({"e99999"})) + fake_bec.active_account = b"e99999" + omny = make_omny() + assert omny._is_superuser_account() is True + + +def test_is_superuser_account_false_when_empty_allowlist(monkeypatch, fake_bec): + monkeypatch.setattr(omny_module, "_SUPERUSER_ACCOUNTS", frozenset()) + fake_bec.active_account = "e99999" + omny = make_omny() + assert omny._is_superuser_account() is False + + +def test_tomo_parameters_forces_type_1_for_non_superuser(monkeypatch, fake_bec): + """A non-superuser account that ends up picking an advanced tomo_type via + the interactive prompt must be forced back to tomo_type 1.""" + monkeypatch.setattr(omny_module, "_SUPERUSER_ACCOUNTS", frozenset({"e99999"})) + fake_bec.active_account = "e00001" # not a superuser + omny = make_omny() + omny.tomo_countingtime = 1.0 + omny.tomo_shellstep = 0.1 + omny.fovx = 20 + omny.fovy = 20 + omny.stitch_x = 1 + omny.stitch_y = 1 + omny.ptycho_reconstruct_foldername = "queue" + omny.tomo_angle_stepsize = 10.0 + omny.tomo_type = 1 + fake_dev = types.SimpleNamespace( + omny_samples=types.SimpleNamespace(get_sample_name_in_samplestage=lambda: "test_sample") + ) + monkeypatch.setattr(omny_module, "dev", fake_dev, raising=False) + omny.OMNYTools = types.SimpleNamespace(yesno=lambda *a, **k: False) + + # canned answers for: ctime, step, fovx, fovy, stitch_x, stitch_y, + # reconstruct_foldername, tomo_type (=2, an advanced/hidden mode), then the + # type-1 fallback's own "Total number of projections" sub-prompt (since + # tomo_type gets forced back to 1 after the superuser check) + canned = iter([1.0, 0.1, 20, 20, 1, 1, "queue", 2, 36]) + monkeypatch.setattr(OMNY, "_get_val", staticmethod(lambda *a, **k: next(canned))) + + omny.tomo_parameters() + + assert omny.tomo_type == 1 + + +def test_tomo_parameters_allows_superuser_to_unlock_8_subtomo_mode(monkeypatch, fake_bec): + """A superuser account can select and unlock tomo_type 5 (8 sub-tomograms) + with the correct authorization code.""" + monkeypatch.setattr(omny_module, "_SUPERUSER_ACCOUNTS", frozenset({"e99999"})) + fake_bec.active_account = "e99999" # superuser + omny = make_omny() + omny.tomo_countingtime = 1.0 + omny.tomo_shellstep = 0.1 + omny.fovx = 20 + omny.fovy = 20 + omny.stitch_x = 1 + omny.stitch_y = 1 + omny.ptycho_reconstruct_foldername = "queue" + omny.tomo_angle_stepsize = 10.0 + omny.tomo_type = 1 + fake_dev = types.SimpleNamespace( + omny_samples=types.SimpleNamespace(get_sample_name_in_samplestage=lambda: "test_sample") + ) + monkeypatch.setattr(omny_module, "dev", fake_dev, raising=False) + omny.OMNYTools = types.SimpleNamespace(yesno=lambda *a, **k: False) + + # canned answers for: ctime, step, fovx, fovy, stitch_x, stitch_y, + # reconstruct_foldername, tomo_type (=5, 8 sub-tomograms), authorization + # code, total number of projections + canned = iter([1.0, 0.1, 20, 20, 1, 1, "queue", 5, "x12sa", 144]) + monkeypatch.setattr(OMNY, "_get_val", staticmethod(lambda *a, **k: next(canned))) + + omny.tomo_parameters() + + assert omny.tomo_type == 5 + # 144 total projections / 8 sub-tomograms -> stepsize 180/(144/8) = 10.0 + assert omny.tomo_angle_stepsize == pytest.approx(10.0)