diff --git a/src/calng/base_correction.py b/src/calng/base_correction.py index 51d6f5cdaefd3f069ff308317b497d91e68880d2..2d2bf6af7dad6888b73ddc5a8991654ea9dbb75a 100644 --- a/src/calng/base_correction.py +++ b/src/calng/base_correction.py @@ -893,6 +893,10 @@ class BaseCorrection(PythonDevice): corrections, processed_buffer, previews = self.kernel_runner.correct( image_data, cell_table, *additional_data ) + # hack: get extra data out from AGIPD until addon API is updated + if isinstance(corrections, list): + corrections, extra_stuff = corrections + data_hash.merge(extra_stuff) data_hash["corrections"] = corrections # write previews first so addons cannot mess with them diff --git a/src/calng/corrections/AgipdCorrection.py b/src/calng/corrections/AgipdCorrection.py index 02a502d39eecc3da8a8ba407397330c505fafa5d..e73f9a0730f13c78046feede6d3f231efb156d00 100644 --- a/src/calng/corrections/AgipdCorrection.py +++ b/src/calng/corrections/AgipdCorrection.py @@ -6,10 +6,12 @@ from karabo.bound import ( DOUBLE_ELEMENT, FLOAT_ELEMENT, KARABO_CLASSINFO, + NDARRAY_ELEMENT, OUTPUT_CHANNEL, OVERWRITE_ELEMENT, STRING_ELEMENT, UINT32_ELEMENT, + Hash, MetricPrefix, Unit, ) @@ -343,11 +345,6 @@ class AgipdBaseRunner(base_kernel_runner.BaseKernelRunner): self.num_pixels_fs, ) - def extra_processing( - self, out_hash, processed_dat,a cell_table, pulse_table, gain_map - ): - out_hash["numPixelsPerGainStage"] = ... - @property def _gm_map_shape(self): return self._map_shape + (3,) # for gain-mapped constants @@ -611,6 +608,15 @@ class AgipdGpuRunner(AgipdBaseRunner): ), name="common_mode_asic", ) + self.gain_count_kernel = self._xp.RawKernel( + code=base_kernel_runner.get_kernel_template("gaincount.cu").render( + ss_dim=self.num_pixels_ss, + fs_dim=self.num_pixels_fs, + block_size=128, # TODO: tune + num_gain_stages=3, + ), + name="count_pixels_per_gain_stage", + ) def _correct(self, flags, image_data, cell_table, processed_data, gain_map): num_frames = self._xp.uint16(image_data.shape[0]) @@ -658,6 +664,59 @@ class AgipdGpuRunner(AgipdBaseRunner): ), ) + def correct(self, image_data, cell_table, *additional_data): + """Temporary override to utilize gain map for monitoring""" + num_frames = image_data.shape[0] + processed_buffers = self._make_output_buffers( + num_frames, self._correction_flag_preview + ) + image_data = self._xp.asarray(image_data) + additional_data = [self._xp.asarray(data) for data in additional_data] + self._correct( + self._correction_flag_preview, + image_data, + cell_table, + *additional_data, + *processed_buffers, + ) + + preview_buffers = self._preview_data_views( + image_data, *additional_data, *processed_buffers + ) + + if self._correction_flag_preview != self._correction_flag_enabled: + processed_buffers = self._make_output_buffers( + num_frames, self._correction_flag_enabled + ) + self._correct( + self._correction_flag_enabled, + image_data, + cell_table, + *additional_data, + *processed_buffers, + ) + num_frames = buffers["image.data"].shape[0] + gain_count_buffer = self._xp.empty( + (num_frames, 3), dtype=np.uint32 + ) + self.gain_count_kernel( + (num_frames, 1, 1), + (1, 128, 1), + ( + buffers["image.gain"], + num_frames, + gain_count_buffer, + ) + ) + return ( + [ + self._correction_applied_hash, + Hash("numPixelsPerGainStage", gain_count_buffer.get()), + ], + processed_buffers, + preview_buffers, + ) + class AgipdCalcatFriend(base_calcat.BaseCalcatFriend): _constant_enum_class = Constants @@ -786,10 +845,20 @@ class AgipdCalcatFriend(base_calcat.BaseCalcatFriend): return res +def custom_output_schema(use_shmem_handles): + res = schemas.xtdf_output_schema(use_shmem_handles) + ( + NDARRAY_ELEMENT(res) + .key("numPixelsPerGainStage") + .commit(), + ) + return res + + @KARABO_CLASSINFO("AgipdCorrection", deviceVersion) class AgipdCorrection(base_correction.BaseCorrection): # subclass *must* set these attributes - _base_output_schema = schemas.xtdf_output_schema + _base_output_schema = custom_output_schema _calcat_friend_class = AgipdCalcatFriend _constant_enum_class = Constants _correction_steps = correction_steps