Fix constant loading, edge cases in kernel

src/calng/JungfrauCorrection.py

@@ -90,7 +90,7 @@ class JungfrauGpuRunner(base_gpu.BaseGpuRunner):
                 "input_data_dtype": utils.np_dtype_to_c_type(self.input_data_dtype),
                 "output_data_dtype": utils.np_dtype_to_c_type(self.output_data_dtype),
                 "corr_enum": utils.enum_to_c_template(CorrectionFlags),
-                "burst_mode": self.burst_mode,
+                "burst_mode": self.memory_cells > 1,
             }
         )
         for i, line in enumerate(kernel_source.split("\n")):
@@ -113,15 +113,6 @@ class JungfrauGpuRunner(base_gpu.BaseGpuRunner):
         if self.burst_mode:
             self.cell_table_gpu.set(cell_table)
 
-    def load_constant(self, constant, constant_data):
-        if constant is JungfrauConstants.Offset10Hz:
-            self.offset_map_gpu.set(constant_data.astype(np.float32))
-        elif constant is JungfrauConstants.RelativeGain10Hz:
-            self.rel_gain_map_gpu.set(constant_data.astype(np.float32))
-        elif constant is JungfrauConstants.BadPixelsDark10Hz:
-            self.bad_pixel_map_gpu.set(constant_data)
-
-
     def correct(self, flags):
         self.correction_kernel(
             self.full_grid,
@@ -240,7 +231,7 @@ class JungfrauCalcatFriend(calcat_utils.BaseCalcatFriend):
         res["Sensor Temperature"] = self._get_param("sensorTemperature")
         res["Gain Setting"] = self._get_param("gainSetting")
         gain_mode = JungfrauGainMode[self._get_param("gainMode")]
-        if gain_mode is not JungfrauGainMode.dynamicgain:
+        if gain_mode != 0: #is not JungfrauGainMode.dynamicgain:
             # TODO: figure out what to set
             res["Gain mode"] = 1
         return res
@@ -392,6 +383,21 @@ class JungfrauCorrection(BaseCorrection):
             )
 
     def _load_constant_to_runner(self, constant, constant_data):
-        self.kernel_runner.load_constant(
-            constant, np.transpose(constant_data, (2, 0, 1, 3))
-        )
+        if constant is JungfrauConstants.Offset10Hz:
+            constant_data = np.transpose(constant_data, (2, 0, 1, 3))
+            self.kernel_runner.offset_map_gpu.set(constant_data.astype(np.float32))
+            if not self.get("corrections.offset.available"):
+                self.set("corrections.offset.available", True)
+        elif constant is JungfrauConstants.RelativeGain10Hz:
+            constant_data = np.transpose(constant_data, (2, 1, 0, 3))
+            self.kernel_runner.rel_gain_map_gpu.set(constant_data.astype(np.float32))
+            if not self.get("corrections.relGain.available"):
+                self.set("corrections.relGain.available", True)
+        elif constant is JungfrauConstants.BadPixelsDark10Hz:
+            constant_data = np.transpose(constant_data, (2, 0, 1, 3))
+            self.kernel_runner.bad_pixel_map_gpu.set(constant_data)
+            if not self.get("corrections.badPixels.available"):
+                self.set("corrections.badPixels.available", True)
+
+        self._update_correction_flags()
+        self.log_status_info(f"Done loading {constant.name} to GPU")

src/calng/kernels/jungfrau_gpu.cu

@@ -25,7 +25,6 @@ extern "C" {
 			return;
 		}
 
-		// note: strides differ from numpy strides because unit here is sizeof(...), not byte
 		const size_t data_stride_x = 1;
 		const size_t data_stride_y = X * data_stride_x;
 		const size_t data_stride_cell = Y * data_stride_y;
@@ -35,44 +34,45 @@ extern "C" {
 		float res = (float)data[data_index];
 
 		// gain mapped constant shape: cell, y, x, gain_level (dim size 3)
+		// note: in fixed gain mode, constant still provides data for three stages
 		const size_t map_stride_gain = 1;
 		const size_t map_stride_x = 3 * map_stride_gain;
 		const size_t map_stride_y = X * map_stride_x;
 		const size_t map_stride_cell = Y * map_stride_y;
 
+		// TODO: warn user about cell_table value of 255 in either mode
+		// note: cell table may contain 255 if data didn't arrive
 		{% if burst_mode %}
+		// burst mode: "cell 255" will get copied
+		// TODO: consider masking "cell 255"
 		const size_t map_cell = cell_table[memory_cell];
 		{% else %}
+		// single cell: "cell 255" will get "corrected"
 		const size_t map_cell = 0;
 		{% endif %}
 
 		if (map_cell < map_memory_cells) {
 			unsigned char gain = gain_stage[data_index];
-
-			{% if burst_mode %}
 			if (gain == 2) {
-				gain = 1;
-			} else if (gain == 3) {
-				gain = 2;
-			}
-			{% else %}
-			if (gain == 3) {
-				gain = 2;
-			}
-			{% endif %}
-
-			const size_t map_index = map_cell * map_stride_cell +
-				y * map_stride_y +
-				x * map_stride_x +
-				gain * map_stride_gain;
-			if ((corr_flags & BPMASK) && bad_pixel_map[map_index]) {
+				// gain should be read as 0, 1, or 3; value of 2 indicates issue
 				res = bad_pixel_mask_value;
 			} else {
-				if (corr_flags & OFFSET) {
-					res -= offset_map[map_index];
+				if (gain == 3) {
+					gain = 2;
 				}
-				if (corr_flags & REL_GAIN) {
-					res /= rel_gain_map[map_index];
+				const size_t map_index = map_cell * map_stride_cell +
+					y * map_stride_y +
+					x * map_stride_x +
+					gain * map_stride_gain;
+				if ((corr_flags & BPMASK) && bad_pixel_map[map_index]) {
+					res = bad_pixel_mask_value;
+				} else {
+					if (corr_flags & OFFSET) {
+						res -= offset_map[map_index];
+					}
+					if (corr_flags & REL_GAIN) {
+						res /= rel_gain_map[map_index];
+					}
 				}
 			}
 		}