it works now?

src/ga_qas.py

@@ -3,6 +3,7 @@
 # "Genetic optimization of ansatz expressibility for enhanced variational quantum algorithm performance"
 
 import random
+import sys
 from typing import Generator, Never
 
 import matplotlib.pyplot as plt
@@ -11,14 +12,6 @@ from qas_flow import Stream
 from quantum_circuit import (Gate, GateType, QuantumCircuit, circ_from_layers,
                              sample_random_generator, single_typ)
 
-DEPTH: int = 6
-QUBITS: int = 6
-GENERATIONS: int = 40
-GENERATION_SIZE: int = 60
-PARENT_AMOUNT: int = 10
-MUTATION_RATE: float = 0.1
-
-
 gate_set: list[GateType] = [
     GateType.H,
     GateType.RX,
@@ -77,47 +70,59 @@ def sample_hyperspace(
         yield point
 
 
-def plot_best_circuits(best_circuits: list[QuantumCircuit]) -> None:
+EXPRESSIBILITY_SAMPLES: int = 2000
-    fig, ax = plt.subplots()
-
-    ax.plot([-circ.expressibility for circ in best_circuits])
-    fig.savefig("best_circuits.png")
 
 
-def main() -> None:
+def run_ga_qas(
-    seed_rng: random.Random = random.Random(1020381)
+    depth: int,
+    qubits: int,
+    generations: int,
+    generation_size: int,
+    parent_amount: int,
+    mutation_rate: float,
+    seed: int,
+) -> list[tuple[int, int, int, int, int, float, float, float]]:
+
+    print(
+        f"running GA QAS {seed} with {qubits} qubits, {depth} depth, for {generations} generations of size {generation_size}, with {parent_amount} parents and {mutation_rate:.3f} mutation rate",
+        file=sys.stderr,
+    )
+
+    seed_rng = random.Random(seed)
     initial_population: list[QuantumCircuit] = (
-        Stream(sample_random_generator(random.Random(101020), QUBITS, DEPTH, gate_set))
+        Stream(
-        .apply(lambda circ: print(circ))
+            sample_random_generator(
+                random.Random(seed_rng.randint(1000, 1000000000)),
+                qubits,
+                depth,
+                gate_set,
+            )
+        )
         .apply(
             lambda circ: circ.expressibility_estimate(
-                2000, seed_rng.randint(1000, 1000000000)
+                EXPRESSIBILITY_SAMPLES, seed_rng.randint(1000, 1000000000)
             )
         )
-        .apply(lambda circ: print(circ))
+        .take(generation_size)
-        .take(GENERATION_SIZE)
         .collect()
     )
 
-    population = initial_population
+    population: list[QuantumCircuit] = initial_population
-
+    main_rng = random.Random(seed_rng.randint(1000, 1000000000))
-    main_rng = random.Random(2837175)
 
     best_circuits: list[QuantumCircuit] = []
-
+    return_data: list[tuple[int, int, int, int, int, float, float, float]] = []
-    for generation in range(GENERATIONS):
+    for generation in range(generations):
-        print(f"starting generation {generation}")
+        print(f"starting generation {generation} for seed {seed}", file=sys.stderr)
         population.sort(key=lambda qc: qc.expressibility, reverse=True)
-        parents = population[:PARENT_AMOUNT]
+        parents: list[QuantumCircuit] = population[:parent_amount]
-        for parent in parents:
+        offspring: list[QuantumCircuit] = []
-            print(parent)
+        for _ in range(generation_size):
-        offspring = []
-        for _ in range(GENERATION_SIZE):
             [p1, p2] = main_rng.sample(parents, 2)
-            crossover_layer = main_rng.randint(1, DEPTH)
+            crossover_layer = main_rng.randint(1, depth)
             child_layers = p1.gates[:crossover_layer] + p2.gates[crossover_layer:]
-            if main_rng.random() < MUTATION_RATE:
+            if main_rng.random() < mutation_rate:
-                layer_idx = main_rng.randrange(DEPTH)
+                layer_idx = main_rng.randrange(depth)
                 layer = child_layers[layer_idx]
                 gate_idx = main_rng.randrange(len(layer))
                 old_gate = child_layers[layer_idx][gate_idx]
@@ -137,21 +142,83 @@ def main() -> None:
                         old_gate.param_idx,
                     )
 
-            child = circ_from_layers(child_layers, QUBITS)
+            child = circ_from_layers(child_layers, qubits)
             child.expressibility_estimate(2000, seed_rng.randint(1000, 1000000000))
             offspring.append(child)
 
         offspring.sort(key=lambda qc: qc.expressibility, reverse=True)
+        return_data.append(
+            (
+                depth,
+                qubits,
+                generation,
+                generation_size,
+                parent_amount,
+                mutation_rate,
+                population[0].expressibility,
+                offspring[0].expressibility,
+            )
+        )
         if population[0].expressibility > offspring[0].expressibility:
-            print(f"best parent > best child")
             best_circuits.append(population[0])
         else:
-            print(f"best child > best parent")
             best_circuits.append(offspring[0])
         population = offspring
+    print(f"finished seed {seed}", file=sys.stderr)
+    return return_data
 
-    plot_best_circuits(best_circuits)
+
-    plt.show()
+def run_from_point(pnt: tuple[tuple[int, int, int, int, float], int]):
+    (point, seed) = pnt
+    return run_ga_qas(
+        point[0],
+        point[1],
+        20,
+        point[2],
+        point[3],
+        point[4],
+        seed,
+    )
+
+
+def print_ret(ret_data):
+    for dat in ret_data:
+        (
+            depth,
+            qubits,
+            generation,
+            generation_size,
+            parent_amount,
+            mutation_rate,
+            best_pop,
+            best_offspring,
+        ) = dat
+        print(
+            f"{depth},{qubits},{generation},{generation_size},{parent_amount},{mutation_rate},{best_pop},{best_offspring}",
+            flush=True,
+        )
+
+
+def main() -> None:
+
+    rng = random.Random(123456789)
+
+    results: list[QuantumCircuit] = (
+        Stream(
+            sample_hyperspace(
+                (1, 40),
+                (1, 10),
+                (1, 100),
+                (1, 20),
+                (0.0, 1.0),
+                seed=rng.randint(1000, 1000000000),
+            )
+        )
+        .map(lambda point: (point, rng.randint(1000, 1000000000)))
+        .par_map(run_from_point)
+        .apply(print_ret)
+        .collect()
+    )
 
 
 if __name__ == "__main__":

src/qas_flow/funcs.py

@@ -1,4 +1,9 @@
-from typing import Callable
+import os
+import sys
+from collections.abc import Generator
+from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
+from typing import Any, Callable, Never
+
 from .stream import Stream, T, U
 
 
@@ -8,20 +13,48 @@ def map(stream: Stream[T], fn: Callable[[T], U]) -> Stream[U]:
     Applies the function `fn` to every element of the stream
     """
 
-    def gen():
+    def gen() -> Generator[U, Any, None]:
         for x in stream:
             yield fn(x)
 
     return Stream(gen())
 
 
+@Stream.extension()
+def par_map(stream: Stream[T], fn: Callable[[T], U], cores: int = 0) -> Stream[U]:
+
+    def gen() -> Generator[U, Any, None]:
+        nprocs = (os.cpu_count() or 2) - 1 if cores <= 0 else cores
+        inflight = max(1, nprocs)
+
+        it = iter(stream)
+
+        with ProcessPoolExecutor(max_workers=nprocs) as ex:
+            pending = set()
+
+            for _ in range(inflight):
+                x = next(it)
+                pending.add(ex.submit(fn, x))
+
+            while True:
+                done, pending = wait(pending, return_when=FIRST_COMPLETED)
+
+                for fut in done:
+                    print(f"yielding: {fut}", file=sys.stderr)
+                    yield fut.result()
+                    x = next(it)
+                    pending.add(ex.submit(fn, x))
+
+    return Stream(gen())
+
+
 @Stream.extension()
 def filter(stream: Stream[T], pred: Callable[[T], bool]) -> Stream[T]:
     """
     Applies the predicate `pred` to every element and only returns elements where the predicate is true
     """
 
-    def gen():
+    def gen() -> Generator[T, Any, None]:
         for x in stream:
             if pred(x):
                 yield x
@@ -35,7 +68,7 @@ def apply(stream: Stream[T], fn: Callable[[T], None]) -> Stream[T]:
     Apply the function `fn` to every element of the stream in-place
     """
 
-    def gen():
+    def gen() -> Generator[T, Any, None]:
         for x in stream:
             fn(x)
             yield x
@@ -49,7 +82,7 @@ def take(stream: Stream[T], n: int) -> Stream[T]:
     Return a stream with at most `n` elements
     """
 
-    def gen():
+    def gen() -> Generator[T, Any, None]:
         c = 0
         for x in stream:
             if c < n:
@@ -67,7 +100,7 @@ def skip(stream: Stream[T], n: int) -> Stream[T]:
     Ignore the first `n` elements of a stream
     """
 
-    def gen():
+    def gen() -> Generator[T, Any, None]:
         c = 0
         for x in stream:
             c += 1
@@ -83,7 +116,7 @@ def batch(stream: Stream[T], n: int) -> Stream[list[T]]:
     Create batches of size `n` from the stream
     """
 
-    def gen():
+    def gen() -> Generator[list[T], Any, None]:
         ls: list[T] = []
         for x in stream:
             ls.append(x)
@@ -100,7 +133,7 @@ def enumerate(stream: Stream[T]) -> Stream[tuple[int, T]]:
     Add an index to each element of the stream
     """
 
-    def gen():
+    def gen() -> Generator[tuple[int, T], Any, None]:
         idx = 0
         for x in stream:
             yield (idx, x)

src/quantum_circuit.py

@@ -2,7 +2,7 @@ import math
 import random
 from dataclasses import dataclass
 from enum import IntEnum
-from typing import override
+from typing import Self, override
 
 import numpy as np
 from qiskit import QuantumCircuit as QiskitCircuit
@@ -167,7 +167,7 @@ class QuantumCircuit:
 
     def expressibility_estimate(
         self, samples: int, seed: int, bins: int = 75, eps: float = 1e-12
-    ) -> "QuantumCircuit":
+    ) -> Self:
         qc, thetas = self.to_qiskit_for_expressibility()
 
         if self.params <= 0: