Add ability to create waterfall plots

hexrdgui/calibration/polarview.py

@@ -55,9 +55,6 @@ def __init__(self, instrument, distortion_instrument=None):
             # Use an image dict with the panel buffers applied.
             # This keeps invalid pixels from bleeding out in the polar view
             self.images_dict = HexrdConfig().images_dict
-            # 0 is a better fill value because it results in fewer nans in
-            # the final image.
-            HexrdConfig().apply_panel_buffer_to_images(self.images_dict, 0)
 
         self.warp_dict = {}
 
@@ -176,8 +173,16 @@ def images_dict(self):
 
     @images_dict.setter
     def images_dict(self, v):
+        # This images_dict sometimes gets modified by external callers,
+        # such as when a waterfall plot is created. So we need to make
+        # sure that everything that needs to be updated gets updated
+        # here.
         self._images_dict = v
 
+        # 0 is a better fill value because it results in fewer nans in
+        # the final image.
+        HexrdConfig().apply_panel_buffer_to_images(self._images_dict, 0)
+
         # Cache the image min and max for later use
         self.min = min(x.min() for x in v.values())
         self.max = max(x.max() for x in v.values())
@@ -240,13 +245,12 @@ def detector_borders(self, det):
     @property
     def all_detector_borders(self):
         borders = {}
-        for key in self.images_dict.keys():
+        for key in self.detectors:
             borders[key] = self.detector_borders(key)
 
         return borders
 
     def create_warp_image(self, det):
-        # lcount = 0
         img = self.images_dict[det]
         panel = self.detectors[det]
 
@@ -508,7 +512,7 @@ def warp_all_images(self):
         self.reset_cached_distortion_fields()
 
         # Create the warped image for each detector
-        for det in self.images_dict.keys():
+        for det in self.detectors:
             self.create_warp_image(det)
 
         # Generate the final image
@@ -540,6 +544,9 @@ def update_detectors(self, detectors):
         self.generate_image()
 
     def reset_cached_distortion_fields(self):
+        # These are only reset so that other parts of the code
+        # will not use them while we are generating new ones.
+        # They are actually still cached elsewhere.
         HexrdConfig().polar_corr_field_polar = None
         HexrdConfig().polar_angular_grid = None
 

hexrdgui/image_canvas.py

@@ -4,8 +4,9 @@
 import sys
 
 from PySide6.QtCore import QThreadPool, QTimer, Signal, Qt
-from PySide6.QtWidgets import QFileDialog, QMessageBox
+from PySide6.QtWidgets import QFileDialog, QMessageBox, QProgressDialog
 
+from matplotlib.axes import Axes
 from matplotlib.backends.backend_qtagg import FigureCanvas
 from matplotlib.figure import Figure
 from matplotlib.lines import Line2D
@@ -20,6 +21,7 @@
 
 from hexrd import distortion as distortion_pkg
 
+from hexrdgui import utils
 from hexrdgui.async_worker import AsyncWorker
 from hexrdgui.blit_manager import BlitManager
 from hexrdgui.calibration.cartesian_plot import cartesian_viewer
@@ -33,12 +35,12 @@
 from hexrdgui.masking.create_polar_mask import create_polar_line_data_from_raw
 from hexrdgui.masking.mask_manager import MaskManager
 from hexrdgui.snip_viewer_dialog import SnipViewerDialog
-from hexrdgui import utils
 from hexrdgui.utils.array import split_array
 from hexrdgui.utils.conversions import (
     angles_to_stereo, cart_to_angles, cart_to_pixels, q_to_tth, tth_to_q,
 )
 from hexrdgui.utils.tth_distortion import apply_tth_distortion_if_needed
+from hexrdgui.waterfall_plot import WaterfallPlotDialog
 
 # Increase these font sizes (compared to the global font) by the specified
 # amounts.
@@ -52,6 +54,8 @@ class ImageCanvas(FigureCanvas):
     norm_modified = Signal()
     transform_modified = Signal()
 
+    _update_waterfall_plot_progress = Signal()
+
     def __init__(self, parent=None, image_names=None):
         self.figure = Figure(tight_layout=True)
         super().__init__(self.figure)
@@ -76,6 +80,7 @@ def __init__(self, parent=None, image_names=None):
         self.raw_view_images_dict = {}
         self._mask_boundary_artists = []
         self._latest_compute_view_worker = None
+        self._waterfall_plot_dialog = None
 
         # Track the current mode so that we can more lazily clear on change.
         self.mode = None
@@ -123,6 +128,14 @@ def setup_connections(self):
         HexrdConfig().panel_distortion_modified.connect(
             self.on_panel_distortion_changed)
 
+        # This *must* be a queued connection, because Mac requires the
+        # progress to be updated on the GUI thread. Otherwise, it will
+        # crash the application on Mac.
+        self._update_waterfall_plot_progress.connect(
+            self._update_waterfall_plot_progress_slot,
+            Qt.QueuedConnection,
+        )
+
     @property
     def thread_pool(self):
         return QThreadPool.globalInstance()
@@ -1157,55 +1170,10 @@ def finish_show_polar(self, iviewer):
                 HexrdConfig().last_unscaled_azimuthal_integral_data = unscaled
 
                 self.azimuthal_integral_axis = axis
-                axis.set_ylabel(r'Azimuthal Average', **self.label_kwargs)
                 self.update_azimuthal_plot_overlays()
                 self.update_wppf_plot()
 
-                # Set up formatting for the x-axis
-                default_formatter = axis.xaxis.get_major_formatter()
-                f = self.format_polar_x_major_ticks
-                formatter = PolarXAxisFormatter(default_formatter, f)
-                axis.xaxis.set_major_formatter(formatter)
-
-                axis.yaxis.set_major_locator(AutoLocator())
-                axis.yaxis.set_minor_locator(AutoMinorLocator())
-
-                axis.xaxis.set_major_locator(PolarXAxisTickLocator(self))
-                self.axis.xaxis.set_minor_locator(
-                    PolarXAxisMinorTickLocator(self)
-                )
-
-                # change property of ticks
-                axis.tick_params(**self.major_tick_kwargs)
-                axis.tick_params(**self.minor_tick_kwargs)
-
-                # add grid lines parallel to x-axis in azimuthal average
-                kwargs = {
-                    'visible': True,
-                    'which': 'major',
-                    'axis': 'y',
-                    'linewidth': 0.25,
-                    'linestyle': '-',
-                    'color': 'k',
-                    'alpha': 0.75,
-                }
-                axis.grid(**kwargs)
-
-                kwargs = {
-                    'visible': True,
-                    'which': 'minor',
-                    'axis': 'y',
-                    'linewidth': 0.075,
-                    'linestyle': '--',
-                    'color': 'k',
-                    'alpha': 0.9,
-                }
-                axis.grid(**kwargs)
-
-                # add grid lines parallel to y-axis
-                kwargs['which'] = 'both'
-                kwargs['axis'] = 'x'
-                axis.grid(**kwargs)
+                self._setup_azimuthal_axis(axis)
             else:
                 self.update_azimuthal_integral_plot()
                 axis = self.azimuthal_integral_axis
@@ -1331,6 +1299,65 @@ def on_beam_energy_modified(self):
         # Update the beam energy on the instrument
         self.iviewer.instr.beam_energy = HexrdConfig().beam_energy
 
+    def _setup_azimuthal_axis(self, axis: Axes):
+        # Set the labels
+        axis.set_xlabel(self.polar_xlabel, **self.label_kwargs)
+        axis.set_ylabel(r'Azimuthal Average', **self.label_kwargs)
+
+        # Set up formatting for the x-axis
+        # This is important in case "Q" is on the x axis instead
+        # of two theta.
+        default_formatter = axis.xaxis.get_major_formatter()
+        f = self.format_polar_x_major_ticks
+        formatter = PolarXAxisFormatter(default_formatter, f)
+        axis.xaxis.set_major_formatter(formatter)
+
+        axis.yaxis.set_major_locator(AutoLocator())
+        axis.yaxis.set_minor_locator(AutoMinorLocator())
+
+        axis.xaxis.set_major_locator(PolarXAxisTickLocator(self))
+        self.axis.xaxis.set_minor_locator(
+            PolarXAxisMinorTickLocator(self)
+        )
+
+        # change property of ticks
+        axis.tick_params(**self.major_tick_kwargs)
+        axis.tick_params(**self.minor_tick_kwargs)
+
+        # Set up the grids
+        # These are default kwargs for the grids.
+        default_kwargs = {
+            'visible': True,
+            'linewidth': 0.075,
+            'linestyle': '--',
+            'color': 'k',
+            'alpha': 0.9,
+        }
+
+        # Grid for minor y tickers
+        axis.grid(**{
+            **default_kwargs,
+            'which': 'minor',
+            'axis': 'y',
+            'linewidth': 0.25,
+            'linestyle': '-',
+            'alpha': 0.75,
+        })
+
+        # Grid for major y tickers
+        axis.grid(**{
+            **default_kwargs,
+            'which': 'major',
+            'axis': 'y',
+        })
+
+        # Grid for all x tickers
+        axis.grid(**{
+            **default_kwargs,
+            'which': 'both',
+            'axis': 'x',
+        })
+
     @property
     def polar_x_axis_type(self):
         return HexrdConfig().polar_x_axis_type
@@ -1510,10 +1537,14 @@ def compute_azimuthal_integral_sum(self, scaled=True):
             pimg = self.scaled_images[0]
         else:
             pimg = self.unscaled_images[0]
+
+        return self._compute_azimuthal_integral_sum(pimg)
+
+    def _compute_azimuthal_integral_sum(self, pimg: np.ndarray) -> np.ndarray:
         # !!! NOTE: visible polar masks have already been applied
         #           in polarview.py
-        masked = np.ma.masked_array(pimg, mask=np.isnan(pimg))
         offset = HexrdConfig().azimuthal_offset
+        masked = np.ma.masked_array(pimg, mask=np.isnan(pimg))
         return masked.sum(axis=0) / np.sum(~masked.mask, axis=0) + offset
 
     def clear_azimuthal_overlay_artists(self):
@@ -1767,6 +1798,133 @@ def export_current_plot(self, filename):
 
         self.iviewer.write_image(filename)
 
+    def create_waterfall_plot(self):
+        if self.mode != ViewType.polar:
+            msg = 'Cannot create waterfall plot if we are not in polar mode'
+            raise Exception(msg)
+
+        if not self.iviewer:
+            msg = 'Cannot create waterfall plot without an iviewer'
+            raise Exception(msg)
+
+        if self._waterfall_plot_dialog is not None:
+            self._waterfall_plot_dialog.hide()
+            self._waterfall_plot_dialog = None
+
+        # Determine the number of lineouts
+        num_lineouts = HexrdConfig().imageseries_length
+
+        # Display a progress dialog indicating that we are
+        # generating intensities...
+        progress = QProgressDialog(
+            'Generating azimuthal lineouts...',
+            None,
+            0,
+            num_lineouts,
+            self,
+        )
+        progress.setWindowTitle('HEXRD')
+        progress.setValue(1)
+
+        # No close button in the corner
+        flags = progress.windowFlags()
+        progress.setWindowFlags(
+            (flags | Qt.CustomizeWindowHint) &
+            ~Qt.WindowCloseButtonHint
+        )
+
+        self._create_waterfall_progress = progress
+
+        # Compute azimuthal lineouts in a background thread
+        worker = AsyncWorker(self._create_waterfall_lineouts)
+        self.thread_pool.start(worker)
+        self._latest_compute_view_worker = worker
+
+        def on_finished():
+            progress.reject()
+
+        # Get the results and close the progress dialog when finished
+        worker.signals.result.connect(self._finish_create_waterfall)
+        worker.signals.finished.connect(on_finished)
+
+        progress.exec()
+
+    def _update_waterfall_plot_progress_slot(self):
+        progress = self._create_waterfall_progress
+        if progress is None:
+            return
+
+        progress.setValue(progress.value() + 1)
+
+    def _create_waterfall_lineouts(self) -> list[np.ndarray]:
+        # Determine the number of lineouts
+        num_lineouts = HexrdConfig().imageseries_length
+        lineouts = [None] * num_lineouts
+
+        # We can already compute the lineout for the current frame
+        current_idx = HexrdConfig().current_imageseries_idx
+        lineouts[current_idx] = self.compute_azimuthal_integral_sum()
+
+        # Make a deep copy of the iviewer, since we will modify it
+        iviewer = copy.deepcopy(self.iviewer)
+
+        # Now generate the lineouts for the other frames
+        for i in range(num_lineouts):
+            if i == current_idx:
+                # We already generated this one
+                continue
+
+            # Create the new imageseries dict
+            HexrdConfig().current_imageseries_idx = i
+            try:
+                new_images_dict = HexrdConfig().images_dict
+            finally:
+                # Always restore the previous index
+                HexrdConfig().current_imageseries_idx = current_idx
+
+            # Now force the image dict to change
+            iviewer.pv.images_dict = new_images_dict
+
+            # Generate the new image
+            iviewer.pv.warp_all_images()
+
+            # Grab the new image
+            polar_img = iviewer.img
+            if HexrdConfig().polar_apply_scaling_to_lineout:
+                # Apply the transform
+                polar_img = self.transform(polar_img)
+
+            # Compute the integration
+            lineouts[i] = self._compute_azimuthal_integral_sum(polar_img)
+
+            # The progress must be updated in the GUI thread. Otherwise,
+            # it will crash on Mac.
+            self._update_waterfall_plot_progress.emit()
+
+        return lineouts
+
+    def _finish_create_waterfall(self, lineouts: list[np.ndarray]):
+        # Now create the waterfall plot dialog with the lineouts
+        # Create a matplotlib figure and set up everything
+        figure = plt.figure()
+        ax = figure.add_subplot()
+
+        # Grab tth
+        angular_grid = self.iviewer.angular_grid
+        tth = np.degrees(angular_grid[1][0])
+        line_data = [(tth, lineout.filled(np.nan)) for lineout in lineouts]
+
+        # Set up the same azimuthal axes parameters as the polar view
+        self._setup_azimuthal_axis(ax)
+
+        # Disable the tick labels
+        ax.set_yticklabels([])
+
+        # Now create and show the waterfall plot
+        dialog = WaterfallPlotDialog(ax, line_data)
+        dialog.show()
+        self._waterfall_plot_dialog = dialog
+
     def export_to_maud(self, filename):
         if self.mode != ViewType.polar:
             msg = 'Must be in polar mode. Cannot export.'