'''
Helper functions for adding tomography routines.
'''

from PulsePrimitives import *

def state_tomo(seq, qubits=None, numPulses=4):

	#Tomography pulses
	if numPulses == 4:
		tomoSet = [Id, X90, Y90, X]
	elif numPulses == 6:
		tomoSet = [Id, X90, X90m, Y90, Y90m, X]
	else:
		raise NameError("Only able to handle numPulses=4 or 6")

	q1, q2 = qubits
	return [seq + [p1(q1)*p2(q2), MEAS(q1,q2)] for p1 in tomoSet for p2 in tomoSet]


@@ -50,16 +50,61 @@
 es.%0A'''%0A
+from itertools import product%0Aimport operator
 %0Afrom Pu
@@ -135,18 +135,19 @@
 def 
-st
+cre
 ate_tomo
 (seq
@@ -146,25 +146,22 @@
 tomo
-(seq, 
+_blocks(
 qubits
-=None
 , nu
@@ -171,13 +171,126 @@
 lses
-=4):%0A
+, alignment='parallel'):%0A%09'''%0A%09Helper function to create the tomography pulse block in either parallel or serial.%0A%09'''
 %0A%09#T
@@ -304,16 +304,20 @@
 hy pulse
+ set
 s%0A%09if nu
@@ -443,11 +443,12 @@
 ise 
-Nam
+Valu
 eErr
@@ -496,102 +496,779 @@
 )%0A%0A%09
-q1, q2 = qubits%0A%09return %5Bseq + %5Bp1(q1)*p2(q2), MEAS(q1,q2)%5D for p1 in tomoSet for p2 in tomoSet
+#Create all combinations of pulses for the number of qubits%0A%09return %5Breduce(operator.mul, %5Bp(q) for p,q in zip(pulseSet, qubits)%5D) for pulseSet in product(tomoSet, repeat=len(qubits))%5D%0A%0Adef state_tomo(seq, qubits=None, numPulses=4):%0A%09'''%0A%09Apply state tomography readout pulses and measurement.%0A%09Expects a single entry list sequence%0A%09'''%0A%09return %5Bseq + %5BtomoBlock,  MEAS(*qubits)%5D%0A%09%09%09%09 for tomoBlock in create_tomo_blocks(qubits, numPulses)%5D%0A%0Adef process_tomo(seq, qubits=None, numPulses=4):%0A%09'''%0A%09Apply process tomography state prep. readout pulses and measurement.%0A%09Expects a single entry list sequence%0A%09'''%0A%09return %5B%5BprepBlock%5D + seq + %5BreadoutBlock,  MEAS(*qubits)%5D%0A%09%09%09%09 for prepBlock, readoutBlock in product(create_tomo_blocks(qubits, numPulses), repeat=len(qubits)) 
 %5D%0A%0A

Generalize tomography helpers.