Enable type-checking for visualization.py

examples/visualization.py

@@ -73,7 +73,7 @@
 edges = [(1, 4), (1, 6), (2, 4), (2, 5), (2, 6), (3, 5), (3, 6)]
 inputs = {1, 2, 3}
 outputs = {4, 5, 6}
-graph = nx.Graph()
+graph: nx.Graph[int] = nx.Graph()
 graph.add_nodes_from(nodes)
 graph.add_edges_from(edges)
 meas_planes = {1: Plane.XY, 2: Plane.XY, 3: Plane.XY}

graphix/generator.py

@@ -79,7 +79,7 @@ def generate_from_graph(
 
     # search for flow first
     f, l_k = graphix.gflow.find_flow(graph, inputs_set, outputs_set, meas_planes=meas_planes)
-    if f is not None:
+    if f is not None and l_k is not None:
         # flow found
         pattern = _flow2pattern(graph, angles, inputs, f, l_k)
         pattern.reorder_output_nodes(outputs)

graphix/gflow.py

@@ -102,7 +102,7 @@ def find_flow(
     iset: set[int],
     oset: set[int],
     meas_planes: dict[int, Plane] | None = None,
-) -> tuple[dict[int, set[int]], dict[int, int]]:
+) -> tuple[dict[int, set[int]], dict[int, int]] | tuple[None, None]:
     """Causal flow finding algorithm.
 
     For open graph g with input, output, and measurement planes, this returns causal flow.
@@ -282,7 +282,7 @@ def find_pauliflow(
     return pf[0], pf[1]
 
 
-def flow_from_pattern(pattern: Pattern) -> tuple[dict[int, set[int]], dict[int, int]]:
+def flow_from_pattern(pattern: Pattern) -> tuple[dict[int, set[int]], dict[int, int]] | tuple[None, None]:
     """Check if the pattern has a valid flow. If so, return the flow and layers.
 
     Parameters
@@ -292,9 +292,11 @@ def flow_from_pattern(pattern: Pattern) -> tuple[dict[int, set[int]], dict[int,
 
     Returns
     -------
-    f: dict
+    None, None:
+        The tuple ``(None, None)`` is returned if the pattern does not have a valid causal flow.
+    f: dict[int, set[int]]
         flow function. g[i] is the set of qubits to be corrected for the measurement of qubit i.
-    l_k: dict
+    l_k: dict[int, int]
         layers obtained by flow algorithm. l_k[d] is a node set of depth d.
     """
     if not pattern.is_standard(strict=True):
@@ -332,7 +334,7 @@ def flow_from_pattern(pattern: Pattern) -> tuple[dict[int, set[int]], dict[int,
     return None, None
 
 
-def gflow_from_pattern(pattern: Pattern) -> tuple[dict[int, set[int]], dict[int, int]]:
+def gflow_from_pattern(pattern: Pattern) -> tuple[dict[int, set[int]], dict[int, int]] | tuple[None, None]:
     """Check if the pattern has a valid gflow. If so, return the gflow and layers.
 
     Parameters
@@ -342,9 +344,11 @@ def gflow_from_pattern(pattern: Pattern) -> tuple[dict[int, set[int]], dict[int,
 
     Returns
     -------
-    g: dict
+    None, None:
+        The tuple ``(None, None)`` is returned if the pattern does not have a valid gflow.
+    g: dict[int, set[int]]
         gflow function. g[i] is the set of qubits to be corrected for the measurement of qubit i.
-    l_k: dict
+    l_k: dict[int, int]
         layers obtained by gflow algorithm. l_k[d] is a node set of depth d.
     """
     if not pattern.is_standard(strict=True):
@@ -404,9 +408,11 @@ def pauliflow_from_pattern(
 
     Returns
     -------
-    p: dict
+    None, None:
+        The tuple ``(None, None)`` is returned if the pattern does not have a valid Pauli flow.
+    p: dict[int, set[int]]
         Pauli flow function. p[i] is the set of qubits to be corrected for the measurement of qubit i.
-    l_k: dict
+    l_k: dict[int, int]
         layers obtained by Pauli flow algorithm. l_k[d] is a node set of depth d.
     """
     if not pattern.is_standard(strict=True):

graphix/pattern.py

@@ -994,11 +994,11 @@ def get_measurement_order_from_flow(self) -> list[int] | None:
             measurement order
         """
         graph = self.extract_graph()
-        vin = set(self.input_nodes) if self.input_nodes is not None else set()
+        vin = set(self.input_nodes)
         vout = set(self.output_nodes)
         meas_planes = self.get_meas_plane()
         f, l_k = find_flow(graph, vin, vout, meas_planes=meas_planes)
-        if f is None:
+        if f is None or l_k is None:
             return None
         depth, layer = get_layers(l_k)
         meas_order: list[int] = []
@@ -1020,7 +1020,7 @@ def get_measurement_order_from_gflow(self) -> list[int]:
         isolated = list(nx.isolates(graph))
         if isolated:
             raise ValueError("The input graph must be connected")
-        vin = set(self.input_nodes) if self.input_nodes is not None else set()
+        vin = set(self.input_nodes)
         vout = set(self.output_nodes)
         meas_planes = self.get_meas_plane()
         flow, l_k = find_gflow(graph, vin, vout, meas_planes=meas_planes)
@@ -1389,10 +1389,9 @@ def draw_graph(
         show_local_clifford: bool = False,
         show_measurement_planes: bool = False,
         show_loop: bool = True,
-        node_distance: tuple[int, int] = (1, 1),
+        node_distance: tuple[float, float] = (1, 1),
         figsize: tuple[int, int] | None = None,
-        save: bool = False,
-        filename: str | None = None,
+        filename: Path | None = None,
     ) -> None:
         """Visualize the underlying graph of the pattern with flow or gflow structure.
 
@@ -1412,13 +1411,12 @@ def draw_graph(
             Distance multiplication factor between nodes for x and y directions.
         figsize : tuple
             Figure size of the plot.
-        save : bool
-            If True, the plot is saved as a png file.
-        filename : str
-            Filename of the saved plot.
+        filename : Path | None
+            If not None, filename of the png file to save the plot. If None, the plot is not saved.
+            Default in None.
         """
         graph = self.extract_graph()
-        vin = self.input_nodes if self.input_nodes is not None else []
+        vin = self.input_nodes
         vout = self.output_nodes
         meas_planes = self.get_meas_plane()
         meas_angles = self.get_angles()
@@ -1435,7 +1433,6 @@ def draw_graph(
                 show_loop=show_loop,
                 node_distance=node_distance,
                 figsize=figsize,
-                save=save,
                 filename=filename,
             )
         else:
@@ -1446,7 +1443,6 @@ def draw_graph(
                 show_loop=show_loop,
                 node_distance=node_distance,
                 figsize=figsize,
-                save=save,
                 filename=filename,
             )