Compute expectation value (#4484)

**Context:**

It would be nice to have a function to calculate the expectation values
of an operator $A$ given a state vector $\vert\psi\rangle$. This can be
also used for computing the fidelity between a mixed and a pure state in
simple way, avoiding eigendecomposition problems.

**Description of the Change:**

Added a overlap calculation between state vectors and density matrices 

**Benefits:**

It is faster then computing $\text{Tr}(A
\vert\psi\rangle\langle\psi\vert)$. No need for eigendecomposition,
therefore, it is a way to avoid issues regarding differentiation as
pointed on issue #4373.

**Possible Drawbacks:**

It does not work if the user pass as input two state vectors or two
matrices. It is tailored to first input as (batched) matrices and the
second one as (batched) state vectors.

**Related GitHub Issues:**
#4373

---------

Co-authored-by: Frederik Wilde <[email protected]>
Co-authored-by: Christina Lee <[email protected]>
Co-authored-by: Matthew Silverman <[email protected]>
Co-authored-by: Romain Moyard <[email protected]>
Co-authored-by: Mudit Pandey <[email protected]>
Co-authored-by: Olivia Di Matteo <[email protected]>
Co-authored-by: Olivia Di Matteo <[email protected]>
Co-authored-by: Jay Soni <[email protected]>
Co-authored-by: Christina Lee <[email protected]>
Co-authored-by: Matthew Silverman <[email protected]>
Co-authored-by: David Wierichs <[email protected]>
Co-authored-by: Edward Jiang <[email protected]>
Co-authored-by: Utkarsh <[email protected]>
Co-authored-by: Korbinian Kottmann <[email protected]>
Co-authored-by: Thomas R. Bromley <[email protected]>
Co-authored-by: Astral Cai <[email protected]>
Co-authored-by: Astral Cai <[email protected]>

doc/releases/changelog-dev.md

@@ -34,6 +34,16 @@
   [1 1 0]
   ```
 
+* `expectation_value` was added to `qml.math` to calculate the expectation value of a matrix for pure states.
+  [(#4484)](https://github.com/PennyLaneAI/pennylane/pull/4484)
+
+  ```pycon
+  >>> state_vector = [1/np.sqrt(2), 0, 1/np.sqrt(2), 0]
+  >>> operator_matrix = qml.matrix(qml.PauliZ(0), wire_order=[0,1])
+  >>> qml.math.expectation_value(operator_matrix, state_vector)
+  tensor(-2.23711432e-17+0.j, requires_grad=True)
+  ```
+
 * The `default.tensor` device is introduced to perform tensor network simulations of quantum circuits using the `mps` (Matrix Product State) method.
   [(#5699)](https://github.com/PennyLaneAI/pennylane/pull/5699)
 
@@ -470,6 +480,7 @@ Tarun Kumar Allamsetty,
 Guillermo Alonso-Linaje,
 Utkarsh Azad,
 Lillian M. A. Frederiksen,
+Ludmila Botelho,
 Gabriel Bottrill,
 Astral Cai,
 Ahmed Darwish,

pennylane/math/__init__.py

@@ -68,6 +68,7 @@
 from .quantum import (
     cov_matrix,
     dm_from_state_vector,
+    expectation_value,
     marginal_prob,
     mutual_info,
     partial_trace,
@@ -151,6 +152,7 @@ def __getattr__(name):
     "dot",
     "einsum",
     "expand_matrix",
+    "expectation_value",
     "eye",
     "fidelity",
     "fidelity_statevector",

pennylane/math/quantum.py

@@ -472,7 +472,10 @@ def reduce_statevector(state, indices, check_state=False, c_dtype="complex128"):
         [ABC[i + 1] for i in sorted(indices)] + [ABC[num_wires + i + 1] for i in sorted(indices)]
     )
     density_matrix = einsum(
-        f"a{indices1},a{indices2}->a{target}", state, np.conj(state), optimize="greedy"
+        f"a{indices1},a{indices2}->a{target}",
+        state,
+        np.conj(state),
+        optimize="greedy",
     )
 
     # Return the reduced density matrix by using numpy tensor product
@@ -514,7 +517,10 @@ def dm_from_state_vector(state, check_state=False, c_dtype="complex128"):
     """
     num_wires = int(np.log2(np.shape(state)[-1]))
     return reduce_statevector(
-        state, indices=list(range(num_wires)), check_state=check_state, c_dtype=c_dtype
+        state,
+        indices=list(range(num_wires)),
+        check_state=check_state,
+        c_dtype=c_dtype,
     )
 
 
@@ -663,7 +669,14 @@ def _compute_vn_entropy(density_matrix, base=None):
 
 
 # pylint: disable=too-many-arguments
-def mutual_info(state, indices0, indices1, base=None, check_state=False, c_dtype="complex128"):
+def mutual_info(
+    state,
+    indices0,
+    indices1,
+    base=None,
+    check_state=False,
+    c_dtype="complex128",
+):
     r"""Compute the mutual information between two subsystems given a state:
 
     .. math::
@@ -719,29 +732,139 @@ def mutual_info(state, indices0, indices1, base=None, check_state=False, c_dtype
         raise ValueError("Subsystems for computing mutual information must not overlap.")
 
     return _compute_mutual_info(
-        state, indices0, indices1, base=base, check_state=check_state, c_dtype=c_dtype
+        state,
+        indices0,
+        indices1,
+        base=base,
+        check_state=check_state,
+        c_dtype=c_dtype,
     )
 
 
 # pylint: disable=too-many-arguments
 def _compute_mutual_info(
-    state, indices0, indices1, base=None, check_state=False, c_dtype="complex128"
+    state,
+    indices0,
+    indices1,
+    base=None,
+    check_state=False,
+    c_dtype="complex128",
 ):
     """Compute the mutual information between the subsystems."""
     all_indices = sorted([*indices0, *indices1])
     vn_entropy_1 = vn_entropy(
-        state, indices=indices0, base=base, check_state=check_state, c_dtype=c_dtype
+        state,
+        indices=indices0,
+        base=base,
+        check_state=check_state,
+        c_dtype=c_dtype,
     )
     vn_entropy_2 = vn_entropy(
-        state, indices=indices1, base=base, check_state=check_state, c_dtype=c_dtype
+        state,
+        indices=indices1,
+        base=base,
+        check_state=check_state,
+        c_dtype=c_dtype,
     )
     vn_entropy_12 = vn_entropy(
-        state, indices=all_indices, base=base, check_state=check_state, c_dtype=c_dtype
+        state,
+        indices=all_indices,
+        base=base,
+        check_state=check_state,
+        c_dtype=c_dtype,
     )
 
     return vn_entropy_1 + vn_entropy_2 - vn_entropy_12
 
 
+def _check_hermitian_operator(operators):
+    """Check the shape, and if the matrix is hermitian."""
+    dim = operators.shape[-1]
+
+    if (
+        len(operators.shape) not in (2, 3)
+        or operators.shape[-2] != dim
+        or not np.log2(dim).is_integer()
+    ):
+        raise ValueError(
+            "Operator matrix must be of shape (2**wires,2**wires) "
+            "or (batch_dim, 2**wires, 2**wires)."
+        )
+
+    if len(operators.shape) == 2:
+        operators = qml.math.stack([operators])
+
+    if not is_abstract(operators):
+        for ops in operators:
+            conj_trans = np.transpose(np.conj(ops))
+            if not allclose(ops, conj_trans):
+                raise ValueError("The matrix is not Hermitian.")
+
+
+def expectation_value(
+    operator_matrix, state_vector, check_state=False, check_operator=False, c_dtype="complex128"
+):
+    r"""Compute the expectation value of an operator with respect to a pure state.
+
+    The expectation value is the probabilistic expected result of an experiment.
+    Given a pure state, i.e., a state which can be represented as a single
+    vector :math:`\ket{\psi}` in the Hilbert space, the expectation value of an
+    operator :math:`A` can computed as
+
+    .. math::
+        \langle A \rangle_\psi = \bra{\psi} A \ket{\psi}
+
+
+    Args:
+        operator_matrix (tensor_like): operator matrix with shape ``(2**N, 2**N)`` or ``(batch_dim, 2**N, 2**N)``.
+        state_vector (tensor_like): state vector with shape ``(2**N)`` or ``(batch_dim, 2**N)``.
+        check_state (bool): if True, the function will check the validity of the state vector
+            via its shape and the norm.
+        check_operator (bool): if True, the function will check the validity of the operator
+            via its shape and whether it is hermitian.
+        c_dtype (str): complex floating point precision type.
+
+    Returns:
+        float: Expectation value of the operator for the state vector.
+
+    **Example**
+
+    The expectation value for any operator can obtained by passing their matrix representation as an argument.
+    For example, for a 2 qubit state, we can compute the expectation value of the operator :math:`Z \otimes I` as
+
+
+    >>> state_vector = [1/np.sqrt(2), 0, 1/np.sqrt(2), 0]
+    >>> operator_matrix = qml.matrix(qml.PauliZ(0), wire_order=[0,1])
+    >>> qml.math.expectation_value(operator_matrix, state_vector)
+    tensor(-2.23711432e-17+0.j, requires_grad=True)
+
+    .. seealso:: :func:`pennylane.math.fidelity`
+
+    """
+    state_vector = cast(state_vector, dtype=c_dtype)
+    operator_matrix = cast(operator_matrix, dtype=c_dtype)
+
+    if check_state:
+        _check_state_vector(state_vector)
+
+    if check_operator:
+        _check_hermitian_operator(operator_matrix)
+
+    if qml.math.shape(operator_matrix)[-1] != qml.math.shape(state_vector)[-1]:
+        raise qml.QuantumFunctionError(
+            "The operator and the state vector must have the same number of wires."
+        )
+
+    # The overlap <psi|A|psi>
+    expval = qml.math.einsum(
+        "...i,...i->...",
+        qml.math.conj(state_vector),
+        qml.math.einsum("...ji,...i->...j", operator_matrix, state_vector, optimize="greedy"),
+        optimize="greedy",
+    )
+    return expval
+
+
 # pylint: disable=too-many-arguments
 def vn_entanglement_entropy(
     state, indices0, indices1, base=None, check_state=False, c_dtype="complex128"
@@ -884,7 +1007,10 @@ def _compute_relative_entropy(rho, sigma, base=None):
     # the matrix of inner products between eigenvectors of rho and eigenvectors
     # of sigma; this is a doubly stochastic matrix
     rel = qml.math.einsum(
-        f"{indices_rho},{indices_sig}->{target}", np.conj(u_rho), u_sig, optimize="greedy"
+        f"{indices_rho},{indices_sig}->{target}",
+        np.conj(u_rho),
+        u_sig,
+        optimize="greedy",
     )
     rel = np.abs(rel) ** 2