Add coordinates of reflex projection to the detector plane, add extraction of...

Add coordinates of reflex projection to the detector plane, add extraction of unit cell file from stream file

src/geomtools/__init__.py

 # flake8: noqa E401
 from .sfx import (
     parse_crystfel_streamfile, read_crystfel_streamfile, extract_geometry,
-    plot_center_shift, plot_cell_parameters, plot_peakogram, plot_powder,
-    ph_en_to_lambda, get_q_from_xyz, get_min_bragg_dist, spacing,
-    gauss2d_fit, ellipse, parse_xwiz_summary, get_peak_position,
-    plot_geoptimiser_errormap, avg_pixel_displacement, rmsd_per_group,
+    extract_cell, plot_center_shift, plot_cell_parameters, plot_peakogram,
+    plot_powder, ph_en_to_lambda, get_q_from_xyz, get_min_bragg_dist,
+    spacing, cell_volume, gauss2d_fit, ellipse, parse_xwiz_summary,
+    get_peak_position, plot_geoptimiser_errormap, avg_pixel_displacement,
+    rmsd_per_group, read_crystfel_cell,
 )
 from .detector import (
     read_crystfel_geom, get_pixel_positions, assemble_data,

src/geomtools/sfx/__init__.py

 # flake8: noqa E401
 from .crystfelio import (
     parse_crystfel_streamfile, read_crystfel_streamfile, extract_geometry,
+    extract_cell,
 )
 from .draw import (
     plot_center_shift, plot_cell_parameters, plot_peakogram, plot_powder,
     plot_geoptimiser_errormap,
 )
 from .lattice import (
-    spacing, ph_en_to_lambda, get_q_from_xyz, get_min_bragg_dist,
+    cell_volume, spacing, ph_en_to_lambda, get_q_from_xyz,
+    get_min_bragg_dist, read_crystfel_cell,
 )
 from .misc import (
     gauss2d_fit, ellipse, get_peak_position, avg_pixel_displacement,

src/geomtools/sfx/crystfelio.py

@@ -13,6 +13,8 @@ from cfelpyutils.crystfel_stream import (
 
 GEOM_START_MARKER = "----- Begin geometry file -----"
 GEOM_END_MARKER = "----- End geometry file -----"
+CELL_START_MARKER = "----- Begin unit cell -----"
+CELL_END_MARKER = "----- End unit cell -----"
 
 PEAK_COLUMN_MAP = {
     "fs/px": "fs", "ss/px": "ss", "(1/d)/nm^-1": "res",
@@ -49,7 +51,8 @@ def match_reflexes_to_peaks(reflexes, peaks, lattices,
     # 2. one reflex to one peak
     # 3. one peak to one reflex in the same crystal
     #    but to can also match refelex in other crystals
-    re = reflexes[['fs', 'ss', 'panel', 'frame', 'cryst', 'reflno']].copy()
+    re = reflexes[['fs', 'ss', 'panel', 'frame', 'cryst', 'reflno',
+                   'xa', 'ya']].copy()
     if group_name != 'panel':
         re = re.join(panels[[group_name]], on='panel')
     re = re.join(get_peak_position(re, panels))
@@ -83,7 +86,8 @@ def match_reflexes_to_peaks(reflexes, peaks, lattices,
 
 def parse_crystfel_streamfile(stream_filename, panels, connected_groups,
                               begin=0, end=None):
-    lattice_params = ['lattice_type', 'centering', 'unique_axis']
+    lattice_params = ['lattice_type', 'centering', 'unique_axis',
+                      'astar', 'bstar', 'cstar']
     lattice_arrays = {
         'Cell parameters/lengths': ['a', 'b', 'c'],
         'Cell parameters/angles': ['alpha', 'beta', 'gamma'],
@@ -143,12 +147,22 @@ def parse_crystfel_streamfile(stream_filename, panels, connected_groups,
                     1. / res_avg, clen_avg, lmd, cell)
 
                 lattices.append(la)
+
                 la_kwargs = dict((name, la[name]) for name in cell_columns)
 
                 re = crystal['reflections']
                 re['res'] = np.sqrt(spacing(
                     re.h.values, re.k.values, re.l.values, **la_kwargs))
 
+                M = np.array([
+                    crystal['astar'], crystal['bstar'], crystal['cstar']
+                ])
+                hkl = np.array([re.h.values, re.k.values, re.l.values])
+                u, v, w = M.T @ hkl
+                s = clen / (1e-9 / lmd + w) / 200e-6
+                re['xa'] = u * s
+                re['ya'] = v * s
+
                 re['frame'] = frame_ix
                 re['cryst'] = cryst_ix
                 reflexes.append(re)
@@ -176,29 +190,63 @@ def parse_crystfel_streamfile(stream_filename, panels, connected_groups,
     return frames, peaks, lattices, reflexes, match
 
 
-def extract_geometry(stream_filename):
-    """Extracts geometry from Crystfel stream file.
+def extract_chunk(stream_filename, start_marker, end_marker):
+    """Extracts text bracketed with markers from Crystfel stream file.
 
     Input
     -----
     stream_filename: str
         Crystfel stream file name
+    start_marker: str
+        String indicated begin of the text
+    end_marker: str
+        String indicated end of the text
 
     Returns
     -------
-        StringIO buffer with geometry file
+        StringIO buffer with text
     """
     with open(stream_filename, 'r') as f:
         line = f.readline()
-        while line and line.strip() != GEOM_START_MARKER:
+        while line and line.strip() != start_marker:
             if line.strip() == CHUNK_START_MARKER:
                 return
             line = f.readline()
-        geom = _buffer_to_line(f, GEOM_END_MARKER)
+        geom = _buffer_to_line(f, end_marker)
 
     return geom
 
 
+def extract_geometry(stream_filename):
+    """Extracts geometry from Crystfel stream file.
+
+    Input
+    -----
+    stream_filename: str
+        Crystfel stream file name
+
+    Returns
+    -------
+        StringIO buffer with geometry file
+    """
+    return extract_chunk(stream_filename, GEOM_START_MARKER, GEOM_END_MARKER)
+
+
+def extract_cell(stream_filename):
+    """Extracts unit cell from Crystfel stream file.
+
+    Input
+    -----
+    stream_filename: str
+        Crystfel stream file name
+
+    Returns
+    -------
+        StringIO buffer with crystfel cell file
+    """
+    return extract_chunk(stream_filename, CELL_START_MARKER, CELL_END_MARKER)
+
+
 def _split_file(filename, nproc=20, partsize=None, nbytes=None):
     """Split file on portions."""
     if nbytes is None: