Merge branch 'feat/detector-layouts' into 'main'

Add detector layout in sfx geometry refinement

See merge request esobolev/geomtools!2

src/geomtools/sfx/optimiser.py

@@ -9,7 +9,7 @@ from ..detector import (
     write_crystfel_geom)
 from .crystfelio import extract_geometry, read_crystfel_streamfile
 from .misc import badpixels_mask
-from .refine import refine_geometry
+from .refine import LAYOUTS, refine_geometry
 
 
 def refine():
@@ -24,6 +24,9 @@ def refine():
     parser.add_argument('-o', '--output', type=pathlib.Path,
                         default=pathlib.Path("refined.geom"),
                         help="File to store refined geometry")
+    parser.add_argument('-l', '--layout', type=str,
+                        choices=list(LAYOUTS.keys()),
+                        help="Detector layout")
 
     args = parser.parse_args()
 
@@ -39,14 +42,34 @@ def refine():
 
     temp_geom_file.close()
 
+    shape = get_detector_shape(panels)
+
+    if args.layout is None:
+        possible_layouts = [name for name, p in LAYOUTS.items()
+                            if p.shape == shape[-2:]]
+        if len(possible_layouts) != 1:
+            print(f"{parser.prog}: error: many layouts fit detector shape, "
+                  "please specify layout option")
+            exit(1)
+
+        layout = possible_layouts[0]
+    else:
+        layout = args.layout
+        if LAYOUTS[layout].shape != shape[-2:]:
+            print(
+                f"{parser.prog}: error: the detector shape {shape} in "
+                f"geometry does not corresponded to the {args.layout} layout")
+            exit(1)
+
+    print("Detector layout:", layout)
+
     fr, pe, la, re, ma = read_crystfel_streamfile(
         stream_filename, panels, args.connected, disp=True)
 
     panels_new, transform, clen = refine_geometry(
-        ma, panels, 1.0, args.min_counts, args.connected)
+        ma, panels, layout, 1.0, args.min_counts, args.connected)
 
     px = badpixels_mask(pe, panels)
-    shape = get_detector_shape(panels)
     badpx = pixels_to_image(shape, px)
 
     geom_fn = args.output

src/geomtools/sfx/refine.py

 import time
+from dataclasses import dataclass
 
 import numpy as np
 import pandas as pd
 from natsort import natsorted
 
 
-def refine_geometry(matches, panels, clen_scale=1., min_counts=50,
-                    rigid_group="modules", disp=True):
+@dataclass
+class Layout:
+    shape: tuple
+    xy: np.ndarray
+    ss_size: int
+    fs_size: int
+    ss_gap: float
+    fs_gap: float
+
+    def __post_init__(self):
+        self.ss_stride = self.ss_size + self.ss_gap
+        self.fs_stride = self.fs_size + self.fs_gap
+
+
+LAYOUTS = {
+    "AGIPD": Layout(
+        (512, 128),
+        np.array([[1, 0], [0, -1]]),
+        64, 128, 2, 0,
+    ),
+    "Jungfrau": Layout(
+        (512, 1024),
+        np.array([[0, 1], [1, 0]]),
+        256, 256, 2, 2,
+    ),
+}
+
+
+def refine_geometry(matches, panels, layout, clen_scale=1.,
+                    min_counts=50, rigid_group="modules", disp=True):
     transform = {}
     clen = panels.clen[0]
     scale_mean = 0
     nmod = 0
 
+    if isinstance(layout, str):
+        p = LAYOUTS[layout]
+    elif isinstance(layout, Layout):
+        p = layout
+    else:
+        raise TypeError("The value of `layout` argument may be `str` or "
+                        "`Layout`")
+
     if disp:
         tm0 = time.monotonic()
         print("Group   Roll [deg]      Scale           Corner offset "
@@ -24,15 +61,18 @@ def refine_geometry(matches, panels, clen_scale=1., min_counts=50,
         if count < min_counts:
             continue
 
-        # AGIPD1M layout
-        ss_p = mp.ss_p + (np.round(mp.ss_p) // 64) * 2
-        fs_p = -mp.fs_p
+        # layout
+        ss_p = mp.ss_p + (mp.ss_p.astype(int) // p.ss_size) * p.ss_gap
+        fs_p = mp.fs_p + (mp.fs_p.astype(int) // p.fs_size) * p.fs_gap
+
+        # probably p.xy should be transposed here
+        xp, yp = p.xy @ np.array([ss_p, fs_p])
 
-        xc_p = ss_p.mean()
-        yc_p = fs_p.mean()
+        xc_p = xp.mean()
+        yc_p = yp.mean()
 
-        xp = ss_p - xc_p
-        yp = fs_p - yc_p
+        xp -= xc_p
+        yp -= yc_p
 
         # reflexes
         xa = mp.xa / clen_scale
@@ -77,10 +117,9 @@ def refine_geometry(matches, panels, clen_scale=1., min_counts=50,
         print(f"Camera length: {clen_new:.8} m")
 
     params = {}
-    xy = np.array([[1, 0], [0, -1]])
     for name, (R, corner, scale) in transform.items():
         # matrix to transform from pixel indices to lab coordinates
-        K = R.T @ xy
+        K = R.T @ p.xy
         # module offsets
         coffset = clen * (scale - scale_mean)
 
@@ -106,9 +145,9 @@ def refine_geometry(matches, panels, clen_scale=1., min_counts=50,
         else:
             K, corner, coffset = position
 
-            # AGIPD layout
-            pane = attr.orig_min_ss // 64
-            cnx, cny = (K.T @ np.array([66 * pane, 0]) + corner)
+            ss = p.ss_stride * (attr.orig_min_ss // p.ss_size)
+            fs = p.fs_stride * (attr.orig_min_fs // p.fs_size)
+            cnx, cny = (K.T @ np.array([ss, fs]) + corner)
 
         panels_new.append({
             'cnx': cnx, 'cny': cny, 'coffset': coffset,
@@ -119,6 +158,7 @@ def refine_geometry(matches, panels, clen_scale=1., min_counts=50,
             'panel': panel, 'modules': module_name,
             'quads': quad_name, 'modno': modno,
             'clen': clen_new, 'data': attr.data, 'mask': attr['mask'],
+            'res': attr['res'],
         })
 
     if disp: