Compute expectation value

pennylane/math/__init__.py

@@ -69,6 +69,7 @@
     dm_from_state_vector,
     fidelity,
     fidelity_statevector,
+    expectation_value,
     marginal_prob,
     mutual_info,
     purity,

pennylane/math/quantum.py

@@ -720,6 +720,63 @@
     return vn_entropy_1 + vn_entropy_2 - vn_entropy_12
 
 
+def expectation_value(operator_matrix, state_vector, check_state=False, c_dtype="complex128"):
+    r"""Compute the expectation value of an operator with respect to a pure state.
+
+    The expectation value of an operator :math:`A` for a pure state given by a state vector :math:`\ket{\psi}`
+    is defined 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
+        c_dtype (str): Complex floating point precision type.
+
+    Returns:
+        float: Expectation value of the operator for the state vector.
+
+    **Example**
+
+    An operator matrix and a state vector can be used as arguments to obtain the expectation value, e.g.:
+
+    >>> state_vector = np.array([0, 1])
+    >>> operator_matrix = np.array([[0,1],[1,0]])
+    >>> qml.math.expectation_value(operator_matrix, state_vector)
+    0.0
+
+    .. seealso:: :func:`pennylane.math.fidelity`
+
+    """
+    # Cast as a c_dtype array
+    state_vector = cast(state_vector, dtype=c_dtype)
+
+    # Cannot be cast_like if jit
+    if not is_abstract(state_vector):
+        operator_matrix = cast_like(operator_matrix, state_vector)
+
+    if check_state:
+        _check_state_vector(state_vector)
+
+    if qml.math.shape(operator_matrix)[-1] != qml.math.shape(state_vector)[-1]:
+        raise qml.QuantumFunctionError("The two states must have the same number of wires.")
+
+    batched0 = len(qml.math.shape(state_vector)) > 1
+    batched1 = len(qml.math.shape(operator_matrix)) > 2
+    # The overlap <psi|A|psi>
+    overlap = qml.math.einsum(
+        f"...j,...ji,...i->...",
+        qml.math.conj(state_vector),
+        operator_matrix,
+        state_vector,
+        optimize="greedy",
+    )
+    return overlap
+
+
 def fidelity(state0, state1, check_state=False, c_dtype="complex128"):
     r"""Compute the fidelity for two states (given as density matrices) acting on quantum
     systems with the same size.

tests/math/test_expectation_value_math.py

@@ -0,0 +1,136 @@
+# Copyright 2023 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.
+"""Unit tests for differentiable matrix-vector expectation values.
+"""
+
+import numpy as onp
+import pytest
+
+import torch
+import jax.numpy as jnp
+import tensorflow as tf
+
+import pennylane as qml
+from pennylane import numpy as np
+
+pytestmark = pytest.mark.all_interfaces
+
+
+class TestExpectationValueMath:
+    """Tests for Expectation value of a operator for a state vector."""
+
+    ops_vs_vecstates = [
+        ([[1, 0], [0, 0]], [1, 0], 1),
+        ([[0, 1], [1, 0]], [0, 1], 0),
+        ([[0.5, 0.5], [0.5, 0.5]], [1, 1] / np.sqrt(2), 1),
+    ]
+
+    array_funcs = [
+        lambda x: x,
+        onp.array,
+        np.array,
+        jnp.array,
+        torch.tensor,
+        tf.Variable,
+        tf.constant,
+    ]
+
+    check_state = [True, False]
+
+    @pytest.mark.parametrize("operator_and_states", ops_vs_vecstates)
+    @pytest.mark.parametrize("check_state", check_state)
+    @pytest.mark.parametrize("func", array_funcs)
+    def test_mat_expectation_value(self, operator_and_states, check_state, func):
+        """Test the expectation value of a operator for a vector state."""
+        ops, state_vectors, expected = operator_and_states
+        ops = func(ops)
+        state_vectors = func(state_vectors)
+        overlap = qml.math.expectation_value(ops, state_vectors, check_state)
+        assert qml.math.allclose(expected, overlap)
+
+    state_wrong_amp = [([[1, 0], [0, 1]], [0.5, 0]), ([[1, 0], [0, 1]], [0, 0.5])]
+
+    @pytest.mark.parametrize("ops,state_vectors", state_wrong_amp)
+    def test_state_vector_wrong_amplitudes(self, ops, state_vectors):
+        """Test that a message is raised when a state does not have right norm"""
+        with pytest.raises(ValueError, match="Sum of amplitudes-squared does not equal one."):
+            qml.math.expectation_value(ops, state_vectors, check_state=True)
+
+    state_wrong_shape = [([[1, 0], [0, 1]], [0, 1, 0]), ([[1, 0], [0, 1]], [0, 0, 0, 0, 1])]
+
+    @pytest.mark.parametrize("ops,state_vectors", state_wrong_shape)
+    def test_state_vector_wrong_shape(self, ops, state_vectors):
+        """Test that a message is raised when the state does not have the right shape."""
+        with pytest.raises(ValueError, match="State vector must be of shape"):
+            qml.math.expectation_value(ops, state_vectors, check_state=True)
+
+    def test_same_number_wires_dm(self):
+        """Test that the two states must act on the same number of wires"""
+        ops = np.diag([0, 1, 0, 0])
+        state_vectors = [1, 0]
+        with pytest.raises(
+            qml.QuantumFunctionError, match="The two states must have the same number of wires"
+        ):
+            qml.math.expectation_value(ops, state_vectors, check_state=True)
+
+    @pytest.mark.parametrize("check_state", check_state)
+    @pytest.mark.parametrize("func", array_funcs)
+    def test_broadcast_op_sv(self, check_state, func):
+        """Test simultaneous broadcasting of operators and state vectors works."""
+        ops = func(
+            [
+                [
+                    [
+                        1,
+                        0,
+                    ],
+                    [0, 0],
+                ],
+                [[0.5, 0.5], [0.5, 0.5]],
+                [[1, 0], [0, 1]],
+            ]
+        )
+        state_vectors = qml.math.stack([func([0, 1]), func([1, 0]), func([1, 1]) / np.sqrt(2)])
+        expected = [0, 0.5, 1]
+
+        overlap = qml.math.expectation_value(ops, state_vectors, check_state)
+        assert qml.math.allclose(overlap, expected)
+
+    @pytest.mark.parametrize("check_state", check_state)
+    @pytest.mark.parametrize("func", array_funcs)
+    def test_broadcast_op_unbatched(self, check_state, func):
+        """Test broadcasting works for expectation values when the operators input is unbatched"""
+        ops = func([[1, 0], [0, 0]])
+        state_vectors = func([[0, 1], [1, 0], [1, 1] / np.sqrt(2)])
+        expected = [0, 1, 0.5]
+
+        overlap = qml.math.expectation_value(ops, state_vectors, check_state)
+        assert qml.math.allclose(overlap, expected)
+
+    @pytest.mark.parametrize("check_state", check_state)
+    @pytest.mark.parametrize("func", array_funcs)
+    def test_broadcast_sv_unbatched(self, check_state, func):
+        """Test broadcasting works for expectation values when the state vector input is unbatched"""
+        ops = func(
+            [
+                [[1, 0], [0, 0]],
+                [[0.5, 0.5], [0.5, 0.5]],
+                [[0, 0], [0, 1]],
+            ]
+        )
+        state_vectors = func([1, 0])
+        expected = [1, 0.5, 0]
+
+        overlap = qml.math.expectation_value(ops, state_vectors, check_state)
+        assert qml.math.allclose(overlap, expected)