From ec3e7792dc4fba547f9812d5e4d921089916185b Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Sat, 18 Jun 2016 16:30:34 +0800
Subject: [PATCH] add MultiScanManager

---
 .../examples/master/repository/multi_scan.py  | 17 +++++++
 artiq/language/scan.py                        | 48 +++++++++++++++++--
 2 files changed, 61 insertions(+), 4 deletions(-)
 create mode 100644 artiq/examples/master/repository/multi_scan.py

diff --git a/artiq/examples/master/repository/multi_scan.py b/artiq/examples/master/repository/multi_scan.py
new file mode 100644
index 000000000..de5bb5fa0
--- /dev/null
+++ b/artiq/examples/master/repository/multi_scan.py
@@ -0,0 +1,17 @@
+from artiq.experiment import *
+
+
+class MultiScan(EnvExperiment):
+    def build(self):
+        self.setattr_argument("a", Scannable(default=LinearScan(0, 10, 4)))
+        self.setattr_argument("b", Scannable(default=LinearScan(0, 10, 4)))
+        self.setattr_argument("c", Scannable(default=LinearScan(0, 10, 4)))
+
+    def run(self):
+        msm = MultiScanManager(
+            ("a", self.a),
+            ("b", self.b),
+            ("c", self.c),
+        )
+        for point in msm:
+            print("a={} b={} c={}".format(point.a, point.b, point.c))
diff --git a/artiq/language/scan.py b/artiq/language/scan.py
index 6d59e1f7d..242e23e8c 100644
--- a/artiq/language/scan.py
+++ b/artiq/language/scan.py
@@ -3,7 +3,8 @@ Implementation and management of scan objects.
 
 A scan object (e.g. :class:`artiq.language.scan.LinearScan`) represents a
 one-dimensional sweep of a numerical range. Multi-dimensional scans are
-constructed by combining several scan objects.
+constructed by combining several scan objects, for example using
+:class:`artiq.language.scan.MultiScanManager`.
 
 Iterate on a scan object to scan it, e.g. ::
 
@@ -14,12 +15,11 @@ Iterating multiple times on the same scan object is possible, with the scan
 yielding the same values each time. Iterating concurrently on the
 same scan object (e.g. via nested loops) is also supported, and the
 iterators are independent from each other.
-
-Scan objects are supported both on the host and the core device.
 """
 
 import random
 import inspect
+from itertools import product
 
 from artiq.language.core import *
 from artiq.language.environment import NoDefault, DefaultMissing
@@ -28,7 +28,7 @@ from artiq.language import units
 
 __all__ = ["ScanObject",
            "NoScan", "LinearScan", "RandomScan", "ExplicitScan",
-           "Scannable"]
+           "Scannable", "MultiScanManager"]
 
 
 class ScanObject:
@@ -222,3 +222,43 @@ class Scannable:
         d["global_max"] = self.global_max
         d["ndecimals"] = self.ndecimals
         return d
+
+
+class MultiScanManager:
+    """
+    Makes an iterator that returns elements from the first scan object until
+    it is exhausted, then proceeds to the next iterable, until all of the
+    scan objects are exhausted. Used for treating consecutive scans as a
+    single scan.
+
+    Scan objects must be passed as a list of tuples (name, scan_object).
+    Íteration produces scan points that have attributes that correspond
+    to the names of the scan objects, and have the last value yielded by
+    that scan object.
+    """
+    def __init__(self, *args):
+        self.names = [a[0] for a in args]
+        self.scan_objects = [a[1] for a in args]
+
+        class ScanPoint:
+            def __init__(self, **kwargs):
+                self.attr = set()
+                for k, v in kwargs.items():
+                    setattr(self, k, v)
+                    self.attr.add(k)
+
+            def __repr__(self):
+                return ("<ScanPoint " +
+                    " ".join("{}={}".format(k, getattr(self, k))
+                             for k in self.attr) +
+                    ">")
+
+        self.scan_point_cls = ScanPoint
+
+    def _gen(self):
+        for values in product(*self.scan_objects):
+            d = {k: v for k, v in zip(self.names, values)}
+            yield self.scan_point_cls(**d)
+
+    def __iter__(self):
+        return self._gen()