benchmark测试，发现瓶颈：路径搜索

benchmark_mps.py

@@ -1,10 +1,13 @@
 """Benchmark: qibojit (reference) vs qibotn/quimb MPS, with error comparison."""
 import time
 import argparse
+import os
 import numpy as np
 import qibo
 from qibo import Circuit, gates
 
+DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
+
 
 def make_circuit(circuit_type, nqubits, nlayers=1):
     c = Circuit(nqubits)
@@ -12,10 +15,11 @@ def make_circuit(circuit_type, nqubits, nlayers=1):
         from qibo.models import QFT
         return QFT(nqubits)
     elif circuit_type == "variational":
-        for _ in range(nlayers):
+        for layer in range(nlayers):
             for q in range(nqubits):
                 c.add(gates.RY(q, theta=np.random.uniform(0, 2 * np.pi)))
-            for q in range(0, nqubits - 1, 2):
+            offset = layer % 2
+            for q in range(offset, nqubits - 1, 2):
                 c.add(gates.CZ(q, q + 1))
     elif circuit_type == "ghz":
         c.add(gates.H(0))
@@ -49,13 +53,18 @@ def run_quimb_mps(circuit, max_bond, svd_cutoff, optimizer):
 
 
 def compare(sv_ref, sv_mps):
-    sv_ref = np.array(sv_ref).flatten()
-    sv_mps = np.array(sv_mps).flatten()
+    sv_ref = np.array(sv_ref, dtype=complex).flatten()
+    sv_mps = np.array(sv_mps, dtype=complex).flatten()
     fidelity = abs(np.dot(sv_ref.conj(), sv_mps)) ** 2
     l2_err = np.linalg.norm(sv_ref - sv_mps)
     return fidelity, l2_err
 
 
+def jit_cache_path(circuit_type, nqubits, nlayers):
+    os.makedirs(DATA_DIR, exist_ok=True)
+    return os.path.join(DATA_DIR, f"jit_{circuit_type}_n{nqubits}_l{nlayers}.npy")
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--nqubits", type=int, default=10)
@@ -66,16 +75,27 @@ def main():
                         help="Max bond dimension for MPS (None = unlimited)")
     parser.add_argument("--svd-cutoff", type=float, default=1e-6)
     parser.add_argument("--optimizer", type=str, default="auto-hq")
+    parser.add_argument("--skip-jit", action="store_true",
+                        help="Skip qibojit run, load cached statevector if available")
     args = parser.parse_args()
 
-    np.random.seed(42)
     print(f"Circuit: {args.circuit}, nqubits={args.nqubits}, nlayers={args.nlayers}")
     print(f"MPS config: max_bond={args.max_bond}, svd_cutoff={args.svd_cutoff}, optimizer={args.optimizer}")
 
-    circuit_ref = make_circuit(args.circuit, args.nqubits, args.nlayers)
-    sv_ref, t_ref = run_qibojit(circuit_ref)
-    print(f"\n[qibojit]  time={t_ref:.4f}s")
+    cache_path = jit_cache_path(args.circuit, args.nqubits, args.nlayers)
+    t_ref = None
 
+    if args.skip_jit and os.path.exists(cache_path):
+        sv_ref = np.load(cache_path)
+        print(f"\n[qibojit]  loaded from cache: {cache_path}")
+    else:
+        np.random.seed(42)
+        circuit_ref = make_circuit(args.circuit, args.nqubits, args.nlayers)
+        sv_ref, t_ref = run_qibojit(circuit_ref)
+        np.save(cache_path, sv_ref)
+        print(f"\n[qibojit]  time={t_ref:.4f}s  (saved to {cache_path})")
+
+    np.random.seed(42)
     circuit_mps = make_circuit(args.circuit, args.nqubits, args.nlayers)
     try:
         sv_mps, t_mps = run_quimb_mps(circuit_mps, args.max_bond, args.svd_cutoff, args.optimizer)
@@ -83,7 +103,8 @@ def main():
         print(f"[quimb MPS] time={t_mps:.4f}s")
         print(f"\nFidelity : {fidelity:.8f}  (1=perfect)")
         print(f"L2 error : {l2_err:.2e}")
-        print(f"Speedup  : {t_ref/t_mps:.2f}x" if t_mps > 0 else "")
+        if t_ref is not None and t_mps > 0:
+            print(f"Speedup  : {t_ref/t_mps:.2f}x")
     except Exception as e:
         print(f"[quimb MPS] FAILED: {e}")
         raise