Enable type-checking for simulator and backends

CHANGELOG.md

@@ -16,6 +16,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - #277: Methods for pretty-printing `Pattern`: `to_ascii`,
   `to_unicode`, `to_latex`.
 
+- #312: The separation between `TensorNetworkBackend` and backends
+  that operate on a full-state representation, such as
+  `StatevecBackend` and `DensityMatrixBackend`, is now clearer with
+  the introduction of the abstract classes `DenseStateBackend` and
+  `DenseState`, which derive from `Backend` and `BackendState`,
+  respectively. `StatevecBackend` and `DensityMatrixBackend` inherit
+  from `DenseStateBackend`, while `Statevec` and `DensityMatrix`
+  inherit from `DenseState`. Note that the class hierarchy of
+  `BackendState` mirrors that of `Backend`.
+
 - #322: Added a new `optimization` module containing:
 
   * a functional version of `standardize` that returns a standardized
@@ -44,6 +54,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   semantics is preserved between circuits, ZX graphs, open graphs and
   patterns.
 
+- #302, #308, #312: `Pattern`, `Circuit`, `PatternSimulator`, and
+  backends are now type-checked.
+
 ### Changed
 
 - #277: The method `Pattern.print_pattern` is now deprecated.
@@ -61,6 +74,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Note: the method `perform_pauli_measurements` still places C
   commands before X and Z commands.
 
+- #312: Backend's `State` has been renamed to `BackendState` to avoid
+  a name conflict with the `State` class defined in `graphix.states`,
+  which represents the state of a single qubit.
+
+- #312: `Backend[StateT_co]` and `DenseStateBackend[DenseStateT_co]`
+  are now parameterized by covariant type variables, allowing
+  subclasses to narrow the type of the state field to match their
+  specific state representation. Covariance is sound in this context
+  because the classes are frozen, and it ensures that
+  `Backend[BackendState]` is a supertype of all backend classes.
+
 ## [0.3.1] - 2025-04-21
 
 ### Added

graphix/channels.py

@@ -13,7 +13,7 @@
 from graphix.ops import Ops
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Iterable, Iterator
 
 _T = TypeVar("_T", bound=np.generic)
 
@@ -147,6 +147,10 @@ def __len__(self) -> int:
         """Return the number of Kraus operators."""
         return len(self.__data)
 
+    def __iter__(self) -> Iterator[KrausData]:
+        """Iterate over Kraus operators."""
+        return iter(self.__data)
+
     @property
     def nqubit(self) -> int:
         """Return the number of qubits."""

graphix/graphsim.py

@@ -9,13 +9,16 @@
 
 from graphix import utils
 from graphix.clifford import Clifford
+from graphix.measurements import outcome
 from graphix.ops import Ops
 from graphix.sim.statevec import Statevec
 
 if TYPE_CHECKING:
     import functools
     from collections.abc import Iterable, Mapping
 
+    from graphix.measurements import Outcome
+
 
 if TYPE_CHECKING:
     Graph = nx.Graph[int]
@@ -373,7 +376,7 @@ def equivalent_fill_node(self, node: int) -> int:
             return 2
         return 0
 
-    def measure_x(self, node: int, choice: int = 0) -> int:
+    def measure_x(self, node: int, choice: Outcome = 0) -> Outcome:
         """Perform measurement in X basis.
 
         According to original paper, we realise X measurement by
@@ -406,7 +409,7 @@ def measure_x(self, node: int, choice: int = 0) -> int:
         self.h(node)
         return self.measure_z(node, choice=choice)
 
-    def measure_y(self, node: int, choice: int = 0) -> int:
+    def measure_y(self, node: int, choice: Outcome = 0) -> Outcome:
         """Perform measurement in Y basis.
 
         According to original paper, we realise Y measurement by
@@ -431,7 +434,7 @@ def measure_y(self, node: int, choice: int = 0) -> int:
         self.h(node)
         return self.measure_z(node, choice=choice)
 
-    def measure_z(self, node: int, choice: int = 0) -> int:
+    def measure_z(self, node: int, choice: Outcome = 0) -> Outcome:
         """Perform measurement in Z basis.
 
         To realize the simple Z measurement on undecorated graph state,
@@ -455,7 +458,7 @@ def measure_z(self, node: int, choice: int = 0) -> int:
         if choice:
             for i in self.neighbors(node):
                 self.flip_sign(i)
-        result = choice if not isolated else int(self.nodes[node]["sign"])
+        result = choice if not isolated else outcome(self.nodes[node]["sign"])
         self.remove_node(node)
         return result
 

graphix/measurements.py

@@ -4,14 +4,23 @@
 
 import dataclasses
 import math
-from typing import NamedTuple, SupportsInt
+from typing import Literal, NamedTuple, SupportsInt
+
+from typing_extensions import TypeAlias  # TypeAlias introduced in Python 3.10
 
 from graphix import utils
 from graphix.fundamentals import Axis, Plane, Sign
 
 # Ruff suggests to move this import to a type-checking block, but dataclass requires it here
 from graphix.parameter import ExpressionOrFloat  # noqa: TC001
 
+Outcome: TypeAlias = Literal[0, 1]
+
+
+def outcome(b: bool) -> Outcome:
+    """Return 1 if True, 0 if False."""
+    return 1 if b else 0
+
 
 @dataclasses.dataclass
 class Domains:

graphix/noise_models/noise_model.py

@@ -13,6 +13,7 @@
 
 if TYPE_CHECKING:
     from graphix.channels import KrausChannel
+    from graphix.measurements import Outcome
     from graphix.simulator import PatternSimulator
 
 
@@ -66,17 +67,17 @@ def measure(self) -> KrausChannel:
         ...
 
     @abc.abstractmethod
-    def confuse_result(self, result: bool) -> bool:
+    def confuse_result(self, result: Outcome) -> Outcome:
         """Return a possibly flipped measurement outcome.
 
         Parameters
         ----------
-        result : bool
+        result : Outcome
             Ideal measurement result.
 
         Returns
         -------
-        bool
+        Outcome
             Possibly corrupted result.
         """
 

graphix/noise_models/noiseless_noise_model.py

@@ -7,12 +7,17 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 import typing_extensions
 
 from graphix.channels import KrausChannel, KrausData
 from graphix.noise_models.noise_model import NoiseModel
 
+if TYPE_CHECKING:
+    from graphix.measurements import Outcome
+
 
 class NoiselessNoiseModel(NoiseModel):
     """Noise model that performs no operation."""
@@ -51,7 +56,7 @@ def measure(self) -> KrausChannel:
         return KrausChannel([KrausData(1.0, np.eye(2))])
 
     @typing_extensions.override
-    def confuse_result(self, result: bool) -> bool:
+    def confuse_result(self, result: Outcome) -> Outcome:
         """Return the unmodified measurement result.
 
         Parameters

graphix/parameter.py

@@ -334,6 +334,26 @@ def __hash__(self) -> int:
 T = TypeVar("T")
 
 
+def check_expression_or_complex(value: object) -> ExpressionOrComplex:
+    """Check that the given object is of type ExpressionOrComplex and return it."""
+    if isinstance(value, Expression):
+        return value
+    if isinstance(value, SupportsComplex):
+        return complex(value)
+    msg = f"ExpressionOrComplex expected, but {type(value)} found."
+    raise TypeError(msg)
+
+
+def check_expression_or_float(value: object) -> ExpressionOrFloat:
+    """Check that the given object is of type ExpressionOrFloat and return it."""
+    if isinstance(value, Expression):
+        return value
+    if isinstance(value, SupportsFloat):
+        return float(value)
+    msg = f"ExpressionOrFloat expected, but {type(value)} found."
+    raise TypeError(msg)
+
+
 @overload
 def subs(value: ExpressionOrFloat, variable: Parameter, substitute: ExpressionOrSupportsFloat) -> ExpressionOrFloat: ...
 

graphix/pattern.py

@@ -12,7 +12,7 @@
 from copy import deepcopy
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, SupportsFloat
+from typing import TYPE_CHECKING, SupportsFloat, TypeVar
 
 import networkx as nx
 from typing_extensions import assert_never, override
@@ -23,7 +23,7 @@
 from graphix.fundamentals import Axis, Plane, Sign
 from graphix.gflow import find_flow, find_gflow, get_layers
 from graphix.graphsim import GraphState
-from graphix.measurements import PauliMeasurement
+from graphix.measurements import Outcome, PauliMeasurement
 from graphix.pretty_print import OutputFormat, pattern_to_str
 from graphix.simulator import PatternSimulator
 from graphix.states import BasicStates
@@ -35,9 +35,10 @@
     from typing import Any, Literal
 
     from graphix.parameter import ExpressionOrFloat, ExpressionOrSupportsFloat, Parameter
-    from graphix.sim.base_backend import Backend
-    from graphix.sim.base_backend import State as BackendState
-    from graphix.sim.density_matrix import Data
+    from graphix.sim import Backend, BackendState, Data
+
+
+_StateT_co = TypeVar("_StateT_co", bound="BackendState", covariant=True)
 
 
 @dataclass(frozen=True)
@@ -81,7 +82,7 @@ class Pattern:
         total number of nodes in the resource state
     """
 
-    results: dict[int, int]
+    results: dict[int, Outcome]
     __seq: list[Command]
 
     def __init__(
@@ -250,7 +251,7 @@ def compose(
 
         mapped_inputs = [mapping_complete[n] for n in other.input_nodes]
         mapped_outputs = [mapping_complete[n] for n in other.output_nodes]
-        mapped_results = {mapping_complete[n]: m for n, m in other.results.items()}
+        mapped_results: dict[int, Outcome] = {mapping_complete[n]: m for n, m in other.results.items()}
 
         merged = mapping_values_set.intersection(self.__output_nodes)
 
@@ -289,7 +290,7 @@ def update_command(cmd: Command) -> Command:
 
         seq = self.__seq + [update_command(c) for c in other]
 
-        results = {**self.results, **mapped_results}
+        results: dict[int, Outcome] = {**self.results, **mapped_results}
         p = Pattern(input_nodes=inputs, output_nodes=outputs, cmds=seq)
         p.results = results
 
@@ -1354,7 +1355,7 @@ def space_list(self) -> list[int]:
         return n_list
 
     def simulate_pattern(
-        self, backend: Backend | str = "statevector", input_state: Data = BasicStates.PLUS, **kwargs: Any
+        self, backend: Backend[_StateT_co] | str = "statevector", input_state: Data = BasicStates.PLUS, **kwargs: Any
     ) -> BackendState:
         """Simulate the execution of the pattern by using :class:`graphix.simulator.PatternSimulator`.
 
@@ -1644,7 +1645,7 @@ def measure_pauli(pattern: Pattern, leave_input: bool, copy: bool = False) -> Pa
     nodes, edges = pattern.get_graph()
     vop_init = pattern.get_vops(conj=False)
     graph_state = GraphState(nodes=nodes, edges=edges, vops=vop_init)
-    results: dict[int, int] = {}
+    results: dict[int, Outcome] = {}
     to_measure, non_pauli_meas = pauli_nodes(pattern, leave_input)
     if not leave_input and len(list(set(pattern.input_nodes) & {i[0].node for i in to_measure})) > 0:
         new_inputs = []
@@ -1684,7 +1685,7 @@ def measure_pauli(pattern: Pattern, leave_input: bool, copy: bool = False) -> Pa
         if basis.sign == Sign.PLUS:
             results[pattern_cmd.node] = measure(pattern_cmd.node, choice=0)
         else:
-            results[pattern_cmd.node] = 1 - measure(pattern_cmd.node, choice=1)
+            results[pattern_cmd.node] = 0 if measure(pattern_cmd.node, choice=1) else 1
 
     # measure (remove) isolated nodes. if they aren't Pauli measurements,
     # measuring one of the results with probability of 1 should not occur as was possible above for Pauli measurements,

graphix/sim/__init__.py

@@ -1 +1,10 @@
 """Simulation backends."""
+
+from __future__ import annotations
+
+from graphix.sim.base_backend import Backend, BackendState
+from graphix.sim.data import Data
+from graphix.sim.density_matrix import DensityMatrix
+from graphix.sim.statevec import Statevec
+
+__all__ = ["Backend", "BackendState", "Data", "DensityMatrix", "Statevec"]