Add virtual image FFT pane

.github/black.yml

@@ -0,0 +1,14 @@
+name: Check code style
+
+on:
+  push:
+    branches: [ "dev" ]
+  pull_request:
+    branches: [ "dev" ]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: psf/black@stable

pyproject.toml

@@ -17,6 +17,7 @@ classifiers = [
 ]
 dependencies = [
   "py4dstem >= 0.14.3",
+  "emdfile >= 0.0.11",
   "numpy >= 1.19",
   "matplotlib >= 3.2.2",
   "PyQt5 >= 5.10",
@@ -28,4 +29,7 @@ py4DGUI = "py4D_browser.runGUI:launch"
 
 [project.urls]
 "Homepage" = "https://github.com/py4dstem/py4D-browser"
-"Bug Tracker" = "https://github.com/py4dstem/py4D-browser/issues"
+"Bug Tracker" = "https://github.com/py4dstem/py4D-browser/issues"
+
+[tool.pyright]
+venv = "py4dstem"

src/py4D_browser/main_window.py

@@ -1,26 +1,14 @@
-from PyQt5.QtCore import Qt
+from PyQt5 import QtCore, QtGui
 from PyQt5.QtWidgets import (
     QApplication,
-    QLabel,
     QMainWindow,
     QWidget,
     QMenu,
     QAction,
-    QFileDialog,
-    QVBoxLayout,
     QHBoxLayout,
-    QFrame,
-    QPushButton,
-    QScrollArea,
-    QCheckBox,
-    QLineEdit,
-    QRadioButton,
-    QButtonGroup,
-    QDesktopWidget,
-    QMessageBox,
+    QSplitter,
     QActionGroup,
 )
-from PyQt5 import QtGui
 
 import pyqtgraph as pg
 import numpy as np
@@ -81,6 +69,10 @@ def __init__(self, argv):
 
         self.show()
 
+        # If a file was passed on the command line, open it
+        if len(argv) > 1:
+            self.load_file(argv[1])
+
     def setup_menus(self):
         self.menu_bar = self.menuBar()
 
@@ -229,7 +221,7 @@ def setup_menus(self):
 
         detector_point_action = QAction("&Point", self)
         detector_point_action.setCheckable(True)
-        detector_point_action.setChecked(True) # Default
+        detector_point_action.setChecked(True)  # Default
         detector_point_action.triggered.connect(self.update_diffraction_detector)
         detector_shape_group.addAction(detector_point_action)
         self.detector_shape_menu.addAction(detector_point_action)
@@ -263,7 +255,9 @@ def setup_views(self):
         self.diffraction_space_widget.addItem(self.diffraction_space_view_text)
 
         # Create virtual detector ROI selector
-        self.virtual_detector_point = pg_point_roi(self.diffraction_space_widget.getView())
+        self.virtual_detector_point = pg_point_roi(
+            self.diffraction_space_widget.getView()
+        )
         self.virtual_detector_point.sigRegionChanged.connect(
             self.update_real_space_view
         )
@@ -300,9 +294,29 @@ def setup_views(self):
         self.diffraction_space_widget.dropEvent = self.dropEvent
         self.real_space_widget.dropEvent = self.dropEvent
 
+        # Set up the FFT window.
+        self.fft_widget = pg.ImageView()
+        self.fft_widget.setImage(np.zeros((512, 512)))
+
+        # Name and return
+        self.fft_widget.setWindowTitle("FFT of Virtual Image")
+        self.fft_widget.addItem(pg.TextItem("FFT", (200, 200, 200), None, (0, 1)))
+
+        self.fft_widget.setAcceptDrops(True)
+        self.fft_widget.dragEnterEvent = self.dragEnterEvent
+        self.fft_widget.dropEvent = self.dropEvent
+
         layout = QHBoxLayout()
         layout.addWidget(self.diffraction_space_widget, 1)
-        layout.addWidget(self.real_space_widget, 1)
+
+        # add a resizeable layout for the vimg and FFT
+        rightside = QSplitter()
+        rightside.addWidget(self.real_space_widget)
+        rightside.addWidget(self.fft_widget)
+        rightside.setOrientation(QtCore.Qt.Vertical)
+        rightside.setStretchFactor(0, 2)
+        layout.addWidget(rightside, 1)
+
         widget = QWidget()
         widget.setLayout(layout)
         self.setCentralWidget(widget)
@@ -319,4 +333,3 @@ def dropEvent(self, event):
         if len(files) == 1:
             print(f"Reieving dropped file: {files[0]}")
             self.load_file(files[0])
-

src/py4D_browser/update_views.py

@@ -2,15 +2,20 @@
 import numpy as np
 import py4DSTEM
 
-from py4D_browser.utils import pg_point_roi
+from py4D_browser.utils import pg_point_roi, make_detector
 
 
 def update_real_space_view(self, reset=False):
     scaling_mode = self.vimg_scaling_group.checkedAction().text().replace("&", "")
     assert scaling_mode in ["Linear", "Log", "Square Root"], scaling_mode
 
     detector_shape = self.detector_shape_group.checkedAction().text().replace("&", "")
-    assert detector_shape in ["Point", "Rectangular", "Circle", "Annulus"], detector_shape
+    assert detector_shape in [
+        "Point",
+        "Rectangular",
+        "Circle",
+        "Annulus",
+    ], detector_shape
 
     detector_mode = self.detector_mode_group.checkedAction().text().replace("&", "")
     assert detector_mode in [
@@ -55,37 +60,34 @@ def update_real_space_view(self, reset=False):
             mask[slice_x, slice_y] = True
 
     elif detector_shape == "Circle":
-        (slice_x, slice_y), _ = self.virtual_detector_roi.getArraySlice(
-            self.datacube.data[0, 0, :, :], self.diffraction_space_widget.getImageItem()
-        )
-        x0 = (slice_x.start + slice_x.stop) / 2.0
-        y0 = (slice_y.start + slice_y.stop) / 2.0
-        R = (slice_y.stop - slice_y.start) / 2.0
+        R = self.virtual_detector_roi.size()[0] / 2.0
 
-        self.diffraction_space_view_text.setText(f"[({x0},{y0}),{R}]")
+        x0 = self.virtual_detector_roi.pos()[0] + R
+        y0 = self.virtual_detector_roi.pos()[1] + R
 
-        mask = py4DSTEM.datacube.virtualimage.DataCubeVirtualImager.make_detector(
+        self.diffraction_space_view_text.setText(f"[({x0:.0f},{y0:.0f}),{R:.0f}]")
+
+        mask = make_detector(
             (self.datacube.Q_Nx, self.datacube.Q_Ny), "circle", ((x0, y0), R)
         )
     elif detector_shape == "Annulus":
-        (slice_x, slice_y), _ = self.virtual_detector_roi_outer.getArraySlice(
-            self.datacube.data[0, 0, :, :], self.diffraction_space_widget.getImageItem()
-        )
-        x0 = (slice_x.start + slice_x.stop) / 2.0
-        y0 = (slice_y.start + slice_y.stop) / 2.0
-        R_outer = (slice_y.stop - slice_y.start) / 2.0
+        inner_pos = self.virtual_detector_roi_inner.pos()
+        inner_size = self.virtual_detector_roi_inner.size()
+        R_inner = inner_size[0] / 2.0
+        x0 = inner_pos[0] + R_inner
+        y0 = inner_pos[1] + R_inner
 
-        (slice_ix, slice_iy), _ = self.virtual_detector_roi_inner.getArraySlice(
-            self.datacube.data[0, 0, :, :], self.diffraction_space_widget.getImageItem()
-        )
-        R_inner = (slice_iy.stop - slice_iy.start) / 2.0
+        outer_size = self.virtual_detector_roi_outer.size()
+        R_outer = outer_size[0] / 2.0
 
-        if R_inner == R_outer:
+        if R_inner <= R_outer:
             R_inner -= 1
 
-        self.diffraction_space_view_text.setText(f"[({x0},{y0}),({R_inner},{R_outer})]")
+        self.diffraction_space_view_text.setText(
+            f"[({x0:.0f},{y0:.0f}),({R_inner:.0f},{R_outer:.0f})]"
+        )
 
-        mask = py4DSTEM.datacube.virtualimage.DataCubeVirtualImager.make_detector(
+        mask = make_detector(
             (self.datacube.Q_Nx, self.datacube.Q_Ny),
             "annulus",
             ((x0, y0), (R_inner, R_outer)),
@@ -99,14 +101,19 @@ def update_real_space_view(self, reset=False):
         # Normalize coordinates
         xc = np.clip(xc, 0, self.datacube.Q_Nx - 1)
         yc = np.clip(yc, 0, self.datacube.Q_Ny - 1)
-        vimg = self.datacube.data[: ,: , xc, yc]
+        vimg = self.datacube.data[:, :, xc, yc]
 
         self.diffraction_space_view_text.setText(f"[{xc},{yc}]")
 
     else:
         raise ValueError("Detector shape not recognized")
 
     if mask is not None:
+        # For debugging masks:
+        # self.diffraction_space_widget.setImage(
+        #     mask.T, autoLevels=True, autoRange=True
+        # )
+        mask = mask.astype(np.float32)
         vimg = np.zeros((self.datacube.R_Nx, self.datacube.R_Ny))
         iterator = py4DSTEM.tqdmnd(self.datacube.R_Nx, self.datacube.R_Ny, disable=True)
 
@@ -155,6 +162,11 @@ def update_real_space_view(self, reset=False):
         raise ValueError("Mode not recognized")
     self.real_space_widget.setImage(new_view.T, autoLevels=True)
 
+    # Update FFT view
+    fft = np.abs(np.fft.fftshift(np.fft.fft2(new_view))) ** 0.5
+    levels = (np.min(fft), np.percentile(fft, 99.9))
+    self.fft_widget.setImage(fft.T, autoLevels=False, levels=levels, autoRange=reset)
+
 
 def update_diffraction_space_view(self, reset=False):
     scaling_mode = self.diff_scaling_group.checkedAction().text().replace("&", "")
@@ -204,7 +216,9 @@ def update_diffraction_detector(self):
 
     # Remove existing detector
     if hasattr(self, "virtual_detector_point"):
-        self.diffraction_space_widget.view.scene().removeItem(self.virtual_detector_point)
+        self.diffraction_space_widget.view.scene().removeItem(
+            self.virtual_detector_point
+        )
     if hasattr(self, "virtual_detector_roi"):
         self.diffraction_space_widget.view.scene().removeItem(self.virtual_detector_roi)
     if hasattr(self, "virtual_detector_roi_inner"):
@@ -218,7 +232,9 @@ def update_diffraction_detector(self):
 
     # Rectangular detector
     if detector_shape == "Point":
-        self.virtual_detector_point = pg_point_roi(self.diffraction_space_widget.getView())
+        self.virtual_detector_point = pg_point_roi(
+            self.diffraction_space_widget.getView()
+        )
         self.virtual_detector_point.sigRegionChanged.connect(
             self.update_real_space_view
         )
@@ -279,11 +295,7 @@ def update_diffraction_detector(self):
         )
 
     else:
-        raise ValueError(
-            "Unknown detector shape! Got: {}".format(
-                detector_shape
-            )
-        )
+        raise ValueError("Unknown detector shape! Got: {}".format(detector_shape))
 
     self.update_real_space_view()
 

src/py4D_browser/utils.py

@@ -1,4 +1,6 @@
 import pyqtgraph as pg
+import numpy as np
+
 
 def pg_point_roi(view_box):
     """
@@ -11,4 +13,48 @@ def pg_point_roi(view_box):
     h.update()
     view_box.addItem(circ_roi)
     circ_roi.removeHandle(0)
-    return circ_roi
+    return circ_roi
+
+
+def make_detector(shape: tuple, mode: str, geometry) -> np.ndarray:
+    match mode, geometry:
+        case ["point", (qx, qy)]:
+            mask = np.zeros(shape, dtype=np.bool_)
+            mask[qx, qy] = True
+        case ["point", geom]:
+            raise ValueError(
+                f"Point detector shape must be specified as (qx,qy), not {geom}"
+            )
+
+        case [("circle" | "circular"), ((qx, qy), r)]:
+            ix, iy = np.indices(shape)
+            mask = np.hypot(ix - qx, iy - qy) <= r
+        case [("circle" | "circular"), geom]:
+            raise ValueError(
+                f"Circular detector shape must be specified as ((qx,qy),r), not {geom}"
+            )
+
+        case [("annulus" | "annular"), ((qx, qy), (ri, ro))]:
+            ix, iy = np.indices(shape)
+            ir = np.hypot(ix - qx, iy - qy)
+            mask = np.logical_and(ir >= ri, ir <= ro)
+        case [("annulus" | "annular"), geom]:
+            raise ValueError(
+                f"Annular detector shape must be specified as ((qx,qy),(ri,ro)), not {geom}"
+            )
+
+        case [("rectangle" | "square" | "rectangular"), (xmin, xmax, ymin, ymax)]:
+            mask = np.zeros(shape, dtype=np.bool_)
+            mask[xmin:xmax, ymin:ymax] = True
+        case [("rectangle" | "square" | "rectangular"), geom]:
+            raise ValueError(
+                f"Rectangular detector shape must be specified as (xmin,xmax,ymin,ymax), not {geom}"
+            )
+
+        case ["mask", mask_arr]:
+            mask = mask_arr
+
+        case unknown:
+            raise ValueError(f"mode and geometry not understood: {unknown}")
+
+    return mask