cockpit

@@ -2,6 +2,8 @@
 
 import os
 import re
+import json
+import glob
 import shlex
 import collections
 import datetime
@@ -22,6 +24,7 @@ def read_info_file(path):
 
     return result
 
+
 def read_edf_id19_header(filename):
     result = {}
 
@@ -43,6 +46,20 @@ def extract_motor_positions(id19_header):
     return {k: float(v) for (k, v) in zip(names, values)}
 
 
+def have_flats(path):
+    large = os.path.join(path, 'fc')
+    small = os.path.join(path, 'fc-small')
+
+    return (os.path.exists(large) and
+            os.path.exists(small) and
+            glob.glob(os.path.join(large, '*.tif')) and
+            glob.glob(os.path.join(small, '*.tif')))
+
+
+def have_params(path):
+    return os.path.exists(os.path.join(path, 'params.json'))
+
+
 class Configuration(object):
 
     def __init__(self, source, destination):
@@ -129,10 +146,19 @@ class LogList(LineList):
     def __init__(self, window, colors):
         self.line_list = LineList(window)
         self.c = colors
+        self.log_file = open('cockpit.log', 'a')
 
     def _log_time(self, s, attr):
         timestamp = datetime.datetime.now().strftime('%H:%M:%S')
-        self.line_list.add_line('[{}] {}'.format(timestamp, s), attr)
+        log = '[{}] {}'.format(timestamp, s)
+        self.line_list.add_line(log, attr)
+        self.log_file.write(log)
+
+        if not log.endswith('\n'):
+            self.log_file.write('\n')
+
+        self.log_file.flush()
+        os.fsync(self.log_file.fileno())
 
     def info(self, s):
         self._log_time(s, self.c.get(Colors.NORMAL))
@@ -196,6 +222,8 @@ class StateMachine(object):
     CLEAN = 3
     FLATCORRECT = 4
     OPTIMIZE = 5
+    QUICK_OPTIMIZE = 6
+    RECONSTRUCT = 7
 
     def __init__(self):
         self.current = StateMachine.START
@@ -207,6 +235,8 @@ class StateMachine(object):
             StateMachine.QUIT: collections.OrderedDict(),
             StateMachine.SYNC: collections.OrderedDict(),
             StateMachine.OPTIMIZE: collections.OrderedDict(),
+            StateMachine.QUICK_OPTIMIZE: collections.OrderedDict(),
+            StateMachine.RECONSTRUCT: collections.OrderedDict(),
         }
 
     def add_action(self, from_state, action):
@@ -259,6 +289,98 @@ class Application(object):
         info_file = os.path.join(self.config.destination, '{}.info'.format(self.prefix))
         return read_info_file(info_file)
 
+    def optimize(self, resize):
+        slices_per_device = 100
+        half_range = 1.0
+
+        info = self.read_info()
+        axis = (float(info['Col_end']) + 1) / 2.0
+        axis_step = 0.25
+
+        fname = os.path.join(self.config.destination, '{}0000.edf'.format(self.prefix))
+        header = read_edf_id19_header(fname)
+        motor_pos = extract_motor_positions(header)
+
+        inclination_angle = motor_pos['rytot']
+        theta = 90.0 - inclination_angle
+        angle_step = 0.025
+        angle_start = theta - half_range
+        angle_stop = theta + half_range
+
+        x_region = 960
+        y_region = 960
+        fc_path = '{prefix}/fc'.format(prefix=self.prefix)
+
+        if resize:
+            x_region /= 2
+            y_region /= 2
+            axis /= 2
+            axis_step /= 2
+            fc_path = '{prefix}/fc-small'.format(prefix=self.prefix)
+
+        axis_start = axis - slices_per_device * axis_step
+        axis_stop = axis + slices_per_device * axis_step
+
+        self.log.info(" Using theta = {}, inclination angle = {}".format(theta, inclination_angle))
+        self.log.info(" Scanning lamino angle within [{}:{}:{}]".format(angle_start, angle_stop, angle_step))
+        self.log.info(" Scanning x axis within [{}:{}:{}]".format(axis_start, axis_stop, axis_step))
+
+        opt_params = ('--num-iterations 2'
+                      ' --axis-range={ax_start},{ax_stop},{ax_step}'
+                      ' --lamino-angle-range={an_start},{an_stop},{an_step}'
+                      ' --metric kurtosis --z-metric kurtosis'
+                      .format(ax_start=axis_start, ax_stop=axis_stop, ax_step=axis_step,
+                              an_start=angle_start, an_stop=angle_stop, an_step=angle_step))
+
+        params = ('--x-region="-{x_region},{x_region},1"'
+                  ' --y-region="-{y_region},{y_region},1"'
+                  ' --overall-angle -360'
+                  ' --pixel-size {pixel_size}e-6'
+                  ' --roll-angle 0'
+                  ' --slices-per-device 100'
+                  .format(pixel_size=info['PixelSize'], x_region=x_region, y_region=y_region))
+
+        cmd = ('optimize-parameters --verbose'
+               ' {fc_path}'
+               ' {opt_params}'
+               ' --reco-params "{params}"'
+               ' --params-filename {prefix}/params.json'
+               .format(opt_params=opt_params, params=params, fc_path=fc_path, prefix=self.prefix))
+
+        return self.run_command(cmd)
+
+    def on_reconstruct(self):
+        self.log.highlight("Reconstructing ...")
+
+        with open(os.path.join(self.prefix, 'params.json')) as f:
+            info = self.read_info()
+            opt = json.load(f)
+
+            lamino_angle = opt['lamino-angle']['value']
+            axis = opt['x-center']['value']
+
+            x_region = 960
+            y_region = 960
+
+            cmd = ('tofu lamino --verbose'
+                   ' --projections {prefix}/fc'
+                   ' --x-region="-{x_region},{x_region},1"'
+                   ' --y-region="-{y_region},{y_region},1"'
+                   ' --lamino-angle {lamino_angle}'
+                   ' --axis {x_axis},{y_axis}'
+                   ' --overall-angle -360'
+                   ' --pixel-size {pixel_size}e-6'
+                   ' --roll-angle 0'
+                   ' --slices-per-device 600'
+                   ' --output {prefix}/Slices'
+                   .format(pixel_size=info['PixelSize'],
+                           x_region=x_region, y_region=y_region,
+                           lamino_angle=lamino_angle,
+                           x_axis=axis, y_axis=float(info['Dim_2']) / 2,
+                           prefix=self.prefix))
+
+        return self.run_command(cmd)
+
     def on_quit(self):
         self.running = False
         return True
@@ -275,7 +397,8 @@ class Application(object):
     def on_flat_correct(self):
         self.log.highlight("Flat field correction ...")
         info = self.read_info()
-        data = dict(path=self.config.destination, prefix=self.prefix, num=info['TOMO_N'], step=1)
+        path = self.config.destination
+        data = dict(path=path, prefix=self.prefix, num=info['TOMO_N'], step=1)
 
         cmd = ('tofu flatcorrect --verbose'
                ' --reduction-mode median'
@@ -283,66 +406,45 @@ class Application(object):
                ' --darks {path}/darkend0000.edf'
                ' --flats {path}/ref*_0000.edf'
                ' --flats2 {path}/ref*_{num}.edf'
-               ' --output {path}/fc/fc-%04i.tif'
                ' --number {num}'
                ' --step {step}'
                ' --absorptivity'
                ' --fix-nan-and-inf'.format(**data))
 
-        return self.run_command(cmd)
+        large_cmd = cmd + ' --output {path}/fc/fc-%04i.tif'.format(path=path)
+        small_cmd = cmd + ' --resize 2 --output {path}/fc-small/fc-%04i.tif'.format(path=path)
 
-    def on_optimize(self):
-        self.log.highlight("Optimizing ...")
+        if not self.run_command(small_cmd):
+            self.log.error("Could not create small flat field corrected projections")
+            return result
 
-        slices_per_device = 100
-        half_range = 1.0
+        return self.run_command(large_cmd)
 
-        info = self.read_info()
-        axis = (float(info['Col_end']) + 1) / 2.0
-        axis_step = 0.25
-        axis_start = axis - slices_per_device * axis_step
-        axis_stop = axis + slices_per_device * axis_step
+    def on_quick_optimize(self):
+        self.log.highlight("Quick optimization ...")
 
-        fname = os.path.join(self.config.destination, '{}0000.edf'.format(self.prefix))
-        header = read_edf_id19_header(fname)
-        motor_pos = extract_motor_positions(header)
+        if self.optimize(True):
+            with open(os.path.join(self.prefix, 'params.json')) as f:
+                opt = json.load(f)
+                self.log.highlight(" Optimal axis: {}".format(opt['lamino-angle']['value']))
+                self.log.highlight(" Optimal center: {}".format(opt['x-center']['value'] * 2.0))
 
-        inclination_angle = motor_pos['rytot']
-        theta = 90.0 - inclination_angle
-        angle_step = 0.025
-        angle_start = theta - half_range
-        angle_stop = theta + half_range
+            return True
 
-        self.log.info(" Using theta = {}, inclination angle = {}".format(theta, inclination_angle))
-        self.log.info(" Scanning angle within [{}:{}:{}]".format(angle_start, angle_stop, angle_step))
-        self.log.info(" Scanning axis within [{}:{}:{}]".format(axis_start, axis_stop, axis_step))
+        return False
 
-        opt_params = ('--num-iterations 2'
-                      ' --axis-range={ax_start},{ax_stop},{ax_step}'
-                      ' --lamino-angle-range={an_start},{an_stop},{an_step}'
-                      ' --metric kurtosis --z-metric kurtosis'
-                      .format(ax_start=axis_start, ax_stop=axis_stop, ax_step=axis_step,
-                              an_start=angle_start, an_stop=angle_stop, an_step=angle_step))
-
-        params = ('--x-region=-960,960,1'
-                  ' --y-region=-960,960,1'
-                  ' --overall-angle -360'
-                  ' --pixel-size {pixel_size}e-6'
-                  ' --roll-angle 0'
-                  ' --slices-per-device 100'
-                  .format(pixel_size=info['PixelSize']))
+    def on_optimize(self):
+        self.log.highlight("Optimizing ...")
 
-        cmd = ('optimize-parameters --verbose'
-               ' {prefix}/fc'
-               ' {opt_params}'
-               ' --reco-params "{params}"'
-               ' --params-filename params.json'
-               .format(opt_params=opt_params, params=params, prefix=self.prefix))
+        if self.optimize(False):
+            with open(os.path.join(self.prefix, 'params.json')) as f:
+                opt = json.load(f)
+                self.log.highlight(" Optimal axis: {}".format(opt['lamino-angle']['value']))
+                self.log.highlight(" Optimal center: {}".format(opt['x-center']['value']))
 
-        with open('log.txt', 'w') as f:
-            f.write(cmd)
+            return True
 
-        return self.run_command(cmd)
+        return False
 
     def do_nothing(self):
         return True
@@ -366,32 +468,42 @@ class Application(object):
 
         machine = StateMachine()
 
-        quit = Action('q', 'Quit', self.on_quit, machine.QUIT)
+        quit = Action('e', 'Exit', self.on_quit, machine.QUIT)
         sync = Action('s', 'Sync', self.on_sync, machine.SYNC)
-        flatcorrect = Action('f', 'Flat correct', self.on_flat_correct, machine.FLATCORRECT)
         clean = Action('c', 'Clean', self.on_clean, machine.CLEAN)
+        quick_optimize = Action('q', 'Quick', self.on_quick_optimize, machine.QUICK_OPTIMIZE)
         optimize = Action('o', 'Optimize', self.on_optimize, machine.OPTIMIZE)
+        reconstruct = Action('r', 'Reconstruct', self.on_reconstruct, machine.RECONSTRUCT)
 
         machine.add_action(machine.START, sync)
-        machine.add_action(machine.START, flatcorrect)
+        machine.add_action(machine.START, quick_optimize)
         machine.add_action(machine.START, optimize)
-        machine.add_action(machine.START, clean)
         machine.add_action(machine.START, quit)
 
-        machine.add_action(machine.SYNC, flatcorrect)
-        machine.add_action(machine.SYNC, clean)
         machine.add_action(machine.SYNC, quit)
 
-        machine.add_action(machine.FLATCORRECT, optimize)
-        machine.add_action(machine.FLATCORRECT, quit)
+        machine.add_action(machine.QUICK_OPTIMIZE, reconstruct)
+        machine.add_action(machine.QUICK_OPTIMIZE, optimize)
+        machine.add_action(machine.QUICK_OPTIMIZE, quit)
 
-        machine.add_action(machine.OPTIMIZE, sync)
-        machine.add_action(machine.OPTIMIZE, clean)
+        machine.add_action(machine.OPTIMIZE, reconstruct)
+        machine.add_action(machine.OPTIMIZE, quick_optimize)
         machine.add_action(machine.OPTIMIZE, quit)
 
         machine.add_action(machine.CLEAN, sync)
         machine.add_action(machine.CLEAN, quit)
 
+        if have_params(self.config.destination):
+            machine.add_action(machine.START, reconstruct)
+
+        if not have_flats(self.config.destination):
+            flatcorrect = Action('f', 'Flat correct', self.on_flat_correct, machine.FLATCORRECT)
+            machine.add_action(machine.START, flatcorrect)
+            machine.add_action(machine.SYNC, flatcorrect)
+            machine.add_action(machine.FLATCORRECT, quick_optimize)
+            machine.add_action(machine.FLATCORRECT, optimize)
+            machine.add_action(machine.FLATCORRECT, quit)
+
         cmd_window.set_actions(machine.actions)
 
         self.log = LogList(right_pane, colors)