aboutsummaryrefslogtreecommitdiff
path: root/scripts/stream2sol.py
blob: 1b081119cc744dd002a7aed698e7904567414a79 (plain)
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()