pdq2.mediator: use channel numbers from subdevices, track segment duration

artiq/devices/pdq2/mediator.py

@@ -4,8 +4,7 @@ from artiq.language import *
 frame_setup = 20*ns
 trigger_duration = 50*ns
 sample_period = 10*ns
-delay_margin_factor = 1.0001
-channels_per_pdq2 = 9
+delay_margin_factor = 1 + 1e-4
 
 
 class FrameActiveError(Exception):
@@ -38,6 +37,7 @@ class _Segment:
         self.segment_number = segment_number
 
         self.lines = []
+        self.duration = 0*s
 
         # for @kernel
         self.core = frame.pdq.core
@@ -48,12 +48,10 @@ class _Segment:
         if self.frame.pdq.armed:
             raise ArmError()
         self.lines.append((dac_divider, duration, channel_data))
+        self.duration += duration*sample_period/dac_divider
 
     def get_duration(self):
-        r = 0*s
-        for dac_divider, duration, _ in self.lines:
-            r += duration*sample_period/dac_divider
-        return r
+        return self.duration
 
     @kernel
     def advance(self):
@@ -96,7 +94,7 @@ class _Frame:
 
     def _arm(self):
         self.segment_delays = [
-            seconds_to_mu(s.get_duration()*delay_margin_factor, self.core)
+            seconds_to_mu(s.duration*delay_margin_factor, self.core)
             for s in self.segments]
 
     def _invalidate(self):
@@ -175,6 +173,9 @@ class CompoundPDQ2:
         for dev in self.pdq2s:
             dev.park()
 
+    def get_program(self):
+        return [f._get_program() for f in self.frames]
+
     def arm(self):
         if self.armed:
             raise ArmError()
@@ -182,8 +183,10 @@ class CompoundPDQ2:
             frame._arm()
         self.armed = True
 
-        full_program = [f._get_program() for f in self.frames]
-        for n, pdq2 in enumerate(self.pdq2s):
+        full_program = self.get_program()
+        n = 0
+        for pdq2 in self.pdq2s:
+            dn = pdq2.get_num_channels()
             program = []
             for full_frame_program in full_program:
                 frame_program = []
@@ -191,13 +194,13 @@ class CompoundPDQ2:
                     line = {
                         "dac_divider": full_line["dac_divider"],
                         "duration": full_line["duration"],
-                        "channel_data": full_line["channel_data"]
-                        [n*channels_per_pdq2:(n+1)*channels_per_pdq2],
+                        "channel_data": full_line["channel_data"][n:n + dn],
                         "trigger": full_line["trigger"],
                     }
                     frame_program.append(line)
                 program.append(frame_program)
             pdq2.program(program)
+            n += dn
         for pdq2 in self.pdq2s:
             pdq2.unpark()