Skip to content
Snippets Groups Projects
Commit fcc08a7c authored by David Hammer's avatar David Hammer
Browse files

DRY preview generation code, AGIPD only preview gain map

parent 39ef1be4
No related branches found
No related tags found
2 merge requests!12Snapshot: field test deployed version as of end of run 202201,!3Base correction device, CalCat interaction, DSSC and AGIPD devices
Hide whitespace changes
Inline Side-by-side
src/calng/AgipdCorrection.py
+ 20
− 42
View file @ fcc08a7c
......@@ -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
......
src/calng/DsscCorrection.py
+ 1
− 1
View file @ fcc08a7c
......@@ -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,
)
......
src/calng/agipd_gpu.py
+ 13
− 47
View file @ fcc08a7c
......@@ -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)
......
src/calng/agipd_gpu_kernels.cpp
+ 3
− 7
View file @ fcc08a7c
......@@ -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 %}
}
}
}
src/calng/base_correction.py
+ 21
− 21
View file @ fcc08a7c
......@@ -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
......
src/calng/base_gpu.py
+ 35
− 25
View file @ fcc08a7c
......@@ -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
......
src/calng/dssc_gpu.py
+ 4
− 4
View file @ fcc08a7c
......@@ -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
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment

Legal Notice | Contacts