Refactor exporting the polar view

hexrdgui/calibration/polar_plot.py

@@ -40,8 +40,8 @@ def raw_img(self):
         return self.pv.raw_img
 
     @property
-    def raw_mask(self):
-        return self.pv.raw_mask
+    def warp_mask(self):
+        return self.pv.warp_mask
 
     @property
     def snipped_img(self):
@@ -59,6 +59,10 @@ def display_img(self):
     def snip_background(self):
         return self.pv.snip_background
 
+    @property
+    def erosion_mask(self):
+        return self.pv.erosion_mask
+
     def update_angular_grid(self):
         self.pv.update_angular_grid()
 
@@ -126,31 +130,48 @@ def write_image(self, filename='polar_image.npz'):
             np.savetxt(filename, azimuthal_integration, delimiter=delimiter)
             return
 
-        intensities = self.raw_img.data
-        intensities[self.raw_mask] = np.nan
+        intensities = self.img.filled(np.nan)
+        raw_intensities = self.raw_img.filled(np.nan)
 
         eta, tth = np.degrees(self.angular_grid)
 
         # Prepare the data to write out
         data = {
-            'tth_coordinates': tth,
-            'eta_coordinates': eta,
-            'q_coordinates': tth_to_q(tth, self.instr.beam_energy),
             'intensities': intensities,
-            'extent': self._extent,
             'azimuthal_integration': azimuthal_integration,
+            'extent': self._extent,
+            'eta_coordinates': eta,
+            'tth_coordinates': tth,
+            'q_coordinates': tth_to_q(tth, self.instr.beam_energy),
+            'raw_intensities': raw_intensities,
+            'warp_mask': self.warp_mask,
+            # FIXME: add intensity corrections too later
         }
 
         if self.snip_background is not None:
             # Add the snip background if it was used
             data['snip_background'] = self.snip_background
+            if self.erosion_mask is not None:
+                # Also add the erosion mask if it was used
+                data['erosion_mask'] = self.erosion_mask
+
+        if HexrdConfig().polar_tth_distortion:
+            # Add the tth distortion correction field
+            data['tth_corr_field'] = self.pv.create_corr_field_polar()
 
-        # Add visible polar mask data if we have any
+        # Add polar mask data if we have any
         for name, mask in MaskManager().masks.items():
-            if mask.type == MaskType.threshold or not mask.visible:
+            if mask.type == MaskType.threshold:
                 continue
 
-            data[f'mask_{name}'] = mask.get_masked_arrays(self.type)
+            if mask.visible:
+                data[f'visible_mask_{name}'] = mask.get_masked_arrays(
+                    self.type,
+                )
+            elif mask.show_border:
+                data[f'border_mask_{name}'] = mask.get_masked_arrays(
+                    self.type,
+                )
 
         # Delete the file if it already exists
         if filename.exists():

hexrdgui/calibration/polarview.py

@@ -70,6 +70,7 @@ def __init__(self, instrument, distortion_instrument=None):
         self._corr_field_polar_cached = None
 
         self.snip_background = None
+        self.erosion_mask = None
 
         self.update_angular_grid()
 
@@ -388,7 +389,7 @@ def generate_image(self):
         self.apply_image_processing()
 
     @property
-    def raw_mask(self):
+    def warp_mask(self):
         return self.raw_img.mask
 
     def apply_image_processing(self):
@@ -418,7 +419,7 @@ def apply_snip(self, img):
             no_nan_methods = [SnipAlgorithmType.Fast_SNIP_1D]
 
             if HexrdConfig().polar_snip1d_algorithm not in no_nan_methods:
-                img[self.raw_mask] = np.nan
+                img[self.warp_mask] = np.nan
 
             # Perform the background subtraction
             self.snip_background = run_snip1d(img)
@@ -434,17 +435,20 @@ def apply_snip(self, img):
                 ))
                 mask = binary_erosion(~self.raw_img.mask, structure)
                 img[~mask] = np.nan
-
+                self.erosion_mask = mask
+            else:
+                self.erosion_mask = None
         else:
             self.snip_background = None
+            self.erosion_mask = None
 
         return img
 
     def apply_visible_masks(self, img):
         # Apply user-specified masks if they are present
         from hexrdgui.masking.mask_manager import MaskManager
         img = img.copy()
-        total_mask = self.raw_mask
+        total_mask = self.warp_mask
         for mask in MaskManager().masks.values():
             if mask.type == MaskType.threshold or not mask.visible:
                 continue
@@ -464,7 +468,7 @@ def apply_boundary_masks(self, img):
         # Apply user-specified masks if they are present
         from hexrdgui.masking.mask_manager import MaskManager
         img = img.copy()
-        total_mask = self.raw_mask
+        total_mask = self.warp_mask
         for mask in MaskManager().masks.values():
             if mask.type == MaskType.threshold or not mask.show_border:
                 continue

hexrdgui/image_canvas.py

@@ -1773,7 +1773,7 @@ def polar_show_snip1d(self):
 
             no_nan_methods = [utils.SnipAlgorithmType.Fast_SNIP_1D]
             if HexrdConfig().polar_snip1d_algorithm not in no_nan_methods:
-                img[self.iviewer.raw_mask] = np.nan
+                img[self.iviewer.warp_mask] = np.nan
 
             background = utils.run_snip1d(img)