Add MPSStateSpace.Sample and its unit test for small case.

lib/BUILD

@@ -674,7 +674,10 @@ cc_library(
 cc_library(
     name = "mps_statespace",
     hdrs = ["mps_statespace.h"],
-    deps = ["@eigen//:eigen3"],
+    deps = [
+        ":util",
+        "@eigen//:eigen3",
+    ],
 )
 
 cc_library(

lib/mps_statespace.h

@@ -28,6 +28,7 @@
 #include <memory>
 
 #include "../eigen/Eigen/Dense"
+#include "util.h"
 
 namespace qsim {
 
@@ -92,6 +93,10 @@ class MPSStateSpace {
     unsigned bond_dim_;
   };
 
+  static uint64_t MinSize(unsigned num_qubits) {
+    return 2 * (uint64_t{1} << num_qubits);
+  };
+
   // Note: ForArgs are currently unused.
   template <typename... ForArgs>
   MPSStateSpace(ForArgs&&... args) : for_(args...) {}
@@ -150,24 +155,85 @@ class MPSStateSpace {
     return true;
   }
 
+  static void SetAllZeros(MPS& state) {
+    auto size = Size(state);
+    memset(state.get(), 0, sizeof(fp_type) * size);
+  }
+
   // Set the MPS to the |0> state.
   static void SetStateZero(MPS& state) {
     auto size = Size(state);
-    memset(state.get(), 0, sizeof(fp_type) * size);
+    SetAllZeros(state);
     auto block_size = 4 * state.bond_dim() * state.bond_dim();
     state.get()[0] = 1.0;
     for (unsigned i = 4 * state.bond_dim(); i < size; i += block_size) {
       state.get()[i] = 1.0;
     }
   }
 
+  static std::complex<fp_type> GetAmpl(const MPS& state, uint64_t i) {
+    uint64_t p = 2 * i;
+    return std::complex<fp_type>(state.get()[p], state.get()[p + 1]);
+  }
+
+  static void SetAmpl(
+      MPS& state, uint64_t i, const std::complex<fp_type>& ampl) {
+    uint64_t p = 2 * i;
+    state.get()[p] = std::real(ampl);
+    state.get()[p + 1] = std::imag(ampl);
+  }
+
+  static void SetAmpl(MPS& state, uint64_t i, fp_type re, fp_type im) {
+    uint64_t p = 2 * i;
+    state.get()[p] = re;
+    state.get()[p + 1] = im;
+  }
+
   // Computes Re{<state1 | state2 >} for two equal sized MPS.
   // Requires: state1.bond_dim() == state2.bond_dim() &&
   //           state1.num_qubits() == state2.num_qubits()
   static fp_type RealInnerProduct(MPS& state1, MPS& state2) {
     return InnerProduct(state1, state2).real();
   }
 
+  // Generates the sampled bitstrings from the final w
+  template <typename DistrRealType = double>
+  std::vector<uint64_t> Sample(
+      const MPS& state, uint64_t num_samples, unsigned seed) const {
+    std::vector<uint64_t> bitstrings;
+
+    if (num_samples > 0) {
+      double norm = 0;
+      uint64_t size = MinSize(state.num_qubits()) / 2;
+
+      const fp_type* p = state.get();
+
+      for (uint64_t k = 0; k < size; ++k) {
+        auto re = p[2 * k];
+        auto im = p[2 * k + 1];
+        norm += re * re + im * im;
+      }
+
+      auto rs = GenerateRandomValues<DistrRealType>(num_samples, seed, norm);
+
+      uint64_t m = 0;
+      double csum = 0;
+      bitstrings.reserve(num_samples);
+
+      for (uint64_t k = 0; k < size; ++k) {
+        auto re = p[2 * k];
+        auto im = p[2 * k + 1];
+        csum += re * re + im * im;
+        while (rs[m] < csum && m < num_samples) {
+          bitstrings.emplace_back(k);
+          ++m;
+        }
+      }
+    }
+
+    return bitstrings;
+  }
+
   // Computes <state1 | state2 > for two equal sized MPS.
   // Requires: state1.bond_dim() == state2.bond_dim() &&
   //           state1.num_qubits() == state2.num_qubits()

tests/mps_statespace_test.cc

@@ -586,6 +586,51 @@ TEST(MPSStateSpaceTest, InnerProduct4) {
   EXPECT_NEAR(f, 0.5524, 1e-4);
 }
 
+TEST(MPSStateSpaceTest, SamplingSmall) {
+  uint64_t num_samples = 2000000;
+  constexpr unsigned num_qubits = 6;
+  constexpr int bond_dim = 4;
+  constexpr uint64_t size = uint64_t{1} << num_qubits;
+
+  auto ss = MPSStateSpace<For, float>(1);
+  auto mps = ss.Create(num_qubits, bond_dim);
+
+  ss.SetAllZeros(mps);
+
+  std::vector<float> ps = {
+    0.0046, 0.0100, 0.0136, 0.0239, 0.0310, 0.0263, 0.0054, 0.0028,
+    0.0129, 0.0128, 0.0319, 0.0023, 0.0281, 0.0185, 0.0284, 0.0175,
+    0.0148, 0.0307, 0.0017, 0.0319, 0.0185, 0.0304, 0.0133, 0.0229,
+    0.0190, 0.0213, 0.0197, 0.0028, 0.0288, 0.0084, 0.0052, 0.0147,
+    0.0164, 0.0084, 0.0154, 0.0093, 0.0215, 0.0059, 0.0115, 0.0039,
+    0.0111, 0.0101, 0.0314, 0.0003, 0.0173, 0.0075, 0.0116, 0.0091,
+    0.0289, 0.0005, 0.0027, 0.0293, 0.0063, 0.0181, 0.0043, 0.0063,
+    0.0169, 0.0101, 0.0295, 0.0227, 0.0275, 0.0218, 0.0131, 0.0172,
+  };
+
+  for (unsigned i = 0; i < ss.Size(mps); ++i) {
+    mps.get()[i] = i;
+  }
+  for (uint64_t i = 0; i < size; ++i) {
+    auto r = std::sqrt(ps[i]);
+    ss.SetAmpl(mps, i, r * std::cos(i), r * std::sin(i));
+  }
+
+  auto samples = ss.Sample(mps, num_samples, 1);
+
+  EXPECT_EQ(samples.size(), num_samples);
+  std::vector<double> bins(size, 0);
+
+  for (auto sample : samples) {
+    ASSERT_LT(sample, size);
+    bins[sample] += 1;
+  }
+
+  for (uint64_t i = 0; i < size; ++i) {
+    EXPECT_NEAR(bins[i] / num_samples, ps[i], 4e-4);
+  }
+}
+
 }  // namespace
 }  // namespace mps
 }  // namespace qsim