1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
|
#!/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()
|
