From 69f526ba8908fb3839b1edf20a330fa809fbf5f1 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 5 Nov 2026 12:18:52 +0000 Subject: [PATCH 1/6] refactor(loader): use separate thread pools - Modified loader.py to replace single pool with a pool for current tomogram slices, one for background read-aheads, and one for opening tomograms. - Modified tests/test_loader.py to update tests to call respective pools, and to add tests covering new split pool functionality. --- src/pickmix/loader.py | 46 +++++++++++++++++++++++++++++++------------ tests/test_loader.py | 35 +++++++++++++++++++++++++++++--- 2 files changed, 65 insertions(+), 16 deletions(-) diff --git a/src/pickmix/loader.py b/src/pickmix/loader.py index c19e2cf..ef9ff14 100644 --- a/src/pickmix/loader.py +++ b/src/pickmix/loader.py @@ -35,8 +35,11 @@ def __init__(self, project: Project, *, cache_slices: int = 96, keep_open: int = self.project = project self.cache_slices = cache_slices self.keep_open = keep_open - self._pool = ThreadPoolExecutor(max_workers=3, thread_name_prefix='pickmix-load') + self._fg = ThreadPoolExecutor(max_workers=2, thread_name_prefix='pickmix-slice') + self._bg = ThreadPoolExecutor(max_workers=2, thread_name_prefix='pickmix-warm') + self._opener = ThreadPoolExecutor(max_workers=2, thread_name_prefix='pickmix-open') self._lock = threading.Lock() + self._inflight: dict[tuple[str, int], Future[np.ndarray]] = {} self._opened: OrderedDict[str, Future[Opened]] = OrderedDict() self._slices: OrderedDict[tuple[str, int], np.ndarray] = OrderedDict() self._warm_generation = 0 @@ -48,7 +51,7 @@ def open(self, name: str) -> Future[Opened]: with self._lock: future = self._opened.get(name) if future is None: - future = self._opened[name] = self._submit(self._open, name) + future = self._opened[name] = self._submit(self._opener, self._open, name) self._opened.move_to_end(name) while len(self._opened) > self.keep_open: self._opened.popitem(last=False) @@ -64,45 +67,62 @@ def cached(self, name: str, z: int) -> np.ndarray | None: # slice: the slice, reading it now if it is not cached; blocks, so call from a worker or when it is cheap def slice(self, name: str, z: int) -> np.ndarray: + key = (name, z) hit = self.cached(name, z) if hit is not None: return hit - opened = self.open(name).result() - with timed(f'slice {name} z={z}'): - data = np.array(opened.vol[z]) with self._lock: - self._slices[(name, z)] = data + pending = self._inflight.get(key) + leader = pending is None + if pending is None: + pending = self._inflight[key] = Future() + if not leader: + return pending.result() + try: + opened = self.open(name).result() + with timed(f'slice {name} z={z}'): + data = np.array(opened.vol[z]) + except BaseException as e: + with self._lock: + del self._inflight[key] + pending.set_exception(e) + raise + with self._lock: + self._slices[key] = data while len(self._slices) > self.cache_slices: self._slices.popitem(last=False) + del self._inflight[key] + pending.set_result(data) return data # fetch: future of a slice, for the GUI to wait on without blocking def fetch(self, name: str, z: int) -> Future[np.ndarray]: - return self._submit(self.slice, name, z) + return self._submit(self._fg, self.slice, name, z) # prefetch: open a tomogram, read its middle slices and warm the page cache for the rest of the file def prefetch(self, name: str) -> None: self._warm_generation += 1 generation = self._warm_generation - self._submit(self._prefetch, name, generation) + self._submit(self._bg, self._prefetch, name, generation) # warm_around: read slices near z so scrolling is served from memory def warm_around(self, name: str, z: int, radius: int = 8) -> None: - self._submit(self._warm_around, name, z, radius) + self._submit(self._bg, self._warm_around, name, z, radius) # shutdown: stop the threads def shutdown(self) -> None: self._closed = True self._warm_generation += 1 - self._pool.shutdown(wait=False, cancel_futures=True) + for pool in (self._fg, self._bg, self._opener): + pool.shutdown(wait=False, cancel_futures=True) - # _submit: run on the pool, or hand back a cancelled future once shut down so late callers fail quietly - def _submit[T](self, fn: Callable[..., T], *args: object) -> Future[T]: + # _submit: run on a pool, or hand back a cancelled future once shut down so late callers fail quietly + def _submit[T](self, pool: ThreadPoolExecutor, fn: Callable[..., T], *args: object) -> Future[T]: if self._closed: future: Future[T] = Future() future.cancel() return future - return self._pool.submit(fn, *args) + return pool.submit(fn, *args) # _tomo: the tomogram record for a name def _tomo(self, name: str) -> Tomogram: diff --git a/tests/test_loader.py b/tests/test_loader.py index 5cc36e5..c6c287e 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -3,12 +3,13 @@ ''' # Import external dependencies -import json, numpy as np +import json, numpy as np, pytest, time from collections.abc import Callable +from concurrent.futures import Future, ThreadPoolExecutor from pathlib import Path # Import internal pickmix objects -from pickmix.loader import Loader +from pickmix.loader import Loader, Opened from pickmix.project import open_project # NAME: the single tomogram of the fixture project @@ -76,7 +77,7 @@ def test_prefetch(self, make_config: Callable[..., Path]) -> None: for _ in range(200): if loader.cached(NAME, 2) is not None: break - loader._pool.submit(lambda: None).result() + loader._bg.submit(lambda: None).result() assert loader.cached(NAME, 2) is not None loader.shutdown() @@ -85,3 +86,31 @@ def test_fetch(self, make_config: Callable[..., Path]) -> None: loader = _loader(make_config) assert loader.fetch(NAME, 3).result().shape == (6, 8) loader.shutdown() + + # test_slice_read_once: concurrent requests for one slice share a single read + def test_slice_read_once(self, make_config: Callable[..., Path]) -> None: + loader = _loader(make_config) + opened = loader.open(NAME).result() + reads = [] + # _Counting: a volume that records which planes are read + class _Counting: + shape = opened.vol.shape + # __getitem__: record the plane, slowly enough for the other threads to arrive + def __getitem__(self, z: int) -> np.ndarray: + reads.append(z) + time.sleep(0.05) + return opened.vol[z] + loader._opened[NAME] = Future() + loader._opened[NAME].set_result(Opened(NAME, _Counting(), opened.limits)) # type: ignore[arg-type] + with ThreadPoolExecutor(max_workers=4) as pool: + results = list(pool.map(lambda _: loader.slice(NAME, 1), range(4))) + assert reads == [1] and all(np.array_equal(r, results[0]) for r in results) + loader.shutdown() + + # test_slice_error_not_cached: a failed read raises for every waiter and is tried again next time + def test_slice_error_not_cached(self, make_config: Callable[..., Path]) -> None: + loader = _loader(make_config) + with pytest.raises(IndexError): + loader.slice(NAME, 99) + assert loader._inflight == {} and loader.slice(NAME, 1).shape == (6, 8) + loader.shutdown() \ No newline at end of file From f60c0fe00146077133641200adbf4cad34d4ab13 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 5 Nov 2026 12:32:58 +0000 Subject: [PATCH 2/6] feat(loader): cancel update fetch priority - Modified loader.py to cancel superseded fetches, and to preferentially fetch slices in the same direction as scroll. - Modified gui/tabs/pick_tab.py to track scroll direction and pass to loader. - Modified tests/test_loader.py to add tests for warm order, fetch cancellation and superseded read-ahead. --- src/pickmix/gui/tabs/pick_tab.py | 7 +++++-- src/pickmix/loader.py | 35 ++++++++++++++++++++++++-------- tests/test_loader.py | 35 ++++++++++++++++++++++++++++++-- 3 files changed, 65 insertions(+), 12 deletions(-) diff --git a/src/pickmix/gui/tabs/pick_tab.py b/src/pickmix/gui/tabs/pick_tab.py index be1602d..15de6a1 100644 --- a/src/pickmix/gui/tabs/pick_tab.py +++ b/src/pickmix/gui/tabs/pick_tab.py @@ -41,6 +41,8 @@ def __init__(self) -> None: self.active: str | None = None self.hidden: set[str] = set() self.z = 0 + # direction of the last z step, so read-ahead goes the way the user is scrolling + self._sign = 1 self._timer: Timer | None = None self._first_slice = False self._range = (0.0, 1.0) @@ -344,6 +346,7 @@ def _show_z(self, z: int) -> None: if self.opened is None or self.loader is None: return name = self.opened.name + self._sign = 1 if z >= self.z else -1 self.z = z total = self.opened.vol.shape[0] # as wide as the longest text it will show, so scrolling never changes the layout @@ -360,7 +363,7 @@ def _show_z(self, z: int) -> None: # _slice_done: a slice read finished; show it if it is still the one wanted def _slice_done(self, name: str, z: int, future: Future[np.ndarray]) -> None: - if self.opened is None or self.opened.name != name or self.z != z: + if future.cancelled() or self.opened is None or self.opened.name != name or self.z != z: return try: self._render(future.result()) @@ -377,7 +380,7 @@ def _render(self, arr: np.ndarray) -> None: self.canvas.fit() if self._timer: self._timer.lap('first slice shown') - self.loader.warm_around(self.opened.name, self.z) + self.loader.warm_around(self.opened.name, self.z, sign=self._sign) # _raw_window: contrast limits from the range slider, before brightness def _raw_window(self) -> tuple[float, float]: diff --git a/src/pickmix/loader.py b/src/pickmix/loader.py index ef9ff14..00ccdfc 100644 --- a/src/pickmix/loader.py +++ b/src/pickmix/loader.py @@ -21,6 +21,15 @@ # _WARM_CHUNK: bytes read per step when warming the page cache _WARM_CHUNK: int = 8 * 1024 * 1024 +# warm_order: slices to read around z, nearest first; sign is the scrolling direction, so slices ahead of it come before the ones behind +def warm_order(z: int, depth: int, *, sign: int, ahead: int, behind: int) -> list[int]: + order = [] + for d in range(1, max(ahead, behind) + 1): + for step, reach in ((sign * d, ahead), (-sign * d, behind)): + if d <= reach and 0 <= z + step < depth: + order.append(z + step) + return order + # Opened: a tomogram ready to show; vol is the read-only (z, y, x) memory map @dataclass(frozen=True) class Opened: @@ -43,6 +52,8 @@ def __init__(self, project: Project, *, cache_slices: int = 96, keep_open: int = self._opened: OrderedDict[str, Future[Opened]] = OrderedDict() self._slices: OrderedDict[tuple[str, int], np.ndarray] = OrderedDict() self._warm_generation = 0 + self._near_generation = 0 + self._fetching: Future[np.ndarray] | None = None self._closed = False self._contrast_path = project.config.out_dir / 'contrast.json' @@ -95,9 +106,13 @@ def slice(self, name: str, z: int) -> np.ndarray: pending.set_result(data) return data - # fetch: future of a slice, for the GUI to wait on without blocking + # fetch: future of a slice, for the GUI to wait on without blocking; the previous fetch is cancelled if it has not started, since the GUI has moved on def fetch(self, name: str, z: int) -> Future[np.ndarray]: - return self._submit(self._fg, self.slice, name, z) + with self._lock: + if self._fetching is not None: + self._fetching.cancel() + future = self._fetching = self._submit(self._fg, self.slice, name, z) + return future # prefetch: open a tomogram, read its middle slices and warm the page cache for the rest of the file def prefetch(self, name: str) -> None: @@ -105,14 +120,16 @@ def prefetch(self, name: str) -> None: generation = self._warm_generation self._submit(self._bg, self._prefetch, name, generation) - # warm_around: read slices near z so scrolling is served from memory - def warm_around(self, name: str, z: int, radius: int = 8) -> None: - self._submit(self._bg, self._warm_around, name, z, radius) + # warm_around: read slices near z so scrolling is served from memory, more of them ahead of the scrolling direction (sign) than behind; a newer call stops an older one + def warm_around(self, name: str, z: int, *, sign: int = 1, ahead: int = 16, behind: int = 4) -> None: + self._near_generation += 1 + self._submit(self._bg, self._warm_around, name, z, self._near_generation, sign, ahead, behind) # shutdown: stop the threads def shutdown(self) -> None: self._closed = True self._warm_generation += 1 + self._near_generation += 1 for pool in (self._fg, self._bg, self._opener): pool.shutdown(wait=False, cancel_futures=True) @@ -162,10 +179,12 @@ def _prefetch(self, name: str, generation: int) -> None: self.slice(name, opened.vol.shape[0] // 2) self._warm_file(volume_path(self._tomo(name), source=self.project.config.source), generation) - # _warm_around: read z - radius .. z + radius - def _warm_around(self, name: str, z: int, radius: int) -> None: + # _warm_around: read the slices warm_order gives, until a newer warm_around starts + def _warm_around(self, name: str, z: int, generation: int, sign: int, ahead: int, behind: int) -> None: depth = self.open(name).result().vol.shape[0] - for zi in range(max(z - radius, 0), min(z + radius + 1, depth)): + for zi in warm_order(z, depth, sign=sign, ahead=ahead, behind=behind): + if self._near_generation != generation: + return self.slice(name, zi) # _warm_file: read a file sequentially so its pages are in the operating system cache; stops when a newer prefetch starts diff --git a/tests/test_loader.py b/tests/test_loader.py index c6c287e..f9cd852 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -9,7 +9,7 @@ from pathlib import Path # Import internal pickmix objects -from pickmix.loader import Loader, Opened +from pickmix.loader import Loader, Opened, warm_order from pickmix.project import open_project # NAME: the single tomogram of the fixture project @@ -113,4 +113,35 @@ def test_slice_error_not_cached(self, make_config: Callable[..., Path]) -> None: with pytest.raises(IndexError): loader.slice(NAME, 99) assert loader._inflight == {} and loader.slice(NAME, 1).shape == (6, 8) - loader.shutdown() \ No newline at end of file + loader.shutdown() + + # test_warm_order: nearest first, and the side the user is scrolling towards gets more + def test_warm_order(self) -> None: + assert warm_order(5, 20, sign=1, ahead=3, behind=1) == [6, 4, 7, 8] + assert warm_order(5, 20, sign=-1, ahead=3, behind=1) == [4, 6, 3, 2] + assert warm_order(1, 4, sign=1, ahead=5, behind=5) == [2, 0, 3] + + # test_fetch_cancels_queued: a fetch that has not started is cancelled by the next one + def test_fetch_cancels_queued(self, make_config: Callable[..., Path]) -> None: + loader = _loader(make_config) + gate = Future() + # two workers held busy, so the next fetches queue + blockers = [loader._fg.submit(gate.result) for _ in range(2)] + first = loader.fetch(NAME, 0) + second = loader.fetch(NAME, 1) + assert first.cancelled() + gate.set_result(None) + assert second.result().shape == (6, 8) + loader.shutdown() + assert all(b.done() for b in blockers) + + # test_warm_stops_when_superseded: a newer warm_around stops an older one + def test_warm_stops_when_superseded(self, make_config: Callable[..., Path]) -> None: + loader = _loader(make_config) + loader.open(NAME).result() + loader._near_generation += 1 + loader._warm_around(NAME, 0, loader._near_generation - 1, 1, 3, 3) + assert all(loader.cached(NAME, z) is None for z in range(4)) + loader._warm_around(NAME, 0, loader._near_generation, 1, 3, 3) + assert [loader.cached(NAME, z) is not None for z in range(4)] == [False, True, True, True] + loader.shutdown() From 5e0a62ac07fbe07c31c39f223c17352ba46dff77 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 5 Nov 2026 12:48:53 +0000 Subject: [PATCH 3/6] feat(loader): define slice cache in config - Modified config.py to add field for cache size to define maximum memory used for caching slices. - Modified loader.py to bound slice cache by bytes rather than count, evicting least recently used slices. - Modified tests/test_loader.py to use byte bound rather than slice bound, and to add tests for byte budget. --- src/pickmix/config.py | 2 ++ src/pickmix/loader.py | 23 ++++++++++++++++++----- tests/test_loader.py | 26 ++++++++++++++++++++++++-- 3 files changed, 44 insertions(+), 7 deletions(-) diff --git a/src/pickmix/config.py b/src/pickmix/config.py index 16822e7..db4f33b 100644 --- a/src/pickmix/config.py +++ b/src/pickmix/config.py @@ -70,6 +70,8 @@ class Config(BaseModel): out_dir: Path = Path('out') # volume read for display, estimation and extraction; falls back to raw when a tomogram has no denoised pair source: Literal['raw', 'denoised'] = 'denoised' + # most memory the slice cache may hold; scrolling is smooth where a whole tomogram fits + cache_mb: int = Field(default=2048, ge=64, title='Slice cache (MB)', description='Most memory the viewer may use to keep slices of tomograms ready for scrolling.') ring: RingParams = Field(default_factory=RingParams) kinds: dict[str, Kind] = Field(default_factory=_default_kinds) # kind chosen by each of the keys 1 to 9, '' for an unused key diff --git a/src/pickmix/loader.py b/src/pickmix/loader.py index 00ccdfc..4a78362 100644 --- a/src/pickmix/loader.py +++ b/src/pickmix/loader.py @@ -18,6 +18,9 @@ # _CONTRAST_HALF: slices either side of the middle used for contrast limits _CONTRAST_HALF: int = 2 +# _MIB: bytes in a mebibyte +_MIB: int = 1024 * 1024 + # _WARM_CHUNK: bytes read per step when warming the page cache _WARM_CHUNK: int = 8 * 1024 * 1024 @@ -40,9 +43,9 @@ class Opened: # Loader: opens tomograms and serves slices class Loader: # __init__: nothing is read until asked - def __init__(self, project: Project, *, cache_slices: int = 96, keep_open: int = 3) -> None: + def __init__(self, project: Project, *, cache_bytes: int | None = None, keep_open: int = 3) -> None: self.project = project - self.cache_slices = cache_slices + self.cache_bytes = project.config.cache_mb * _MIB if cache_bytes is None else cache_bytes self.keep_open = keep_open self._fg = ThreadPoolExecutor(max_workers=2, thread_name_prefix='pickmix-slice') self._bg = ThreadPoolExecutor(max_workers=2, thread_name_prefix='pickmix-warm') @@ -51,6 +54,7 @@ def __init__(self, project: Project, *, cache_slices: int = 96, keep_open: int = self._inflight: dict[tuple[str, int], Future[np.ndarray]] = {} self._opened: OrderedDict[str, Future[Opened]] = OrderedDict() self._slices: OrderedDict[tuple[str, int], np.ndarray] = OrderedDict() + self._bytes = 0 self._warm_generation = 0 self._near_generation = 0 self._fetching: Future[np.ndarray] | None = None @@ -99,9 +103,7 @@ def slice(self, name: str, z: int) -> np.ndarray: pending.set_exception(e) raise with self._lock: - self._slices[key] = data - while len(self._slices) > self.cache_slices: - self._slices.popitem(last=False) + self._store(key, data) del self._inflight[key] pending.set_result(data) return data @@ -141,6 +143,17 @@ def _submit[T](self, pool: ThreadPoolExecutor, fn: Callable[..., T], *args: obje return future return pool.submit(fn, *args) + # _store: add a slice to the cache and drop the least recently used until it fits the budget, but never the slice just added; call with the lock held + def _store(self, key: tuple[str, int], data: np.ndarray) -> None: + old = self._slices.pop(key, None) + if old is not None: + self._bytes -= old.nbytes + self._slices[key] = data + self._bytes += data.nbytes + while self._bytes > self.cache_bytes and len(self._slices) > 1: + _, dropped = self._slices.popitem(last=False) + self._bytes -= dropped.nbytes + # _tomo: the tomogram record for a name def _tomo(self, name: str) -> Tomogram: return self.project.tomos[name] diff --git a/tests/test_loader.py b/tests/test_loader.py index f9cd852..4e3dfb2 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -35,9 +35,9 @@ def test_cached(self, make_config: Callable[..., Path]) -> None: assert loader.cached(NAME, 1) is not None loader.shutdown() - # test_lru_eviction: the oldest slice goes first and a touched one stays + # test_lru_eviction: the oldest slice goes first and a touched one stays; a slice here is 6 x 8 float32, 192 bytes def test_lru_eviction(self, make_config: Callable[..., Path]) -> None: - loader = _loader(make_config, cache_slices=2) + loader = _loader(make_config, cache_bytes=2 * 192) loader.slice(NAME, 0) loader.slice(NAME, 1) loader.cached(NAME, 0) @@ -145,3 +145,25 @@ def test_warm_stops_when_superseded(self, make_config: Callable[..., Path]) -> N loader._warm_around(NAME, 0, loader._near_generation, 1, 3, 3) assert [loader.cached(NAME, z) is not None for z in range(4)] == [False, True, True, True] loader.shutdown() + + # test_cache_budget: the cache holds the budget in bytes and tracks it as slices come and go + def test_cache_budget(self, make_config: Callable[..., Path]) -> None: + loader = _loader(make_config, cache_bytes=3 * 192) + for z in range(4): + loader.slice(NAME, z) + assert loader._bytes == 3 * 192 and loader.cached(NAME, 0) is None and loader.cached(NAME, 3) is not None + loader.shutdown() + + # test_cache_budget_keeps_newest: a budget smaller than one slice still keeps the slice just read + def test_cache_budget_keeps_newest(self, make_config: Callable[..., Path]) -> None: + loader = _loader(make_config, cache_bytes=1) + loader.slice(NAME, 0) + loader.slice(NAME, 1) + assert loader.cached(NAME, 0) is None and loader.cached(NAME, 1) is not None and loader._bytes == 192 + loader.shutdown() + + # test_default_budget: the budget comes from cache_mb in the config + def test_default_budget(self, make_config: Callable[..., Path]) -> None: + loader = _loader(make_config) + assert loader.cache_bytes == 2048 * 1024 * 1024 + loader.shutdown() From e72c0f716fbda93d926ded6f66621f4f22746ea3 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 5 Nov 2026 12:54:45 +0000 Subject: [PATCH 4/6] feat(gui): add slice cache size to settings pane - Modified gui/panes/settings_pane.py to add spin box for slice cache. - Modified tests/gui/test_panes.py to add test for autosaved slice cache changes. --- src/pickmix/gui/panes/settings_pane.py | 14 +++++++++++--- tests/gui/test_panes.py | 8 ++++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/src/pickmix/gui/panes/settings_pane.py b/src/pickmix/gui/panes/settings_pane.py index 0fea03d..f8a1055 100644 --- a/src/pickmix/gui/panes/settings_pane.py +++ b/src/pickmix/gui/panes/settings_pane.py @@ -8,7 +8,7 @@ from PySide6.QtWidgets import QComboBox, QDoubleSpinBox, QFormLayout, QGroupBox, QHBoxLayout, QLabel, QSpinBox, QVBoxLayout, QWidget # Import internal pickmix objects -from pickmix.config import HOTKEYS, RingParams +from pickmix.config import HOTKEYS, Config, RingParams from pickmix.gui.components.collapsible import Collapsible from pickmix.gui.components.toast import notify from pickmix.gui.style import form_label, heading @@ -80,11 +80,17 @@ def __init__(self) -> None: self._loading = False self.source_box = QComboBox() self.source_box.addItems(['denoised', 'raw']) + self.cache_spin = QSpinBox() + self.cache_spin.setRange(64, 1_000_000) + self.cache_spin.setSingleStep(256) + self.cache_spin.setKeyboardTracking(False) self.hotkey_boxes = [QComboBox() for _ in range(HOTKEYS)] self.ring_form = ParamsForm(RingParams()) general = QGroupBox('General') form = QFormLayout(general) form.addRow('Volume source', self.source_box) + cache = Config.model_fields['cache_mb'] + form.addRow(form_label(cache.title or 'cache_mb', cache.description or ''), self.cache_spin) for i, box in enumerate(self.hotkey_boxes): form.addRow(f'Hotkey {i + 1}', box) box.currentTextChanged.connect(lambda text, i=i: self._hotkey_changed(i, text)) @@ -100,6 +106,7 @@ def __init__(self) -> None: layout.addWidget(estimators) layout.addStretch() self.source_box.currentTextChanged.connect(self._save) + self.cache_spin.valueChanged.connect(self._save) self.ring_form.changed.connect(self._save) # set_project: show a project's settings without saving them back @@ -107,6 +114,7 @@ def set_project(self, project: Project) -> None: self.project = project self._loading = True self.source_box.setCurrentText(project.config.source) + self.cache_spin.setValue(project.config.cache_mb) self.ring_form.set(project.config.ring) self._loading = False self.refresh_hotkeys() @@ -123,12 +131,12 @@ def refresh_hotkeys(self) -> None: box.setCurrentText(chosen or NO_KIND) self._loading = False - # _save: write source and ring parameters; a value the config refuses is not written + # _save: write source, slice cache and ring parameters; a value the config refuses is not written def _save(self) -> None: if self._loading or self.project is None: return try: - self.project.update(source=self.source_box.currentText(), ring=self.ring_form.get()) + self.project.update(source=self.source_box.currentText(), cache_mb=self.cache_spin.value(), ring=self.ring_form.get()) except ValidationError as e: notify(f'Not saved: {e.errors()[0]["msg"]}') return diff --git a/tests/gui/test_panes.py b/tests/gui/test_panes.py index 3fc28f9..9187e5c 100644 --- a/tests/gui/test_panes.py +++ b/tests/gui/test_panes.py @@ -342,6 +342,14 @@ def test_source_autosaves(self, qtbot: QtBot, make_config: Callable[..., Path]) pane.source_box.setCurrentText('raw') assert load_config(path).source == 'raw' + # test_cache_autosaves: a committed slice cache size is written, and the project's own is shown + def test_cache_autosaves(self, qtbot: QtBot, make_config: Callable[..., Path]) -> None: + pane, path = _settings_pane(qtbot, make_config) + assert pane.cache_spin.value() == 2048 + with qtbot.waitSignal(pane.config_saved): + pane.cache_spin.setValue(3072) + assert load_config(path).cache_mb == 3072 + # test_ring_value_autosaves: a committed ring value is written, the others are kept def test_ring_value_autosaves(self, qtbot: QtBot, make_config: Callable[..., Path]) -> None: pane, path = _settings_pane(qtbot, make_config) From 03a6a97807aa69e531b32295a4c15ea8f226182e Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 5 Nov 2026 13:03:58 +0000 Subject: [PATCH 5/6] feat(loader): fill cache with tomogram on idle - Modified loader.py to continue reading slices outwards from current slice until complete or out-of-cache. - Modified tests/test_loader.py to add tests for filling cache within budget and waiting on idle viewer. --- src/pickmix/loader.py | 51 +++++++++++++++++++++++++++++++++++++------ tests/test_loader.py | 23 +++++++++++++++++++ 2 files changed, 67 insertions(+), 7 deletions(-) diff --git a/src/pickmix/loader.py b/src/pickmix/loader.py index 4a78362..4ba2c2a 100644 --- a/src/pickmix/loader.py +++ b/src/pickmix/loader.py @@ -3,7 +3,7 @@ ''' # Import external dependencies -import json, threading, numpy as np +import json, threading, time, numpy as np from collections import OrderedDict from collections.abc import Callable from concurrent.futures import Future, ThreadPoolExecutor @@ -21,6 +21,9 @@ # _MIB: bytes in a mebibyte _MIB: int = 1024 * 1024 +# _IDLE_S: seconds after the last slice the viewer waited for before read-ahead of far slices and whole-file warming carry on +_IDLE_S: float = 0.25 + # _WARM_CHUNK: bytes read per step when warming the page cache _WARM_CHUNK: int = 8 * 1024 * 1024 @@ -58,6 +61,8 @@ def __init__(self, project: Project, *, cache_bytes: int | None = None, keep_ope self._warm_generation = 0 self._near_generation = 0 self._fetching: Future[np.ndarray] | None = None + self._fg_running = 0 + self._fg_stamp = 0.0 self._closed = False self._contrast_path = project.config.out_dir / 'contrast.json' @@ -113,7 +118,8 @@ def fetch(self, name: str, z: int) -> Future[np.ndarray]: with self._lock: if self._fetching is not None: self._fetching.cancel() - future = self._fetching = self._submit(self._fg, self.slice, name, z) + self._fg_stamp = time.monotonic() + future = self._fetching = self._submit(self._fg, self._fetch_task, name, z) return future # prefetch: open a tomogram, read its middle slices and warm the page cache for the rest of the file @@ -154,6 +160,27 @@ def _store(self, key: tuple[str, int], data: np.ndarray) -> None: _, dropped = self._slices.popitem(last=False) self._bytes -= dropped.nbytes + # _fetch_task: read a slice for the viewer, counted so read-ahead can stand aside while it runs + def _fetch_task(self, name: str, z: int) -> np.ndarray: + with self._lock: + self._fg_running += 1 + try: + return self.slice(name, z) + finally: + with self._lock: + self._fg_running -= 1 + self._fg_stamp = time.monotonic() + + # _wait_idle: wait until no slice is being read for the viewer (and none was for grace seconds); False if stale() says the work is no longer wanted + def _wait_idle(self, stale: Callable[[], bool], grace: float) -> bool: + while not stale(): + with self._lock: + idle = self._fg_running == 0 and time.monotonic() - self._fg_stamp >= grace + if idle: + return True + time.sleep(0.02) + return False + # _tomo: the tomogram record for a name def _tomo(self, name: str) -> Tomogram: return self.project.tomos[name] @@ -192,16 +219,26 @@ def _prefetch(self, name: str, generation: int) -> None: self.slice(name, opened.vol.shape[0] // 2) self._warm_file(volume_path(self._tomo(name), source=self.project.config.source), generation) - # _warm_around: read the slices warm_order gives, until a newer warm_around starts + # _warm_around: read the nearby slices warm_order gives, then the rest of the tomogram outwards from z while it fits the cache and the viewer is idle; a newer warm_around stops it def _warm_around(self, name: str, z: int, generation: int, sign: int, ahead: int, behind: int) -> None: - depth = self.open(name).result().vol.shape[0] - for zi in warm_order(z, depth, sign=sign, ahead=ahead, behind=behind): - if self._near_generation != generation: + vol = self.open(name).result().vol + depth, plane = vol.shape[0], vol[0].nbytes + stale = lambda: self._near_generation != generation + near = warm_order(z, depth, sign=sign, ahead=ahead, behind=behind) + for zi in near: + if not self._wait_idle(stale, 0.0): + return + self.slice(name, zi) + for zi in warm_order(z, depth, sign=sign, ahead=depth, behind=depth): + if zi in near: + continue + if self._bytes + plane > self.cache_bytes or not self._wait_idle(stale, _IDLE_S): return self.slice(name, zi) # _warm_file: read a file sequentially so its pages are in the operating system cache; stops when a newer prefetch starts def _warm_file(self, path: Path, generation: int) -> None: with timed(f'warm {path.name}'), open(path, 'rb') as handle: - while self._warm_generation == generation and handle.read(_WARM_CHUNK): + stale = lambda: self._warm_generation != generation + while self._wait_idle(stale, _IDLE_S) and handle.read(_WARM_CHUNK): pass diff --git a/tests/test_loader.py b/tests/test_loader.py index 4e3dfb2..10ba0c6 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -135,6 +135,29 @@ def test_fetch_cancels_queued(self, make_config: Callable[..., Path]) -> None: loader.shutdown() assert all(b.done() for b in blockers) + # test_warm_fills_cache: after the nearby slices the rest of the tomogram is read outwards while the budget allows + def test_warm_fills_cache(self, make_config: Callable[..., Path]) -> None: + loader = _loader(make_config) + loader.slice(NAME, 1) + loader._warm_around(NAME, 1, loader._near_generation, 1, 1, 1) + assert all(loader.cached(NAME, z) is not None for z in range(4)) + small = Loader(loader.project, cache_bytes=3 * 192) + small.slice(NAME, 0) + small._warm_around(NAME, 0, small._near_generation, 1, 1, 1) + assert [small.cached(NAME, z) is not None for z in range(4)] == [True, True, True, False] + loader.shutdown() + small.shutdown() + + # test_wait_idle: read-ahead waits while the viewer is reading, and gives up once it is no longer wanted + def test_wait_idle(self, make_config: Callable[..., Path]) -> None: + loader = _loader(make_config) + assert loader._wait_idle(lambda: False, 0.0) + loader._fg_running = 1 + wanted = iter([False, False, True]) + assert not loader._wait_idle(lambda: next(wanted), 0.0) + loader._fg_running = 0 + loader.shutdown() + # test_warm_stops_when_superseded: a newer warm_around stops an older one def test_warm_stops_when_superseded(self, make_config: Callable[..., Path]) -> None: loader = _loader(make_config) From b581555551a533bcdf64a730b2554f8121243139 Mon Sep 17 00:00:00 2001 From: Sam Holland Date: Thu, 5 Nov 2026 13:04:04 +0000 Subject: [PATCH 6/6] chore: update pickmix version to v0.3.2 --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 292940b..03fe524 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pickmix" -version = "0.3.1" +version = "0.3.2" description = "Multi-kind particle picking GUI for cryo-ET sub-tomogram averaging workflows." readme = "README.md" authors = [ diff --git a/uv.lock b/uv.lock index 7f3e8a7..9d722e0 100644 --- a/uv.lock +++ b/uv.lock @@ -137,7 +137,7 @@ wheels = [ [[package]] name = "pickmix" -version = "0.3.1" +version = "0.3.2" source = { editable = "." } dependencies = [ { name = "imageio" },