Add writing geometry in CrystFEL format

src/geomtools/__init__.py

@@ -8,6 +8,6 @@ from .sfx import (
     rmsd_per_group, read_crystfel_cell,
 )
 from .detector import (
-    read_crystfel_geom, get_pixel_positions, assemble_data,
-    plot_detector_layout, plot_data_on_detector
+    read_crystfel_geom, write_crystfel_geom, get_pixel_positions,
+    assemble_data, plot_detector_layout, plot_data_on_detector
 )

src/geomtools/detector/__init__.py

 # flake8: noqa E401
-from .geom import (
-    read_crystfel_geom, get_pixel_positions, assemble_data
-)
+from .crystfel_frm import read_crystfel_geom, write_crystfel_geom
+from .geom import get_pixel_positions, assemble_data
 from .draw import plot_detector_layout, plot_data_on_detector

src/geomtools/detector/crystfel_frm.py

+import numpy as np
+import pandas as pd
+
+from cfelpyutils.geometry import load_crystfel_geometry
+from natsort import natsorted
+
+
+HEADER_TEMPLATE = """\
+; {detector} geometry file written by geomtools-optimiser-sfx
+; You may need to edit this file to add:
+; - data and mask locations in the file
+; - mask_good & mask_bad values to interpret the mask
+; - adu_per_eV & photon_energy
+; - clen (detector distance)
+;
+; See: http://www.desy.de/~twhite/crystfel/manual-crystfel_geometry.html
+
+{paths}
+dim{frame_dim} = %
+res = {resolution} ; pixels per metre"
+
+; Beam energy in eV
+photon_energy = {photon_energy}
+
+; Camera length, aka detector distance
+clen = {clen}
+
+; Analogue Digital Units per eV
+adu_per_eV = {adu_per_ev}
+
+"""
+
+
+def _format_vec(vec):
+    return ' '.join([
+        f'{np.format_float_positional(value, sign=True)}{axis}'
+        for value, axis in zip(vec, 'xyz')
+        if not (axis == 'z' and value == 0)
+    ])
+
+
+def write_crystfel_geom(file, panels, beam, rigid_groups):
+    """Writes geometry in CrystFEL format in file."""
+    # header
+    clen = panels.clen[0]
+    paths = {}
+    paths['data'] = (
+        panels.data[0] if panels.data[0]
+        else "/entry_1/data_1/data"
+    )
+    mask = panels['mask'][0]
+    if mask:
+        paths['mask'] = mask
+    else:
+        paths['mask'] = "/entry_1/instrument_1/detector_1/mask"
+    path_str = '\n'.join('{} = {} ;'.format(i, j) for i, j in paths.items())
+    header = HEADER_TEMPLATE.format(
+        detector="AGIPD1M", resolution=1/200e-6, frame_dim=0,
+        clen=clen, photon_energy=beam["photon_energy"],
+        adu_per_ev=1, paths=path_str,
+    )
+
+    # rigid groups
+    rigid_groups_text = []
+    for group in rigid_groups:
+        items = panels.reset_index().groupby(group).agg(
+            {'panel': lambda x: ','.join(natsorted(x))})
+        items = items.reindex(natsorted(items.index))
+
+        rigid_groups_text += [
+            f"rigid_group_{name} = {attr.panel}"
+            for name, attr in items.iterrows()
+        ]
+        rigid_groups_text += [
+            f"rigid_group_collection_{group} = " + ",".join(items.index),
+            "",
+        ]
+    rigid_groups_text.append("")
+
+    # panels
+    panels_text = []
+    for name, row in panels.iterrows():
+        panels_text += [
+            name + "/dim1 = " + str(row.modno),
+            name + "/dim2 = ss",
+            name + "/dim3 = fs",
+            name + "/min_fs = " + str(row.min_fs),
+            name + "/max_fs = " + str(row.max_fs),
+            name + "/min_ss = " + str(row.min_ss),
+            name + "/max_ss = " + str(row.max_ss),
+            name + "/fs = " + _format_vec(np.float32([row.fsx, row.fsy])),
+            name + "/ss = " + _format_vec(np.float32([row.ssx, row.ssy])),
+            name + "/corner_x = " + np.format_float_positional(
+                np.float32(row.cnx)),
+            name + "/corner_y = " + np.format_float_positional(
+                np.float32(row.cny)),
+            name + "/coffset = " + np.format_float_positional(
+                np.float32(row.coffset)),
+            "",
+        ]
+
+    # write to file
+    file.write(header)
+    file.write("\n".join(rigid_groups_text))
+    file.write("\n".join(panels_text))
+
+
+def append_rigid_group(panels, geom, group_name):
+    """Appends the rigid group name to the panels."""
+    groups = dict(
+        (rg, geom.detector['rigid_groups'][rg])
+        for rg in geom.detector['rigid_group_collections'][group_name]
+    )
+
+    ix, pa = [], []
+    for m, p in groups.items():
+        ix += [m] * len(p)
+        pa += p
+
+    panels = panels.join(
+        pd.DataFrame(data={group_name: ix, 'panel': pa}).set_index('panel'))
+
+    return panels
+
+
+def read_crystfel_geom(filename, indexes=dict()):
+    """Reads Crystfel geometry file."""
+    geom = load_crystfel_geometry(filename)
+
+    panels = pd.DataFrame(
+        geom.detector['panels'].values(),
+        index=geom.detector['panels'].keys()
+    )
+    panels.index.name = 'panel'
+
+    for group_name in geom.detector['rigid_group_collections']:
+        panels = append_rigid_group(panels, geom, group_name)
+
+    for column, dimno in indexes.items():
+        panels[column] = panels.dim_structure.apply(lambda x: x[dimno])
+
+    return panels, geom.beam

src/geomtools/detector/geom.py

 import numpy as np
-import pandas as pd
-
-from cfelpyutils.geometry import load_crystfel_geometry
-
-
-def append_rigid_group(panels, geom, group_name):
-    """Appends the rigid group name to the panels."""
-    groups = dict(
-        (rg, geom.detector['rigid_groups'][rg])
-        for rg in geom.detector['rigid_group_collections'][group_name]
-    )
-
-    ix, pa = [], []
-    for m, p in groups.items():
-        ix += [m] * len(p)
-        pa += p
-
-    panels = panels.join(
-        pd.DataFrame(data={group_name: ix, 'panel': pa}).set_index('panel'))
-
-    return panels
-
-
-def read_crystfel_geom(filename, indexes=dict()):
-    """Reads Crystfel geometry file."""
-    geom = load_crystfel_geometry(filename)
-
-    panels = pd.DataFrame(
-        geom.detector['panels'].values(),
-        index=geom.detector['panels'].keys()
-    )
-    panels.index.name = 'panel'
-
-    for group_name in geom.detector['rigid_group_collections']:
-        panels = append_rigid_group(panels, geom, group_name)
-
-    for column, dimno in indexes.items():
-        panels[column] = panels.dim_structure.apply(lambda x: x[dimno])
-
-    return panels, geom.beam
 
 
 def get_pixel_positions(panels):