qibotn/tools/baseline_mps_expectation.py

"""MPS expectation benchmark for qmatchatea and Vidal backends."""

import argparse
import json
import logging
import os
import socket
import time

import numpy as np

from qibotn.benchmark_cases import (
    build_circuit as build_benchmark_circuit,
    exact_pauli_sum,
    observable_terms,
    terms_to_dict,
)
from qibotn.backends.qmatchatea import QMatchaTeaBackend
from qibotn.backends.vidal_tebd import run_vidal_ring_xz


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 build_circuit(nqubits, nlayers, seed):
    return build_benchmark_circuit("brickwall_cnot", nqubits, nlayers, seed)


def build_observable(nqubits):
    return terms_to_dict(observable_terms("ring_xz", nqubits))


def exact_expectation(circuit, nqubits):
    return exact_pauli_sum(circuit, observable_terms("ring_xz", nqubits), nqubits)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--nqubits", type=int, default=40)
    parser.add_argument("--nlayers", type=int, default=30)
    parser.add_argument("--bond", "--bonds", dest="bond", type=optional_int, default=512)
    parser.add_argument("--cut-ratio", type=optional_float, default=1e-12)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--tensor-module", choices=("numpy", "torch"), default="torch")
    parser.add_argument("--torch-threads", type=int, default=32)
    parser.add_argument(
        "--executor",
        choices=("qmatchatea", "vidal", "vidal-mpi"),
        default="qmatchatea",
    )
    parser.add_argument("--mpi-ct", action="store_true")
    parser.add_argument("--mpi-barriers", type=int, default=-1)
    parser.add_argument("--mpi-isometrization", type=int, default=-1)
    parser.add_argument("--exact", action="store_true")
    parser.add_argument("--exact-max-qubits", type=int, default=24)
    parser.add_argument("--reference-file")
    parser.add_argument(
        "--mpi-rank-map",
        action="store_true",
        help="Print MPI rank, host, pid, and torch thread placement metadata.",
    )
    args = parser.parse_args()
    logging.getLogger("qibo.config").setLevel(logging.ERROR)
    logging.getLogger("qtealeaves").setLevel(logging.ERROR)
    import torch

    torch.set_num_threads(args.torch_threads)
    rank = 0
    size = 1
    if args.mpi_ct:
        from mpi4py import MPI

        rank = MPI.COMM_WORLD.Get_rank()
        size = MPI.COMM_WORLD.Get_size()
        if args.mpi_rank_map:
            rank_info = {
                "rank": rank,
                "size": size,
                "host": socket.gethostname(),
                "pid": os.getpid(),
                "torch_threads": args.torch_threads,
                "omp_num_threads": os.environ.get("OMP_NUM_THREADS", ""),
                "mkl_num_threads": os.environ.get("MKL_NUM_THREADS", ""),
            }
            rank_infos = MPI.COMM_WORLD.gather(rank_info, root=0)
            if rank == 0:
                print("mpi_rank_map")
                for item in sorted(rank_infos, key=lambda row: row["rank"]):
                    print(
                        "rank={rank} size={size} host={host} pid={pid} "
                        "torch_threads={torch_threads} "
                        "OMP_NUM_THREADS={omp_num_threads} "
                        "MKL_NUM_THREADS={mkl_num_threads}".format(**item)
                    )

    circuit = build_circuit(args.nqubits, args.nlayers, args.seed)
    observable = build_observable(args.nqubits)
    exact = None
    if args.reference_file:
        with open(args.reference_file, "r", encoding="utf-8") as f:
            exact = float(json.load(f)["expectation"])
    elif args.exact:
        if args.nqubits > args.exact_max_qubits:
            raise ValueError(
                f"--exact is limited to {args.exact_max_qubits} qubits by default."
            )
        exact = exact_expectation(circuit, args.nqubits)

    if rank == 0:
        if args.mpi_ct and args.executor in ("vidal", "vidal-mpi"):
            mpi_label = f"VidalSegment/{size}"
        else:
            mpi_label = f"MPIMPS/{size}" if args.mpi_ct else "SR"
        print(
            f"nqubits={args.nqubits} nlayers={args.nlayers} "
            f"bond={format_optional(args.bond)} "
            f"cut_ratio={format_optional(args.cut_ratio)} seed={args.seed} "
            f"tensor_module={args.tensor_module} svd_control=E! "
            f"compile_circuit=True mpi={mpi_label} executor={args.executor}"
        )
        if exact is not None:
            print(f"exact={exact:.16e}")
        print("expval abs_error rel_error seconds")

    start = time.perf_counter()
    timings = None
    if args.executor in ("vidal", "vidal-mpi"):
        if args.executor == "vidal-mpi" and not args.mpi_ct:
            raise ValueError("--executor vidal-mpi requires --mpi-ct.")
        if args.mpi_ct:
            from qibotn.backends.vidal_mpi_segment import run_segment_vidal_mpi_ring_xz

            value, timings = run_segment_vidal_mpi_ring_xz(
                circuit,
                max_bond=args.bond,
                cut_ratio=args.cut_ratio,
                tensor_module=args.tensor_module,
                comm=MPI.COMM_WORLD,
            )
        else:
            value = run_vidal_ring_xz(
                circuit,
                max_bond=args.bond,
                cut_ratio=args.cut_ratio,
                tensor_module=args.tensor_module,
            )
    else:
        backend = QMatchaTeaBackend()
        backend.configure_tn_simulation(
            ansatz="MPS",
            max_bond_dimension=args.bond,
            cut_ratio=args.cut_ratio,
            svd_control="E!",
            tensor_module=args.tensor_module,
            compile_circuit=True,
            track_memory=False,
            mpi_approach="CT" if args.mpi_ct else "SR",
            mpi_num_procs=size,
            mpi_where_barriers=args.mpi_barriers if args.mpi_ct else -1,
            mpi_isometrization=args.mpi_isometrization,
        )
        value = backend.expectation(
            circuit,
            observable,
            preprocess=False,
            compile_circuit=True,
        )
    max_timings = None
    if timings:
        max_timings = {
            key: MPI.COMM_WORLD.reduce(local_value, op=MPI.MAX, root=0)
            for key, local_value in timings.items()
        }
    if rank != 0:
        return
    value = float(np.real(value))
    elapsed = time.perf_counter() - start
    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)
    print(f"{value:.16e} {abs_error:.6e} {rel_error:.6e} {elapsed:.3f}")
    if max_timings:
        print("timing_section max_seconds")
        for key, max_value in max_timings.items():
            print(f"{key} {max_value:.6f}")


if __name__ == "__main__":
    main()