From e011a14df523e6ecc5407cdae2d3d89c976c222e Mon Sep 17 00:00:00 2001 From: Noa Aarts Date: Mon, 12 Jan 2026 14:04:12 +0100 Subject: [PATCH] do the ground-truth check as well --- src/tf-qas.py | 442 +++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 398 insertions(+), 44 deletions(-) diff --git a/src/tf-qas.py b/src/tf-qas.py index 54a2348..eaf9bd0 100755 --- a/src/tf-qas.py +++ b/src/tf-qas.py @@ -1,19 +1,32 @@ #!/usr/bin/env python +from __future__ import annotations +import argparse +from itertools import repeat import json +import math +import random from dataclasses import dataclass from enum import IntEnum from multiprocessing import Pool -import random -import numpy as np from typing import override -from qiskit.circuit import ParameterVector -from tqdm import tqdm +import numpy as np from qiskit import QuantumCircuit as QiskitCircuit, transpile +from qiskit.circuit import ParameterVector, ParameterVectorElement from qiskit_aer import AerSimulator +from tqdm import tqdm from qas_flow import Stream +# ---------------------------- +# Paper hyperparameters/defaults +# ---------------------------- + TWO_QUBIT_GATE_PROBABILITY = 0.5 +CHEMICAL_ACCURACY = 1.6e-3 # Ha, success if E - E0 <= this + +# ---------------------------- +# Circuit representation +# ---------------------------- class GateType(IntEnum): @@ -32,7 +45,11 @@ class Gate: param_idx: int def to_json(self): - return {"type": self.type, "qubits": self.qubits, "param_idx": self.param_idx} + return { + "type": int(self.type), + "qubits": self.qubits, + "param_idx": self.param_idx, + } def haar_fidelity_pdf(F: np.ndarray, d: int) -> np.ndarray: @@ -55,16 +72,22 @@ class QuantumCircuit: paths: int = 0 expressibility: float = float("-inf") + # ground-truth fields (filled later) + gt_energy: float | None = None + gt_error: float | None = None + gt_steps: int | None = None + gt_success: bool | None = None + 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: + if gate.type <= 3 or isinstance(gate.qubits, int): continue (q1, q2) = gate.qubits + # same logic as your existing file path_counts[q1] = path_counts[q1] + path_counts[q2] path_counts[q2] = path_counts[q1] - self.paths = sum(path_counts) def to_json(self): @@ -74,9 +97,14 @@ class QuantumCircuit: "params": self.params, "paths": self.paths, "expressibility": self.expressibility, + "gt_energy": self.gt_energy, + "gt_error": self.gt_error, + "gt_steps": self.gt_steps, + "gt_success": self.gt_success, } - def to_qiskit(self): + def to_qiskit_ansatz(self) -> tuple[QiskitCircuit, ParameterVector]: + """Return *just* the parametrized ansatz defined by this circuit.""" qc = QiskitCircuit(self.qubits) thetas = ParameterVector("theta", self.params) @@ -95,27 +123,43 @@ class QuantumCircuit: qc.ryy(theta, *gate.qubits) elif gate.type == GateType.ZZ: qc.rzz(theta, *gate.qubits) + else: + raise ValueError(f"Unknown gate type: {gate.type}") + return qc, thetas + + def to_qiskit_for_expressibility(self) -> tuple[QiskitCircuit, ParameterVector]: + """Expressibility uses |0...0> input (no TFIM-specific H layer).""" + qc, thetas = self.to_qiskit_ansatz() + qc.save_statevector() + return qc, thetas + + def to_qiskit_for_tfim_vqe(self) -> tuple[QiskitCircuit, ParameterVector]: + """TFIM VQE begins with a layer of Hadamards on all qubits.""" + ans, thetas = self.to_qiskit_ansatz() + qc = QiskitCircuit(self.qubits) + qc.h(range(self.qubits)) + qc.compose(ans, inplace=True) 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() + ) -> "QuantumCircuit": + qc, thetas = self.to_qiskit_for_expressibility() if self.params <= 0: - return float("inf") + self.expressibility = float("-inf") + return self rng = random.Random(seed) d = 1 << self.qubits backend = AerSimulator(method="statevector", seed_simulator=seed) - tqc = transpile(qc, backend, optimization_level=0) nexp = 2 * samples - - binds = [ + # vectorized binds: one circuit, many parameter values + binds: list[dict[ParameterVectorElement, list[float]]] = [ {param: [rng.random() for _ in range(nexp)] for param in thetas.params} ] @@ -134,14 +178,10 @@ class QuantumCircuit: ) F = (inners.conjugate() * inners).real - # 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 = hist.astype(np.float64) + 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 @@ -160,7 +200,6 @@ class QuantumCircuit: strs[i] += "---" idx = 0 - for gate in layer: match gate.type: case GateType.RX: @@ -184,12 +223,21 @@ class QuantumCircuit: strs[q1] = strs[q1][:-2] + f"Z{idx}" strs[q2] = strs[q2][:-2] + f"Z{idx}" idx += 1 + gt = "" + if self.gt_energy is not None: + gt = f"\nGT: E={self.gt_energy:.12f}, err={self.gt_error:.3e}, steps={self.gt_steps}, success={self.gt_success}" return ( "\n".join(strs) + f"\npaths: {self.paths}, expressibility: {self.expressibility}" + + gt ) +# ---------------------------- +# Search space sampling (layerwise) +# ---------------------------- + + def even_parity(qubits: int): return [(x, x + 1) for x in range(0, qubits, 2)] @@ -211,6 +259,7 @@ def sample_circuit(rng: random.Random, qubits: int, depth: int) -> QuantumCircui 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] @@ -226,6 +275,7 @@ def sample_circuit(rng: random.Random, qubits: int, depth: int) -> QuantumCircui ) total_double += len(gate_locations) params += 1 + return QuantumCircuit(qubits, gates, total_single, total_double, params) @@ -238,43 +288,347 @@ def more_single_than_double(qc: QuantumCircuit) -> bool: return qc.single_qubit_gates >= qc.two_qubit_gates -if __name__ == "__main__": - rng = random.Random() - qubits = 4 - depth = 15 - sample_amount = 50000 - expressibility_samples = 2000 - proxy_pass_amount = 5000 +# ---------------------------- +# TFIM ground-truth VQE +# ---------------------------- + + +def _pauli_x(): + return np.array([[0.0, 1.0], [1.0, 0.0]], dtype=np.float64) + + +def _pauli_z(): + return np.array([[1.0, 0.0], [0.0, -1.0]], dtype=np.float64) + + +def _eye(): + return np.eye(2, dtype=np.float64) + + +def kron_n(ops: list[np.ndarray]) -> np.ndarray: + out = ops[0] + for op in ops[1:]: + out = np.kron(out, op) + return out + + +def tfim_hamiltonian_matrix(n: int, periodic: bool = True) -> np.ndarray: + """ + HTFIM = sum_{i=1}^n Z_i Z_{i+1} + X_i (paper definition) :contentReference[oaicite:5]{index=5} + with Z_{n+1}=Z_1 if periodic. + """ + X = _pauli_x() + Z = _pauli_z() + I = _eye() + dim = 1 << n + H = np.zeros((dim, dim), dtype=np.float64) + + # Xi terms + for i in range(n): + ops = [I] * n + ops[i] = X + H += kron_n(ops) + + # Zi Zi+1 terms + for i in range(n): + j = (i + 1) % n + if (not periodic) and (j == 0): + continue + ops = [I] * n + ops[i] = Z + ops[j] = Z + H += kron_n(ops) + + return H + + +def exact_ground_energy(H: np.ndarray) -> float: + # H is real symmetric here + w = np.linalg.eigvalsh(H) + return float(w[0]) + + +def statevector_from_bound_circuit( + backend: AerSimulator, tqc: QiskitCircuit, bind: dict +) -> np.ndarray: + res = backend.run([tqc], parameter_binds=[bind]).result() + sv = np.asarray(res.get_statevector(0), dtype=np.complex128) + return sv + + +def energy_expectation_from_sv(H: np.ndarray, sv: np.ndarray) -> float: + # E = + # H real symmetric; sv complex + return float(np.real(np.vdot(sv, H @ sv))) + + +def adam_optimize_tfim_energy( + circuit: QuantumCircuit, + H: np.ndarray, + seed: int, + lr: float = 0.01, # paper :contentReference[oaicite:6]{index=6} + max_steps: int = 2000, + tol: float = 1e-7, + beta1: float = 0.9, + beta2: float = 0.999, + eps: float = 1e-8, + param_shift: float = math.pi / 2, +) -> tuple[float, int]: + """ + Optimize parameters with Adam using parameter-shift gradients. + Returns: (best_energy, steps_taken) + """ + qc, thetas = circuit.to_qiskit_for_tfim_vqe() + + p = circuit.params + rng = np.random.default_rng(seed) + theta = rng.uniform( + -2 * math.pi, 2 * math.pi, size=p + ) # paper init range :contentReference[oaicite:7]{index=7} + + backend = AerSimulator(method="statevector", seed_simulator=seed) + tqc = transpile(qc, backend, optimization_level=0) + + def E(th: np.ndarray) -> float: + bind = {param: [val] for param, val in zip(thetas.params, th)} + sv = statevector_from_bound_circuit(backend, tqc, bind) + return energy_expectation_from_sv(H, sv) + + def grad(th: np.ndarray) -> np.ndarray: + # parameter-shift: d/dθ f(θ) = 0.5*(f(θ+π/2)-f(θ-π/2)) + g = np.zeros_like(th) + for i in range(p): + plus = th.copy() + minus = th.copy() + plus[i] += param_shift + minus[i] -= param_shift + g[i] = 0.5 * (E(plus) - E(minus)) + return g + + m = np.zeros_like(theta) + v = np.zeros_like(theta) + + best_E = float("inf") + prev_E = None + + for t in range(1, max_steps + 1): + g = grad(theta) + m = beta1 * m + (1 - beta1) * g + v = beta2 * v + (1 - beta2) * (g * g) + mhat = m / (1 - beta1**t) + vhat = v / (1 - beta2**t) + theta = theta - lr * mhat / (np.sqrt(vhat) + eps) + + cur_E = E(theta) + if cur_E < best_E: + best_E = cur_E + + if prev_E is not None and abs(prev_E - cur_E) < tol: + return best_E, t + prev_E = cur_E + + return best_E, max_steps + + +def ground_truth_tfim( + circ: QuantumCircuit, + H: np.ndarray, + E0: float, + seed: int, + lr: float, + max_steps: int, + tol: float, +) -> QuantumCircuit: + best_E, steps = adam_optimize_tfim_energy( + circ, H, seed=seed, lr=lr, max_steps=max_steps, tol=tol + ) + err = best_E - E0 + circ.gt_energy = best_E + circ.gt_error = err + circ.gt_steps = steps + circ.gt_success = err <= CHEMICAL_ACCURACY + return circ + + +# ---------------------------- +# Main +# ---------------------------- + + +def parse_args(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--qubits", + type=int, + default=6, + help="TFIM qubit count (paper uses 6). :contentReference[oaicite:8]{index=8}", + ) + ap.add_argument( + "--depth", + type=int, + default=10, + help="Max depth for TFIM layerwise (paper max depth 10). :contentReference[oaicite:9]{index=9}", + ) + ap.add_argument( + "--sample_amount", + type=int, + default=50000, + help="Number of sampled circuits (paper uses 50,000). :contentReference[oaicite:10]{index=10}", + ) + ap.add_argument( + "--expressibility_samples", + type=int, + default=2000, + help="Fidelity samples for expressibility (paper uses 2000). :contentReference[oaicite:11]{index=11}", + ) + ap.add_argument( + "--proxy_pass_amount", + type=int, + default=5000, + help="R: top circuits kept by paths (paper uses 5000). :contentReference[oaicite:12]{index=12}", + ) + ap.add_argument( + "--topk", + type=int, + default=100, + help="K: how many top circuits to output/store.", + ) + ap.add_argument( + "--seed", type=int, default=0, help="RNG seed for sampling/reproducibility." + ) + ap.add_argument( + "--periodic", + action="store_true", + help="Use periodic TFIM boundary (default true if set).", + ) + ap.add_argument( + "--no-periodic", + dest="periodic", + action="store_false", + help="Use open boundary TFIM.", + ) + ap.set_defaults(periodic=True) + + # ground-truth options + ap.add_argument( + "--do_ground_truth", + action="store_true", + help="Also evaluate ground-truth TFIM VQE performance.", + ) + ap.add_argument( + "--gt_budget", + type=int, + default=100, + help="Max number of circuits to query for ground-truth (paper often reports within 100 queries). :contentReference[oaicite:13]{index=13}", + ) + ap.add_argument( + "--gt_lr", + type=float, + default=0.01, + help="Adam learning rate (paper uses 0.01). :contentReference[oaicite:14]{index=14}", + ) + ap.add_argument( + "--gt_max_steps", + type=int, + default=2000, + help="Max Adam steps per circuit (practical cap).", + ) + ap.add_argument( + "--gt_tol", + type=float, + default=1e-7, + help="Convergence tolerance for |E_t - E_{t-1}|.", + ) + ap.add_argument("--dump", type=str, default="dump.json", help="Output JSON file.") + + return ap.parse_args() + + +def expr_worker(payload): + circ, seed, samples = payload + return circ.expressibility_estimate(samples, seed) + + +def main(): + args = parse_args() + rng = random.Random(args.seed) + + # --- Stage 0: sample circuits circuits = ( - Stream(sample_generator(rng, qubits, depth)) + Stream(sample_generator(rng, args.qubits, args.depth)) .filter(more_single_than_double) - .take(sample_amount) + .take(args.sample_amount) .collect() ) - for circuit in tqdm(circuits): - circuit.calculate_paths() - circuits.sort(key=lambda qc: qc.paths, reverse=True) - seeds = [rng.randint(0, 1000000000) for _ in range(proxy_pass_amount)] - def starmap_func(args): - (circ, seed, idx) = args - res = circ.expressibility_estimate(expressibility_samples, seed) - return res + # --- Stage 1: paths + for c in tqdm(circuits, desc="paths"): + c.calculate_paths() + circuits.sort(key=lambda qc: qc.paths, reverse=True) + candidates = circuits[: args.proxy_pass_amount] + + # --- Stage 2: expressibility + seeds = [rng.randint(0, 1_000_000_000) for _ in range(len(candidates))] final_circuits: list[QuantumCircuit] = [] with Pool() as p: for circ in tqdm( p.imap_unordered( - starmap_func, - zip(circuits[:proxy_pass_amount], seeds, range(proxy_pass_amount)), + expr_worker, zip(candidates, seeds, repeat(args.expressibility_samples)) ), - total=proxy_pass_amount, + total=len(candidates), + desc="expressibility", ): final_circuits.append(circ) final_circuits.sort(key=lambda qc: qc.expressibility, reverse=True) - for i, circuit in enumerate(final_circuits[::-1]): - print(f"circuit {i}:\n{circuit}") - with open("dump.json", "w+") as fp: - json.dump([c.to_json() for c in final_circuits], fp, indent=4) + # --- Ground-truth TFIM (optional): query from best to worst until success or budget + success_idx = None + if args.do_ground_truth: + H = tfim_hamiltonian_matrix(args.qubits, periodic=args.periodic) + E0 = exact_ground_energy(H) + + print(f"\nTFIM exact ground energy E0 = {E0:.12f} (periodic={args.periodic})") + + gt_seed_stream = random.Random(args.seed + 1234567) + queried = 0 + for i, circ in enumerate(tqdm(final_circuits, desc="ground-truth queries")): + if queried >= args.gt_budget: + break + seed = gt_seed_stream.randint(0, 1_000_000_000) + ground_truth_tfim( + circ, + H, + E0, + seed=seed, + lr=args.gt_lr, + max_steps=args.gt_max_steps, + tol=args.gt_tol, + ) + queried += 1 + if circ.gt_success: + success_idx = i + break + + if success_idx is None: + print(f"\nNo circuit hit chemical accuracy within budget={args.gt_budget}.") + else: + print(f"\nFirst success at rank {success_idx} (0-based) within budget.") + print(final_circuits[success_idx]) + + # Print bottom few (worst) and top few (best) for sanity + print("\nTop circuits (best expressibility):") + for i, c in enumerate(final_circuits[: min(args.topk, 10)]): + print(f"\nrank {i}:\n{c}") + + # Dump top-K (or all if you want) + out = final_circuits[: args.topk] + with open(args.dump, "w", encoding="utf-8") as fp: + json.dump([c.to_json() for c in out], fp, indent=2) + + print(f"\nWrote {len(out)} circuits to {args.dump}.") + + +if __name__ == "__main__": + main()