赛前稳定版

tools/run_tn_custom.py

@@ -0,0 +1,243 @@
+#!/usr/bin/env python
+"""Run TN expectation for a user-provided circuit and observable.
+
+The case module should define:
+
+    def build_circuit(nqubits, nlayers, seed): ...
+    def build_observable(nqubits, seed): ...
+
+``build_observable`` may return a Qibo SymbolicHamiltonian/form or the qibotn
+dict form:
+
+    {"terms": [
+        {"coefficient": 1.0, "operators": [("X", 0), ("Z", 1)]},
+    ]}
+
+For a single repeated Pauli string, pass ``--pauli-pattern`` instead of
+defining ``build_observable``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import inspect
+import json
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+SRC = ROOT / "src"
+if str(SRC) not in sys.path:
+    sys.path.insert(0, str(SRC))
+
+from qibotn.expectation_runner import (  # noqa: E402
+    ExpectationConfig,
+    exact_for_observable,
+    run_cpu_expectation,
+)
+
+
+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 load_module(path):
+    path = Path(path).resolve()
+    spec = importlib.util.spec_from_file_location(path.stem, path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Cannot import case module from {path}.")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def call_builder(fn, **kwargs):
+    sig = inspect.signature(fn)
+    if any(p.kind == p.VAR_KEYWORD for p in sig.parameters.values()):
+        return fn(**kwargs)
+    accepted = {
+        name: value
+        for name, value in kwargs.items()
+        if name in sig.parameters
+    }
+    return fn(**accepted)
+
+
+def load_observable(args, module):
+    if args.pauli_pattern:
+        return {"pauli_string_pattern": args.pauli_pattern}
+    if args.observable_json:
+        with Path(args.observable_json).open() as f:
+            return json.load(f)
+    if hasattr(module, "build_observable"):
+        return call_builder(
+            module.build_observable,
+            nqubits=args.nqubits,
+            nlayers=args.nlayers,
+            seed=args.seed,
+        )
+    if hasattr(module, "OBSERVABLE"):
+        return module.OBSERVABLE
+    raise ValueError(
+        "No observable supplied. Define build_observable/OBSERVABLE in the case "
+        "module, or pass --pauli-pattern / --observable-json."
+    )
+
+
+def build_parallel_opts(args):
+    slicing_opts = {}
+    if args.tn_target_slices is not None:
+        slicing_opts["target_slices"] = args.tn_target_slices
+    if args.tn_target_size is not None:
+        slicing_opts["target_size"] = args.tn_target_size
+
+    opts = {
+        "slicing_opts": slicing_opts or None,
+        "search_workers": args.tn_search_workers or args.torch_threads,
+        "max_repeats": args.tn_search_repeats,
+        "max_time": args.tn_search_time,
+        "print_stats": not args.no_tn_stats,
+    }
+    if args.tn_search_backend is not None:
+        opts["search_backend"] = args.tn_search_backend
+    if args.dask_address is not None:
+        opts["dask_address"] = args.dask_address
+    if args.dask_close_workers:
+        opts["dask_close_workers"] = True
+    if args.tn_save_tree is not None:
+        opts["save_tree_path"] = args.tn_save_tree
+    if args.tn_load_tree is not None:
+        opts["load_tree_path"] = args.tn_load_tree
+    if args.tn_search_only:
+        opts["search_only"] = True
+    return opts
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Run CPU TN expectation for a custom qibo circuit module."
+    )
+    parser.add_argument("case_module", help="Python file defining build_circuit.")
+    parser.add_argument("--nqubits", type=int, required=True)
+    parser.add_argument("--nlayers", type=int, default=0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--mpi", action="store_true")
+    parser.add_argument("--exact", action="store_true")
+    parser.add_argument("--exact-max-qubits", type=int, default=24)
+    parser.add_argument("--bond", "--bonds", dest="bond", type=optional_int, default=1024)
+    parser.add_argument("--cut-ratio", type=optional_float, default=1e-12)
+    parser.add_argument("--torch-threads", type=int, default=8)
+    parser.add_argument("--quimb-backend", choices=("numpy", "torch"), default="torch")
+    parser.add_argument("--dtype", choices=("complex128", "complex64"), default="complex128")
+    parser.add_argument("--pauli-pattern")
+    parser.add_argument("--observable-json")
+    parser.add_argument("--tn-target-slices", type=int)
+    parser.add_argument("--tn-target-size", type=int, default=2**32)
+    parser.add_argument("--tn-search-workers", type=int)
+    parser.add_argument("--tn-search-repeats", type=int, default=128)
+    parser.add_argument("--tn-search-time", type=float, default=60.0)
+    parser.add_argument("--tn-search-backend", choices=("processpool", "dask"))
+    parser.add_argument("--dask-address")
+    parser.add_argument("--dask-close-workers", action="store_true")
+    parser.add_argument("--tn-save-tree")
+    parser.add_argument("--tn-load-tree")
+    parser.add_argument("--tn-search-only", action="store_true")
+    parser.add_argument("--no-tn-stats", action="store_true")
+    args = parser.parse_args()
+
+    rank = 0
+    if args.mpi:
+        from mpi4py import MPI
+
+        rank = MPI.COMM_WORLD.Get_rank()
+
+    module = load_module(args.case_module)
+    if not hasattr(module, "build_circuit"):
+        raise ValueError("case_module must define build_circuit.")
+
+    circuit = call_builder(
+        module.build_circuit,
+        nqubits=args.nqubits,
+        nlayers=args.nlayers,
+        seed=args.seed,
+    )
+    observable = load_observable(args, module)
+
+    config = ExpectationConfig(
+        ansatz="tn",
+        mpi=args.mpi,
+        bond=args.bond,
+        cut_ratio=args.cut_ratio,
+        tensor_module="torch",
+        quimb_backend=args.quimb_backend,
+        dtype=args.dtype,
+        torch_threads=args.torch_threads,
+        parallel_opts=build_parallel_opts(args),
+    )
+
+    if rank == 0:
+        mode = "MPI" if args.mpi else "serial"
+        print(
+            f"backend=cpu ansatz=TN mode={mode} case={Path(args.case_module).name} "
+            f"nqubits={args.nqubits} nlayers={args.nlayers} seed={args.seed} "
+            f"quimb_backend={args.quimb_backend} dtype={args.dtype} "
+            f"torch_threads={args.torch_threads}",
+            flush=True,
+        )
+        print("observable exact value abs_error rel_error seconds", flush=True)
+
+    exact = None
+    if args.exact and rank == 0:
+        if args.nqubits > args.exact_max_qubits:
+            raise ValueError(
+                f"--exact is limited to {args.exact_max_qubits} qubits by default."
+            )
+        exact = exact_for_observable(circuit, observable, args.nqubits)
+
+    result = run_cpu_expectation(circuit, observable, config)
+    if args.mpi and result.rank != 0:
+        return
+
+    abs_error = float("nan") if exact is None else abs(result.value - exact)
+    rel_error = float("nan") if exact is None else abs_error / max(abs(exact), 1e-15)
+    exact_text = "nan" if exact is None else f"{exact:.16e}"
+    print(
+        f"custom {exact_text} {result.value:.16e} "
+        f"{abs_error:.6e} {rel_error:.6e} {result.seconds:.3f}",
+        flush=True,
+    )
+
+    for stat in result.parallel_stats or ():
+        cost = stat["path_cost"]
+        search_stats = stat.get("search_stats", {})
+        print(
+            "tn_term_summary "
+            f"term={stat.get('term_index', 0)} "
+            f"search_seconds={stat.get('search_seconds', float('nan')):.3f} "
+            f"contract_seconds={stat.get('contract_seconds', float('nan')):.3f} "
+            f"completed_trials={search_stats.get('completed_trials', 'na')} "
+            f"finite_trials={search_stats.get('finite_trials', 'na')} "
+            f"failed_trials={search_stats.get('failed_trials', 'na')} "
+            f"requested_trials={search_stats.get('requested_trials', 'na')} "
+            f"best_score={search_stats.get('best_score', float('nan')):.6g} "
+            f"slices={cost.get('slices')} "
+            f"log10_flops={cost.get('log10_flops', float('nan')):.3f} "
+            f"log10_write={cost.get('log10_write', float('nan')):.3f} "
+            f"log2_size={cost.get('log2_size', float('nan')):.3f} "
+            f"peak_memory_gib={cost.get('peak_memory_gib', float('nan')):.3g} "
+            f"rank_slices={stat.get('rank_slices')}",
+            flush=True,
+        )
+
+
+if __name__ == "__main__":
+    main()