diff --git a/src/calng/AgipdCorrection.py b/src/calng/AgipdCorrection.py
index a9c5ecc623a31c6fda3dddaa3bd5f453d7a7e84a..09810d4df3ca670d0d04078ec6d79fcec432d0b6 100644
--- a/src/calng/AgipdCorrection.py
+++ b/src/calng/AgipdCorrection.py
@@ -3,19 +3,17 @@ from karabo.bound import (
BOOL_ELEMENT,
FLOAT_ELEMENT,
KARABO_CLASSINFO,
- NDARRAY_ELEMENT,
NODE_ELEMENT,
OUTPUT_CHANNEL,
STRING_ELEMENT,
VECTOR_STRING_ELEMENT,
- Schema,
)
from karabo.common.states import State
from . import shmem_utils, utils
from ._version import version as deviceVersion
from .agipd_gpu import AgipdGainMode, AgipdGpuRunner, BadPixelValues, CorrectionFlags
-from .base_correction import BaseCorrection, add_correction_step_schema
+from .base_correction import BaseCorrection, add_correction_step_schema, preview_schema
from .calcat_utils import AgipdCalcatFriend, AgipdConstants
@@ -55,14 +53,7 @@ class AgipdCorrection(BaseCorrection):
.defaultValue("ADAPTIVE_GAIN")
.options("ADAPTIVE_GAIN,FIXED_HIGH_GAIN,FIXED_MEDIUM_GAIN,FIXED_LOW_GAIN")
.commit(),
- BOOL_ELEMENT(expected)
- .key("sendGainMap")
- .displayedName("Send gain map on dataOutput")
- .assignmentOptional()
- .defaultValue(False)
- .commit(),
)
- # TODO: make sendGainMap reconfigurable
(
STRING_ELEMENT(expected)
@@ -72,14 +63,14 @@ class AgipdCorrection(BaseCorrection):
.defaultValue("")
.commit()
)
- preview_schema = Schema()
- (
- NODE_ELEMENT(preview_schema).key("data").commit(),
- NDARRAY_ELEMENT(preview_schema).key("data.adc").dtype("FLOAT").commit(),
- )
+
(
OUTPUT_CHANNEL(expected)
- .key("preview.outputGain")
+ .key("preview.outputRawGain")
+ .dataSchema(preview_schema)
+ .commit(),
+ OUTPUT_CHANNEL(expected)
+ .key("preview.outputGainMap")
.dataSchema(preview_schema)
.commit(),
)
@@ -207,7 +198,6 @@ class AgipdCorrection(BaseCorrection):
"gain_mode": self.gain_mode,
"bad_pixel_mask_value": self.bad_pixel_mask_value,
"g_gain_value": config.get("corrections.relGainXray.gGainValue"),
- "output_gain_map": config.get("sendGainMap"),
}
self._shmem_buffer_gain_map = None
@@ -285,11 +275,9 @@ class AgipdCorrection(BaseCorrection):
(
preview_raw,
preview_corrected,
- ) = self.gpu_runner.compute_preview(preview_slice_index)
- if self._schema_cache["sendGainMap"]:
- preview_gain = self.gpu_runner.compute_preview_gain(
- preview_slice_index
- )
+ preview_raw_gain,
+ preview_gain_map,
+ ) = self.gpu_runner.compute_previews(preview_slice_index)
# reusing input data hash for sending
data_hash.set("image.data", buffer_handle)
@@ -315,26 +303,16 @@ class AgipdCorrection(BaseCorrection):
self._write_output(data_hash, metadata)
if do_generate_preview:
- if self._schema_cache["sendGainMap"]:
- self._write_combiner_previews(
- (
- ("preview.outputRaw", preview_raw),
- ("preview.outputCorrected", preview_corrected),
- ("preview.outputGain", preview_gain),
- ),
- train_id,
- source,
- )
- # TODO: DRY
- else:
- self._write_combiner_previews(
- (
- ("preview.outputRaw", preview_raw),
- ("preview.outputCorrected", preview_corrected),
- ),
- train_id,
- source,
- )
+ self._write_combiner_previews(
+ (
+ ("preview.outputRaw", preview_raw),
+ ("preview.outputCorrected", preview_corrected),
+ ("preview.outputRawGain", preview_raw_gain),
+ ("preview.outputGainMap", preview_gain_map),
+ ),
+ train_id,
+ source,
+ )
def _load_constant_to_gpu(self, constant, constant_data):
# TODO: encode correction / constant dependencies in a clever way
diff --git a/src/calng/DsscCorrection.py b/src/calng/DsscCorrection.py
index df287d78725e4ba96c84abfbbc3d25daec6c55b3..30c99f7eaebf0923a589bf6cbbe2353638e97243 100644
--- a/src/calng/DsscCorrection.py
+++ b/src/calng/DsscCorrection.py
@@ -102,7 +102,7 @@ class DsscCorrection(BaseCorrection):
pulse_table,
warn_func=self.log_status_warn,
)
- preview_raw, preview_corrected = self.gpu_runner.compute_preview(
+ preview_raw, preview_corrected = self.gpu_runner.compute_previews(
preview_slice_index,
)
diff --git a/src/calng/agipd_gpu.py b/src/calng/agipd_gpu.py
index ff212c6befaaa3c95565cc20d846f80f422c225c..df1bc811e4f804cde296ba578d9023ee6968a659 100644
--- a/src/calng/agipd_gpu.py
+++ b/src/calng/agipd_gpu.py
@@ -39,7 +39,6 @@ class AgipdGpuRunner(base_gpu.BaseGpuRunner):
bad_pixel_mask_value=cupy.nan,
gain_mode=AgipdGainMode.ADAPTIVE_GAIN,
g_gain_value=1,
- output_gain_map=False,
):
self.gain_mode = gain_mode
if self.gain_mode is AgipdGainMode.ADAPTIVE_GAIN:
@@ -48,7 +47,6 @@ class AgipdGpuRunner(base_gpu.BaseGpuRunner):
self.default_gain = cupy.uint8(gain_mode - 1)
self.input_shape = (memory_cells, 2, pixels_x, pixels_y)
self.processed_shape = (memory_cells, pixels_x, pixels_y)
- self.output_gain_map = output_gain_map
super().__init__(
pixels_x,
pixels_y,
@@ -57,12 +55,10 @@ class AgipdGpuRunner(base_gpu.BaseGpuRunner):
input_data_dtype,
output_data_dtype,
)
- if self.output_gain_map:
- self.gain_map_gpu = cupy.empty(self.processed_shape, dtype=cupy.uint8)
- self.preview_gain = np.empty(self.preview_shape, dtype=np.float32)
- else:
- # TODO: don't even set in this case
- self.gain_map_gpu = cupy.empty(0, dtype=cupy.uint8)
+ self.gain_map_gpu = cupy.empty(self.processed_shape, dtype=cupy.float32)
+ self.preview_buffer_getters.extend(
+ [self._get_raw_gain_for_preview, self._get_gain_map_for_preview]
+ )
self.map_shape = (self.constant_memory_cells, self.pixels_x, self.pixels_y)
self.gm_map_shape = self.map_shape + (3,) # for gain-mapped constants
@@ -82,39 +78,19 @@ class AgipdGpuRunner(base_gpu.BaseGpuRunner):
self.update_block_size((1, 1, 64))
- def compute_preview_gain(self, preview_index):
- assert self.output_gain_map
- # TODO: abstract most of this in base_gpu to DRY
- if preview_index < -4:
- raise ValueError(f"No statistic with code {preview_index} defined")
- elif preview_index >= self.memory_cells:
- raise ValueError(f"Memory cell index {preview_index} out of range")
-
- if preview_index >= 0:
- self.gain_map_gpu[preview_index].astype(np.float32).get(
- out=self.preview_gain
- )
- elif preview_index == -1:
- # TODO: confirm that max is pixel and not integrated intensity
- # separate from next case because dtype not applicable here
- cupy.max(self.gain_map_gpu, axis=0).astype(cupy.float32).get(
- out=self.preview_gain
- )
- elif preview_index in (-2, -3, -4):
- stat_fun = {-1: cupy.max, -2: cupy.mean, -3: cupy.sum, -4: cupy.std}[
- preview_index
- ]
- stat_fun(self.gain_map_gpu, axis=0, dtype=cupy.float32).get(
- out=self.preview_gain
- )
- return self.preview_gain
-
- def _preview_preprocess_raw(self):
+ def _get_raw_for_preview(self):
return self.input_data_gpu[:, 0]
- def _preview_preprocess_corr(self):
+ def _get_corrected_for_preview(self):
return self.processed_data_gpu
+ # special to AGIPD
+ def _get_raw_gain_for_preview(self):
+ return self.input_data_gpu[:, 1]
+
+ def _get_gain_map_for_preview(self):
+ return self.gain_map_gpu
+
def load_thresholds(self, threshold_map):
# shape: y, x, memory cell, thresholds and gain values
# note: the gain values are something like means used to derive thresholds
@@ -277,15 +253,6 @@ class AgipdGpuRunner(base_gpu.BaseGpuRunner):
),
)
- def get_gain_map(self, output_order, out=None):
- assert self.output_gain_map
- return cupy.ascontiguousarray(
- cupy.transpose(
- self.gain_map_gpu,
- utils.transpose_order(self._corrected_axis_order, output_order),
- )
- ).get(out=out)
-
def _init_kernels(self):
kernel_source = self._kernel_template.render(
{
@@ -296,7 +263,6 @@ class AgipdGpuRunner(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),
- "output_gain_map": self.output_gain_map,
}
)
self.source_module = cupy.RawModule(code=kernel_source)
diff --git a/src/calng/agipd_gpu_kernels.cpp b/src/calng/agipd_gpu_kernels.cpp
index 0561e84bc2b4f8efd44c4ec7dce475f972eab991..d5d5ca02394633510b5d394806b7ea8a961f6113 100644
--- a/src/calng/agipd_gpu_kernels.cpp
+++ b/src/calng/agipd_gpu_kernels.cpp
@@ -23,7 +23,7 @@ extern "C" {
const float g_gain_value,
const unsigned int* bad_pixel_map,
const float bad_pixel_mask_value,
- unsigned char* gain_map,
+ float* gain_map, // TODO: more compact yet plottable representation
{{output_data_dtype}}* output) {
const size_t X = {{pixels_x}};
const size_t Y = {{pixels_y}};
@@ -99,6 +99,7 @@ extern "C" {
gain = 2;
}
}
+ gain_map[output_index] = (float)gain;
const size_t map_index = map_cell * map_stride_cell +
y * map_stride_y +
@@ -111,6 +112,7 @@ extern "C" {
if ((corr_flags & BPMASK) && bad_pixel_map[gm_map_index]) {
corrected = bad_pixel_mask_value;
+ gain_map[output_index] = bad_pixel_mask_value;
} else {
if (corr_flags & OFFSET) {
corrected -= offset_map[gm_map_index];
@@ -132,10 +134,6 @@ extern "C" {
{% else %}
output[output_index] = ({{output_data_dtype}})corrected;
{% endif %}
-
- {% if output_gain_map %}
- gain_map[output_index] = gain;
- {% endif %}
} else {
// TODO: decide what to do when we cannot threshold
{% if output_data_dtype == "half" %}
@@ -144,9 +142,7 @@ extern "C" {
output[data_index] = ({{output_data_dtype}})corrected;
{% endif %}
- {% if output_gain_map %}
gain_map[data_index] = 255;
- {% endif %}
}
}
}
diff --git a/src/calng/base_correction.py b/src/calng/base_correction.py
index c8a64eb44dea00bb86a13d53e657a2fe70f0331e..49279fde4806810c1ab059a7c7eb4aba5f3a4f36 100644
--- a/src/calng/base_correction.py
+++ b/src/calng/base_correction.py
@@ -40,6 +40,18 @@ from ._version import version as deviceVersion
PROCESSING_STATE_TIMEOUT = 10
+preview_schema = Schema()
+(
+ NODE_ELEMENT(preview_schema).key("image").commit(),
+ NDARRAY_ELEMENT(preview_schema).key("image.data").dtype("FLOAT").commit(),
+ UINT64_ELEMENT(preview_schema)
+ .key("image.trainId")
+ .displayedName("Train ID")
+ .assignmentOptional()
+ .defaultValue(0)
+ .commit(),
+)
+
@KARABO_CLASSINFO("BaseCorrection", deviceVersion)
class BaseCorrection(PythonDevice):
@@ -224,8 +236,8 @@ class BaseCorrection(PythonDevice):
.metricPrefix(MetricPrefix.GIGA)
.description(
"Corrected trains are written to shared memory locations. These are "
- "pre-allocated and re-used (circular buffer). This parameter determines "
- "much memory to set aside for the buffer."
+ "pre-allocated and re-used (circular buffer). This parameter "
+ "determines how much memory to set aside for the buffer."
)
.assignmentOptional()
.defaultValue(10)
@@ -310,10 +322,10 @@ class BaseCorrection(PythonDevice):
.key("dataFormat.outputAxisOrder")
.displayedName("Output axis order")
.description(
- "Axes of main data output can be reordered after correction. Axis order "
- "is specified as string consisting of 'x', 'y', and 'c', with the "
- "latter indicating the memory cell axis. The default value of 'cxy' "
- "puts pixels along the fast axes."
+ "Axes of main data output can be reordered after correction. Axis "
+ "order is specified as string consisting of 'x', 'y', and 'c', with "
+ "the latter indicating the memory cell axis. The default value of "
+ "'cxy' puts pixels on the fast axes."
)
.assignmentOptional()
.defaultValue("cxy")
@@ -349,17 +361,6 @@ class BaseCorrection(PythonDevice):
.commit()
)
- preview_schema = Schema()
- (
- NODE_ELEMENT(preview_schema).key("image").commit(),
- NDARRAY_ELEMENT(preview_schema).key("image.data").dtype("FLOAT").commit(),
- UINT64_ELEMENT(preview_schema)
- .key("image.trainId")
- .displayedName("Train ID")
- .assignmentOptional()
- .defaultValue(0)
- .commit(),
- )
(
NODE_ELEMENT(expected).key("preview").displayedName("Preview").commit(),
OUTPUT_CHANNEL(expected)
@@ -377,7 +378,6 @@ class BaseCorrection(PythonDevice):
.defaultValue(True)
.reconfigurable()
.commit(),
- # TODO: Split into AGIPD-specific or see if others like cell ID over pulse ID
INT32_ELEMENT(expected)
.key("preview.index")
.displayedName("Index (or stat) for preview")
@@ -518,7 +518,7 @@ class BaseCorrection(PythonDevice):
if parse_version(karaboVersion) >= parse_version("2.11"):
# TODO: the CalCatFriend could add these for us
- # note: overly complicated slot function creation necessary for closure to work
+ # note: overly complicated for closure to work
def make_wrapper_capturing_constant(constant):
def aux():
self.calcat_friend.get_specific_constant_version_and_call_me_back(
@@ -672,7 +672,7 @@ class BaseCorrection(PythonDevice):
timestamp = Timestamp(Epochstamp(), Trainstamp(train_id))
metadata = ChannelMetaData(source, timestamp)
for channel_name, data in channel_data_pairs:
- preview_hash.set("data.adc", data[..., np.newaxis])
+ preview_hash.set("image.data", data)
channel = self.signalSlotable.getOutputChannel(channel_name)
channel.write(preview_hash, metadata, False)
channel.update()
@@ -809,7 +809,7 @@ class BaseCorrection(PythonDevice):
)
self.set("dataFormat.memoryCells", image_data.shape[0])
with self._buffer_lock:
- # TODO: re-validate frame filter against new number of cells
+ self._update_frame_filter()
self._update_buffers()
# DataAggregator typically tells us the wrong axis order
diff --git a/src/calng/base_gpu.py b/src/calng/base_gpu.py
index 4850dcbaa8c4108369daffc85719bd4184bd32ee..5a4821a0f954166587b5a1543efbf94570f5eb64 100644
--- a/src/calng/base_gpu.py
+++ b/src/calng/base_gpu.py
@@ -69,20 +69,24 @@ class BaseGpuRunner:
self.processed_shape, dtype=output_data_dtype
)
self.reshaped_data_gpu = None # currently not reusing buffer
- self.preview_raw = np.empty(self.preview_shape, dtype=np.float32)
- self.preview_corrected = np.empty(self.preview_shape, dtype=np.float32)
+
+ # default preview layers: raw and corrected (subclass can extend)
+ self.preview_buffer_getters = [
+ self._get_raw_for_preview,
+ self._get_corrected_for_preview,
+ ]
# to get data from respective buffers to cell, x, y shape for preview computation
- def _preview_preprocess_raw(self):
- """Should return view of self.input_data_gpu with shape (cell, x, y)"""
+ def _get_raw_for_preview(self):
+ """Should return view of self.input_data_gpu with shape (cell, x/y, x/y)"""
raise NotImplementedError()
- def _preview_preprocess_corr(self):
- """Should return view of self.processed_data_gpu with shape (cell, x, y)"""
+ def _get_corrected_for_preview(self):
+ """Should return view of self.processed_data_gpu with shape (cell, x/y, x/y)"""
raise NotImplementedError()
def flush_buffers(self):
- """Optional reset GPU buffers (implement in appropriate subclasses"""
+ """Optional reset GPU buffers (implement in subclasses which need this)"""
pass
def correct(self, flags):
@@ -127,7 +131,7 @@ class BaseGpuRunner:
def load_cell_table(self, cell_table):
self.cell_table_gpu.set(cell_table)
- def compute_preview(self, preview_index):
+ def compute_previews(self, preview_index):
"""Generate single slice or reduction preview of raw and corrected data
Special values of preview_index are -1 for max, -2 for mean, -3 for sum, and
@@ -147,23 +151,29 @@ class BaseGpuRunner:
raise ValueError(f"Memory cell index {preview_index} out of range")
# TODO: enum around reduction type
- for (preprocces, output_buffer) in (
- (self._preview_preprocess_raw, self.preview_raw),
- (self._preview_preprocess_corr, self.preview_corrected),
- ):
- image_data = preprocces()
- if preview_index >= 0:
- image_data[preview_index].astype(np.float32).get(out=output_buffer)
- elif preview_index == -1:
- # TODO: confirm that max is pixel and not integrated intensity
- # separate from next case because dtype not applicable here
- cupy.max(image_data, axis=0).astype(cupy.float32).get(out=output_buffer)
- elif preview_index in (-2, -3, -4):
- stat_fun = {-1: cupy.max, -2: cupy.mean, -3: cupy.sum, -4: cupy.std}[
- preview_index
- ]
- stat_fun(image_data, axis=0, dtype=cupy.float32).get(out=output_buffer)
- return self.preview_raw, self.preview_corrected
+ return tuple(
+ self._compute_a_preview(image_data=getter(), preview_index=preview_index)
+ for getter in self.preview_buffer_getters
+ )
+
+ def _compute_a_preview(self, image_data, preview_index):
+ """image_data must have cells on first axis; X and Y order is not important
+ here for now (and can differ between AGIPD and DSSC)"""
+ if preview_index >= 0:
+ # TODO: reuse pinned buffers for this
+ return image_data[preview_index].astype(np.float32).get()
+ elif preview_index == -1:
+ # TODO: confirm that max is pixel and not integrated intensity
+ # separate from next case because dtype not applicable here
+ return cupy.nanmax(image_data, axis=0).astype(cupy.float32).get()
+ elif preview_index in (-2, -3, -4):
+ stat_fun = {
+ -1: cupy.nanmax,
+ -2: cupy.nanmean,
+ -3: cupy.nansum,
+ -4: cupy.nanstd,
+ }[preview_index]
+ return stat_fun(image_data, axis=0, dtype=cupy.float32).get()
def update_block_size(self, full_block, target_shape=None):
"""Compute grid such that thread block grid covers target shape
diff --git a/src/calng/dssc_gpu.py b/src/calng/dssc_gpu.py
index 2b4bac53f025f8bcbe3dde08355a35294319713d..138e301647f3ed138422eeadc42cfc4438a27d6d 100644
--- a/src/calng/dssc_gpu.py
+++ b/src/calng/dssc_gpu.py
@@ -44,11 +44,11 @@ class DsscGpuRunner(base_gpu.BaseGpuRunner):
self.update_block_size((1, 1, 64))
- def _preview_preprocess_raw(self):
- return cupy.transpose(self.input_data_gpu, (0, 2, 1))
+ def _get_raw_for_preview(self):
+ return self.input_data_gpu
- def _preview_preprocess_corr(self):
- return cupy.transpose(self.processed_data_gpu, (0, 2, 1))
+ def _get_corrected_for_preview(self):
+ return self.processed_data_gpu
def load_offset_map(self, offset_map):
# can have an extra dimension for some reason