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()