pcb.qaoa

  1"""Implementation of the QAOA algorithm."""
  2
  3import numpy as np
  4from loguru import logger as logging
  5from qiskit import QuantumCircuit
  6from qiskit.circuit.library import QAOAAnsatz
  7from qiskit.primitives import PrimitiveJob
  8from qiskit.providers import BackendV2 as Backend
  9from qiskit.quantum_info import SparsePauliOp
 10from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
 11from qiskit_aer.primitives import EstimatorV2 as AerEstimator
 12from qiskit_ibm_runtime import EstimatorV2 as IBMEstimator
 13from qiskit_ibm_runtime import RuntimeJobV2 as IBMRuntimeJob
 14from scipy.optimize import minimize
 15
 16
 17def _energy(
 18    x: np.ndarray,
 19    ansatz: QuantumCircuit,
 20    operator: SparsePauliOp,
 21    estimator: IBMEstimator | AerEstimator,
 22    records: list[tuple[np.ndarray, float, dict]] | None = None,
 23) -> float:
 24    """
 25    Evaluates a set of parameters for the QAOA ansatz. Returns an array of all
 26    resulting energies and a dictionary containing data about the job submitted
 27    to IBMQ or Aer. See `_job_to_dict` for the structure of the dictionary.
 28
 29    Args:
 30        parameters (np.ndarray): A `(N, n_qaoa_params)` array of real parameters
 31        ansatz (QuantumCircuit):
 32        operator (SparsePauliOp):
 33        estimator (IBMEstimator | AerEstimator):
 34        records (list, optional): If provided, results of this evaluation will
 35            be appended.
 36    """
 37    if x.ndim != 1:
 38        raise ValueError(
 39            f"Parameters must be a 1D array, got shape: {x.shape}"
 40        )
 41    pubs = [(ansatz, [operator], [x])]
 42    job = estimator.run(pubs=pubs)
 43    results = job.result()
 44    e = float(results[0].data.evs[0])
 45    if records is not None:
 46        records.append((x, e, _job_to_dict(job)))
 47    return e
 48
 49
 50def _job_to_dict(job: IBMRuntimeJob | PrimitiveJob) -> dict:
 51    """
 52    Note:
 53        The keys will be prefixed with `ibmq_` even if the job is an Aer job.
 54    """
 55    return {
 56        "ibmq_jid": job.job_id(),
 57        "ibmq_sid": getattr(job, "session_id", None),
 58    }
 59
 60
 61def qaoa(
 62    operator: SparsePauliOp,
 63    backend: Backend,
 64    estimator: IBMEstimator | AerEstimator,
 65    cost_qc: QuantumCircuit | None = None,
 66    n_qaoa_steps: int = 2,
 67    pm_optimization_level: int = 3,
 68    max_iter: int = 128,
 69) -> tuple[
 70    tuple[np.ndarray, float], tuple[np.ndarray, np.ndarray], list[dict]
 71]:
 72    """
 73    Runs QAOA on the given operator using the given backend. The Trotterization
 74    of the cost operator can be given explicitely as a `QuantumCircuit`.
 75    Otherwise, it is automatically Trotterized.
 76
 77    This method submits PUBs in batches, meaning that at every iteration,
 78    `batch_size` parameters are tested.
 79
 80    Warning:
 81        This method tries to *MAXIMIZE* the energy of the operator. If you want
 82        to minimize, flip the sign of the weights in `operator`.
 83
 84    Args:
 85        operator (SparsePauliOp): The operator whose energy to maximize
 86        backend (Backend):
 87        estimator (IBMEstimator | AerEstimator):
 88        cost_qc (QuantumCircuit | None, optional):
 89        n_qaoa_steps (int, optional):
 90        pm_optimization_level (int, optional):
 91
 92    Returns:
 93        1. A tuple containing the optimal parameters and the optimal energy.
 94        2. A tuple containing array of all parameters that were tried and the
 95           array of all resulting energies.
 96        3. A list of dict decribing the results of each job. An element looks
 97           like:
 98
 99            {
100                "energy": 0.006450488764709524,
101                "ibmq_job_id": "3c94c73c-385b-4e71-8675-e3e1ad76aa4e",
102                "ibmq_session_id": "cz4h80039f40008scarg",
103                "batch": 3,
104            }
105    """
106    ansatz = QAOAAnsatz(
107        cost_operator=cost_qc if cost_qc is not None else operator,
108        reps=n_qaoa_steps,
109    )
110    pm = generate_preset_pass_manager(
111        target=backend.target, optimization_level=pm_optimization_level
112    )
113    ansatz_isa = pm.run(ansatz)
114    operator_isa = operator.apply_layout(ansatz_isa.layout)
115    logging.debug(
116        "ISA ansatz depth/size: {}/{}",
117        ansatz_isa.depth(),
118        ansatz_isa.size(),
119    )
120    x0 = np.array(
121        ([np.pi / 2] * (len(ansatz_isa.parameters) // 2))  # β's
122        + ([np.pi] * (len(ansatz_isa.parameters) // 2))  # γ's
123    )
124    bounds = (
125        ([(0, np.pi)] * (len(ansatz_isa.parameters) // 2))  # β's
126        + ([(0, 2 * np.pi)] * (len(ansatz_isa.parameters) // 2))  # γ's
127    )
128    records: list[tuple[np.ndarray, float, dict]] = []
129    minimize(
130        lambda *args: -1 * _energy(*args),  # energy maximization
131        x0=x0,
132        args=(ansatz_isa, operator_isa, estimator, records),
133        method="cobyla",
134        bounds=bounds,
135        options={"maxiter": max_iter},
136    )
137    all_x = np.stack([x for x, _, _ in records])
138    all_e = np.array([e for _, e, _ in records])
139    results = [
140        {"energy": float(e), "step": i, **m}
141        for i, (_, e, m) in enumerate(records)
142    ]
143    j = all_e.argmax()
144    return (all_x[j], all_e[j]), (all_x, all_e), results