-
Notifications
You must be signed in to change notification settings - Fork 8
/
Copy pathadd-beam-params
executable file
·130 lines (104 loc) · 3.49 KB
/
add-beam-params
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
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Add beam parameters to stream
#
# Copyright © 2014-2020 Deutsches Elektronen-Synchrotron DESY,
# a research centre of the Helmholtz Association.
#
# Author:
# 2014,2017 Thomas White <[email protected]>
#
import sys
import optparse
from collections import deque
def count_crystals(f, g, start_after, stop_after):
n_crystals_seen = 0
n_crystals_written = 0
in_crystal = 0
in_header = 1
in_chunk_prelim = 0
need_end_chunk = 0
chunk_so_far = deque()
while True:
fline = f.readline()
if not fline:
break
if fline.find("Begin chunk") != -1:
in_header = 0
in_chunk_prelim = 1
chunk_so_far.clear()
if in_chunk_prelim:
chunk_so_far.append(fline)
if fline.find("End chunk") != -1:
in_chunk = 0
if need_end_chunk:
g.write(fline)
need_end_chunk = 0
if (stop_after != 0) and (n_crystals_written == stop_after):
break
if in_crystal or in_header:
g.write(fline)
if fline.find("Begin crystal") != -1:
in_chunk_prelim = 0
if ( (n_crystals_seen >= start_after)
and ((stop_after == 0) or (n_crystals_written < stop_after)) ):
in_crystal = 1
for line in chunk_so_far:
g.write(line)
chunk_so_far.clear()
if fline.find("End crystal") != -1:
n_crystals_seen += 1
if in_crystal:
n_crystals_written += 1
in_crystal = 0
need_end_chunk = 1
print("Wrote {} crystals to {}".format(n_crystals_written, opt.ofn))
def count_chunks(f, g, start_after, stop_after):
n_chunks_seen = 0
n_chunks_written = 0
in_chunk = 0
in_header = 1
while True:
fline = f.readline()
if not fline:
break
if fline.find("Begin chunk") != -1:
in_header = 0
if ( n_chunks_seen >= start_after ):
in_chunk = 1
if in_chunk or in_header:
g.write(fline)
if fline.find("End chunk") != -1:
n_chunks_seen += 1
if in_chunk:
n_chunks_written += 1
in_chunk = 0
if n_chunks_written == stop_after:
break
print("Wrote {} chunks to {}".format(n_chunks_written, opt.ofn))
op = optparse.OptionParser()
op.add_option('', '--input', action='store', type='string', dest='ifn',
help="Input stream")
op.add_option('', '--output', action='store', type='string', dest='ofn',
help="Output stream")
op.add_option('', '--start-after', action='store', type='int', dest='start',
help="Start after this many crystals", default=0)
op.add_option('', '--stop-after', action='store', type='int', dest='stop',
help="Stop after this many crystals (0=never stop)", default=0)
op.add_option('', '--chunks', action='store_true', dest='chunks',
help="Count chunks instead of crystals")
opt,arg = op.parse_args(sys.argv)
if not (opt.ifn and opt.ofn):
print("You need at least --input and --output")
exit(1)
f = open(opt.ifn, 'r')
g = open(opt.ofn, 'w')
start_after = opt.start
stop_after = opt.stop
if opt.chunks:
count_chunks(f, g, start_after, stop_after)
else:
count_crystals(f, g, start_after, stop_after)
f.close()
g.close()