aboutsummaryrefslogtreecommitdiff
path: root/scripts/stream2sol.py
diff options
context:
space:
mode:
Diffstat (limited to 'scripts/stream2sol.py')
-rwxr-xr-xscripts/stream2sol.py267
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