From cfe24551b10f1ec80278ecd97bccc0eb1bf0b57f Mon Sep 17 00:00:00 2001 From: Amintor Dusko <87949283+AmintorDusko@users.noreply.github.com> Date: Wed, 31 Jan 2024 10:56:03 -0500 Subject: [PATCH] Update lightning qubit memory management (#601) * update dev version * update clean * expand .gitignore * expand lightning qubit class-specific Python bindings * add some statevector manipulation methods to the statevector managed simulator * update the lightning qubit Python class * add c++ unit tests * add comments to test * update lightning_qubit.py * remove obsolete tests * rename statevector instance * return some important tests * implement PR review suggestions * add review suggestion * Add semi-colon. * Fix Projector obs in L-Qubit and add Proj support in L-Kokkos. * Implement previous commit fix for None shots only. * Add tests for Proj expval/var * Auto update version * Trigger CI * remove comment * add projector support to LGPU * Format * revert changes for LGPU * skip tests for Projector observable not supported * expand tests for lightning.qubit * update changelog * Auto update version * Trigger CI * add some review suggestions * remove identities * update LKokkos and LQubit _apply_state_vector * Update tests/test_apply.py --------- Co-authored-by: Vincent Michaud-Rioux --- .github/CHANGELOG.md | 6 + .gitignore | 1 + Makefile | 1 + pennylane_lightning/core/_version.py | 2 +- .../StateVectorLQubitManaged.hpp | 42 ++- .../bindings/LQubitBindings.hpp | 72 ++++- .../tests/Test_StateVectorLQubitManaged.cpp | 59 +++- .../lightning_kokkos/lightning_kokkos.py | 18 +- .../lightning_qubit/lightning_qubit.py | 170 +++++------ tests/test_adjoint_jacobian.py | 6 +- tests/test_apply.py | 282 +----------------- tests/test_expval.py | 26 +- tests/test_var.py | 24 ++ 13 files changed, 311 insertions(+), 398 deletions(-) diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md index a81cc96ea8..be57b94646 100644 --- a/.github/CHANGELOG.md +++ b/.github/CHANGELOG.md @@ -9,6 +9,12 @@ ### Improvements +* Decouple LightningQubit memory ownership from numpy and migrate it to Lightning-Qubit managed state-vector class. + [(#601)](https://github.com/PennyLaneAI/pennylane-lightning/pull/601) + +* Expand support for Projector observables on Lightning-Kokkos. + [(#601)](https://github.com/PennyLaneAI/pennylane-lightning/pull/601) + * Split Docker build cron job into two jobs: master and latest. This is mainly for reporting in the `plugin-test-matrix` repo. [(#600)](https://github.com/PennyLaneAI/pennylane-lightning/pull/600) diff --git a/.gitignore b/.gitignore index 7c5a86618e..df49993fc2 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ doc/code/api/ PennyLane_Lightning.egg-info/ PennyLane_Lightning_Kokkos.egg-info/ build/ +build_lightning_*/ Build/ BuildCov/ BuildGBench/ diff --git a/Makefile b/Makefile index 7ad687c448..0e9031f427 100644 --- a/Makefile +++ b/Makefile @@ -59,6 +59,7 @@ help: clean: find . -type d -name '__pycache__' -exec rm -r {} \+ rm -rf build Build BuildTests BuildTidy BuildGBench + rm -rf build_* rm -rf .coverage coverage_html_report/ rm -rf pennylane_lightning/*_ops* diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py index 3676e2aa5a..3c07426591 100644 --- a/pennylane_lightning/core/_version.py +++ b/pennylane_lightning/core/_version.py @@ -16,4 +16,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "0.35.0-dev9" +__version__ = "0.35.0-dev10" diff --git a/pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubitManaged.hpp b/pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubitManaged.hpp index add76efeb0..727b431b63 100644 --- a/pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubitManaged.hpp +++ b/pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubitManaged.hpp @@ -19,6 +19,7 @@ #pragma once +#include // fill #include #include @@ -71,12 +72,12 @@ class StateVectorLQubitManaged final * @param memory_model Memory model the statevector will use */ explicit StateVectorLQubitManaged( - size_t num_qubits, Threading threading = Threading::SingleThread, + std::size_t num_qubits, Threading threading = Threading::SingleThread, CPUMemoryModel memory_model = bestCPUMemoryModel()) : BaseType{num_qubits, threading, memory_model}, data_{exp2(num_qubits), ComplexT{0.0, 0.0}, getAllocator(this->memory_model_)} { - data_[0] = {1, 0}; + setBasisState(0U); } /** @@ -102,7 +103,7 @@ class StateVectorLQubitManaged final * @param threading Threading option the statevector to use * @param memory_model Memory model the statevector will use */ - StateVectorLQubitManaged(const ComplexT *other_data, size_t other_size, + StateVectorLQubitManaged(const ComplexT *other_data, std::size_t other_size, Threading threading = Threading::SingleThread, CPUMemoryModel memory_model = bestCPUMemoryModel()) : BaseType(log2PerfectPower(other_size), threading, memory_model), @@ -139,6 +140,39 @@ class StateVectorLQubitManaged final ~StateVectorLQubitManaged() = default; + /** + * @brief Prepares a single computational basis state. + * + * @param index Index of the target element. + */ + void setBasisState(const std::size_t index) { + std::fill(data_.begin(), data_.end(), 0); + data_[index] = {1, 0}; + } + + /** + * @brief Set values for a batch of elements of the state-vector. + * + * @param values Values to be set for the target elements. + * @param indices Indices of the target elements. + */ + void setStateVector(const std::vector &indices, + const std::vector &values) { + for (std::size_t n = 0; n < indices.size(); n++) { + data_[indices[n]] = values[n]; + } + } + + /** + * @brief Reset the data back to the \f$\ket{0}\f$ state. + * + */ + void resetStateVector() { + if (this->getLength() > 0) { + setBasisState(0U); + } + } + [[nodiscard]] auto getData() -> ComplexT * { return data_.data(); } [[nodiscard]] auto getData() const -> const ComplexT * { @@ -164,7 +198,7 @@ class StateVectorLQubitManaged final * @param new_data data pointer to new data. * @param new_size size of underlying data storage. */ - void updateData(const ComplexT *new_data, size_t new_size) { + void updateData(const ComplexT *new_data, std::size_t new_size) { PL_ASSERT(data_.size() == new_size); std::copy(new_data, new_data + new_size, data_.data()); } diff --git a/pennylane_lightning/core/src/simulators/lightning_qubit/bindings/LQubitBindings.hpp b/pennylane_lightning/core/src/simulators/lightning_qubit/bindings/LQubitBindings.hpp index 147e35c3aa..74352a08aa 100644 --- a/pennylane_lightning/core/src/simulators/lightning_qubit/bindings/LQubitBindings.hpp +++ b/pennylane_lightning/core/src/simulators/lightning_qubit/bindings/LQubitBindings.hpp @@ -26,7 +26,7 @@ #include "GateOperation.hpp" #include "MeasurementsLQubit.hpp" #include "ObservablesLQubit.hpp" -#include "StateVectorLQubitRaw.hpp" +#include "StateVectorLQubitManaged.hpp" #include "TypeList.hpp" #include "VectorJacobianProduct.hpp" @@ -36,7 +36,7 @@ using namespace Pennylane::Bindings; using namespace Pennylane::LightningQubit::Algorithms; using namespace Pennylane::LightningQubit::Measures; using namespace Pennylane::LightningQubit::Observables; -using Pennylane::LightningQubit::StateVectorLQubitRaw; +using Pennylane::LightningQubit::StateVectorLQubitManaged; } // namespace /// @endcond @@ -44,8 +44,8 @@ namespace py = pybind11; namespace Pennylane::LightningQubit { using StateVectorBackends = - Pennylane::Util::TypeList, - StateVectorLQubitRaw, void>; + Pennylane::Util::TypeList, + StateVectorLQubitManaged, void>; /** * @brief Get a gate kernel map for a statevector. @@ -162,12 +162,68 @@ void registerControlledGate(PyClass &pyclass) { */ template void registerBackendClassSpecificBindings(PyClass &pyclass) { + using PrecisionT = + typename StateVectorT::PrecisionT; // Statevector's precision + using ComplexT = typename StateVectorT::ComplexT; + using ParamT = PrecisionT; // Parameter's data precision + using np_arr_c = py::array_t, + py::array::c_style | py::array::forcecast>; + registerGatesForStateVector(pyclass); registerControlledGate(pyclass); - pyclass.def("applyControlledMatrix", &applyControlledMatrix, - "Apply controlled operation"); - pyclass.def("kernel_map", &svKernelMap, - "Get internal kernels for operations"); + + pyclass + .def(py::init([](std::size_t num_qubits) { + return new StateVectorT(num_qubits); + })) + .def("resetStateVector", &StateVectorT::resetStateVector) + .def( + "setBasisState", + [](StateVectorT &sv, const size_t index) { + sv.setBasisState(index); + }, + "Create Basis State.") + .def( + "setStateVector", + [](StateVectorT &sv, const std::vector &indices, + const np_arr_c &state) { + const auto buffer = state.request(); + std::vector state_in; + if (buffer.size) { + const auto ptr = static_cast(buffer.ptr); + state_in = std::vector{ptr, ptr + buffer.size}; + } + sv.setStateVector(indices, state_in); + }, + "Set State Vector with values and their corresponding indices") + .def( + "getState", + [](const StateVectorT &sv, np_arr_c &state) { + py::buffer_info numpyArrayInfo = state.request(); + auto *data_ptr = + static_cast *>(numpyArrayInfo.ptr); + if (state.size()) { + std::copy(sv.getData(), sv.getData() + sv.getLength(), + data_ptr); + } + }, + "Copy StateVector data into a Numpy array.") + .def( + "UpdateData", + [](StateVectorT &device_sv, const np_arr_c &state) { + const py::buffer_info numpyArrayInfo = state.request(); + auto *data_ptr = static_cast(numpyArrayInfo.ptr); + const auto length = + static_cast(numpyArrayInfo.shape[0]); + if (length) { + device_sv.updateData(data_ptr, length); + } + }, + "Copy StateVector data into a Numpy array.") + .def("applyControlledMatrix", &applyControlledMatrix, + "Apply controlled operation") + .def("kernel_map", &svKernelMap, + "Get internal kernels for operations"); } /** diff --git a/pennylane_lightning/core/src/simulators/lightning_qubit/tests/Test_StateVectorLQubitManaged.cpp b/pennylane_lightning/core/src/simulators/lightning_qubit/tests/Test_StateVectorLQubitManaged.cpp index 127cfe3df3..45b97d1b4f 100644 --- a/pennylane_lightning/core/src/simulators/lightning_qubit/tests/Test_StateVectorLQubitManaged.cpp +++ b/pennylane_lightning/core/src/simulators/lightning_qubit/tests/Test_StateVectorLQubitManaged.cpp @@ -21,7 +21,8 @@ #include -#include "LinearAlgebra.hpp" //randomUnitary +#include "CPUMemoryModel.hpp" // getBestAllocator +#include "LinearAlgebra.hpp" //randomUnitary #include "StateVectorLQubitManaged.hpp" #include "StateVectorLQubitRaw.hpp" #include "TestHelpers.hpp" // createRandomStateVectorData, TestVector @@ -107,4 +108,58 @@ TEMPLATE_TEST_CASE("StateVectorLQubitManaged::StateVectorLQubitManaged", REQUIRE(sv.getDataVector() == approx(st_data)); } -} \ No newline at end of file +} + +TEMPLATE_TEST_CASE("StateVectorLQubitManaged::setBasisState", + "[StateVectorLQubitManaged]", float, double) { + using PrecisionT = TestType; + using ComplexT = std::complex; + using TestVectorT = TestVector; + + const std::size_t num_qubits = 3; + + SECTION("Prepares a single computational basis state.") { + TestVectorT init_state = + createRandomStateVectorData(re, num_qubits); + + TestVectorT expected_state(size_t{1U} << num_qubits, 0.0, + getBestAllocator()); + std::size_t index = GENERATE(0, 1, 2, 3, 4, 5, 6, 7); + expected_state[index] = {1.0, 0.0}; + StateVectorLQubitManaged sv(init_state); + sv.setBasisState(index); + + REQUIRE(sv.getDataVector() == approx(expected_state)); + } +} + +TEMPLATE_TEST_CASE("StateVectorLQubitManaged::SetStateVector", + "[StateVectorLQubitManaged]", float, double) { + using PrecisionT = TestType; + using ComplexT = std::complex; + using TestVectorT = TestVector; + + const std::size_t num_qubits = 3; + + SECTION("Set state vector with values and indices") { + TestVectorT init_state = + createRandomStateVectorData(re, num_qubits); + + auto expected_state = init_state; + + for (size_t i = 0; i < Pennylane::Util::exp2(num_qubits - 1); i++) { + std::swap(expected_state[i * 2], expected_state[i * 2 + 1]); + } + + std::vector indices = {0, 2, 4, 6, 1, 3, 5, 7}; + + std::vector values = { + init_state[1], init_state[3], init_state[5], init_state[7], + init_state[0], init_state[2], init_state[4], init_state[6]}; + + StateVectorLQubitManaged sv{num_qubits}; + sv.setStateVector(indices, values); + + REQUIRE(sv.getDataVector() == approx(expected_state)); + } +} diff --git a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py index f3153dcfe9..401a6d6de0 100644 --- a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py +++ b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py @@ -150,6 +150,7 @@ def _kokkos_configuration(): "Hadamard", "Hermitian", "Identity", + "Projector", "SparseHamiltonian", "Hamiltonian", "Sum", @@ -196,7 +197,7 @@ def __init__( shots=None, batch_obs=False, kokkos_args=None, - ): # pylint: disable=unused-argument + ): # pylint: disable=unused-argument, too-many-arguments super().__init__(wires, shots=shots, c_dtype=c_dtype) if kokkos_args is None: @@ -213,10 +214,6 @@ def __init__( if not LightningKokkos.kokkos_config: LightningKokkos.kokkos_config = _kokkos_configuration() - # Create the initial state. Internally, we store the - # state as an array of dimension [2]*wires. - self._pre_rotated_state = _kokkos_dtype(c_dtype)(self.num_wires) - @staticmethod def _asarray(arr, dtype=None): arr = np.asarray(arr) # arr is not copied @@ -347,9 +344,8 @@ def _apply_state_vector(self, state, device_wires): """ if isinstance(state, self._kokkos_state.__class__): - state_data = np.zeros(state.size, dtype=self.C_DTYPE) - state_data = self._asarray(state_data, dtype=self.C_DTYPE) - state.DeviceToHost(state_data.ravel(order="C")) + state_data = allocate_aligned_array(state.size, np.dtype(self.C_DTYPE), True) + state.DeviceToHost(state_data) state = state_data ravelled_indices, state = self._preprocess_state_vector(state, device_wires) @@ -468,6 +464,9 @@ def expval(self, observable, shot_range=None, bin_size=None): if observable.name in [ "Projector", ]: + if self.shots is None: + qs = qml.tape.QuantumScript([], [qml.expval(observable)]) + self.apply(self._get_diagonalizing_gates(qs)) return super().expval(observable, shot_range=shot_range, bin_size=bin_size) if self.shots is not None: @@ -528,6 +527,9 @@ def var(self, observable, shot_range=None, bin_size=None): if observable.name in [ "Projector", ]: + if self.shots is None: + qs = qml.tape.QuantumScript([], [qml.var(observable)]) + self.apply(self._get_diagonalizing_gates(qs)) return super().var(observable, shot_range=shot_range, bin_size=bin_size) if self.shots is not None: diff --git a/pennylane_lightning/lightning_qubit/lightning_qubit.py b/pennylane_lightning/lightning_qubit/lightning_qubit.py index 860f9bb2dd..f6ff7d2fea 100644 --- a/pennylane_lightning/lightning_qubit/lightning_qubit.py +++ b/pennylane_lightning/lightning_qubit/lightning_qubit.py @@ -75,6 +75,11 @@ VectorJacobianProductC128, ) + def _state_dtype(dtype): + if dtype not in [np.complex128, np.complex64]: # pragma: no cover + raise ValueError(f"Data type is not supported for state-vector computation: {dtype}") + return StateVectorC128 if dtype == np.complex128 else StateVectorC64 + allowed_operations = { "Identity", "BasisState", @@ -226,10 +231,7 @@ def __init__( # pylint: disable=too-many-arguments # Create the initial state. Internally, we store the # state as an array of dimension [2]*wires. - self._state = self._create_basis_state(0) - self._pre_rotated_state = self._state - self._c_dtype = c_dtype - + self._qubit_state = _state_dtype(c_dtype)(self.num_wires) self._batch_obs = batch_obs self._mcmc = mcmc if self._mcmc: @@ -274,24 +276,16 @@ def _asarray(arr, dtype=None): def _create_basis_state(self, index): """Return a computational basis state over all wires. Args: - index (int): integer representing the computational basis state - Returns: - array[complex]: complex array of shape ``[2]*self.num_wires`` - representing the statevector of the basis state - Note: This function does not support broadcasted inputs yet. + index (int): integer representing the computational basis state. """ - state = allocate_aligned_array(2**self.num_wires, np.dtype(self.C_DTYPE), True) - state[index] = 1 - return self._reshape(state, [2] * self.num_wires) + self._qubit_state.setBasisState(index) def reset(self): """Reset the device""" super().reset() # init the state vector to |00..0> - if not self.state[0] == 1.0 + 0j: - self._state = self._create_basis_state(0) - self._pre_rotated_state = self._state + self._qubit_state.resetStateVector() @property def create_ops_list(self): @@ -301,25 +295,32 @@ def create_ops_list(self): @property def measurements(self): """Returns a Measurements object matching ``use_csingle`` precision.""" - ket = np.ravel(self._state) - state_vector = StateVectorC64(ket) if self.use_csingle else StateVectorC128(ket) return ( - MeasurementsC64(state_vector) + MeasurementsC64(self.state_vector) if self.use_csingle - else MeasurementsC128(state_vector) + else MeasurementsC128(self.state_vector) ) @property def state(self): - """Returns the flattened state vector.""" - shape = (1 << self.num_wires,) - return self._reshape(self._pre_rotated_state, shape) + """Copy the state vector data to a numpy array. + + **Example** + + >>> dev = qml.device('lightning.kokkos', wires=1) + >>> dev.apply([qml.PauliX(wires=[0])]) + >>> print(dev.state) + [0.+0.j 1.+0.j] + """ + state = np.zeros(2**self.num_wires, dtype=self.C_DTYPE) + state = self._asarray(state, dtype=self.C_DTYPE) + self._qubit_state.getState(state) + return state @property def state_vector(self): - """Returns a handle to a StateVector object matching ``use_csingle`` precision.""" - ket = np.ravel(self._state) - return StateVectorC64(ket) if self.use_csingle else StateVectorC128(ket) + """Returns a handle to the statevector.""" + return self._qubit_state def _apply_state_vector(self, state, device_wires): """Initialize the internal state vector in a specified state. @@ -328,6 +329,12 @@ def _apply_state_vector(self, state, device_wires): or broadcasted state of shape ``(batch_size, 2**len(wires))`` device_wires (Wires): wires that get initialized in the state """ + + if isinstance(state, self._qubit_state.__class__): + state_data = allocate_aligned_array(state.size, np.dtype(self.C_DTYPE), True) + self._qubit_state.getState(state_data) + state = state_data + ravelled_indices, state = self._preprocess_state_vector(state, device_wires) # translate to wire labels used by device @@ -336,12 +343,11 @@ def _apply_state_vector(self, state, device_wires): if len(device_wires) == self.num_wires and Wires(sorted(device_wires)) == device_wires: # Initialize the entire device state with the input state - self._state = self._reshape(state, output_shape) + state = self._reshape(state, output_shape).ravel(order="C") + self._qubit_state.UpdateData(state) return - state = self._scatter(ravelled_indices, state, [2**self.num_wires]) - state = self._reshape(state, output_shape) - self._state = self._asarray(state, dtype=self.C_DTYPE) + self._qubit_state.setStateVector(ravelled_indices, state) # this operation on device def _apply_basis_state(self, state, wires): """Initialize the state vector in a specified computational basis state. @@ -355,22 +361,23 @@ def _apply_basis_state(self, state, wires): Note: This function does not support broadcasted inputs yet. """ num = self._get_basis_state_index(state, wires) - self._state = self._create_basis_state(num) + self._create_basis_state(num) - def _apply_lightning_controlled(self, sim, operation): + def _apply_lightning_controlled(self, operation): """Apply an arbitrary controlled operation to the state tensor. Args: - sim (StateVectorC64, StateVectorC128): a state vector simulator operation (~pennylane.operation.Operation): operation to apply Returns: array[complex]: the output state tensor """ + state = self.state_vector + basename = "PauliX" if operation.name == "MultiControlledX" else operation.base.name if basename == "Identity": return - method = getattr(sim, f"{basename}", None) + method = getattr(state, f"{basename}", None) control_wires = self.wires.indices(operation.control_wires) control_values = ( [bool(int(i)) for i in operation.hyperparameters["control_values"]] @@ -386,7 +393,7 @@ def _apply_lightning_controlled(self, sim, operation): param = operation.parameters method(control_wires, control_values, target_wires, inv, param) else: # apply gate as an n-controlled matrix - method = getattr(sim, "applyControlledMatrix") + method = getattr(state, "applyControlledMatrix") target_wires = self.wires.indices(operation.target_wires) try: method( @@ -402,27 +409,24 @@ def _apply_lightning_controlled(self, sim, operation): operation.base.matrix, control_wires, control_values, target_wires, False ) - def apply_lightning(self, state, operations): + def apply_lightning(self, operations): """Apply a list of operations to the state tensor. Args: - state (array[complex]): the input state tensor operations (list[~pennylane.operation.Operation]): operations to apply Returns: array[complex]: the output state tensor """ - state_vector = np.ravel(state) - sim = ( - StateVectorC64(state_vector) if self.use_csingle else StateVectorC128(state_vector) - ) + state = self.state_vector # Skip over identity operations instead of performing # matrix multiplication with it. for operation in operations: - if operation.name == "Identity": + name = operation.name + if name == "Identity": continue - method = getattr(sim, operation.name, None) + method = getattr(state, name, None) wires = self.wires.indices(operation.wires) if method is not None: # apply specialized gate @@ -430,23 +434,21 @@ def apply_lightning(self, state, operations): param = operation.parameters method(wires, inv, param) elif ( - operation.name[0:2] == "C(" - or operation.name == "ControlledQubitUnitary" - or operation.name == "MultiControlledX" + name[0:2] == "C(" + or name == "ControlledQubitUnitary" + or name == "MultiControlledX" ): # apply n-controlled gate - self._apply_lightning_controlled(sim, operation) + self._apply_lightning_controlled(operation) else: # apply gate as a matrix # Inverse can be set to False since qml.matrix(operation) is already in # inverted form - method = getattr(sim, "applyMatrix") + method = getattr(state, "applyMatrix") try: method(qml.matrix(operation), wires, False) except AttributeError: # pragma: no cover # To support older versions of PL method(operation.matrix, wires, False) - return np.reshape(state_vector, state.shape) - # pylint: disable=unused-argument def apply(self, operations, rotations=None, **kwargs): """Applies operations to the state vector.""" @@ -468,15 +470,7 @@ def apply(self, operations, rotations=None, **kwargs): f"Operations have already been applied on a {self.short_name} device." ) - if operations: - self._pre_rotated_state = self.apply_lightning(self._state, operations) - else: - self._pre_rotated_state = self._state - - if rotations: - self._state = self.apply_lightning(np.copy(self._pre_rotated_state), rotations) - else: - self._state = self._pre_rotated_state + self.apply_lightning(operations) # pylint: disable=protected-access def expval(self, observable, shot_range=None, bin_size=None): @@ -494,9 +488,11 @@ def expval(self, observable, shot_range=None, bin_size=None): Expectation value of the observable """ if observable.name in [ - "Identity", "Projector", ]: + if self.shots is None: + qs = qml.tape.QuantumScript([], [qml.expval(observable)]) + self.apply(self._get_diagonalizing_gates(qs)) return super().expval(observable, shot_range=shot_range, bin_size=bin_size) if self.shots is not None: @@ -505,14 +501,10 @@ def expval(self, observable, shot_range=None, bin_size=None): samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size) return np.squeeze(np.mean(samples, axis=0)) - # Initialization of state - ket = np.ravel(self._pre_rotated_state) - - state_vector = StateVectorC64(ket) if self.use_csingle else StateVectorC128(ket) measurements = ( - MeasurementsC64(state_vector) + MeasurementsC64(self.state_vector) if self.use_csingle - else MeasurementsC128(state_vector) + else MeasurementsC128(self.state_vector) ) if observable.name == "SparseHamiltonian": csr_hamiltonian = observable.sparse_matrix(wire_order=self.wires).tocsr(copy=False) @@ -552,9 +544,11 @@ def var(self, observable, shot_range=None, bin_size=None): Variance of the observable """ if observable.name in [ - "Identity", "Projector", ]: + if self.shots is None: + qs = qml.tape.QuantumScript([], [qml.var(observable)]) + self.apply(self._get_diagonalizing_gates(qs)) return super().var(observable, shot_range=shot_range, bin_size=bin_size) if self.shots is not None: @@ -563,14 +557,10 @@ def var(self, observable, shot_range=None, bin_size=None): samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size) return np.squeeze(np.var(samples, axis=0)) - # Initialization of state - ket = np.ravel(self._pre_rotated_state) - - state_vector = StateVectorC64(ket) if self.use_csingle else StateVectorC128(ket) measurements = ( - MeasurementsC64(state_vector) + MeasurementsC64(self.state_vector) if self.use_csingle - else MeasurementsC128(state_vector) + else MeasurementsC128(self.state_vector) ) if observable.name == "SparseHamiltonian": @@ -603,14 +593,10 @@ def generate_samples(self): 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) measurements = ( - MeasurementsC64(state_vector) + MeasurementsC64(self.state_vector) if self.use_csingle - else MeasurementsC128(state_vector) + else MeasurementsC128(self.state_vector) ) if self._mcmc: return measurements.generate_mcmc_samples( @@ -630,13 +616,22 @@ def probability_lightning(self, wires): Returns: array[float]: list of the probabilities """ - state_vector = self.state_vector return ( - MeasurementsC64(state_vector) + MeasurementsC64(self.state_vector) if self.use_csingle - else MeasurementsC128(state_vector) + else MeasurementsC128(self.state_vector) ).probs(wires) + # pylint: disable=attribute-defined-outside-init + def sample(self, observable, shot_range=None, bin_size=None, counts=False): + """Return samples of an observable.""" + if observable.name != "PauliZ": + self.apply_lightning(observable.diagonalizing_gates()) + self._samples = self.generate_samples() + return super().sample( + observable, shot_range=shot_range, bin_size=bin_size, counts=counts + ) + @staticmethod def _check_adjdiff_supported_measurements( measurements: List[MeasurementProcess], @@ -700,14 +695,11 @@ def _init_process_jacobian_tape(self, tape, starting_state, use_device_state): "The number of qubits of starting_state must be the same as " "that of the device." ) - ket = self._asarray(starting_state, dtype=self.C_DTYPE) - else: - if not use_device_state: - self.reset() - self.apply(tape.operations) - ket = self._pre_rotated_state - ket = ket.reshape(-1) - return StateVectorC64(ket) if self.use_csingle else StateVectorC128(ket) + self._apply_state_vector(starting_state, self.wires) + elif not use_device_state: + self.reset() + self.apply(tape.operations) + return self.state_vector def adjoint_jacobian(self, tape, starting_state=None, use_device_state=False): """Computes and returns the Jacobian with the adjoint method.""" diff --git a/tests/test_adjoint_jacobian.py b/tests/test_adjoint_jacobian.py index 58a156aa4a..86c0a9e246 100644 --- a/tests/test_adjoint_jacobian.py +++ b/tests/test_adjoint_jacobian.py @@ -583,11 +583,9 @@ def test_provide_starting_state(self, tol, dev): dM1 = dev.adjoint_jacobian(tape) - if device_name == "lightning.kokkos": - dev._pre_rotated_state = dev.state_vector # necessary for lightning.kokkos - + if device_name in ["lightning.kokkos", "lightning.qubit"]: qml.execute([tape], dev, None) - dM2 = dev.adjoint_jacobian(tape, starting_state=dev._pre_rotated_state) + dM2 = dev.adjoint_jacobian(tape, starting_state=dev.state_vector) assert np.allclose(dM1, dM2, atol=tol, rtol=0) else: diff --git a/tests/test_apply.py b/tests/test_apply.py index a8075a8176..c889eb6e18 100644 --- a/tests/test_apply.py +++ b/tests/test_apply.py @@ -23,6 +23,7 @@ import pennylane as qml from pennylane import DeviceError from pennylane.operation import Operation +import functools class TestApply: @@ -57,7 +58,6 @@ def test_apply_operation_single_wire_no_parameters( ): """Tests that applying an operation yields the expected output state for single wire operations that have no parameters.""" - from pennylane.wires import Wires dev = qubit_device(wires=1) _state = np.array(input).astype(dev.C_DTYPE) @@ -1186,286 +1186,6 @@ def circuit(): assert np.allclose(res_probs, expected_prob, atol=tol, rtol=0) -@pytest.mark.skipif( - device_name == "lightning.kokkos" or device_name == "lightning.gpu", - reason="lightning.kokkos/gpu does not support apply with rotations.", -) -@pytest.mark.parametrize("theta,phi,varphi", list(zip(THETA, PHI, VARPHI))) -class TestTensorExpval: - """Test tensor expectation values""" - - def test_paulix_pauliy(self, theta, phi, varphi, qubit_device, tol): - """Test that a tensor product involving PauliX and PauliY works correctly""" - dev = qubit_device(wires=3) - dev.reset() - - obs = qml.PauliX(0) @ qml.PauliY(2) - - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.expval(obs) - - expected = np.sin(theta) * np.sin(phi) * np.sin(varphi) - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_pauliz_identity(self, theta, phi, varphi, qubit_device, tol): - """Test that a tensor product involving PauliZ and Identity works correctly""" - dev = qubit_device(wires=3) - dev.reset() - - obs = qml.PauliZ(0) @ qml.Identity(1) @ qml.PauliZ(2) - - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.expval(obs) - - expected = np.cos(varphi) * np.cos(phi) - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_pauliz_hadamard(self, theta, phi, varphi, qubit_device, tol): - """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly""" - dev = qubit_device(wires=3) - obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2) - - dev.reset() - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.expval(obs) - - expected = -(np.cos(varphi) * np.sin(phi) + np.sin(varphi) * np.cos(theta)) / np.sqrt(2) - - assert np.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.skipif( - device_name == "lightning.kokkos" or device_name == "lightning.gpu", - reason="lightning.kokkos/gpu does not support apply with rotations.", -) -@pytest.mark.parametrize("theta, phi, varphi", list(zip(THETA, PHI, VARPHI))) -class TestTensorVar: - """Tests for variance of tensor observables""" - - def test_paulix_pauliy(self, theta, phi, varphi, qubit_device, tol): - """Test that a tensor product involving PauliX and PauliY works correctly""" - dev = qubit_device(wires=3) - obs = qml.PauliX(0) @ qml.PauliY(2) - - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.var(obs) - - expected = ( - 8 * np.sin(theta) ** 2 * np.cos(2 * varphi) * np.sin(phi) ** 2 - - np.cos(2 * (theta - phi)) - - np.cos(2 * (theta + phi)) - + 2 * np.cos(2 * theta) - + 2 * np.cos(2 * phi) - + 14 - ) / 16 - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_pauliz_hadamard(self, theta, phi, varphi, qubit_device, tol): - """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly""" - dev = qubit_device(wires=3) - obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2) - - dev.reset() - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.var(obs) - - expected = ( - 3 - + np.cos(2 * phi) * np.cos(varphi) ** 2 - - np.cos(2 * theta) * np.sin(varphi) ** 2 - - 2 * np.cos(theta) * np.sin(phi) * np.sin(2 * varphi) - ) / 4 - - assert np.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.skipif( - device_name == "lightning.kokkos" or device_name == "lightning.gpu", - reason="lightning.kokkos/lightning.gpu does not support apply with rotations.", -) -@pytest.mark.parametrize("theta, phi, varphi", list(zip(THETA, PHI, VARPHI))) -@pytest.mark.parametrize("shots", [None, 100000]) -class TestTensorSample: - """Test sampling tensor the tensor product of observables""" - - def test_paulix_pauliy(self, theta, phi, varphi, shots, tol): - """Test that a tensor product involving PauliX and PauliY works correctly""" - tolerance = tol if shots is None else TOL_STOCHASTIC - dev = qml.device(device_name, wires=3, shots=shots) - - obs = qml.PauliX(0) @ qml.PauliY(2) - - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - dev._wires_measured = {0, 1, 2} - dev._samples = dev.generate_samples() if shots is not None else None - - s1 = qml.eigvals(obs) - p = dev.probability(wires=obs.wires) - - # s1 should only contain 1 and -1 - assert np.allclose(s1**2, 1, atol=tolerance, rtol=0) - - mean = s1 @ p - expected = np.sin(theta) * np.sin(phi) * np.sin(varphi) - assert np.allclose(mean, expected, atol=tolerance, rtol=0) - - var = (s1**2) @ p - (s1 @ p).real ** 2 - expected = ( - 8 * np.sin(theta) ** 2 * np.cos(2 * varphi) * np.sin(phi) ** 2 - - np.cos(2 * (theta - phi)) - - np.cos(2 * (theta + phi)) - + 2 * np.cos(2 * theta) - + 2 * np.cos(2 * phi) - + 14 - ) / 16 - assert np.allclose(var, expected, atol=tolerance, rtol=0) - - def test_pauliz_hadamard(self, theta, phi, varphi, shots, qubit_device, tol): - """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly""" - tolerance = tol if shots is None else TOL_STOCHASTIC - dev = qubit_device(wires=3) - obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2) - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - dev._wires_measured = {0, 1, 2} - dev._samples = dev.generate_samples() if dev.shots is not None else None - - s1 = qml.eigvals(obs) - p = dev.marginal_prob(dev.probability(), wires=obs.wires) - - # s1 should only contain 1 and -1 - assert np.allclose(s1**2, 1, atol=tol, rtol=0) - - mean = s1 @ p - expected = -(np.cos(varphi) * np.sin(phi) + np.sin(varphi) * np.cos(theta)) / np.sqrt(2) - assert np.allclose(mean, expected, atol=tol, rtol=0) - - var = (s1**2) @ p - (s1 @ p).real ** 2 - expected = ( - 3 - + np.cos(2 * phi) * np.cos(varphi) ** 2 - - np.cos(2 * theta) * np.sin(varphi) ** 2 - - 2 * np.cos(theta) * np.sin(phi) * np.sin(2 * varphi) - ) / 4 - assert np.allclose(var, expected, atol=tolerance, rtol=0) - - def test_qubitunitary_rotation_hadamard(self, theta, phi, varphi, shots, qubit_device, tol): - """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly""" - tolerance = tol if shots is None else TOL_STOCHASTIC - dev = qubit_device(wires=3) - obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2) - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - [ - qml.QubitUnitary( - qml.matrix(obs.diagonalizing_gates()[0]), - wires=obs.diagonalizing_gates()[0].wires, - ), - *obs.diagonalizing_gates()[1:], - ], - ) - - dev._wires_measured = {0, 1, 2} - dev._samples = dev.generate_samples() if dev.shots is not None else None - - s1 = qml.eigvals(obs) - p = dev.marginal_prob(dev.probability(), wires=obs.wires) - - # s1 should only contain 1 and -1 - assert np.allclose(s1**2, 1, atol=tol, rtol=0) - - mean = s1 @ p - expected = -(np.cos(varphi) * np.sin(phi) + np.sin(varphi) * np.cos(theta)) / np.sqrt(2) - assert np.allclose(mean, expected, atol=tol, rtol=0) - - var = (s1**2) @ p - (s1 @ p).real ** 2 - expected = ( - 3 - + np.cos(2 * phi) * np.cos(varphi) ** 2 - - np.cos(2 * theta) * np.sin(varphi) ** 2 - - 2 * np.cos(theta) * np.sin(phi) * np.sin(2 * varphi) - ) / 4 - assert np.allclose(var, expected, atol=tolerance, rtol=0) - - class TestApplyLightningMethod: """Unit tests for the apply_lightning method.""" diff --git a/tests/test_expval.py b/tests/test_expval.py index e64ef2a329..561a9ed235 100644 --- a/tests/test_expval.py +++ b/tests/test_expval.py @@ -106,9 +106,33 @@ def test_hadamard_expectation(self, theta, phi, qubit_device, tol): ) / np.sqrt(2) assert np.allclose(res, expected, tol) + def test_projector_expectation(self, theta, phi, qubit_device, tol): + """Test that Projector variance value is correct""" + n_qubits = 2 + dev_def = qml.device("default.qubit", wires=n_qubits) + dev = qubit_device(wires=n_qubits) + + if "Projector" not in dev.observables: + pytest.skip("Device does not support the Projector observable.") + + init_state = np.random.rand(2**n_qubits) + 1j * np.random.rand(2**n_qubits) + init_state /= np.sqrt(np.dot(np.conj(init_state), init_state)) + obs = qml.Projector(np.array([0, 1, 0, 0]) / np.sqrt(2), wires=[0, 1]) + + def circuit(): + qml.StatePrep(init_state, wires=range(n_qubits)) + qml.RY(theta, wires=[0]) + qml.RY(phi, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(obs) + + circ = qml.QNode(circuit, dev) + circ_def = qml.QNode(circuit, dev_def) + assert np.allclose(circ(), circ_def(), tol) + @pytest.mark.parametrize("n_wires", range(1, 7)) def test_hermitian_expectation(self, n_wires, theta, phi, qubit_device, tol): - """Test that Hadamard expectation value is correct""" + """Test that Hermitian expectation value is correct""" n_qubits = 7 dev_def = qml.device("default.qubit", wires=n_qubits) dev = qubit_device(wires=n_qubits) diff --git a/tests/test_var.py b/tests/test_var.py index 76946705d0..ccee5cda20 100644 --- a/tests/test_var.py +++ b/tests/test_var.py @@ -47,6 +47,30 @@ def test_var(self, theta, phi, qubit_device, tol): assert np.allclose(var, expected, tol) + def test_projector_var(self, theta, phi, qubit_device, tol): + """Test that Projector variance value is correct""" + n_qubits = 2 + dev_def = qml.device("default.qubit", wires=n_qubits) + dev = qubit_device(wires=n_qubits) + + if "Projector" not in dev.observables: + pytest.skip("Device does not support the Projector observable.") + + init_state = np.random.rand(2**n_qubits) + 1j * np.random.rand(2**n_qubits) + init_state /= np.sqrt(np.dot(np.conj(init_state), init_state)) + obs = qml.Projector(np.array([0, 1, 0, 0]) / np.sqrt(2), wires=[0, 1]) + + def circuit(): + qml.StatePrep(init_state, wires=range(n_qubits)) + qml.RY(theta, wires=[0]) + qml.RY(phi, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(obs) + + circ = qml.QNode(circuit, dev) + circ_def = qml.QNode(circuit, dev_def) + assert np.allclose(circ(), circ_def(), tol) + @pytest.mark.parametrize("theta, phi, varphi", list(zip(THETA, PHI, VARPHI))) class TestTensorVar: