diff options
Diffstat (limited to 'scripts')
| -rwxr-xr-x | scripts/stream2sol.py | 87 |
1 files changed, 43 insertions, 44 deletions
diff --git a/scripts/stream2sol.py b/scripts/stream2sol.py index d12ea2ba..1b081119 100755 --- a/scripts/stream2sol.py +++ b/scripts/stream2sol.py @@ -21,13 +21,13 @@ 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', +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) @@ -39,17 +39,17 @@ class Crystal: @property def initialized(self): - return all([x is not None + return all([x is not None for x in [*self.astar, *self.bstar, *self.cstar, - *self.det_shift, self.profile_radius, + *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 + return None else: - cs = ' '.join(['{0[0]} {0[1]} {0[2]}'.format(vec) + 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 @@ -65,21 +65,21 @@ class Chunk: self.start_line = line self.x_shift = 0 self.y_shift = 0 - + @property def n_cryst(self): return len(self.crystals) - - @property + + @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)) @@ -92,7 +92,7 @@ class Chunk: for ii, cryst in enumerate(self.crystals)]) -def parse_stream(stream, sol=None, return_meta=True, +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): @@ -109,18 +109,18 @@ def parse_stream(stream, sol=None, return_meta=True, 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 @@ -130,10 +130,10 @@ def parse_stream(stream, sol=None, return_meta=True, 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)) @@ -141,10 +141,10 @@ def parse_stream(stream, sol=None, return_meta=True, 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() @@ -156,49 +156,49 @@ def parse_stream(stream, sol=None, return_meta=True, elif l.startswith(BEGIN_CRYSTAL): if not curr_chunk.initialized: - raise RuntimeError('Crystal for incomplete chunk in ' + str(ln)) + 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 @@ -209,30 +209,30 @@ def parse_stream(stream, sol=None, return_meta=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 - + 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)') @@ -241,27 +241,26 @@ def main(): 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, + 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, + + 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 |