Merge pull request #886 from HEXRD/preprocess-eiger-stream-v2

Add ability to preprocess eiger-stream-v2 files

Author: psavery

hexrd/hedm/preprocess/preprocessors.py

@@ -1,26 +1,27 @@
 import logging
+import os
+import time
+from typing import Any, Optional, Union, Sequence, cast
+
+import h5py
+from numpy.typing import NDArray
+from numpy import float32
 
+from hexrd.core import imageseries
 from hexrd.core.imageseries.baseclass import ImageSeries
 from hexrd.core.imageseries.omega import OmegaWedges
+from hexrd.core.imageseries.process import ProcessedImageSeries
 from hexrd.hedm.preprocess.profiles import (
     Eiger_Arguments,
     Dexelas_Arguments,
     HexrdPPScript_Arguments,
 )
-from hexrd.core import imageseries
-from hexrd.core.imageseries.process import ProcessedImageSeries
-import os
-import time
-from typing import Any, Optional, Union, Sequence, cast
-from numpy.typing import NDArray
-from numpy import float32
 
 logger = logging.getLogger(__name__)
 
 
 class PP_Base(object):
     PROCFMT: Optional[str] = None
-    RAWFMT: Optional[str] = None
 
     def __init__(
         self,
@@ -90,7 +91,6 @@ class PP_Eiger(PP_Base):
     """PP_Eiger"""
 
     PROCFMT = "frame-cache"
-    RAWFMT = "eiger-stream-v1"
 
     def __init__(
         self,
@@ -99,6 +99,8 @@ def __init__(
         panel_opts: list = [],
         frame_start: int = 0,
         style: str = "npz",
+        eiger_stream_v2_threshold: str = "threshold_1",
+        eiger_stream_v2_multiplier: float = 1.0,
     ) -> None:
         super().__init__(
             fname=fname,
@@ -107,19 +109,52 @@ def __init__(
             frame_start=frame_start,
             style=style,
         )
-        self.raw = imageseries.open(self.fname, format=self.RAWFMT)
+        self.raw = imageseries.open(self.fname, format=self.rawfmt)
         self.use_frame_list = self.nframes != len(self.raw)
+        self.eiger_stream_v2_threshold = eiger_stream_v2_threshold
+        self.eiger_stream_v2_multiplier = eiger_stream_v2_multiplier
         logger.info(
             f"On Init:\n\t{self.fname}, {self.nframes} frames, "
             f"{self.omwedges.nframes} omw, {len(self.raw)} total"
         )
 
+        if self.is_eiger_stream_v2:
+            logger.info(
+                "Eiger stream v2 file format detected. Threshold setting is: "
+                + self.eiger_stream_v2_threshold
+            )
+            if self.eiger_stream_v2_threshold == "man_diff":
+                logger.info(
+                    f"Eiger stream v2 multiplier is: {self.eiger_stream_v2_multiplier}"
+                )
+
+            self.raw.set_option('threshold_setting', self.eiger_stream_v2_threshold)
+            self.raw.set_option('multiplier', self.eiger_stream_v2_multiplier)
+
+    @property
+    def rawfmt(self) -> str:
+        # Open the file and check if it is eiger-stream-v1 or eiger-stream-v2
+        # Only check the format once. Cache it.
+        if not hasattr(self, '_rawfmt'):
+            v2_str = 'CHESS_EIGER_STREAM_V2'
+            fmt = 'eiger-stream-v1'
+            with h5py.File(self.fname, 'r') as f:
+                if f.attrs.get('version', None) == v2_str:
+                    fmt = 'eiger-stream-v2'
+
+            self._raw_fmt = fmt
+
+        return self._raw_fmt
+
+    @property
+    def is_eiger_stream_v2(self) -> bool:
+        return self.rawfmt == 'eiger-stream-v2'
+
 
 class PP_Dexela(PP_Base):
     """PP_Dexela"""
 
     PROCFMT = "frame-cache"
-    RAWFMT = "hdf5"
 
     RAWPATH = "/imageseries"
     DARKPCTILE = 50
@@ -143,10 +178,12 @@ def __init__(
             style=style,
         )
 
+        self.rawfmt = "hdf5"
+
         self._panel_id = panel_id
         # TODO is this logic applicable also for Eiger ?
         if raw_format.lower() == "hdf5":
-            self.raw = imageseries.open(self.fname, self.RAWFMT, path=self.RAWPATH)
+            self.raw = imageseries.open(self.fname, self.rawfmt, path=self.RAWPATH)
         else:
             self.raw = imageseries.open(self.fname, raw_format.lower())
         self._dark = dark
@@ -195,6 +232,8 @@ def preprocess(args: HexrdPPScript_Arguments) -> None:
             omw=omw,
             frame_start=args.start_frame,
             style=args.style,
+            eiger_stream_v2_threshold=args.eiger_stream_v2_threshold,
+            eiger_stream_v2_multiplier=args.eiger_stream_v2_multiplier,
         )
         ppe.save_processed(args.output, args.threshold)
     elif type(args) == Dexelas_Arguments:

hexrd/hedm/preprocess/profiles.py

@@ -128,10 +128,14 @@ class Eiger_Arguments(Chess_Arguments):
     profile_name = "eiger"
     # fields
     absolute_path: Optional[str] = None
+    eiger_stream_v2_threshold: str = 'threshold_1'
+    eiger_stream_v2_multiplier: float = 1.0
 
     help_messages = {
         **Chess_Arguments.help_messages,
         "absolute_path": "absolute path to image file",
+        "eiger_stream_v2_threshold": "Threshold to use for eiger-stream-v2 input file. Options are 'threshold_1', 'threshold_2', or 'man_diff', which is defined as `threshold_1 - multiplier * threshold_2`",
+        "eiger_stream_v2_multiplier": "Multiplier to use for threshold setting 'man_diff'. Unused otherwise.",
     }
 
     short_switches = {

tests/test_preprocess.py

@@ -1,10 +1,14 @@
+from pathlib import Path
+import tempfile
+
+import pytest
+
 from hexrd.hedm.preprocess.profiles import (
     Eiger_Arguments,
     Dexelas_Arguments,
     HexrdPPScript_Arguments,
 )
-import pytest
-from pathlib import Path
+from hexrd.hedm.preprocess.preprocessors import preprocess
 
 
 # FIXME: this is not the appropriate file to test the preprocess scrips since
@@ -20,6 +24,11 @@ def ceria_examples_path(eiger_examples_path: Path) -> Path:
     return eiger_examples_path / 'first_ceria' / 'ff_000_data_000001.h5'
 
 
+@pytest.fixture
+def eiger_stream_v2_examples_path(eiger_examples_path: Path) -> Path:
+    return eiger_examples_path / 'eiger_stream_v2/eiger_stream_v2_test_dataset.h5'
+
+
 # test load/safe with defaults
 def test_save_load_defaults_eiger():
     eargs = Eiger_Arguments()
@@ -41,14 +50,38 @@ def test_save_load_eiger(ceria_examples_path):
     buffer = eargs.dump_config()
     args = HexrdPPScript_Arguments.load_from_config(buffer)
     assert eargs == args
-    with pytest.raises(
-        RuntimeError
-    ):  # required argument 'absolute_path' is missing
+    with pytest.raises(RuntimeError):  # required argument 'absolute_path' is missing
         args.validate_arguments()
     args.absolute_path = str(ceria_examples_path)
     args.validate_arguments()
 
 
+def test_save_load_eiger_stream_v2(eiger_stream_v2_examples_path: Path):
+    eargs = Eiger_Arguments()
+    buffer = eargs.dump_config()
+    args = HexrdPPScript_Arguments.load_from_config(buffer)
+    assert eargs == args
+    with pytest.raises(RuntimeError):  # required argument 'absolute_path' is missing
+        args.validate_arguments()
+    args.absolute_path = str(eiger_stream_v2_examples_path)
+    args.eiger_stream_v2_threshold = 'man_diff'
+    args.eiger_stream_v2_multiplier = 0.75
+    args.num_frames = 1
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        # Write the output file in the temporary directory
+        output_file = Path(tmpdirname) / 'test'
+        args.output = output_file
+
+        args.validate_arguments()
+
+        # Write out the example file!
+        preprocess(args)
+
+        # Verify that there is an output file
+        assert len(list(Path(tmpdirname).glob(f"*.npz"))) > 0
+
+
 def test_save_load_dexelas():
     eargs = Dexelas_Arguments()
     eargs.base_dir = "/data"