feat: stack-viewer v2 using ndv as baseclass

.pre-commit-config.yaml

@@ -27,5 +27,5 @@ repos:
       - id: mypy
         files: "^src/"
         additional_dependencies:
-          - pymmcore-plus >=0.9.0
+          - pymmcore-plus >=0.10.0
           - useq-schema >=0.4.7

examples/mda_viewer.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from pymmcore_plus import CMMCorePlus, configure_logging
+from qtpy import QtWidgets
+from useq import MDASequence
+
+from pymmcore_widgets._stack_viewer_v2._mda_viewer import MDAViewer
+
+configure_logging(stderr_level="WARNING")
+
+mmcore = CMMCorePlus.instance()
+mmcore.loadSystemConfiguration()
+mmcore.defineConfig("Channel", "DAPI", "Camera", "Mode", "Artificial Waves")
+mmcore.defineConfig("Channel", "DAPI", "Camera", "StripeWidth", "1")
+mmcore.defineConfig("Channel", "FITC", "Camera", "Mode", "Artificial Waves")
+mmcore.defineConfig("Channel", "FITC", "Camera", "StripeWidth", "4")
+
+sequence = MDASequence(
+    channels=({"config": "DAPI", "exposure": 1}, {"config": "FITC", "exposure": 1}),
+    stage_positions=[(0, 0), (1, 1)],
+    z_plan={"range": 9, "step": 0.4},
+    time_plan={"interval": 0.2, "loops": 4},
+    # grid_plan={"rows": 2, "columns": 1},
+)
+
+
+qapp = QtWidgets.QApplication([])
+v = MDAViewer()
+v.show()
+
+mmcore.run_mda(sequence, output=v.data)
+qapp.exec()

examples/mda_viewer_queue.py

@@ -0,0 +1,24 @@
+import sys
+from queue import Queue
+
+from pymmcore_plus import CMMCorePlus
+from qtpy import QtWidgets
+from useq import MDAEvent
+
+from pymmcore_widgets._stack_viewer_v2._mda_viewer import MDAViewer
+
+app = QtWidgets.QApplication(sys.argv)
+mmcore = CMMCorePlus.instance()
+mmcore.loadSystemConfiguration()
+
+canvas = MDAViewer()
+canvas.show()
+
+q = Queue()
+mmcore.run_mda(iter(q.get, None), output=canvas.data)
+for i in range(10):
+    for c in range(2):
+        q.put(MDAEvent(index={"t": i, "c": c}, exposure=1))
+q.put(None)
+
+app.exec()

examples/stack_viewer.py → examples/stack_viewer_v1.py

@@ -6,7 +6,7 @@
 from qtpy import QtWidgets
 from useq import MDASequence
 
-from pymmcore_widgets.experimental import StackViewer
+from pymmcore_widgets._stack_viewer_v1 import StackViewer
 
 size = 1028
 
@@ -20,7 +20,7 @@
 
 sequence = MDASequence(
     channels=(
-        # {"config": "DAPI", "exposure": 10},
+        {"config": "DAPI", "exposure": 10},
         # {"config": "FITC", "exposure": 1},
         {"config": "Cy5", "exposure": 1},
     ),

pyproject.toml

@@ -50,14 +50,23 @@ dependencies = [
     'fonticon-materialdesignicons6',
     'pymmcore-plus[cli] >=0.10.2',
     'qtpy >=2.0',
-    'superqt[quantity] >=0.5.3',
+    'superqt[quantity] >=0.6.5',
     'useq-schema >=0.4.7',
+    'ndv >=0.0.3',
 ]
 
 # extras
 # https://peps.python.org/pep-0621/#dependencies-optional-dependencies
 [project.optional-dependencies]
-test = ["pytest>=6.0", "pytest-cov", "pytest-qt", "PyYAML", "vispy", "cmap", "zarr"]
+test = [
+    "pytest>=6.0",
+    "pytest-cov",
+    "pytest-qt",
+    "PyYAML",
+    "vispy",
+    "cmap",
+    "zarr",
+]
 pyqt5 = ["PyQt5"]
 pyside2 = ["PySide2"]
 pyqt6 = ["PyQt6"]
@@ -129,6 +138,7 @@ docstring-code-format = true
 [tool.pytest.ini_options]
 minversion = "6.0"
 testpaths = ["tests"]
+markers = ["allow_leaks: mark test to allow widget leaks"]
 filterwarnings = [
     "error",
     "ignore:distutils Version classes are deprecated",
@@ -174,4 +184,4 @@ ignore = [
 ]
 
 [tool.typos.default]
-extend-ignore-identifiers-re = ["(?i)ome"]
+extend-ignore-identifiers-re = ["(?i)nd2?.*", "(?i)ome"]

src/pymmcore_widgets/_stack_viewer/_stack_viewer.py → src/pymmcore_widgets/_stack_viewer_v1/_stack_viewer.py

@@ -464,11 +464,9 @@ def _reload_position(self) -> None:
         self.cmap_names = self.qt_settings.value("cmaps", ["gray", "cyan", "magenta"])
 
     def _collapse_view(self) -> None:
+        w, h = self.img_size
         view_rect = (
-            (
-                self.view_rect[0][0] - self.img_size[0] / 2,
-                self.view_rect[0][1] + self.img_size[1] / 2,
-            ),
+            (self.view_rect[0][0] - w / 2, self.view_rect[0][1] + h / 2),
             self.view_rect[1],
         )
         self.view.camera.rect = view_rect

src/pymmcore_widgets/_stack_viewer_v2/__init__.py

@@ -0,0 +1,3 @@
+from ._mda_viewer import MDAViewer
+
+__all__ = ["MDAViewer"]

src/pymmcore_widgets/_stack_viewer_v2/_data_wrapper.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import warnings
+from contextlib import suppress
+from typing import TYPE_CHECKING, cast
+
+from ndv import DataWrapper
+from pymmcore_plus.mda.handlers import TensorStoreHandler
+
+if TYPE_CHECKING:
+    from collections.abc import Hashable, Mapping
+    from pathlib import Path
+    from typing import Any, TypeGuard
+
+    import numpy as np
+    from ndv import Indices
+    from pymmcore_plus.mda.handlers._5d_writer_base import _5DWriterBase
+
+
+class MMTensorStoreWrapper(DataWrapper["TensorStoreHandler"]):
+    def sizes(self) -> Mapping[Hashable, int]:
+        with suppress(Exception):
+            return self.data.current_sequence.sizes  # type: ignore
+        return {}
+
+    def guess_channel_axis(self) -> Hashable | None:
+        return "c"
+
+    @classmethod
+    def supports(cls, obj: Any) -> TypeGuard[TensorStoreHandler]:
+        return isinstance(obj, TensorStoreHandler)
+
+    def isel(self, indexers: Indices) -> np.ndarray:
+        return self.data.isel({str(k): v for k, v in indexers.items()})  # type: ignore [no-any-return]
+
+    def save_as_zarr(self, save_loc: str | Path) -> None:
+        if (store := self.data.store) is None:
+            return
+        import tensorstore as ts
+
+        new_spec = store.spec().to_json()
+        new_spec["kvstore"] = {"driver": "file", "path": str(save_loc)}
+        new_ts = ts.open(new_spec, create=True).result()
+        new_ts[:] = store.read().result()
+
+
+class MM5DWriter(DataWrapper["_5DWriterBase"]):
+    def guess_channel_axis(self) -> Hashable | None:
+        return "c"
+
+    @classmethod
+    def supports(cls, obj: Any) -> TypeGuard[_5DWriterBase]:
+        try:
+            from pymmcore_plus.mda.handlers._5d_writer_base import _5DWriterBase
+        except ImportError:
+            from pymmcore_plus.mda.handlers import OMETiffWriter, OMEZarrWriter
+
+            _5DWriterBase = (OMETiffWriter, OMEZarrWriter)  # type: ignore
+        if isinstance(obj, _5DWriterBase):
+            return True
+        return False
+
+    def save_as_zarr(self, save_loc: str | Path) -> None:
+        import zarr
+        from pymmcore_plus.mda.handlers import OMEZarrWriter
+
+        if isinstance(self._data, OMEZarrWriter):
+            zarr.copy_store(self._data.group.store, zarr.DirectoryStore(save_loc))
+        raise NotImplementedError(f"Cannot save {type(self._data)} data to Zarr.")
+
+    def isel(self, indexers: Indices) -> np.ndarray:
+        p_index = indexers.get("p", 0)
+        if isinstance(p_index, slice):
+            warnings.warn("Cannot slice over position index", stacklevel=2)  # TODO
+            p_index = p_index.start
+        p_index = cast(int, p_index)
+
+        try:
+            sizes = [*list(self._data.position_sizes[p_index]), "y", "x"]
+        except IndexError as e:
+            raise IndexError(
+                f"Position index {p_index} out of range for "
+                f"{len(self._data.position_sizes)}"
+            ) from e
+
+        data = self._data.position_arrays[self._data.get_position_key(p_index)]
+        full = slice(None, None)
+        index = tuple(indexers.get(k, full) for k in sizes)
+        return data[index]  # type: ignore [no-any-return]

src/pymmcore_widgets/_stack_viewer_v2/_mda_viewer.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Mapping
+
+import superqt
+import useq
+from ndv import DataWrapper, NDViewer
+from pymmcore_plus.mda.handlers import TensorStoreHandler
+from qtpy.QtWidgets import QFileDialog, QPushButton, QWidget
+from superqt.iconify import QIconifyIcon
+
+# this import is necessary so that ndv can find our custom DataWrapper
+from . import _data_wrapper  # noqa: F401
+
+if TYPE_CHECKING:
+    from pymmcore_plus.mda.handlers._5d_writer_base import _5DWriterBase
+    from qtpy.QtWidgets import QWidget
+
+
+class MDAViewer(NDViewer):
+    """StackViewer specialized for pymmcore-plus MDA acquisitions."""
+
+    _data: _5DWriterBase
+
+    def __init__(
+        self,
+        datastore: _5DWriterBase | TensorStoreHandler | None = None,
+        *,
+        parent: QWidget | None = None,
+    ):
+        if datastore is None:
+            datastore = TensorStoreHandler()
+
+        # patch the frameReady method to call the superframeReady method
+        # AFTER handling the event
+        self._superframeReady = getattr(datastore, "frameReady", None)
+        if callable(self._superframeReady):
+            datastore.frameReady = self._patched_frame_ready  # type: ignore
+        else:  # pragma: no cover
+            warnings.warn(
+                "MDAViewer: datastore does not have a frameReady method to patch, "
+                "are you sure this is a valid data handler?",
+                stacklevel=2,
+            )
+
+        super().__init__(datastore, parent=parent, channel_axis="c")
+        self._save_btn = SaveButton(self._data_wrapper)
+        self._btns.addWidget(self._save_btn)
+        self.dims_sliders.set_locks_visible(True)
+        self._channel_names: dict[int, str] = {}
+
+    def _patched_frame_ready(self, *args: Any) -> None:
+        self._superframeReady(*args)  # type: ignore
+        if len(args) >= 2 and isinstance(e := args[1], useq.MDAEvent):
+            self._on_frame_ready(e)
+
+    @superqt.ensure_main_thread  # type: ignore
+    def _on_frame_ready(self, event: useq.MDAEvent) -> None:
+        c = event.index.get(self._channel_axis)
+        if c not in self._channel_names and c is not None and event.channel:
+            self._channel_names[c] = event.channel.config
+        self.set_current_index(event.index)
+
+    def _get_channel_name(self, index: Mapping) -> str:
+        if self._channel_axis in index:
+            if name := self._channel_names.get(index[self._channel_axis]):
+                return name
+        c = index.get(self._channel_axis, 0)
+        return f"Ch {c}"
+
+
+class SaveButton(QPushButton):
+    def __init__(
+        self,
+        data_wrapper: DataWrapper,
+        parent: QWidget | None = None,
+    ):
+        super().__init__(parent=parent)
+        self.setIcon(QIconifyIcon("mdi:content-save"))
+        self.clicked.connect(self._on_click)
+
+        self._data_wrapper = data_wrapper
+        self._last_loc = str(Path.home())
+
+    def _on_click(self) -> None:
+        self._last_loc, _ = QFileDialog.getSaveFileName(
+            self, "Choose destination", str(self._last_loc), ""
+        )
+        suffix = Path(self._last_loc).suffix
+        if suffix in (".zarr", ".ome.zarr", ""):
+            self._data_wrapper.save_as_zarr(self._last_loc)
+        else:
+            raise ValueError(f"Unsupported file format: {self._last_loc}")

src/pymmcore_widgets/experimental.py

@@ -1,3 +1,4 @@
-from ._stack_viewer import StackViewer
+from ._stack_viewer_v1 import StackViewer
+from ._stack_viewer_v2 import MDAViewer
 
-__all__ = ["StackViewer"]
+__all__ = ["StackViewer", "MDAViewer"]