[Jungfrau][Correct] New A1256 JF Strixel

Description

This MR addresses the integration for a new JF strixel at SCS

The main purpose of this MR is to integrate the new detector. Additionally, to update the strixel_configuration to be able to update the jfstrixel functions to switch between the different jf strixels we have now at HED(A0123) and SCS(A1256)

As I did previously with epix100 and GH2, I saw this opportunity to collect the Jungfrau modules in a dedicated directory. I have also included the 2 new .npz files which @hammerd produced at calng!76 (merged). This enables us to avoid creating the LUT for every train. This led to removing some functions from jfstrixel.

I updated the dark notebook as well as the test_jungfraulib file to adapt the library path update.

How Has This Been Tested?

• I used the reference test code to test jfstrixel HED correction compared to with master.
• I ran as well the correction for jfstrixel with xfel-calibrate.

Relevant Documents (optional)

jfstrixel_A1256.pdf

jfstrixel_A0123.pdf

Types of changes

• New feature (non-breaking change which adds functionality)
• Refactor (refactoring code with no functionality changes)

Checklist:

Reviewers

@hammerd @kluyvert

src/cal_tools/jungfrau/jfstrixel.py

 
-import numpy as np
+from functools import lru_cache
+from pathlib import Path
 
+import numpy as np
 
 REGULAR_SHAPE = (512, 1024)
-STRIXEL_SHAPE = (86, 3090)
-
-
-def _normal_indices():
-    """Build normal size pixel indices."""
-    
-    # Normal pixels
-    yin = np.arange(256)
-    xin = np.arange(1024)
-
-    Yin, Xin = np.meshgrid(yin, xin)
-    Yout, Xout = np.meshgrid(yin // 3, (xin // 256 * 774) + (xin % 256) * 3)
-    Xout += (yin % 3).astype(int)[None, :]
-    
-    return Yout, Xout, Yin, Xin
-
-
-def _gap_indices(in_gap_offset=0, out_gap_offset=0,
-                 xout_factor=+1, yout_offset=0):
-    """Build one half of double size gap pixel indices."""
-    
-    igap = np.arange(3)
-    yin = np.arange(256)
-
-    Yin, Xin = np.meshgrid(yin, igap * 256 + 255 + in_gap_offset)
-    Yout, Xout = np.meshgrid(yin // 6 * 2, igap * 774 + 765 + out_gap_offset)
-    Xout += xout_factor * (yin % 6).astype(int)[None, :]
-    Yout += yout_offset
-    
-    return Yout, Xout, Yin, Xin
-
-
-def transformation_indices2d():
-    """Build 2D strixel transformation index arrays."""
-    
-    # Each of this index sets contains four 2D index arrays
-    # Yout, Xout, Yin, Xin from different parts constituting the full
-    # strixel frame. They are each concatenated across these parts into
-    # four final index arrays to be used for translating between the
-    # regular frame and the strixel frame.
-    index_sets = [
-        _normal_indices(),
-        
-        # Left gap
-        _gap_indices(0, 0, +1, 0), _gap_indices(0, 0, +1, 1),
-        
-        # Right gap
-        _gap_indices(1, 11, -1, 0), _gap_indices(1, 11, -1, 1)
-    ]
-    
-    # Yout, Xout, Yin, Xin
-    # Casting to int64 improves indexing performance by up to 30%.
-    return [np.concatenate(index_set).astype(np.int64)
-            for index_set in zip(*index_sets)]
-
-
-def transformation_indices1d():
-    """Build 1D strixel transformation index arrays.
-    
-    Internally this function reduces the 2D index arrays to a single
-    dimension to operate on raveled data arrays. This improves the
-    transformation performance substantially by up to 3x.
+DIR_PATH = package_directory = Path(__file__).resolve().parent
+
+
+@lru_cache
+def get_strixel_parameters(kind):
+    """Returns a dictionary of strixel parameters stored in .npz file
+    based on the given kind.
+
+    Args:
+        kind (str): Specifies the type of strixel parameters to retrieve.
+            There is two possible values: "A0123" or "A1256"
+    Returns:
+        (dict): Dictionary contating the strixel parameters.
     """
+    strx_parameters = {}
 
-    Yout, Xout, Yin, Xin = transformation_indices2d()
-     
-    regular_pixel_idx = np.arange(np.prod(REGULAR_SHAPE), dtype=np.uint32) \
-        .reshape(REGULAR_SHAPE)
-    strixel_pixel_idx = np.empty(STRIXEL_SHAPE, dtype=np.int64)
-    strixel_pixel_idx.fill(-1)
-    strixel_pixel_idx[Yout, Xout] = regular_pixel_idx[Yin, Xin]
+    if kind == "A0123":
+        file_path = DIR_PATH / "strixel_cols_A0123-lut_mask.npz"
+    elif kind == "A1256":
+        file_path = DIR_PATH / "strixel_rows_A1256-lut_mask.npz"
 
-    Iout = np.where(strixel_pixel_idx.ravel() != -1)[0].astype(np.int64)
-    Iin = strixel_pixel_idx.ravel()[Iout].astype(np.int64)
-    
-    return Iout, Iin
+    with np.load(file_path) as data:
+        for k in data.files:
+            strx_parameters[k] = data[k]
 
+    return strx_parameters
 
-def double_pixel_indices():
+
+def store_double_pixel_indices():
     """Build index arrays for double-size pixels.
 
-    In raw data, the entire columns 255, 256, 511, 512, 767 and 768
+    In raw data for A0123 strixel detector,
+    the entire columns 255, 256, 511, 512, 767 and 768
     are double-size pixels. After strixelation, these end up in columns
     765-776, 1539-1550 and 2313-2324 on rows 0-85 or 0-83, with a set
     of four columns with 86 rows followed by a set of 84 and 86 again.
 
     This function builds the index arrays for double pixels after
-    strixelation.
-
-    Returns: 
-        (ndarray, ndarray) 2D index arrays for double pixel Y and X.
+    strixelation and stores it in the available A0123 .npz file.
     """
 
-    Ydouble = []
-    Xdouble = []
+    ydouble = []
+    xdouble = []
+    file_path = DIR_PATH / "strixel_cols_A0123-lut_mask.npz"
 
-    for double_col in [765, 1539, 2313]:
-        for col in range(double_col, double_col+12):
-            for row in range(84 if ((col-double_col) // 4) == 1 else 86):
-                Ydouble.append(row)
-                Xdouble.append(col)
+    with np.load(file_path) as data:
+        for double_col in [765, 1539, 2313]:
+            for col in range(double_col, double_col+12):
+                for row in range(84 if ((col-double_col) // 4) == 1 else 86):
+                    ydouble.append(row)
+                    xdouble.append(col)
+        np.savez(file_path, **data, ydouble=ydouble, xdouble=xdouble)
 
-    return np.array(Ydouble), np.array(Xdouble)
 
-
-def to_strixel(data, out=None):
+def to_strixel(data, out=None, kind="A0123"):
     """Transform from regular to strixel geometry.
 
     Only the last two axes are considered for transformation, input data
     may have any number of additional axes in front.
-    
+
     Args:
         data (array_like): Data in regular geometry.
         out (array_like, optional): Buffer for transformed output, a new
@@ -124,16 +75,22 @@ def to_strixel(data, out=None):
         (array_like) Data in strixel geometry.
     """
 
+    if kind is None:
+        return data
+
+    strx = get_strixel_parameters(kind)
+
     if out is None:
-        out = np.zeros((*data.shape[:-2], *STRIXEL_SHAPE), dtype=data.dtype)
+        out = np.zeros(
+            (*data.shape[:-2], *strx["frame_shape"]), dtype=data.dtype)
 
-    out.reshape(*out.shape[:-2], -1)[..., Iout] = data.reshape(
-        *data.shape[:-2], -1)[..., Iin]
+    out.reshape(*out.shape[:-2], -1)[..., ~strx["mask"]] = data.reshape(
+        *data.shape[:-2], -1)[..., strx["lut"]]
 
     return out
 
 
-def from_strixel(data, out=None):
+def from_strixel(data, out=None, kind="A0123"):
     """Transform from strixel to regular geometry.
 
     Only the last two axes are considered for transformation, input data
@@ -149,13 +106,15 @@ def from_strixel(data, out=None):
         (array_like): Data in regular geometry.
     """
 
+    if kind is None:
+        return data
+
+    strx = get_strixel_parameters(kind)
+
     if out is None:
         out = np.zeros((*data.shape[:-2], *REGULAR_SHAPE), dtype=data.dtype)
 
-    out.reshape(*out.shape[:-2], -1)[..., Iin] = data.reshape(
-        *data.shape[:-2], -1)[..., Iout]
+    out.reshape(*out.shape[:-2], -1)[..., strx["lut"]] = data.reshape(
+        *data.shape[:-2], -1)[..., strx["mask"]]
 
     return out
-
-
-Iout, Iin = transformation_indices1d()