1# Copyright 2019 The Cirq Developers
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14from typing import Sequence
15
16import numpy as np
17import pytest
18
19import cirq
20
21Q0, Q1, Q2, Q3 = cirq.LineQubit.range(4)
22
23
24def test_state_tomography_diagonal():
25    n = 2
26    qubits = cirq.LineQubit.range(n)
27    for state in range(2 ** n):
28        circuit = cirq.Circuit()
29        for i, q in enumerate(qubits):
30            bit = state & (1 << (n - i - 1))
31            if bit:
32                circuit.append(cirq.X(q))
33        res = cirq.experiments.state_tomography(
34            cirq.Simulator(seed=87539319),
35            qubits,
36            circuit,
37            repetitions=1000,
38            prerotations=[(0, 0), (0, 0.5), (0.5, 0.5)],
39        )
40        should_be = np.zeros((2 ** n, 2 ** n))
41        should_be[state, state] = 1
42        assert np.allclose(res.data, should_be, atol=0.05)
43
44
45def test_state_tomography_ghz_state():
46    circuit = cirq.Circuit()
47    circuit.append(cirq.H(cirq.LineQubit(0)))
48    circuit.append(cirq.CNOT(cirq.LineQubit(0), cirq.LineQubit(1)))
49    circuit.append(cirq.CNOT(cirq.LineQubit(0), cirq.LineQubit(2)))
50    res = cirq.experiments.state_tomography(
51        cirq.Simulator(seed=87539319),
52        [cirq.LineQubit(0), cirq.LineQubit(1), cirq.LineQubit(2)],
53        circuit,
54        repetitions=1000,
55    )
56    should_be = np.zeros((8, 8))
57    should_be[0, 0] = 0.5
58    should_be[7, 7] = 0.5
59    should_be[0, 7] = 0.5
60    should_be[7, 0] = 0.5
61    assert np.allclose(res.data, should_be, atol=0.05)
62
63
64def test_make_experiment_no_rots():
65    exp = cirq.experiments.StateTomographyExperiment(
66        [cirq.LineQubit(0), cirq.LineQubit(1), cirq.LineQubit(2)]
67    )
68    assert len(exp.rot_sweep) > 0
69
70
71def compute_density_matrix(circuit: cirq.Circuit, qubits: Sequence[cirq.Qid]) -> np.ndarray:
72    """Computes density matrix prepared by circuit based on its unitary."""
73    u = circuit.unitary(qubit_order=qubits)
74    phi = u[:, 0]
75    rho = np.outer(phi, np.conjugate(phi))
76    return rho
77
78
79@pytest.mark.parametrize(
80    'circuit, qubits',
81    (
82        (cirq.Circuit(cirq.X(Q0) ** 0.25), (Q0,)),
83        (cirq.Circuit(cirq.CNOT(Q0, Q1) ** 0.25), (Q0, Q1)),
84        (cirq.Circuit(cirq.CNOT(Q0, Q1) ** 0.25), (Q1, Q0)),
85        (cirq.Circuit(cirq.TOFFOLI(Q0, Q1, Q2)), (Q1, Q0, Q2)),
86        (
87            cirq.Circuit(
88                cirq.H(Q0),
89                cirq.H(Q1),
90                cirq.CNOT(Q0, Q2),
91                cirq.CNOT(Q1, Q3),
92                cirq.X(Q0),
93                cirq.X(Q1),
94                cirq.CNOT(Q1, Q0),
95            ),
96            (Q1, Q0, Q2, Q3),
97        ),
98    ),
99)
100def test_density_matrix_from_state_tomography_is_correct(circuit, qubits):
101    sim = cirq.Simulator(seed=87539319)
102    tomography_result = cirq.experiments.state_tomography(sim, qubits, circuit, repetitions=5000)
103    actual_rho = tomography_result.data
104    expected_rho = compute_density_matrix(circuit, qubits)
105    error_rho = actual_rho - expected_rho
106    assert np.linalg.norm(error_rho) < 0.05
107    assert np.max(np.abs(error_rho)) < 0.05
108
109
110@pytest.mark.parametrize(
111    'circuit',
112    (
113        cirq.Circuit(cirq.CNOT(Q0, Q1) ** 0.3),
114        cirq.Circuit(cirq.H(Q0), cirq.CNOT(Q0, Q1)),
115        cirq.Circuit(cirq.X(Q0) ** 0.25, cirq.ISWAP(Q0, Q1)),
116    ),
117)
118def test_agrees_with_two_qubit_state_tomography(circuit):
119    qubits = (Q0, Q1)
120    sim = cirq.Simulator(seed=87539319)
121    tomography_result = cirq.experiments.state_tomography(sim, qubits, circuit, repetitions=5000)
122    actual_rho = tomography_result.data
123
124    two_qubit_tomography_result = cirq.experiments.two_qubit_state_tomography(
125        sim, qubits[0], qubits[1], circuit, repetitions=5000
126    )
127    expected_rho = two_qubit_tomography_result.data
128
129    error_rho = actual_rho - expected_rho
130
131    assert np.linalg.norm(error_rho) < 0.06
132    assert np.max(np.abs(error_rho)) < 0.05
133