change use of measurements to records in tomography

cirq-core/cirq/experiments/qubit_characterizations.py

@@ -29,10 +29,10 @@
 import cirq.vis.histogram as cirq_histogram
 from cirq import circuits, ops, protocols
 from cirq.devices import grid_qubit
+from cirq.protocols import measurement_key_names
 
 if TYPE_CHECKING:
     from mpl_toolkits import mplot3d
-
     import cirq
 
 
@@ -527,18 +527,24 @@ def single_qubit_state_tomography(
     Returns:
         A TomographyResult object that stores and plots the density matrix.
     """
-    circuit_z = circuit + circuits.Circuit(ops.measure(qubit, key='z'))
+    keys = measurement_key_names(circuit)
+    tomo_key = "tomo_key"
+    while tomo_key in keys:
+        tomo_key = f"tomo_key{uuid.uuid4().hex}"
+
+    circuit_z = circuit + circuits.Circuit(ops.measure(qubit, key=tomo_key))
+
     results = sampler.run(circuit_z, repetitions=repetitions)
-    rho_11 = np.mean(results.measurements['z'])
+    rho_11 = np.mean(results.records[tomo_key][:, -1, :])
     rho_00 = 1.0 - rho_11
 
-    circuit_x = circuits.Circuit(circuit, ops.X(qubit) ** 0.5, ops.measure(qubit, key='z'))
+    circuit_x = circuits.Circuit(circuit, ops.X(qubit) ** 0.5, ops.measure(qubit, key=tomo_key))
     results = sampler.run(circuit_x, repetitions=repetitions)
-    rho_01_im = np.mean(results.measurements['z']) - 0.5
+    rho_01_im = np.mean(results.records[tomo_key][:, -1, :]) - 0.5
 
-    circuit_y = circuits.Circuit(circuit, ops.Y(qubit) ** -0.5, ops.measure(qubit, key='z'))
+    circuit_y = circuits.Circuit(circuit, ops.Y(qubit) ** -0.5, ops.measure(qubit, key=tomo_key))
     results = sampler.run(circuit_y, repetitions=repetitions)
-    rho_01_re = 0.5 - np.mean(results.measurements['z'])
+    rho_01_re = 0.5 - np.mean(results.records[tomo_key][:, -1, :])
 
     rho_01 = rho_01_re + 1j * rho_01_im
     rho_10 = np.conj(rho_01)

cirq-core/cirq/experiments/qubit_characterizations_test.py

@@ -171,6 +171,21 @@ def test_single_qubit_state_tomography():
     np.testing.assert_almost_equal(act_rho_2, tar_rho_2, decimal=1)
     np.testing.assert_almost_equal(act_rho_3, tar_rho_3, decimal=1)
 
+def test_single_qubit_state_tomography_unique_key():
+    # Checks if the key 'z' is already being used in the circuit, and if so picks a unique key
+    sim = cirq.Simulator()
+    qubits = cirq.LineQubit.range(2)
+
+    circuit_1 = cirq.Circuit(
+        cirq.H(qubits[0]), cirq.H(qubits[1]), cirq.measure(qubits[0], qubits[1], key='z')
+    )
+
+    result = single_qubit_state_tomography(sim, qubits[0], circuit_1, repetitions=1000)
+    act_rho_1 = result.data
+    tar_rho_1 = np.array([[0.527 + 0.0j, -0.001 - 0.02j], [-0.001 + 0.02j, 0.473 + 0.0j]])
+
+    np.testing.assert_almost_equal(act_rho_1, tar_rho_1, decimal=1)
+
 
 def test_two_qubit_state_tomography():
     # Check that the density matrices of the four Bell states closely match