Add MidMeasureMP in LightningQubit (cpp only) (#645)

* rng to MeasurementBase with setSeed() * collapse() and normalize() in StateVector * Auto update version * trigger ci * Auto update version * trigger ci * CHANGELOG.md Co-authored-by: Vincent Michaud-Rioux <[email protected]> --------- Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com> Co-authored-by: Vincent Michaud-Rioux <[email protected]> Co-authored-by: Vincent Michaud-Rioux <[email protected]>
PennyLaneAI · Mar 18, 2024 · 5b4ef77 · 5b4ef77
1 parent a3071d3
commit 5b4ef77
Show file tree

Hide file tree

Showing 7 changed files with 116 additions and 10 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* Add `collapse` and `normalize` methods to the `StateVectorLQubit` classes, enabling "branching" of the wavefunction. Add methods to create and seed an RNG in the `Measurements` modules.
+  [(#645)](https://github.com/PennyLaneAI/pennylane-lightning/pull/645)
+
 * Add two new python classes (LightningStateVector and LightningMeasurements) to support `lightning.qubit2`.
   [(#613)](https://github.com/PennyLaneAI/pennylane-lightning/pull/613)
 

diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0-dev10"
+__version__ = "0.36.0-dev11"
diff --git a/pennylane_lightning/core/src/measurements/MeasurementsBase.hpp b/pennylane_lightning/core/src/measurements/MeasurementsBase.hpp
@@ -17,6 +17,7 @@
  */
 #pragma once
 
+#include <random>
 #include <string>
 #include <vector>
 
@@ -55,6 +56,7 @@ template <class StateVectorT, class Derived> class MeasurementsBase {
 #else
     const StateVectorT &_statevector;
 #endif
+    std::mt19937 rng;
 
   public:
 #ifdef _ENABLE_PLGPU
@@ -65,6 +67,23 @@ template <class StateVectorT, class Derived> class MeasurementsBase {
         : _statevector{statevector} {};
 #endif
 
+    /**
+     * @brief Set the seed of the internal random generator
+     *
+     * @param seed Seed
+     */
+    void setSeed(const size_t seed) { rng.seed(seed); }
+
+    /**
+     * @brief Randomly set the seed of the internal random generator
+     *
+     * @param seed Seed
+     */
+    void setRandomSeed() {
+        std::random_device rd;
+        setSeed(rd());
+    }
+
     /**
      * @brief Calculate the expectation value for a general Observable.
      *

diff --git a/pennylane_lightning/core/src/simulators/lightning_gpu/measurements/MeasurementsGPU.hpp b/pennylane_lightning/core/src/simulators/lightning_gpu/measurements/MeasurementsGPU.hpp
@@ -224,10 +224,10 @@ class Measurements final
             data_type = CUDA_C_32F;
         }
 
-        std::mt19937 gen(std::random_device{}());
+        this->setRandomSeed();
         std::uniform_real_distribution<PrecisionT> dis(0.0, 1.0);
         for (size_t n = 0; n < num_samples; n++) {
-            rand_nums[n] = dis(gen);
+            rand_nums[n] = dis(this->rng);
         }
         std::vector<size_t> samples(num_samples * num_qubits, 0);
         std::unordered_map<size_t, size_t> cache;

diff --git a/pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubit.hpp b/pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubit.hpp
@@ -36,6 +36,7 @@ namespace {
 using Pennylane::LightningQubit::Util::Threading;
 using Pennylane::Util::CPUMemoryModel;
 using Pennylane::Util::exp2;
+using Pennylane::Util::squaredNorm;
 using namespace Pennylane::LightningQubit::Gates;
 } // namespace
 /// @endcond
@@ -632,5 +633,53 @@ class StateVectorLQubit : public StateVectorBase<PrecisionT, Derived> {
 
         applyMatrix(matrix.data(), wires, inverse);
     }
+
+    /**
+     * @brief Collapse the state vector as after having measured one of the
+     * qubits.
+     *
+     * The branch parameter imposes the measurement result on the given wire.
+     *
+     * @param wire Wire to collapse.
+     * @param branch Branch 0 or 1.
+     */
+    void collapse(const std::size_t wire, const bool branch) {
+        auto *arr = this->getData();
+        const size_t stride = pow(2, this->num_qubits_ - (1 + wire));
+        const size_t vec_size = pow(2, this->num_qubits_);
+        const auto section_size = vec_size / stride;
+        const auto half_section_size = section_size / 2;
+
+        // zero half the entries
+        // the "half" entries depend on the stride
+        // *_*_*_*_ for stride 1
+        // **__**__ for stride 2
+        // ****____ for stride 4
+        const size_t k = branch ? 0 : 1;
+        for (size_t idx = 0; idx < half_section_size; idx++) {
+            const size_t offset = stride * (k + 2 * idx);
+            for (size_t ids = 0; ids < stride; ids++) {
+                arr[offset + ids] = {0., 0.};
+            }
+        }
+
+        normalize();
+    }
+
+    /**
+     * @brief Normalize vector (to have norm 1).
+     */
+    void normalize() {
+        auto *arr = this->getData();
+        PrecisionT norm = std::sqrt(squaredNorm(arr, this->getLength()));
+
+        PL_ABORT_IF(norm < std::numeric_limits<PrecisionT>::epsilon() * 1e2,
+                    "vector has norm close to zero and can't be normalized");
+
+        std::complex<PrecisionT> inv_norm = 1. / norm;
+        for (size_t k = 0; k < this->getLength(); k++) {
+            arr[k] *= inv_norm;
+        }
+    }
 };
-} // namespace Pennylane::LightningQubit
+} // namespace Pennylane::LightningQubit
diff --git a/pennylane_lightning/core/src/simulators/lightning_qubit/measurements/MeasurementsLQubit.hpp b/pennylane_lightning/core/src/simulators/lightning_qubit/measurements/MeasurementsLQubit.hpp
@@ -466,11 +466,10 @@ class Measurements final
     generate_samples_metropolis(const std::string &kernelname,
                                 size_t num_burnin, size_t num_samples) {
         size_t num_qubits = this->_statevector.getNumQubits();
-        std::random_device rd;
-        std::mt19937 gen(rd());
         std::uniform_real_distribution<PrecisionT> distrib(0.0, 1.0);
         std::vector<size_t> samples(num_samples * num_qubits, 0);
         std::unordered_map<size_t, size_t> cache;
+        this->setRandomSeed();
 
         TransitionKernelType transition_kernel = TransitionKernelType::Local;
         if (kernelname == "NonZeroRandom") {
@@ -484,13 +483,14 @@ class Measurements final
 
         // Burn In
         for (size_t i = 0; i < num_burnin; i++) {
-            idx = metropolis_step(this->_statevector, tk, gen, distrib,
+            idx = metropolis_step(this->_statevector, tk, this->rng, distrib,
                                   idx); // Burn-in.
         }
 
         // Sample
         for (size_t i = 0; i < num_samples; i++) {
-            idx = metropolis_step(this->_statevector, tk, gen, distrib, idx);
+            idx = metropolis_step(this->_statevector, tk, this->rng, distrib,
+                                  idx);
 
             if (cache.contains(idx)) {
                 size_t cache_id = cache[idx];
@@ -562,9 +562,9 @@ class Measurements final
         auto &&probabilities = probs();
 
         std::vector<size_t> samples(num_samples * num_qubits, 0);
-        std::mt19937 generator(std::random_device{}());
         std::uniform_real_distribution<PrecisionT> distribution(0.0, 1.0);
         std::unordered_map<size_t, size_t> cache;
+        this->setRandomSeed();
 
         const size_t N = probabilities.size();
         std::vector<double> bucket(N);
@@ -611,7 +611,7 @@ class Measurements final
 
         // Pick samples
         for (size_t i = 0; i < num_samples; i++) {
-            PrecisionT pct = distribution(generator) * N;
+            PrecisionT pct = distribution(this->rng) * N;
             auto idx = static_cast<size_t>(pct);
             if (pct - idx > bucket[idx]) {
                 idx = bucket_partner[idx];

diff --git a/pennylane_lightning/core/src/simulators/lightning_qubit/tests/Test_StateVectorLQubit.cpp b/pennylane_lightning/core/src/simulators/lightning_qubit/tests/Test_StateVectorLQubit.cpp
@@ -216,3 +216,38 @@ TEMPLATE_PRODUCT_TEST_CASE("StateVectorLQubit::applyOperations",
             LightningException, "must all be equal"); // invalid parameters
     }
 }
+
+TEMPLATE_PRODUCT_TEST_CASE("StateVectorLQubit::collapse", "[StateVectorLQubit]",
+                           (StateVectorLQubitManaged, StateVectorLQubitRaw),
+                           (float, double)) {
+    using StateVectorT = TestType;
+    using PrecisionT = typename StateVectorT::PrecisionT;
+    using ComplexT = typename StateVectorT::ComplexT;
+    using TestVectorT = TestVector<ComplexT>;
+
+    const std::size_t num_qubits = 3;
+
+    SECTION("Collapse the state vector as after having measured one of the "
+            "qubits.") {
+        TestVectorT init_state = createPlusState<PrecisionT>(num_qubits);
+
+        const ComplexT coef{0.5, PrecisionT{0.0}};
+        const ComplexT zero{PrecisionT{0.0}, PrecisionT{0.0}};
+
+        std::vector<std::vector<std::vector<ComplexT>>> expected_state = {
+            {{coef, coef, coef, coef, zero, zero, zero, zero},
+             {coef, coef, zero, zero, coef, coef, zero, zero},
+             {coef, zero, coef, zero, coef, zero, coef, zero}},
+            {{zero, zero, zero, zero, coef, coef, coef, coef},
+             {zero, zero, coef, coef, zero, zero, coef, coef},
+             {zero, coef, zero, coef, zero, coef, zero, coef}},
+        };
+
+        std::size_t wire = GENERATE(0, 1, 2);
+        std::size_t branch = GENERATE(0, 1);
+        StateVectorLQubitManaged<PrecisionT> sv(init_state);
+        sv.collapse(wire, branch);
+
+        REQUIRE(sv.getDataVector() == approx(expected_state[branch][wire]));
+    }
+}