Update main to version 0.99999

pyproject.toml

@@ -4,9 +4,9 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "py4D_browser"
-version = "0.9999"
+version = "0.99999"
 authors = [
-  { name="Steven Zeltmann", email="steven.zeltmann@berkeley.edu" },
+  { name="Steven Zeltmann", email="steven.zeltmann@lbl.gov" },
 ]
 description = "A 4D-STEM data browser built on py4DSTEM."
 readme = "README.md"
@@ -18,6 +18,7 @@ classifiers = [
 dependencies = [
   "py4dstem >= 0.14.3",
   "emdfile >= 0.0.11",
+  "h5py",
   "numpy >= 1.19",
   "matplotlib >= 3.2.2",
   "PyQt5 >= 5.10",

src/py4D_browser/empad2_reader.py

@@ -0,0 +1,70 @@
+import empad2
+from PyQt5.QtWidgets import QFileDialog, QMessageBox
+import numpy as np
+
+
+def set_empad2_sensor(self, sensor_name):
+    self.empad2_calibrations = empad2.load_calibration_data(sensor=sensor_name)
+
+
+def load_empad2_background(self):
+    if self.empad2_calibrations is not None:
+        filename = raw_file_dialog(self)
+        self.empad2_background = empad2.load_background(
+            filepath=filename, calibration_data=self.empad2_calibrations
+        )
+    else:
+        QMessageBox.warning(
+            self, "No calibrations loaded!", "Please select a sensor first"
+        )
+
+
+def load_empad2_dataset(self):
+    if self.empad2_calibrations is not None:
+        dummy_data = False
+        if self.empad2_background is None:
+            continue_wo_bkg = QMessageBox.question(
+                self,
+                "Load without background?",
+                "Background data has not been loaded. Do you want to continue loading data?",
+            )
+            if continue_wo_bkg == QMessageBox.No:
+                return
+            else:
+                self.empad2_background = {
+                    "even": np.zeros((128, 128), dtype=np.float32),
+                    "odd": np.zeros((128, 128), dtype=np.float32),
+                }
+                dummy_data = True
+
+        filename = raw_file_dialog(self)
+        self.datacube = empad2.load_dataset(
+            filename, self.empad2_background, self.empad2_calibrations
+        )
+
+        if dummy_data:
+            self.empad2_background = None
+
+        self.update_diffraction_space_view(reset=True)
+        self.update_real_space_view(reset=True)
+
+        self.setWindowTitle(filename)
+
+    else:
+        QMessageBox.warning(
+            self, "No calibrations loaded!", "Please select a sensor first"
+        )
+
+
+def raw_file_dialog(browser):
+    filename = QFileDialog.getOpenFileName(
+        browser,
+        "Open EMPAD-G2 Data",
+        "",
+        "EMPAD-G2 Data (*.raw);;Any file(*)",
+    )
+    if filename is not None and len(filename[0]) > 0:
+        return filename[0]
+    else:
+        print("File was invalid, or something?")
+        raise ValueError("Could not read file")

src/py4D_browser/main_window.py

@@ -15,6 +15,7 @@
 
 from functools import partial
 from pathlib import Path
+import importlib
 
 from py4D_browser.utils import pg_point_roi
 
@@ -34,16 +35,28 @@ class DataViewer(QMainWindow):
         load_data_bin,
         load_data_mmap,
         show_file_dialog,
+        get_savefile_name,
+        export_datacube,
+        export_virtual_image,
     )
 
     from py4D_browser.update_views import (
         update_diffraction_space_view,
         update_real_space_view,
+        update_realspace_detector,
         update_diffraction_detector,
         update_annulus_pos,
         update_annulus_radii,
     )
 
+    HAS_EMPAD2 = importlib.util.find_spec("empad2") is not None
+    if HAS_EMPAD2:
+        from py4D_browser.empad2_reader import (
+            set_empad2_sensor,
+            load_empad2_background,
+            load_empad2_dataset,
+        )
+
     def __init__(self, argv):
         super().__init__()
         # Define this as the QApplication object
@@ -80,6 +93,10 @@ def setup_menus(self):
         self.file_menu = QMenu("&File", self)
         self.menu_bar.addMenu(self.file_menu)
 
+        import_label = QAction("Import", self)
+        import_label.setDisabled(True)
+        self.file_menu.addAction(import_label)
+
         self.load_auto_action = QAction("&Load Data...", self)
         self.load_auto_action.triggered.connect(self.load_data_auto)
         self.file_menu.addAction(self.load_auto_action)
@@ -92,6 +109,63 @@ def setup_menus(self):
         self.load_binned_action.triggered.connect(self.load_data_bin)
         self.file_menu.addAction(self.load_binned_action)
 
+        self.file_menu.addSeparator()
+
+        export_label = QAction("Export", self)
+        export_label.setDisabled(True)
+        self.file_menu.addAction(export_label)
+
+        # Submenu to export datacube
+        datacube_export_menu = QMenu("Export Datacube", self)
+        self.file_menu.addMenu(datacube_export_menu)
+        for method in ["Raw float32", "py4DSTEM HDF5", "Plain HDF5"]:
+            menu_item = datacube_export_menu.addAction(method)
+            menu_item.triggered.connect(partial(self.export_datacube, method))
+
+        # Submenu to export virtual image
+        vimg_export_menu = QMenu("Export Virtual Image", self)
+        self.file_menu.addMenu(vimg_export_menu)
+        for method in ["PNG", "TIFF", "TIFF (raw)"]:
+            menu_item = vimg_export_menu.addAction(method)
+            menu_item.triggered.connect(
+                partial(self.export_virtual_image, method, "image")
+            )
+
+        # Submenu to export diffraction
+        vdiff_export_menu = QMenu("Export Diffraction Pattern", self)
+        self.file_menu.addMenu(vdiff_export_menu)
+        for method in ["PNG", "TIFF", "TIFF (raw)"]:
+            menu_item = vdiff_export_menu.addAction(method)
+            menu_item.triggered.connect(
+                partial(self.export_virtual_image, method, "diffraction")
+            )
+
+        # EMPAD2 menu
+        if self.HAS_EMPAD2:
+            self.empad2_calibrations = None
+            self.empad2_background = None
+
+            self.empad2_menu = QMenu("&EMPAD-G2", self)
+            self.menu_bar.addMenu(self.empad2_menu)
+
+            sensor_menu = self.empad2_menu.addMenu("&Sensor")
+            calibration_action_group = QActionGroup(self)
+            calibration_action_group.setExclusive(True)
+            from empad2 import SENSORS
+
+            for name, sensor in SENSORS.items():
+                menu_item = sensor_menu.addAction(sensor["display-name"])
+                calibration_action_group.addAction(menu_item)
+                menu_item.setCheckable(True)
+                menu_item.triggered.connect(partial(self.set_empad2_sensor, name))
+
+            self.empad2_menu.addAction("Load &Background...").triggered.connect(
+                self.load_empad2_background
+            )
+            self.empad2_menu.addAction("Load &Dataset...").triggered.connect(
+                self.load_empad2_dataset
+            )
+
         # Scaling Menu
         self.scaling_menu = QMenu("&Scaling", self)
         self.menu_bar.addMenu(self.scaling_menu)
@@ -205,11 +279,11 @@ def setup_menus(self):
         detector_mode_group.addAction(detector_CoM_angle)
         self.detector_menu.addAction(detector_CoM_angle)
 
-        # detector_iCoM = QAction("i&CoM", self)
-        # detector_iCoM.setCheckable(True)
-        # detector_iCoM.triggered.connect(partial(self.update_real_space_view, True))
-        # detector_mode_group.addAction(detector_iCoM)
-        # self.detector_menu.addAction(detector_iCoM)
+        detector_iCoM = QAction("i&CoM", self)
+        detector_iCoM.setCheckable(True)
+        detector_iCoM.triggered.connect(partial(self.update_real_space_view, True))
+        detector_mode_group.addAction(detector_iCoM)
+        self.detector_menu.addAction(detector_iCoM)
 
         # Detector Shape Menu
         self.detector_shape_menu = QMenu("Detector &Shape", self)
@@ -219,6 +293,10 @@ def setup_menus(self):
         detector_shape_group.setExclusive(True)
         self.detector_shape_group = detector_shape_group
 
+        diffraction_detector_separator = QAction("Diffraction", self)
+        diffraction_detector_separator.setDisabled(True)
+        self.detector_shape_menu.addAction(diffraction_detector_separator)
+
         detector_point_action = QAction("&Point", self)
         detector_point_action.setCheckable(True)
         detector_point_action.setChecked(True)  # Default
@@ -245,6 +323,29 @@ def setup_menus(self):
         detector_shape_group.addAction(detector_annulus_action)
         self.detector_shape_menu.addAction(detector_annulus_action)
 
+        self.detector_shape_menu.addSeparator()
+
+        diffraction_detector_separator = QAction("Real Space", self)
+        diffraction_detector_separator.setDisabled(True)
+        self.detector_shape_menu.addAction(diffraction_detector_separator)
+
+        rs_detector_shape_group = QActionGroup(self)
+        rs_detector_shape_group.setExclusive(True)
+        self.rs_detector_shape_group = rs_detector_shape_group
+
+        rs_detector_point_action = QAction("Poin&t", self)
+        rs_detector_point_action.setCheckable(True)
+        rs_detector_point_action.setChecked(True)  # Default
+        rs_detector_point_action.triggered.connect(self.update_realspace_detector)
+        rs_detector_shape_group.addAction(rs_detector_point_action)
+        self.detector_shape_menu.addAction(rs_detector_point_action)
+
+        detector_rectangle_action = QAction("Rectan&gular", self)
+        detector_rectangle_action.setCheckable(True)
+        detector_rectangle_action.triggered.connect(self.update_realspace_detector)
+        rs_detector_shape_group.addAction(detector_rectangle_action)
+        self.detector_shape_menu.addAction(detector_rectangle_action)
+
     def setup_views(self):
         # Set up the diffraction space window.
         self.diffraction_space_widget = pg.ImageView()
@@ -261,11 +362,6 @@ def setup_views(self):
         self.virtual_detector_point.sigRegionChanged.connect(
             self.update_real_space_view
         )
-        # self.virtual_detector_roi = pg.RectROI([5, 5], [20, 20], pen=(3, 9))
-        # self.diffraction_space_widget.getView().addItem(self.virtual_detector_roi)
-        # self.virtual_detector_roi.sigRegionChangeFinished.connect(
-        #     partial(self.update_real_space_view, False)
-        # )
 
         # Name and return
         self.diffraction_space_widget.setWindowTitle("Diffraction Space")