text: """Implements direct fidelity estimation.
  
  Fidelity between the desired pure state rho and the actual state sigma is
  defined as:
  F(rho, sigma) = Tr (rho sigma)
  
  It is a unit-less measurement between 0.0 and 1.0. The following two papers
  independently described a faster way to estimate its value:
  
  Direct Fidelity Estimation from Few Pauli Measurements
  https://arxiv.org/abs/1104.4695
  
  Practical characterization of quantum devices without tomography
  https://arxiv.org/abs/1104.3835
  
  This code implements the algorithm proposed for an example circuit (defined in
  the function build_circuit()) and a noise (defines in the variable noise).
  """
  
  import argparse
  import asyncio
  import itertools
  from typing import cast
  from typing import List
  from typing import Optional
  from typing import Tuple
  import sys
  import numpy as np
  import cirq
  
  
  def build_circuit() -> Tuple[cirq.Circuit, List[cirq.Qid]]:
      # Builds an arbitrary circuit to test. Do not include a measurement gate.
      # The circuit need not be Clifford, but if it is, simulations will be
      # faster.
      qubits: List[cirq.Qid] = cast(List[cirq.Qid], cirq.LineQubit.range(3))
      circuit: cirq.Circuit = cirq.Circuit(cirq.CNOT(qubits[0], qubits[2]),
                                           cirq.Z(qubits[0]), cirq.H(qubits[2]),
                                           cirq.CNOT(qubits[2], qubits[1]),
                                           cirq.X(qubits[0]), cirq.X(qubits[1]),
                                           cirq.CNOT(qubits[0], qubits[2]))
      print('Circuit used:')
      print(circuit)
      return circuit, qubits
  
  
  def compute_characteristic_function(circuit: cirq.Circuit,
                                      pauli_string: cirq.PauliString,
                                      qubits: List[cirq.Qid],
                                      density_matrix: np.ndarray):
      n_qubits = len(qubits)
      d = 2**n_qubits
  
      qubit_map = dict(zip(qubits, range(n_qubits)))
      # rho_i or sigma_i in https://arxiv.org/abs/1104.3835
      trace = pauli_string.expectation_from_density_matrix(
          density_matrix, qubit_map)
      assert np.isclose(trace.imag, 0.0, atol=1e-6)
      trace = trace.real
  
      prob = trace * trace / d  # Pr(i) in https://arxiv.org/abs/1104.3835
  
      return trace, prob
  
  
  async def estimate_characteristic_function(
          circuit: cirq.Circuit, pauli_string: cirq.PauliString,
          qubits: List[cirq.Qid], simulator: cirq.DensityMatrixSimulator,
          samples_per_term: int):
      """
      Estimates the characteristic function using a (noisy) circuit simulator by
      sampling the results.
  
      Args:
          circuit: The circuit to run the simulation on.
          pauli_string: The Pauli string.
          qubits: The list of qubits.
          simulator: The (noisy) simulator.
          samples_per_term: An integer greater than 0, the number of samples.
  
      Returns:
          The estimated characteristic function.
      """
      p = cirq.PauliSumCollector(circuit=circuit,
                                 observable=pauli_string,
                                 samples_per_term=samples_per_term)
  
      await p.collect_async(sampler=simulator)
  
      sigma_i = p.estimated_energy()
      assert np.isclose(sigma_i.imag, 0.0, atol=1e-6)
      sigma_i = sigma_i.real
  
      return sigma_i
  
  
  def _estimate_pauli_traces_clifford(n_qubits: int,
                                      clifford_state: cirq.CliffordState,
                                      n_clifford_trials: int):
      """
      Estimates the Pauli traces in case the circuit is Clifford. When we have a
      Clifford circuit, there are 2**n Pauli traces that have probability 1/2**n
      and all the other traces have probability 0. In addition, there is a fast
      way to compute find out what the traces are. See the documentation of
      cirq.CliffordState for more detail. This function uses the speedup to sample
      the Pauli states with non-zero probability.
  
      Args:
          n_qubits: An integer that is the number of qubits.
          clifford_state: The basis of the Pauli states with non-zero probability.
          n_clifford_trials: An integer that is the number of Pauli states to
              sample.
  
      Returns:
          A list of Pauli states (represented as tuples of Pauli string, rho_i,
              and probability.
      """
  
      # When the circuit consists of Clifford gates only, we can sample the
      # Pauli states more efficiently as described on page 4 of:
      # https://arxiv.org/abs/1104.4695
  
      d = 2**n_qubits
  
      # The stabilizers_basis variable only contains basis vectors. For
      # example, if we have n=3 qubits, then we should have 2**n=8 Pauli
      # states that we can sample, but the basis will still have 3 entries. We
      # must flip a coin for each, whether or not to include them.
      stabilizer_basis = clifford_state.stabilizers()
  
      pauli_traces = []
      for _ in range(n_clifford_trials):
          # Build the Pauli string as a random sample of the basis elements.
          dense_pauli_string = cirq.DensePauliString.eye(n_qubits)
          for stabilizer in stabilizer_basis:
              if np.random.randint(2) == 1:
                  dense_pauli_string *= stabilizer
  
          # The code below is equivalent to calling
          # clifford_state.wave_function() and then calling
          # compute_characteristic_function() on the results (albeit with a
          # wave function instead of a density matrix). It is, however,
          # unncessary to do so. Instead we directly obtain the scalar rho_i.
          rho_i = dense_pauli_string.coefficient
  
          assert np.isclose(rho_i.imag, 0.0, atol=1e-6)
          rho_i = rho_i.real
  
          dense_pauli_string *= rho_i
  
          assert np.isclose(abs(rho_i), 1.0, atol=1e-6)
          Pr_i = 1.0 / d
  
          pauli_traces.append({
              'P_i': dense_pauli_string.sparse(),
              'rho_i': rho_i,
              'Pr_i': Pr_i
          })
      return pauli_traces
  
  
  def _estimate_pauli_traces_general(qubits: List[cirq.Qid],
                                     circuit: cirq.Circuit):
      """
      Estimates the Pauli traces in case the circuit is not Clifford. In this case
      we cannot use the speedup implemented in the function
      _estimate_pauli_traces_clifford() above, and so do a slow, density matrix
      simulation.
  
      Args:
          qubits: The list of qubits.
          circuit: The (non Clifford) circuit.
  
      Returns:
          A list of Pauli states (represented as tuples of Pauli string, rho_i,
              and probability.
      """
  
      n_qubits = len(qubits)
  
      dense_simulator = cirq.DensityMatrixSimulator()
      # rho in https://arxiv.org/abs/1104.3835
      clean_density_matrix = cast(
          cirq.DensityMatrixTrialResult,
          dense_simulator.simulate(circuit)).final_density_matrix
  
      pauli_traces = []
      for P_i in itertools.product([cirq.I, cirq.X, cirq.Y, cirq.Z],
                                   repeat=n_qubits):
          pauli_string = cirq.PauliString(dict(zip(qubits, P_i)))
          rho_i, Pr_i = compute_characteristic_function(circuit, pauli_string,
                                                        qubits,
                                                        clean_density_matrix)
          pauli_traces.append({'P_i': pauli_string, 'rho_i': rho_i, 'Pr_i': Pr_i})
      return pauli_traces
  
  
  def direct_fidelity_estimation(circuit: cirq.Circuit, qubits: List[cirq.Qid],
                                 noise: cirq.NoiseModel, n_trials: int,
                                 n_clifford_trials: int, samples_per_term: int):
      """
      Implementation of direct fidelity estimation, as per 'Direct Fidelity
      Estimation from Few Pauli Measurements' https://arxiv.org/abs/1104.4695 and
      'Practical characterization of quantum devices without tomography'
      https://arxiv.org/abs/1104.3835.
  
      Args:
          circuit: The circuit to run the simulation on.
          qubits: The list of qubits.
          noise: The noise model when doing a simulation.
          n_trial: The total number of Pauli measurements.
          n_clifford_trials: In case the circuit is Clifford, we specify the
              number of trials to estimate the noise-free pauli traces.
          samples_per_term: is set to 0, we use the 'noise' parameter above and
              simulate noise in the circuit. If greater than 0, we ignore the
              'noise' parameter above and instead run an estimation of the
              characteristic function.
  
      Returns:
          The estimated fidelity.
      """
      # n_trials is upper-case N in https://arxiv.org/abs/1104.3835
  
      # Number of qubits, lower-case n in https://arxiv.org/abs/1104.3835
      n_qubits = len(qubits)
      d = 2**n_qubits
  
      clifford_circuit = True
      clifford_state: Optional[cirq.CliffordState] = None
      try:
          clifford_state = cirq.CliffordState(
              qubit_map={qubits[i]: i for i in range(len(qubits))})
          for gate in circuit.all_operations():
              clifford_state.apply_unitary(gate)
      except ValueError:
          clifford_circuit = False
  
      # Computes for every \hat{P_i} of https://arxiv.org/abs/1104.3835
      # estimate rho_i and Pr(i). We then collect tuples (rho_i, Pr(i), \hat{Pi})
      # inside the variable 'pauli_traces'.
      if clifford_circuit:
          print('Circuit is Clifford')
          assert clifford_state is not None
          pauli_traces = _estimate_pauli_traces_clifford(
              n_qubits, cast(cirq.CliffordState, clifford_state),
              n_clifford_trials)
      else:
          print('Circuit is not Clifford')
          pauli_traces = _estimate_pauli_traces_general(qubits, circuit)
  
      p = np.asarray([x['Pr_i'] for x in pauli_traces])
  
      if not clifford_circuit:
          # For Clifford circuits, we do a Monte Carlo simulations, and thus there
          # is no guarantee that it adds up to 1.0 (but it should to the limit).
          assert np.isclose(np.sum(p), 1.0, atol=1e-6)
  
      # The package np.random.choice() is quite sensitive to probabilities not
      # summing up to 1.0. Even an absolute difference below 1e-6 (as checked just
      # above) does bother it, so we re-normalize the probs.
      p /= np.sum(p)
  
      noisy_simulator = cirq.DensityMatrixSimulator(noise=noise)
      fidelity = 0.0
  
      if samples_per_term == 0:
          # sigma in https://arxiv.org/abs/1104.3835
          noisy_density_matrix = cast(
              cirq.DensityMatrixTrialResult,
              noisy_simulator.simulate(circuit)).final_density_matrix
  
      for _ in range(n_trials):
          # Randomly sample as per probability.
          i = np.random.choice(len(pauli_traces), p=p)
  
          Pr_i = pauli_traces[i]['Pr_i']
          measure_pauli_string: cirq.PauliString = pauli_traces[i]['P_i']
          rho_i = pauli_traces[i]['rho_i']
  
          if samples_per_term > 0:
              sigma_i = asyncio.get_event_loop().run_until_complete(
                  estimate_characteristic_function(circuit, measure_pauli_string,
                                                   qubits, noisy_simulator,
                                                   samples_per_term))
          else:
              sigma_i, _ = compute_characteristic_function(
                  circuit, measure_pauli_string, qubits, noisy_density_matrix)
  
          fidelity += Pr_i * sigma_i / rho_i
  
      return fidelity / n_trials * d
  
  
  def parse_arguments(args):
      """Helper function that parses the given arguments."""
      parser = argparse.ArgumentParser('Direct fidelity estimation.')
  
      parser.add_argument('--n_trials',
                          default=10,
                          type=int,
                          help='Number of trials to run.')
  
      # TODO(#2802): Offer some guidance on how to set this flag. Maybe have an
      # option to do an exhaustive sample and do numerical studies to know which
      # choice is the best.
      parser.add_argument('--n_clifford_trials',
                          default=3,
                          type=int,
                          help='Number of trials for Clifford circuits. This is '
                          'in effect when the circuit is Clifford. In this '
                          'case, we randomly sample the Pauli traces with '
                          'non-zero probabilities. The higher the number, '
                          'the more accurate the overall fidelity '
                          'estimation, at the cost of extra computing and '
                          'measurements.')
  
      parser.add_argument('--samples_per_term',
                          default=0,
                          type=int,
                          help='Number of samples per trial or 0 if no sampling.')
  
      return vars(parser.parse_args(args))
  
  
  def main(*, n_trials: int, n_clifford_trials: int, samples_per_term: int):
      circuit, qubits = build_circuit()
  
      noise = cirq.ConstantQubitNoiseModel(cirq.depolarize(0.1))
      print('Noise model: %s' % (noise))
  
      estimated_fidelity = direct_fidelity_estimation(
          circuit,
          qubits,
          noise,
          n_trials=n_trials,
          n_clifford_trials=n_clifford_trials,
          samples_per_term=samples_per_term)
      print('Estimated fidelity: %f' % (estimated_fidelity))
  
  
  if __name__ == '__main__':
      main(**parse_arguments(sys.argv[1:]))
  