v0.3.2 release

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 = [

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

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,13 +106,15 @@ 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
     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

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]:

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
@@ -18,9 +18,24 @@
 # _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
+
+# _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
 
+# 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:
@@ -31,15 +46,23 @@ 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._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._bytes = 0
         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'
 
@@ -48,7 +71,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 +87,99 @@ 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
-            while len(self._slices) > self.cache_slices:
-                self._slices.popitem(last=False)
+            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._store(key, data)
+            del self._inflight[key]
+        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.slice, name, z)
+        with self._lock:
+            if self._fetching is not None:
+                self._fetching.cancel()
+            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
     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)
+    # 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._pool.shutdown(wait=False, cancel_futures=True)
+        self._near_generation += 1
+        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)
+
+    # _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
+
+    # _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:
@@ -142,14 +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 z - radius .. z + radius
-    def _warm_around(self, name: str, z: int, radius: int) -> None:
-        depth = self.open(name).result().vol.shape[0]
-        for zi in range(max(z - radius, 0), min(z + radius + 1, depth)):
+    # _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:
+        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

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)