cockpit

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 import os
+import re
 import shlex
 import collections
 import datetime
@@ -11,6 +12,37 @@ import subprocess
 import curses
 
 
+def read_info_file(path):
+    result = {}
+
+    with open(path) as f:
+        for line in f:
+            key, value = line.split('=')
+            result[key] = value.strip()
+
+    return result
+
+def read_edf_id19_header(filename):
+    result = {}
+
+    with open(filename) as f:
+        data = f.read(1024)
+        pattern = re.compile(r'(.*) = (.*) ;')
+
+        for line in data.split('\n'):
+            m = pattern.match(line)
+            if m:
+                result[m.group(1).strip()] = m.group(2).strip()
+
+    return result
+
+
+def extract_motor_positions(id19_header):
+    names = id19_header['motor_mne'].split(' ')
+    values = id19_header['motor_pos'].split(' ')
+    return {k: float(v) for (k, v) in zip(names, values)}
+
+
 class Configuration(object):
 
     def __init__(self, source, destination):
@@ -105,6 +137,9 @@ class LogList(LineList):
     def info(self, s):
         self._log_time(s, self.c.get(Colors.NORMAL))
 
+    def highlight(self, s):
+        self._log_time(s, self.c.get(Colors.HIGHLIGHT))
+
     def success(self, s):
         self._log_time(s, self.c.get(Colors.SUCCESS))
 
@@ -116,7 +151,7 @@ class LogList(LineList):
 
 
 class CommandList(LineList):
-    
+
     def __init__(self, window, colors):
         self.line_list = LineList(window)
         self.normal = colors.get(Colors.NORMAL)
@@ -197,15 +232,14 @@ class Application(object):
         try:
             p = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
-            stdout, stderr = p.communicate()
+            for line in iter(p.stdout.readline, ''):
+                self.log.info(line)
 
-            def output_lines(pipe, log_func):
-                for line in pipe.split(os.linesep):
-                    if line:
-                        log_func(line)
+            for line in iter(p.stderr.readline, ''):
+                self.log.error(line)
 
-            output_lines(stdout, self.log.info)
-            output_lines(stderr, self.log.error)
+            while p.poll() is None:
+                pass
 
             if p.returncode == 0:
                 self.log.success("done")
@@ -213,44 +247,101 @@ class Application(object):
 
             self.log.error("Command returned {}".format(p.returncode))
             return False
-
-        except OSError as e:
+        except Exception as e:
             self.log.error("{}".format(e))
             return False
 
+    @property
+    def prefix(self):
+        return os.path.basename(self.config.destination)
+
+    def read_info(self):
+        info_file = os.path.join(self.config.destination, '{}.info'.format(self.prefix))
+        return read_info_file(info_file)
+
     def on_quit(self):
         self.running = False
         return True
 
     def on_sync(self):
-        self.log.info("Syncing data ...")
+        self.log.highlight("Syncing data ...")
         cmd = 'bash sync.sh {} {}'.format(self.config.source, self.config.destination)
         return self.run_command(cmd)
 
     def on_clean(self):
-        self.log.info("Cleaning {}".format(self.config.destination))
+        self.log.highlight("Cleaning {}".format(self.config.destination))
         return True
 
     def on_flat_correct(self):
-        num = 0
-        data = dict(p=self.config.destination, num=num, step=123)
+        self.log.highlight("Compute flat field correction, please wait ...")
+        info = self.read_info()
+        data = dict(path=self.config.destination, prefix=self.prefix, num=info['TOMO_N'], step=1)
+
         cmd = ('tofu flatcorrect --verbose'
                ' --reduction-mode median'
-               ' --projections "{p}/foo*.edf"'
-               ' --darks {p}/darkend0000.edf'
-               ' --flats {p}/ref*_0000.edf'
-               ' --flats2 {p}/ref*_{num}.edf'
-               ' --output {p}/fc/fc-%04i.tif'
+               ' --projections "{path}/{prefix}*.edf"'
+               ' --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)
 
     def on_optimize(self):
-        self.log.info("Optimizing ...")
+        self.log.highlight("Optimizing ...")
+
+        slices_per_device = 100
+        half_range = 1.0
+
+        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
+
+        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
+
+        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))
+
+        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']))
+
         cmd = ('optimize-parameters --verbose'
-               ' --reco-params "--x-region=-960,960,1 --y-region=-960,960,1 --overall-angle -360')
+               ' {prefix}/fc'
+               ' {opt_params}'
+               ' --reco-params "{params}"'
+               ' --final-reco-params "{params}"'
+               .format(opt_params=opt_params, params=params, prefix=self.prefix))
+
+        with open('log.txt', 'w') as f:
+            f.write(cmd)
+
         return self.run_command(cmd)
 
     def do_nothing(self):
@@ -282,6 +373,8 @@ class Application(object):
         optimize = Action('o', 'Optimize', self.on_optimize, machine.OPTIMIZE)
 
         machine.add_action(machine.START, sync)
+        machine.add_action(machine.START, flatcorrect)
+        machine.add_action(machine.START, optimize)
         machine.add_action(machine.START, clean)
         machine.add_action(machine.START, quit)