Compatibility with cfelpyutils 2.x & read/write bad regions from .geom file

.github/dependabot/constraints.txt

@@ -1,4 +1,4 @@
-cfelpyutils==1.0.1
+cfelpyutils==2.0.6
 coverage<6.4
 cycler==0.11.0
 extra_data==1.10.0

extra_geom/base.py

@@ -1,8 +1,9 @@
 from functools import lru_cache
 from itertools import chain
+from warnings import warn
 
 import numpy as np
-from cfelpyutils.crystfel_utils import load_crystfel_geometry
+from cfelpyutils.geometry import load_crystfel_geometry
 
 from .crystfel_fmt import write_crystfel_geom
 from .snapped import GridGeometryFragment, SnappedGeometry
@@ -32,8 +33,8 @@ def from_panel_dict(cls, d):
         corner_pos = np.array([d['cnx']/res, d['cny']/res, d['coffset']])
         ss_vec = np.array([d['ssx'], d['ssy'], d['ssz']]) / res
         fs_vec = np.array([d['fsx'], d['fsy'], d['fsz']]) / res
-        ss_pixels = d['max_ss'] - d['min_ss'] + 1
-        fs_pixels = d['max_fs'] - d['min_fs'] + 1
+        ss_pixels = d['orig_max_ss'] - d['orig_min_ss'] + 1
+        fs_pixels = d['orig_max_fs'] - d['orig_min_fs'] + 1
         return cls(corner_pos, ss_vec, fs_vec, ss_pixels, fs_pixels)
 
     def corners(self):
@@ -252,15 +253,16 @@ def _cfel_panels_by_data_coord(cls, panels: dict):
             if len(ix_dims) > 1:
                 raise ValueError(f"Too many index dimensions for {pname}: {dims}")
 
-            min_ss = info['min_ss']
+            min_ss = info['orig_min_ss']
             if ix_dims:
                 # Geometry for 3D data, modules stacked along separate axis
                 modno = ix_dims[0]
             else:
                 # Geometry for 2D data, modules concatenated along slow-scan axis
                 modno, min_ss = divmod(min_ss, cls.expected_data_shape[1])
 
-            res[(modno, min_ss, info['min_fs'])] = info
+            info['panel_name'] = pname
+            res[(modno, min_ss, info['orig_min_fs'])] = info
 
         return res
 
@@ -270,21 +272,25 @@ def from_crystfel_geom(cls, filename):
 
         Returns a new geometry object.
         """
-        geom_dict = load_crystfel_geometry(filename)
-        panels_by_data_coord = cls._cfel_panels_by_data_coord(geom_dict['panels'])
+        cfel_geom = load_crystfel_geometry(filename)
+        panels_by_data_coord = cls._cfel_panels_by_data_coord(
+            cfel_geom.detector['panels']
+        )
         n_modules = cls.n_modules
         if n_modules == 0:
             # Detector type with varying number of modules (e.g. JUNGFRAU)
             n_modules = max(c[0] for c in panels_by_data_coord) + 1
 
         modules = []
+        panel_names_to_pNaM = {}
         for p in range(n_modules):
             tiles = []
             modules.append(tiles)
             for a in range(cls.n_tiles_per_module):
                 ss_slice, fs_slice = cls._tile_slice(a)
                 d = panels_by_data_coord[p, ss_slice.start, fs_slice.start]
                 tiles.append(GeometryFragment.from_panel_dict(d))
+                panel_names_to_pNaM[d['panel_name']] = f'p{p}a{a}'
 
         # Store some extra fields to write if we create another .geom file.
         # It's possible for these to have different values for different panels,
@@ -293,7 +299,32 @@ def from_crystfel_geom(cls, filename):
         cfel_md_keys = ('data', 'mask', 'adu_per_eV', 'clen')
         d1 = panels_by_data_coord[0, 0, 0]
         metadata = {'crystfel': {k: d1.get(k) for k in cfel_md_keys}}
-        # TODO: photon_energy (not returned with cfelpyutils 1.0)
+        metadata['crystfel']['photon_energy'] = cfel_geom.beam['photon_energy']
+
+        # Normalise description of bad regions, so we can output it correctly.
+        # - Change panel names to uniform pNaM (panel N asic M) format
+        # - If the file has a 2D layout (modules arranged along the slow-scan
+        #   axis), convert slow-scan coordinates to a 3D layout.
+        file_geom_is_2d = not any(isinstance(d, int) for d in d1['dim_structure'])
+        adjusted_bad_regions = {}
+        for bad_name, bad_d in cfel_geom.detector['bad'].items():
+            panel_name = bad_d['panel']
+            if panel_name:
+                try:
+                    bad_d['panel'] = panel_names_to_pNaM[panel_name]
+                except KeyError:
+                    warn("Discarding {bad_name}, no such panel {panel_name!r}")
+                    continue
+            if bad_d['is_fsss']:
+                if not panel_name:
+                    warn("Discarding {bad_name}, ss/fs region without panel name")
+                    continue
+                if file_geom_is_2d:
+                    bad_d['min_ss'] %= cls.expected_data_shape[1]
+                    bad_d['max_ss'] %= cls.expected_data_shape[1]
+            adjusted_bad_regions[bad_name] = bad_d
+
+        metadata['crystfel']['bad'] = adjusted_bad_regions
 
         return cls(modules, filename=filename, metadata=metadata)
 
@@ -350,6 +381,7 @@ def write_crystfel_geom(self, filename, *,
 
         write_crystfel_geom(
             self, filename, data_path=data_path, mask_path=mask_path, dims=dims,
+            bad_regions=cfelmeta.get('bad', {}),
             nquads=nquads, adu_per_ev=adu_per_ev, clen=clen,
             photon_energy=photon_energy,
         )

extra_geom/crystfel_fmt.py

@@ -1,5 +1,6 @@
 """Write geometry in CrystFEL format.
 """
+import re
 from itertools import product
 
 import numpy as np
@@ -77,11 +78,12 @@ def frag_to_crystfel(fragment, p, a, ss_slice, fs_slice, dims, pixel_size):
         coffset=fragment.corner_pos[2],
     )
 
-def write_crystfel_geom(self, filename, *,
-                        data_path='/entry_1/instrument_1/detector_1/data',
-                        mask_path=None, dims=('frame', 'modno', 'ss', 'fs'),
-                        nquads=4, adu_per_ev=None, clen=None,
-                        photon_energy=None):
+def write_crystfel_geom(
+        self, filename, *,
+        data_path='/entry_1/instrument_1/detector_1/data',
+        mask_path=None, dims=('frame', 'modno', 'ss', 'fs'), bad_regions={},
+        nquads=4, adu_per_ev=None, clen=None, photon_energy=None,
+):
     """Write this geometry to a CrystFEL format (.geom) geometry file.
     """
     from . import __version__
@@ -150,11 +152,45 @@ def write_crystfel_geom(self, filename, *,
             clen=clen_str,
             photon_energy=photon_energy_str
         ))
+        f.write(format_bad_regions(
+            bad_regions,
+            mod_ss_pixels=self.expected_data_shape[1],
+            layout_2d=('modno' not in dims)
+        ))
         rigid_groups = get_rigid_groups(self, nquads=nquads)
         f.write(rigid_groups)
         for chunk in panel_chunks:
             f.write(chunk)
 
+def format_bad_regions(bad_regions: dict, mod_ss_pixels: int, layout_2d=False):
+    lines = []
+    for name, d in bad_regions.items():
+        if d['is_fsss']:
+            if layout_2d:
+                modno = int(re.match("p(\d+)a\d+", d['panel'])[1])
+                mod_offset = modno * mod_ss_pixels
+                min_ss = d['min_ss'] + mod_offset
+                max_ss = d['max_ss'] + mod_offset
+            else:
+                min_ss, max_ss = d['min_ss'], d['max_ss']
+            lines += [
+                f"{name}/panel = {d['panel']}",
+                f"{name}/min_ss = {min_ss}",
+                f"{name}/max_ss = {max_ss}",
+                f"{name}/min_fs = {d['min_fs']}",
+                f"{name}/max_fs = {d['max_fs']}",
+                ""
+            ]
+        else:
+            lines += [
+                f"{name}/min_x = {d['min_x']}",
+                f"{name}/max_x = {d['max_x']}",
+                f"{name}/min_y = {d['min_y']}",
+                f"{name}/max_y = {d['max_y']}",
+                ""
+            ]
+    return "\n".join(lines)
+
 def get_rigid_groups(geom, nquads=4):
     """Create string for rigid groups definition."""
 

extra_geom/tests/test_agipd500k2g_geometry.py

@@ -5,7 +5,7 @@
 import pyFAI.detectors
 import pytest
 import xarray as xr
-from cfelpyutils.crystfel_utils import load_crystfel_geometry
+from cfelpyutils.geometry import load_crystfel_geometry
 from extra_data.stacking import stack_detector_data
 from matplotlib.axes import Axes
 
@@ -81,7 +81,7 @@ def test_write_read_crystfel_file(tmpdir):
     np.testing.assert_allclose(loaded.modules[0][0].fs_vec, geom.modules[0][0].fs_vec)
 
     # Load the geometry file with cfelpyutils and test the rigid groups
-    geom_dict = load_crystfel_geometry(path)
+    geom_dict = load_crystfel_geometry(path).detector
     quad_gr0 = [  # quadrant: p0a0 ... p7a7
         'p{}a{}'.format(p, a) for p, a in product(range(8), range(8))
     ]
@@ -90,10 +90,10 @@ def test_write_read_crystfel_file(tmpdir):
     assert geom_dict['rigid_groups']['q0'] == quad_gr0
     assert geom_dict['panels']['p0a0']['res'] == 5000  # 5000 pixels/metre
     p3a7 = geom_dict['panels']['p3a7']
-    assert p3a7['min_ss'] == 448
-    assert p3a7['max_ss'] == 511
-    assert p3a7['min_fs'] == 0
-    assert p3a7['max_fs'] == 127
+    assert p3a7['orig_min_ss'] == 448
+    assert p3a7['orig_max_ss'] == 511
+    assert p3a7['orig_min_fs'] == 0
+    assert p3a7['orig_max_fs'] == 127
 
 
 def test_write_read_crystfel_file_2d(tmpdir):
@@ -109,14 +109,14 @@ def test_write_read_crystfel_file_2d(tmpdir):
     np.testing.assert_allclose(loaded.modules[0][0].fs_vec, geom.modules[0][0].fs_vec)
 
     # Load the geometry file with cfelpyutils and check some values
-    geom_dict = load_crystfel_geometry(path)
+    geom_dict = load_crystfel_geometry(path).detector
 
     p3a7 = geom_dict['panels']['p3a7']
     assert p3a7['dim_structure'] == ['%', 'ss', 'fs']
-    assert p3a7['min_ss'] == (3 * 512) + 448
-    assert p3a7['max_ss'] == (3 * 512) + 511
-    assert p3a7['min_fs'] == 0
-    assert p3a7['max_fs'] == 127
+    assert p3a7['orig_min_ss'] == (3 * 512) + 448
+    assert p3a7['orig_max_ss'] == (3 * 512) + 511
+    assert p3a7['orig_min_fs'] == 0
+    assert p3a7['orig_max_fs'] == 127
 
 
 def test_inspect():

extra_geom/tests/test_agipd_geometry.py

@@ -4,7 +4,7 @@
 import numpy as np
 import pyFAI.detectors
 import pytest
-from cfelpyutils.crystfel_utils import load_crystfel_geometry
+from cfelpyutils.geometry import load_crystfel_geometry
 from matplotlib.axes import Axes
 
 from extra_geom import AGIPD_1MGeometry, agipd_asic_seams
@@ -69,21 +69,43 @@ def test_assemble_symmetric():
         geom.position_modules_symmetric(stacked_data, out=img[:-1, :-1])
 
 
+bad_xy = {
+    'is_fsss': False,
+    'min_x': -20., 'max_x': 20., 'min_y': -100., 'max_y': 100,
+    'min_fs': 0, 'max_fs': 0, 'min_ss': 0, 'max_ss': 0, 'panel': '',
+}
+bad_fsss = {
+    'is_fsss': True,
+    'min_x': None, 'max_x': None, 'min_y': None, 'max_y': None,
+    'min_fs': 10, 'max_fs': 100, 'min_ss': 450, 'max_ss': 500, 'panel': 'p3a7',
+}
+def assert_bad_region_like(actual, expected):
+    actual = {k: None if isinstance(v, float) and np.isnan(v) else v
+              for (k, v) in actual.items()}
+    assert actual == expected
+
+
 def test_write_read_crystfel_file(tmpdir):
     geom = AGIPD_1MGeometry.from_quad_positions(
         quad_pos=[(-525, 625), (-550, -10), (520, -160), (542.5, 475)]
     )
     # Use the z dimension (coffset in .geom)
     geom = geom.offset((0, 0, 0.001), modules=np.s_[8:12])
+
+    # Add some bad regions in CrystFEL file
+    geom.metadata['crystfel'] = {'bad': {'bad_xy': bad_xy, 'bad_fsss': bad_fsss}}
+
     path = str(tmpdir / 'test.geom')
     geom.write_crystfel_geom(filename=path, photon_energy=9000,
                              adu_per_ev=0.0075, clen=0.2)
 
     loaded = AGIPD_1MGeometry.from_crystfel_geom(path)
     assert_geom_close(loaded, geom)
+    assert_bad_region_like(loaded.metadata['crystfel']['bad']['bad_xy'], bad_xy)
+    assert_bad_region_like(loaded.metadata['crystfel']['bad']['bad_fsss'], bad_fsss)
 
     # Load the geometry file with cfelpyutils and test the rigid groups
-    geom_dict = load_crystfel_geometry(path)
+    geom_dict = load_crystfel_geometry(path).detector
     quad_gr0 = [  # 1st quadrant: p0a0 ... p3a7
         'p{}a{}'.format(p, a) for p, a in product(range(4), range(8))
     ]
@@ -92,35 +114,45 @@ def test_write_read_crystfel_file(tmpdir):
     assert geom_dict['rigid_groups']['q0'] == quad_gr0
     assert geom_dict['panels']['p0a0']['res'] == 5000  # 5000 pixels/metre
     p3a7 = geom_dict['panels']['p3a7']
-    assert p3a7['min_ss'] == 448
-    assert p3a7['max_ss'] == 511
-    assert p3a7['min_fs'] == 0
-    assert p3a7['max_fs'] == 127
+    assert p3a7['orig_min_ss'] == 448
+    assert p3a7['orig_max_ss'] == 511
+    assert p3a7['orig_min_fs'] == 0
+    assert p3a7['orig_max_fs'] == 127
+
+    print(geom_dict['bad'])
+    assert_bad_region_like(geom_dict['bad']['bad_xy'], bad_xy)
+    assert_bad_region_like(geom_dict['bad']['bad_fsss'], bad_fsss)
 
 
 def test_write_read_crystfel_file_2d(tmpdir):
     geom = AGIPD_1MGeometry.from_quad_positions(
         quad_pos=[(-525, 625), (-550, -10), (520, -160), (542.5, 475)]
     )
     path = str(tmpdir / 'test.geom')
+
+    # Add some bad regions in CrystFEL file
+    geom.metadata['crystfel'] = {'bad': {'bad_xy': bad_xy, 'bad_fsss': bad_fsss}}
+
     geom.write_crystfel_geom(filename=path, dims=('frame', 'ss', 'fs'),
                              adu_per_ev=0.0075, clen=0.2)
 
     loaded = AGIPD_1MGeometry.from_crystfel_geom(path)
-    np.testing.assert_allclose(
-        loaded.modules[0][0].corner_pos, geom.modules[0][0].corner_pos
-    )
-    np.testing.assert_allclose(loaded.modules[0][0].fs_vec, geom.modules[0][0].fs_vec)
+    assert_geom_close(loaded, geom)
+    assert_bad_region_like(loaded.metadata['crystfel']['bad']['bad_xy'], bad_xy)
+    assert_bad_region_like(loaded.metadata['crystfel']['bad']['bad_fsss'], bad_fsss)
 
     # Load the geometry file with cfelpyutils and check some values
-    geom_dict = load_crystfel_geometry(path)
+    geom_dict = load_crystfel_geometry(path).detector
 
     p3a7 = geom_dict['panels']['p3a7']
     assert p3a7['dim_structure'] == ['%', 'ss', 'fs']
-    assert p3a7['min_ss'] == (3 * 512) + 448
-    assert p3a7['max_ss'] == (3 * 512) + 511
-    assert p3a7['min_fs'] == 0
-    assert p3a7['max_fs'] == 127
+    assert p3a7['orig_min_ss'] == (3 * 512) + 448
+    assert p3a7['orig_max_ss'] == (3 * 512) + 511
+    assert p3a7['orig_min_fs'] == 0
+    assert p3a7['orig_max_fs'] == 127
+
+    assert geom_dict['bad']['bad_fsss']['min_ss'] == (3 * 512) + 450
+    assert geom_dict['bad']['bad_fsss']['max_ss'] == (3 * 512) + 500
 
 
 def test_quad_positions():

extra_geom/tests/test_epix_geometry.py

@@ -1,7 +1,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import pytest
-from cfelpyutils.crystfel_utils import load_crystfel_geometry
+from cfelpyutils.geometry import load_crystfel_geometry
 
 import extra_geom
 from extra_geom import Epix10KGeometry, Epix100Geometry
@@ -132,14 +132,14 @@ def test_write_read_crystfel_file(args, tmpdir):
     assert_geom_close(loaded, epix)
 
     # Load the geometry file with cfelpyutils and test the rigid groups
-    geom_dict = load_crystfel_geometry(path)
+    geom_dict = load_crystfel_geometry(path).detector
     assert geom_dict['panels']['p0a0']['res'] == 1 / pxsz
     assert len(geom_dict['panels']) == 4
     p0a0 = geom_dict['panels']['p0a0']
-    assert p0a0['max_ss'] == nrow - 1
-    assert p0a0['min_ss'] == 0
-    assert p0a0['max_fs'] == ncol - 1
-    assert p0a0['min_fs'] == 0
+    assert p0a0['orig_max_ss'] == nrow - 1
+    assert p0a0['orig_min_ss'] == 0
+    assert p0a0['orig_max_fs'] == ncol - 1
+    assert p0a0['orig_min_fs'] == 0
 
 
 @pytest.mark.parametrize('args', ['epix100'], indirect=True)