feat(omny): add 4/8 sub-tomogram tomo modes, hide advanced modes from regular users
CI for csaxs_bec / test (push) Successful in 2m26s

- 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).
This commit is contained in:
x01dc
2026-09-01 14:26:10 +02:00
parent 446b201302
commit e6928f295e
2 changed files with 390 additions and 34 deletions
+131 -34
View File
@@ -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
@@ -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)