Support for expected values in the XIR interpreter (#49)

* Add OperatorStmt to XIR exports

* Add Scale decorator to model observable prefactors

* Add support for operators and expval statements

* Update changelog

* Fix NumPy dtype docstrings

* Rename Python test file from TBCC to TBC

.github/CHANGELOG.md

@@ -18,6 +18,8 @@ This release contains contributions from (in alphabetical order):
 
 ### New features since last release
 
+* The Jet interpreter for XIR scripts now supports an `expval` output. [(#49)](https://github.com/XanaduAI/jet/pull/49)
+
 * The Jet interpreter for XIR scripts now accepts a `dimension` option for CV circuits. [(#47)](https://github.com/XanaduAI/jet/pull/47)
 
 * The Jet interpreter for XIR programs now supports a `probabilities` output. [(#44)](https://github.com/XanaduAI/jet/pull/44)

python/jet/circuit.py

@@ -181,7 +181,7 @@ def tensor_network(self, dtype: np.dtype = np.complex128) -> TensorNetworkType:
         """Returns the tensor network representation of this circuit.
 
         Args:
-            dtype (type): Data type of the tensor network.
+            dtype (np.dtype): Data type of the tensor network.
 
         Returns:
             TensorNetworkType: Tensor network representation of this circuit.

python/jet/gate.py

@@ -19,6 +19,7 @@
     "GateFactory",
     # Decorator gates
     "Adjoint",
+    "Scale",
     # CV Fock gates
     "FockGate",
     "Displacement",
@@ -170,7 +171,7 @@ def _data(self) -> np.ndarray:
         """Returns the matrix representation of this gate."""
         pass
 
-    def tensor(self, dtype: type = np.complex128) -> TensorType:
+    def tensor(self, dtype: np.dtype = np.complex128) -> TensorType:
         """Returns the tensor representation of this gate.
 
         Args:
@@ -198,7 +199,9 @@ class GateFactory:
     registry: Dict[str, type] = {}
 
     @staticmethod
-    def create(name: str, *params: float, adjoint: bool = False, **kwargs) -> Gate:
+    def create(
+        name: str, *params: float, adjoint: bool = False, scalar: float = 1, **kwargs
+    ) -> Gate:
         """Constructs a gate by name.
 
         Raises:
@@ -207,6 +210,8 @@ def create(name: str, *params: float, adjoint: bool = False, **kwargs) -> Gate:
         Args:
             name (str): Registered name of the desired gate.
             params (float): Parameters to pass to the gate constructor.
+            adjoint (bool): Whether to take the adjoint of the gate.
+            scalar (float): Scaling factor to apply to the gate.
             kwargs: Keyword arguments to pass to the gate constructor.
 
         Returns:
@@ -219,7 +224,10 @@ def create(name: str, *params: float, adjoint: bool = False, **kwargs) -> Gate:
         gate = subclass(*params, **kwargs)
 
         if adjoint:
-            gate = Adjoint(gate)
+            gate = Adjoint(gate=gate)
+
+        if scalar != 1:
+            gate = Scale(gate=gate, scalar=scalar)
 
         return gate
 
@@ -273,6 +281,7 @@ class Adjoint(Gate):
 
     def __init__(self, gate: Gate):
         self._gate = gate
+
         super().__init__(
             name=gate.name, num_wires=gate.num_wires, dim=gate.dimension, params=gate.params
         )
@@ -284,6 +293,29 @@ def _validate_dimension(self, dim):
         self._gate._validate_dimension(dim)
 
 
+class Scale(Gate):
+    """Scale is a decorator which linearly scales an existing ``Gate``.
+
+    Args:
+        gate (Gate): Gate to scale.
+        scalar (float): Scaling factor.
+    """
+
+    def __init__(self, gate: Gate, scalar: float):
+        self._gate = gate
+        self._scalar = scalar
+
+        super().__init__(
+            name=gate.name, num_wires=gate.num_wires, dim=gate.dimension, params=gate.params
+        )
+
+    def _data(self):
+        return self._scalar * self._gate._data()
+
+    def _validate_dimension(self, dim):
+        self._gate._validate_dimension(dim)
+
+
 ####################################################################################################
 # Continuous variable Fock gates
 ####################################################################################################

python/jet/interpreter.py

@@ -6,11 +6,11 @@
 
 import numpy as np
 
-from xir import Statement, XIRProgram
+from xir import OperatorStmt, Statement, XIRProgram
 from xir import parse_script as parse_xir_script
 
 from .bindings import PathInfo
-from .circuit import Circuit
+from .circuit import Circuit, Operation
 from .factory import TaskBasedContractor, TensorNetworkType, TensorType
 from .gate import FockGate, GateFactory
 from .state import Qudit
@@ -178,6 +178,34 @@ def run_xir_program(program: XIRProgram) -> List[Union[np.number, np.ndarray]]:
             output = _compute_probabilities(circuit=circuit)
             result.append(output)
 
+        elif stmt.name in ("Expval", "expval"):
+            if not isinstance(stmt.params, dict) or "observable" not in stmt.params:
+                raise ValueError(f"Statement '{stmt}' is missing an 'observable' parameter.")
+
+            operator = stmt.params["observable"]
+
+            if operator not in program.operators:
+                raise ValueError(
+                    f"Statement '{stmt}' has an 'observable' parameter which "
+                    f"references an undefined operator."
+                )
+
+            elif program.operators[operator]["params"]:
+                raise ValueError(
+                    f"Statement '{stmt}' has an 'observable' parameter which "
+                    f"references a parameterized operator."
+                )
+
+            elif stmt.wires != tuple(range(num_wires)):
+                raise ValueError(f"Statement '{stmt}' must be applied to [0 .. {num_wires - 1}].")
+
+            observable = _generate_observable_from_operation_statements(
+                program.operators[operator]["statements"]
+            )
+
+            output = _compute_expected_value(circuit=circuit, observable=observable)
+            result.append(output)
+
         else:
             raise ValueError(f"Statement '{stmt}' is not supported.")
 
@@ -354,15 +382,41 @@ def _bind_statement_wires(gate_signature_map: Dict[str, GateSignature], stmt: St
     return {name: have_wires[i] for (i, name) in enumerate(want_wires)}
 
 
+def _generate_observable_from_operation_statements(
+    stmts: Iterator[OperatorStmt],
+) -> Iterator[Operation]:
+    """Generates an observable from a series of operator statements.
+
+    Args:
+        stmts (Iterator[OperatorStmt]): Operator statements defining the observable.
+
+    Returns:
+        Iterator[Operation]: Iterator over a sequence of ``Operation`` objects
+            which implement the observable given by the operator statements.
+    """
+    for stmt in stmts:
+        try:
+            scalar = float(stmt.pref)
+        except ValueError:
+            raise ValueError(
+                f"Operator statement '{stmt}' has a prefactor ({stmt.pref}) "
+                f"which cannot be converted to a floating-point number."
+            )
+
+        for gate_name, wire_id in stmt.terms:
+            gate = GateFactory.create(gate_name, scalar=scalar)
+            yield Operation(part=gate, wire_ids=[wire_id])
+
+
 def _compute_amplitude(
-    circuit: Circuit, state: List[int], dtype: type = np.complex128
+    circuit: Circuit, state: List[int], dtype: np.dtype = np.complex128
 ) -> np.number:
     """Computes the amplitude of a state at the end of a circuit.
 
     Args:
         circuit (Circuit): Circuit to apply the amplitude measurement to.
         state (list[int]): State to measure the amplitude of.
-        dtype (type): Data type of the amplitude.
+        dtype (np.dtype): Data type of the amplitude.
 
     Returns:
         Number: NumPy number representing the amplitude of the given state.
@@ -376,35 +430,57 @@ def _compute_amplitude(
         qudit = Qudit(dim=circuit.dimension, data=data)
         circuit.append_state(qudit, wire_ids=[i])
 
-    result = _simulate(circuit=circuit, dtype=dtype)
-    return dtype(result.scalar)
+    amplitude = _simulate(circuit=circuit, dtype=dtype)
+    return dtype(amplitude.scalar)
 
 
-def _compute_probabilities(circuit: Circuit, dtype: type = np.complex128) -> np.ndarray:
+def _compute_probabilities(circuit: Circuit, dtype: np.dtype = np.complex128) -> np.ndarray:
     """Computes the probability distribution at the end of a circuit.
 
     Args:
         circuit (Circuit): Circuit to produce the probability distribution for.
-        dtype (type): Data type of the probability computation.
+        dtype (np.dtype): Data type of the probability computation.
 
     Returns:
         Array: NumPy array representing the probability of measuring each basis state.
     """
-    result = _simulate(circuit=circuit, dtype=dtype)
+    tensor = _simulate(circuit=circuit, dtype=dtype)
 
     # Arrange the indices in increasing order of wire ID.
-    state = result.transpose([wire.index for wire in circuit.wires])
+    state = tensor.transpose([wire.index for wire in circuit.wires])
 
     amplitudes = np.array(state.data).flatten()
     return amplitudes.conj() * amplitudes
 
 
-def _simulate(circuit: Circuit, dtype: type = np.complex128) -> TensorType:
+def _compute_expected_value(
+    circuit: Circuit, observable: Iterator[Operation], dtype: np.dtype = np.complex128
+) -> np.number:
+    """Computes the expected value of an observable with respect to a circuit.
+
+    Args:
+        circuit (Circuit): Circuit to apply the expectation measurement to.
+        observable (Observable): Observable to take the expected value of.
+        dtype (np.dtype): Data type of the expected value.
+
+    Returns:
+        Number: NumPy number representing the expected value.
+    """
+    # Do not modify the original circuit.
+    circuit = deepcopy(circuit)
+
+    circuit.take_expected_value(observable)
+
+    expval = _simulate(circuit=circuit, dtype=dtype)
+    return dtype(expval.scalar)
+
+
+def _simulate(circuit: Circuit, dtype: np.dtype = np.complex128) -> TensorType:
     """Simulates a circuit using the task-based contractor.
 
     Args:
         circuit (Circuit): Circuit to simulate.
-        dtype (type): Data type of the tensor network to contract.
+        dtype (np.dtype): Data type of the tensor network to contract.
 
     Returns:
         Tensor: Result of the simulation.

python/tests/jet/test_gate.py

@@ -173,6 +173,43 @@ def test_data(self, gate, state, want_tensor):
         assert have_tensor.indices == want_tensor.indices
 
 
+class TestScale:
+    @pytest.mark.parametrize("gate", [jet.Hadamard(), jet.U2(1, 2)])
+    def test_attributes(self, gate):
+        """Tests that Scale returns the same attributes as the decorated gate."""
+        scaled_gate = jet.Scale(gate, scalar=1)
+        assert scaled_gate.name == gate.name
+        assert scaled_gate.num_wires == gate.num_wires
+        assert scaled_gate.params == gate.params
+
+    @pytest.mark.parametrize(
+        ["gate", "state", "scalar", "want_tensor"],
+        [
+            pytest.param(
+                jet.PauliX(),
+                jet.Tensor(indices=["1"], shape=[2], data=[1, 0]),
+                2,
+                jet.Tensor(indices=["0"], shape=[2], data=[0, 2]),
+                id="2 * X|0>",
+            ),
+            pytest.param(
+                jet.PauliX(),
+                jet.Tensor(indices=["1"], shape=[2], data=[0, 1]),
+                -3j,
+                jet.Tensor(indices=["0"], shape=[2], data=[-3j, 0]),
+                id="-3i * X|1>",
+            ),
+        ],
+    )
+    def test_data(self, gate, state, scalar, want_tensor):
+        """Tests that Scale linearly scales the tensor of a ``PauliX`` gate."""
+        scaled_gate = jet.Scale(gate, scalar=scalar)
+        have_tensor = jet.contract_tensors(scaled_gate.tensor(), state)
+        assert have_tensor.data == pytest.approx(want_tensor.data)
+        assert have_tensor.shape == want_tensor.shape
+        assert have_tensor.indices == want_tensor.indices
+
+
 class TestFockGate:
     @pytest.fixture
     def gate(self) -> MockFockGate: