Add shot expval in C++ layer (#556)

* init commit with PauliX,Y,Z, Hardmard support

* add unit tests & sample

* add PauliX, Y, Z & Hadamard support

* move expval shots to base class

* Auto update version

* add MPI support

* add samples_to_counts to measurement base class

* tidy up tests

* move tests to measurement base

* move mpi tests to measurement base class

* make format

* add python layer

* add mpi tests

* remove else

* add tensor product support

* add more tests for tensor prod obs shots

* make format

* add Hamiltonian support

* make format

* quick fix

* add hamiltonian tests

* add more py tests for tensorprod

* remove changes made in py layer

* revert more changes in py layer

* tidy up code

* update measurement base and kokkos

* make format

* refactor and add _sample_state method

* for loop sum to accumulate sum

* refactor and add _preprocess_state() method

* add more unit tests for LK backend

* Auto update version

* tidy up docstring

* remove samples() and samples2counts method

* fix typo

* quick fix

* update LQRaw object creation

* fix for multiple backends

* update format

* format fix

* format update

* make format & trigger MPI CI

* add changelog and codecov for expval(obs) method

* add shot_range coverage

* tidy up code

* add more unit tests

* fix typos

* add MPI Hamiltonian tests for shots

* add HamBase tests for applyInPlaceShots

* update based on comments

* add more unit tests

* move tests for non-distributed backends to base

* move all unit tests into the base class

* move getTotalNumQubits to base and tidy up code

* tidy up code

* revert changes in Test_MeasurementLQ

* fix typo

.github/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* Add shots support for expectation value calculation for given observables (`NamedObs`, `TensorProd` and `Hamiltonian`) based on Pauli words, `Identity` and `Hadamard` in the C++ layer by adding `measure_with_samples` to the measurement interface. All Lightning backends support this support feature.
+[(#556)](https://github.com/PennyLaneAI/pennylane-lightning/pull/556)
+
 * `qml.QubitUnitary` operators can be included in a circuit differentiated with the adjoint method. Lightning handles circuits with arbitrary non-differentiable `qml.QubitUnitary` operators. 1,2-qubit `qml.QubitUnitary` operators with differentiable parameters can be differentiated using decomposition.
   [(#540)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/540)
 

mpitests/test_expval.py

@@ -31,8 +31,6 @@ class TestExpval:
     def test_identity_expectation(self, theta, phi, tol):
         """Test that identity expectation value (i.e. the trace) is 1"""
         dev = qml.device(device_name, mpi=True, wires=3)
-        if device_name == "lightning.gpu" and dev.R_DTYPE == np.float32:
-            pytest.skip("Skipped FP32 tests for expval in lightning.gpu")
 
         O1 = qml.Identity(wires=[0])
         O2 = qml.Identity(wires=[1])
@@ -49,9 +47,6 @@ def test_pauliz_expectation(self, theta, phi, tol):
         """Test that PauliZ expectation value is correct"""
         dev = qml.device(device_name, mpi=True, wires=3)
 
-        if device_name == "lightning.gpu" and dev.R_DTYPE == np.float32:
-            pytest.skip("Skipped FP32 tests for expval in lightning.gpu")
-
         O1 = qml.PauliZ(wires=[0])
         O2 = qml.PauliZ(wires=[1])
 
@@ -67,9 +62,6 @@ def test_paulix_expectation(self, theta, phi, tol):
         """Test that PauliX expectation value is correct"""
         dev = qml.device(device_name, mpi=True, wires=3)
 
-        if device_name == "lightning.gpu" and dev.R_DTYPE == np.float32:
-            pytest.skip("Skipped FP32 tests for expval in lightning.gpu")
-
         O1 = qml.PauliX(wires=[0])
         O2 = qml.PauliX(wires=[1])
 
@@ -89,9 +81,6 @@ def test_pauliy_expectation(self, theta, phi, tol):
         """Test that PauliY expectation value is correct"""
         dev = qml.device(device_name, mpi=True, wires=3)
 
-        if device_name == "lightning.gpu" and dev.R_DTYPE == np.float32:
-            pytest.skip("Skipped FP32 tests for expval in lightning.gpu")
-
         O1 = qml.PauliY(wires=[0])
         O2 = qml.PauliY(wires=[1])
 
@@ -130,8 +119,6 @@ def test_hermitian_expectation(self, n_wires, theta, phi, tol):
         n_qubits = 7
         dev_def = qml.device("default.qubit", wires=n_qubits)
         dev = qml.device(device_name, mpi=True, wires=n_qubits)
-        if device_name == "lightning.gpu" and dev.R_DTYPE == np.float32:
-            pytest.skip("Skipped FP32 tests for expval in lightning.gpu")
         comm = MPI.COMM_WORLD
 
         m = 2**n_wires

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.34.0-dev7"
+__version__ = "0.34.0-dev8"

pennylane_lightning/core/src/measurements/MeasurementsBase.hpp

@@ -17,10 +17,13 @@
  */
 #pragma once
 
+#include <string>
 #include <vector>
 
 #include "Observables.hpp"
 
+#include "CPUMemoryModel.hpp"
+
 /// @cond DEV
 namespace {
 using namespace Pennylane::Observables;
@@ -111,6 +114,190 @@ template <class StateVectorT, class Derived> class MeasurementsBase {
     auto generate_samples(size_t num_samples) -> std::vector<size_t> {
         return static_cast<Derived *>(this)->generate_samples(num_samples);
     };
-};
 
-} // namespace Pennylane::Measures
+    /**
+     * @brief Calculate the expectation value for a general Observable.
+     *
+     * @param obs Observable.
+     * @param num_shots Number of shots used to generate samples
+     * @param shot_range The range of samples to use. All samples are used
+     * by default.
+     *
+     * @return Expectation value with respect to the given observable.
+     */
+    auto expval(const Observable<StateVectorT> &obs, const size_t &num_shots,
+                const std::vector<size_t> &shot_range = {}) -> PrecisionT {
+        PrecisionT result{0.0};
+
+        if (obs.getObsName().find("SparseHamiltonian") != std::string::npos) {
+            // SparseHamiltonian does not support samples in pennylane.
+            PL_ABORT("For SparseHamiltonian Observables, expval calculation is "
+                     "not supported by shots");
+        } else if (obs.getObsName().find("Hermitian") != std::string::npos) {
+            // TODO support. This support requires an additional method to solve
+            // eigenpair and unitary matrices, and the results of eigenpair and
+            // unitary matrices data need to be added to the Hermitian class and
+            // public methods are need to access eigen values. Note the
+            // assumption that eigen values are -1 and 1 in the
+            // `measurement_with_sample` method should be updated as well.
+            PL_ABORT("For Hermitian Observables, expval calculation is not "
+                     "supported by shots");
+        } else if (obs.getObsName().find("Hamiltonian") != std::string::npos) {
+            auto coeffs = obs.getCoeffs();
+            for (size_t obs_term_idx = 0; obs_term_idx < coeffs.size();
+                 obs_term_idx++) {
+                auto obs_samples = measure_with_samples(
+                    obs, num_shots, shot_range, obs_term_idx);
+                PrecisionT result_per_term = std::accumulate(
+                    obs_samples.begin(), obs_samples.end(), 0.0);
+
+                result +=
+                    coeffs[obs_term_idx] * result_per_term / obs_samples.size();
+            }
+        } else {
+            auto obs_samples = measure_with_samples(obs, num_shots, shot_range);
+            result =
+                std::accumulate(obs_samples.begin(), obs_samples.end(), 0.0);
+            result /= obs_samples.size();
+        }
+        return result;
+    }
+
+    /**
+     * @brief Calculate the expectation value for a general Observable.
+     *
+     * @param obs Observable.
+     * @param num_shots Number of shots used to generate samples
+     * @param shot_range The range of samples to use. All samples are used
+     * by default.
+     * @param term_idx Index of a Hamiltonian term
+     *
+     * @return Expectation value with respect to the given observable.
+     */
+    auto measure_with_samples(const Observable<StateVectorT> &obs,
+                              const size_t &num_shots,
+                              const std::vector<size_t> &shot_range,
+                              size_t term_idx = 0) {
+        const size_t num_qubits = _statevector.getTotalNumQubits();
+        std::vector<size_t> obs_wires;
+        std::vector<size_t> identity_wires;
+
+        auto sub_samples = _sample_state(obs, num_shots, shot_range, obs_wires,
+                                         identity_wires, term_idx);
+
+        size_t num_samples = shot_range.empty() ? num_shots : shot_range.size();
+
+        std::vector<PrecisionT> obs_samples(num_samples, 0);
+
+        size_t num_identity_obs = identity_wires.size();
+        if (!identity_wires.empty()) {
+            size_t identity_obs_idx = 0;
+            for (size_t i = 0; i < obs_wires.size(); i++) {
+                if (identity_wires[identity_obs_idx] == obs_wires[i]) {
+                    std::swap(obs_wires[identity_obs_idx], obs_wires[i]);
+                    identity_obs_idx++;
+                }
+            }
+        }
+
+        for (size_t i = 0; i < num_samples; i++) {
+            std::vector<size_t> local_sample(obs_wires.size());
+            size_t idx = 0;
+            for (auto &obs_wire : obs_wires) {
+                local_sample[idx] = sub_samples[i * num_qubits + obs_wire];
+                idx++;
+            }
+
+            if (num_identity_obs != obs_wires.size()) {
+                // eigen values are `1` and `-1` for PauliX, PauliY, PauliZ,
+                // Hadamard gates the eigen value for a eigen vector |00001> is
+                // -1 since sum of the value at each bit position is odd
+                size_t bitSum = static_cast<size_t>(
+                    std::accumulate(local_sample.begin() + num_identity_obs,
+                                    local_sample.end(), 0));
+                if ((bitSum & size_t{1}) == 1) {
+                    obs_samples[i] = -1;
+                } else {
+                    obs_samples[i] = 1;
+                }
+            } else {
+                // eigen value for Identity gate is `1`
+                obs_samples[i] = 1;
+            }
+        }
+        return obs_samples;
+    }
+
+  private:
+    /**
+     * @brief Return preprocess state with a observable
+     *
+     * @param obs The observable to sample
+     * @param obs_wires Observable wires.
+     * @param identity_wires Wires of Identity gates
+     * @param term_idx Index of a Hamiltonian term. For other observables, its
+     * value is 0, which is set as default.
+     *
+     * @return a StateVectorT object
+     */
+    auto _preprocess_state(const Observable<StateVectorT> &obs,
+                           std::vector<size_t> &obs_wires,
+                           std::vector<size_t> &identity_wires,
+                           const size_t &term_idx = 0) {
+        if constexpr (std::is_same_v<
+                          typename StateVectorT::MemoryStorageT,
+                          Pennylane::Util::MemoryStorageLocation::External>) {
+            StateVectorT sv(_statevector.getData(), _statevector.getLength());
+            sv.updateData(_statevector.getData(), _statevector.getLength());
+            obs.applyInPlaceShots(sv, identity_wires, obs_wires, term_idx);
+            return sv;
+        } else {
+            StateVectorT sv(_statevector);
+            obs.applyInPlaceShots(sv, identity_wires, obs_wires, term_idx);
+            return sv;
+        }
+    }
+
+    /**
+     * @brief Return samples of a observable
+     *
+     * @param obs The observable to sample
+     * @param num_shots Number of shots used to generate samples
+     * @param shot_range The range of samples to use. All samples are used by
+     * default.
+     * @param obs_wires Observable wires.
+     * @param identity_wires Wires of Identity gates
+     * @param term_idx Index of a Hamiltonian term. For other observables, its
+     * value is 0, which is set as default.
+     *
+     * @return std::vector<size_t> samples in std::vector
+     */
+    auto _sample_state(const Observable<StateVectorT> &obs,
+                       const size_t &num_shots,
+                       const std::vector<size_t> &shot_range,
+                       std::vector<size_t> &obs_wires,
+                       std::vector<size_t> &identity_wires,
+                       const size_t &term_idx = 0) {
+        const size_t num_qubits = _statevector.getTotalNumQubits();
+        auto sv = _preprocess_state(obs, obs_wires, identity_wires, term_idx);
+        Derived measure(sv);
+        auto samples = measure.generate_samples(num_shots);
+
+        if (!shot_range.empty()) {
+            std::vector<size_t> sub_samples(shot_range.size() * num_qubits);
+            // Get a slice of samples based on the shot_range vector
+            size_t shot_idx = 0;
+            for (const auto &i : shot_range) {
+                for (size_t j = i * num_qubits; j < (i + 1) * num_qubits; j++) {
+                    // TODO some extra work to make it cache-friendly
+                    sub_samples[shot_idx * num_qubits + j - i * num_qubits] =
+                        samples[j];
+                }
+                shot_idx++;
+            }
+            return sub_samples;
+        }
+        return samples;
+    }
+};
+} // namespace Pennylane::Measures