Change use of measurements to records in tomography [continued]

cirq-core/cirq/experiments/qubit_characterizations.py

@@ -17,6 +17,7 @@
 import dataclasses
 import functools
 import itertools
+import uuid
 from typing import Any, cast, Iterator, Mapping, Sequence, TYPE_CHECKING
 
 import attrs
@@ -631,18 +632,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 = protocols.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

@@ -157,24 +157,34 @@ def test_two_qubit_randomized_benchmarking():
 def test_single_qubit_state_tomography():
     # Check that the density matrices of the output states of X/2, Y/2 and
     # H + Y gates closely match the ideal cases.
+    # Checks that unique tomography keys are generated
     simulator = sim.Simulator()
-    qubit = GridQubit(0, 0)
+    q_0 = GridQubit(0, 0)
+    q_1 = GridQubit(0, 1)
 
-    circuit_1 = circuits.Circuit(ops.X(qubit) ** 0.5)
-    circuit_2 = circuits.Circuit(ops.Y(qubit) ** 0.5)
-    circuit_3 = circuits.Circuit(ops.H(qubit), ops.Y(qubit))
+    circuit_1 = circuits.Circuit(ops.X(q_0) ** 0.5)
+    circuit_2 = circuits.Circuit(ops.Y(q_0) ** 0.5)
+    circuit_3 = circuits.Circuit(ops.H(q_0), ops.Y(q_0))
+    circuit_4 = circuits.Circuit(ops.H(q_0), ops.Y(q_0), cirq.measure(q_1, key='z'))
+    circuit_5 = circuits.Circuit(ops.H(q_0), ops.Y(q_0), cirq.measure(q_1, key='tomo_key'))
 
-    act_rho_1 = single_qubit_state_tomography(simulator, qubit, circuit_1, 1000).data
-    act_rho_2 = single_qubit_state_tomography(simulator, qubit, circuit_2, 1000).data
-    act_rho_3 = single_qubit_state_tomography(simulator, qubit, circuit_3, 1000).data
+    act_rho_1 = single_qubit_state_tomography(simulator, q_0, circuit_1, 1000).data
+    act_rho_2 = single_qubit_state_tomography(simulator, q_0, circuit_2, 1000).data
+    act_rho_3 = single_qubit_state_tomography(simulator, q_0, circuit_3, 1000).data
+    act_rho_4 = single_qubit_state_tomography(simulator, q_0, circuit_4, 1000).data
+    act_rho_5 = single_qubit_state_tomography(simulator, q_0, circuit_5, 1000).data
 
     tar_rho_1 = np.array([[0.5, 0.5j], [-0.5j, 0.5]])
     tar_rho_2 = np.array([[0.5, 0.5], [0.5, 0.5]])
     tar_rho_3 = np.array([[0.5, -0.5], [-0.5, 0.5]])
+    tar_rho_4 = np.array([[0.5, -0.5], [-0.5, 0.5]])
+    tar_rho_5 = np.array([[0.5, -0.5], [-0.5, 0.5]])
 
     np.testing.assert_almost_equal(act_rho_1, tar_rho_1, decimal=1)
     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)
+    np.testing.assert_almost_equal(act_rho_4, tar_rho_4, decimal=1)
+    np.testing.assert_almost_equal(act_rho_5, tar_rho_5, decimal=1)
 
 
 def test_two_qubit_state_tomography():