Initial Jungfrau device boilerplate

setup.py

@@ -26,6 +26,7 @@ setup(name='calng',
           'karabo.bound_device': [
               'AgipdCorrection = calng.AgipdCorrection:AgipdCorrection',
               'DsscCorrection = calng.DsscCorrection:DsscCorrection',
+              'JungfrauCorrection = calng.JungfrauCorrection:JungfrauCorrection',
               'ModuleStacker = calng.ModuleStacker:ModuleStacker',
               'ShmemToZMQ = calng.ShmemToZMQ:ShmemToZMQ',
           ],

src/calng/JungfrauCorrection.py

+import enum
+import timeit
+
+import cupy
+import numpy as np
+from karabo.bound import (
+    DOUBLE_ELEMENT,
+    KARABO_CLASSINFO,
+    OVERWRITE_ELEMENT,
+    VECTOR_STRING_ELEMENT,
+)
+from karabo.common.states import State
+
+from . import base_gpu, calcat_utils, utils
+from ._version import version as deviceVersion
+from .base_correction import BaseCorrection, add_correction_step_schema
+
+
+class JungfrauConstants(enum.Enum):
+    Offset10Hz = enum.auto()
+    BadPixelsDark10Hz = enum.auto()
+
+
+class JungfrauGainMode(enum.IntEnum):
+    # TODO: coordinate with pycalibration and check what is saved by control device
+    # gain stages: G0, G1, G2 (plus secret HG0 - high CDS)
+    dynamicgain = enum.auto()  # default; uses G0-G2
+    dynamichg0 = enum.auto()  # uses HG0, G1, G2
+    fixgain1 = enum.auto()  # fix, use only G1
+    fixgain2 = enum.auto()  # fix, use only G2
+    forceswitchg1 = enum.auto()  # only for darks
+    forceswitchg2 = enum.auto()  # only for darks
+
+
+class CorrectionFlags(enum.IntFlag):
+    NONE = 0
+    OFFSET = 1
+    REL_GAIN = 2
+
+
+class JungfrauGpuRunner(base_gpu.BaseGpuRunner):
+    _kernel_source_filename = "jungfrau_gpu_kernels.cpp"
+    _corrected_axis_order = ...  # TODO: get specs for jungfrau data
+
+    def __init__(
+        self,
+        pixels_x,
+        pixels_y,
+        memory_cells,
+        constant_memory_cells,
+        input_data_dtype=cupy.uint16,
+        output_data_dtype=cupy.float32,
+    ):
+        self.input_shape = ...
+        self.processed_shape = ...
+        super().__ini__(
+            pixels_x,
+            pixels_y,
+            memory_cells,
+            constant_memory_cells,
+            input_data_dtype,
+            output_data_dtype,
+        )
+        self.map_shape = ...
+        # is jungfrau stuff gain mapped?
+        self.offset_map_gpu = cupy.zeros(..., dtype=cupy.float32)
+        self.rel_gain_map_gpu = cupy.ones(..., dtype=cupy.float32)
+
+
+class JungfrauCalcatFriend(calcat_utils.BaseCalcatFriend):
+    _constant_enum_class = JungfrauConstants
+
+    def __init__(self, device, *args, **kwargs):
+        super().__init__(device, *args, **kwargs)
+        self._constants_need_conditions = {
+            JungfrauConstants.Offset10Hz: self.dark_condition,
+            JungfrauConstants.BadPixelsDark10Hz: self.dark_condition,
+        }
+
+    @staticmethod
+    def add_schema(
+        schema,
+        managed_keys,
+        param_prefix="constantParameters",
+        status_prefix="foundConstants",
+    ):
+        super(JungfrauCalcatFriend, JungfrauCalcatFriend).add_schema(
+            schema, managed_keys, "jungfrau-Type", param_prefix, status_prefix
+        )
+
+        # set some defaults for common parameters
+        (
+            OVERWRITE_ELEMENT(schema)
+            .key(f"{param_prefix}.pixelsX")
+            .setNewDefaultValue(1024)
+            .commit(),
+            OVERWRITE_ELEMENT(schema)
+            .key(f"{param_prefix}.pixelsY")
+            .setNewDefaultValue(512)
+            .commit(),
+            OVERWRITE_ELEMENT(schema)
+            .key(f"{param_prefix}.memoryCells")
+            .setNewDefaultValue(1)
+            .commit(),
+            OVERWRITE_ELEMENT(schema)
+            .key(f"{param_prefix}.biasVoltage")
+            .setNewDefaultValue(90)
+            .commit(),
+        )
+
+        # add extra parameters
+        (
+            DOUBLE_ELEMENT(schema)
+            .key(f"{param_prefix}.integrationTime")
+            .displayedName("Integration time")
+            .description("Integration time in ms")
+            .assignmentOptional()
+            .defaultValue(350)
+            .reconfigurable()
+            .commit(),
+            DOUBLE_ELEMENT(schema)
+            .key(f"{param_prefix}.sensorTemperature")
+            .displayedName("Sensor temperature")
+            .description("Sensor temperature in K")
+            .assignmentOptional()
+            .defaultValue(291)
+            .reconfigurable()
+            .commit(),
+            DOUBLE_ELEMENT(schema)
+            .key(f"{param_prefix}.gainSetting")
+            .displayedName("Gain setting")
+            .description("Feedback capacitor setting; 0 is default, 1 is HG0")
+            .assignmentOptional()
+            .defaultValue(0)
+            .reconfigurable()
+            .commit(),
+        )
+        managed_keys.add(f"{param_prefix}.integrationTime")
+        managed_keys.add(f"{param_prefix}.sensorTemperature")
+        managed_keys.add(f"{param_prefix}.gainSetting")
+
+        calcat_utils.add_status_schema_from_enum(
+            schema, status_prefix, JungfrauConstants
+        )
+
+    def dark_condition(self):
+        res = calcat_utils.OperatingConditions()
+        res["Memory cells"] = self._get_param("memoryCells")
+        res["Sensor Bias Voltage"] = self._get_param("biasVoltage")
+        res["Pixels X"] = self._get_param("pixelsX")
+        res["Pixels Y"] = self._get_param("pixelsY")
+        res["Integration Time"] = self._get_param("integrationTime")
+        res["Sensor Temperature"] = self._get_param("sensorTemperature")
+        res["Gain Setting"] = self._get_param("gainSetting")
+        return res
+
+
+@KARABO_CLASSINFO("JungfrauCorrection", deviceVersion)
+class JungfrauCorrection(BaseCorrection):
+    _correction_flag_class = CorrectionFlags
+    _correction_field_names = (
+        ("offset", CorrectionFlags.OFFSET),
+        ("relGain", CorrectionFlags.REL_GAIN),
+    )
+    _kernel_runner_class = JungfrauGpuRunner
+    _calcat_friend_class = JungfrauCalcatFriend
+    _constant_enum_class = JungfrauConstants
+
+    @staticmethod
+    def expectedParameters(expected):
+        super(JungfrauCorrection, JungfrauCorrection).expectedParameters(expected)
+        (
+            OVERWRITE_ELEMENT(expected)
+            .key("dataFormat.memoryCells")
+            .setNewDefaultValue(1)
+            .commit(),
+            OVERWRITE_ELEMENT(expected)
+            .key("preview.selectionMode")
+            .setNewDefaultValue("frame")
+            .commit(),
+        )
+
+        JungfrauCalcatFriend.add_schema(expected, JungfrauCorrection._managed_keys)
+        add_correction_step_schema(
+            expected,
+            JungfrauCorrection._managed_keys,
+            JungfrauCorrection._correction_field_names,
+        )
+
+        # mandatory: manager needs this in schema
+        (
+            VECTOR_STRING_ELEMENT(expected)
+            .key("managedKeys")
+            .assignmentOptional()
+            .defaultValue(list(JungfrauCorrection._managed_keys))
+            .commit()
+        )
+
+    @property
+    def input_data_shape(self):
+        # TODO: check up on this
+        return (
+            self._schema_cache["dataFormat.memoryCells"],
+            1,
+            self._schema_cache["dataFormat.pixelsX"],
+            self._schema_cache["dataFormat.pixelsY"],
+        )
+
+    def __init__(self, config):
+        super().__init__(config)
+        # TODO: gain mode as constant parameter and / or device configuration
+        self.gain_mode = JungfrauGainMode[config.get("constantParameters.gainMode")]
+        # TODO: rest of this
+
+    def process_data(
+        self,
+        data_hash,
+        metadata,
+        source,
+        train_id,
+        image_data,
+        cell_table,
+        do_generate_preview,
+    ):
+        # TODO
+        ...