Add a possibility to override UB matrix in hdf viewer

Fix #70

pyzebra/app/panel_hdf_viewer.py

@@ -30,6 +30,7 @@ from bokeh.models import (
     Slider,
     Spacer,
     Spinner,
+    TextInput,
     TableColumn,
     TabPanel,
     Tabs,
@@ -49,9 +50,36 @@ def create():
     doc = curdoc()
     log = doc.logger
     dataset = []
+    ub_matrix_global = None
 
     num_formatter = NumberFormatter(format="0.00", nan_format="")
 
+    def ub_matrix_callback(_attr, _old, new):
+        nonlocal ub_matrix_global
+        if new == "":
+            ub_matrix_global = None
+            return
+
+        try:
+            ub_values = list(map(float, new.strip().split()))
+            if len(ub_values) != 9:
+                raise ValueError("UB matrix must have 9 values")
+            ub_matrix_global = np.array(ub_values).reshape(3, 3)
+
+        except Exception as e:
+            log.error(f"Invalid UB matrix: {e}")
+            return
+
+        if dataset:
+            scan = _get_selected_scan()
+            scan["ub"] = ub_matrix_global
+            _update_image()
+
+    ub_matrix = TextInput(
+        title="Apply UB Matrix:", placeholder="Enter 9 space-separated values", width=600
+    )
+    ub_matrix.on_change("value", ub_matrix_callback)
+
     def file_select_update():
         proposal_path = proposal_textinput.name
         if proposal_path:
@@ -74,7 +102,7 @@ def create():
     def upload_hdf_button_callback(_attr, _old, new):
         nonlocal dataset
         try:
-            scan = pyzebra.read_detector_data(io.BytesIO(base64.b64decode(new)))
+            scan = pyzebra.read_detector_data(io.BytesIO(base64.b64decode(new)), ub_matrix_global)
         except Exception as e:
             log.exception(e)
             return
@@ -107,7 +135,7 @@ def create():
         new_data = []
         for f_path in file_select.value:
             try:
-                file_data = [pyzebra.read_detector_data(f_path)]
+                file_data = [pyzebra.read_detector_data(f_path, ub_matrix_global)]
             except Exception as e:
                 log.exception(e)
                 continue
@@ -130,7 +158,7 @@ def create():
         file_data = []
         for f_path in file_select.value:
             try:
-                file_data = [pyzebra.read_detector_data(f_path)]
+                file_data = [pyzebra.read_detector_data(f_path, ub_matrix_global)]
             except Exception as e:
                 log.exception(e)
                 continue
@@ -295,7 +323,6 @@ def create():
             x=np.mean(current_image, axis=1), y=np.arange(0, IMAGE_H) + 0.5
         )
 
-        image_source.data.update(h=[np.zeros((1, 1))], k=[np.zeros((1, 1))], l=[np.zeros((1, 1))])
         image_source.data.update(image=[current_image])
 
         if main_auto_checkbox.active:
@@ -347,6 +374,9 @@ def create():
             gamma=[gamma_c], nu=[nu_c], omega=[omega_c], chi=[chi_c], phi=[phi_c]
         )
 
+        h, k, l = calculate_hkl(scan, index)
+        image_source.data.update(h=[h], k=[k], l=[l])
+
     def _update_proj_plots():
         scan = _get_selected_scan()
         counts = scan["counts"]
@@ -465,16 +495,6 @@ def create():
     plot.image(source=image_source, image="nu", global_alpha=0)
     plot.image(source=image_source, image="omega", global_alpha=0)
 
-    # calculate hkl-indices of first mouse entry
-    def mouse_enter_callback(_event):
-        if dataset and np.array_equal(image_source.data["h"][0], np.zeros((1, 1))):
-            scan = _get_selected_scan()
-            index = index_spinner.value
-            h, k, l = calculate_hkl(scan, index)
-            image_source.data.update(h=[h], k=[k], l=[l])
-
-    plot.on_event(MouseEnter, mouse_enter_callback)
-
     # Single frame projection plots
     proj_v = figure(
         x_range=plot.x_range,
@@ -892,7 +912,7 @@ def create():
 
     tab_layout = row(
         column(import_layout, colormap_layout),
-        column(row(scan_layout, layout_proj), layout_controls),
+        column(ub_matrix, row(scan_layout, layout_proj), layout_controls),
         column(roi_avg_plot, layout_image),
     )
 

pyzebra/h5.py

@@ -59,14 +59,15 @@ def parse_h5meta(file):
     return content
 
 
-def read_detector_data(filepath):
+def read_detector_data(filepath, ub_matrix=None):
     """Read detector data and angles from an h5 file.
 
     Args:
         filepath (str): File path of an h5 file.
+        ub_matrix (np.ndarray, optional): UB matrix to use. If None, will be read from the file.
 
     Returns:
-        ndarray: A 3D array of data, omega, gamma, nu.
+        dict: A dictionary containing the detector data and associated metadata.
     """
     with h5py.File(filepath, "r") as h5f:
         counts = h5f["/entry1/area_detector2/data"][:].astype(float)
@@ -123,6 +124,11 @@ def read_detector_data(filepath):
             scan["ub"] = np.eye(3) * 0.177
         else:
             scan["ub"] = h5f["/entry1/sample/UB"][:].reshape(3, 3)
+
+        # override UB matrix if provided as an argument
+        if ub_matrix is not None:
+            scan["ub"] = ub_matrix
+
         scan["name"] = h5f["/entry1/sample/name"][0].decode()
         scan["cell"] = h5f["/entry1/sample/cell"][:]