Reset a register subset

Workaround to specify all the multi-qubit branches even when only a
subset is measured and reset

QGL/BasicSequences/Feedback.py

@@ -9,9 +9,16 @@
 
 
 @qfunction
-def qreset(qubits, signVec, measDelay, buf):
-    # for each qubit, build the set of feedback actions to perform when
-    # receiving a zero or one in the comparison register
+def qreset(qubits, signVec, measDelay, buf, reg_size=None, TDM_map=None):
+    """
+    for each qubit, build the set of feedback actions to perform when
+    receiving a zero or one in the comparison register
+    reg_size, TDM_map: optional arguments to reset a subset of the qubit register (see Reset)
+    """
+    if not reg_size:
+        reg_size = len(qubits)
+        TDM_map = np.arange(reg_size,0,-1)
+
     FbGates = []
     for ct, q in enumerate(qubits):
         if signVec[ct] == 0:
@@ -23,9 +30,12 @@ def qreset(qubits, signVec, measDelay, buf):
     # load register
     seq = [Id(qubits[0], measDelay), qwait(kind='CMP'), Id(qubits[0], buf)]
     # create a branch for each possible comparison value
-    for ct in range(2**len(qubits)):
-        seq += qif(ct, [FbSeq[ct]])
-
+    for ct in range(2**reg_size):
+        # duplicate branches for the irrelevant results if reg_size > len(TDM_map)
+        meas_result = [(ct & TDM_bit)>0 for TDM_bit in 2**(TDM_map-1)]
+        branch_idx = sum([t*2**ind for ind,t in enumerate((meas_result))])
+        seq += qif(ct, [FbSeq[branch_idx]])
+
     return seq
 
 
@@ -37,38 +47,42 @@ def Reset(qubits,
           showPlot=False,
           measChans=None,
           docals=True,
-          calRepeats=2):
+          calRepeats=2,
+          reg_size=None,
+          TDM_map=None):
     """
 
-	Variable amplitude Rabi nutation experiment for an arbitrary number of qubits simultaneously
-
-	Parameters
-	----------
-	qubits : tuple of logical channels to implement sequence (LogicalChannel)
-	measDelay : delay between end of measuerement and LOADCMP
-	signVec : conditions for feedback. Tuple of 0 (flip if signal is above threshold) and 1 (flip if below) for each qubit. Default = 0 for all qubits
-	doubleRound : if true, double round of feedback
-	showPlot : whether to plot (boolean)
-	measChans : tuble of qubits to be measured (LogicalChannel)
-	docals, calRepeats: enable calibration sequences, repeated calRepeats times
+    Variable amplitude Rabi nutation experiment for an arbitrary number of qubits simultaneously
 
-	Returns
-	-------
-	plotHandle : handle to plot window to prevent destruction
-	"""
+    Parameters
+    ----------
+    qubits : tuple of logical channels to implement sequence (LogicalChannel)
+    measDelay : delay between end of measuerement and LOADCMP
+    signVec : conditions for feedback. Tuple of 0 (flip if signal is above threshold) and 1 (flip if below) for each qubit. Default = 0 for all qubits
+    doubleRound : if true, double round of feedback
+    showPlot : whether to plot (boolean)
+    measChans : tuble of qubits to be measured (LogicalChannel)
+    docals, calRepeats: enable calibration sequences, repeated calRepeats times
+    reg_size: total number of qubits, including those that are not reset. Default set to len(qubits)
+    TDM_map: map each qubit to a TDM digital input. Default: np.array(qN, qN-1, ..., q1) from MSB to LSB.
+    
+    Returns
+    -------
+    plotHandle : handle to plot window to prevent destruction
+    """
     if measChans is None:
         measChans = qubits
 
     if signVec == None:
         signVec = (0, ) * len(qubits)
 
-    seqs = [prep + [qreset(qubits, signVec, measDelay, buf)]
+    seqs = [prep + [qreset(qubits, signVec, measDelay, buf, reg_size=reg_size, TDM_map=TDM_map)]
             for prep in create_cal_seqs(qubits, 1)]
     measBlock = reduce(operator.mul, [MEAS(q) for q in qubits])
     if doubleRound:
         for seq in seqs:
             seq += [measBlock]
-            seq.append(qreset(qubits, signVec, measDelay, buf))
+            seq.append(qreset(qubits, signVec, measDelay, buf, reg_size=reg_size, TDM_map=TDM_map))
 
     # add final measurement
     for seq in seqs:

QGL/ControlFlow.py

@@ -37,9 +37,9 @@ def qfunction(func):
     target = {}
 
     @wraps(func)
-    def crfunc(*args):
+    def crfunc(*args, **kwargs):
         if args not in target:
-            seq = func(*args) + [Return()]
+            seq = func(*args, **kwargs) + [Return()]
             target[args] = label(seq)
             qfunction_seq[label(seq)] = seq
         return Call(target[args])