code: # Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Test parallel Cirq simulations."""
  import numpy as np
  import tensorflow as tf
  from absl.testing import parameterized
  from scipy import stats
  import cirq
  
  from tensorflow_quantum.core.ops import batch_util
  from tensorflow_quantum.python import util
  
  BATCH_SIZE = 12
  N_QUBITS = 5
  PAULI_LENGTH = 3
  SYMBOLS = ['alpha', 'beta', 'gamma']
  
  
  def _get_mixed_batch(qubits, symbols, size):
      circuit1, resolver1 = util.random_circuit_resolver_batch(qubits, size // 2)
      circuit2, resolver2 = util.random_symbol_circuit_resolver_batch(
          qubits, symbols, size // 2)
      return circuit1 + circuit2, resolver1 + resolver2
  
  
  def _pad_state(sim, state, n):
      if isinstance(sim, cirq.sim.sparse_simulator.Simulator):
          state = state.final_state
      if isinstance(sim, cirq.DensityMatrixSimulator):
          state = state.final_density_matrix
      return np.pad(state, (0, (1 << n) - state.shape[-1]),
                    'constant',
                    constant_values=-2)
  
  
  def _expectation_helper(sim, circuit, params, op):
      if isinstance(sim, cirq.sim.sparse_simulator.Simulator):
          state = sim.simulate(circuit, params).final_state.astype(np.complex128)
          return [
              op.expectation_from_wavefunction(
                  state,
                  dict(
                      zip(sorted(circuit.all_qubits()),
                          (j for j in range(len(circuit.all_qubits())))))).real
          ]
      if isinstance(sim, cirq.DensityMatrixSimulator):
          state = sim.simulate(circuit, params).final_density_matrix
          return [
              sum(
                  x._expectation_from_density_matrix_no_validation(
                      state,
                      dict(
                          zip(sorted(circuit.all_qubits()), (
                              j
                              for j in range(len(circuit.all_qubits()))))))
                  for x in op)
          ]
  
      return NotImplemented
  
  
  def _sample_helper(sim, state, n_qubits, n_samples):
      if isinstance(sim, cirq.sim.sparse_simulator.Simulator):
          return cirq.sample_state_vector(state.final_state,
                                          list(range(n_qubits)),
                                          repetitions=n_samples)
      if isinstance(sim, cirq.DensityMatrixSimulator):
          return cirq.sample_density_matrix(state.final_density_matrix,
                                            list(range(n_qubits)),
                                            repetitions=n_samples)
  
      return NotImplemented
  
  
  class BatchUtilTest(tf.test.TestCase, parameterized.TestCase):
      """Test cases for BatchUtils main functions."""
  
      @parameterized.parameters([{
          'sim': cirq.DensityMatrixSimulator()
      }, {
          'sim': cirq.sim.sparse_simulator.Simulator()
      }])
      def test_batch_simulate_state(self, sim):
          """Test variable sized wavefunction output."""
          circuit_batch, resolver_batch = _get_mixed_batch(
              cirq.GridQubit.rect(1, N_QUBITS), SYMBOLS, BATCH_SIZE)
          results = batch_util.batch_calculate_state(circuit_batch,
                                                     resolver_batch, sim)
  
          for circuit, resolver, result in zip(circuit_batch, resolver_batch,
                                               results):
              r = _pad_state(sim, sim.simulate(circuit, resolver), N_QUBITS)
              self.assertAllClose(r, result, rtol=1e-5, atol=1e-5)
  
          self.assertDTypeEqual(results, np.complex64)
  
      @parameterized.parameters([{
          'sim': cirq.DensityMatrixSimulator()
      }, {
          'sim': cirq.sim.sparse_simulator.Simulator()
      }])
      def test_batch_expectation(self, sim):
          """Test expectation."""
          qubits = cirq.GridQubit.rect(1, N_QUBITS)
          circuit_batch, resolver_batch = _get_mixed_batch(
              qubits + [cirq.GridQubit(9, 9)], SYMBOLS, BATCH_SIZE)
          ops = util.random_pauli_sums(qubits, PAULI_LENGTH, BATCH_SIZE)
  
          results = batch_util.batch_calculate_expectation(
              circuit_batch, resolver_batch, [[x] for x in ops], sim)
  
          for circuit, resolver, result, op in zip(circuit_batch, resolver_batch,
                                                   results, ops):
              r = _expectation_helper(sim, circuit, resolver, op)
              self.assertAllClose(r, result, rtol=1e-5, atol=1e-5)
  
          self.assertDTypeEqual(results, np.float32)
  
      @parameterized.parameters([{
          'sim': cirq.DensityMatrixSimulator()
      }, {
          'sim': cirq.sim.sparse_simulator.Simulator()
      }])
      def test_batch_sampled_expectation(self, sim):
          """Test expectation."""
          qubits = cirq.GridQubit.rect(1, N_QUBITS)
          circuit_batch, resolver_batch = _get_mixed_batch(
              qubits + [cirq.GridQubit(9, 9)], SYMBOLS, BATCH_SIZE)
  
          ops = util.random_pauli_sums(qubits, PAULI_LENGTH, BATCH_SIZE)
          n_samples = [[1000] for _ in range(len(ops))]
  
          results = batch_util.batch_calculate_sampled_expectation(
              circuit_batch, resolver_batch, [[x] for x in ops], n_samples, sim)
  
          for circuit, resolver, result, op in zip(circuit_batch, resolver_batch,
                                                   results, ops):
              r = _expectation_helper(sim, circuit, resolver, op)
              self.assertAllClose(r, result, rtol=1.0, atol=1e-1)
  
          self.assertDTypeEqual(results, np.float32)
  
      @parameterized.parameters([{
          'sim': cirq.DensityMatrixSimulator()
      }, {
          'sim': cirq.sim.sparse_simulator.Simulator()
      }])
      def test_batch_sample_basic(self, sim):
          """Test sampling."""
          n_samples = 1
          n_qubits = 8
          qubits = cirq.GridQubit.rect(1, n_qubits)
          circuit = cirq.Circuit(*cirq.Z.on_each(*qubits[:n_qubits // 2]),
                                 *cirq.X.on_each(*qubits[n_qubits // 2:]))
  
          test_results = batch_util.batch_sample([circuit],
                                                 [cirq.ParamResolver({})],
                                                 n_samples, sim)
  
          state = sim.simulate(circuit, cirq.ParamResolver({}))
          expected_results = _sample_helper(sim, state, len(qubits), n_samples)
  
          self.assertAllEqual(expected_results, test_results[0])
          self.assertDTypeEqual(test_results, np.int32)
  
      @parameterized.parameters([{
          'sim': cirq.DensityMatrixSimulator()
      }, {
          'sim': cirq.sim.sparse_simulator.Simulator()
      }])
      def test_batch_sample(self, sim):
          """Test sampling."""
          n_samples = 2000 * (2**N_QUBITS)
  
          circuit_batch, resolver_batch = _get_mixed_batch(
              cirq.GridQubit.rect(1, N_QUBITS), SYMBOLS, BATCH_SIZE)
  
          results = batch_util.batch_sample(circuit_batch, resolver_batch,
                                            n_samples, sim)
  
          tfq_histograms = []
          for r in results:
              tfq_histograms.append(
                  np.histogram(r.dot(1 << np.arange(r.shape[-1] - 1, -1, -1)),
                               range=(0, 2**N_QUBITS),
                               bins=2**N_QUBITS)[0])
  
          cirq_histograms = []
          for circuit, resolver in zip(circuit_batch, resolver_batch):
              state = sim.simulate(circuit, resolver)
              r = _sample_helper(sim, state, len(circuit.all_qubits()), n_samples)
              cirq_histograms.append(
                  np.histogram(r.dot(1 << np.arange(r.shape[-1] - 1, -1, -1)),
                               range=(0, 2**N_QUBITS),
                               bins=2**N_QUBITS)[0])
  
          for a, b in zip(tfq_histograms, cirq_histograms):
              self.assertLess(stats.entropy(a + 1e-8, b + 1e-8), 0.005)
  
          self.assertDTypeEqual(results, np.int32)
  
      @parameterized.parameters([{
          'sim': cirq.DensityMatrixSimulator()
      }, {
          'sim': cirq.sim.sparse_simulator.Simulator()
      }])
      def test_empty_circuits(self, sim):
          """Test functions with empty circuits."""
          # Common preparation
          resolver_batch = [cirq.ParamResolver({}) for _ in range(BATCH_SIZE)]
          circuit_batch = [cirq.Circuit() for _ in range(BATCH_SIZE)]
          qubits = cirq.GridQubit.rect(1, N_QUBITS)
          ops = util.random_pauli_sums(qubits, PAULI_LENGTH, BATCH_SIZE)
          n_samples = [[1000] for _ in range(len(ops))]
          # If there is no op on a qubit, the expectation answer is -2.0
          true_expectation = (-2.0,)
  
          # (1) Test expectation
          results = batch_util.batch_calculate_expectation(
              circuit_batch, resolver_batch, [[x] for x in ops], sim)
  
          for _, _, result, _ in zip(circuit_batch, resolver_batch, results, ops):
              self.assertAllClose(true_expectation, result, rtol=1e-5, atol=1e-5)
  
          self.assertDTypeEqual(results, np.float32)
  
          # (2) Test sampled_expectation
          results = batch_util.batch_calculate_sampled_expectation(
              circuit_batch, resolver_batch, [[x] for x in ops], n_samples, sim)
  
          for _, _, result, _ in zip(circuit_batch, resolver_batch, results, ops):
              self.assertAllClose(true_expectation, result, rtol=1.0, atol=1e-1)
  
          self.assertDTypeEqual(results, np.float32)
  
          # (3) Test state
          results = batch_util.batch_calculate_state(circuit_batch,
                                                     resolver_batch, sim)
  
          for circuit, resolver, result in zip(circuit_batch, resolver_batch,
                                               results):
              r = _pad_state(sim, sim.simulate(circuit, resolver), 0)
              self.assertAllClose(r, result, rtol=1e-5, atol=1e-5)
  
          self.assertDTypeEqual(results, np.complex64)
  
          # (4) Test sampling
          n_samples = 2000 * (2**N_QUBITS)
          results = batch_util.batch_sample(circuit_batch, resolver_batch,
                                            n_samples, sim)
  
          for circuit, resolver, a in zip(circuit_batch, resolver_batch, results):
              state = sim.simulate(circuit, resolver)
              r = _sample_helper(sim, state, len(circuit.all_qubits()), n_samples)
              self.assertAllClose(r, a, atol=1e-5)
  
          self.assertDTypeEqual(results, np.int32)
  
  
  if __name__ == '__main__':
      tf.test.main()
  
apis: numpy.arange numpy.pad scipy.stats.entropy tensorflow.test.main