stream2sol.py: Initial import

1 files changed, 267 insertions, 0 deletions

diff --git a/scripts/stream2sol.py b/scripts/stream2sol.py
new file mode 100755
index 00000000..d12ea2ba
--- /dev/null
+++ b/scripts/stream2sol.py
@@ -0,0 +1,267 @@
+#!/usr/bin/env python
+from io import StringIO
+import re
+from warnings import warn
+from copy import deepcopy
+
+BEGIN_GEOM = '----- Begin geometry file -----'
+END_GEOM = '----- End geometry file -----'
+BEGIN_CELL = '----- Begin unit cell -----'
+END_CELL = '----- End unit cell -----'
+BEGIN_CHUNK = '----- Begin chunk -----'
+END_CHUNK = '----- End chunk -----'
+BEGIN_CRYSTAL = '--- Begin crystal'
+END_CRYSTAL = '--- End crystal'
+BEGIN_PEAKS = 'Peaks from peak search'
+END_PEAKS = 'End of peak list'
+BEGIN_REFLECTIONS = 'Reflections measured after indexing'
+END_REFLECTIONS = 'End of reflections'
+HEAD = 'CrystFEL stream format {}.{}'.format(2, 3)
+GENERATOR = 'Generated by diffractem StreamParser'
+PEAK_COLUMNS = ['fs/px', 'ss/px', '(1/d)/nm^-1', 'Intensity', 'Panel']
+REFLECTION_COLUMNS = ['h', 'k', 'l', 'I', 'Sigma(I)', 'Peak', 'Background', 'fs/px', 'ss/px', 'Panel']
+ID_FIELDS = ['file', 'Event', 'serial']
+CRYSTAL_DATA_FIELS = ['astar', 'bstar', 'cstar', 'predict_refine/det_shift', 
+                      'profile_radius', 'diffraction_resolution_limit']
+
+args = None
+
+class Crystal:
+    
+    def __init__(self, line):
+        self.astar = (None, None, None)
+        self.bstar = (None, None, None)
+        self.cstar = (None, None, None)
+        self.det_shift = (None, None)
+        self.profile_radius = None
+        self.resolution = None
+        self.start_line = line
+
+    @property
+    def initialized(self):
+        return all([x is not None 
+                    for x in [*self.astar, *self.bstar, *self.cstar,
+                              *self.det_shift, self.profile_radius, 
+                              self.resolution]])
+    
+    def __str__(self):
+        if not self.initialized:
+            warn('Trying to get string from non-initialized crystal from line {}.'.format(self.start_line))
+            return None        
+        else:
+            cs = ' '.join(['{0[0]} {0[1]} {0[2]}'.format(vec) 
+                            for vec in [self.astar, self.bstar, self.cstar]])
+            cs += ' {0[0]} {0[1]}'.format(self.det_shift)
+            if args.include_pars: # this is a bit dirty but will become obsolete one day
+                cs += ' {0} {1}'.format(self.profile_radius, self.resolution)
+            return cs
+
+class Chunk:
+
+    def __init__(self, line):
+        self.file = None
+        self.Event = None
+        self.crystals = []
+        self.start_line = line
+        self.x_shift = 0
+        self.y_shift = 0
+        
+    @property
+    def n_cryst(self):
+        return len(self.crystals)
+    
+    @property 
+    def initialized(self):
+        return (self.file is not None) and (self.Event is not None)
+        
+    def add_crystal(self, crystal):
+        if (not crystal.initialized) or (crystal is None):
+            raise RuntimeError('Trying to add non-initialied crystal to chunk from line {}.'.format(self.start_line))
+        self.crystals.append(deepcopy(crystal))
+        # print(crystal)
+        
+    def __str__(self):
+        if not self.initialized:
+            warn('Trying to get string from non-initialized chunk from line {}.'.format(self.start_line))
+            return None
+        else:
+            # return '\n'.join([' '.join([self.file, *self.Event.split('//'), str(cryst)])
+            #             for ii, cryst in enumerate(self.crystals)])
+            # new-style (not working yet)
+            return '\n'.join([' '.join([self.file, self.Event, str(cryst)])
+                        for ii, cryst in enumerate(self.crystals)])
+
+
+def parse_stream(stream, sol=None, return_meta=True, 
+                 file_label='Image filename', event_label='Event',
+                 x_shift_label=None, y_shift_label=None, shift_factor=1):
+
+    curr_chunk = None
+    curr_cryst = None
+    geom = ''
+    cell = ''
+    command = ''
+    parsing_geom = False
+    parsing_cell = False
+    parsing_peaks = False
+    have_cell = False
+    have_geom = False
+    have_command = False
+    parsing_reflections = False
+    parse_vec = lambda l: tuple(float(k) for k in re.findall(r'[+-]?\d*\.\d*', l))
+    
+    with open(stream, 'r') as fh_in, (StringIO() if sol is None else open(sol,'w')) as fh_out:
+
+        for ln, l in enumerate(fh_in):
+            
+            if parsing_reflections:
+                if l.startswith(END_REFLECTIONS):
+                    parsing_reflections = False
+                else:
+                    # here, any reflection parsing would go
+                    pass
+                    
+            elif parsing_peaks:
+                if l.startswith(END_PEAKS):
+                    parsing_peaks = False
+                else:
+                    # here, any peak parsing would go
+                    pass
+
+            elif l.startswith(BEGIN_CHUNK):
+                curr_chunk = Chunk(ln)
+                
+            elif (curr_chunk is not None) and (curr_cryst is None):
+                # parsing chunks (= events = shots) _outside_ crystals
+                
+                if l.startswith(END_CHUNK):
+                    if not curr_chunk.initialized:
+                        raise RuntimeError('Incomplete chunk found before line ' + str(ln))
+                    if curr_chunk.n_cryst:
+                        fh_out.write(str(curr_chunk) + '\n')
+                        # print(str(curr_chunk))
+                    curr_chunk = None
+                    
+                elif l.startswith(file_label):
+                    curr_chunk.file = l.split(' ')[-1].strip()
+                    
+                elif l.startswith(event_label):
+                    curr_chunk.Event = l.split(' ')[-1].strip()
+
+                elif x_shift_label and l.startswith(x_shift_label):
+                    curr_chunk.x_shift = float(l.split(' ')[-1].strip())
+
+                elif y_shift_label and l.startswith(y_shift_label):
+                    curr_chunk.y_shift = float(l.split(' ')[-1].strip())
+
+                elif l.startswith(BEGIN_CRYSTAL):
+                    if not curr_chunk.initialized:
+                        raise RuntimeError('Crystal for incomplete chunk in ' + str(ln))                
+                    curr_cryst = Crystal(ln)
+                
+            elif curr_cryst is not None:
+                # parsing a (single) crystal
+                
+                if l.startswith(END_CRYSTAL):
+                    curr_chunk.add_crystal(curr_cryst)
+                    curr_cryst = None
+                        
+                elif l.startswith('astar'):
+                    curr_cryst.astar = parse_vec(l)
+                    
+                elif l.startswith('bstar'):
+                    curr_cryst.bstar = parse_vec(l)
+                    
+                elif l.startswith('cstar'):
+                    curr_cryst.cstar = parse_vec(l)
+                    
+                elif l.startswith('profile_radius'):
+                    curr_cryst.profile_radius = parse_vec(l)[0]
+                    
+                elif l.startswith('predict_refine/det_shift'):
+                    curr_cryst.det_shift = parse_vec(l)
+                    curr_cryst.det_shift = (curr_cryst.det_shift[0] + curr_chunk.x_shift,
+                                            curr_cryst.det_shift[1] + curr_chunk.y_shift)
+                    
+                elif l.startswith('diffraction_resolution_limit'):
+                    curr_cryst.resolution = parse_vec(l)[0]
+                    
+            elif  l.startswith(BEGIN_GEOM) and not have_geom:
+                parsing_geom = True
+                
+            elif parsing_geom:
+                if not l.startswith(END_GEOM):
+                    geom += l
+                else:
+                    parsing_geom = False
+                    have_geom = True
+                
+            elif l.startswith(BEGIN_CELL) and not have_cell:
+                parsing_cell = True
+                
+            elif parsing_cell:
+                if not l.startswith(END_CELL):
+                    cell += l
+                else:
+                    parsing_cell = False
+                    have_cell = True
+
+            elif ('indexamajig' in l) and not have_command:
+                command = l
+                have_command = True
+                                    
+            elif l.startswith(BEGIN_PEAKS):
+                parsing_peaks = True
+                
+            elif l.startswith(BEGIN_REFLECTIONS):
+                parsing_reflections = True                    
+                
+        if sol is None:
+            out = fh_out.getvalue()
+            if return_meta:
+                return out, (command, geom, cell)
+            else:
+                return out
+        
+        else:
+            if return_meta:
+                return command, geom, cell
+
+def main():
+    global args
+    
+    from argparse import ArgumentParser
+    parser = ArgumentParser(description='Conversion tool from stream to solution file(s) for re-integration/-refinement.')
+    
+    parser.add_argument('-i', '--input', type=str, help='Input stream file', required=True)
+    parser.add_argument('-o', '--output', type=str, help='Output solution file', required=True)
+    parser.add_argument('-g', '--geometry-out', type=str, help='Output geometry file (optional)')
+    parser.add_argument('-p', '--cell-out', type=str, help='Output cell file (optional)')
+    parser.add_argument('--file-field', type=str, help='Field in chunks for image filename', default='Image filename')
+    parser.add_argument('--event-field', type=str, help='Field in chunk for event identifier', default='Event')
+    parser.add_argument('--x-shift-field', type=str, help='Field in chunk for x-shift identifier', default='')
+    parser.add_argument('--y-shift-field', type=str, help='Field in chunk for y-shift identifier', default='')
+    parser.add_argument('--shift-factor', type=float, 
+                        help='Pre-factor for shifts, typically the pixel size in mm if the shifts are in pixel', default=1)
+    parser.add_argument('--include-pars', help='Include profile radius and resolution estimate into sol file', action='store_true')
+
+    args = parser.parse_args()
+    
+    meta = parse_stream(args.input, args.output, return_meta=True, 
+                        file_label=args.file_field, event_label=args.event_field,
+                        x_shift_label=args.x_shift_field, y_shift_label=args.y_shift_field)
+    # print('Original indexamajig call was: \n' + meta[0])
+    if args.geometry_out:
+        with open(args.geometry_out, 'w') as fh:
+            fh.write(meta[1])
+    
+    if args.cell_out:
+        if not meta[1]:
+            print('No cell found in stream file. Not writing cell file.')
+        else:
+            with open(args.cell_out, 'w') as fh:
+                fh.write(meta[2])
+     
+if __name__ == '__main__':
+    main()
+    
\ No newline at end of file