It works? but slowly

.gitignore

@@ -1,3 +1,5 @@
 .direnv
 
+**/__pycache__
+
 *.pdf

flake.lock

@@ -2,11 +2,11 @@
   "nodes": {
     "crane": {
       "locked": {
-        "lastModified": 1755993354,
+        "lastModified": 1767461147,
-        "narHash": "sha256-FCRRAzSaL/+umLIm3RU3O/+fJ2ssaPHseI2SSFL8yZU=",
+        "narHash": "sha256-TH/xTeq/RI+DOzo+c+4F431eVuBpYVwQwBxzURe7kcI=",
         "owner": "ipetkov",
         "repo": "crane",
-        "rev": "25bd41b24426c7734278c2ff02e53258851db914",
+        "rev": "7d59256814085fd9666a2ae3e774dc5ee216b630",
         "type": "github"
       },
       "original": {
@@ -24,11 +24,11 @@
         "rust-analyzer-src": "rust-analyzer-src"
       },
       "locked": {
-        "lastModified": 1755585599,
+        "lastModified": 1767596244,
-        "narHash": "sha256-tl/0cnsqB/Yt7DbaGMel2RLa7QG5elA8lkaOXli6VdY=",
+        "narHash": "sha256-P4NRZUjYbeuzv4hGrXxfdg0QpdGVoeNn0CMmzIyr398=",
         "owner": "nix-community",
         "repo": "fenix",
-        "rev": "6ed03ef4c8ec36d193c18e06b9ecddde78fb7e42",
+        "rev": "eedfb5a27900e82ec0390acc83d4d226ce86e714",
         "type": "github"
       },
       "original": {
@@ -42,11 +42,11 @@
         "nixpkgs-lib": "nixpkgs-lib"
       },
       "locked": {
-        "lastModified": 1754487366,
+        "lastModified": 1767609335,
-        "narHash": "sha256-pHYj8gUBapuUzKV/kN/tR3Zvqc7o6gdFB9XKXIp1SQ8=",
+        "narHash": "sha256-feveD98mQpptwrAEggBQKJTYbvwwglSbOv53uCfH9PY=",
         "owner": "hercules-ci",
         "repo": "flake-parts",
-        "rev": "af66ad14b28a127c5c0f3bbb298218fc63528a18",
+        "rev": "250481aafeb741edfe23d29195671c19b36b6dca",
         "type": "github"
       },
       "original": {
@@ -57,11 +57,11 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1765270179,
+        "lastModified": 1767364772,
-        "narHash": "sha256-g2a4MhRKu4ymR4xwo+I+auTknXt/+j37Lnf0Mvfl1rE=",
+        "narHash": "sha256-fFUnEYMla8b7UKjijLnMe+oVFOz6HjijGGNS1l7dYaQ=",
         "owner": "nixos",
         "repo": "nixpkgs",
-        "rev": "677fbe97984e7af3175b6c121f3c39ee5c8d62c9",
+        "rev": "16c7794d0a28b5a37904d55bcca36003b9109aaa",
         "type": "github"
       },
       "original": {
@@ -73,11 +73,11 @@
     },
     "nixpkgs-lib": {
       "locked": {
-        "lastModified": 1753579242,
+        "lastModified": 1765674936,
-        "narHash": "sha256-zvaMGVn14/Zz8hnp4VWT9xVnhc8vuL3TStRqwk22biA=",
+        "narHash": "sha256-k00uTP4JNfmejrCLJOwdObYC9jHRrr/5M/a/8L2EIdo=",
         "owner": "nix-community",
         "repo": "nixpkgs.lib",
-        "rev": "0f36c44e01a6129be94e3ade315a5883f0228a6e",
+        "rev": "2075416fcb47225d9b68ac469a5c4801a9c4dd85",
         "type": "github"
       },
       "original": {
@@ -95,11 +95,11 @@
     "rust-analyzer-src": {
       "flake": false,
       "locked": {
-        "lastModified": 1755504847,
+        "lastModified": 1767551763,
-        "narHash": "sha256-VX0B9hwhJypCGqncVVLC+SmeMVd/GAYbJZ0MiiUn2Pk=",
+        "narHash": "sha256-lcA/e3++3aZQSj6xCsBi2VpYyC3Q+oO/oukgfHiL+Ts=",
         "owner": "rust-lang",
         "repo": "rust-analyzer",
-        "rev": "a905e3b21b144d77e1b304e49f3264f6f8d4db75",
+        "rev": "6a1246b69ca761480b9278df019f717b549cface",
         "type": "github"
       },
       "original": {
@@ -145,18 +145,36 @@
           "nixpkgs"
         ],
         "rust-manifest": "rust-manifest",
-        "systems": "systems"
+        "systems": "systems",
+        "typst": "typst_2"
       },
       "locked": {
-        "lastModified": 1765380218,
+        "lastModified": 1767654408,
-        "narHash": "sha256-rMeYAuKaQ9U3n5yPc0XoGUOBS6u+kakGSZk+RBhM/WU=",
+        "narHash": "sha256-7YhRyF2u7CuhBLMpKnHZy9D5Y2L6k9LRLGf3SE0AUYc=",
         "owner": "typst",
-        "repo": "typst",
+        "repo": "typst-flake",
-        "rev": "4c1db395be2ea8090658b172b8493adfa4f55cf7",
+        "rev": "ec2389e45e1014c746a3a1d6a222cde352c56886",
         "type": "github"
       },
       "original": {
         "owner": "typst",
+        "repo": "typst-flake",
+        "type": "github"
+      }
+    },
+    "typst_2": {
+      "flake": false,
+      "locked": {
+        "lastModified": 1765535432,
+        "narHash": "sha256-KCNFCl7vFWHGmQPrie1BoncQtu5G/rN5qf45jPiYrT4=",
+        "owner": "typst",
+        "repo": "typst",
+        "rev": "b33de9de113c91c184214b299bd7a8eb3070c3ab",
+        "type": "github"
+      },
+      "original": {
+        "owner": "typst",
+        "ref": "0.14",
         "repo": "typst",
         "type": "github"
       }

flake.nix

@@ -2,7 +2,7 @@
   inputs = {
     nixpkgs.url = "github:nixos/nixpkgs/nixpkgs-unstable";
     typst = {
-      url = "github:typst/typst";
+      url = "github:typst/typst-flake";
       inputs.nixpkgs.follows = "nixpkgs";
     };
   };
@@ -21,7 +21,7 @@
           system:
           f {
             inherit system;
-            typst = typst.packages.${system};
+            typ = typst.packages.${system};
             pkgs = import nixpkgs { inherit system; };
           }
         );
@@ -30,17 +30,19 @@
       devShells = forAllSystems (
         {
           pkgs,
-          typst,
+          typ,
           system,
           ...
         }:
         {
           default = pkgs.mkShellNoCC {
             packages = [
-              pkgs.typst
+              typ.default
               (pkgs.python3.withPackages (ppkgs: [
                 ppkgs.numpy
                 ppkgs.tqdm
+                ppkgs.qiskit
+                ppkgs.qiskit-aer
               ]))
             ];
           };

planning/plan.typ

@@ -1,4 +1,3 @@
-
 #title("Quality-Diversity Quantum Architecture Search")
 
 = Introduction

src/qas_flow/__init__.py

@@ -0,0 +1,4 @@
+from .stream import Stream
+from .funcs import map, filter, take, skip, batch, enumerate, collect
+
+__all__ = ["Stream", "map", "filter", "take", "skip", "batch", "enumerate", "collect"]

src/qas_flow/funcs.py

@@ -0,0 +1,76 @@
+from typing import Callable
+from .stream import Stream, T, U
+
+
+@Stream.extension()
+def map(stream: Stream[T], fn: Callable[[T], U]) -> Stream[U]:
+    def gen():
+        for x in stream:
+            yield fn(x)
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def filter(stream: Stream[T], pred: Callable[[T], bool]) -> Stream[T]:
+    def gen():
+        for x in stream:
+            if pred(x):
+                yield x
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def take(stream: Stream[T], n: int) -> Stream[T]:
+    def gen():
+        c = 0
+        for x in stream:
+            if c < n:
+                c += 1
+                yield x
+            else:
+                return
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def skip(stream: Stream[T], n: int) -> Stream[T]:
+    def gen():
+        c = 0
+        for x in stream:
+            c += 1
+            if c > n:
+                yield x
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def batch(stream: Stream[T], n: int) -> Stream[list[T]]:
+    def gen():
+        ls: list[T] = []
+        for x in stream:
+            ls.append(x)
+            if len(ls) == n:
+                yield ls
+                ls = []
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def enumerate(stream: Stream[T]) -> Stream[tuple[int, T]]:
+    def gen():
+        idx = 0
+        for x in stream:
+            yield (idx, x)
+            idx += 1
+
+    return Stream(gen())
+
+
+@Stream.extension()
+def collect(stream: Stream[T]) -> list[T]:
+    return [v for v in stream]

src/qas_flow/stream.py

@@ -0,0 +1,39 @@
+from collections.abc import Iterator
+from typing import Any, Callable, Generic, TypeVar, final
+
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+Op = Callable[["Stream[T]"], "Stream[U]"]
+
+
+@final
+class Stream(Generic[T]):
+    _extensions: dict[str, Callable[..., Any]] = {}
+
+    def __init__(self, it: Iterator[T]) -> None:
+        self._it = it
+
+    def __iter__(self) -> Iterator[T]:
+        return self._it
+
+    @classmethod
+    def extension(cls, name: str | None = None):
+        """Register a function as Stream.<name>(...). First arg will be the stream."""
+
+        def deco(fn: Callable[..., Any]):
+            cls._extensions[name or fn.__name__] = fn
+            return fn
+
+        return deco
+
+    def __getattr__(self, attr: str):
+        fn = self._extensions.get(attr)
+        if fn is None:
+            raise AttributeError(attr)
+
+        def bound(*args, **kwargs):
+            return fn(self, *args, **kwargs)
+
+        return bound

src/tf-qas.py

@@ -0,0 +1,252 @@
+#!/usr/bin/env python
+from dataclasses import dataclass
+from enum import IntEnum
+import math
+import random
+import numpy as np
+from typing import assert_type, override
+from qiskit.circuit import ParameterVector
+from tqdm import tqdm
+from qiskit import QuantumCircuit as QiskitCircuit, transpile
+from qiskit_aer import AerSimulator
+
+from qas_flow import Stream
+
+TWO_QUBIT_GATE_PROBABILITY = 0.5
+
+
+class GateType(IntEnum):
+    RX = 1
+    RY = 2
+    RZ = 3
+    XX = 4
+    YY = 5
+    ZZ = 6
+
+
+@dataclass(frozen=True)
+class Gate:
+    type: GateType
+    qubits: int | tuple[int, int]
+    param_idx: int
+
+
+def haar_fidelity_pdf(F: np.ndarray, d: int) -> np.ndarray:
+    # p(F) = (d-1) * (1-F)^(d-2), for F in [0,1]
+    return (d - 1) * np.power(1.0 - F, d - 2)
+
+
+def kl_divergence(p: np.ndarray, q: np.ndarray) -> float:
+    # assumes p, q are normalized and strictly positive
+    return float(np.sum(p * np.log(p / q)))
+
+
+@dataclass()
+class QuantumCircuit:
+    qubits: int
+    gates: list[list[Gate]]
+    single_qubit_gates: int
+    two_qubit_gates: int
+    params: int
+    paths: int = 0
+    expressibility: float = -10000.0
+
+    def calculate_paths(self):
+        path_counts = [1 for _ in range(self.qubits)]
+        for layer in self.gates:
+            for gate in layer:
+                if gate.type <= 3 or type(gate.qubits) is int:
+                    continue
+                (q1, q2) = gate.qubits
+                path_counts[q1] = path_counts[q1] + path_counts[q2]
+                path_counts[q2] = path_counts[q1]
+
+        self.paths = sum(path_counts)
+
+    def to_qiskit(self):
+        qc = QiskitCircuit(self.qubits)
+        thetas = ParameterVector("theta", self.params)
+
+        for layer in self.gates:
+            for gate in layer:
+                theta = thetas[gate.param_idx]
+                if gate.type == GateType.RX:
+                    qc.rx(theta, gate.qubits)
+                elif gate.type == GateType.RY:
+                    qc.ry(theta, gate.qubits)
+                elif gate.type == GateType.RZ:
+                    qc.rz(theta, gate.qubits)
+                elif gate.type == GateType.XX:
+                    qc.rxx(theta, *gate.qubits)
+                elif gate.type == GateType.YY:
+                    qc.ryy(theta, *gate.qubits)
+                elif gate.type == GateType.ZZ:
+                    qc.rzz(theta, *gate.qubits)
+        qc.save_statevector()
+        return qc, thetas
+
+    def expressibility_estimate(
+        self, samples: int, seed: int, bins: int = 75, eps: float = 1e-12
+    ):
+        qc, thetas = self.to_qiskit()
+
+        if self.params <= 0:
+            return float("inf")
+
+        rng = random.Random(seed)
+        d = 1 << self.qubits
+
+        backend = AerSimulator(method="statevector", seed_simulator=seed)
+
+        tqc = transpile(qc, backend)
+
+        # 1) build 2*samples parameterized circuits (left+right)
+        binds = []
+        for _ in range(2 * samples):
+            binds.append(
+                {thetas: [rng.random() for _ in range(self.params)]}
+            )  # SAFE binding
+
+        # 2) compute statevectors (no backend/job overhead)
+        job = backend.run(
+            [tqc.assign_parameters(bind, inplace=False) for bind in binds]
+        )
+        result = job.result()
+        sv = [result.get_statevector(i) for i in range(len(binds))]
+        left = sv[:samples]
+        right = sv[samples:]
+
+        # 3) fidelities F = |<ψ|φ>|^2
+        inners = np.array(
+            [l.inner(r) for l, r in zip(left, right)], dtype=np.complex128
+        )
+        F = (inners * inners.conjugate()).real
+        F = np.clip(F, 0.0, 1.0)  # numerical safety
+
+        # 4) empirical histogram (as a probability mass over bins)
+        hist, edges = np.histogram(F, bins=bins, range=(0.0, 1.0), density=False)
+        p = hist.astype(np.float64)
+        p = p + eps
+        p = p / p.sum()
+
+        # 5) Haar distribution mass per bin (integrate PDF over each bin)
+        # approximate via midpoint rule
+        mids = 0.5 * (edges[:-1] + edges[1:])
+        widths = edges[1:] - edges[:-1]
+        q = haar_fidelity_pdf(mids, d) * widths
+        q = q + eps
+        q = q / q.sum()
+
+        kl = kl_divergence(p, q)
+        self.expressibility = kl
+        return kl
+
+    @override
+    def __str__(self):
+        strs = ["-" for _ in range(self.qubits)]
+        for layer in self.gates:
+            for i in range(self.qubits):
+                strs[i] += "---"
+
+            idx = 0
+
+            for gate in layer:
+                match gate.type:
+                    case GateType.RX:
+                        strs[gate.qubits] = strs[gate.qubits][:-2] + "RX"
+                    case GateType.RY:
+                        strs[gate.qubits] = strs[gate.qubits][:-2] + "RY"
+                    case GateType.RZ:
+                        strs[gate.qubits] = strs[gate.qubits][:-2] + "RZ"
+                    case GateType.XX:
+                        (q1, q2) = gate.qubits
+                        strs[q1] = strs[q1][:-2] + f"X{idx}"
+                        strs[q2] = strs[q2][:-2] + f"X{idx}"
+                        idx += 1
+                    case GateType.YY:
+                        (q1, q2) = gate.qubits
+                        strs[q1] = strs[q1][:-2] + f"Y{idx}"
+                        strs[q2] = strs[q2][:-2] + f"Y{idx}"
+                        idx += 1
+                    case GateType.ZZ:
+                        (q1, q2) = gate.qubits
+                        strs[q1] = strs[q1][:-2] + f"Z{idx}"
+                        strs[q2] = strs[q2][:-2] + f"Z{idx}"
+                        idx += 1
+        return (
+            "\n".join(strs)
+            + f"\npaths: {self.paths}, expressibility: {self.expressibility}"
+        )
+
+
+def even_parity(qubits: int):
+    return [(x, x + 1) for x in range(0, qubits, 2)]
+
+
+def odd_parity(qubits: int):
+    return [(x, x + 1) for x in range(1, qubits, 2)]
+
+
+def sample_circuit(rng: random.Random, qubits: int, depth: int) -> QuantumCircuit:
+    even = even_parity(qubits)
+    odd = odd_parity(qubits)
+
+    total_single = 0
+    total_double = 0
+    params = 0
+
+    gates: list[list[Gate]] = []
+    for _ in range(depth):
+        gate_type_offset = 3 if rng.random() < TWO_QUBIT_GATE_PROBABILITY else 0
+        gate_type = rng.randint(1, 3)
+        gate_locations = even if rng.random() < 0.5 else odd
+        if gate_type_offset == 0:
+            gates.append(
+                [Gate(GateType(gate_type), x, params) for (x, _) in gate_locations]
+            )
+            total_single += len(gate_locations)
+        else:
+            gates.append(
+                [
+                    Gate(GateType(gate_type + gate_type_offset), xy, params)
+                    for xy in gate_locations
+                    if xy[1] != qubits
+                ]
+            )
+            total_double += len(gate_locations)
+        params += 1
+    return QuantumCircuit(qubits, gates, total_single, total_double, params)
+
+
+def sample_generator(rng: random.Random, qubits: int, depth: int):
+    while True:
+        yield sample_circuit(rng, qubits, depth)
+
+
+def more_single_than_double(qc: QuantumCircuit) -> bool:
+    return qc.single_qubit_gates >= qc.two_qubit_gates
+
+
+if __name__ == "__main__":
+    rng = random.Random()
+    qubits = 6
+    depth = 10
+    sample_amount = 50000
+    expressibility_samples = 2000
+    proxy_pass_amount = 5000
+    circuits = (
+        Stream(sample_generator(rng, qubits, depth))
+        .filter(more_single_than_double)
+        .take(sample_amount)
+        .collect()
+    )
+    for circuit in tqdm(circuits):
+        circuit.calculate_paths()
+    circuits.sort(key=lambda qc: qc.paths, reverse=True)
+    for circuit in tqdm(circuits[:proxy_pass_amount]):
+        circuit.expressibility_estimate(
+            expressibility_samples, rng.randint(0, 100000000)
+        )
+    circuits.sort(key=lambda qc: qc.expressibility)
+    for i, circuit in enumerate(circuits[:proxy_pass_amount]):
+        print(f"circuit {i}:\n{circuit}")