diff --git a/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb b/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
index 29b213b34e5a7e4d1842f512f8c532e6569548b8..b313a0b0c3fa20b31fcfacda366c67441d732163 100644
--- a/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
+++ b/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
@@ -216,13 +216,9 @@
"metadata": {},
"outputs": [],
"source": [
- "strixel_transform = None\n",
- "output_frame_shape = None\n",
- "\n",
"if strixel_sensor:\n",
- " from cal_tools.jfalgs import strixel_transform, strixel_shape, strixel_double_pixels\n",
- " output_frame_shape = strixel_shape()\n",
- " Ydouble, Xdouble = strixel_double_pixels()\n",
+ " from cal_tools.jfstrixel import STRIXEL_SHAPE as strixel_frame_shape, double_pixel_indices, to_strixel\n",
+ " Ydouble, Xdouble = double_pixel_indices()\n",
" print('Strixel sensor transformation enabled')"
]
},
@@ -360,8 +356,8 @@
" g = gain[index]\n",
" \n",
" # Copy gain over first to keep it at the original 3 for low gain.\n",
- " if strixel_transform is not None:\n",
- " strixel_transform(g, out=gain_corr[index, ...])\n",
+ " if strixel_sensor:\n",
+ " to_strixel(g, out=gain_corr[index, ...])\n",
" else:\n",
" gain_corr[index, ...] = g\n",
"\n",
@@ -405,11 +401,10 @@
"\n",
" msk = np.choose(g, np.moveaxis(mask_cell, -1, 0))\n",
"\n",
- " if strixel_transform is not None:\n",
- " strixel_transform(d, out=data_corr[index, ...])\n",
+ " if strixel_sensor:\n",
+ " to_strixel(d, out=data_corr[index, ...])\n",
" data_corr[index, :, Ydouble, Xdouble] /= strixel_double_norm\n",
- "\n",
- " strixel_transform(msk, out=mask_corr[index, ...])\n",
+ " to_strixel(msk, out=mask_corr[index, ...])\n",
" else:\n",
" data_corr[index, ...] = d\n",
" mask_corr[index, ...] = msk"
@@ -525,8 +520,8 @@
" offset_map, mask, gain_map = constants[local_karabo_da]\n",
" \n",
" # Determine total output shape.\n",
- " if output_frame_shape is not None:\n",
- " oshape = (*ishape[:-2], *output_frame_shape)\n",
+ " if strixel_sensor:\n",
+ " oshape = (*ishape[:-2], *strixel_frame_shape)\n",
" else:\n",
" oshape = ishape\n",
"\n",
diff --git a/setup.py b/setup.py
index 0e39c2bf29538423d82f8ec6ccbe2c2e5c8ed96d..c9f0722aa2d75076598cfcf57cf4bed7c71236e3 100644
--- a/setup.py
+++ b/setup.py
@@ -24,12 +24,6 @@ ext_modules = [
'-ftree-vectorize', '-frename-registers'],
extra_link_args=['-fopenmp'],
),
- Extension(
- "cal_tools.jfalgs",
- ['src/cal_tools/jfalgs.pyx'],
- extra_compile_args=['-O3', '-march=native', '-ftree-vectorize',
- '-frename-registers']
- ),
Extension(
"cal_tools.gotthard2.gotthard2algs",
["src/cal_tools/gotthard2/gotthard2algs.pyx"],
diff --git a/src/cal_tools/jfalgs.pyx b/src/cal_tools/jfalgs.pyx
deleted file mode 100644
index fd1476918d8a8f165150abe1c5bbeb6175721913..0000000000000000000000000000000000000000
--- a/src/cal_tools/jfalgs.pyx
+++ /dev/null
@@ -1,93 +0,0 @@
-
-from cython cimport boundscheck, wraparound, cdivision
-from cython.view cimport contiguous
-
-import numpy as np
-
-
-ctypedef fused jf_data_t:
- unsigned short # raw pixel data
- float # corrected pixel data
- unsigned int # mask data
- unsigned char # gain data
-
-
-DEF STRIXEL_Y = 86
-DEF STRIXEL_X = 1024 * 3 + 18
-
-
-def strixel_shape():
- return STRIXEL_Y, STRIXEL_X
-
-
-def strixel_double_pixels():
- """Build index arrays for double-size pixels.
-
- In raw data, 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 after strixelation.
- """
-
- Ydouble = []
- Xdouble = []
-
- 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)
-
- return np.array(Ydouble), np.array(Xdouble)
-
-
-@boundscheck(False)
-@wraparound(False)
-@cdivision(True)
-def strixel_transform(jf_data_t[:, :, ::contiguous] data,
- jf_data_t[:, :, ::contiguous] out = None):
- """Reorder raw data to physical strixel sensor layout. """
-
- if data.shape[1] < 256 or data.shape[2] < 256:
- raise ValueError('Pixel shape of data may not be below (256, 256)')
-
- if out is None:
- import numpy as np
- out = np.zeros((data.shape[0], STRIXEL_Y, STRIXEL_X), dtype=np.float32)
- elif data.shape[0] > out.shape[0]:
- raise ValueError('Cell shape of data exceeds out')
- elif out.shape[1] < STRIXEL_Y or out.shape[2] < STRIXEL_X:
- raise ValueError(f'Pixel shape of out may not be below '
- f'({STRIXEL_Y}, {STRIXEL_X})')
-
- cdef int cell, yin, xin, xout, yout, igap
-
- for cell in range(data.shape[0]):
- # Normal pixels.
- for yin in range(256):
- yout = yin // 3
-
- for xin in range(1024) :
- xout = 774 * (xin // 256) + 3 * (xin % 256) + yin % 3
- out[cell, yout, xout] = data[cell, yin, xin]
-
- # Gap pixels.
- for yin in range(256):
- yout = 2 * (yin // 6)
-
- for igap in range(3) :
- # Left side of the gap.
- xin = igap * 256 + 255
- xout = igap * 774 + 765 + yin % 6
- out[cell, yout, xout] = data[cell, yin, xin]
- out[cell, yout+1, xout] = data[cell, yin, xin]
-
- # Right side of the gap.
- xin = igap * 256 + 255 + 1
- xout = igap * 774 + 765 + 11 - yin % 6
- out[cell, yout, xout] = data[cell, yin, xin]
- out[cell, yout+1, xout] = data[cell, yin, xin]
-
- return out
diff --git a/src/cal_tools/jfstrixel.py b/src/cal_tools/jfstrixel.py
new file mode 100644
index 0000000000000000000000000000000000000000..189b036f6cddf4aaa8ba802e908a999afb6cdfe3
--- /dev/null
+++ b/src/cal_tools/jfstrixel.py
@@ -0,0 +1,161 @@
+
+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.
+ """
+
+ 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]
+
+ Iout = np.where(strixel_pixel_idx.ravel() != -1)[0].astype(np.int64)
+ Iin = strixel_pixel_idx.ravel()[Iout].astype(np.int64)
+
+ return Iout, Iin
+
+
+def double_pixel_indices():
+ """Build index arrays for double-size pixels.
+
+ In raw data, 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.
+ """
+
+ Ydouble = []
+ Xdouble = []
+
+ 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)
+
+ return np.array(Ydouble), np.array(Xdouble)
+
+
+def to_strixel(data, out=None):
+ """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
+ one is allocated if omitted. Must match all non-frame axes
+ of input data and able to hold strixel frame.
+
+ Returns:
+ (array_like) Data in strixel geometry.
+ """
+
+ if out is None:
+ out = np.zeros((*data.shape[:-2], *STRIXEL_SHAPE), dtype=data.dtype)
+
+ out.reshape(*out.shape[:-2], -1)[..., Iout] = data.reshape(
+ *data.shape[:-2], -1)[..., Iin]
+
+ return out
+
+
+def from_strixel(data, out=None):
+ """Transform from strixel to regular 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 strixel geometry.
+ out (array_like, optional): Buffer for transformed output, a new
+ one is allocated if omitted. Must match all non-frame axes
+ of input data and able to hold regular frame.
+
+ Returns:
+ (array_like): Data in regular geometry.
+ """
+
+ 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]
+
+ return out
+
+
+Iout, Iin = transformation_indices1d()