Implement layer inference

.gitignore

@@ -187,3 +187,4 @@ cobertura.xml
 docs/build
 !docs/build/.nojekyll
 uv.lock
+lcov.info

docs/source/swiflow.rst

@@ -6,7 +6,7 @@ swiflow.common module
 
 .. automodule:: swiflow.common
    :members:
-   :exclude-members: Plane, PPlane, V, P
+   :exclude-members: Plane, PPlane
 
    .. autoclass:: swiflow.common.Plane
 

examples/flow.py

@@ -14,7 +14,7 @@
 # 1 - 3 - 5
 #     |
 # 2 - 4 - 6
-g = nx.Graph([(1, 3), (2, 4), (3, 5), (4, 6)])
+g = nx.Graph([(1, 3), (2, 4), (3, 5), (4, 6), (3, 4)])
 iset = {1, 2}
 oset = {5, 6}
 

examples/gflow.py

@@ -22,7 +22,7 @@
 oset = {4, 5}
 planes = {0: Plane.XY, 1: Plane.XY, 2: Plane.XZ, 3: Plane.YZ}
 
-result = gflow.find(g, iset, oset, planes)
+result = gflow.find(g, iset, oset, plane=planes)
 
 # Found
 assert result is not None

examples/pflow.py

@@ -20,7 +20,7 @@
 oset = {4}
 pplanes = {0: PPlane.Z, 1: PPlane.Z, 2: PPlane.Y, 3: PPlane.Y}
 
-result = pflow.find(g, iset, oset, pplanes)
+result = pflow.find(g, iset, oset, pplane=pplanes)
 
 # Found
 assert result is not None

pyproject.toml

@@ -65,7 +65,7 @@ python-source = "python"
 [tool.mypy]
 python_version = "3.9"
 strict = true
-files = ["docs/source/conf.py", "python", "tests"]
+files = ["docs/source/conf.py", "examples", "python", "tests"]
 
 [tool.pyright]
 reportUnknownArgumentType = "information"
@@ -123,6 +123,10 @@ required-imports = ["from __future__ import annotations"]
 [tool.ruff.lint.pydocstyle]
 convention = "numpy"
 
+[tool.ruff.lint.pylint]
+max-positional-args = 4
+max-args = 12
+
 [tool.ruff.lint.per-file-ignores]
 "docs/**/*.py" = [
   "D1",  # undocumented-XXX

python/swiflow/__init__.py

@@ -1 +1 @@
-"""Initialize the swiflow package."""
+"""swiflow: Rust binding of generalized and pauli flow finding algorithms."""

python/swiflow/_common.py

@@ -2,17 +2,21 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Iterable, Mapping
+import itertools
+from collections.abc import Callable, Hashable, Iterable, Mapping
 from collections.abc import Set as AbstractSet
-from typing import Generic
+from typing import Generic, TypeVar
 
 import networkx as nx
+from typing_extensions import ParamSpec
 
 from swiflow._impl import FlowValidationMessage
-from swiflow.common import P, S, T, V
+from swiflow.common import PPlane
 
+_V = TypeVar("_V", bound=Hashable)
 
-def check_graph(g: nx.Graph[V], iset: AbstractSet[V], oset: AbstractSet[V]) -> None:
+
+def check_graph(g: nx.Graph[_V], iset: AbstractSet[_V], oset: AbstractSet[_V]) -> None:
     """Check if `(g, iset, oset)` is a valid open graph for MBQC.
 
     Raises
@@ -46,7 +50,7 @@ def check_graph(g: nx.Graph[V], iset: AbstractSet[V], oset: AbstractSet[V]) -> N
         raise ValueError(msg)
 
 
-def check_planelike(vset: AbstractSet[V], oset: AbstractSet[V], plike: Mapping[V, P]) -> None:
+def check_planelike(vset: AbstractSet[_V], oset: AbstractSet[_V], plike: Mapping[_V, _P]) -> None:
     r"""Check if measurement description is valid.
 
     Parameters
@@ -80,13 +84,33 @@ def check_planelike(vset: AbstractSet[V], oset: AbstractSet[V], plike: Mapping[V
         raise ValueError(msg)
 
 
-class IndexMap(Generic[V]):
+def check_layer(layer: Mapping[_V, int]) -> None:
+    """Check if layer range is compatible with the current implementation."""
+    if min(layer.values(), default=0) != 0:
+        msg = "Minimum layer must be 0."
+        raise ValueError(msg)
+
+
+def odd_neighbors(g: nx.Graph[_V], kset: AbstractSet[_V]) -> set[_V]:
+    """Compute odd neighbors of `kset` in `g`."""
+    ret: set[_V] = set()
+    for k in kset:
+        ret.symmetric_difference_update(g.neighbors(k))
+    return ret
+
+
+_T = TypeVar("_T")
+_P = TypeVar("_P")
+_S = ParamSpec("_S")
+
+
+class IndexMap(Generic[_V]):
     """Map between `V` and 0-based indices."""
 
-    __v2i: dict[V, int]
-    __i2v: list[V]
+    __v2i: dict[_V, int]
+    __i2v: list[_V]
 
-    def __init__(self, vset: AbstractSet[V]) -> None:
+    def __init__(self, vset: AbstractSet[_V]) -> None:
         """Initialize the map from `vset`.
 
         Parameters
@@ -105,7 +129,7 @@ def __init__(self, vset: AbstractSet[V]) -> None:
             self.__i2v = list(vset)
         self.__v2i = {v: i for i, v in enumerate(self.__i2v)}
 
-    def encode(self, v: V) -> int:
+    def encode(self, v: _V) -> int:
         """Encode `v` to the index.
 
         Returns
@@ -124,7 +148,7 @@ def encode(self, v: V) -> int:
             raise ValueError(msg)
         return ind
 
-    def encode_graph(self, g: nx.Graph[V]) -> list[set[int]]:
+    def encode_graph(self, g: nx.Graph[_V]) -> list[set[int]]:
         """Encode graph.
 
         Returns
@@ -133,11 +157,11 @@ def encode_graph(self, g: nx.Graph[V]) -> list[set[int]]:
         """
         return [self.encode_set(g[v].keys()) for v in self.__i2v]
 
-    def encode_set(self, vset: AbstractSet[V]) -> set[int]:
+    def encode_set(self, vset: AbstractSet[_V]) -> set[int]:
         """Encode set."""
         return {self.encode(v) for v in vset}
 
-    def encode_dictkey(self, mapping: Mapping[V, P]) -> dict[int, P]:
+    def encode_dictkey(self, mapping: Mapping[_V, _P]) -> dict[int, _P]:
         """Encode dict key.
 
         Returns
@@ -146,7 +170,7 @@ def encode_dictkey(self, mapping: Mapping[V, P]) -> dict[int, P]:
         """
         return {self.encode(k): v for k, v in mapping.items()}
 
-    def encode_flow(self, f: Mapping[V, V]) -> dict[int, int]:
+    def encode_flow(self, f: Mapping[_V, _V]) -> dict[int, int]:
         """Encode flow.
 
         Returns
@@ -155,7 +179,7 @@ def encode_flow(self, f: Mapping[V, V]) -> dict[int, int]:
         """
         return {self.encode(i): self.encode(j) for i, j in f.items()}
 
-    def encode_gflow(self, f: Mapping[V, AbstractSet[V]]) -> dict[int, set[int]]:
+    def encode_gflow(self, f: Mapping[_V, AbstractSet[_V]]) -> dict[int, set[int]]:
         """Encode gflow.
 
         Returns
@@ -164,7 +188,7 @@ def encode_gflow(self, f: Mapping[V, AbstractSet[V]]) -> dict[int, set[int]]:
         """
         return {self.encode(i): self.encode_set(si) for i, si in f.items()}
 
-    def encode_layer(self, layer: Mapping[V, int]) -> list[int]:
+    def encode_layer(self, layer: Mapping[_V, int]) -> list[int]:
         """Encode layer.
 
         Returns
@@ -181,7 +205,7 @@ def encode_layer(self, layer: Mapping[V, int]) -> list[int]:
             msg = "Layers must be specified for all nodes."
             raise ValueError(msg) from None
 
-    def decode(self, i: int) -> V:
+    def decode(self, i: int) -> _V:
         """Decode the index.
 
         Returns
@@ -200,11 +224,11 @@ def decode(self, i: int) -> V:
             raise ValueError(msg) from None
         return v
 
-    def decode_set(self, iset: AbstractSet[int]) -> set[V]:
+    def decode_set(self, iset: AbstractSet[int]) -> set[_V]:
         """Decode set."""
         return {self.decode(i) for i in iset}
 
-    def decode_flow(self, f_: Mapping[int, int]) -> dict[V, V]:
+    def decode_flow(self, f_: Mapping[int, int]) -> dict[_V, _V]:
         """Decode MBQC flow.
 
         Returns
@@ -213,7 +237,7 @@ def decode_flow(self, f_: Mapping[int, int]) -> dict[V, V]:
         """
         return {self.decode(i): self.decode(j) for i, j in f_.items()}
 
-    def decode_gflow(self, f_: Mapping[int, AbstractSet[int]]) -> dict[V, set[V]]:
+    def decode_gflow(self, f_: Mapping[int, AbstractSet[int]]) -> dict[_V, set[_V]]:
         """Decode MBQC gflow.
 
         Returns
@@ -222,7 +246,7 @@ def decode_gflow(self, f_: Mapping[int, AbstractSet[int]]) -> dict[V, set[V]]:
         """
         return {self.decode(i): self.decode_set(si) for i, si in f_.items()}
 
-    def decode_layer(self, layer_: Iterable[int]) -> dict[V, int]:
+    def decode_layer(self, layer_: Iterable[int]) -> dict[_V, int]:
         """Decode MBQC layer.
 
         Returns
@@ -268,9 +292,100 @@ def decode_err(self, err: ValueError) -> ValueError:
             raise TypeError  # pragma: no cover
         return ValueError(msg)
 
-    def ecatch(self, f: Callable[S, T], *args: S.args, **kwargs: S.kwargs) -> T:
+    def ecatch(self, f: Callable[_S, _T], *args: _S.args, **kwargs: _S.kwargs) -> _T:
         """Wrap binding call to decode raw error messages."""
         try:
             return f(*args, **kwargs)
         except ValueError as e:
             raise self.decode_err(e) from None
+
+
+def _infer_layers_impl(gd: nx.DiGraph[_V]) -> Mapping[_V, int]:
+    """Fix flow layers one by one depending on order constraints."""
+    pred = {u: set(gd.predecessors(u)) for u in gd.nodes}
+    work = {u for u, pu in pred.items() if not pu}
+    ret: dict[_V, int] = {}
+    for l_now in itertools.count():
+        if not work:
+            break
+        next_work: set[_V] = set()
+        for u in work:
+            ret[u] = l_now
+            for v in gd.successors(u):
+                ent = pred[v]
+                ent.discard(u)
+                if not ent:
+                    next_work.add(v)
+        work = next_work
+    if len(ret) != len(gd):
+        msg = "Failed to determine layer for all nodes."
+        raise ValueError(msg)
+    return ret
+
+
+def _is_special(
+    pp: PPlane | None,
+    in_fu: bool,  # noqa: FBT001
+    in_fu_odd: bool,  # noqa: FBT001
+) -> bool:
+    if pp == PPlane.X:
+        return in_fu
+    if pp == PPlane.Y:
+        return in_fu and in_fu_odd
+    if pp == PPlane.Z:
+        return in_fu_odd
+    return False
+
+
+def _special_edges(
+    g: nx.Graph[_V],
+    anyflow: Mapping[_V, _V | AbstractSet[_V]],
+    pplane: Mapping[_V, PPlane] | None,
+) -> set[tuple[_V, _V]]:
+    """Compute special edges that can bypass partial order constraints in Pauli flow."""
+    ret: set[tuple[_V, _V]] = set()
+    if pplane is None:
+        return ret
+    for u, fu_ in anyflow.items():
+        fu = fu_ if isinstance(fu_, AbstractSet) else {fu_}
+        fu_odd = odd_neighbors(g, fu)
+        for v in itertools.chain(fu, fu_odd):
+            if u == v:
+                continue
+            if _is_special(pplane.get(v), v in fu, v in fu_odd):
+                ret.add((u, v))
+    return ret
+
+
+def infer_layers(
+    g: nx.Graph[_V],
+    anyflow: Mapping[_V, _V | AbstractSet[_V]],
+    pplane: Mapping[_V, PPlane] | None = None,
+) -> Mapping[_V, int]:
+    """Infer layer from flow/gflow using greedy algorithm.
+
+    Parameters
+    ----------
+    g : `networkx.Graph`
+        Simple graph representing MBQC pattern.
+    anyflow : `tuple` of flow-like/layer
+        Flow to verify. Compatible with both flow and generalized flow.
+    pplane : `collections.abc.Mapping`, optional
+        Measurement plane or Pauli index.
+
+    Notes
+    -----
+    This function operates in Pauli flow mode only when :py:obj`pplane` is explicitly given.
+    """
+    gd: nx.DiGraph[_V] = nx.DiGraph()
+    gd.add_nodes_from(g.nodes)
+    special = _special_edges(g, anyflow, pplane)
+    for u, fu_ in anyflow.items():
+        fu = fu_ if isinstance(fu_, AbstractSet) else {fu_}
+        fu_odd = odd_neighbors(g, fu)
+        for v in itertools.chain(fu, fu_odd):
+            if u == v or (u, v) in special:
+                continue
+            gd.add_edge(u, v)
+    gd = gd.reverse()
+    return _infer_layers_impl(gd)

python/swiflow/_impl/flow.pyi

@@ -1,2 +1,8 @@
 def find(g: list[set[int]], iset: set[int], oset: set[int]) -> tuple[dict[int, int], list[int]] | None: ...
-def verify(flow: tuple[dict[int, int], list[int]], g: list[set[int]], iset: set[int], oset: set[int]) -> None: ...
+def verify(
+    flow: tuple[dict[int, int], list[int]],
+    g: list[set[int]],
+    iset: set[int],
+    oset: set[int],
+    optimal: bool,
+) -> None: ...

python/swiflow/_impl/gflow.pyi

@@ -12,4 +12,5 @@ def verify(
     iset: set[int],
     oset: set[int],
     plane: dict[int, Plane],
+    optimal: bool,
 ) -> None: ...

python/swiflow/_impl/pflow.pyi

@@ -15,4 +15,5 @@ def verify(
     iset: set[int],
     oset: set[int],
     pplane: dict[int, PPlane],
+    optimal: bool,
 ) -> None: ...