Get rid of dataFormat, use shape of data or look at geometry

src/calng/CrystfelRunner.py

@@ -111,28 +111,8 @@ class CrystfelRunner(PythonDevice):
             .allowedStates(State.STARTED)
             .commit(),
 
-            # TODO: get these from geometry
-            NODE_ELEMENT(expected)
-            .key("dataFormat")
-            .commit(),
-
-            UINT32_ELEMENT(expected)
-            .key("dataFormat.pixelsSS")
-            .assignmentMandatory()
-            .commit(),
-
-            UINT32_ELEMENT(expected)
-            .key("dataFormat.pixelsFS")
-            .assignmentMandatory()
-            .commit(),
-
-            UINT32_ELEMENT(expected)
-            .key("dataFormat.numModules")
-            .assignmentMandatory()
-            .commit(),
-
             STRING_ELEMENT(expected)
-            .key("dataFormat.dataKey")
+            .key("imageDataKey")
             .assignmentMandatory()
             .commit(),
 
@@ -357,32 +337,7 @@ class CrystfelRunner(PythonDevice):
         self.frameQueue_reset()
         self.KARABO_ON_INPUT("input", self.input_handler)
 
-        # note: loading with cfelpyutils.geometry does not tell us peak_list value
-        # extra note: we need to avoid taking ';' and ' ' into the path!
-        # (Hence the large regexp on the right-hand side)
-        self._geom_data_path = None
-        self._geom_peak_path = None
-        with open(self.get("crystfelArgs.geometryPath"), "rt") as fd:
-            for line in fd:
-                if (match := re.match(r"data\s*=\s*([a-zA-Z0-9/\._]*)", line)):
-                    self._geom_data_path = match.group(1)
-                elif (match := re.match(r"peak_list\s*=\s*([a-zA-Z0-9/\._]*)", line)):
-                    self._geom_peak_path = match.group(1)
-                if None not in (self._geom_data_path, self._geom_peak_path):
-                    break
-        if self.get("crystfelArgs.doPeakfinding"):
-            if self._geom_data_path is None:
-                raise Exception(
-                    "We want to do peakfinding, but the geometry doesn't specify"
-                    "image data path"
-                )
-        else:
-            if self._geom_peak_path is None:
-                raise Exception(
-                    "We want to provide indexamajig with peaks, but the geometry "
-                    "doesn't specify peak list path"
-                )
-
+        self._get_parameters_from_geometry()
         self._serve_to_crystfel()
         self.updateState(State.STOPPED)
         self.startCrystfel()
@@ -439,66 +394,51 @@ class CrystfelRunner(PythonDevice):
 
         if self.unsafe_get("crystfelArgs.doPeakfinding"):
             # TODO: maybe support passing data even when there are already peaks
-            key = self.get("dataFormat.dataKey")
-            images = data_hash.get(key)
-            num_frames = images.shape[0]
-            num_modules = images.shape[1]
             # TODO: apply mask to data
-            for frame_index in range(num_frames):
+            # data should be frame, module, ss, fs
+            for frame in data_hash[self.unsafe_get("imageDataKey")]:
                 payload = {
-                    #"timestamp": time.time(),  # TODO: decide on some timestamp to pass
                     # TODO: send image data if we want to
-                    #"/entry_1/data_1/data": frame.reshape(16 * 256, 256),
-                    self._geom_data_path: images[frame_index].reshape(
-                        self.unsafe_get("dataFormat.numModules") * self.unsafe_get("dataFormat.pixelsSS"),
-                        self.unsafe_get("dataFormat.pixelsFS")
+                    self._geom_data_path: frame.reshape(
+                        frame.shape[0] * frame.shape[1], *frame.shape[2:]
                     ),
-                    "native_data_shape": (
-                        self.unsafe_get("dataFormat.pixelsSS"),
-                        self.unsafe_get("dataFormat.pixelsFS"),
-                    ),  # TODO: figure out what this should be
                 }
                 try:
                     self._crystfel_queue.put_nowait(msgpack.packb([payload]))
                 except queue.Full:
                     ...
         else:
-            peak_count = data_hash.get("peakfinding.numPeaks")  # shape: (frame, module)
+            peak_count = data_hash["peakfinding.numPeaks"]  # shape: (frame, module)
             # TODO: probably rename to ss/ff all around
-            peak_ss = data_hash.get("peakfinding.peakX") # shape: (frame, module, max_peaks)
-            peak_fs = data_hash.get("peakfinding.peakY")
-            peak_intensity = data_hash.get("peakfinding.peakIntensity")
-            num_frames = peak_count.shape[0]
+            peak_ss = data_hash["peakfinding.peakX"] # shape: (frame, module, max_peaks)
+            peak_fs = data_hash["peakfinding.peakY"]
+            peak_intensity = data_hash["peakfinding.peakIntensity"]
             num_modules = peak_count.shape[1]
             # slabbifying the ss coordinates
             peak_ss = peak_ss + (
-                MODULE_SIZE_SS * np.arange(num_modules)[np.newaxis, :, np.newaxis]
+                self._module_ss_size * np.arange(num_modules)[np.newaxis, :, np.newaxis]
             )
-            assert all(thing.shape[0] == num_frames for thing in (peak_ss, peak_fs, peak_intensity))
-            for frame_index in range(num_frames):
+            assert all(
+                thing.shape[0] == peak_count.shape[0]
+                for thing in (peak_ss, peak_fs, peak_intensity)
+            )
+            for frame_count, frame_ss, frame_fs, frame_intensity in zip(
+                peak_count, peak_ss, peak_fs, peak_intensity
+            ):
                 ss, fs, intensity = [], [], []
-                # TODO: configurable
-                if np.sum(peak_count[frame_index]) < 10:
-                    continue
-                for module_index in range(num_modules):
-                    peak_count_local = peak_count[frame_index][module_index]
-                    ss.append(peak_ss[frame_index][module_index][:peak_count_local])
-                    fs.append(peak_fs[frame_index][module_index][:peak_count_local])
-                    intensity.append(peak_intensity[frame_index][module_index][:peak_count_local])
+                for mod_index, (mod_count, mod_ss, mod_fs, mod_intensity) in enumerate(
+                    zip(frame_count, frame_ss, frame_fs, frame_intensity)
+                ):
+                    ss.extend(mod_ss[:mod_count])  # TODO: maybe multiply by mod_index * ss size
+                    fs.extend(mod_fs[:mod_count])
+                    intensity.extend(mod_intensity[:mod_count])
                 payload = {
-                    #"timestamp": time.time(),  # TODO: decide on some timestamp to pass
-                    # TODO: send image data if we want to
-                    #"/entry_1/data_1/data": frame.reshape(16 * 256, 256),
                     self._geom_peak_path: {
                         # TODO: slabbify
-                        "ss": list(itertools.chain.from_iterable(ss)),
-                        "fs": list(itertools.chain.from_iterable(fs)),
-                        "intensity": list(itertools.chain.from_iterable(intensity)),
+                        "ss": list(map(float(ss))),
+                        "fs": list(map(float(fs))),
+                        "intensity": list(map(float(intensity))),
                     },
-                    "native_data_shape": (
-                        MODULE_SIZE_SS,
-                        MODULE_SIZE_FS,
-                    ),  # TODO: figure out what this should be
                 }
                 self._crystfel_queue.put_nowait(msgpack.packb([payload]))
         self.set(
@@ -506,6 +446,36 @@ class CrystfelRunner(PythonDevice):
             self._crystfel_queue.qsize() / self._crystfel_queue.maxsize * 100
         )
 
+    def _get_parameters_from_geometry(self):
+        # note: loading with cfelpyutils.geometry does not tell us peak_list value
+        # extra note: we need to avoid taking ';' and ' ' into the path!
+        # (Hence the large regexp on the right-hand side)
+        self._geom_data_path = None
+        self._geom_peak_path = None
+        self._module_ss_size = None
+        seen_max_ss = set()
+        with open(self.get("crystfelArgs.geometryPath"), "rt") as fd:
+            for line in fd:
+                if (match := re.match(r"data\s*=\s*([a-zA-Z0-9/\._]*)", line)):
+                    self._geom_data_path = match.group(1)
+                elif (match := re.match(r"peak_list\s*=\s*([a-zA-Z0-9/\._]*)", line)):
+                    self._geom_peak_path = match.group(1)
+                elif (match := re.match(r"p0a\d+/max_ss\s*=\s*(\d+)", line)):
+                    seen_max_ss.add(int(match.group(1)))
+        self._module_ss_size = max(seen_max_ss) + 1
+        if self.get("crystfelArgs.doPeakfinding"):
+            if self._geom_data_path is None:
+                raise Exception(
+                    "We want to do peakfinding, but the geometry doesn't specify"
+                    "image data path"
+                )
+        else:
+            if self._geom_peak_path is None:
+                raise Exception(
+                    "We want to provide indexamajig with peaks, but the geometry "
+                    "doesn't specify peak list path"
+                )
+
     def _serve_to_crystfel(self):
         def server():
             self.log.INFO("Starting ZMQ server to feed CrystFEL")