Merge pull request #2910 from sldesnoo-Delft/Tektronix_5014_performance

Significant performance improvement in Tektronix5014 driver
microsoft · Apr 21, 2021 · bd6b67c · bd6b67c
2 parents 5a2ea57 + 8c16444
commit bd6b67c
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/qcodes/instrument_drivers/tektronix/AWG5014.py b/qcodes/instrument_drivers/tektronix/AWG5014.py
@@ -876,7 +876,12 @@ def _pack_record(
                 record_data = value.encode('ASCII')
             else:
                 assert isinstance(value, (abc.Sequence, np.ndarray))
-                record_data = struct.pack('<' + dtype, *value)
+                if dtype[-1] == 'H' and isinstance(value, np.ndarray):
+                    # numpy conversion is fast
+                    record_data = value.astype('<u2').tobytes()
+                else:
+                    # argument unpacking is slow
+                    record_data = struct.pack('<' + dtype, *value)
 
         # the zero byte at the end the record name is the "(Include NULL.)"
         record_name = name.encode('ASCII') + b'\x00'
@@ -1683,11 +1688,11 @@ def _pack_waveform(
             raise TypeError('Marker 2 contains invalid values.' +
                             ' Only 0 and 1 are allowed')
 
-        wflen = len(wf)
-        packed_wf = np.zeros(wflen, dtype=np.uint16)
-        packed_wf += np.uint16(np.round(wf * 8191) + 8191 +
-                               np.round(16384 * m1) +
-                               np.round(32768 * m2))
+        # Note: we use np.trunc here rather than np.round
+        # as it is an order of magnitude faster
+        packed_wf = np.trunc(16384 * m1 + 32768 * m2
+                             + wf * 8191 + 8191.5).astype(np.uint16)
+
         if len(np.where(packed_wf == -1)[0]) > 0:
             print(np.where(packed_wf == -1))
         return packed_wf