ARTIQ scan manager ideas

floppingTRT2.py

from collections import OrderedDict
from itertools import product

from artiq import *


class ScanMe(HasEnvironment):
    def build(self):
        self.attr_argument("foo1", Scannable(LinearScan(1000, 2000, 100)))
        self.attr_argument("foo2", Scannable(LinearScan(1000, 2000, 100)))
        self.attr_argument("foo3", Scannable(LinearScan(1000, 2000, 100)))
        
        self.attr_argument("name")

    def do(self):
        print(self.name, self.foo1.current_value, self.foo2.current_value, self.foo3.current_value)

class ScanManager:
    def __init__(self, parent):
        self.parent = parent
        self.scans = OrderedDict()

    def iterate(self, names=None):
        if names is None:
            names = self.scans.keys()
        scans = [self.scans[name] for name in names]
        for values in product(*scans):
            for scan, value in zip(scans, values):
                scan.current_value = value
            yield

    def __iter__(self):
        return self.iterate()

    def get_argument(self, key, processor=None, group=None):
        if key in self.scans:
            return self.scans[key]
        value = self.parent.get_argument(key, processor, group)
        if isinstance(value, ScanObject):
            self.scans[key] = value
        return value

    def attr_argument(self, key, processor=None, group=None):
        setattr(self, key, self.get_argument(key, processor, group))


class FloppingF2(EnvExperiment):
    """ScanManager test"""

    def build(self):
        self.scan_manager = ScanManager(self)
        self.se1 = ScanMe(parent=self.scan_manager, name="se1")
        self.se2 = ScanMe(parent=self.scan_manager, name="se2")

    def run(self):
        for _ in self.scan_manager:
            self.se1.do()
            self.se2.do()