From dd262ab6cd386d79e6f65daec2279e870673f957 Mon Sep 17 00:00:00 2001
From: Ivan Usov <ivan.usov@psi.ch>
Date: Wed, 25 Jan 2023 20:40:46 +0100
Subject: [PATCH] Optimize ang2hkl functions to accept inverted UB directly

---
 pyzebra/app/panel_hdf_viewer.py |  4 ++--
 pyzebra/app/plot_hkl.py         | 12 ++++++------
 pyzebra/xtal.py                 | 10 ++++------
 3 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/pyzebra/app/panel_hdf_viewer.py b/pyzebra/app/panel_hdf_viewer.py
index 3b08d50..835c2da 100644
--- a/pyzebra/app/panel_hdf_viewer.py
+++ b/pyzebra/app/panel_hdf_viewer.py
@@ -975,7 +975,7 @@ def calculate_hkl(scan, index):
     gammad = scan["gamma"][index]
     om = scan["omega"][index]
     nud = scan["nu"]
-    ub = scan["ub"]
+    ub_inv = np.linalg.inv(scan["ub"])
     geometry = scan["zebra_mode"]
 
     if geometry == "bi":
@@ -990,7 +990,7 @@ def calculate_hkl(scan, index):
     for xi in np.arange(IMAGE_W):
         for yi in np.arange(IMAGE_H):
             h[yi, xi], k[yi, xi], l[yi, xi] = pyzebra.ang2hkl(
-                wave, ddist, gammad, om, chi, phi, nud, ub, xi, yi
+                wave, ddist, gammad, om, chi, phi, nud, ub_inv, xi, yi
             )
 
     return h, k, l
diff --git a/pyzebra/app/plot_hkl.py b/pyzebra/app/plot_hkl.py
index a83fcec..ed4ad9c 100644
--- a/pyzebra/app/plot_hkl.py
+++ b/pyzebra/app/plot_hkl.py
@@ -138,7 +138,7 @@ class PlotHKL:
                     chi = scan["chi"]
                     phi = scan["phi"]
                     nud = 0  # 1d detector
-                    ub = scan["ub"]
+                    ub_inv = np.linalg.inv(scan["ub"])
                     counts = scan["counts"]
                     wave = scan["wavelength"]
 
@@ -153,18 +153,18 @@ class PlotHKL:
                     res_x = []
                     res_y = []
                     for _om in np.linspace(om[0], om[-1], num=res_N):
-                        expr1 = ang2hkl_1d(wave, gammad, _om + res / 2, chi, phi, nud, ub)
-                        expr2 = ang2hkl_1d(wave, gammad, _om - res / 2, chi, phi, nud, ub)
+                        expr1 = ang2hkl_1d(wave, gammad, _om + res / 2, chi, phi, nud, ub_inv)
+                        expr2 = ang2hkl_1d(wave, gammad, _om - res / 2, chi, phi, nud, ub_inv)
                         hkl_temp = M @ (np.abs(expr1 - expr2) / 2)
                         res_x.append(hkl_temp[0])
                         res_y.append(hkl_temp[1])
 
                     # Get first and final hkl
-                    hkl1 = ang2hkl_1d(wave, gammad, om[0], chi, phi, nud, ub)
-                    hkl2 = ang2hkl_1d(wave, gammad, om[-1], chi, phi, nud, ub)
+                    hkl1 = ang2hkl_1d(wave, gammad, om[0], chi, phi, nud, ub_inv)
+                    hkl2 = ang2hkl_1d(wave, gammad, om[-1], chi, phi, nud, ub_inv)
 
                     # Get hkl at best intensity
-                    hkl_m = ang2hkl_1d(wave, gammad, om[np.argmax(counts)], chi, phi, nud, ub)
+                    hkl_m = ang2hkl_1d(wave, gammad, om[np.argmax(counts)], chi, phi, nud, ub_inv)
 
                     # Estimate intensity for marker size scaling
                     y_bkg = [counts[0], counts[-1]]
diff --git a/pyzebra/xtal.py b/pyzebra/xtal.py
index e4baab9..3d21c51 100644
--- a/pyzebra/xtal.py
+++ b/pyzebra/xtal.py
@@ -361,21 +361,19 @@ def angtohkl(wave, ddist, gammad, om, ch, ph, nud, x, y):
     )
 
 
-def ang2hkl(wave, ddist, gammad, om, ch, ph, nud, ub, x, y):
+def ang2hkl(wave, ddist, gammad, om, ch, ph, nud, ub_inv, x, y):
     """Calculate hkl-indices of a reflection from its position (x,y,angles) at the 2d-detector"""
     ga, nu = det2pol(ddist, gammad, nud, x, y)
     z1 = z1frmd(wave, ga, om, ch, ph, nu)
-    ubinv = np.linalg.inv(ub)
-    hkl = ubinv @ z1
+    hkl = ub_inv @ z1
 
     return hkl
 
 
-def ang2hkl_1d(wave, ga, om, ch, ph, nu, ub):
+def ang2hkl_1d(wave, ga, om, ch, ph, nu, ub_inv):
     """Calculate hkl-indices of a reflection from its position (angles) at the 1d-detector"""
     z1 = z1frmd(wave, ga, om, ch, ph, nu)
-    ubinv = np.linalg.inv(ub)
-    hkl = ubinv @ z1
+    hkl = ub_inv @ z1
 
     return hkl