diff --git a/.github/black.yml b/.github/black.yml new file mode 100644 index 0000000..09b2a0f --- /dev/null +++ b/.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 \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 7573393..7e88e45 100644 --- a/pyproject.toml +++ b/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" \ No newline at end of file +"Bug Tracker" = "https://github.com/py4dstem/py4D-browser/issues" + +[tool.pyright] +venv = "py4dstem" \ No newline at end of file diff --git a/src/py4D_browser/main_window.py b/src/py4D_browser/main_window.py index 79a5918..3717388 100644 --- a/src/py4D_browser/main_window.py +++ b/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]) - diff --git a/src/py4D_browser/update_views.py b/src/py4D_browser/update_views.py index d40508d..c4dba23 100644 --- a/src/py4D_browser/update_views.py +++ b/src/py4D_browser/update_views.py @@ -2,7 +2,7 @@ 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): @@ -10,7 +10,12 @@ def update_real_space_view(self, reset=False): 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,7 +101,7 @@ 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}]") @@ -107,6 +109,11 @@ def update_real_space_view(self, reset=False): 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() diff --git a/src/py4D_browser/utils.py b/src/py4D_browser/utils.py index 50bb5b3..ed95364 100644 --- a/src/py4D_browser/utils.py +++ b/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 \ No newline at end of file + 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