Add add-beam-params and find-pairs

3 files changed, 204 insertions, 1 deletions

diff --git a/Makefile.am b/Makefile.am
index f460daf3..35986f75 100644
--- a/Makefile.am
+++ b/Makefile.am
@@ -179,7 +179,7 @@ script_DATA = scripts/alternate-stream scripts/cell-please \
               scripts/split-indexed scripts/stream_grep scripts/zone-axes \
               scripts/Rsplit_surface scripts/Rsplit_surface.py \
               scripts/clean-stream.py scripts/fg-graph scripts/truncate-stream \
-              scripts/gen-sfs-expand
+              scripts/gen-sfs-expand scripts/add-beam-params scripts/find-pairs
 
 EXTRA_DIST += $(script_DATA)
 
diff --git a/scripts/add-beam-params b/scripts/add-beam-params
new file mode 100755
index 00000000..23c9caaf
--- /dev/null
+++ b/scripts/add-beam-params
@@ -0,0 +1,130 @@
+#!/usr/bin/env python
+
+#
+# Add beam parameters to stream
+#
+# Copyright (c) 2014 Deutsches Elektronen-Synchrotron DESY,
+#                    a research centre of the Helmholtz Association.
+#
+# Author:
+#    2014 Thomas White <taw@physics.org>
+#
+
+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 %i crystals to %s" % (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 %i chunks to %s" % (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()
+
+
diff --git a/scripts/find-pairs b/scripts/find-pairs
new file mode 100755
index 00000000..4629767d
--- /dev/null
+++ b/scripts/find-pairs
@@ -0,0 +1,73 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+#
+# Find pairs of observations from the same pattern
+#
+# Copyright © 2014 Deutsches Elektronen-Synchrotron DESY,
+#                  a research centre of the Helmholtz Association.
+#
+# Author:
+#    2014 Thomas White <taw@physics.org>
+#
+
+import re as regexp
+
+fn1 = "test-0.5.0-satfix.stream"
+fn2 = "test-curr.stream"
+
+def find_405(f):
+    lines = dict()
+    fn = ""
+    prog = regexp.compile("^\s+([\d-]+)\s+([\d-]+)\s+([\d-]+)\s+")
+    while True:
+        fline = f.readline()
+        if not fline:
+            break
+        if fline.find("Image filename: ") != -1:
+            fn = fline.split(': ')[1].rstrip("\r\n")
+            continue
+        match = prog.match(fline)
+        if not match:
+            continue
+        h = int(match.group(1))
+        k = int(match.group(2))
+        l = int(match.group(3))
+        if (h==0) and (k==0) and (l==6):
+            lines[fn] = fline
+        if (h==0) and (k==0) and (l==-6):
+            lines[fn] = fline
+    return lines
+
+f = open(fn1, 'r')
+flines = find_405(f)
+print len(flines),"measurements in",fn1
+
+g = open(fn2, 'r')
+glines = find_405(g)
+print len(glines),"measurements in",fn2
+
+print "\nThe common measurements:\n"
+for fn in flines.keys():
+    if fn in glines:
+        print fn
+        print flines[fn].rstrip("\r\n")
+        print glines[fn]
+        print
+        del flines[fn]
+        del glines[fn]
+
+print "\nThe measurements only in",fn1,":\n"
+
+for fn in flines.keys():
+    print fn
+    print flines[fn].rstrip("\r\n")
+    print
+
+print "\nThe measurements only in",fn2,":\n"
+
+for fn in glines.keys():
+    print fn
+    print glines[fn].rstrip("\r\n")
+    print
+