Merge pull request #504 from quantumlib/qsim-repeat-samples

Use qsim for repeated sampling

Author: 95-martin-orion

pybind_interface/pybind_main.cpp

@@ -525,6 +525,15 @@ class SimulatorHelper {
     return helper.release_state_to_python();
   }
 
+  static std::vector<uint64_t> sample_final_state(
+      const py::dict &options, bool is_noisy, uint64_t num_samples) {
+    auto helper = SimulatorHelper(options, is_noisy);
+    if (!helper.is_valid || !helper.simulate(0)) {
+      return {};
+    }
+    return helper.sample(num_samples);
+  }
+
   template <typename StateType>
   static std::vector<std::complex<double>> simulate_expectation_values(
       const py::dict &options,
@@ -753,6 +762,11 @@ class SimulatorHelper {
     return result;
   }
 
+  std::vector<uint64_t> sample(uint64_t num_samples) {
+    StateSpace state_space = factory.CreateStateSpace();
+    return state_space.Sample(state, num_samples, seed);
+  }
+
   py::array_t<float> release_state_to_python() {
     StateSpace state_space = factory.CreateStateSpace();
     state_space.InternalToNormalOrder(state);
@@ -931,6 +945,16 @@ qtrajectory_simulate_moment_expectation_values(
 
 // Methods for sampling.
 
+std::vector<uint64_t> qsim_sample_final(
+    const py::dict &options, uint64_t num_samples) {
+  return SimulatorHelper::sample_final_state(options, false, num_samples);
+}
+
+std::vector<uint64_t> qtrajectory_sample_final(
+    const py::dict &options, uint64_t num_samples) {
+  return SimulatorHelper::sample_final_state(options, true, num_samples);
+}
+
 std::vector<unsigned> qsim_sample(const py::dict &options) {
   Circuit<Cirq::GateCirq<float>> circuit;
   try {

pybind_interface/pybind_main.h

@@ -90,6 +90,8 @@ py::array_t<float> qsim_simulate_fullstate(
       const py::dict &options, const py::array_t<float> &input_vector);
 
 std::vector<unsigned> qsim_sample(const py::dict &options);
+std::vector<uint64_t> qsim_sample_final(
+  const py::dict &options, uint64_t num_samples);
 
 // Methods for simulating noisy circuits.
 std::vector<std::complex<float>> qtrajectory_simulate(const py::dict &options);
@@ -100,6 +102,8 @@ py::array_t<float> qtrajectory_simulate_fullstate(
       const py::dict &options, const py::array_t<float> &input_vector);
 
 std::vector<unsigned> qtrajectory_sample(const py::dict &options);
+std::vector<uint64_t> qtrajectory_sample_final(
+  const py::dict &options, uint64_t num_samples);
 
 // As above, but returning expectation values instead.
 std::vector<std::complex<double>> qsim_simulate_expectation_values(
@@ -200,8 +204,12 @@ std::vector<std::complex<float>> qsimh_simulate(const py::dict &options);
                                                                                       \
       /* Methods for returning samples */                                             \
       m.def("qsim_sample", &qsim_sample, "Call the qsim sampler");                    \
+      m.def("qsim_sample_final", &qsim_sample_final,                                  \
+            "Call the qsim final-state sampler");                                     \
       m.def("qtrajectory_sample", &qtrajectory_sample,                                \
             "Call the qtrajectory sampler");                                          \
+      m.def("qtrajectory_sample_final", &qtrajectory_sample_final,                    \
+            "Call the qtrajectory final-state sampler");                              \
                                                                                       \
       using GateCirq = qsim::Cirq::GateCirq<float>;                                   \
       using OpString = qsim::OpString<GateCirq>;                                      \
@@ -400,8 +408,12 @@ std::vector<std::complex<float>> qsimh_simulate(const py::dict &options);
                                                                                       \
       /* Methods for returning samples */                                             \
       m.def("qsim_sample", &qsim_sample, "Call the qsim sampler");                    \
+      m.def("qsim_sample_final", &qsim_sample_final,                                  \
+            "Call the qsim final-state sampler");                                     \
       m.def("qtrajectory_sample", &qtrajectory_sample,                                \
             "Call the qtrajectory sampler");                                          \
+      m.def("qtrajectory_sample_final", &qtrajectory_sample_final,                    \
+            "Call the qtrajectory final-state sampler");                              \
                                                                                       \
       using GateCirq = qsim::Cirq::GateCirq<float>;                                   \
       using OpString = qsim::OpString<GateCirq>;                                      \

qsimcirq/qsim_simulator.py

@@ -349,19 +349,7 @@ def _sample_measure_results(
             )
 
         noisy = _needs_trajectories(program)
-        if noisy:
-            translator_fn_name = "translate_cirq_to_qtrajectory"
-            sampler_fn = self._sim_module.qtrajectory_sample
-        else:
-            translator_fn_name = "translate_cirq_to_qsim"
-            sampler_fn = self._sim_module.qsim_sample
-
-        if (
-            not noisy
-            and program.are_all_measurements_terminal()
-            and repetitions > 1
-            and num_qubits <= 32  # max length of ndarray.shape
-        ):
+        if not noisy and program.are_all_measurements_terminal() and repetitions > 1:
             # Measurements must be replaced with identity gates to sample properly.
             # Simply removing them may omit qubits from the circuit.
             for i in range(len(program.moments)):
@@ -378,12 +366,12 @@ def _sample_measure_results(
                 cirq.QubitOrder.DEFAULT,
             )
             options["s"] = self.get_seed()
-            final_state = self._sim_module.qsim_simulate_fullstate(options, 0)
-            full_results = cirq.sample_state_vector(
-                final_state.view(np.complex64),
-                range(num_qubits),
-                repetitions=repetitions,
-                seed=self._prng,
+            raw_results = self._sim_module.qsim_sample_final(options, repetitions)
+            full_results = np.array(
+                [
+                    [bool(result & (1 << q)) for q in reversed(range(num_qubits))]
+                    for result in raw_results
+                ]
             )
 
             for key, op in meas_ops.items():
@@ -393,6 +381,13 @@ def _sample_measure_results(
                 results[key] = full_results[:, meas_indices] ^ invert_mask
 
         else:
+            if noisy:
+                translator_fn_name = "translate_cirq_to_qtrajectory"
+                sampler_fn = self._sim_module.qtrajectory_sample
+            else:
+                translator_fn_name = "translate_cirq_to_qsim"
+                sampler_fn = self._sim_module.qsim_sample
+
             options["c"], _ = self._translate_circuit(
                 program,
                 translator_fn_name,