1import numpy as np
2import pytest
3import cirq
4import examples.direct_fidelity_estimation as dfe
5
6
7def test_direct_fidelity_estimation_no_noise_clifford():
8    qubits = cirq.LineQubit.range(3)
9    circuit = cirq.Circuit(cirq.Z(qubits[0]), cirq.X(qubits[1]), cirq.X(qubits[2]))
10
11    no_noise_simulator = cirq.DensityMatrixSimulator()
12
13    estimated_fidelity, _ = dfe.direct_fidelity_estimation(
14        circuit, qubits, no_noise_simulator, n_measured_operators=3, samples_per_term=0
15    )
16    assert np.isclose(estimated_fidelity, 1.0, atol=0.01)
17
18
19def test_direct_fidelity_estimation_no_noise_non_clifford():
20    qubits = cirq.LineQubit.range(3)
21    circuit = cirq.Circuit(cirq.Z(qubits[0]) ** 0.123, cirq.X(qubits[1]), cirq.X(qubits[2]))
22
23    no_noise_simulator = cirq.DensityMatrixSimulator()
24
25    estimated_fidelity, _ = dfe.direct_fidelity_estimation(
26        circuit, qubits, no_noise_simulator, n_measured_operators=64, samples_per_term=0
27    )
28    assert np.isclose(estimated_fidelity, 1.0, atol=0.01)
29
30
31def test_direct_fidelity_estimation_with_noise_clifford():
32    qubits = cirq.LineQubit.range(3)
33    circuit = cirq.Circuit(cirq.Z(qubits[0]), cirq.X(qubits[1]), cirq.X(qubits[2]))
34
35    noise = cirq.ConstantQubitNoiseModel(cirq.depolarize(0.1))
36    noisy_simulator = cirq.DensityMatrixSimulator(noise=noise)
37
38    estimated_fidelity, _ = dfe.direct_fidelity_estimation(
39        circuit, qubits, noisy_simulator, n_measured_operators=None, samples_per_term=100
40    )
41    assert estimated_fidelity >= -1.0 and estimated_fidelity <= 1.0
42
43
44def test_direct_fidelity_estimation_with_noise_non_clifford():
45    qubits = cirq.LineQubit.range(3)
46    circuit = cirq.Circuit(
47        cirq.Z(qubits[0]) ** 0.25,  # T-Gate, non Clifford.
48        cirq.X(qubits[1]) ** 0.123,
49        cirq.X(qubits[2]) ** 0.456,
50    )
51
52    noise = cirq.ConstantQubitNoiseModel(cirq.depolarize(0.1))
53    noisy_simulator = cirq.DensityMatrixSimulator(noise=noise)
54
55    estimated_fidelity, _ = dfe.direct_fidelity_estimation(
56        circuit, qubits, noisy_simulator, n_measured_operators=None, samples_per_term=100
57    )
58    assert estimated_fidelity >= -1.0 and estimated_fidelity <= 1.0
59
60
61def test_incorrect_sampler_raises_exception():
62    qubits = cirq.LineQubit.range(1)
63    circuit = cirq.Circuit(cirq.X(qubits[0]))
64    sampler_incorrect_type = cirq.ZerosSampler
65
66    with pytest.raises(TypeError):
67        dfe.direct_fidelity_estimation(
68            circuit, qubits, sampler_incorrect_type, n_measured_operators=3, samples_per_term=0
69        )
70
71
72def test_direct_fidelity_estimation_clifford_all_trials():
73    qubits = cirq.LineQubit.range(2)
74    circuit = cirq.Circuit(cirq.Z(qubits[0]), cirq.X(qubits[1]))
75
76    no_noise_simulator = cirq.DensityMatrixSimulator()
77
78    for n_measured_operators in [1, 2, 3, 4, None]:
79        estimated_fidelity, _ = dfe.direct_fidelity_estimation(
80            circuit,
81            qubits,
82            no_noise_simulator,
83            n_measured_operators=n_measured_operators,
84            samples_per_term=0,
85        )
86        assert np.isclose(estimated_fidelity, 1.0, atol=0.01)
87
88
89def test_same_pauli_traces_clifford():
90    n_qubits = 4
91
92    qubits = cirq.LineQubit.range(n_qubits)
93    circuit_clifford = cirq.Circuit(
94        cirq.X(qubits[3]),
95    )
96
97    circuit_general = cirq.Circuit(
98        cirq.CCX(qubits[0], qubits[1], qubits[2]),
99        circuit_clifford,
100    )
101
102    def _run_dfe(circuit):
103        class NoiseOnLastQubitOnly(cirq.NoiseModel):
104            def __init__(self):
105                self.qubit_noise_gate = cirq.amplitude_damp(1.0)
106
107            def noisy_moment(self, moment, system_qubits):
108                return [
109                    moment,
110                    cirq.ops.Moment(
111                        [
112                            self.qubit_noise_gate(q).with_tags(cirq.ops.VirtualTag())
113                            for q in system_qubits[-1:]
114                        ]
115                    ),
116                ]
117
118        np.random.seed(0)
119        noise = NoiseOnLastQubitOnly()
120        noisy_simulator = cirq.DensityMatrixSimulator(noise=noise)
121
122        _, intermediate_results = dfe.direct_fidelity_estimation(
123            circuit, qubits, noisy_simulator, n_measured_operators=None, samples_per_term=1
124        )
125        return intermediate_results.pauli_traces, intermediate_results.clifford_tableau is not None
126
127    # Run both algos
128    pauli_traces_clifford, clifford_is_clifford = _run_dfe(circuit_clifford)
129    pauli_traces_general, general_is_clifford = _run_dfe(circuit_general)
130
131    assert clifford_is_clifford
132    assert not general_is_clifford
133
134    assert len(pauli_traces_clifford) == 2 ** n_qubits
135    for pauli_trace_clifford in pauli_traces_clifford:
136        pauli_trace_general = [x for x in pauli_traces_general if x.P_i == pauli_trace_clifford.P_i]
137        assert len(pauli_trace_general) == 1
138        pauli_trace_general = pauli_trace_general[0]
139
140        # The code itself checks that the rho_i is either +1 or -1, so here we
141        # simply test that the sign is the same.
142        assert np.isclose(pauli_trace_general.rho_i, pauli_trace_clifford.rho_i, atol=0.01)
143
144
145def test_direct_fidelity_estimation_intermediate_results():
146    qubits = cirq.LineQubit.range(1)
147    circuit = cirq.Circuit(cirq.I(qubits[0]))
148    no_noise_simulator = cirq.DensityMatrixSimulator()
149
150    _, intermediate_result = dfe.direct_fidelity_estimation(
151        circuit, qubits, no_noise_simulator, n_measured_operators=1, samples_per_term=0
152    )
153    # We only test a few fields to be sure that they are set properly. In
154    # particular, some of them are random, and so we don't test them.
155    assert str(intermediate_result.clifford_tableau) == "+ Z "
156
157    np.testing.assert_equal(len(intermediate_result.pauli_traces), 1)
158    assert np.isclose(intermediate_result.pauli_traces[0].rho_i, 1.0)
159    assert np.isclose(intermediate_result.pauli_traces[0].Pr_i, 0.5)
160
161    np.testing.assert_equal(len(intermediate_result.trial_results), 1)
162    assert np.isclose(intermediate_result.trial_results[0].sigma_i, 1.0)
163
164    assert np.isclose(intermediate_result.std_dev_estimate, 0.0)
165    assert np.isclose(intermediate_result.std_dev_bound, 0.5)
166
167
168def test_parsing_args():
169    dfe.parse_arguments(['--samples_per_term=10'])
170
171
172def test_calling_main():
173    dfe.main(n_measured_operators=3, samples_per_term=0)
174    dfe.main(n_measured_operators=3, samples_per_term=10)
175