Add mcmc sampler (#384)

Co-authored-by: Lee James O'Riordan <mlxd@users.noreply.github.com>
Co-authored-by: Amintor Dusko <87949283+AmintorDusko@users.noreply.github.com>
Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Shuli Shu <08cnbj@gmail.com>
Co-authored-by: Vincent Michaud-Rioux <vincentm@nanoacademic.com>
Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

.github/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* Add MCMC sampler.
+[(#384)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/384)
+
 ### Breaking changes
 
 ### Improvements
@@ -68,7 +71,7 @@ Amintor Dusko, Vincent Michaud-Rioux, Lee James O'Riordan, Chae-Yeun Park
 
 This release contains contributions from (in alphabetical order):
 
-Amintor Dusko
+Amintor Dusko, Shuli Shu, Trevor Vincent
 
 ---
 

doc/devices.rst

@@ -98,3 +98,26 @@ If you are computing a large number of expectation values, or if you are using a
     os.environ["OMP_NUM_THREADS"] = 4
     import pennylane as qml
     dev = qml.device("lightning.qubit", wires=2, batch_obs=True)
+
+.. raw:: html
+
+    </div>
+
+**Markov Chain Monte Carlo sampling support:**
+
+The ``lightning.qubit`` device allows users to use the Markov Chain Monte Carlo (MCMC) sampling method to generate approximate samples. To enable the MCMC sampling method for sample generation, initialize a ``lightning.qubit`` device with the ``mcmc=True`` keyword argument, as:
+
+.. code-block:: python
+
+    import pennylane as qml
+    dev = qml.device("lightning.qubit", wires=2, shots=1000, mcmc=True)
+
+By default, the ``kernel_name`` is ``"Local"`` and ``num_burnin`` is ``100``. The local kernel conducts a bit-flip local transition between states. The local kernel generates a random qubit site and then generates a random number to determine  the new bit at that qubit site. 
+
+The ``lightning.qubit`` device also supports a ``"NonZeroRandom"`` kernel. This kernel randomly transits between states that have nonzero probability. It can be enabled by initializing the device as:
+
+.. code-block:: python
+
+    import pennylane as qml
+    dev = qml.device("lightning.qubit", wires=2, shots=1000, mcmc=True, kernel_name="NonZeroRandom", num_burnin=200)
+

pennylane_lightning/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.30.0-dev"
+__version__ = "0.30.0-dev1"

pennylane_lightning/lightning_qubit.py

@@ -160,6 +160,13 @@ class LightningQubit(QubitDevice):
             the expectation values. Defaults to ``None`` if not specified. Setting
             to ``None`` results in computing statistics like expectation values and
             variances analytically.
+        mcmc (bool): Determine whether to use the approximate Markov Chain Monte Carlo sampling method when generating samples.
+        kernel_name (str): name of transition kernel. The current version supports two kernels: ``"Local"`` and ``"NonZeroRandom"``.
+            The local kernel conducts a bit-flip local transition between states. The local kernel generates a
+            random qubit site and then generates a random number to determine the new bit at that qubit site. The ``"NonZeroRandom"`` kernel
+            randomly transits between states that have nonzero probability.
+        num_burnin (int): number of steps that will be dropped. Increasing this value will
+        result in a closer approximation but increased runtime.
         batch_obs (bool): Determine whether we process observables parallelly when computing the
             jacobian. This value is only relevant when the lightning qubit is built with OpenMP.
     """
@@ -173,7 +180,18 @@ class LightningQubit(QubitDevice):
     operations = allowed_operations
     observables = allowed_observables
 
-    def __init__(self, wires, *, c_dtype=np.complex128, shots=None, batch_obs=False, analytic=None):
+    def __init__(
+        self,
+        wires,
+        *,
+        c_dtype=np.complex128,
+        shots=None,
+        mcmc=False,
+        kernel_name="Local",
+        num_burnin=100,
+        batch_obs=False,
+        analytic=None,
+    ):
         if c_dtype is np.complex64:
             r_dtype = np.float32
             self.use_csingle = True
@@ -190,6 +208,20 @@ def __init__(self, wires, *, c_dtype=np.complex128, shots=None, batch_obs=False,
         self._state = self._create_basis_state(0)
         self._pre_rotated_state = self._state
 
+        self._mcmc = mcmc
+        if self._mcmc:
+            if kernel_name not in [
+                "Local",
+                "NonZeroRandom",
+            ]:
+                raise NotImplementedError(
+                    f"The {kernel_name} is not supported and currently only 'Local' and 'NonZeroRandom' kernels are supported."
+                )
+            if num_burnin >= shots:
+                raise ValueError("Shots should be greater than num_burnin.")
+            self._kernel_name = kernel_name
+            self._num_burnin = num_burnin
+
     @property
     def stopping_condition(self):
         """.BooleanFn: Returns the stopping condition for the device. The returned
@@ -776,14 +808,17 @@ def generate_samples(self):
         Returns:
             array[int]: array of samples in binary representation with shape ``(dev.shots, dev.num_wires)``
         """
-
         # Initialization of state
         ket = np.ravel(self._state)
 
         state_vector = StateVectorC64(ket) if self.use_csingle else StateVectorC128(ket)
         M = MeasuresC64(state_vector) if self.use_csingle else MeasuresC128(state_vector)
-
-        return M.generate_samples(len(self.wires), self.shots).astype(int, copy=False)
+        if self._mcmc:
+            return M.generate_mcmc_samples(
+                len(self.wires), self._kernel_name, self._num_burnin, self.shots
+            ).astype(int, copy=False)
+        else:
+            return M.generate_samples(len(self.wires), self.shots).astype(int, copy=False)
 
     def expval(self, observable, shot_range=None, bin_size=None):
         """Expectation value of the supplied observable.

pennylane_lightning/src/bindings/Bindings.cpp

@@ -135,6 +135,27 @@ void lightning_class_bindings(py::module_ &m) {
                      strides /* strides for each axis     */
                      ));
              })
+        .def("generate_mcmc_samples",
+             [](Measures<PrecisionT> &M, size_t num_wires,
+                const std::string &kernelname, size_t num_burnin,
+                size_t num_shots) {
+                 std::vector<size_t> &&result = M.generate_samples_metropolis(
+                     kernelname, num_burnin, num_shots);
+
+                 const size_t ndim = 2;
+                 const std::vector<size_t> shape{num_shots, num_wires};
+                 constexpr auto sz = sizeof(size_t);
+                 const std::vector<size_t> strides{sz * num_wires, sz};
+                 // return 2-D NumPy array
+                 return py::array(py::buffer_info(
+                     result.data(), /* data as contiguous array  */
+                     sz,            /* size of one scalar        */
+                     py::format_descriptor<size_t>::format(), /* data type */
+                     ndim,   /* number of dimensions      */
+                     shape,  /* shape of the matrix       */
+                     strides /* strides for each axis     */
+                     ));
+             })
         .def("var",
              [](Measures<PrecisionT> &M, const std::string &operation,
                 const std::vector<size_t> &wires) {

pennylane_lightning/src/simulator/CMakeLists.txt

@@ -1,6 +1,6 @@
 project(lightning_simulator)
 
-set(SIMULATOR_FILES StateVectorRawCPU.cpp Observables.cpp StateVectorManagedCPU.cpp Measures.cpp CACHE INTERNAL "" FORCE)
+set(SIMULATOR_FILES StateVectorRawCPU.cpp Observables.cpp StateVectorManagedCPU.cpp TransitionKernels.cpp Measures.cpp CACHE INTERNAL "" FORCE)
 
 add_library(lightning_simulator STATIC ${SIMULATOR_FILES})
 

pennylane_lightning/src/simulator/Measures.hpp

@@ -33,6 +33,7 @@
 #include "Observables.hpp"
 #include "StateVectorManagedCPU.hpp"
 #include "StateVectorRawCPU.hpp"
+#include "TransitionKernels.hpp"
 
 namespace Pennylane::Simulators {
 
@@ -331,6 +332,106 @@ class Measures {
         return expected_value_list;
     };
 
+    /**
+     * @brief Complete a single Metropolis-Hastings step.
+     *
+     * @param sv state vector
+     * @param tk User-defined functor for producing transitions
+     * between metropolis states.
+     * @param gen Random number generator.
+     * @param distrib Random number distribution.
+     * @param init_idx Init index of basis state.
+     */
+    size_t metropolis_step(const SVType &sv,
+                           const std::unique_ptr<TransitionKernel<fp_t>> &tk,
+                           std::mt19937 &gen,
+                           std::uniform_real_distribution<fp_t> &distrib,
+                           size_t init_idx) {
+        auto init_plog = std::log(
+            (sv.getData()[init_idx] * std::conj(sv.getData()[init_idx]))
+                .real());
+
+        auto init_qratio = tk->operator()(init_idx);
+
+        // transition kernel outputs these two
+        auto &trans_idx = init_qratio.first;
+        auto &trans_qratio = init_qratio.second;
+
+        auto trans_plog = std::log(
+            (sv.getData()[trans_idx] * std::conj(sv.getData()[trans_idx]))
+                .real());
+
+        auto alph =
+            std::min<fp_t>(1., trans_qratio * std::exp(trans_plog - init_plog));
+        auto ran = distrib(gen);
+
+        if (ran < alph) {
+            return trans_idx;
+        }
+        return init_idx;
+    }
+
+    /**
+     * @brief Generate samples using the Metropolis-Hastings method.
+     * Reference: Numerical Recipes, NetKet paper
+     *
+     * @param transition_kernel User-defined functor for producing transitions
+     * between metropolis states.
+     * @param num_burnin Number of Metropolis burn-in steps.
+     * @param num_samples The number of samples to generate.
+     * @return 1-D vector of samples in binary, each sample is
+     * separated by a stride equal to the number of qubits.
+     */
+    std::vector<size_t>
+    generate_samples_metropolis(const std::string &kernelname,
+                                size_t num_burnin, size_t num_samples) {
+        size_t num_qubits = original_statevector.getNumQubits();
+        std::random_device rd;
+        std::mt19937 gen(rd());
+        std::uniform_real_distribution<fp_t> distrib(0.0, 1.0);
+        std::vector<size_t> samples(num_samples * num_qubits, 0);
+        std::unordered_map<size_t, size_t> cache;
+
+        TransitionKernelType transition_kernel =
+            Pennylane::TransitionKernelType::Local;
+        if (kernelname == "NonZeroRandom") {
+            transition_kernel = Pennylane::TransitionKernelType::NonZeroRandom;
+        }
+
+        auto tk =
+            kernel_factory(transition_kernel, original_statevector.getData(),
+                           original_statevector.getNumQubits());
+        size_t idx = 0;
+
+        // Burn In
+        for (size_t i = 0; i < num_burnin; i++) {
+            idx = metropolis_step(original_statevector, tk, gen, distrib,
+                                  idx); // Burn-in.
+        }
+
+        // Sample
+        for (size_t i = 0; i < num_samples; i++) {
+            idx = metropolis_step(original_statevector, tk, gen, distrib, idx);
+
+            if (cache.contains(idx)) {
+                size_t cache_id = cache[idx];
+                auto it_temp = samples.begin() + cache_id * num_qubits;
+                std::copy(it_temp, it_temp + num_qubits,
+                          samples.begin() + i * num_qubits);
+            }
+
+            // If not cached, compute
+            else {
+                for (size_t j = 0; j < num_qubits; j++) {
+                    samples[i * num_qubits + (num_qubits - 1 - j)] =
+                        (idx >> j) & 1U;
+                }
+                cache[idx] = i;
+            }
+        }
+        return samples;
+    }
+
     /**
      * @brief Variance of a Sparse Hamiltonian.
      *

pennylane_lightning/src/simulator/TransitionKernels.cpp

@@ -0,0 +1,9 @@
+#include "TransitionKernels.hpp"
+
+// explicit instantiation
+template class Pennylane::TransitionKernel<float>;
+template class Pennylane::TransitionKernel<double>;
+template class Pennylane::LocalTransitionKernel<float>;
+template class Pennylane::LocalTransitionKernel<double>;
+template class Pennylane::NonZeroRandomTransitionKernel<float>;
+template class Pennylane::NonZeroRandomTransitionKernel<double>;