diff --git a/README.md b/README.md index 0f5f384..be1c97e 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,9 @@ Open the Pick tab. To the left is a a list of tomograms in the project, clicking ![Pick tab](docs/images/pick-tab.png) +#### 3.1 Template matching +Once several picks exist for a kind, the Match group within the settings pane can be used to identify potential candidates. Click `Scan` to generate an average pick and search unvisited tomograms for matches, which will appear as dotted markers on tomograms and in the pick table with a blue status icon. Click a candidate to select it, then press either `Accept` (or hotkey `a`) to convert it into a pick, or `Reject` (or either hotkey `delete`/`backspace`) to hide it. Currently this search is not rotation aware, so best suits symmetric particles. Setting for template matching can be edited in the project tab within the settings tab. + ### 4. Review Open the Review tab. Picks can be filtered by kind, tomogram, minimum SNR and whether they were flagged, with options to delete the selected pick or all currently shown. All picks can be queued for re-estimation here if required as well. Single-click a pick to view XY/XZ/YZ previews, or double-click to open that pick in the Pick tab. @@ -76,5 +79,5 @@ project/ config.toml settings and kinds tomograms.star merged RELION tomogram table tilt_series/ one star file per tomogram - out/ picks, estimates, averages, RELION & .csv exports + out/ picks, candidates, estimates, averages, RELION & .csv exports ``` diff --git a/pyproject.toml b/pyproject.toml index 03fe524..2d76a35 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pickmix" -version = "0.3.2" +version = "0.4.0" description = "Multi-kind particle picking GUI for cryo-ET sub-tomogram averaging workflows." readme = "README.md" authors = [ diff --git a/src/pickmix/candidates.py b/src/pickmix/candidates.py new file mode 100644 index 0000000..b61be2b --- /dev/null +++ b/src/pickmix/candidates.py @@ -0,0 +1,149 @@ +''' +pickmix: candidate picks found by template matching, saved to out_dir/candidates.csv +''' + +# Import external dependencies +import csv, itertools, os, threading +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from pathlib import Path +from pydantic import BaseModel +from typing import Literal + +# Import internal pickmix objects +from pickmix.project import Project + +# Hit: candidate centre in reconstruction voxels and its score +@dataclass(frozen=True) +class Hit: + x: float + y: float + z: float + score: float + +# Candidate: region matching a template +@dataclass +class Candidate: + id: int + tomogram: str + kind: str + x: float + y: float + z: float + score: float + status: Literal['new', 'rejected'] = 'new' + +# _Row: one candidates.csv row +class _Row(BaseModel): + tomogram: str + kind: str + x: float + y: float + z: float + score: float + status: Literal['new', 'rejected'] = 'new' + +# CandidateStore: every candidate, saved after each change +class CandidateStore: + # __init__: load candidates.csv if there is one + def __init__(self, project: Project) -> None: + self.project = project + self.path: Path = project.config.out_dir / 'candidates.csv' + # called from whichever thread changed the store + self.listeners: list[Callable[[], None]] = [] + self._lock = threading.Lock() + self._ids = itertools.count(1) + self._items: dict[int, Candidate] = {} + if self.path.exists(): + with open(self.path, newline='') as handle: + for row in map(_Row.model_validate, csv.DictReader(handle)): + candidate = Candidate(next(self._ids), **row.model_dump()) + self._items[candidate.id] = candidate + + # get: the candidate with this id + def get(self, candidate_id: int) -> Candidate: + with self._lock: + return self._items[candidate_id] + + # candidates: the new candidates in a tomogram, optionally of one kind and at or above a score, best first + def candidates(self, tomogram: str, kind: str | None = None, *, min_score: float = 0.0) -> list[Candidate]: + with self._lock: + found = [c for c in self._items.values() if c.status == 'new' and c.tomogram == tomogram and (kind is None or c.kind == kind) and c.score >= min_score] + return sorted(found, key=lambda c: -c.score) + + # count: number of new candidates of a kind across all tomograms, or in one tomogram + def count(self, kind: str, tomogram: str | None = None, *, min_score: float = 0.0) -> int: + with self._lock: + return sum(c.status == 'new' and c.kind == kind and (tomogram is None or c.tomogram == tomogram) and c.score >= min_score for c in self._items.values()) + + # replace: the new candidates of a kind in a tomogram become these hits, except any within reach voxels of a rejected one + def replace(self, tomogram: str, kind: str, hits: Iterable[Hit], *, reach: float) -> None: + with self._lock: + rejected = [c for c in self._items.values() if c.status == 'rejected' and c.tomogram == tomogram and c.kind == kind] + for stale in [c.id for c in self._items.values() if c.status == 'new' and c.tomogram == tomogram and c.kind == kind]: + del self._items[stale] + for h in hits: + if all((h.x - r.x) ** 2 + (h.y - r.y) ** 2 + (h.z - r.z) ** 2 > reach**2 for r in rejected): + candidate = Candidate(next(self._ids), tomogram, kind, h.x, h.y, h.z, h.score) + self._items[candidate.id] = candidate + self._changed() + + # reject: mark candidates as not wanted + def reject(self, *candidate_ids: int) -> None: + with self._lock: + for candidate_id in candidate_ids: + self._items[candidate_id].status = 'rejected' + self._changed() + + # remove: forget candidates, as when they become picks + def remove(self, *candidate_ids: int) -> None: + with self._lock: + for candidate_id in candidate_ids: + del self._items[candidate_id] + self._changed() + + # clear: forget the candidates of a kind, or all, rejected ones included + def clear(self, kind: str | None = None) -> None: + with self._lock: + self._items = {i: c for i, c in self._items.items() if kind is not None and c.kind != kind} + self._changed() + + # prune: forget every candidate for a tomogram + def prune(self, tomograms: Iterable[str]) -> None: + gone = set(tomograms) + with self._lock: + kept = {i: c for i, c in self._items.items() if c.tomogram not in gone} + changed = len(kept) != len(self._items) + self._items = kept + if changed: + self._changed() + + # rename_kind: move every candidate of a kind to a new name + def rename_kind(self, old: str, new: str) -> None: + with self._lock: + for candidate in self._items.values(): + if candidate.kind == old: + candidate.kind = new + self._changed() + + # remove_kind: forget every candidate of a kind + def remove_kind(self, kind: str) -> None: + self.clear(kind) + + # save: write CSV atomically + def save(self) -> None: + with self._lock: + rows = [_Row(**{k: v for k, v in vars(c).items() if k != 'id'}).model_dump() for c in self._items.values()] + self.path.parent.mkdir(parents=True, exist_ok=True) + tmp = self.path.with_suffix('.csv.tmp') + with open(tmp, 'w', newline='') as handle: + writer = csv.DictWriter(handle, fieldnames=list(_Row.model_fields)) + writer.writeheader() + writer.writerows(rows) + os.replace(tmp, self.path) + + # _changed: save, then tell listeners + def _changed(self) -> None: + self.save() + for listener in self.listeners: + listener() diff --git a/src/pickmix/config.py b/src/pickmix/config.py index db4f33b..89cde38 100644 --- a/src/pickmix/config.py +++ b/src/pickmix/config.py @@ -51,6 +51,21 @@ class RingParams(BaseModel): # box edge as a multiple of the ring diameter when the kind takes its box from the estimator box_margin: float = Field(default=1.25, title='Box margin (× diameter)', description='Box edge as a multiple of the ring diameter, for kinds that take their box from the estimator.') +# MatchParams: settings of the template matching scan; title and description feed the settings form +class MatchParams(BaseModel): + # template is binned until its edge is about this many pixels, which sets the speed of the scan + template_px: int = Field(default=32, ge=8, title='Template edge (px)', description='The tomogram and template are binned until the template is about this many pixels across. Smaller is faster and coarser.') + # candidates scoring below this are not kept + min_score: float = Field(default=0.2, gt=0, lt=1, title='Minimum score', description='Candidates with a normalised cross-correlation below this are not kept.') + # nearest two candidates may be, as a multiple of the kind's marker diameter + separation: float = Field(default=0.75, gt=0, title='Separation (× marker diameter)', description='Closest two candidates, or a candidate and a pick, may be.') + # mask radius as a multiple of the marker radius + mask_margin: float = Field(default=1.2, gt=0, title='Mask radius (× marker radius)', description='Radius of the spherical mask applied to the template, as a multiple of the marker radius.') + # best candidates kept per tomogram + max_per_tomogram: int = Field(default=200, ge=1, title='Most per tomogram', description='Most candidates kept in one tomogram, best first.') + # threads the correlation may use + workers: int = Field(default=2, ge=1, title='Threads', description='Threads the correlation may use. More is faster but slows the rest of the program.') + # _default_kinds: the single cage kind used when a config has no kinds table def _default_kinds() -> dict[str, Kind]: return {'cage': Kind(marker_a=600.0, estimator='ring')} @@ -73,6 +88,7 @@ class Config(BaseModel): # 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) + match: MatchParams = Field(default_factory=MatchParams) kinds: dict[str, Kind] = Field(default_factory=_default_kinds) # kind chosen by each of the keys 1 to 9, '' for an unused key hotkeys: list[str] = Field(default_factory=list) diff --git a/src/pickmix/gui/components/assets/image_search.svg b/src/pickmix/gui/components/assets/image_search.svg new file mode 100644 index 0000000..70cbc30 --- /dev/null +++ b/src/pickmix/gui/components/assets/image_search.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/src/pickmix/gui/components/canvas.py b/src/pickmix/gui/components/canvas.py index 2a023f0..56cc458 100644 --- a/src/pickmix/gui/components/canvas.py +++ b/src/pickmix/gui/components/canvas.py @@ -45,6 +45,9 @@ class Marker: fit: FitView | None = None # membrane normal projected on the slice, as a unit-length-or-shorter (x, y) vector normal: tuple[float, float] | None = None + # template matching candidate rather than a pick + candidate: bool = False + tip: str = '' # SliceCanvas: one slice plus markers; emits requests, never edits picks itself class SliceCanvas(QGraphicsView): @@ -151,9 +154,10 @@ def set_markers(self, markers: list[Marker]) -> None: ring = QGraphicsEllipseItem(-m.radius_px, -m.radius_px, 2 * m.radius_px, 2 * m.radius_px) ring.setPos(m.x + 0.5, m.y + 0.5) ring.setZValue(1) + ring.setToolTip(m.tip) dot = QGraphicsEllipseItem(-2, -2, 4, 4, ring) dot.setFlag(QGraphicsEllipseItem.GraphicsItemFlag.ItemIgnoresTransformations) - dot.setBrush(QColor(m.colour)) + dot.setBrush(Qt.BrushStyle.NoBrush if m.candidate else QColor(m.colour)) dot.setPen(QPen(Qt.PenStyle.NoPen)) if m.normal is not None: reach = max(4 * m.radius_px, 16.0) @@ -212,6 +216,9 @@ def mouseMoveEvent(self, event: QMouseEvent) -> None: self.verticalScrollBar().setValue(self.verticalScrollBar().value() - int(delta.y())) elif self._drag is not None: pick_id, start = self._drag + # don't drag a candidate + if self._markers[pick_id].candidate: + return pos = self.mapToScene(event.position().toPoint()) delta = pos - start if self._drag_moved or math.hypot(delta.x(), delta.y()) * self.transform().m11() > _DRAG_PX: @@ -275,6 +282,8 @@ def _restyle(self) -> None: pen.setWidth(3 if marker_id == self.selected_id else 2) if m.flagged: pen.setStyle(Qt.PenStyle.DashLine) + elif m.candidate: + pen.setStyle(Qt.PenStyle.DotLine) item.setPen(pen) for arrow in self._arrows.values(): arrow.setVisible(self._normals) diff --git a/src/pickmix/gui/components/status.py b/src/pickmix/gui/components/status.py index a86830e..2bb14cf 100644 --- a/src/pickmix/gui/components/status.py +++ b/src/pickmix/gui/components/status.py @@ -12,8 +12,8 @@ from pickmix.gui.components.icons import tinted from pickmix.gui.style import ERROR, OK, RUNNING, WARN -# Status: ok is a green tick, warn a yellow warning sign, error a red error sign, running a blue spinner -type Status = Literal['ok', 'warn', 'error', 'running'] +# Status: ok is a green tick, warn a yellow warning sign, error a red error sign, running a blue spinner, match a blue image search for a template match candidate +type Status = Literal['ok', 'warn', 'error', 'running', 'match'] # STATUS_STYLE: icon stem, colour and the text used when the icon file is missing STATUS_STYLE: dict[Status, tuple[str, str, str]] = { @@ -21,6 +21,7 @@ 'warn': ('warning', WARN, '!'), 'error': ('error', ERROR, '✖'), 'running': ('refresh', RUNNING, '…'), + 'match': ('image_search', RUNNING, 'tm'), } # _SIZE: edge in pixels of a status icon diff --git a/src/pickmix/gui/main_window.py b/src/pickmix/gui/main_window.py index ffd8866..a89ed6c 100644 --- a/src/pickmix/gui/main_window.py +++ b/src/pickmix/gui/main_window.py @@ -16,6 +16,7 @@ from pickmix.gui.tabs.project_tab import ProjectTab from pickmix.gui.tabs.review_tab import ReviewTab from pickmix.loader import Loader +from pickmix.matcher import MatchService from pickmix.model import Session from pickmix.project import Project @@ -39,6 +40,7 @@ def __init__(self) -> None: self.session: Session | None = None self.loader: Loader | None = None self.service: EstimateService | None = None + self.matcher: MatchService | None = None self.tabs = QTabWidget() self.project_tab = ProjectTab() self.kinds_pane = self.project_tab.kinds_pane @@ -70,6 +72,7 @@ def __init__(self) -> None: self.toast = install(self) for i in range(1, 10): QShortcut(QKeySequence(str(i)), self, context=Qt.ShortcutContext.WindowShortcut).activated.connect(lambda i=i: self._hotkey(i - 1)) + QShortcut(QKeySequence('A'), self, context=Qt.ShortcutContext.WindowShortcut).activated.connect(self._accept_candidate) for key in ('Delete', 'Backspace'): QShortcut(QKeySequence(key), self, context=Qt.ShortcutContext.WindowShortcut).activated.connect(self._delete_pick) self.resize(*WELCOME_SIZE) @@ -107,6 +110,11 @@ def _hotkey(self, i: int) -> None: if self._picking(): self.pick_tab.pick_kind(i) + # _accept_candidate: A converts selected candidate to pick + def _accept_candidate(self) -> None: + if self._picking(): + self.pick_tab.accept_selected() + # _delete_pick: Delete or Backspace removes the selected pick, from wherever focus is def _delete_pick(self) -> None: if self._picking(): @@ -162,12 +170,13 @@ def _on_project(self, project: Project) -> None: if project.tomos: self.session, self.loader = Session(project), Loader(project) self.service = EstimateService(self.session, self.loader) - self.pick_tab.set_session(self.session, self.loader, self.service) + self.matcher = MatchService(self.session, self.service.store) + self.pick_tab.set_session(self.session, self.loader, self.service, self.matcher) self.review_tab.set_session(self.session, self.loader, self.service) self.export_tab.set_session(self.session, self.loader, self.service) self.kinds_pane.set_session(self.session, self.service) else: - self.session = self.loader = self.service = None + self.session = self.loader = self.service = self.matcher = None self.kinds_pane.set_session(None, None) self.export_tab.clear() self._set_enabled(session=self.session is not None) @@ -176,9 +185,11 @@ def _on_project(self, project: Project) -> None: def _shutdown(self) -> None: if self.service is not None: self.service.shutdown() + if self.matcher is not None: + self.matcher.shutdown() if self.loader is not None: self.loader.shutdown() - self.session = self.loader = self.service = None + self.session = self.loader = self.service = self.matcher = None # _on_closed: no project; lock everything but the Project tab def _on_closed(self) -> None: diff --git a/src/pickmix/gui/panes/kinds_pane.py b/src/pickmix/gui/panes/kinds_pane.py index 7ae43bd..25f555d 100644 --- a/src/pickmix/gui/panes/kinds_pane.py +++ b/src/pickmix/gui/panes/kinds_pane.py @@ -186,6 +186,7 @@ def _apply(self, name: str, kind: Kind, *, old: str | None = None) -> None: self.service.store.rename_kind(old, name) if self.session is not None: self.session.store.rename_kind(old, name) + self.session.candidates.rename_kind(old, name) self._saved() # _delete: remove a kind and its picks, asking first if it has any; refused with a reason if the config would not be valid @@ -204,6 +205,7 @@ def _delete(self, name: str) -> None: self.project.update(**changes) if self.session is not None: self.session.store.remove_kind(name) + self.session.candidates.remove_kind(name) self._saved() # _confirm: ask a yes or no question diff --git a/src/pickmix/gui/panes/match_pane.py b/src/pickmix/gui/panes/match_pane.py new file mode 100644 index 0000000..44f17d6 --- /dev/null +++ b/src/pickmix/gui/panes/match_pane.py @@ -0,0 +1,110 @@ +''' +pickmix: template matching pane +''' + +# Import external dependencies +from PySide6.QtCore import Qt, Signal +from PySide6.QtWidgets import QGroupBox, QHBoxLayout, QLabel, QPushButton, QSlider, QVBoxLayout + +# Import internal pickmix objects +from pickmix.matcher import MatchService + +# _STEPS: positions on the score slider, one per hundredth +_STEPS: int = 100 + +# MatchPane: scan for a kind's template, set how good a match to show, and accept or reject the candidates +class MatchPane(QGroupBox): + # scan_requested, stop_requested: the buttons + scan_requested = Signal() + stop_requested = Signal() + # accept_requested, reject_requested, accept_all_requested: act on the selected candidate, or on all those shown in this tomogram + accept_requested = Signal() + reject_requested = Signal() + accept_all_requested = Signal() + # threshold_changed: the lowest score shown was changed + threshold_changed = Signal() + + # __init__: build the controls + def __init__(self) -> None: + super().__init__('Match') + self.matcher: MatchService | None = None + self.scan_button = QPushButton('Scan') + self.stop_button = QPushButton('Stop') + self.score_slider = QSlider(Qt.Orientation.Horizontal) + self.score_slider.setRange(1, _STEPS - 1) + self.score_label = QLabel() + self.status = QLabel() + self.status.setWordWrap(True) + self.accept_button = QPushButton('Accept (A)') + self.reject_button = QPushButton('Reject') + self.accept_all_button = QPushButton('Accept all shown') + for button in (self.scan_button, self.stop_button, self.accept_button, self.reject_button, self.accept_all_button): + button.setFocusPolicy(Qt.FocusPolicy.NoFocus) + self.score_slider.setFocusPolicy(Qt.FocusPolicy.NoFocus) + self.scan_button.setToolTip("Average this kind's picks into a template and look for it in every tomogram") + self.accept_button.setToolTip('Turn the selected candidate into a pick') + self.reject_button.setToolTip('Hide the selected candidate; later scans will not offer it again (Delete)') + self.accept_all_button.setToolTip('Turn every candidate shown in this tomogram into a pick') + self.score_slider.setToolTip('Lowest correlation score shown') + scan_row = QHBoxLayout() + scan_row.addWidget(self.scan_button) + scan_row.addWidget(self.stop_button) + pick_row = QHBoxLayout() + pick_row.addWidget(self.accept_button) + pick_row.addWidget(self.reject_button) + layout = QVBoxLayout(self) + layout.addLayout(scan_row) + layout.addWidget(self.status) + layout.addWidget(self.score_label) + layout.addWidget(self.score_slider) + layout.addLayout(pick_row) + layout.addWidget(self.accept_all_button) + self.scan_button.clicked.connect(self.scan_requested) + self.stop_button.clicked.connect(self.stop_requested) + self.accept_button.clicked.connect(self.accept_requested) + self.reject_button.clicked.connect(self.reject_requested) + self.accept_all_button.clicked.connect(self.accept_all_requested) + self.score_slider.valueChanged.connect(self._moved) + self.set_matcher(None, 0.2) + + # set_matcher: use a scan service and start the score slider at its minimum, or none + def set_matcher(self, matcher: MatchService | None, min_score: float) -> None: + self.matcher = matcher + self.score_slider.blockSignals(True) + self.score_slider.setValue(round(min_score * _STEPS)) + self.score_slider.blockSignals(False) + self.refresh(kind=None, picks=0, shown=0, total=0, selected=False) + + # min_score: the lowest score shown + def min_score(self) -> float: + return self.score_slider.value() / _STEPS + + # refresh: enable what applies and say what the scan is doing; shown and total count the kind's candidates in this tomogram and everywhere, picks the kind's picks + def refresh(self, *, kind: str | None, picks: int, shown: int, total: int, selected: bool) -> None: + busy = self.matcher is not None and self.matcher.busy + self.score_label.setText(f'Score at least {self.min_score():.2f}') + self.scan_button.setEnabled(self.matcher is not None and kind is not None and picks > 0 and not busy) + self.stop_button.setEnabled(busy) + self.accept_button.setEnabled(selected) + self.reject_button.setEnabled(selected) + self.accept_all_button.setEnabled(shown > 0) + self.status.setText(self._status(kind, picks, shown, total, busy)) + + # _status: status (under buttons) + def _status(self, kind: str | None, picks: int, shown: int, total: int, busy: bool) -> str: + matcher = self.matcher + if matcher is None or kind is None: + return '' + if busy: + return f'Scanning {kind}: {matcher.done} of {matcher.total} tomograms' + if matcher.last_error: + return f'Scan failed: {matcher.last_error}' + if picks == 0: + return f'Pick some {kind} to make a template' + note = f' ({len(matcher.skipped)} skipped, other pixel size)' if matcher.skipped else '' + return f'{shown} shown here, {total} in all{note}' + + # _moved: the slider moved + def _moved(self, _value: int) -> None: + self.score_label.setText(f'Score at least {self.min_score():.2f}') + self.threshold_changed.emit() diff --git a/src/pickmix/gui/panes/settings_pane.py b/src/pickmix/gui/panes/settings_pane.py index f8a1055..53ab5ba 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, Config, RingParams +from pickmix.config import HOTKEYS, Config, MatchParams, RingParams from pickmix.gui.components.collapsible import Collapsible from pickmix.gui.components.toast import notify from pickmix.gui.style import form_label, heading @@ -86,6 +86,7 @@ def __init__(self) -> None: self.cache_spin.setKeyboardTracking(False) self.hotkey_boxes = [QComboBox() for _ in range(HOTKEYS)] self.ring_form = ParamsForm(RingParams()) + self.match_form = ParamsForm(MatchParams()) general = QGroupBox('General') form = QFormLayout(general) form.addRow('Volume source', self.source_box) @@ -98,16 +99,22 @@ def __init__(self) -> None: estimator_box = QVBoxLayout(estimators) self.ring_section = Collapsible('Ring', self.ring_form) estimator_box.addWidget(self.ring_section) - self.setMinimumWidth(self.ring_form.sizeHint().width() + 48) + matching = QGroupBox('Template matching') + matching_box = QVBoxLayout(matching) + self.match_section = Collapsible('Scan', self.match_form) + matching_box.addWidget(self.match_section) + self.setMinimumWidth(max(self.ring_form.sizeHint().width(), self.match_form.sizeHint().width()) + 48) layout = QVBoxLayout(self) layout.setContentsMargins(0, 0, 0, 0) layout.addWidget(heading('Settings')) layout.addWidget(general) layout.addWidget(estimators) + layout.addWidget(matching) layout.addStretch() self.source_box.currentTextChanged.connect(self._save) self.cache_spin.valueChanged.connect(self._save) self.ring_form.changed.connect(self._save) + self.match_form.changed.connect(self._save) # set_project: show a project's settings without saving them back def set_project(self, project: Project) -> None: @@ -116,6 +123,7 @@ def set_project(self, project: Project) -> None: self.source_box.setCurrentText(project.config.source) self.cache_spin.setValue(project.config.cache_mb) self.ring_form.set(project.config.ring) + self.match_form.set(project.config.match) self._loading = False self.refresh_hotkeys() @@ -131,12 +139,12 @@ def refresh_hotkeys(self) -> None: box.setCurrentText(chosen or NO_KIND) self._loading = False - # _save: write source, slice cache and ring parameters; a value the config refuses is not written + # _save: write source, slice cache and ring & match 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(), cache_mb=self.cache_spin.value(), ring=self.ring_form.get()) + self.project.update(source=self.source_box.currentText(), cache_mb=self.cache_spin.value(), ring=self.ring_form.get(), match=self.match_form.get()) except ValidationError as e: notify(f'Not saved: {e.errors()[0]["msg"]}') return diff --git a/src/pickmix/gui/tabs/pick_tab.py b/src/pickmix/gui/tabs/pick_tab.py index 15de6a1..ab54023 100644 --- a/src/pickmix/gui/tabs/pick_tab.py +++ b/src/pickmix/gui/tabs/pick_tab.py @@ -16,11 +16,13 @@ from pickmix.geometry import Parent, match_parents from pickmix.gui.components.canvas import FitView, Marker, SliceCanvas from pickmix.gui.components.range_slider import RangeSlider -from pickmix.gui.components.status import StatusBadge, fit_status +from pickmix.gui.components.status import StatusBadge, fit_status, status_item +from pickmix.gui.panes.match_pane import MatchPane from pickmix.gui.panes.tomo_pane import TomoPane from pickmix.gui.workers import post -from pickmix.gui.style import bold_safe, eye_toggle, heading, hint, icon_button, quiet_table, tidy, title, view_toggle, white +from pickmix.gui.style import CELL_PADDING, bold_safe, eye_toggle, heading, hint, icon_button, quiet_table, tidy, title, view_toggle, white from pickmix.loader import Loader, Opened +from pickmix.matcher import MatchService from pickmix.model import Session from pickmix.timing import Timer @@ -35,6 +37,7 @@ def __init__(self) -> None: self.session: Session | None = None self.loader: Loader | None = None self.service: EstimateService | None = None + self.matcher: MatchService | None = None self._want_select: int | None = None self._matched: dict[str, dict[int, Parent]] = {} self.opened: Opened | None = None @@ -50,6 +53,7 @@ def __init__(self) -> None: self.title = title() self.help = hint('Enter: next tomogram. Shift+Enter: previous tomogram.') self.status = QLabel() + self.match_pane = MatchPane() self.tomo = TomoPane() self.tomo.setMinimumWidth(200) self.tomo_toggle = icon_button('left_panel_close', '<', 'Hide the tomogram list', size=20) @@ -85,7 +89,8 @@ def __init__(self) -> None: header.setSectionResizeMode(QHeaderView.ResizeMode.Stretch) header.setSectionResizeMode(0, QHeaderView.ResizeMode.Fixed) header.setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents) - self.table.setColumnWidth(0, 28) + # 16px icon plus the cell padding either side, or a status item's icon is clipped + self.table.setColumnWidth(0, 16 + 2 * CELL_PADDING) self.delete_button = icon_button('delete', 'Delete', 'Delete the selected pick (Delete)') self.delete_button.setEnabled(False) for button in (self.fit_button, self.fill_button, self.delete_button): @@ -117,6 +122,7 @@ def __init__(self) -> None: settings_layout.addWidget(image) settings_layout.addWidget(view) settings_layout.addWidget(kinds) + settings_layout.addWidget(self.match_pane) settings_layout.addStretch() # vertical Z scrollbar between the canvas and the settings z_column = QVBoxLayout() @@ -179,11 +185,22 @@ def __init__(self) -> None: self.canvas.marker_clicked.connect(self._clicked) self.canvas.z_scrolled.connect(lambda steps: self.z_slider.setValue(self.z_slider.value() + steps)) self.table.cellClicked.connect(self._row_clicked) + self.match_pane.scan_requested.connect(self._scan) + self.match_pane.stop_requested.connect(self._stop_scan) + self.match_pane.accept_requested.connect(self.accept_selected) + self.match_pane.reject_requested.connect(self.delete_selected) + self.match_pane.accept_all_requested.connect(self.accept_shown) + self.match_pane.threshold_changed.connect(self._on_picks) # set_session: show a session and follow its picks and navigation - def set_session(self, session: Session, loader: Loader, service: EstimateService) -> None: - self.session, self.loader, self.service = session, loader, service + def set_session(self, session: Session, loader: Loader, service: EstimateService, matcher: MatchService | None = None) -> None: + self.session, self.loader, self.service, self.matcher = session, loader, service, matcher + self.match_pane.set_matcher(matcher, session.project.config.match.min_score) session.store.listeners.append(self._on_picks) + # candidates and scan progress arrive on the scanning thread + session.candidates.listeners.append(lambda: post(self._on_picks)) + if matcher is not None: + matcher.listeners.append(lambda: post(self._on_picks)) # estimates arrive on the service's thread service.store.listeners.append(lambda: post(self._on_picks)) session.goto_listeners.append(self._on_goto) @@ -267,15 +284,53 @@ def set_tomo_visible(self, shown: bool) -> None: self.tomo_toggle.setIcon(panel) self.tomo_toggle.setToolTip('Hide the tomogram list' if shown else 'Show the tomogram list') - # delete_selected: remove the selected pick + # delete_selected: remove selected pick or reject selected candidate def delete_selected(self) -> None: - if self.session is not None and self.canvas.selected_id is not None: + if self.session is None or self.canvas.selected_id is None: + return + if self.canvas.selected_id < 0: + self.session.candidates.reject(-self.canvas.selected_id) + else: self.session.store.remove(self.canvas.selected_id) + # accept_selected: convert selected candidate into pick + def accept_selected(self) -> None: + if self.session is None or self.canvas.selected_id is None or self.canvas.selected_id >= 0: + return + picks = self.session.accept(-self.canvas.selected_id) + if picks: + self._select(picks[0].id) + + # accept_shown: convert all candidates of kind in tomogram to picks + def accept_shown(self) -> None: + if self.session is not None and self.active is not None: + found = self.session.candidates.candidates(self.session.current, self.active, min_score=self.match_pane.min_score()) + if found: + self.session.accept(*(c.id for c in found)) + # _set_active: kind new picks are made as def _set_active(self, name: str) -> None: self.active = name self.kind_buttons[name].setChecked(True) + self._refresh_match() + + # _scan: scan for active kind's template in every tomogram + def _scan(self) -> None: + if self.matcher is not None and self.active is not None: + self.matcher.start(self.active) + + # _stop_scan: end matching scan + def _stop_scan(self) -> None: + if self.matcher is not None: + self.matcher.stop() + + # _refresh_match: refresh match pane for active kind/tomogram + def _refresh_match(self) -> None: + if self.session is None: + return + kind, current = self.active, self.session.current + shown = len(self.session.candidates.candidates(current, kind, min_score=self.match_pane.min_score())) if kind else 0 + self.match_pane.refresh(kind=kind, picks=self.session.store.total(kind) if kind else 0, shown=shown, total=self.session.candidates.count(kind) if kind else 0, selected=self.canvas.selected_id is not None and self.canvas.selected_id < 0) # pick_kind: choose the kind on the i-th hotkey def pick_kind(self, i: int) -> None: @@ -471,25 +526,35 @@ def _markers(self) -> list[Marker]: continue for p in self.session.store.picks(name, kind_name): fit = self.service.store.get(p) if self.service else None - view = None if fit is None else FitView( - fit.x, fit.y, fit.z, fit.radius_a / px if kind.estimator != 'none' else None, - fit.box_a / px if fit.box_a > 0 else None, fit.z_lo, fit.z_hi, - ) + view = None if fit is None else FitView(fit.x, fit.y, fit.z, fit.radius_a / px if kind.estimator != 'none' else None, fit.box_a / px if fit.box_a > 0 else None, fit.z_lo, fit.z_hi) parent = self._matched.get(kind_name, {}).get(p.id) unmatched = kind.parent is not None and fit is not None and parent is None normal = None if parent is None else (parent.normal[0], parent.normal[1]) out.append(Marker(p.id, p.x, p.y, p.z, kind.marker_a / 2 / px, kind.colour, bool(fit and fit.flag) or unmatched, view, normal)) + for c in self.session.candidates.candidates(name, kind_name, min_score=self.match_pane.min_score()): + out.append(Marker(-c.id, c.x, c.y, c.z, kind.marker_a / 2 / px, kind.colour, candidate=True, tip=f'{kind_name} candidate, score {c.score:.2f}')) return out - # _fill_table: the current tomogram's picks of known kinds, each with its estimate status, and the selected one highlighted + # _fill_table: the current tomogram's picks of known kinds, each with its estimate status, then its shown candidates in italics with their score (negative ids); the selected row highlighted def _fill_table(self) -> None: assert self.session is not None name = self.session.current kinds = self.session.project.config.kinds picks = [p for k in kinds for p in self.session.store.picks(name, k)] self._matched = self._match() - self._row_ids = [p.id for p in picks] - self.table.setRowCount(len(picks)) + candidates = [c for k in kinds if k not in self.hidden for c in self.session.candidates.candidates(name, k, min_score=self.match_pane.min_score())] + self._row_ids = [p.id for p in picks] + [-c.id for c in candidates] + self.table.setRowCount(len(self._row_ids)) + for r, c in enumerate(candidates, start=len(picks)): + # an item, not a cell widget: a replaced cell widget can be left with no width + self.table.removeCellWidget(r, 0) + self.table.setItem(r, 0, status_item('match', f'Template match candidate, score {c.score:.2f}')) + for col, text in enumerate((c.kind, f'{c.x:.1f}', f'{c.y:.1f}', f'{c.z:.1f}', '', f'score {c.score:.2f}'), start=1): + item = QTableWidgetItem(text) + font = item.font() + font.setItalic(True) + item.setFont(font) + self.table.setItem(r, col, item) for r, p in enumerate(picks): fit = self.service.store.get(p) if self.service else None snr, flag = ('', '') if fit is None else (f'{fit.snr:.2f}', fit.flag) @@ -514,6 +579,7 @@ def _select(self, pick_id: int | None) -> None: # _sync_delete: Delete is only live while a pick is selected def _sync_delete(self) -> None: self.delete_button.setEnabled(self.canvas.selected_id is not None) + self._refresh_match() # _add: new pick of the active kind on the current slice def _add(self, x: float, y: float) -> None: @@ -528,11 +594,12 @@ def _move(self, pick_id: int, x: float, y: float) -> None: self.session.store.move(pick_id, x, y) self._select(pick_id) - # _clicked: pick selected on the canvas; show the slice it was picked on + # _clicked: pick or candidate (negative id) selected on canvas; show the slice it's on def _clicked(self, pick_id: int) -> None: assert self.session is not None self._select(pick_id) - self.z_slider.setValue(round(self.session.store.get(pick_id).z)) + z = self.session.candidates.get(-pick_id).z if pick_id < 0 else self.session.store.get(pick_id).z + self.z_slider.setValue(round(z)) # _row_clicked: pick chosen in the list def _row_clicked(self, row: int, _col: int) -> None: diff --git a/src/pickmix/match.py b/src/pickmix/match.py new file mode 100644 index 0000000..b06e4ca --- /dev/null +++ b/src/pickmix/match.py @@ -0,0 +1,158 @@ +''' +pickmix: template building and normalised cross-correlation search +''' + +# Import external dependencies +import math, numpy as np, scipy.fft +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from scipy import ndimage, signal + +# Import internal pickmix objects +from pickmix.candidates import Hit +from pickmix.config import Config, Kind +from pickmix.extract import crop_box +from pickmix.utils.io import Tomogram, load_volume + +# _SLAB: planes of a tomogram read and binned at a time +_SLAB: int = 32 + +# _CANDIDATE_POOL: peaks ranked per tomogram before the closest are thinned out +_CANDIDATE_POOL: int = 5000 + +# Cancelled: a scan was stopped part way +class Cancelled(Exception): + pass + +# Template: binned, masked mean of boxes around picks of a kind +@dataclass(frozen=True) +class Template: + kind: str + # mask times the zero-mean template (correlation kernel) + kernel: np.ndarray + # soft spherical weights of the same shape + mask: np.ndarray + # reconstruction voxels per template voxel + bin: int + # edge in binned voxels (odd so centre is a voxel) + edge: int + # reconstruction pixel size + pixel_size: float + # boxes averaged + count: int + +# template_geometry: bin/edge for a kind +def template_geometry(kind: Kind, config: Config, pixel_size: float) -> tuple[int, int]: + box_a = kind.box_a or kind.marker_a * config.ring.box_margin + box_px = max(box_a / pixel_size, 1.0) + factor = max(math.ceil(box_px / config.match.template_px), 1) + edge = max(round(box_px / factor), 3) | 1 + return factor, edge + +# soft_sphere: weight 1 inside radius, falling to 0 over edge voxels with a cosine, in cube of side n +def soft_sphere(n: int, radius: float, edge: float = 2.0) -> np.ndarray: + c = (n - 1) / 2 + z, y, x = np.ogrid[:n, :n, :n] + r = np.sqrt((z - c) ** 2 + (y - c) ** 2 + (x - c) ** 2) + t = np.clip((radius - r) / edge + 0.5, 0, 1) + return (0.5 - 0.5 * np.cos(np.pi * t)).astype(np.float32) + +# bin_cube: mean of factor x factor x factor blocks of a cube whose side is a multiple of factor +def bin_cube(cube: np.ndarray, factor: int) -> np.ndarray: + n = cube.shape[0] // factor + return cube.reshape(n, factor, n, factor, n, factor).mean(axis=(1, 3, 5)) + +# bin_volume: mean of factor-sized blocks of a (z, y, x) volume, read a slab at a time +def bin_volume(vol: np.ndarray, factor: int, *, stop: Callable[[], bool] = lambda: False) -> np.ndarray: + nz, ny, nx = (n // factor * factor for n in vol.shape) + out = np.empty((nz // factor, ny // factor, nx // factor), dtype=np.float32) + step = max(_SLAB // factor, 1) * factor + for z in range(0, nz, step): + if stop(): + raise Cancelled + end = min(z + step, nz) + block = np.asarray(vol[z:end, :ny, :nx], dtype=np.float32) + out[z // factor : end // factor] = block.reshape((end - z) // factor, factor, ny // factor, factor, nx // factor, factor).mean(axis=(1, 3, 5)) + return out + +# to_voxel: reconstruction voxel coordinate of the centre of a binned voxel +def to_voxel(index: float, factor: int) -> float: + return index * factor + (factor - 1) / 2 + +# build_template: mean of the unit-variance boxes cut at each position (tomogram, x, y, z in voxels), binned, masked and made zero-mean under the mask +def build_template( + name: str, kind: Kind, positions: list[tuple[str, float, float, float]], tomos: dict[str, Tomogram], config: Config, +) -> Template: + if not positions: + raise ValueError(f'kind {name!r}: nothing to build a template from') + sizes = {tomos[t].pixel_size for t, *_ in positions} + if len(sizes) > 1: + raise ValueError(f'kind {name!r}: one template needs one pixel size, got {sorted(sizes)!r}') + pixel_size = sizes.pop() + factor, edge = template_geometry(kind, config, pixel_size) + box = edge * factor + total = np.zeros((edge, edge, edge), dtype=np.float64) + for tomogram in dict.fromkeys(t for t, *_ in positions): + vol = load_volume(tomos[tomogram], source=config.source) + for t, x, y, z in positions: + if t == tomogram: + cube = crop_box(vol, (round(z), round(y), round(x)), box) + total += bin_cube((cube - cube.mean()) / (cube.std() or 1.0), factor) + mean = total / len(positions) + radius = kind.marker_a / 2 / (pixel_size * factor) * config.match.mask_margin + mask = soft_sphere(edge, min(radius, edge / 2)) + centred = mean - (mean * mask).sum() / mask.sum() + return Template(name, (centred * mask).astype(np.float32), mask, factor, edge, pixel_size, len(positions)) + +# score_volume: normalised cross-correlation of binned volume with template under its mask +def score_volume(vol: np.ndarray, template: Template, *, workers: int = 2) -> np.ndarray: + centred = vol - vol.mean() + n = float(template.mask.sum()) + energy = float((template.kernel**2 / np.maximum(template.mask, 1e-6)).sum()) + kernel, mask = (k[::-1, ::-1, ::-1] for k in (template.kernel, template.mask)) + with scipy.fft.set_workers(workers): + top = signal.fftconvolve(centred, kernel, mode='same') + local_sum = signal.fftconvolve(centred, mask, mode='same') + local_square = signal.fftconvolve(centred * centred, mask, mode='same') + variance = local_square - local_sum * local_sum / n + live = variance > 1e-2 * n * float(centred.var()) + score = np.zeros(centred.shape, dtype=np.float32) + score[live] = top[live] / np.sqrt(variance[live] * energy) + return score + +# find_peaks: up to limit local maxima of a score volume at or above min_score, best first, none closer than separation voxels and none within margin voxels of a face; (z, y, x) indices +def find_peaks(score: np.ndarray, *, min_score: float, separation: float, margin: int, limit: int) -> list[tuple[tuple[int, int, int], float]]: + if min(score.shape) <= 2 * margin: + return [] + inner = np.zeros_like(score) + core = tuple(slice(margin, n - margin) for n in score.shape) + inner[core] = score[core] + reach = max(round(separation), 1) + # a cheap box filter finds the maxima, then the closest are thinned out best first + maxima = (inner == ndimage.maximum_filter(inner, size=2 * reach + 1, mode='constant')) & (inner >= min_score) + found = np.argwhere(maxima) + values = inner[maxima] + order = np.argsort(-values)[:_CANDIDATE_POOL] + kept: list[tuple[tuple[int, int, int], float]] = [] + for o in order: + centre = found[o] + if all(((centre - np.array(k)) ** 2).sum() > separation**2 for k, _ in kept): + kept.append(((int(centre[0]), int(centre[1]), int(centre[2])), float(values[o]))) + if len(kept) >= limit: + break + return kept + +# scan: candidates in a binned volume, as Hits in reconstruction voxels, left out where within the separation of a known position (x, y, z in voxels) +def scan( + binned: np.ndarray, template: Template, kind: Kind, config: Config, known: Sequence[tuple[float, float, float]] = (), +) -> list[Hit]: + params = config.match + separation = params.separation * kind.marker_a / (template.pixel_size * template.bin) + peaks = find_peaks( + score_volume(binned, template, workers=params.workers), + min_score=params.min_score, separation=separation, margin=max(template.edge // 4, 1), limit=params.max_per_tomogram + len(known), + ) + hits = [Hit(to_voxel(x, template.bin), to_voxel(y, template.bin), to_voxel(z, template.bin), s) for (z, y, x), s in peaks] + near = separation * template.bin + free = [h for h in hits if all((h.x - kx) ** 2 + (h.y - ky) ** 2 + (h.z - kz) ** 2 > near**2 for kx, ky, kz in known)] + return free[: params.max_per_tomogram] diff --git a/src/pickmix/matcher.py b/src/pickmix/matcher.py new file mode 100644 index 0000000..47fc785 --- /dev/null +++ b/src/pickmix/matcher.py @@ -0,0 +1,132 @@ +''' +pickmix: background template matching over the tomograms of a project +''' + +# Import external dependencies +import numpy as np, os, threading +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor + +# Import internal pickmix objects +from pickmix.config import Config +from pickmix.estimates import EstimateStore +from pickmix.match import Cancelled, bin_volume, build_template, scan +from pickmix.model import Session +from pickmix.utils.io import Tomogram, load_volume, volume_path + +# MatchService: builds a template from picks and scans unvisited tomograms for matches +class MatchService: + # __init__: nothing runs until start + def __init__(self, session: Session, estimates: EstimateStore | None = None) -> None: + self.session = session + self.estimates = estimates + # called from the scanning thread after each tomogram and when a scan ends + self.listeners: list[Callable[[], None]] = [] + self.kind: str | None = None + self.total = 0 + self.done = 0 + self.last_error = '' + # tomograms left out because their pixel size differs from the template's + self.skipped: list[str] = [] + self._pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix='pickmix-match') + self._lock = threading.Lock() + self._generation = 0 + self._active = 0 + self._closed = False + + # busy: whether a scan is queued or running + @property + def busy(self) -> bool: + with self._lock: + return self._active > 0 + + # start: scan unvisited tomograms for matches + def start(self, kind: str) -> None: + if self._closed: + return + with self._lock: + self._generation += 1 + generation = self._generation + self._active += 1 + self.kind, self.total, self.done, self.last_error, self.skipped = kind, 0, 0, '', [] + self._notify() + self._pool.submit(self._run, kind, generation) + + # stop: end the scan after the step it is on + def stop(self) -> None: + with self._lock: + self._generation += 1 + + # shutdown: stop the thread + def shutdown(self) -> None: + self._closed = True + self.stop() + self._pool.shutdown(wait=False, cancel_futures=True) + + # _run: one scan; a failure is kept in last_error and a stop is not a failure + def _run(self, kind: str, generation: int) -> None: + try: + self._scan(kind, lambda: self._generation != generation) + except Cancelled: + pass + except Exception as e: + self.last_error = f'{kind}: {type(e).__name__}: {e}' + finally: + with self._lock: + self._active -= 1 + self._notify() + + # _positions: where the template is cut for each pick of a kind, at the fit when it has a current one + def _positions(self, kind: str) -> list[tuple[str, float, float, float]]: + out = [] + for pick in self.session.store.all(kind): + if pick.tomogram in self.session.project.tomos: + fit = self.estimates.get(pick) if self.estimates else None + out.append((pick.tomogram, *((fit.x, fit.y, fit.z) if fit else (pick.x, pick.y, pick.z)))) + return out + + # _scan: build template then scan unvisited tomograms + def _scan(self, name: str, stop: Callable[[], bool]) -> None: + project = self.session.project + config, tomos = project.config, project.tomos + kind = config.kinds[name] + template = build_template(name, kind, self._positions(name), tomos, config) + names = [n for n in self.session.names if n not in self.session.visited] + self.total = len(names) + self._notify() + for tomogram in names: + if stop(): + raise Cancelled + tomo = tomos[tomogram] + if tomogram in self.session.visited: + pass + elif abs(tomo.pixel_size / template.pixel_size - 1) > 0.01: + self.skipped.append(tomogram) + else: + binned = _binned(tomo, template.bin, config, stop) + known = [(p.x, p.y, p.z) for p in self.session.store.picks(tomogram, name)] + hits = scan(binned, template, kind, config, known) + if stop(): + raise Cancelled + reach = config.match.separation * kind.marker_a / tomo.pixel_size + self.session.candidates.replace(tomogram, name, hits, reach=reach) + self.done += 1 + self._notify() + + # _notify: tell listeners + def _notify(self) -> None: + for listener in self.listeners: + listener() + +# _binned: the tomogram binned by a factor, from out_dir/cache/binned when that is newer than the volume, else made and saved there +def _binned(tomo: Tomogram, factor: int, config: Config, stop: Callable[[], bool]) -> np.ndarray: + path = volume_path(tomo, source=config.source) + cached = config.out_dir / 'cache' / 'binned' / f'{path.stem}_x{factor}.npy' + if cached.exists() and cached.stat().st_mtime >= path.stat().st_mtime: + return np.load(cached) + binned = bin_volume(load_volume(tomo, source=config.source), factor, stop=stop) + cached.parent.mkdir(parents=True, exist_ok=True) + partial = cached.with_name(cached.stem + '.partial.npy') + np.save(partial, binned) + os.replace(partial, cached) + return binned diff --git a/src/pickmix/model.py b/src/pickmix/model.py index 898864e..397c87d 100644 --- a/src/pickmix/model.py +++ b/src/pickmix/model.py @@ -10,6 +10,7 @@ from pydantic import BaseModel # Import internal pickmix objects +from pickmix.candidates import CandidateStore from pickmix.project import Project # Pick: one hand-picked centre; x, y, z in reconstruction voxels; id lives for the session only @@ -136,12 +137,13 @@ def _changed(self) -> None: for listener in self.listeners: listener() -# Session: the store plus which tomogram is shown; navigation clamps and never wraps +# Session: pick & candidate stores and which tomogram is shown; navigation clamps and never wraps class Session: # __init__: resume at the tomogram shown last time if it is still usable def __init__(self, project: Project) -> None: self.project = project self.store = PickStore(project) + self.candidates = CandidateStore(project) self.names = list(project.tomos) self.progress_path = project.config.out_dir / 'progress.json' self.goto_listeners: list[Callable[[], None]] = [] @@ -150,16 +152,28 @@ def __init__(self, project: Project) -> None: self.index = self.names.index(saved.last) if saved.last in self.names else 0 if self.names: self._remember() + # don't include seen tomograms in template matching candidates + self.candidates.prune(self.visited - {self.current}) + + # accept: turn candidates into picks of their kind, at the candidate's position + def accept(self, *candidate_ids: int) -> list[Pick]: + found = [self.candidates.get(i) for i in candidate_ids] + picks = [self.store.add(c.tomogram, c.kind, c.x, c.y, c.z) for c in found if c.kind in self.project.config.kinds] + self.candidates.remove(*candidate_ids) + return picks # current: name of the tomogram shown @property def current(self) -> str: return self.names[self.index] - # go: show tomogram i, clamped; saves progress and notifies + # go: show tomogram i, clamped; saves progress, forgets the candidates of the tomogram left, and notifies def go(self, i: int) -> None: + left = self.current self.index = min(max(i, 0), len(self.names) - 1) self._remember() + if left != self.current: + self.candidates.prune([left]) for listener in self.goto_listeners: listener() diff --git a/tests/gui/test_components.py b/tests/gui/test_components.py index 1845510..05ff2b2 100644 --- a/tests/gui/test_components.py +++ b/tests/gui/test_components.py @@ -106,6 +106,21 @@ def test_drag_moves(self, qtbot: QtBot) -> None: pick_id, x, y = blocker.args assert pick_id == 3 and abs(x - 25) < 1.5 and abs(y - 12) < 1.5 + # test_candidate_marker: a candidate is dotted, shows its tip, and is only selected not dragged + def test_candidate_marker(self, qtbot: QtBot) -> None: + canvas = _canvas(qtbot) + canvas.set_markers([Marker(-3, 20, 15, 0, 4.0, 'yellow', candidate=True, tip='score 0.42')]) + item = canvas._items[-3] + assert item.pen().style() == Qt.PenStyle.DotLine and item.toolTip() == 'score 0.42' + start, end = _view_pos(canvas, 20, 15), _view_pos(canvas, 25, 12) + QTest.mousePress(canvas.viewport(), Qt.MouseButton.LeftButton, pos=start) + QTest.mouseMove(canvas.viewport(), end) + moved = [] + canvas.move_requested.connect(lambda *args: moved.append(args)) + with qtbot.waitSignal(canvas.marker_clicked) as blocker: + QTest.mouseRelease(canvas.viewport(), Qt.MouseButton.LeftButton, pos=end) + assert blocker.args == [-3] and moved == [] and canvas.selected_id == -3 + # test_wheel_scrolls_z: a wheel notch emits one z step def test_wheel_scrolls_z(self, qtbot: QtBot) -> None: canvas = _canvas(qtbot) diff --git a/tests/gui/test_main_window.py b/tests/gui/test_main_window.py index 54a6c36..c0fd225 100644 --- a/tests/gui/test_main_window.py +++ b/tests/gui/test_main_window.py @@ -12,6 +12,7 @@ from pytestqt.qtbot import QtBot # Import internal pickmix objects +from pickmix.candidates import Hit from pickmix.config import Kind from pickmix.gui.main_window import PROJECT_SIZE, WELCOME_SIZE, MainWindow from pickmix.project import open_project @@ -58,7 +59,6 @@ def test_no_usable_tomograms(self, qtbot: QtBot, make_config: Callable[..., Path assert window.tabs.isTabEnabled(0) and not any(window.tabs.isTabEnabled(i) for i in (1, 2, 3)) assert window.kinds_pane.kind_names() == ['cage'] - # test_review_tab_wired: the Review tab is built for an opened project and opening a pick switches tabs def test_review_tab_wired(self, qtbot: QtBot, make_config: Callable[..., Path]) -> None: window = MainWindow() @@ -80,6 +80,23 @@ def test_average_pane_wired(self, qtbot: QtBot, make_config: Callable[..., Path] window.service.shutdown() window.loader.shutdown() + # test_matcher_wired: the Pick tab gets a match service, A accepts the selected candidate, and closing the project stops the service + def test_matcher_wired(self, qtbot: QtBot, make_config: Callable[..., Path]) -> None: + window = MainWindow() + qtbot.addWidget(window) + window.project_tab.set_project(open_project(make_config())) + matcher = window.matcher + assert matcher is not None and window.pick_tab.matcher is matcher and window.pick_tab.match_pane.matcher is matcher + window.tabs.setCurrentWidget(window.pick_tab) + qtbot.waitUntil(lambda: window.pick_tab.canvas.has_image, timeout=5000) + window.session.candidates.replace('Position_1_2', 'cage', [Hit(2, 2, 1, 0.6)], reach=1) + qtbot.waitUntil(lambda: window.pick_tab.canvas.marker_ids() == [-1], timeout=5000) + window.pick_tab._clicked(-1) + window._accept_candidate() + assert window.session.store.total('cage') == 1 + window.project_tab.close_button.click() + assert window.matcher is None and matcher._closed + # test_close_project: closing re-locks the other tabs and drops the session def test_close_project(self, qtbot: QtBot, make_config: Callable[..., Path]) -> None: window = MainWindow() @@ -208,7 +225,7 @@ def test_pick_keys_are_window_wide(self, qtbot: QtBot) -> None: window = MainWindow() qtbot.addWidget(window) keys = {s.key().toString(): s for s in window.findChildren(QShortcut)} - assert all(str(i) in keys for i in range(1, 10)) and 'Del' in keys and 'Backspace' in keys + assert all(str(i) in keys for i in range(1, 10)) and 'Del' in keys and 'Backspace' in keys and 'A' in keys assert all(s.context() == Qt.ShortcutContext.WindowShortcut for s in keys.values()) # test_hotkey_and_delete_from_other_focus: with focus off the canvas, a hotkey still picks the kind and delete still removes the pick diff --git a/tests/gui/test_panes.py b/tests/gui/test_panes.py index 9187e5c..b489f9e 100644 --- a/tests/gui/test_panes.py +++ b/tests/gui/test_panes.py @@ -358,6 +358,22 @@ def test_ring_value_autosaves(self, qtbot: QtBot, make_config: Callable[..., Pat config = load_config(path) assert config.ring.radius_step_a == 50.0 and config.ring.radius_range_a == (150.0, 500.0) + # test_match_value_autosaves: a committed match setting is written, the others are kept + def test_match_value_autosaves(self, qtbot: QtBot, make_config: Callable[..., Path]) -> None: + pane, path = _settings_pane(qtbot, make_config) + assert pane.match_form.get().min_score == 0.2 and pane.match_form.isHidden() + with qtbot.waitSignal(pane.config_saved): + pane.match_form._boxes['min_score'][0].setValue(0.35) + config = load_config(path) + assert config.match.min_score == 0.35 and config.match.template_px == 32 + + # test_match_value_refused: a match setting the config refuses is not written + def test_match_value_refused(self, qtbot: QtBot, make_config: Callable[..., Path]) -> None: + pane, path = _settings_pane(qtbot, make_config) + with qtbot.assertNotEmitted(pane.config_saved): + pane.match_form._boxes['min_score'][0].setValue(1.5) + assert load_config(path).match.min_score == 0.2 + # test_no_save_button: there is nothing to press def test_no_save_button(self, qtbot: QtBot) -> None: pane = SettingsPane() diff --git a/tests/gui/test_pick_tab.py b/tests/gui/test_pick_tab.py index a947c0f..3c7003a 100644 --- a/tests/gui/test_pick_tab.py +++ b/tests/gui/test_pick_tab.py @@ -11,13 +11,18 @@ from pytestqt.qtbot import QtBot # Import internal pickmix objects +from pickmix.candidates import Hit from pickmix.config import Kind from pickmix.estimates import EstimateService from pickmix.gui.tabs.pick_tab import PickTab from pickmix.loader import Loader +from pickmix.matcher import MatchService from pickmix.model import Session from pickmix.project import open_project +# NAME: the first tomogram of the fixture project +NAME: str = 'Position_1_2' + # THREE: tomograms for navigation tests THREE: tuple[str, ...] = ('Position_1_2', 'Position_1_3', 'Other_9') @@ -294,3 +299,66 @@ def test_fit_and_fill_buttons(self, setup: Callable[..., tuple[PickTab, Session] fit_scale = tab.canvas.transform().m11() tab.fill_button.click() assert tab.canvas.transform().m11() >= fit_scale + + # test_candidates_drawn: candidates of shown kinds are drawn with negative ids, best-score filter applied + def test_candidates_drawn(self, qtbot: QtBot, setup: Callable[..., tuple[PickTab, Session]]) -> None: + tab, session = setup() + session.candidates.replace(NAME, 'cage', [Hit(2, 2, 1, 0.3), Hit(5, 4, 2, 0.6)], reach=1) + qtbot.waitUntil(lambda: len(tab.canvas.marker_ids()) == 2, timeout=5000) + assert tab.canvas.marker_ids() == [-2, -1] + tab.match_pane.score_slider.setValue(50) + assert tab.canvas.marker_ids() == [-2] + assert 'score 0.60' in tab.canvas._items[-2].toolTip() + + # test_hidden_kind_hides_candidates: hiding a kind hides its candidates too + def test_hidden_kind_hides_candidates(self, qtbot: QtBot, setup: Callable[..., tuple[PickTab, Session]]) -> None: + tab, session = setup() + session.candidates.replace(NAME, 'cage', [Hit(2, 2, 1, 0.6)], reach=1) + qtbot.waitUntil(lambda: tab.canvas.marker_ids() == [-1], timeout=5000) + tab.kind_toggles['cage'].setChecked(False) + assert tab.canvas.marker_ids() == [] + + # test_accept_selected: accepting a selected candidate converts it to a pick of its kind and selects it + def test_accept_selected(self, qtbot: QtBot, setup: Callable[..., tuple[PickTab, Session]]) -> None: + tab, session = setup() + session.candidates.replace(NAME, 'cage', [Hit(2, 2, 1, 0.6)], reach=1) + qtbot.waitUntil(lambda: tab.canvas.marker_ids() == [-1], timeout=5000) + tab._clicked(-1) + assert tab.canvas.selected_id == -1 and tab.z_slider.value() == 1 and tab.match_pane.accept_button.isEnabled() + tab.accept_selected() + (pick,) = session.store.all('cage') + assert (pick.x, pick.y, pick.z) == (2, 2, 1) and tab.canvas.selected_id == pick.id + assert session.candidates.count('cage') == 0 and tab.canvas.marker_ids() == [pick.id] + + # test_delete_rejects_candidate: delete on a selected candidate rejects it and a later scan leaves it out + def test_delete_rejects_candidate(self, qtbot: QtBot, setup: Callable[..., tuple[PickTab, Session]]) -> None: + tab, session = setup() + session.candidates.replace(NAME, 'cage', [Hit(2, 2, 1, 0.6)], reach=1) + qtbot.waitUntil(lambda: tab.canvas.marker_ids() == [-1], timeout=5000) + tab._clicked(-1) + tab.delete_selected() + assert tab.canvas.marker_ids() == [] and session.store.total() == 0 + session.candidates.replace(NAME, 'cage', [Hit(2, 2, 1, 0.6)], reach=1) + assert session.candidates.count('cage') == 0 + + # test_accept_shown: test accept all converts all candidates of a kind in active tomogram to pick (if above score) + def test_accept_shown(self, qtbot: QtBot, setup: Callable[..., tuple[PickTab, Session]]) -> None: + tab, session = setup() + session.candidates.replace(NAME, 'cage', [Hit(2, 2, 1, 0.3), Hit(5, 4, 2, 0.6)], reach=1) + qtbot.waitUntil(lambda: len(tab.canvas.marker_ids()) == 2, timeout=5000) + tab.match_pane.score_slider.setValue(50) + assert tab.match_pane.accept_all_button.isEnabled() + tab.accept_shown() + assert [(p.x, p.y) for p in session.store.all('cage')] == [(5, 4)] and session.candidates.count('cage') == 1 + + # test_scan_needs_picks: scan waits for a pick of active kind + def test_scan_needs_picks(self, qtbot: QtBot, setup: Callable[..., tuple[PickTab, Session]]) -> None: + tab, session = setup() + matcher = MatchService(session) + tab.matcher = matcher + tab.match_pane.set_matcher(matcher, 0.2) + tab._set_active('cage') + assert not tab.match_pane.scan_button.isEnabled() and 'Pick some cage' in tab.match_pane.status.text() + session.store.add(NAME, 'cage', 3, 3, 2) + assert tab.match_pane.scan_button.isEnabled() + matcher.shutdown() diff --git a/tests/test_candidates.py b/tests/test_candidates.py new file mode 100644 index 0000000..36f7d14 --- /dev/null +++ b/tests/test_candidates.py @@ -0,0 +1,123 @@ +''' +pickmix: tests for the candidate store +''' + +# Import external dependencies +from collections.abc import Callable +from pathlib import Path + +# Import internal pickmix objects +from pickmix.candidates import CandidateStore +from pickmix.candidates import Hit +from pickmix.model import Session +from pickmix.project import open_project + +# NAME: the single tomogram of the fixture project +NAME: str = 'Position_1_2' + +# THREE: three tomograms for navigation tests +THREE: tuple[str, ...] = ('Position_1_2', 'Position_1_3', 'Other_9') + +# _store: a store over a fresh project +def _store(make_config: Callable[..., Path]) -> CandidateStore: + return CandidateStore(open_project(make_config())) + +class TestCandidates: + # test_replace_and_query: candidates come back best first, filtered by kind and score + def test_replace_and_query(self, make_config: Callable[..., Path]) -> None: + store = _store(make_config) + heard: list[int] = [] + store.listeners.append(lambda: heard.append(1)) + store.replace(NAME, 'cage', [Hit(1, 1, 1, 0.3), Hit(5, 5, 1, 0.6)], reach=2) + assert [c.score for c in store.candidates(NAME)] == [0.6, 0.3] and heard == [1] + assert [c.score for c in store.candidates(NAME, min_score=0.5)] == [0.6] + assert store.candidates(NAME, 'other') == [] and store.count('cage') == 2 and store.count('cage', NAME, min_score=0.5) == 1 + + # test_replace_swaps_old_for_new: a new scan replaces the old candidates of that kind and tomogram only + def test_replace_swaps_old_for_new(self, make_config: Callable[..., Path]) -> None: + store = _store(make_config) + store.replace(NAME, 'cage', [Hit(1, 1, 1, 0.3)], reach=2) + store.replace(NAME, 'yellow', [Hit(2, 2, 2, 0.4)], reach=2) + store.replace(NAME, 'cage', [Hit(7, 7, 1, 0.5)], reach=2) + assert [(c.kind, c.x) for c in store.candidates(NAME)] == [('cage', 7), ('yellow', 2)] or {(c.kind, c.x) for c in store.candidates(NAME)} == {('cage', 7), ('yellow', 2)} + + # test_reject_is_remembered: a rejected candidate is hidden and a later scan does not offer its place again + def test_reject_is_remembered(self, make_config: Callable[..., Path]) -> None: + store = _store(make_config) + store.replace(NAME, 'cage', [Hit(1, 1, 1, 0.3), Hit(6, 6, 1, 0.4)], reach=2) + first = store.candidates(NAME)[1] + store.reject(first.id) + assert [c.x for c in store.candidates(NAME)] == [6] + store.replace(NAME, 'cage', [Hit(1.5, 1, 1, 0.35), Hit(6, 6, 1, 0.4)], reach=2) + assert [c.x for c in store.candidates(NAME)] == [6] + + # test_save_and_load: candidates, with rejected ones, survive a restart + def test_save_and_load(self, make_config: Callable[..., Path]) -> None: + store = _store(make_config) + store.replace(NAME, 'cage', [Hit(1, 1, 1, 0.3), Hit(6, 6, 1, 0.4)], reach=2) + store.reject(store.candidates(NAME)[0].id) + again = CandidateStore(store.project) + assert [(c.x, c.score) for c in again.candidates(NAME)] == [(1, 0.3)] + assert again.count('cage') == 1 and sum(c.status == 'rejected' for c in again._items.values()) == 1 + + # test_rename_remove_clear: kind renames follow, kinds can be dropped and everything can be cleared + def test_rename_remove_clear(self, make_config: Callable[..., Path]) -> None: + store = _store(make_config) + store.replace(NAME, 'cage', [Hit(1, 1, 1, 0.3)], reach=2) + store.replace(NAME, 'yellow', [Hit(2, 2, 2, 0.4)], reach=2) + store.rename_kind('cage', 'big') + assert store.count('big') == 1 and store.count('cage') == 0 + store.remove_kind('big') + assert store.count('big') == 0 and store.count('yellow') == 1 + store.clear() + assert store.count('yellow') == 0 + + # test_accept_makes_picks: accepted candidates become picks of their kind and leave the store + def test_accept_makes_picks(self, make_config: Callable[..., Path]) -> None: + session = Session(open_project(make_config())) + session.candidates.replace(NAME, 'cage', [Hit(1, 2, 3, 0.5)], reach=2) + (candidate,) = session.candidates.candidates(NAME) + (pick,) = session.accept(candidate.id) + assert (pick.tomogram, pick.kind, pick.x, pick.y, pick.z) == (NAME, 'cage', 1, 2, 3) + assert session.candidates.count('cage') == 0 and session.store.total('cage') == 1 + + # test_accept_skips_unknown_kind: a candidate whose kind has gone is dropped, not turned into a pick + def test_accept_skips_unknown_kind(self, make_config: Callable[..., Path]) -> None: + session = Session(open_project(make_config())) + session.candidates.replace(NAME, 'gone', [Hit(1, 2, 3, 0.5)], reach=2) + assert session.accept(session.candidates.candidates(NAME)[0].id) == [] and session.store.total() == 0 + + # test_prune: every candidate in the given tomograms goes, rejected ones too, and listeners hear it once + def test_prune(self, make_config: Callable[..., Path]) -> None: + store = CandidateStore(open_project(make_config(names=THREE))) + store.replace(THREE[0], 'cage', [Hit(1, 1, 1, 0.3), Hit(6, 6, 1, 0.4)], reach=2) + store.reject(store.candidates(THREE[0])[0].id) + store.replace(THREE[1], 'cage', [Hit(2, 2, 2, 0.5)], reach=2) + heard: list[int] = [] + store.listeners.append(lambda: heard.append(1)) + store.prune([THREE[0]]) + store.prune([THREE[2]]) + assert store.candidates(THREE[0]) == [] and store.count('cage') == 1 and heard == [1] + assert len(CandidateStore(store.project)._items) == 1 + + # test_leaving_a_tomogram_prunes_it: moving on forgets the candidates of the tomogram left, and keeps the others + def test_leaving_a_tomogram_prunes_it(self, make_config: Callable[..., Path]) -> None: + session = Session(open_project(make_config(names=THREE))) + session.candidates.replace(THREE[0], 'cage', [Hit(1, 1, 1, 0.3)], reach=2) + session.candidates.replace(THREE[1], 'cage', [Hit(2, 2, 2, 0.5)], reach=2) + session.go(1) + assert session.candidates.candidates(THREE[0]) == [] and len(session.candidates.candidates(THREE[1])) == 1 + session.go(1) + assert len(session.candidates.candidates(THREE[1])) == 1 + + # test_visited_pruned_on_start: reopening a project forgets candidates of visited tomograms except the one resumed at + def test_visited_pruned_on_start(self, make_config: Callable[..., Path]) -> None: + project = open_project(make_config(names=THREE)) + first = Session(project) + first.go(1) + first.candidates.replace(THREE[0], 'cage', [Hit(1, 1, 1, 0.3)], reach=2) + first.candidates.replace(THREE[1], 'cage', [Hit(2, 2, 2, 0.5)], reach=2) + first.candidates.replace(THREE[2], 'cage', [Hit(3, 3, 3, 0.6)], reach=2) + again = Session(project) + assert again.current == THREE[1] + assert [c.tomogram for c in again.candidates._items.values()] == [THREE[1], THREE[2]] diff --git a/tests/test_config.py b/tests/test_config.py index 47ffd0d..00f104f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -8,7 +8,7 @@ from pydantic import ValidationError # Import internal pickmix objects -from pickmix.config import Config, Kind, RingParams, load_config, save_config +from pickmix.config import Config, Kind, MatchParams, RingParams, load_config, save_config # _config: a config with absolute paths, so it survives a save and load unchanged, and the given kinds def _config(**kinds: Kind) -> Config: @@ -72,6 +72,16 @@ def test_ring_params_documented(self) -> None: for name, field in RingParams.model_fields.items(): assert field.title and field.description, name + # test_match_params_documented: every match setting has a title and a description for the settings form + def test_match_params_documented(self) -> None: + assert all(f.title and f.description for f in MatchParams.model_fields.values()) + + # test_match_round_trip: match settings survive a save and load + def test_match_round_trip(self, tmp_path: Path) -> None: + config = _config().replace(match=MatchParams(template_px=24, min_score=0.3)) + save_config(config, tmp_path / 'config.toml') + assert load_config(tmp_path / 'config.toml').match == MatchParams(template_px=24, min_score=0.3) + # test_default_hotkeys: no hotkeys means the first nine kinds in table order, padded with unused keys def test_default_hotkeys(self) -> None: assert _config().hotkeys == ['cage', *[''] * 8] diff --git a/tests/test_match.py b/tests/test_match.py new file mode 100644 index 0000000..802e499 --- /dev/null +++ b/tests/test_match.py @@ -0,0 +1,138 @@ +''' +pickmix: tests for template building and correlation search +''' + +# Import external dependencies +import mrcfile, pytest, numpy as np +from pathlib import Path + +# Import internal pickmix objects +from pickmix.candidates import Hit +from pickmix.config import Config, Kind +from pickmix.match import Cancelled, bin_cube, bin_volume, build_template, find_peaks, scan, score_volume, soft_sphere, template_geometry, to_voxel +from pickmix.utils.io import Tomogram + +# PIXEL: Angstrom per voxel of the test tomograms +PIXEL: float = 10.0 + +# _blob: a smooth 3D bump of the given radius in voxels, in a cube of odd side +def _blob(radius: float, side: int) -> np.ndarray: + c = side // 2 + z, y, x = np.ogrid[:side, :side, :side] + return np.exp(-(((z - c) ** 2 + (y - c) ** 2 + (x - c) ** 2) / (2 * (radius / 2) ** 2))).astype(np.float32) + +# _tomogram: tomogram with planted blobs +def _tomogram(tmp_path: Path, name: str, centres: list[tuple[int, int, int]], *, seed: int = 0) -> Tomogram: + rng = np.random.default_rng(seed) + vol = rng.normal(0, 1, (48, 96, 128)).astype(np.float32) + blob = _blob(8, 25) + for x, y, z in centres: + vol[z - 12 : z + 13, y - 12 : y + 13, x - 12 : x + 13] += 4 * blob + path = tmp_path / f'{name}.mrc' + with mrcfile.new(path) as mrc: + mrc.set_data(vol) + mrc.voxel_size = PIXEL + return Tomogram(name, name, path, None, PIXEL, vol.shape, PIXEL, (128, 96, 48)) + +# _matches: whether every centre has a hit within tol voxels and no hit is left over +def _matches(hits: list[Hit], centres: list[tuple[int, int, int]], tol: float = 2.0) -> bool: + close = lambda h, c: max(abs(h.x - c[0]), abs(h.y - c[1]), abs(h.z - c[2])) <= tol + return len(hits) == len(centres) and all(any(close(h, c) for h in hits) for c in centres) + +# _config: example config with kind that has 160 Angstrom marker +def _config(tmp_path: Path) -> tuple[Config, Kind]: + kind = Kind(marker_a=160, box_a=240) + return Config(raw_dir=tmp_path, tomograms_star=tmp_path / 't.star', out_dir=tmp_path / 'out', kinds={'blob': kind}), kind + +class TestMatch: + # test_template_geometry: template is binned until it is about template_px across, with an odd edge + def test_template_geometry(self, tmp_path: Path) -> None: + config, kind = _config(tmp_path) + assert template_geometry(kind, config, PIXEL) == (1, 25) + big = kind.model_copy(update={'box_a': 1500.0}) + factor, edge = template_geometry(big, config, PIXEL) + assert factor == 5 and edge % 2 == 1 and abs(edge * factor - 150) <= factor + + # test_template_geometry_default_box: a kind with no box takes the marker diameter times the ring margin + def test_template_geometry_default_box(self, tmp_path: Path) -> None: + config, kind = _config(tmp_path) + assert template_geometry(kind.model_copy(update={'box_a': 0.0}), config, PIXEL) == (1, 21) + + # test_soft_sphere: one inside the radius, zero outside, symmetric + def test_soft_sphere(self) -> None: + m = soft_sphere(21, 6) + assert m[10, 10, 10] == 1 and m[0, 0, 0] == 0 and np.allclose(m, m[::-1, ::-1, ::-1]) + + # test_bin_cube: blocks are averaged + def test_bin_cube(self) -> None: + cube = np.arange(64, dtype=np.float32).reshape(4, 4, 4) + assert bin_cube(cube, 2).shape == (2, 2, 2) and bin_cube(cube, 2)[0, 0, 0] == cube[:2, :2, :2].mean() + + # test_bin_volume: blocks are averaged, remainders dropped, and a stop request raises Cancelled + def test_bin_volume(self) -> None: + vol = np.arange(5 * 6 * 7, dtype=np.float32).reshape(5, 6, 7) + out = bin_volume(vol, 2) + assert out.shape == (2, 3, 3) and out[1, 2, 2] == vol[2:4, 4:6, 4:6].mean() + with pytest.raises(Cancelled): + bin_volume(vol, 2, stop=lambda: True) + + # test_to_voxel: centre of a binned voxel in reconstruction voxels + def test_to_voxel(self) -> None: + assert to_voxel(0, 1) == 0 and to_voxel(2, 4) == 9.5 + + # test_template_is_masked_and_centred: kernel is zero outside the mask and sums to zero + def test_template_is_masked_and_centred(self, tmp_path: Path) -> None: + config, kind = _config(tmp_path) + tomo = _tomogram(tmp_path, 'a', [(40, 40, 24)]) + template = build_template('blob', kind, [('a', 40, 40, 24)], {'a': tomo}, config) + assert template.kernel.shape == (25, 25, 25) and template.count == 1 and template.bin == 1 + assert abs(float(template.kernel.sum())) < 1e-3 and template.kernel[0, 0, 0] == 0 + + # test_build_needs_picks_and_one_pixel_size: no positions or mixed pixel sizes are refused + def test_build_needs_picks_and_one_pixel_size(self, tmp_path: Path) -> None: + config, kind = _config(tmp_path) + a, b = _tomogram(tmp_path, 'a', [(40, 40, 24)]), _tomogram(tmp_path, 'b', []) + with pytest.raises(ValueError, match='nothing'): + build_template('blob', kind, [], {'a': a}, config) + odd = Tomogram('b', 'b', b.raw_path, None, 5.0, b.shape, 5.0, b.size_xyz) + with pytest.raises(ValueError, match='one pixel size'): + build_template('blob', kind, [('a', 40, 40, 24), ('b', 40, 40, 24)], {'a': a, 'b': odd}, config) + + # test_scan_finds_planted_blobs: a template from one tomogram finds blobs planted elsewhere, at the right places + def test_scan_finds_planted_blobs(self, tmp_path: Path) -> None: + config, kind = _config(tmp_path) + source = _tomogram(tmp_path, 'a', [(40, 40, 24), (80, 60, 24)]) + target = _tomogram(tmp_path, 'b', [(30, 30, 24), (90, 50, 20), (60, 70, 28)], seed=1) + template = build_template('blob', kind, [('a', 40, 40, 24), ('a', 80, 60, 24)], {'a': source}, config) + with mrcfile.mmap(target.raw_path, mode='r') as mrc: + hits = scan(np.asarray(mrc.data), template, kind, config) + assert _matches(hits[:3], [(30, 30, 24), (90, 50, 20), (60, 70, 28)]) and all(h.score > 0.3 for h in hits[:3]) + # the three true sites outscore everything else + assert len(hits) == 3 or hits[3].score < hits[2].score + + # test_scan_skips_known: positions already picked are not offered again + def test_scan_skips_known(self, tmp_path: Path) -> None: + config, kind = _config(tmp_path) + source = _tomogram(tmp_path, 'a', [(40, 40, 24)]) + target = _tomogram(tmp_path, 'b', [(30, 30, 24), (90, 50, 20)], seed=1) + template = build_template('blob', kind, [('a', 40, 40, 24)], {'a': source}, config) + with mrcfile.mmap(target.raw_path, mode='r') as mrc: + hits = scan(np.asarray(mrc.data), template, kind, config, [(30.0, 30.0, 24.0)]) + assert _matches(hits[:1], [(90, 50, 20)]) + + # test_score_flat_volume: a flat volume scores zero everywhere + def test_score_flat_volume(self, tmp_path: Path) -> None: + config, kind = _config(tmp_path) + tomo = _tomogram(tmp_path, 'a', [(40, 40, 24)]) + template = build_template('blob', kind, [('a', 40, 40, 24)], {'a': tomo}, config) + assert not score_volume(np.ones((30, 40, 50), dtype=np.float32), template).any() + + # test_find_peaks: best first, separated, away from the faces, above the minimum + def test_find_peaks(self) -> None: + score = np.zeros((20, 20, 20), dtype=np.float32) + score[10, 10, 10], score[10, 10, 12], score[5, 5, 5], score[0, 0, 0], score[15, 15, 15] = 0.9, 0.8, 0.5, 0.99, 0.1 + peaks = find_peaks(score, min_score=0.2, separation=4, margin=2, limit=10) + assert [p for p, _ in peaks] == [(10, 10, 10), (5, 5, 5)] + assert [round(s, 1) for _, s in peaks] == [0.9, 0.5] + assert len(find_peaks(score, min_score=0.2, separation=4, margin=2, limit=1)) == 1 + assert find_peaks(score, min_score=0.2, separation=4, margin=12, limit=10) == [] diff --git a/tests/test_matcher.py b/tests/test_matcher.py new file mode 100644 index 0000000..82002b7 --- /dev/null +++ b/tests/test_matcher.py @@ -0,0 +1,135 @@ +''' +pickmix: tests for background template matching service +''' + +# Import external dependencies +import mrcfile, time, numpy as np, pytest +from collections.abc import Callable, Iterator +from pathlib import Path + +# Import internal pickmix objects +from pickmix.config import Config, Kind, save_config +from pickmix.matcher import MatchService +from pickmix.model import Session +from pickmix.project import open_project + +# STAR: tomograms.star for tomograms of 128 x 96 x 48 voxels at 5.36 Angstrom, binned 4 from 1.34 +STAR = '''data_global + +loop_ +_rlnTomoName #1 +_rlnTomoTiltSeriesPixelSize #2 +_rlnTomoTomogramBinning #3 +_rlnTomoSizeX #4 +_rlnTomoSizeY #5 +_rlnTomoSizeZ #6 +_rlnTomoReconstructedTomogram #7 +''' + +# _blob: a smooth 3D bump in a cube of odd side +def _blob(side: int) -> np.ndarray: + c = side // 2 + z, y, x = np.ogrid[:side, :side, :side] + return np.exp(-(((z - c) ** 2 + (y - c) ** 2 + (x - c) ** 2) / (2 * 4.0**2))).astype(np.float32) + +# blobs: session with two tomograms containing planted blobs +@pytest.fixture +def blobs(tmp_path: Path) -> Iterator[tuple[Session, dict[str, list[tuple[int, int, int]]]]]: + (tmp_path / 'raw').mkdir() + planted = {'A': [(40, 40, 24), (80, 60, 24)], 'B': [(30, 30, 24), (90, 50, 20)]} + rows = [] + for i, (name, centres) in enumerate(planted.items()): + vol = np.random.default_rng(i).normal(0, 1, (48, 96, 128)).astype(np.float32) + for x, y, z in centres: + vol[z - 12 : z + 13, y - 12 : y + 13, x - 12 : x + 13] += 4 * _blob(25) + with mrcfile.new(tmp_path / 'raw' / f'nt_raw4_{name}_stack_Vol.mrc') as mrc: + mrc.set_data(vol) + mrc.voxel_size = 5.36 + rows.append(f'{name} 1.34 4.0 512 384 192 Tomograms/job006/tomograms/{name}_stack_Vol.mrc') + (tmp_path / 't.star').write_text(STAR + '\n'.join(rows) + '\n') + config = Config(raw_dir=tmp_path / 'raw', tomograms_star=tmp_path / 't.star', out_dir=tmp_path / 'out', kinds={'blob': Kind(marker_a=100, box_a=150)}) + save_config(config, tmp_path / 'config.toml') + yield Session(open_project(tmp_path / 'config.toml')), planted + +# _wait: block until service is idle +def _wait(service: MatchService) -> None: + deadline = time.time() + 20 + while service.busy and time.time() < deadline: + time.sleep(0.02) + assert not service.busy + +class TestMatcher: + # test_scan_finds_blobs: picks in visited tomogram lead to candidates at the planted blobs of the unvisited one + def test_scan_finds_blobs(self, blobs: tuple[Session, dict]) -> None: + session, planted = blobs + for x, y, z in planted['A']: + session.store.add('A', 'blob', x, y, z) + service = MatchService(session) + service.start('blob') + _wait(service) + assert service.last_error == '' and service.done == service.total == 1 and service.skipped == [] + found = session.candidates.candidates('B', 'blob', min_score=0.3) + assert len(found) == 2 + for x, y, z in planted['B']: + assert any(max(abs(c.x - x), abs(c.y - y), abs(c.z - z)) <= 2 for c in found) + # A is visited and not scanned + assert session.candidates.count('blob', 'A') == 0 + service.shutdown() + + # test_binned_volumes_cached: the binned volume is written once and reused + def test_binned_volumes_cached(self, blobs: tuple[Session, dict]) -> None: + session, planted = blobs + session.store.add('A', 'blob', *planted['A'][0]) + service = MatchService(session) + service.start('blob') + _wait(service) + files = sorted((session.project.config.out_dir / 'cache' / 'binned').glob('*.npy')) + assert len(files) == 1 and not any('partial' in f.name for f in files) + stamp = files[0].stat().st_mtime_ns + service.start('blob') + _wait(service) + assert files[0].stat().st_mtime_ns == stamp + service.shutdown() + + # test_no_picks_is_an_error: a kind with no picks reports why instead of scanning + def test_no_picks_is_an_error(self, blobs: tuple[Session, dict]) -> None: + session, _ = blobs + service = MatchService(session) + service.start('blob') + _wait(service) + assert 'nothing to build a template from' in service.last_error and session.candidates.count('blob') == 0 + service.shutdown() + + # test_nothing_to_scan_when_all_visited: with every tomogram visited the scan finishes with no candidates + def test_nothing_to_scan_when_all_visited(self, blobs: tuple[Session, dict]) -> None: + session, planted = blobs + session.store.add('A', 'blob', *planted['A'][0]) + session.go(1) + service = MatchService(session) + service.start('blob') + _wait(service) + assert service.last_error == '' and service.total == 0 and session.candidates.count('blob') == 0 + service.shutdown() + + # test_stop: stopping ends the scan without an error and leaves what was found + def test_stop(self, blobs: tuple[Session, dict]) -> None: + session, planted = blobs + session.store.add('A', 'blob', *planted['A'][0]) + service = MatchService(session) + service.start('blob') + service.stop() + _wait(service) + assert service.last_error == '' and service.done <= service.total + service.shutdown() + + # test_progress_is_announced: listeners hear the start, each tomogram and the end + def test_progress_is_announced(self, blobs: tuple[Session, dict]) -> None: + session, planted = blobs + session.store.add('A', 'blob', *planted['A'][0]) + service = MatchService(session) + seen: list[tuple[int, int]] = [] + service.listeners.append(lambda: seen.append((service.done, service.total))) + service.start('blob') + _wait(service) + assert seen[0] == (0, 0) and (1, 1) in seen + service.shutdown() diff --git a/uv.lock b/uv.lock index 9d722e0..ad5b491 100644 --- a/uv.lock +++ b/uv.lock @@ -137,7 +137,7 @@ wheels = [ [[package]] name = "pickmix" -version = "0.3.2" +version = "0.4.0" source = { editable = "." } dependencies = [ { name = "imageio" },