Computing expval and var with LightningTensor using quimb MPS (…

…#686)

* empty commit (triggering CI)

* Auto update version

* Definition of the two front-end classes

* adding the `lightning_tensor` string to the supported backends in `setup.py`

* adding `__init__.py` file to directly import the `lightning_tensor` device class

* re-naming file

* Auto update version

* Creating the first prototype of initial MPS state tensor using `quimb`

* providing the `backend`, `method` parameters and making `wires` optional

* Changing names and structure

* Auto update version

* adding method required by the new device API design

* Auto update version

* using the `kwargs` parameter in `LightningTensor` and `CircuitMPS` in `quimb`

* taking some further inputs from the new device API

* Perhaps decided the overall structure of `LIghtningTensor`

* Auto update version

* adding docs to methods

* temporary changes so that `pylint` does not complain at this stage

* running `isort`

* re-running formatter after `isort`

* re-running formatter after `isort`

* Applying suggested formatting change from CI

* adding tmp unit tests

* Adding `quimb` in `requirements.txt`

* runing `isort` on mps test

* removing `quimb` from requirement and deleting unit tests for `lightning.tensor`

* Auto update version

* re-inserting unit tests with an additional `yml` file

* running isort on quimb test

* changing name of yml file

* preventing error in import

* updating yml file

* inserting `quimb` package in requirements-dev

* strange error with `quimb`

* strange error with `quimb`

* specifying scipy version

* removing installation of scipy from yml file

* removing the new `yml` file

* testing if tests are tested

* Covering all lines in tests

* forgot final line for formatter

* Python formatter on CI complaints

* covering missing lines

* formatter on CI complaints

* Trying not to skip test if Cpp is enabled

* skipping tests if Cpp is enabled

* removing the only line not covered by tests so far

* Auto update version

* Applying suggestions from code review and making the `state` attribute private (new API design)

* Python formatter

* removing params from `QuimbMPS`

* Auto update version

* removing `**kwargs` from `QuimbMPS`

* removing unnecessary param at this stage

* covering test line

* formatter...

* removing param description

* Making `pylint` happy

* forgot new arg in test

* Updating base class and `preprocess` function

* empty commit

* Core structure (TODO: add tests)

* Running `isort`

* Updating `LightningTensor` class with new names from more advanced PR

* Adding unit tests for the `expval` calculation

* Auto update version

* Auto update version

* Triggering CI

* Adding Hamiltonian expval (`pylint` keeps complaining about my formatter)

* Auto update version

* Trying to remove pin from `quimb` in `requirements.dev`

* Auto update version

* Auto update version

* Removing infos on derivatives and using config options to pass parameters to interface

* Usual `pylint` failures

* Trying to solve formatting errors

* Style and format update

* typo in docstring

* Sunday update: improved docstrings and structure

* Support for expval and var plus some tests

* Removing comments from lines in tests

* Removing method that was supposed to be in next PR

* removing old TODO comment

* Adding API skip lines in test

* removing old TODO comment

* Forgot function from forked repo

* Removing lines in tests

* Removing changes from the `setup.py` file

* restoring previous format to `setup.py`

* Auto update version from '0.36.0-dev34' to '0.36.0-dev41'

* Auto update version from '0.36.0-dev40' to '0.36.0-dev41'

* Removing kwargs as suggested from code review

* Addressing comments from CR

* Skipping tests if CPP binary is available

* -am "[ci skip]"

* -am "[ci skip]"

* -am "[ci skip]"

* -am "[ci skip]"

* Removing the `draw` function

* Added changelog

* Added statement in docstring

* Renaming description in files and removing comments

* Typo in docstring

* Merging from master (now that the first PR has been merged)

* Auto update version from '0.36.0-dev45' to '0.36.0-dev46'

* Suggestions from CodeFactor (slipped during merging from master)

* Suggestions from code review

* Removing `Arg` from the `preprocess` function in the interface

* applying suggestion by calling `expval` in `var`

* Modified changelog and removed variations from supported operations

* Switching to `quimb 1.8.1`, adding `contract` option to kwargs, removing recursive decomposition for 3+ qubit gates

* removing `recursive decomposition` from the source dode

* Suggestions from CR

* Trying to make `Codecov` happy by inserting # pragma: no cover (although the lines are tested)

* Auto update version from '0.36.0-dev46' to '0.36.0-dev47'

* Usual formatting check that I forgot

* "[skip ci]"

* Auto update version from '0.36.0-dev47' to '0.36.0-dev48'

* Trying to make Codecov happy with # pragma: no cover, even though lines are tested indeed!

* Removing # pragma - no cover since Codecov is not happy anyway

* Auto update version from '0.37.0-dev2' to '0.37.0-dev3'

* Updating changelog

---------

Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com>
Co-authored-by: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai>

.github/CHANGELOG.md

@@ -12,6 +12,9 @@
 * Update Linux wheels to use manylinux_2_28 images.
   [(#667)](https://github.com/PennyLaneAI/pennylane-lightning/pull/667)
 
+* Add support for `qml.expval` and `qml.var` in the `lightning.tensor` device for the `quimb` interface and the MPS method.
+  [(#686)](https://github.com/PennyLaneAI/pennylane-lightning/pull/686)  
+
 ### Documentation
 
 ### Bug fixes
@@ -20,7 +23,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Amintor Dusko, Vincent Michaud-Rioux
+Amintor Dusko, Pietropaolo Frisoni, Vincent Michaud-Rioux
 
 ---
 
@@ -182,7 +185,7 @@ Amintor Dusko, Vincent Michaud-Rioux
 
 This release contains contributions from (in alphabetical order):
 
-Ali Asadi, Amintor Dusko, Thomas Germain, Christina Lee, Erick Ochoa Lopez, Vincent Michaud-Rioux, Rashid N H M, Lee James O'Riordan, Mudit Pandey, Shuli Shu
+Ali Asadi, Amintor Dusko, Pietropaolo Frisoni, Thomas Germain, Christina Lee, Erick Ochoa Lopez, Vincent Michaud-Rioux, Rashid N H M, Lee James O'Riordan, Mudit Pandey, Shuli Shu
 
 ---
 

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.37.0-dev2"
+__version__ = "0.37.0-dev3"

pennylane_lightning/lightning_tensor/backends/quimb/_mps.py

@@ -14,16 +14,104 @@
 """
 Class implementation for the Quimb MPS interface for simulating quantum circuits while keeping the state always in MPS form.
 """
+import copy
+from typing import Callable, Sequence, Union
 
-import numpy as np
 import pennylane as qml
 import quimb.tensor as qtn
+from pennylane import numpy as np
+from pennylane.devices import DefaultExecutionConfig, ExecutionConfig
+from pennylane.devices.preprocess import (
+    decompose,
+    validate_device_wires,
+    validate_measurements,
+    validate_observables,
+)
+from pennylane.measurements import ExpectationMP, MeasurementProcess, StateMeasurement, VarianceMP
+from pennylane.tape import QuantumScript, QuantumTape
+from pennylane.transforms.core import TransformProgram
+from pennylane.typing import Result, ResultBatch, TensorLike
 from pennylane.wires import Wires
 
-_operations = frozenset({})  # pragma: no cover
+Result_or_ResultBatch = Union[Result, ResultBatch]
+QuantumTapeBatch = Sequence[QuantumTape]
+QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch]
+PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch]
+
+_operations = frozenset(
+    {
+        "Identity",
+        "QubitUnitary",
+        "ControlledQubitUnitary",
+        "MultiControlledX",
+        "DiagonalQubitUnitary",
+        "PauliX",
+        "PauliY",
+        "PauliZ",
+        "MultiRZ",
+        "GlobalPhase",
+        "Hadamard",
+        "S",
+        "T",
+        "SX",
+        "CNOT",
+        "SWAP",
+        "ISWAP",
+        "PSWAP",
+        "SISWAP",
+        "SQISW",
+        "CSWAP",
+        "Toffoli",
+        "CY",
+        "CZ",
+        "PhaseShift",
+        "ControlledPhaseShift",
+        "CPhase",
+        "RX",
+        "RY",
+        "RZ",
+        "Rot",
+        "CRX",
+        "CRY",
+        "CRZ",
+        "CRot",
+        "IsingXX",
+        "IsingYY",
+        "IsingZZ",
+        "IsingXY",
+        "SingleExcitation",
+        "SingleExcitationPlus",
+        "SingleExcitationMinus",
+        "DoubleExcitation",
+        "QubitCarry",
+        "QubitSum",
+        "OrbitalRotation",
+        "QFT",
+        "ECR",
+        "BlockEncode",
+    }
+)
 # The set of supported operations.
 
-_observables = frozenset({})  # pragma: no cover
+
+_observables = frozenset(
+    {
+        "PauliX",
+        "PauliY",
+        "PauliZ",
+        "Hadamard",
+        "Hermitian",
+        "Identity",
+        "Projector",
+        "SparseHamiltonian",
+        "Hamiltonian",
+        "LinearCombination",
+        "Sum",
+        "SProd",
+        "Prod",
+        "Exp",
+    }
+)
 # The set of supported observables.
 
 
@@ -57,15 +145,15 @@ def __init__(self, num_wires, interf_opts, dtype=np.complex128):
         self._wires = Wires(range(num_wires))
         self._dtype = dtype
 
-        self._init_state_ops = {
+        self._init_state_opts = {
             "binary": "0" * max(1, len(self._wires)),
             "dtype": self._dtype.__name__,
             "tags": [str(l) for l in self._wires.labels],
         }
 
         self._gate_opts = {
-            "contract": "swap+split",
             "parametrize": None,
+            "contract": interf_opts["contract"],
             "cutoff": interf_opts["cutoff"],
             "max_bond": interf_opts["max_bond_dim"],
         }
@@ -79,22 +167,205 @@ def __init__(self, num_wires, interf_opts, dtype=np.complex128):
         self._circuitMPS = qtn.CircuitMPS(psi0=self._initial_mps())
 
     @property
-    def state(self):
-        """Current MPS handled by the interface."""
+    def interface_name(self) -> str:
+        """The name of this interface."""
+        return "QuimbMPS interface"
+
+    @property
+    def state(self) -> qtn.MatrixProductState:
+        """Return the current MPS handled by the interface."""
         return self._circuitMPS.psi
 
     def state_to_array(self) -> np.ndarray:
         """Contract the MPS into a dense array."""
         return self._circuitMPS.to_dense()
 
+    def _reset_state(self) -> None:
+        """Reset the MPS."""
+        self._circuitMPS = qtn.CircuitMPS(psi0=self._initial_mps())
+
     def _initial_mps(self) -> qtn.MatrixProductState:
         r"""
-        Returns an initial state to :math:`\ket{0}`.
+        Return an initial state to :math:`\ket{0}`.
 
         Internally, it uses `quimb`'s `MPS_computational_state` method.
 
         Returns:
             MatrixProductState: The initial MPS of a circuit.
         """
+        return qtn.MPS_computational_state(**self._init_state_opts)
+
+    def preprocess(self) -> TransformProgram:
+        """This function defines the device transform program to be applied for this interface.
+
+        Returns:
+            TransformProgram: A transform program that when called returns :class:`~.QuantumTape`'s that the
+            device can natively execute as well as a postprocessing function to be called after execution.
+
+        This interface:
+
+        * Supports any one or two-qubit operations that provide a matrix.
+        * Supports any three or four-qubit operations that provide a decomposition method.
+        * Currently does not support finite shots.
+        """
+
+        program = TransformProgram()
+
+        program.add_transform(validate_measurements, name=self.interface_name)
+        program.add_transform(validate_observables, accepted_observables, name=self.interface_name)
+        program.add_transform(validate_device_wires, self._wires, name=self.interface_name)
+        program.add_transform(
+            decompose,
+            stopping_condition=stopping_condition,
+            skip_initial_state_prep=True,
+            name=self.interface_name,
+        )
+        program.add_transform(qml.transforms.broadcast_expand)
+
+        return program
+
+    # pylint: disable=unused-argument
+    def execute(
+        self,
+        circuits: QuantumTape_or_Batch,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ) -> Result_or_ResultBatch:
+        """Execute a circuit or a batch of circuits and turn it into results.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the quantum circuits to be executed
+            execution_config (ExecutionConfig): a datastructure with additional information required for execution
+
+        Returns:
+            TensorLike, tuple[TensorLike], tuple[tuple[TensorLike]]: A numeric result of the computation.
+        """
+
+        results = []
+        for circuit in circuits:
+            circuit = circuit.map_to_standard_wires()
+            results.append(self.simulate(circuit))
+
+        return tuple(results)
+
+    def simulate(self, circuit: QuantumScript) -> Result:
+        """Simulate a single quantum script. This function assumes that all operations provide matrices.
+
+        Args:
+            circuit (QuantumScript): The single circuit to simulate.
+
+        Returns:
+            Tuple[TensorLike]: The results of the simulation.
+        """
+
+        self._reset_state()
+
+        for op in circuit.operations:
+            self._apply_operation(op)
+
+        if not circuit.shots:
+            if len(circuit.measurements) == 1:
+                return self.measurement(circuit.measurements[0])
+            return tuple(self.measurement(mp) for mp in circuit.measurements)
+
+        raise NotImplementedError  # pragma: no cover
+
+    def _apply_operation(self, op: qml.operation.Operator) -> None:
+        """Apply a single operator to the circuit, keeping the state always in a MPS form.
+
+        Internally it uses `quimb`'s `apply_gate` method.
+
+        Args:
+            op (Operator): The operation to apply.
+        """
+
+        self._circuitMPS.apply_gate(op.matrix().astype(self._dtype), *op.wires, **self._gate_opts)
+
+    def measurement(self, measurementprocess: MeasurementProcess) -> TensorLike:
+        """Measure the measurement required by the circuit over the MPS.
+
+        Args:
+            measurementprocess (MeasurementProcess): measurement to apply to the state.
+
+        Returns:
+            TensorLike: the result of the measurement.
+        """
+
+        return self._get_measurement_function(measurementprocess)(measurementprocess)
+
+    def _get_measurement_function(
+        self, measurementprocess: MeasurementProcess
+    ) -> Callable[[MeasurementProcess, TensorLike], TensorLike]:
+        """Get the appropriate method for performing a measurement.
+
+        Args:
+            measurementprocess (MeasurementProcess): measurement process to apply to the state
+
+        Returns:
+            Callable: function that returns the measurement result
+        """
+        if isinstance(measurementprocess, StateMeasurement):
+            if isinstance(measurementprocess, ExpectationMP):
+                return self.expval
+
+            if isinstance(measurementprocess, VarianceMP):
+                return self.var
+
+        raise NotImplementedError  # pragma: no cover
+
+    def expval(self, measurementprocess: MeasurementProcess) -> float:
+        """Expectation value of the supplied observable contained in the MeasurementProcess.
+
+        Args:
+            measurementprocess (StateMeasurement): measurement to apply to the MPS.
+
+        Returns:
+            Expectation value of the observable.
+        """
+
+        obs = measurementprocess.obs
+
+        result = self._local_expectation(obs.matrix(), tuple(obs.wires))
+
+        return result
+
+    def var(self, measurementprocess: MeasurementProcess) -> float:
+        """Variance of the supplied observable contained in the MeasurementProcess.
+
+        Args:
+            measurementprocess (StateMeasurement): measurement to apply to the MPS.
+
+        Returns:
+            Variance of the observable.
+        """
+
+        obs = measurementprocess.obs
+
+        obs_mat = obs.matrix()
+        expect_op = self.expval(measurementprocess)
+        expect_squar_op = self._local_expectation(obs_mat @ obs_mat.conj().T, tuple(obs.wires))
+
+        return expect_squar_op - np.square(expect_op)
+
+    def _local_expectation(self, matrix, wires) -> float:
+        """Compute the local expectation value of a matrix on the MPS.
+
+        Internally, it uses `quimb`'s `local_expectation` method.
+
+        Args:
+            matrix (array): the matrix to compute the expectation value of.
+            wires (tuple[int]): the wires the matrix acts on.
+
+        Returns:
+            Local expectation value of the matrix on the MPS.
+        """
+
+        # We need to copy the MPS to avoid modifying the original state
+        qc = copy.deepcopy(self._circuitMPS)
+
+        exp_val = qc.local_expectation(
+            matrix,
+            wires,
+            **self._expval_opts,
+        )
 
-        return qtn.MPS_computational_state(**self._init_state_ops)
+        return float(np.real(exp_val))