Add MidMeasureMP C++ implementation

.github/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* `lightning.qubit` supports mid-circuit measurements.
+  [(#621)](https://github.com/PennyLaneAI/pennylane-lightning/pull/621)
+
 * Add two new python classes (LightningStateVector and LightningMeasurements) to support `lightning.qubit2`.
   [(#613)](https://github.com/PennyLaneAI/pennylane-lightning/pull/613)
 
@@ -32,7 +35,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Ali Asadi, Amintor Dusko, Vincent Michaud-Rioux
+Ali Asadi, Amintor Dusko, Thomas Germain, Vincent Michaud-Rioux
 
 ---
 

.github/workflows/tests_gpu_kokkos.yml

@@ -325,7 +325,7 @@ jobs:
           OMP_PROC_BIND: false
         run: |
           cd main/
-          PL_DEVICE=lightning.qubit python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=lightning.qubit python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
           pl-device-test --device lightning.qubit --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.qubit --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           PL_DEVICE=lightning.kokkos python -m pytest tests/ $COVERAGE_FLAGS

.github/workflows/tests_linux.yml

@@ -106,7 +106,7 @@ jobs:
       matrix:
         os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
-    timeout-minutes: 30
+    timeout-minutes: 60
     name: Python tests
     runs-on: ${{ matrix.os }}
 
@@ -184,7 +184,8 @@ jobs:
         run: |
           cd main/
           DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
-          PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
+          OMP_NUM_THREADS=1 PL_DEVICE=${DEVICENAME} python -m pytest -n auto tests/test_native_mcm.py $COVERAGE_FLAGS
           pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}
@@ -392,7 +393,7 @@ jobs:
         run: |
           cd main/
           DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
-          PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
           pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}
@@ -590,7 +591,7 @@ jobs:
         run: |
           cd main/
           DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
-          PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
           pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}
@@ -609,10 +610,10 @@ jobs:
         if: ${{ matrix.pl_backend == 'all' }}
         run: |
           cd main/
-          PL_DEVICE=lightning.qubit python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=lightning.qubit python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
           pl-device-test --device lightning.qubit --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.qubit --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
-          PL_DEVICE=lightning.kokkos python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=lightning.kokkos python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
           pl-device-test --device lightning.kokkos --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.kokkos --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}

.github/workflows/tests_without_binary.yml

@@ -106,7 +106,7 @@ jobs:
         run: |
           cd main/
           DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
-          PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
           pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
 

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0-dev9"
+__version__ = "0.36.0-dev10"

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.
      *

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;

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

pennylane_lightning/core/src/simulators/lightning_qubit/bindings/LQubitBindings.hpp

@@ -220,6 +220,10 @@ void registerBackendClassSpecificBindings(PyClass &pyclass) {
                 }
             },
             "Copy StateVector data into a Numpy array.")
+        .def("collapse", &StateVectorT::collapse,
+             "Collapse the statevector onto the 0 or 1 branch of a given wire.")
+        .def("normalize", &StateVectorT::normalize,
+             "Normalizes the statevector to norm 1.")
         .def("applyControlledMatrix", &applyControlledMatrix<StateVectorT>,
              "Apply controlled operation")
         .def("kernel_map", &svKernelMap<StateVectorT>,

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];

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]));
+    }
+}

pennylane_lightning/lightning_kokkos/lightning_kokkos.py

@@ -64,7 +64,7 @@
     from pennylane.ops.op_math import Adjoint
     from pennylane.wires import Wires
 
-    # pylint: disable=import-error, no-name-in-module, ungrouped-imports
+    # pylint: disable=ungrouped-imports
     from pennylane_lightning.lightning_kokkos_ops.algorithms import (
         AdjointJacobianC64,
         AdjointJacobianC128,