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

.gitignore

@@ -2,7 +2,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
-
+data/
 # C extensions
 *.so
 

benchmark_mps.py

@@ -0,0 +1,114 @@
+"""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)
+    if circuit_type == "qft":
+        from qibo.models import QFT
+        return QFT(nqubits)
+    elif circuit_type == "variational":
+        for layer in range(nlayers):
+            for q in range(nqubits):
+                c.add(gates.RY(q, theta=np.random.uniform(0, 2 * np.pi)))
+            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))
+        for q in range(nqubits - 1):
+            c.add(gates.CNOT(q, q + 1))
+    else:
+        raise ValueError(f"Unknown circuit: {circuit_type}")
+    return c
+
+
+def run_qibojit(circuit):
+    qibo.set_backend("qibojit", platform="numba")
+    t0 = time.time()
+    result = circuit()
+    elapsed = time.time() - t0
+    sv = result.state()
+    return sv, elapsed
+
+
+def run_quimb_mps(circuit, max_bond, svd_cutoff, optimizer):
+    qibo.set_backend("qibotn", platform="quimb")
+    b = qibo.get_backend()
+    b.configure_tn_simulation(ansatz="mps", max_bond_dimension=max_bond, svd_cutoff=svd_cutoff)
+    b.contractions_optimizer = optimizer
+
+    t0 = time.time()
+    result = b.execute_circuit(circuit, return_array=True)
+    elapsed = time.time() - t0
+    sv = result.state()
+    return sv, elapsed
+
+
+def compare(sv_ref, sv_mps):
+    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)
+    parser.add_argument("--circuit", type=str, default="ghz",
+                        choices=["qft", "variational", "ghz"])
+    parser.add_argument("--nlayers", type=int, default=3)
+    parser.add_argument("--max-bond", type=int, default=None,
+                        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()
+
+    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}")
+
+    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)
+        fidelity, l2_err = compare(sv_ref, sv_mps)
+        print(f"[quimb MPS] time={t_mps:.4f}s")
+        print(f"\nFidelity : {fidelity:.8f}  (1=perfect)")
+        print(f"L2 error : {l2_err:.2e}")
+        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
+
+
+if __name__ == "__main__":
+    main()

log

@@ -0,0 +1,11 @@
+[qibojit] loaded from cache: /home/yx/qibotn/data/jit_variational_n32_l5.npy
+
+  bond     time(s)      fidelity      l2_err
+----------------------------------------------
+     1    157.4587    0.00000280    9.99e-01
+     8     61.9126    0.99999014    2.22e-03
+    16     63.4902    0.99999014    2.22e-03
+    32     58.3594    0.99999014    2.22e-03
+    64     59.7043    0.99999014    2.22e-03
+   128     64.6368    0.99999014    2.22e-03
+   256     64.9058    0.99999014    2.22e-03

poetry.lock

@@ -1733,14 +1733,14 @@ files = [
 
 [[package]]
 name = "mako"
-version = "1.3.10"
+version = "1.3.11"
 description = "A super-fast templating language that borrows the best ideas from the existing templating languages."
 optional = false
 python-versions = ">=3.8"
 groups = ["main"]
 files = [
-    {file = "mako-1.3.10-py3-none-any.whl", hash = "sha256:baef24a52fc4fc514a0887ac600f9f1cff3d82c61d4d700a1fa84d597b88db59"},
-    {file = "mako-1.3.10.tar.gz", hash = "sha256:99579a6f39583fa7e5630a28c3c1f440e4e97a414b80372649c0ce338da2ea28"},
+    {file = "mako-1.3.11-py3-none-any.whl", hash = "sha256:e372c6e333cf004aa736a15f425087ec977e1fcbd2966aae7f17c8dc1da27a77"},
+    {file = "mako-1.3.11.tar.gz", hash = "sha256:071eb4ab4c5010443152255d77db7faa6ce5916f35226eb02dc34479b6858069"},
 ]
 
 [package.dependencies]

src/qibotn/backends/quimb.py

@@ -167,7 +167,7 @@ def execute_circuit(
         raise_error(ValueError, "Initial state not None supported only for MPS ansatz.")
 
     circ_quimb = self.circuit_ansatz.from_openqasm2_str(
-        circuit.to_qasm(), psi0=initial_state
+        circuit.to_qasm(), psi0=initial_state, gate_opts={"max_bond": self.max_bond_dimension, "cutoff": self.svd_cutoff}
     )
 
     if nshots:

src/qibotn/result.py

@@ -58,7 +58,7 @@ class TensorNetworkResult:
 
     def state(self):
         """Return the statevector if the number of qubits is less than 20."""
-        if self.nqubits < 20:
+        if self.nqubits < 35:
             return self.statevector
         raise_error(
             NotImplementedError,

sweep_bond_32q.py

@@ -0,0 +1,39 @@
+"""Bond dimension sweep for 32-qubit variational circuit."""
+import os
+import sys
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(__file__))
+from benchmark_mps import make_circuit, run_qibojit, run_quimb_mps, compare, jit_cache_path, DATA_DIR
+
+NQUBITS = 32
+NLAYERS = 5
+BOND_VALUES = [1, 8, 16, 32, 64, 128, 256]
+SVD_CUTOFF = 1e-6
+OPTIMIZER = "auto-hq"
+
+if __name__ == "__main__":
+    cache_path = jit_cache_path("variational", NQUBITS, NLAYERS)
+
+    if os.path.exists(cache_path):
+        sv_ref = np.load(cache_path)
+        print(f"[qibojit] loaded from cache: {cache_path}\n")
+    else:
+        np.random.seed(42)
+        circuit_ref = make_circuit("variational", NQUBITS, NLAYERS)
+        sv_ref, t_ref = run_qibojit(circuit_ref)
+        np.save(cache_path, sv_ref)
+        print(f"[qibojit] time={t_ref:.4f}s  (saved to {cache_path})\n")
+
+    print(f"{'bond':>6}  {'time(s)':>10}  {'fidelity':>12}  {'l2_err':>10}")
+    print("-" * 46)
+
+    for bond in BOND_VALUES:
+        np.random.seed(42)
+        circuit_mps = make_circuit("variational", NQUBITS, NLAYERS)
+        try:
+            sv_mps, t_mps = run_quimb_mps(circuit_mps, bond, SVD_CUTOFF, OPTIMIZER)
+            fidelity, l2_err = compare(sv_ref, sv_mps)
+            print(f"{bond:>6}  {t_mps:>10.4f}  {fidelity:>12.8f}  {l2_err:>10.2e}")
+        except Exception as e:
+            print(f"{bond:>6}  FAILED: {e}")