void ComputeLarge()

in tensorflow_quantum/core/ops/noise/tfq_noisy_expectation.cc [155:250]
80 lines of code
14 McCabe index (conditional complexity)

  void ComputeLarge(const std::vector<int>& num_qubits,
                    const std::vector<NoisyQsimCircuit>& ncircuits,
                    const std::vector<std::vector<PauliSum>>& pauli_sums,
                    const std::vector<std::vector<int>>& num_samples,
                    tensorflow::OpKernelContext* context,
                    tensorflow::TTypes<float, 1>::Matrix* output_tensor) {
    // Instantiate qsim objects.
    const auto tfq_for = tfq::QsimFor(context);
    using Simulator = qsim::Simulator<const tfq::QsimFor&>;
    using StateSpace = Simulator::StateSpace;
    using QTSimulator =
        qsim::QuantumTrajectorySimulator<qsim::IO, QsimGate,
                                         qsim::MultiQubitGateFuser, Simulator>;

    // Begin simulation.
    int largest_nq = 1;
    Simulator sim = Simulator(tfq_for);
    StateSpace ss = StateSpace(tfq_for);
    auto sv = ss.Create(largest_nq);
    auto scratch = ss.Create(largest_nq);

    tensorflow::GuardedPhiloxRandom random_gen;
    int max_n_shots = 1;
    for (int i = 0; i < num_samples.size(); i++) {
      for (int j = 0; j < num_samples[i].size(); j++) {
        max_n_shots = std::max(max_n_shots, num_samples[i][j]);
      }
    }
    random_gen.Init(tensorflow::random::New64(), tensorflow::random::New64());
    auto local_gen =
        random_gen.ReserveSamples128(ncircuits.size() * max_n_shots + 1);
    tensorflow::random::SimplePhilox rand_source(&local_gen);

    // Simulate programs one by one. Parallelizing over state vectors
    // we no longer parallelize over circuits. Each time we encounter a
    // a larger circuit we will grow the Statevector as necessary.
    for (int i = 0; i < ncircuits.size(); i++) {
      int nq = num_qubits[i];

      // (#679) Just ignore empty program
      if (ncircuits[i].channels.size() == 0) {
        for (int j = 0; j < pauli_sums[i].size(); j++) {
          (*output_tensor)(i, j) = -2.0;
        }
        continue;
      }

      if (nq > largest_nq) {
        largest_nq = nq;
        sv = ss.Create(largest_nq);
        scratch = ss.Create(largest_nq);
      }
      QTSimulator::Parameter param;
      param.collect_kop_stat = false;
      param.collect_mea_stat = false;
      param.normalize_before_mea_gates = true;
      std::vector<uint64_t> unused_stats;
      // Track op-wise stats.
      std::vector<int> run_samples(num_samples[i].size(), 0);
      std::vector<double> rolling_sums(num_samples[i].size(), 0.0);

      while (1) {
        ss.SetStateZero(sv);

        QTSimulator::RunOnce(param, ncircuits[i], rand_source.Rand64(), ss, sim,
                             scratch, sv, unused_stats);

        // Use this trajectory as a source for all expectation calculations.
        for (int j = 0; j < pauli_sums[i].size(); j++) {
          if (run_samples[j] >= num_samples[i][j]) {
            continue;
          }
          float exp_v = 0.0;
          OP_REQUIRES_OK(context,
                         ComputeExpectationQsim(pauli_sums[i][j], sim, ss, sv,
                                                scratch, &exp_v));
          rolling_sums[j] += static_cast<double>(exp_v);
          run_samples[j]++;
        }
        bool break_loop = true;
        for (int j = 0; j < num_samples[i].size(); j++) {
          if (run_samples[j] < num_samples[i][j]) {
            break_loop = false;
            break;
          }
        }
        if (break_loop) {
          for (int j = 0; j < num_samples[i].size(); j++) {
            rolling_sums[j] /= num_samples[i][j];
            (*output_tensor)(i, j) = static_cast<float>(rolling_sums[j]);
          }
          break;
        }
      }
    }
  }