nit

cirq-core/cirq/experiments/z_phase_calibration_test.py

@@ -93,6 +93,13 @@ def test_calibrate_z_phases(angles, error):
     initial_unitary = cirq.unitary(original_gate)
     final_unitary = cirq.unitary(calibrated_gate)
     target_unitary = cirq.unitary(actual_gate)
-    assert _trace_distance(final_unitary, target_unitary) < _trace_distance(
-        initial_unitary, target_unitary
-    )
+    maximally_mixed_state = np.eye(4) / 2
+    dm_initial = initial_unitary @ maximally_mixed_state @ initial_unitary.T.conj()
+    dm_final = final_unitary @ maximally_mixed_state @ final_unitary.T.conj()
+    dm_target = target_unitary @ maximally_mixed_state @ target_unitary.T.conj()
+
+    original_dist = _trace_distance(dm_initial, dm_target)
+    new_dist = _trace_distance(dm_final, dm_target)
+
+    # Either we reduced the error or the error is small enough.
+    assert new_dist < original_dist or new_dist < 1e-6