Quantum circuit and state models for IR integration

.github/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ### New features since last release
 
+* Quantum circuit and state models have been added to the `jet` Python package. [(#21)](https://github.com/XanaduAI/jet/pull/21)
+
 * Quantum gate models have been added to the `jet` Python package. [(#16)](https://github.com/XanaduAI/jet/pull/16)
 
 * Python bindings are now available for the `TaskBasedCpuContractor` class. [(#19)](https://github.com/XanaduAI/jet/pull/19)

python/jet/__init__.py

@@ -2,8 +2,10 @@
 from .bindings import *
 
 # The rest of the modules control their exports using `__all__`.
+from .circuit import *
 from .factory import *
 from .gate import *
+from .state import *
 
 # Grab the current Jet version from the C++ headers.
 __version__ = version()

python/jet/circuit.py

@@ -0,0 +1,152 @@
+from dataclasses import dataclass
+from functools import wraps
+from typing import Callable, Iterator, List, Sequence, Tuple, Union
+
+import numpy as np
+
+from .factory import TensorNetwork, TensorNetworkType
+from .gate import Gate
+from .state import Qudit, State
+
+__all__ = [
+    "Wire",
+    "Circuit",
+]
+
+
+@dataclass
+class Wire:
+    """Wire represents a collection of tensor indices that are directly or
+    transitively associated with a qudit of a quantum circuit.
+    """
+
+    # Position of the wire in the circuit.
+    id_: int
+    # Number of gates applied to this wire.
+    depth: int = 0
+    # Whether this wire has been terminated with a state.
+    closed: bool = False
+
+    @property
+    def index(self) -> str:
+        """Returns the current index label of this wire."""
+        return f"{self.id_}-{self.depth}"
+
+
+class Circuit:
+    def __init__(self, num_wires: int, dim: int = 2):
+        """Constructs a quantum circuit. Each wire is initialized with a qudit
+        of the specified dimension in the vacuum state.
+
+        Args:
+            num_wires (int): number of wires in the circuit.
+            dim (int): dimension of each wire.
+        """
+        self._wires = [Wire(i) for i in range(num_wires)]
+        self._parts = [Qudit(dim=dim) for _ in range(num_wires)]
+
+        for (wire, part) in zip(self._wires, self._parts):
+            part.indices = [wire.index]
+
+    @property
+    def parts(self) -> Iterator[Union[Gate, State]]:
+        """Returns the gates and states that comprise this circuit.  The first
+        ``self.num_wires`` parts are the qudits that begin each wire; other
+        parts appear in the order they were appended to the circuit.
+        """
+        return iter(self._parts)
+
+    @property
+    def wires(self) -> Iterator[Wire]:
+        """Returns the wires of this circuit in increasing order of wire ID."""
+        return iter(self._wires)
+
+    def _append_validator(append_fn: Callable) -> Callable:
+        """Decorator which validates the arguments to an append function.
+
+        Args:
+            append_fn (Callable): append function to be validated.
+
+        Returns:
+            Function that validates the passed wire IDs before invoking the
+            decorated append function with the given arguments.
+        """
+
+        @wraps(append_fn)
+        def validator(self, part: Union[Gate, State], wire_ids: Sequence[int]) -> None:
+            if len(wire_ids) != part.num_wires:
+                raise ValueError(
+                    f"Number of wire IDs ({len(wire_ids)}) must match the number of "
+                    f"wires connected to the circuit component ({part.num_wires})."
+                )
+
+            num_wires = len(self._wires)
+
+            for wire_id in wire_ids:
+                if not 0 <= wire_id < num_wires:
+                    raise ValueError(f"Wire ID {wire_id} falls outside the range [0, {num_wires}).")
+
+                elif wire_ids.count(wire_id) > 1:
+                    raise ValueError(f"Wire ID {wire_id} is specified more than once.")
+
+                elif self._wires[wire_id].closed:
+                    raise ValueError(f"Wire {wire_id} is closed.")
+
+            append_fn(self, part, wire_ids)
+
+        return validator
+
+    @_append_validator
+    def append_gate(self, gate: Gate, wire_ids: Sequence[int]) -> None:
+        """Applies a gate along the specified wires.
+
+        Args:
+            gate (Gate): gate to be applied.
+            wire_ids (Sequence[int]): IDs of the wires the gate is applied to.
+        """
+        input_indices = self.indices(wire_ids)
+
+        for i in wire_ids:
+            self._wires[i].depth += 1
+
+        output_indices = self.indices(wire_ids)
+
+        gate.indices = output_indices + input_indices
+        self._parts.append(gate)
+
+    @_append_validator
+    def append_state(self, state: State, wire_ids: Sequence[int]) -> None:
+        """Terminates the specified wires with a quantum state.
+
+        Args:
+            state (State): state to be used for termination.
+            wire_ids (Sequence[int]): IDs of the wires the state terminates.
+        """
+        for i in wire_ids:
+            self._wires[i].closed = True
+
+        state.indices = self.indices(wire_ids)
+        self._parts.append(state)
+
+    def indices(self, wire_ids: Sequence[int]) -> List[str]:
+        """Returns the index labels associated with a sequence of wire IDs.
+
+        Args:
+            wire_ids (Sequence[int]): IDs of the wires to get the index labels for.
+
+        Returns:
+            List of index labels.
+        """
+        return [self._wires[i].index for i in wire_ids]
+
+    def tensor_network(self, dtype: type = np.complex128) -> TensorNetworkType:
+        """Returns the tensor network representation of this circuit.
+
+        Args:
+            dtype (type): data type of the tensor network.
+        """
+        tn = TensorNetwork(dtype=dtype)
+        for part in self._parts:
+            tensor = part.tensor(dtype=dtype)
+            tn.add_tensor(tensor)
+        return tn

python/jet/gate.py

@@ -113,7 +113,7 @@ def indices(self, indices: Optional[Sequence[str]]) -> None:
         # Check that `indices` has the correct length.
         elif len(indices) != 2 * self._num_wires:
             raise ValueError(
-                f"Gates must have two indices per wire; received {len(indices)}"
+                f"Gates must have two indices per wire; received {len(indices)} "
                 f"indices for {self._num_wires} wires."
             )
 

python/jet/state.py

@@ -0,0 +1,175 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional, Sequence
+
+import numpy as np
+
+from .factory import Tensor, TensorType
+
+__all__ = [
+    "State",
+    "Qudit",
+    "QuditRegister",
+    "Qubit",
+    "QubitRegister",
+]
+
+
+class State(ABC):
+    def __init__(self, name: str, num_wires: int):
+        """Constructs a quantum state.
+
+        Args:
+            name (str): name of the state.
+            num_wires (str): number of wires the state is connected to.
+        """
+        self.name = name
+
+        self._indices = None
+        self._num_wires = num_wires
+
+    @property
+    def indices(self) -> Optional[List[str]]:
+        """Returns the indices of this state. An index is a label associated with
+        an axis of the tensor representation of a state; the indices of a tensor
+        determine its connectivity in the context of a tensor network.
+        """
+        return self._indices
+
+    @indices.setter
+    def indices(self, indices: Optional[Sequence[str]]) -> None:
+        """Sets the indices of this state. The ``indices`` property of a state
+        is used to construct its tensor representation (unless ``indices`` is
+        None). See @indices.getter for more information about tensor indices.
+
+        Raises:
+            ValueError if the given indices are not a sequence of unique strings
+            or the number of provided indices is invalid.
+
+        Args:
+            indices (Sequence[str] or None): new indices of the state.
+        """
+        # Skip the sequence property checks if `indices` is None.
+        if indices is None:
+            pass
+
+        # Check that `indices` is a sequence of unique strings.
+        elif (
+            not isinstance(indices, Sequence)
+            or not all(isinstance(idx, str) for idx in indices)
+            or len(set(indices)) != len(indices)
+        ):
+            raise ValueError("Indices must be a sequence of unique strings.")
+
+        # Check that `indices` has the correct length (or is None).
+        elif len(indices) != self.num_wires:
+            raise ValueError(
+                f"States must have one index per wire. "
+                f"Received {len(indices)} indices for {self.num_wires} wires."
+            )
+
+        self._indices = indices
+
+    @property
+    def num_wires(self) -> int:
+        """Returns the number of wires connected to this state."""
+        return self._num_wires
+
+    def __eq__(self, other) -> bool:
+        """Reports whether this state is equivalent to the given state."""
+        return np.all(self._data() == other._data())
+
+    def __ne__(self, other) -> bool:
+        """Reports whether this state is not equivalent to the given state."""
+        return not (self == other)
+
+    @abstractmethod
+    def _data(self) -> np.ndarray:
+        """Returns the vector representation of this state."""
+        pass
+
+    def tensor(self, dtype: type = np.complex128, adjoint: bool = False) -> TensorType:
+        """Returns the tensor representation of this state.
+
+        Args:
+            dtype (type): data type of the tensor.
+            adjoint (bool): whether to take the adjoint of the tensor.
+        """
+        if adjoint:
+            data = np.conj(self._data())
+        else:
+            data = self._data()
+
+        indices = self.indices
+        if indices is None:
+            indices = list(map(str, range(self.num_wires)))
+
+        dimension = int(round(len(data) ** (1 / len(indices))))
+        shape = [dimension] * len(indices)
+
+        return Tensor(indices=indices, shape=shape, data=data, dtype=dtype)
+
+
+class Qudit(State):
+    def __init__(self, dim: int, data: Optional[np.ndarray] = None):
+        """Constructs a qudit state.
+
+        Args:
+            dim (int): dimension of the qudit.
+            data (np.ndarray or None): optional state vector.
+        """
+        name = "Qubit" if dim == 2 else f"Qudit(d={dim})"
+        super().__init__(name=name, num_wires=1)
+
+        if data is None:
+            self._state_vector = (np.arange(dim) == 0).astype(np.complex128)
+        else:
+            self._state_vector = data.flatten()
+
+    def _data(self) -> np.ndarray:
+        return self._state_vector
+
+
+class QuditRegister(State):
+    def __init__(self, dim: int, size: int, data: Optional[np.ndarray] = None):
+        """Constructs a qudit register state.
+
+        Args:
+            dim (int): dimension of the qudits.
+            size (int): number of qudits.
+            data (np.ndarray or None): optional state vector.
+        """
+        name = f"Qubit[{size}]" if dim == 2 else f"Qudit(d={dim})[{size}]"
+        super().__init__(name=name, num_wires=size)
+
+        if data is None:
+            self._state_vector = (np.arange(dim ** size) == 0).astype(np.complex128)
+        else:
+            self._state_vector = data.flatten()
+
+    def _data(self) -> np.ndarray:
+        return self._state_vector
+
+
+def Qubit(data: Optional[np.ndarray] = None) -> Qudit:
+    """Constructs a qubit state using an optional state vector.
+
+    Args:
+        data (np.ndarray or None): optional state vector.
+
+    Returns:
+        Qudit instance constructed using the specified state vector.
+    """
+    return Qudit(dim=2, data=data)
+
+
+def QubitRegister(size: int, data: Optional[np.ndarray] = None) -> QuditRegister:
+    """Constructs a qubit register state with the given size and optional state vector.
+
+    Args:
+        size (int): number of qubits.
+        data (np.ndarray or None): optional state vector.
+
+    Returns:
+        QuditRegister instance constructed using the specified state vector.
+    """
+    return QuditRegister(dim=2, size=size, data=data)