Add first_modno input argument for components.JUNGFRAU

extra_data/components.py

@@ -1446,6 +1446,9 @@ class JUNGFRAU(MultimodDetectorBase):
     n_modules: int
       Number of detector modules in the experiment setup. Default is
       None, in which case it will be estimated from the available data.
+    first_modno: int
+      The module number in the source name for the first detector module.
+      e.g. FXE_XAD_JF500K/DET/JNGFR03:daqOutput should have first_modno = 3
     """
     # We appear to have a few different formats for source names:
     # SPB_IRDA_JNGFR/DET/MODULE_1:daqOutput  (e.g. in p 2566, r 61)
@@ -1461,14 +1464,23 @@ class JUNGFRAU(MultimodDetectorBase):
     module_shape = (512, 1024)
 
     def __init__(self, data: DataCollection, detector_name=None, modules=None,
-                 *, min_modules=1, n_modules=None):
+                 *, min_modules=1, n_modules=None, first_modno=1):
         super().__init__(data, detector_name, modules, min_modules=min_modules)
 
+        self.modno_to_source = {}
+        # Overwrite modno based on given starting module number and update
+        # source_to_modno and modno_to_source.
+        for source in self.source_to_modno.keys():
+            # JUNGFRAU modno is expected (e.g. extra_geom) to start with 1.
+            modno = int(self._source_re.search(source)['modno']) - first_modno + 1
+            self.source_to_modno[source] = modno
+            self.modno_to_source[modno] = source
+
         if n_modules is not None:
             self.n_modules = int(n_modules)
         else:
             # For JUNGFRAU modules are indexed from 1
-            self.n_modules = max(modno for (_, modno) in self._source_matches(
+            self.n_modules = max(modno - first_modno + 1 for (_, modno) in self._source_matches(
                 data, self.detector_name
             ))
 

extra_data/tests/conftest.py

@@ -141,6 +141,11 @@ def mock_jungfrau_run():
         make_examples.make_jungfrau_run(td)
         yield td
 
+@pytest.fixture(scope='session')
+def mock_fxe_jungfrau_run():
+    with TemporaryDirectory() as td:
+        make_examples.make_fxe_jungfrau_run(td)
+        yield td
 
 @pytest.fixture(scope='session')
 def mock_scs_run():

extra_data/tests/make_examples.py

@@ -404,6 +404,24 @@ def make_jungfrau_run(dir_path):
         JUNGFRAUPower('SPB_IRDA_JF4M/MDL/POWER'),
     ], ntrains=100, chunksize=1, format_version='1.0')
 
+def make_fxe_jungfrau_run(dir_path):
+    # Naming based on /gpfs/exfel/exp/FXE/202101/p002478/raw/
+    for modno in range(1, 3):
+        path = osp.join(dir_path, f'RAW-R0012-JNGFR{modno:02}-S00000.h5')
+        write_file(path, [
+            JUNGFRAUModule(f'FXE_XAD_JF1M/DET/JNGFR{modno:02}')
+        ], ntrains=100, chunksize=1, format_version='1.0')
+
+    path = osp.join(dir_path, f'RAW-R0052-JNGFR03-S00000.h5')
+    write_file(path, [
+        JUNGFRAUModule(f'FXE_XAD_JF500K/DET/JNGFR03')
+    ], ntrains=100, chunksize=1, format_version='1.0')
+
+    write_file(osp.join(dir_path, f'RAW-R0052-JNGFRCTRL00-S00000.h5'), [
+        JUNGFRAUControl('FXE_XAD_JF1M/DET/CONTROL'),
+        JUNGFRAUControl('FXE_XAD_JF500K/DET/CONTROL'),
+    ], ntrains=100, chunksize=1, format_version='1.0')
+
 def make_remi_run(dir_path):
     write_file(osp.join(dir_path, f'CORR-R0210-REMI01-S00000.h5'), [
         ReconstructedDLD6('SQS_REMI_DLD6/DET/TOP'),

extra_data/tests/test_components.py

@@ -235,6 +235,32 @@ def test_get_array_jungfrau(mock_jungfrau_run):
     assert arr.dtype == np.float32
 
 
+def test_jungfraus_first_modno(mock_jungfrau_run, mock_fxe_jungfrau_run):
+
+    # Test SPB_IRDA_JF4M component by setting the first_modno to the default value 1.
+    run = RunDirectory(mock_jungfrau_run)
+    jf = JUNGFRAU(run.select_trains(by_index[:2]), first_modno=1)
+    assert jf.detector_name == 'SPB_IRDA_JF4M'
+    assert jf.n_modules == 8
+
+    arr = jf.get_array('data.adc')
+    assert np.all(arr['module'] == list(range(1, 9)))
+
+    # Test FXE_XAD_JF500K component with and without setting first_modno to 3.
+    for first_modno, modno in zip([1, 3], [3, 1]):
+        run = RunDirectory(mock_fxe_jungfrau_run)
+        jf = JUNGFRAU(
+            run.select_trains(by_index[:2]),
+            detector_name='FXE_XAD_JF500K',
+            first_modno=first_modno,
+        )
+        assert jf.detector_name == 'FXE_XAD_JF500K'
+        assert jf.n_modules == modno
+
+        arr = jf.get_array('data.adc')
+        assert np.all(arr['module'] == [modno])
+
+
 def test_get_dask_array(mock_fxe_raw_run):
     run = RunDirectory(mock_fxe_raw_run)
     det = LPD1M(run)