Add native implementation for LPD correction loop

setup.py

@@ -18,6 +18,13 @@ ext_modules = [
         extra_compile_args=["-fopenmp", "-march=native"],
         extra_link_args=["-fopenmp"],
     ),
+    Extension(
+        'cal_tools.lpdalgs',
+        ['src/cal_tools/lpdalgs.pyx'],
+        include_dirs=[numpy.get_include()],
+        extra_compile_args=['-O3', '-fopenmp', '-march=native'],
+        extra_link_args=['-fopenmp'],
+    )
 ]
 
 

src/cal_tools/lpdalgs.pyx

+from cython cimport boundscheck, wraparound, cdivision
+from cython.view cimport contiguous
+from cython.parallel cimport prange
+
+from cal_tools.enums import BadPixels
+
+ctypedef unsigned short cell_t
+ctypedef unsigned short data_t
+ctypedef float pixel_t
+ctypedef unsigned char gain_t
+ctypedef unsigned int mask_t
+
+cdef mask_t WRONG_GAIN_VALUE = BadPixels.WRONG_GAIN_VALUE, \
+    VALUE_OUT_OF_RANGE = BadPixels.VALUE_OUT_OF_RANGE, \
+    VALUE_IS_NAN = BadPixels.VALUE_IS_NAN
+
+
+@boundscheck(False)
+@wraparound(False)
+@cdivision(True)
+def correct_lpd_frames(
+    # (pulse, x, y)
+    data_t[:, :, ::contiguous] in_data,
+    cell_t[::contiguous] in_cell,
+
+    # (pulse, x, y)
+    pixel_t[:, :, ::contiguous] out_pixels,
+    gain_t[:, :, ::contiguous] out_gain,
+    mask_t[:, :, ::contiguous] out_mask,
+
+    # (x, y, cell, gain)
+    float[:, :, :, ::contiguous] ccv_offset,
+    float[:, :, :, ::contiguous] ccv_gain,
+    mask_t[:, :, :, ::contiguous] ccv_mask,
+
+    int num_threads=1,
+):
+    cdef int frame, ss, fs
+    cdef cell_t cell
+    cdef pixel_t pixel
+    cdef gain_t gain
+    cdef mask_t mask
+
+    for frame in prange(in_data.shape[0], nogil=True, num_threads=num_threads):
+        for ss in range(in_data.shape[1]):
+            for fs in range(in_data.shape[2]):
+                cell = in_cell[frame]
+                pixel = <pixel_t>(in_data[frame, ss, fs] & 0b0000111111111111)
+                gain = in_data[frame, ss, fs] & 0b0000000000000011
+
+                if gain <= 2:
+                    mask = ccv_mask[ss, fs, cell, gain]
+                else:
+                    pixel = 0.0
+                    gain = 0
+                    mask = WRONG_GAIN_VALUE
+
+                pixel = pixel - ccv_offset[ss, fs, cell, gain]
+
+                if ccv_gain[ss, fs, cell, gain] != 0.0:
+                    pixel = pixel * ccv_gain[ss, fs, cell, gain]
+                else:
+                    pixel = 0.0
+                    mask = mask | VALUE_IS_NAN
+
+                if pixel > 1e7 or pixel < -1e7:
+                    pixel = 0.0
+                    mask = mask | VALUE_OUT_OF_RANGE
+
+                out_pixels[frame, ss, fs] = pixel
+                out_gain[frame, ss, fs] = gain
+                out_mask[frame, ss, fs] = mask