赛前稳定版

tools/vidal_mpi_contest_runner.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import argparse
+import math
+import time
+
+import numpy as np
+from mpi4py import MPI
+from qibo import Circuit, gates, hamiltonians
+from qibo.symbols import X, Y, Z
+
+from qibotn.backends.vidal import VidalBackend
+
+
+def optional_int(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return int(text)
+
+
+def optional_float(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return float(text)
+
+
+def format_optional(value, fmt="g"):
+    return "None" if value is None else format(value, fmt)
+
+
+def set_torch_threads(nthreads):
+    try:
+        import torch
+
+        torch.set_num_threads(nthreads)
+    except Exception:
+        pass
+
+
+def build_circuit(kind, nqubits, nlayers, seed):
+    rng = np.random.default_rng(seed)
+    circuit = Circuit(nqubits)
+
+    for layer in range(nlayers):
+        for q in range(nqubits):
+            circuit.add(gates.RY(q, theta=rng.uniform(-math.pi, math.pi)))
+            circuit.add(gates.RZ(q, theta=rng.uniform(-math.pi, math.pi)))
+            if kind in ("rxx_rzz", "scramble"):
+                circuit.add(gates.RX(q, theta=rng.uniform(-math.pi, math.pi)))
+
+        if kind == "reversed_cnot":
+            for q in range(0, nqubits - 1, 2):
+                circuit.add(gates.CNOT(q + 1, q) if layer % 2 else gates.CNOT(q, q + 1))
+            for q in range(1, nqubits - 1, 2):
+                circuit.add(gates.CNOT(q + 1, q) if layer % 2 == 0 else gates.CNOT(q, q + 1))
+        elif kind == "rxx_rzz":
+            for q in range(layer % 2, nqubits - 1, 2):
+                circuit.add(gates.RXX(q, q + 1, theta=rng.uniform(-0.9, 0.9)))
+                circuit.add(gates.RZZ(q, q + 1, theta=rng.uniform(-0.9, 0.9)))
+        elif kind == "scramble":
+            for q in range(layer % 2, nqubits - 1, 2):
+                circuit.add(gates.RXX(q, q + 1, theta=rng.uniform(-0.8, 0.8)))
+                circuit.add(gates.RZZ(q, q + 1, theta=rng.uniform(-0.8, 0.8)))
+                if layer % 5 == 4:
+                    circuit.add(gates.SWAP(q, q + 1))
+        else:
+            raise ValueError(f"Unknown circuit kind {kind!r}.")
+
+    return circuit
+
+
+def ring_xz(nqubits):
+    form = 0
+    for q in range(nqubits):
+        form += 0.5 * X(q) * Z((q + 1) % nqubits)
+    return hamiltonians.SymbolicHamiltonian(form=form)
+
+
+def open_zz(nqubits):
+    form = 0
+    for q in range(nqubits - 1):
+        form += (1.0 / (nqubits - 1)) * Z(q) * Z(q + 1)
+    return hamiltonians.SymbolicHamiltonian(form=form)
+
+
+def range2_xx(nqubits):
+    form = 0
+    for q in range(nqubits - 2):
+        form += (1.0 / (nqubits - 2)) * X(q) * X(q + 2)
+    return hamiltonians.SymbolicHamiltonian(form=form)
+
+
+def dense_observable(nqubits, qubits, seed, dim):
+    rng = np.random.default_rng(seed)
+    raw = rng.normal(size=(dim, dim)) + 1j * rng.normal(size=(dim, dim))
+    matrix = (raw + raw.conj().T) / 2.0
+    matrix = matrix / np.linalg.norm(matrix)
+    return {"matrix": matrix, "qubits": list(qubits)}
+
+
+def observables_for_case(nqubits, seed):
+    q1 = nqubits // 4
+    q2 = nqubits // 2
+    q3 = (3 * nqubits) // 4
+    last = nqubits - 1
+
+    return [
+        ("boundary_ZZ_q1", hamiltonians.SymbolicHamiltonian(form=Z(q1 - 1) * Z(q1))),
+        ("boundary_ZZ_q2", hamiltonians.SymbolicHamiltonian(form=Z(q2 - 1) * Z(q2))),
+        ("boundary_ZZ_q3", hamiltonians.SymbolicHamiltonian(form=Z(q3 - 1) * Z(q3))),
+        (
+            "long_Z_5_sites",
+            hamiltonians.SymbolicHamiltonian(form=Z(0) * Z(q1) * Z(q2) * Z(q3) * Z(last)),
+        ),
+        (
+            "mixed_XZYZX",
+            hamiltonians.SymbolicHamiltonian(form=X(0) * Z(q1) * Y(q2) * Z(q3) * X(last)),
+        ),
+        ("ring_xz", ring_xz(nqubits)),
+        ("open_zz", open_zz(nqubits)),
+        ("range2_xx", range2_xx(nqubits)),
+        ("complex_iZ0", hamiltonians.SymbolicHamiltonian(form=1.0j * Z(0))),
+        ("dense2_mid", dense_observable(nqubits, (q2 - 1, q2), seed + 101, 4)),
+        ("dense3_spread", dense_observable(nqubits, (q1, q2, q3), seed + 202, 8)),
+    ]
+
+
+def run_case(args):
+    set_torch_threads(args.torch_threads)
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    size = comm.Get_size()
+
+    circuit = build_circuit(args.kind, args.nqubits, args.nlayers, args.seed)
+    observables = observables_for_case(args.nqubits, args.seed)
+    if args.obs_filter:
+        wanted = set(args.obs_filter.split(","))
+        observables = [(name, obs) for name, obs in observables if name in wanted]
+        if not observables:
+            raise ValueError(f"OBS_FILTER matched no observables: {args.obs_filter!r}")
+
+    if rank == 0:
+        print("=" * 88, flush=True)
+        print(
+            "case "
+            f"label={args.label} kind={args.kind} ranks={size} "
+            f"nqubits={args.nqubits} nlayers={args.nlayers} gates={len(circuit.queue)} "
+            f"bond={format_optional(args.bond)} "
+            f"cut_ratio={format_optional(args.cut_ratio)} "
+            f"torch_threads={args.torch_threads} seed={args.seed} "
+            f"obs_filter={args.obs_filter or 'all'}",
+            flush=True,
+        )
+        print(
+            "observable value seconds trunc_sum trunc_max status",
+            flush=True,
+        )
+
+    for obs_name, observable in observables:
+        backend = VidalBackend()
+        backend.configure_tn_simulation(
+            max_bond_dimension=args.bond,
+            cut_ratio=args.cut_ratio,
+            tensor_module="torch",
+            mpi_approach="CT",
+            mpi_num_procs=size,
+            fallback=False,
+        )
+
+        comm.Barrier()
+        start = time.perf_counter()
+        try:
+            value = backend.expectation(
+                circuit,
+                observable,
+                preprocess=True,
+                compile_circuit=False,
+            )
+            status = "ok"
+        except Exception as exc:  # pragma: no cover - printed for manual runs
+            value = np.nan
+            status = type(exc).__name__ + ":" + str(exc).split("\n", 1)[0]
+        seconds = time.perf_counter() - start
+
+        if rank == 0:
+            print(
+                f"{obs_name} {value!r} {seconds:.3f} "
+                f"{backend.last_truncation_error:.6e} "
+                f"{backend.last_max_truncation_error:.6e} {status}",
+                flush=True,
+            )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--label", required=True)
+    parser.add_argument("--kind", choices=("reversed_cnot", "rxx_rzz", "scramble"), required=True)
+    parser.add_argument("--nqubits", type=int, required=True)
+    parser.add_argument("--nlayers", type=int, required=True)
+    parser.add_argument("--bond", type=optional_int, required=True)
+    parser.add_argument("--cut-ratio", type=optional_float, required=True)
+    parser.add_argument("--seed", type=int, required=True)
+    parser.add_argument("--torch-threads", type=int, required=True)
+    parser.add_argument("--obs-filter", default="")
+    run_case(parser.parse_args())
+
+
+if __name__ == "__main__":
+    main()