Computing expval and var with LightningTensor using quimb MPS

.github/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* 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)
+
 * Add Python class for the `lightning.tensor` device which uses the new device API and the interface for `quimb` based on the MPS method.
   [(#671)](https://github.com/PennyLaneAI/pennylane-lightning/pull/671)
 

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0-dev47"
+__version__ = "0.36.0-dev48"

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))