Consider tth distortion when saving picks

We previously would not consider the tth distortion when saving
picks, so the cartesian coordinates would be incorrect.

Now we reverse the tth distortion if needed. Also, the two calibration
workflows (Composite and Fast Powder) write the coordinates directly
from the overlays, which is more accurate (some precision is lost when
converting back and forth between cartesian and polar when the tth
distortion is applied, because it is a field correction).

Signed-off-by: Patrick Avery <[email protected]>

Author: psavery

hexrdgui/calibration/auto/powder_runner.py

@@ -242,7 +242,7 @@ def on_finished():
     def save_picks_to_file(self, selected_file):
         # Reuse the same logic from the HKLPicksTreeViewDialog
         dialog = self.create_hkl_picks_tree_view_dialog()
-        dialog.export_picks(selected_file)
+        dialog.export_picks_from_overlays(selected_file, self.overlays)
 
     def load_picks_from_file(self, selected_file):
         # Reuse the same logic from the HKLPicksTreeViewDialog

hexrdgui/calibration/calibration_runner.py

@@ -1045,7 +1045,8 @@ def on_finished():
 
     def save_picks_to_file(self, selected_file):
         # Reuse the same logic from the HKLPicksTreeViewDialog
-        self.edit_picks_dialog.export_picks(selected_file)
+        d = self.edit_picks_dialog
+        d.export_picks_from_overlays(selected_file, self.overlays)
 
     def load_picks_from_file(self, selected_file):
         # Reuse the same logic from the HKLPicksTreeViewDialog

hexrdgui/calibration/hkl_picks_tree_view_dialog.py

@@ -16,6 +16,7 @@
 from hexrdgui.ui_loader import UiLoader
 from hexrdgui.utils.conversions import angles_to_cart, cart_to_angles
 from hexrdgui.utils.dicts import ndarrays_to_lists
+from hexrdgui.utils.tth_distortion import apply_tth_distortion_if_needed
 
 
 class HKLPicksTreeViewDialog:
@@ -92,18 +93,28 @@ def export_picks_clicked(self):
             return self.export_picks(selected_file)
 
     def export_picks(self, filename):
+        return self._export_dict_to_file(filename, {
+            'angular': self.dictionary,
+            'cartesian': self.dict_with_cart_coords,
+        })
+
+    def export_picks_from_overlays(self, filename, overlays):
+        # Export picks from overlays using the same export logic as
+        # the regular dictionary.
+        return self._export_dict_to_file(filename, {
+            'angular': overlays_to_tree_format(overlays, polar=True),
+            'cartesian': overlays_to_tree_format(overlays, polar=False),
+        })
+
+    def _export_dict_to_file(self, filename: str, export_data: dict):
         filename = Path(filename)
 
         if filename.exists():
             filename.unlink()
 
-        # unwrap_dict_to_h5 unfortunately modifies the data
-        # make a deep copy to avoid the modification.
-        export_data = {
-            'angular': copy.deepcopy(self.dictionary),
-            'cartesian': self.dict_with_cart_coords,
-        }
-
+        # unwrap_dict_to_h5 unfortunately modifies the data.
+        # Make a deep copy to avoid the modification.
+        export_data = copy.deepcopy(export_data)
         with h5py.File(filename, 'w') as wf:
             unwrap_dict_to_h5(wf, export_data)
 
@@ -178,7 +189,7 @@ def button_box_visible(self, b):
         self.ui.button_box.setVisible(b)
 
 
-def convert_picks(picks, conversion_function, **kwargs):
+def convert_picks(picks, conversion_function):
     instr = create_hedm_instrument()
     ret = copy.deepcopy(picks)
     for name, detectors in ret.items():
@@ -191,25 +202,43 @@ def convert_picks(picks, conversion_function, **kwargs):
                         # Avoid the runtime warning
                         hkls[hkl] = [np.nan, np.nan]
                     else:
-                        hkls[hkl] = conversion_function([spot], panel,
-                                                        **kwargs)[0]
+                        hkls[hkl] = conversion_function([spot], panel)[0]
                 continue
 
             # Must be powder
             for hkl, line in hkls.items():
                 if len(line) != 0:
-                    hkls[hkl] = conversion_function(line, panel, **kwargs)
+                    hkls[hkl] = conversion_function(line, panel)
 
     return ret
 
 
 def picks_angles_to_cartesian(picks):
-    return convert_picks(picks, angles_to_cart)
+    # Create the conversion function
+    def func(angs, panel):
+        # Reverse the tth distortion first
+        angs = apply_tth_distortion_if_needed(
+            angs,
+            in_degrees=True,
+            reverse=True,
+        )
+        # Now convert to cart
+        return angles_to_cart(angs, panel)
+
+    return convert_picks(picks, func)
 
 
 def picks_cartesian_to_angles(picks):
-    kwargs = {'eta_period': HexrdConfig().polar_res_eta_period}
-    return convert_picks(picks, cart_to_angles, **kwargs)
+    # Create the conversion function
+    eta_period = HexrdConfig().polar_res_eta_period
+
+    def func(xys, panel):
+        angs = cart_to_angles(xys, panel, eta_period=eta_period)
+
+        # Apply tth distortion now as well
+        return apply_tth_distortion_if_needed(angs, in_degrees=True)
+
+    return convert_picks(picks, func)
 
 
 def generate_picks_results(overlays, polar=True):
@@ -249,8 +278,8 @@ def generate_picks_results(overlays, polar=True):
     return pick_results
 
 
-def overlays_to_tree_format(overlays):
-    picks = generate_picks_results(overlays)
+def overlays_to_tree_format(overlays, polar=True):
+    picks = generate_picks_results(overlays, polar=polar)
     return picks_to_tree_format(picks)