Merge pull request #25 from sezelt/dev

More intuitive tooltips and manual tcBF

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "py4D_browser"
-version = "1.1.2"
+version = "1.1.3"
 authors = [
   { name="Steven Zeltmann", email="[email protected]" },
 ]

src/py4D_browser/dialogs.py

@@ -1,4 +1,4 @@
-from py4DSTEM import DataCube, data
+from py4DSTEM import DataCube, data, tqdmnd
 import pyqtgraph as pg
 import numpy as np
 from PyQt5.QtWidgets import QFrame, QPushButton, QApplication, QLabel
@@ -16,7 +16,7 @@
     QGridLayout,
     QCheckBox,
 )
-from py4D_browser.utils import make_detector
+from py4D_browser.utils import make_detector, StatusBarWriter
 
 
 class ResizeDialog(QDialog):
@@ -301,19 +301,23 @@ def __init__(self, parent):
         params_box.setLayout(params_layout)
 
         params_layout.addWidget(QLabel("Rotation [deg]"), 0, 0, Qt.AlignRight)
-        rotation_box = QLineEdit()
-        rotation_box.setValidator(QDoubleValidator())
-        params_layout.addWidget(rotation_box, 0, 1)
+        self.rotation_box = QLineEdit()
+        self.rotation_box.setValidator(QDoubleValidator())
+        params_layout.addWidget(self.rotation_box, 0, 1)
 
         params_layout.addWidget(QLabel("Transpose x/y"), 1, 0, Qt.AlignRight)
-        transpose_box = QCheckBox()
-        params_layout.addWidget(transpose_box, 1, 1)
+        self.transpose_box = QCheckBox()
+        params_layout.addWidget(self.transpose_box, 1, 1)
 
         params_layout.addWidget(QLabel("Max Shift [px]"), 2, 0, Qt.AlignRight)
         self.max_shift_box = QLineEdit()
         self.max_shift_box.setValidator(QDoubleValidator())
         params_layout.addWidget(self.max_shift_box, 2, 1)
 
+        params_layout.addWidget(QLabel("Pad Images"), 3, 0, Qt.AlignRight)
+        self.pad_checkbox = QCheckBox()
+        params_layout.addWidget(self.pad_checkbox, 3, 1)
+
         button_layout = QHBoxLayout()
         button_layout.addStretch()
         cancel_button = QPushButton("Cancel")
@@ -371,7 +375,7 @@ def reconstruct(self):
             self.parent.statusBar().showMessage("Max Shift must be specified")
             return
 
-        rotation = float(self.rotation_box.text() or 0.0)
+        rotation = np.radians(float(self.rotation_box.text() or 0.0))
         transpose = self.transpose_box.checkState()
         max_shift = float(self.max_shift_box.text())
 
@@ -389,7 +393,6 @@ def reconstruct(self):
         # unrotated shifts in scan pixels
         shifts_pix_x = pix_coord_x / np.max(q_pix * mask) * max_shift
         shifts_pix_y = pix_coord_y / np.max(q_pix * mask) * max_shift
-        # shifts_pix = np.
 
         R = np.array(
             [
@@ -401,3 +404,60 @@ def reconstruct(self):
 
         if transpose:
             R = T @ R
+
+        shifts_pix = np.stack([shifts_pix_x, shifts_pix_y], axis=2) @ R
+        shifts_pix_x, shifts_pix_y = shifts_pix[..., 0], shifts_pix[..., 1]
+
+        # generate image to accumulate reconstruction
+        pad = self.pad_checkbox.checkState()
+        pad_width = int(
+            np.maximum(np.abs(shifts_pix_x).max(), np.abs(shifts_pix_y).max())
+        )
+
+        reconstruction = (
+            np.zeros((datacube.R_Nx + 2 * pad_width, datacube.R_Ny + 2 * pad_width))
+            if pad
+            else np.zeros((datacube.R_Nx, datacube.R_Ny))
+        )
+
+        qx = np.fft.fftfreq(reconstruction.shape[0])
+        qy = np.fft.fftfreq(reconstruction.shape[1])
+
+        qx_operator, qy_operator = np.meshgrid(qx, qy, indexing="ij")
+        qx_operator = qx_operator * -2.0j * np.pi
+        qy_operator = qy_operator * -2.0j * np.pi
+
+        # loop over images and shift
+        for mx, my in tqdmnd(
+            *mask.shape,
+            desc="Shifting images",
+            file=StatusBarWriter(self.parent.statusBar()),
+            mininterval=1.0,
+        ):
+            if mask[mx, my]:
+                img_raw = datacube.data[:, :, mx, my]
+
+                if pad:
+                    img = np.zeros_like(reconstruction) + img_raw.mean()
+                    img[
+                        pad_width : img_raw.shape[0] + pad_width,
+                        pad_width : img_raw.shape[1] + pad_width,
+                    ] = img_raw
+                else:
+                    img = img_raw
+
+                reconstruction += np.real(
+                    np.fft.ifft2(
+                        np.fft.fft2(img)
+                        * np.exp(
+                            qx_operator * shifts_pix_x[mx, my]
+                            + qy_operator * shifts_pix_y[mx, my]
+                        )
+                    )
+                )
+
+        # crop away padding so the image lines up with the original
+        if pad:
+            reconstruction = reconstruction[pad_width:-pad_width, pad_width:-pad_width]
+
+        self.parent.set_virtual_image(reconstruction, reset=True)

src/py4D_browser/main_window.py

@@ -10,6 +10,7 @@
     QActionGroup,
     QLabel,
     QToolTip,
+    QPushButton,
 )
 
 from matplotlib.backend_bases import tools
@@ -331,6 +332,10 @@ def setup_menus(self):
         self.detector_menu = QMenu("&Detector Response", self)
         self.menu_bar.addMenu(self.detector_menu)
 
+        detector_mode_separator = QAction("Diffraction", self)
+        detector_mode_separator.setDisabled(True)
+        self.detector_menu.addAction(detector_mode_separator)
+
         detector_mode_group = QActionGroup(self)
         detector_mode_group.setExclusive(True)
         self.detector_mode_group = detector_mode_group
@@ -372,6 +377,33 @@ def setup_menus(self):
         detector_mode_group.addAction(detector_iCoM)
         self.detector_menu.addAction(detector_iCoM)
 
+        # Detector Response for realspace selector
+        self.detector_menu.addSeparator()
+        rs_detector_mode_separator = QAction("Virtual Image", self)
+        rs_detector_mode_separator.setDisabled(True)
+        self.detector_menu.addAction(rs_detector_mode_separator)
+
+        realspace_detector_mode_group = QActionGroup(self)
+        realspace_detector_mode_group.setExclusive(True)
+        self.realspace_detector_mode_group = realspace_detector_mode_group
+
+        detector_integrating_action = QAction("&Integrating", self)
+        detector_integrating_action.setCheckable(True)
+        detector_integrating_action.setChecked(True)
+        detector_integrating_action.triggered.connect(
+            partial(self.update_diffraction_space_view, True)
+        )
+        realspace_detector_mode_group.addAction(detector_integrating_action)
+        self.detector_menu.addAction(detector_integrating_action)
+
+        detector_maximum_action = QAction("&Maximum", self)
+        detector_maximum_action.setCheckable(True)
+        detector_maximum_action.triggered.connect(
+            partial(self.update_diffraction_space_view, True)
+        )
+        realspace_detector_mode_group.addAction(detector_maximum_action)
+        self.detector_menu.addAction(detector_maximum_action)
+
         # Detector Shape Menu
         self.detector_shape_menu = QMenu("Detector &Shape", self)
         self.menu_bar.addMenu(self.detector_shape_menu)
@@ -472,11 +504,11 @@ def setup_menus(self):
         tcBF_action_manual = QAction("tcBF (Manual)...", self)
         tcBF_action_manual.triggered.connect(self.reconstruct_tcBF_manual)
         self.processing_menu.addAction(tcBF_action_manual)
-        tcBF_action_manual.setEnabled(False)
 
         tcBF_action_auto = QAction("tcBF (Automatic)", self)
         tcBF_action_auto.triggered.connect(self.reconstruct_tcBF_auto)
         self.processing_menu.addAction(tcBF_action_auto)
+        # tcBF_action_auto.setEnabled(False)
 
         # Help menu
         self.help_menu = QMenu("&Help", self)
@@ -490,7 +522,6 @@ def setup_views(self):
         # Set up the diffraction space window.
         self.diffraction_space_widget = pg.ImageView()
         self.diffraction_space_widget.setImage(np.zeros((512, 512)))
-        self.diffraction_space_view_text = QLabel("Slice")
 
         self.diffraction_space_widget.setMouseTracking(True)
 
@@ -515,7 +546,6 @@ def setup_views(self):
         # Set up the real space window.
         self.real_space_widget = pg.ImageView()
         self.real_space_widget.setImage(np.zeros((512, 512)))
-        self.real_space_view_text = QLabel("Scan Position")
 
         # Add point selector connected to displayed diffraction pattern
         self.real_space_point_selector = pg_point_roi(self.real_space_widget.getView())
@@ -576,12 +606,35 @@ def setup_views(self):
         self.fft_widget.getView().setMenuEnabled(False)
 
         # Setup Status Bar
-        self.realspace_statistics_text = QLabel("Image Stats")
-        self.diffraction_statistics_text = QLabel("Diffraction Stats")
-        self.statusBar().addPermanentWidget(VLine())
-        self.statusBar().addPermanentWidget(self.realspace_statistics_text)
+        self.stats_button = QPushButton("Statistics")
+        self.stats_menu = QMenu()
+
+        self.realspace_title = QAction("Virtual Image")
+        self.realspace_title.setDisabled(False)
+        self.stats_menu.addAction(self.realspace_title)
+        self.realspace_statistics_actions = [QAction("") for i in range(5)]
+        for a in self.realspace_statistics_actions:
+            self.stats_menu.addAction(a)
+
+        self.stats_menu.addSeparator()
+
+        self.diffraction_title = QAction("Diffraction")
+        self.diffraction_title.setDisabled(False)
+        self.stats_menu.addAction(self.diffraction_title)
+        self.diffraction_statistics_actions = [QAction("") for i in range(5)]
+        for a in self.diffraction_statistics_actions:
+            self.stats_menu.addAction(a)
+
+        self.stats_button.setMenu(self.stats_menu)
+
+        self.cursor_value_text = QLabel("")
+        self.diffraction_space_view_text = QLabel("Slice")
+        self.real_space_view_text = QLabel("Scan Position")
+
+        # self.statusBar().addPermanentWidget(VLine())
+        self.statusBar().addPermanentWidget(self.cursor_value_text)
         self.statusBar().addPermanentWidget(VLine())
-        self.statusBar().addPermanentWidget(self.diffraction_statistics_text)
+        self.statusBar().addPermanentWidget(self.stats_button)
         self.statusBar().addPermanentWidget(VLine())
         self.statusBar().addPermanentWidget(self.diffraction_space_view_text)
         self.statusBar().addPermanentWidget(VLine())

src/py4D_browser/menu_actions.py

@@ -280,6 +280,7 @@ def reconstruct_tcBF_auto(self):
 
     # do tcBF!
     self.statusBar().showMessage("Reconstructing... (This may take a while)")
+    self.app.processEvents()
 
     tcBF = py4DSTEM.process.phase.Parallax(
         energy=300e3,