#!/usr/bin/env python """Contest-style multi-node Vidal/MPS expectation runner.""" from __future__ import annotations import argparse import math import sys import time from dataclasses import dataclass from pathlib import Path import numpy as np from mpi4py import MPI from qibo import Circuit, gates, hamiltonians from qibo.symbols import X, Y, Z ROOT = Path(__file__).resolve().parents[1] SRC = ROOT / "src" if str(SRC) not in sys.path: sys.path.insert(0, str(SRC)) from qibotn.backends.vidal import VidalBackend # noqa: E402 from qibotn.expectation_runner import exact_for_observable # noqa: E402 @dataclass(frozen=True) class CaseSpec: circuit_kind: str observables: tuple[str, ...] nqubits: int nlayers: int bond: int | None seed: int CASES = { "main1": CaseSpec( circuit_kind="reversed_cnot", observables=("ring_xz",), nqubits=128, nlayers=24, bond=512, seed=31001, ), "main2": CaseSpec( circuit_kind="rxx_rzz", observables=("open_zz", "range2_xx", "mixed_local"), nqubits=128, nlayers=32, bond=1024, seed=31002, ), "strong": CaseSpec( circuit_kind="scramble", observables=("ring_xz", "long_z_string", "dense3_spread"), nqubits=256, nlayers=48, bond=2048, seed=41001, ), } 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 add_single_qubit_layer(circuit, nqubits, rng, include_rx=False): for qubit in range(nqubits): circuit.add(gates.RY(qubit, theta=rng.uniform(-math.pi, math.pi))) circuit.add(gates.RZ(qubit, theta=rng.uniform(-math.pi, math.pi))) if include_rx: circuit.add(gates.RX(qubit, theta=rng.uniform(-math.pi, math.pi))) def build_circuit(kind, nqubits, nlayers, seed): rng = np.random.default_rng(seed) circuit = Circuit(nqubits) for layer in range(nlayers): if kind == "reversed_cnot": add_single_qubit_layer(circuit, nqubits, rng) for qubit in range(0, nqubits - 1, 2): gate = gates.CNOT(qubit + 1, qubit) if layer % 2 else gates.CNOT(qubit, qubit + 1) circuit.add(gate) for qubit in range(1, nqubits - 1, 2): gate = gates.CNOT(qubit + 1, qubit) if layer % 2 == 0 else gates.CNOT(qubit, qubit + 1) circuit.add(gate) elif kind == "rxx_rzz": add_single_qubit_layer(circuit, nqubits, rng, include_rx=True) for qubit in range(layer % 2, nqubits - 1, 2): circuit.add(gates.RXX(qubit, qubit + 1, theta=rng.uniform(-0.9, 0.9))) circuit.add(gates.RZZ(qubit, qubit + 1, theta=rng.uniform(-0.9, 0.9))) elif kind == "scramble": add_single_qubit_layer(circuit, nqubits, rng, include_rx=True) for qubit in range(layer % 2, nqubits - 1, 2): circuit.add(gates.RXX(qubit, qubit + 1, theta=rng.uniform(-0.8, 0.8))) circuit.add(gates.RZZ(qubit, qubit + 1, theta=rng.uniform(-0.8, 0.8))) if layer % 5 == 4: circuit.add(gates.SWAP(qubit, qubit + 1)) else: raise ValueError(f"Unknown circuit kind {kind!r}.") return circuit def dense_observable(nqubits, qubits, seed, dim): del nqubits 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 observable(kind, nqubits, seed): q1 = nqubits // 4 q2 = nqubits // 2 q3 = (3 * nqubits) // 4 last = nqubits - 1 if kind == "boundary_ZZ_q1": return hamiltonians.SymbolicHamiltonian(form=Z(q1 - 1) * Z(q1)) if kind == "boundary_ZZ_q2": return hamiltonians.SymbolicHamiltonian(form=Z(q2 - 1) * Z(q2)) if kind == "boundary_ZZ_q3": return hamiltonians.SymbolicHamiltonian(form=Z(q3 - 1) * Z(q3)) if kind == "long_Z_5_sites": return hamiltonians.SymbolicHamiltonian(form=Z(0) * Z(q1) * Z(q2) * Z(q3) * Z(last)) if kind == "mixed_XZYZX": return hamiltonians.SymbolicHamiltonian(form=X(0) * Z(q1) * Y(q2) * Z(q3) * X(last)) if kind == "ring_xz": form = 0 for qubit in range(nqubits): form += 0.5 * X(qubit) * Z((qubit + 1) % nqubits) return hamiltonians.SymbolicHamiltonian(form=form) if kind == "open_zz": form = 0 for qubit in range(nqubits - 1): form += (1.0 / max(1, nqubits - 1)) * Z(qubit) * Z(qubit + 1) return hamiltonians.SymbolicHamiltonian(form=form) if kind == "range2_xx": form = 0 for qubit in range(nqubits - 2): form += (1.0 / max(1, nqubits - 2)) * X(qubit) * X(qubit + 2) return hamiltonians.SymbolicHamiltonian(form=form) if kind == "mixed_local": form = 0.25 * X(0) - 0.5 * Z(last) + 0.125 * X(q1) * Z(q2) * Y(q3) return hamiltonians.SymbolicHamiltonian(form=form) if kind == "complex_iZ0": return hamiltonians.SymbolicHamiltonian(form=1.0j * Z(0)) if kind == "dense2_mid": return dense_observable(nqubits, (q2 - 1, q2), seed + 101, 4) if kind == "dense3_spread": return dense_observable(nqubits, (q1, q2, q3), seed + 202, 8) raise ValueError(f"Unknown observable kind {kind!r}.") def selected_observables(args, case): if args.observables: return tuple(args.observables) if args.obs_filter: return tuple(x.strip() for x in args.obs_filter.split(",") if x.strip()) return case.observables def apply_case_defaults(args): case = CASES[args.case] if args.nqubits is None: args.nqubits = case.nqubits if args.nlayers is None: args.nlayers = case.nlayers if args.bond == "case-default": args.bond = case.bond if args.seed is None: args.seed = case.seed args.observables = selected_observables(args, case) def run_case(args): set_torch_threads(args.torch_threads) comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() case = CASES[args.case] circuit = build_circuit(case.circuit_kind, args.nqubits, args.nlayers, args.seed) if rank == 0: print("=" * 88, flush=True) print( "backend=vidal_mps " f"case={args.case} circuit={case.circuit_kind} ranks={size} " f"nqubits={args.nqubits} nlayers={args.nlayers} gates={len(circuit.queue)} " f"bond={format_optional(args.bond)} cut_ratio={format_optional(args.cut_ratio)} " f"torch_threads={args.torch_threads} seed={args.seed} " f"observables={','.join(args.observables)}", flush=True, ) print("observable exact value abs_error rel_error seconds trunc_sum trunc_max status", flush=True) for obs_name in args.observables: obs = observable(obs_name, args.nqubits, args.seed) 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, obs, args.nqubits) 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, obs, preprocess=True, compile_circuit=False, ) status = "ok" except Exception as exc: value = np.nan status = type(exc).__name__ + ":" + str(exc).split("\n", 1)[0] seconds = time.perf_counter() - start if rank == 0: abs_error = float("nan") if exact is None else abs(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"{obs_name} {exact_text} {value!r} " f"{abs_error:.6e} {rel_error:.6e} {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("mode", choices=("run", "validate", "list")) parser.add_argument("--case", choices=sorted(CASES), default="main1") parser.add_argument("--observables", nargs="+") parser.add_argument("--obs-filter", default="") parser.add_argument("--nqubits", type=int) parser.add_argument("--nlayers", type=int) parser.add_argument("--bond", "--bonds", dest="bond", default="case-default") parser.add_argument("--cut-ratio", type=optional_float, default=1e-12) parser.add_argument("--seed", type=int) parser.add_argument("--torch-threads", type=int, default=8) parser.add_argument("--exact", action="store_true") parser.add_argument("--exact-max-qubits", type=int, default=24) args = parser.parse_args() if args.mode == "list": for name, case in CASES.items(): print( f"{name}: circuit={case.circuit_kind} " f"observables={','.join(case.observables)} " f"nqubits={case.nqubits} nlayers={case.nlayers} " f"bond={case.bond} seed={case.seed}" ) return apply_case_defaults(args) if isinstance(args.bond, str): args.bond = optional_int(args.bond) if args.mode == "validate": args.exact = True args.nqubits = min(args.nqubits, args.exact_max_qubits) run_case(args) if __name__ == "__main__": main()