diff --git a/src/calng/base_correction.py b/src/calng/base_correction.py
index 3a66fe56e263bac1f926e7f74feaebbec8b05b61..d414ffb74bead6d759bf53af44c34fbba902bf75 100644
--- a/src/calng/base_correction.py
+++ b/src/calng/base_correction.py
@@ -756,17 +756,14 @@ class BaseCorrection(PythonDevice):
self.KARABO_ON_EOS("dataInput", self.handle_eos)
self._enabled_addons = [
- addon_initializer()
+ (node_name, addon_initializer(self._parameters[f"addons.{node_name}"]))
for node_name, addon_initializer in self._available_addons
if self.get(f"addons.{node_name}.enable")
]
if self._enabled_addons:
schema_override = Schema()
output_schema_override = self._base_output_schema
- for addon in self._enabled_addons:
- addon.extend_device_schema(
- schema_override, self._managed_keys, "addons"
- )
+ for _, addon in self._enabled_addons:
addon.extend_output_schema(output_schema_override)
(
# add channel to schema override to re-inject with new output schema
@@ -849,9 +846,17 @@ class BaseCorrection(PythonDevice):
):
self._update_correction_flags()
+ # TODO: only send subhash of reconfiguration (like with addons)
if update.has("preview"):
self._preview_friend.reconfigure(update)
+ if update.has("addons"):
+ # note: can avoid iterating, but it's not that costly
+ for node_name, addon in self._enabled_addons:
+ full_path = f"addons.{node_name}"
+ if update.has(full_path):
+ addon.reconfigure(update[full_path])
+
def _lock_and_update(self, method, background=True):
# TODO: securely handle errors (postReconfigure may succeed, device state not)
def runner():
@@ -1428,6 +1433,25 @@ def add_preview_outputs(schema, channels):
preview_utils.PreviewFriend.add_schema(schema, output_channels=channels)
+def add_addon_nodes(schema, device_class, prefix="addons"):
+ for node_name, addon_class in device_class._available_addons:
+ (
+ NODE_ELEMENT(schema)
+ .key(f"{prefix}.{node_name}")
+ .commit(),
+
+ BOOL_ELEMENT(schema)
+ .key(f"{prefix}.{node_name}.enable")
+ .assignmentOptional()
+ .defaultValue(False)
+ .commit(),
+ )
+ device_class._managed_keys.add(f"{prefix}.{node_name}.enable")
+ addon_class.extend_device_schema(
+ schema, device_class._managed_keys, f"{prefix}.{node_name}"
+ )
+
+
def get_bad_pixel_field_selection(self):
selection = 0
for field in utils.BadPixelValues:
diff --git a/src/calng/correction_addons.py b/src/calng/correction_addons.py
index 4f503adcf79ab43e1df596d7fd7223ee2104753e..304d052e53a122b495e494edbcf8ac41dd2c35d6 100644
--- a/src/calng/correction_addons.py
+++ b/src/calng/correction_addons.py
@@ -1,4 +1,9 @@
-from karabo.bound import NDARRAY_ELEMENT
+from karabo.bound import (
+ DOUBLE_ELEMENT,
+ NDARRAY_ELEMENT,
+ NODE_ELEMENT,
+ Unit,
+)
import numpy as np
@@ -11,16 +16,23 @@ def maybe_get(a):
class BaseCorrectionAddon:
@staticmethod
- def extend_device_schema(schema, managed_keys, prefix="addons"):
+ def extend_device_schema(schema, managed_keys, prefix):
+ """Prefix should be something like 'addons.nodeName'"""
pass
@staticmethod
def extend_output_schema(schema):
pass
+ def __init__(self, config):
+ pass
+
def post_correction(self, processed_data, cell_table, pulse_table, output_hash):
pass
+ def reconfigure(self, changed_config):
+ pass
+
class IntegratedIntensityAddon(BaseCorrectionAddon):
@staticmethod
@@ -38,3 +50,43 @@ class IntegratedIntensityAddon(BaseCorrectionAddon):
output_hash["image.integratedIntensity"] = maybe_get(
np.nansum(processed_data, axis=(1, 2))
)
+
+
+class RandomFramesAddon(BaseCorrectionAddon):
+ @staticmethod
+ def extend_device_schema(schema, managed_keys, prefix):
+ (
+ DOUBLE_ELEMENT(schema)
+ .key(f"{prefix}.probability")
+ .unit(Unit.PERCENT)
+ .assignmentOptional()
+ .defaultValue(50)
+ .reconfigurable()
+ .commit(),
+ )
+ managed_keys.add(f"{prefix}.probability")
+
+ @staticmethod
+ def extend_output_schema(schema):
+ (
+ NODE_ELEMENT(schema)
+ .key("data")
+ .commit(), # TODO: check if rest of node needs rebuilding if exists
+
+ NDARRAY_ELEMENT(schema)
+ .key("data.dataFramePattern")
+ .dtype("UINT8")
+ .commit(),
+ )
+
+ def __init__(self, config):
+ self._probability = config["probability"] / 100
+
+ def post_correction(self, processed_data, cell_table, pulse_table, output_hash):
+ output_hash["data.dataFramePattern"] = (
+ np.random.random(cell_table.size) < self._probability
+ ).astype(np.uint8)
+
+ def reconfigure(self, changed_config):
+ if changed_config.has("probability"):
+ self._probability = changed_config["probability"] / 100
diff --git a/src/calng/corrections/AgipdCorrection.py b/src/calng/corrections/AgipdCorrection.py
index 009bc14a8094bcc3fe51a7645672ab34f6c3b4c3..3356129eaccaebc46bb954173e97a4434a10f486 100644
--- a/src/calng/corrections/AgipdCorrection.py
+++ b/src/calng/corrections/AgipdCorrection.py
@@ -497,7 +497,8 @@ class AgipdCalcatFriend(base_calcat.BaseCalcatFriend):
class AgipdCorrection(base_correction.BaseCorrection):
# subclass *must* set these attributes
_available_addons = [
- ("integratedIntensity", correction_addons.IntegratedIntensityAddon)
+ ("integratedIntensity", correction_addons.IntegratedIntensityAddon),
+ ("randomFrames", correction_addons.RandomFramesAddon),
]
_base_output_schema = schemas.xtdf_output_schema()
_correction_flag_class = CorrectionFlags
@@ -556,6 +557,7 @@ class AgipdCorrection(base_correction.BaseCorrection):
# this is not automatically done by superclass for complicated class reasons
base_correction.add_preview_outputs(expected, AgipdCorrection._preview_outputs)
+ base_correction.add_addon_nodes(expected, AgipdCorrection)
AgipdCalcatFriend.add_schema(expected, AgipdCorrection._managed_keys)
(
# support both CPU and GPU kernels
@@ -778,7 +780,7 @@ class AgipdCorrection(base_correction.BaseCorrection):
buffer_handle, buffer_array = self._shmem_buffer.next_slot()
self.kernel_runner.correct(self._correction_flag_enabled)
- for addon in self._enabled_addons:
+ for _, addon in self._enabled_addons:
addon.post_correction(
self.kernel_runner.processed_data, cell_table, pulse_table, data_hash
)