Initial Python infrastructure to build more complex processing pipelines and the corresponding changes in ROOF filters

6 files changed, 358 insertions, 0 deletions

diff --git a/tests/roof/__init__.py b/tests/roof/__init__.py
new file mode 100644
index 0000000..b8023d8
--- /dev/null
+++ b/tests/roof/__init__.py
@@ -0,0 +1 @@
+__version__ = '0.0.1'
diff --git a/tests/roof/arguments.py b/tests/roof/arguments.py
new file mode 100644
index 0000000..22ea42b
--- /dev/null
+++ b/tests/roof/arguments.py
@@ -0,0 +1,32 @@
+import argparse
+from roof.defaults import roof_data_types
+
+def roof_get_args():
+    data_types = []
+    for stage in roof_data_types:
+        data_types += roof_data_types[stage].keys()
+    data_types = set(data_types)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-c', '--config',   dest="config",          default="roof.json",                                                                help="ROOF configuration (JSON)")
+
+    # Modes
+    parser.add_argument('-s', '--simulate',     dest="simulate",        default=None,   action="store_true",                                            help="Simulation mode, read data from files instead of network")
+    parser.add_argument('-b', '--benchmark',    dest="benchmark",       default=None,   action="store_true",                                            help="Bencmarking mode, writes to /dev/null")
+    parser.add_argument('-g', '--gui',          dest='gui',             default=False,  action="store_true",                                            help="Visualize data")
+    parser.add_argument('-t', '--track',        dest='track',           default=False,  action="store_true",                                            help="Track & control experiment")
+    
+    parser.add_argument(      '--no-roof',      dest="noroof",          default=False,  action="store_true",                                            help="Disable ROOF, only network testing (no sinogram building, store linearly)")
+
+    # I/O
+    #parser.add_argument('-i', '--input',       dest="input",           default=None,                                                                   help="Reconstruct from sinograms")
+    parser.add_argument('-o', '--output',       dest="output",          default=None,                                                                   help="Output file(s)")
+    parser.add_argument('-r', '--read',         dest="read",            default=None,   choices=data_types,     nargs='?', const="raw_sinograms",       help="Read recorded sinograms instead of listening on the network")
+    parser.add_argument('-w', '--write',        dest="write",           default=None,   choices=data_types,     nargs='?', const="slices",              help="Only generate flat-fields, dark-fields, or sinograms (default)")
+    parser.add_argument(      '--format',       dest="format",          default=None,                                                                   help="Override default storage format")
+
+    # Limits & Filtering
+    parser.add_argument('-n', '--number',       dest="number",          default=None,   type=int,                                                       help="Specify number of frames to capture (limits number of captured frames irrespective of further filtering)")
+    parser.add_argument('-p', '--plane',        dest="plane",           default=None,   type=int,                                                       help="Only process the specified detector plane (indexed from 1)")
+
+    return parser.parse_args()
diff --git a/tests/roof/config.py b/tests/roof/config.py
new file mode 100644
index 0000000..e085ed8
--- /dev/null
+++ b/tests/roof/config.py
@@ -0,0 +1,67 @@
+import re
+import json
+
+from roof.arguments import roof_get_args
+from roof.defaults import roof_default_paths, roof_raw_data_types
+
+class RoofConfig:
+    def __init__(self, config=None):
+        self.args = roof_get_args()
+        self.config_file = self.get_arg('config', 'roof.json') if config is None else config
+        with open(self.config_file) as json_file:
+            self.cfg = json.load(json_file)
+
+        self.path = self.get_opt('data', 'base_path', './')
+        self.planes = self.get_opt('hardware', 'planes', 1)
+        self.modules = self.get_opt('hardware', 'modules', None)
+        self.streams = self.get_opt('network', 'streams', 1 if self.modules is None else self.modules)
+        self.bit_depth = self.get_opt('hardware', 'bit_depth', 8)
+
+        if self.args.number is None: self.args.number = 0 if self.args.benchmark else self.planes
+
+        # Consistency and default mode
+        if (self.args.plane is not None) and (self.args.plane > self.planes):
+            raise ValueError("Only {} planes in configuration, but the plane {} is requested".format(self.planes, self.args.plane))
+        
+        n_modes = (int(self.args.gui) + int(self.args.track) + int(0 if self.args.write is None else 1))
+        if n_modes  > 1:
+            raise ValueError("GUI, Control, and Write modes are mutualy incompatible")
+        elif n_modes == 0:
+            self.args.write = "raw_sinograms"
+
+
+    def get_arg(self, arg, default = None):
+        ret = getattr(self.args, arg)
+        return ret if ret is not None else default
+
+    def get_opt(self, group, item, default = None):
+        if self.cfg.get(group, {}).get(item) != None:
+            return self.cfg[group][item]
+        else:
+            return default
+
+    def get_roof_path(self, data_type):
+        subpath = self.get_opt('data', data_type)
+        if subpath is None: subpath = roof_default_paths[data_type] 
+        if subpath is None: raise "Unknown data type %s is requested" % subpath
+        return subpath if subpath.startswith('/') else  self.path + '/' + subpath
+
+    def get_writer_type(self):
+        return None if self.args.benchmark else self.args.write if self.args.write else 'raw_sinograms'
+    
+    def get_writer_path(self):
+        data_type = self.get_writer_type()
+        if data_type is not None:
+            path = self.args.output if self.args.output is not None else self.get_roof_path(data_type)
+            if self.args.format: path = re.sub('\.([^.]+)$', '.' + self.args.format, path)
+            return path
+        return None
+            
+    def check_writer_type_is_raw_or_none(self):
+        data_type = self.get_writer_type()
+        data_path = self.get_writer_path()
+        return (data_type is None) or ((data_type in roof_raw_data_types) and re.search('\.raw$', data_path))
+
+    def check_writer_type_is_raw(self):
+        data_type = self.get_writer_type()
+        return (data_type is not None) and self.check_writer_type_is_raw_or_none()
diff --git a/tests/roof/defaults.py b/tests/roof/defaults.py
new file mode 100644
index 0000000..eed3fe5
--- /dev/null
+++ b/tests/roof/defaults.py
@@ -0,0 +1,42 @@
+roof_default_paths = {
+ 'flat_fields':                 "flats/flat_%04u.raw",
+ 'dark_fields':                 "darks/dark_%04u.raw",
+ 'raw_sinograms':               "raw/sino_%04u.raw",
+ 'fan_sinograms':               "fan/sino_%04u.tif",
+ 'parallel_sinograms':          "par/sino_%04u.tif",
+ 'filtered_sinograms':          "flt/sino_%04u.tif",
+ 'slices':                      "slices/slice_%04u.raw"
+}
+
+#roof_default_simulation_paths = {
+# 'data':                        "sim/data_%02u.dat",
+# 'flat_fields':                 "sim/flat_%02u.dat",
+# 'dark_fields':                 "sim/dark_%02u.dat"
+#}
+
+roof_filters = {
+ 'correction':                  [ "flat-field-correct" ],
+ 'reconstruction':              [ "roof-fan2par", "fft", "filter", "ifft", "backproject" ],
+ 'control':                     [ ],
+ 'visualization':               [ ]
+}
+
+roof_data_types = {
+ 'correction': {
+    'flat_fields':              'flat-field-correct',
+    'dark_fields':              'flat-field-correct',
+    'raw_sinograms':            'flat-field-correct',
+    'fan_sinograms':            None    
+  },
+  'reconstruction': {
+    'fan_sinograms':            'roof-fan2par',
+    'parallel_sinograms':       'fft',
+    'filtered_sinograms':       'backproject',
+    'slices':                    None
+  },
+  'control': {
+  }
+}
+
+roof_raw_data_types = [k for k, v in roof_data_types['correction'].items() if v is not None ]
+roof_aux_data_types = [v for v in roof_raw_data_types if 'sino' not in v ]
diff --git a/tests/roof/graph.py b/tests/roof/graph.py
new file mode 100644
index 0000000..c34a3ed
--- /dev/null
+++ b/tests/roof/graph.py
@@ -0,0 +1,203 @@
+import re
+import gi
+
+gi.require_version('Ufo', '0.0') 
+from gi.repository import Ufo
+from gi.repository import GObject
+
+from roof.config import RoofConfig
+from roof.defaults import roof_filters, roof_data_types, roof_raw_data_types, roof_aux_data_types
+from roof.utils import get_filenames
+
+class RoofGraph(RoofConfig):
+    def __init__(self, config=None):
+        self.pm  = Ufo.PluginManager()
+        self.graph = Ufo.TaskGraph()
+        self.scheduler = Ufo.Scheduler()
+        self.tasks = {}
+
+        super(RoofGraph, self).__init__()
+
+    def get_task(self, name, **kwargs):
+        task = self.pm.get_task(name)
+        task.set_properties(name, **kwargs)
+        return task
+
+    def save_task(self, stage, alias, task):
+        if stage is None: stage = "general"
+        if stage not in self.tasks: self.tasks[stage] = {}
+        self.tasks[stage][alias if alias is not None else name] = task
+        return task
+
+    def get_roof_task(self, name, **kwargs):
+        kwargs.update(config = self.config_file)
+        return self.get_task(name, **kwargs)
+    
+    def get_processor_task(self, stage, name, **kwargs):
+        extra_args = self.get_opt(stage, name + '-options')
+        if extra_args is not None: kwargs.update(extra_args)
+        if (re.compile('roof').match(name)): kwargs.update(config = self.config_file)
+        return self.save_task(stage, name, self.get_task(name, **kwargs))
+    
+    def get_reader(self):
+        first = self.get_opt('data', 'first_file_number', 1)
+        if self.args.read:
+            # Reconstruction from standard UFO files
+            path = self.get_roof_path(self.args.read)
+            step = 1
+            if (self.args.plane is not None) and (self.args.plane > 0): 
+                first += self.args.plane - 1;
+                step = self.planes
+
+            params = { 'path': path, 'first': first, 'step': step }
+            if self.args.number: 
+                params['number'] = self.args.number
+            
+            print ("Reading {} data from {}".format(self.args.read,path))
+            return self.get_task('read', **params)
+        else:
+            path = None
+            if self.args.simulate:
+                first = self.get_opt('simulation', 'first_file_number', first)
+                base_path = self.get_opt('simulation', 'base_path', self.path)
+                read_path = self.get_opt('simulation', self.args.write if self.args.write and self.args.write in roof_aux_data_types else 'data')
+                path = read_path if read_path.startswith('/') else  base_path + '/' + read_path
+                print ("Simulating packets from {}".format(path))
+
+            # Reconstruction from network or simulated data (also generation of flat/dark-fields)
+            build_type = "raw" if self.args.noroof else "sino" if self.check_writer_type_is_raw() else "ufo"
+            build = self.get_roof_task('roof-build', simulate = self.args.simulate, number = self.args.number, build = build_type)
+            for id in range(self.streams):
+                read = self.get_roof_task('roof-read', id = id, simulate = self.args.simulate, path = path, first_file_number = first)
+                self.graph.connect_nodes(read, build)
+                build.bind_property('stop', read, 'stop', GObject.BindingFlags.DEFAULT)
+
+            return build
+
+    def get_writer(self):
+        path = self.get_writer_path()
+        if path is None:
+            print ("Starting ROOF using NULL writter")
+            write = self.get_task('null')
+        else:
+            # FIXME: If writting non raw data, we may need to generate all-0-frames if something broken/corrupted.
+            print ("Starting ROOF streaming to {}".format(path))
+            write = self.get_task('write', filename=path)
+        return write
+
+    def get_correction_flat_field_correct(self, head):
+        # Standard UFO reconstruction stack distinguish flat/dark-fields recorded before and after experiment. We only do 'before experiment' part.
+        darks = self.get_roof_path('dark_fields')
+        n_darks = len(get_filenames(darks))
+        if n_darks == 0: raise FileNotFoundError("Dark fields are not found in {}".format(darks))
+        flats = self.get_roof_path('falt_fields')
+        n_flats = len(get_filenames(flats))
+        if n_flats == 0: raise FileNotFoundError("Flat fields are not found in {}".format(flats))
+        dark_reader = self.get_task('read', path = darks)
+        flat_reader = self.get_task('read', path = flats)
+
+        # We are using standard get_task here because this is too generic plugin to allow config-based customization
+        mode = self.get_opt('correction', 'aggregation', 'average')
+        if mode == 'median':
+            dark_stack = self.get_task('stack', number = n_darks)
+            dark_reduced = self.get_task('flatten', mode = 'median')
+            flat_stack = self.get_task('stack', number = n_flats)
+            flat_reduced = self.get_task('flatten', mode = 'median')
+
+            self.graph.connect_nodes(dark_reader, dark_stack)
+            self.graph.connect_nodes(dark_stack, dark_reduced)
+            self.graph.connect_nodes(flat_reader, flat_stack)
+            self.graph.connect_nodes(flat_stack, flat_reduced)
+        elif mode == 'average':
+            dark_reduced = self.get_task('average')
+            flat_reduced = self.get_task('average')
+            self.graph.connect_nodes(dark_reader, dark_reduced)
+            self.graph.connect_nodes(flat_reader, flat_reduced)
+        else:
+            raise ValueError('Invalid reduction mode')
+
+        ffc = self.get_task('flat-field-correct') # dark_scale=args.dark_scale, absorption_correct=args.absorptivity, fix_nan_and_inf=args.fix_nan_and_inf)
+        self.graph.connect_nodes_full(head, ffc, 0)
+        self.graph.connect_nodes_full(dark_reduced, ffc, 1)
+        self.graph.connect_nodes_full(flat_reduced, ffc, 2)
+        return ffc
+
+    def get_processor(self, head, stage, writer = None):
+        # skip (but not if not already skipped in previous processor)
+        # how to connect readers to ffc?
+        
+        filters = self.get_opt(stage, 'filters', roof_filters[stage])
+        read_here = self.args.read and self.args.read in roof_data_types[stage].keys()
+        write_here = self.args.write and self.args.write in roof_data_types[stage].keys()
+        
+        start_pos = 0
+        if read_here:
+            start_filter = roof_data_types[stage][self.args.read]
+            start_pos = filters.index(start_filter)
+
+        last_pos = len(filters)
+        if write_here:
+            stop_filter = roof_data_types[stage][self.args.write]
+            if stop_filter: last_pos = filters.index(stop_filter)
+
+        # Will just execute empty range if we start reading from the end (e.g. 'fan-sinograms' in correction)
+        for i in range(start_pos, last_pos):
+            method = 'get_' + stage + '_' + filters[i].replace('-','_')
+            if method in dir(self):
+                f = getattr(self, method)(head)
+            else:
+                f = self.get_processor_task(stage, filters[pos])
+                graph.connect_nodes(head, f)
+            head = f 
+
+        if write_here and writer:
+            self.graph.connect_nodes(head, writer)
+                        
+        return None if write_here else head
+
+
+    def get(self):
+        reader = self.get_reader()
+        writer = self.get_writer()
+
+        # We support following operation modes (defined by modifiers -w -c -g ]
+        # - Record mode: Writting raw data (raw-sinograms, flat-fields, dark-fields)                                                    [ no modified or -w <...> ]
+        # - Write mode: The reconstruction is performed and data is written after the specified step (default)                          [ -w <all other data types> ]
+        # - Control mode: Control branch and raw data writting                                                                          [ -c ]
+        # - GUI mode: Visualization in GUI + raw_sinograms are written when enabled in GUI + some control tasks (also when enabled)     [ -g ]
+    
+        head = reader
+        # Check if we are branching here
+        if (self.args.track or self.args.gui) and (self.get_data_type() is not None):
+            # FIXME: In GUI mode we can add here a 'write filter' to pause/resume writting. Alternative is to pass gobject flaga to fastwriter (this will be limited to fastwriter, then, which is likely OK)
+            # FIXME: we may need to convert in the end if we are writing raw data and the data is comming from net/simulation
+            # In other case (non branch), either we have already converted (in reader) or we don't need to convert (writing raw data). Small performance penalty if we convert before filter, but ....
+            copy = Ufo.CopyTask()
+            self.graph.connect_nodes(reader, copy)
+            self.graph.connect_nodes(copy, writer)
+            head = copy
+
+        # Sinograms are already filtered in the reader
+        if not self.args.read:
+            main_filter = self.get_task('roof-filter', plane = self.args.plane) if self.args.plane else None
+            if main_filter:
+                self.graph.connect_nodes(head, main_filter)
+                head = main_filter
+
+        class finish(Exception): pass
+        try:
+            if not self.args.read or self.args.read in roof_data_types['correction'].keys():
+                head = self.get_processor(head, 'correction', writer)
+                if not head: raise finish()
+
+            if head != reader or self.args.read in roof_data_types['reconstruction'].keys():
+                head = self.get_processor(head, 'reconstruction', writer)
+                if not head: raise finish()
+
+            # if head split to 3 branches.... Otherwise, continue with control branch...
+        except finish:
+            pass
+        
+    def run(self):
+        self.scheduler.run(self.graph)
+        
\ No newline at end of file
diff --git a/tests/roof/utils.py b/tests/roof/utils.py
new file mode 100644
index 0000000..eb389ed
--- /dev/null
+++ b/tests/roof/utils.py
@@ -0,0 +1,13 @@
+import glob
+import logging
+import math
+import os
+
+def get_filenames(path):
+    """Get all filenams from *path*, which could be a directory or a pattern
+    for matching files in a directory.
+    """
+    if os.path.isdir(path):
+        path = os.path.join(path, '*')
+
+    return sorted(glob.glob(path))