Add bindings to C++ expval, var, probs in Lightning (#214)

* Add pybind for PL::Measures::*

* Add python tests

* Update probs, expval, and var bindings w.r.t. vec-val QNode examples

* Update interface

* Add python tests

* Update tests

* Update methods and tests

* Integrate probs

* Integrate expval and var

* Update bindings

* Complete integration and tests

* Update bindings

* Update python tests

* Add more tests

* Fix codecov issues

* Apply some minor changes

* Update

* Try to fix the issue with Wheel::Windows

* Update

* Update

* Update for MSVC

* Update measures pybind

* Apply code review suggestions

* Remove using namespace Util in Measures.hpp

* Update CHANGELOG

Co-authored-by: Chae-Yeun Park <[email protected]>

Author: maliasadi

.github/CHANGELOG.md

@@ -7,6 +7,9 @@
 * Add C++ only benchmark for a given list of gates.
 [(#199)](https://github.com/PennyLaneAI/pennylane-lightning/pull/199)
 
+* Add bindings to C++ expval, var, probs.
+[(#214)](https://github.com/PennyLaneAI/pennylane-lightning/pull/214)
+
 ### Breaking changes
 
 ### Improvements
@@ -34,7 +37,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Amintor Dusko, Chae-Yeun Park, Lee James O'Riordan
+Ali Asadi, Amintor Dusko, Chae-Yeun Park, Lee James O'Riordan
 
 ---
 

pennylane_lightning/lightning_qubit.py

@@ -15,6 +15,7 @@
 This module contains the :class:`~.LightningQubit` class, a PennyLane simulator device that
 interfaces with C++ for fast linear algebra calculations.
 """
+from typing import List
 from warnings import warn
 
 import numpy as np
@@ -29,14 +30,17 @@
 import pennylane as qml
 from pennylane.devices import DefaultQubit
 from pennylane.operation import Expectation
+from pennylane.wires import Wires
 
 from ._version import __version__
 
 try:
     from .lightning_qubit_ops import (
+        MeasuresC64,
         StateVectorC64,
         AdjointJacobianC64,
         VectorJacobianProductC64,
+        MeasuresC128,
         StateVectorC128,
         AdjointJacobianC128,
         VectorJacobianProductC128,
@@ -493,6 +497,157 @@ def batch_vjp(
 
         return jacs, vjps
 
+    def probability(self, wires=None, shot_range=None, bin_size=None):
+        """Return the probability of each computational basis state.
+
+        Devices that require a finite number of shots always return the
+        estimated probability.
+
+        Args:
+            wires (Iterable[Number, str], Number, str, Wires): wires to return
+                marginal probabilities for. Wires not provided are traced out of the system.
+            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+                to use. If not specified, all samples are used.
+            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+                returns the measurement statistic separately over each bin. If not
+                provided, the entire shot range is treated as a single bin.
+
+        Returns:
+            array[float]: list of the probabilities
+        """
+        if self.shots is not None:
+            return self.estimate_probability(wires=wires, shot_range=shot_range, bin_size=bin_size)
+
+        wires = wires or self.wires
+        wires = Wires(wires)
+
+        # translate to wire labels used by device
+        device_wires = self.map_wires(wires)
+
+        # To support np.complex64 based on the type of self._state
+        dtype = self._state.dtype
+        if dtype == np.complex64:
+            use_csingle = True
+        elif dtype == np.complex128:
+            use_csingle = False
+        else:
+            raise TypeError(f"Unsupported complex Type: {dtype}")
+
+        # Initialization of state
+        ket = np.ravel(self._state)
+
+        if use_csingle:
+            ket = ket.astype(np.complex64)
+
+        state_vector = StateVectorC64(ket) if use_csingle else StateVectorC128(ket)
+        M = MeasuresC64(state_vector) if use_csingle else MeasuresC128(state_vector)
+
+        return M.probs(device_wires)
+
+    def expval(self, observable, shot_range=None, bin_size=None):
+        """Expectation value of the supplied observable.
+
+        Args:
+            observable: A PennyLane observable.
+            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+                to use. If not specified, all samples are used.
+            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+                returns the measurement statistic separately over each bin. If not
+                provided, the entire shot range is treated as a single bin.
+
+        Returns:
+            Expectation value of the observable
+        """
+        if isinstance(observable.name, List) or observable.name in [
+            "Identity",
+            "Projector",
+            "Hermitian",
+            "Hamiltonian",
+            "SparseHamiltonian",
+        ]:
+            # TODO: requires backend support
+            return super().expval(observable, shot_range=shot_range, bin_size=bin_size)
+
+        if self.shots is not None:
+            # estimate the expectation value
+            # TODO: Lightning support for sampling
+            samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
+            return np.squeeze(np.mean(samples, axis=0))
+
+        # To support np.complex64 based on the type of self._state
+        dtype = self._state.dtype
+        if dtype == np.complex64:
+            use_csingle = True
+        elif dtype == np.complex128:
+            use_csingle = False
+        else:
+            raise TypeError(f"Unsupported complex Type: {dtype}")
+
+        # Initialization of state
+        ket = np.ravel(self._pre_rotated_state)
+
+        if use_csingle:
+            ket = ket.astype(np.complex64)
+
+        state_vector = StateVectorC64(ket) if use_csingle else StateVectorC128(ket)
+        M = MeasuresC64(state_vector) if use_csingle else MeasuresC128(state_vector)
+
+        # translate to wire labels used by device
+        observable_wires = self.map_wires(observable.wires)
+
+        return M.expval(observable.name, observable_wires)
+
+    def var(self, observable, shot_range=None, bin_size=None):
+        """Variance of the supplied observable.
+
+        Args:
+            observable: A PennyLane observable.
+            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+                to use. If not specified, all samples are used.
+            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+                returns the measurement statistic separately over each bin. If not
+                provided, the entire shot range is treated as a single bin.
+
+        Returns:
+            Variance of the observable
+        """
+        if isinstance(observable.name, List) or observable.name in [
+            "Identity",
+            "Projector",
+            "Hermitian",
+        ]:
+            # TODO: requires backend support
+            return super().var(observable, shot_range=shot_range, bin_size=bin_size)
+
+        if self.shots is not None:
+            # estimate the var
+            # TODO: Lightning support for sampling
+            samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
+            return np.squeeze(np.var(samples, axis=0))
+
+        # To support np.complex64 based on the type of self._state
+        dtype = self._state.dtype
+        if dtype == np.complex64:
+            use_csingle = True
+        elif dtype == np.complex128:
+            use_csingle = False
+        else:
+            raise TypeError(f"Unsupported complex Type: {dtype}")
+
+        # Initialization of state
+        ket = np.ravel(self._pre_rotated_state)
+
+        if use_csingle:
+            ket = ket.astype(np.complex64)
+
+        state_vector = StateVectorC64(ket) if use_csingle else StateVectorC128(ket)
+        M = MeasuresC64(state_vector) if use_csingle else MeasuresC128(state_vector)
+
+        # translate to wire labels used by device
+        observable_wires = self.map_wires(observable.wires)
+
+        return M.var(observable.name, observable_wires)
+
 
 if not CPP_BINARY_AVAILABLE:
 

pennylane_lightning/src/bindings/Bindings.cpp

@@ -291,6 +291,30 @@ void lightning_class_bindings(py::module &m) {
             return py::make_tuple(py::array_t<ParamT>(py::cast(jac)),
                                   py::array_t<ParamT>(py::cast(vjp_res)));
         });
+
+    //***********************************************************************//
+    //                              Measures
+    //***********************************************************************//
+
+    class_name = "MeasuresC" + bitsize;
+    py::class_<Measures<PrecisionT>>(m, class_name.c_str())
+        .def(py::init<const StateVectorRaw<PrecisionT> &>())
+        .def("probs",
+             [](Measures<PrecisionT> &M, const std::vector<size_t> &wires) {
+                 if (wires.empty()) {
+                     return py::array_t<ParamT>(py::cast(M.probs()));
+                 }
+                 return py::array_t<ParamT>(py::cast(M.probs(wires)));
+             })
+        .def("expval",
+             [](Measures<PrecisionT> &M, const std::string &operation,
+                const std::vector<size_t> &wires) {
+                 return M.expval(operation, wires);
+             })
+        .def("var", [](Measures<PrecisionT> &M, const std::string &operation,
+                       const std::vector<size_t> &wires) {
+            return M.var(operation, wires);
+        });
 }
 
 /**

pennylane_lightning/src/bindings/Bindings.hpp

@@ -20,7 +20,9 @@
 #include "AdjointDiff.hpp"
 #include "IndicesUtil.hpp"
 #include "JacobianProd.hpp"
+#include "Measures.hpp"
 #include "StateVectorRaw.hpp"
+
 #include "pybind11/complex.h"
 #include "pybind11/numpy.h"
 #include "pybind11/pybind11.h"

pennylane_lightning/src/simulator/DynamicDispatcher.hpp

@@ -61,10 +61,11 @@ template <typename fp_t> class DynamicDispatcher {
     std::unordered_map<std::string, KernelType> kernel_map_;
 
     std::unordered_map<std::pair<std::string, KernelType>, Func,
-                       Internal::PairHash>
+                       Pennylane::Internal::PairHash>
         gates_;
 
     DynamicDispatcher() {
+        namespace Internal = Pennylane::Internal;
         for (const auto &[gate_op, n_wires] : Constant::gate_wires) {
             gate_wires_.emplace(lookup(Constant::gate_names, gate_op), n_wires);
         }
@@ -234,6 +235,7 @@ constexpr auto gateOpToFunctor() {
     return [](std::complex<PrecisionT> *data, size_t num_qubits,
               const std::vector<size_t> &wires, bool inverse,
               const std::vector<PrecisionT> &params) {
+        namespace Internal = Pennylane::Internal;
         constexpr size_t num_params =
             static_lookup<gate_op>(Constant::gate_num_params);
         constexpr auto func_ptr = static_lookup<gate_op>(

pennylane_lightning/src/simulator/Measures.cpp

@@ -0,0 +1,19 @@
+// Copyright 2021 Xanadu Quantum Technologies Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//     http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "Measures.hpp"
+
+// explicit instantiation
+template class Pennylane::Measures<float, Pennylane::StateVectorRaw<float>>;
+template class Pennylane::Measures<double, Pennylane::StateVectorRaw<double>>;