"""This module contains the ``QAOASolver`` class."""
from __future__ import annotations
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, SupportsInt
import numpy as np
import torch
from numpy.random import RandomState
from numpy.typing import ArrayLike, NDArray
from tno.quantum.devices import DeviceConfig
from tno.quantum.optimization.qubo.components import QUBO, Solver
from tno.quantum.utils import OptimizerConfig
from tno.quantum.utils.validation import (
check_arraylike,
check_bool,
check_int,
check_random_state,
check_string,
)
from tqdm import trange
from tno.quantum.optimization.qubo.solvers._dqo._qaoa_result import QAOAResult
from tno.quantum.optimization.qubo.solvers._dqo._qnode_factory import QNodeFactory
if TYPE_CHECKING:
import pennylane as qml
from torch.optim.optimizer import Optimizer
[docs]
class QAOASolver(Solver[QAOAResult]):
"""Solver implementing a basic Quantum Approximate Optimization Algorithm (QAOA).
The Mixer Hamiltonian is a Pauli-X Hamiltonian and the cost Hamiltonian is the
Lenz-Ising Hamiltonian retrieved from the QUBO problem.
Example:
>>> from tno.quantum.optimization.qubo.components import QUBO
>>> from tno.quantum.optimization.qubo.solvers import QAOASolver
>>>
>>> qubo = QUBO([[1, 2, 3], [4, -50, 6], [7, 8, 9]])
>>> solver = QAOASolver(
... num_layers=4,
... num_iter=100,
... )
>>> result = solver.solve(qubo)
>>> result.best_bitvector
BitVector(010)
"""
[docs]
def __init__( # noqa: PLR0913
self,
*,
random_state: int | RandomState | None = None,
num_layers: SupportsInt = 1,
num_iter: SupportsInt = 100,
init_beta: ArrayLike | None = None,
init_gamma: ArrayLike | None = None,
training_device: DeviceConfig | Mapping[str, Any] | None = None,
evaluation_device: DeviceConfig | Mapping[str, Any] | None = None,
optimizer: OptimizerConfig | Mapping[str, Any] | None = None,
diff_method: str | None = "best",
verbose: bool = False,
) -> None:
r"""Init :py:class:`QAOASolver`.
Args:
random_state: Random state to seed. The seed is used when creating random
parameters (gamma and beta). Note that this seed *does not* seed the
devices or optimizers. These should be seeded separately.
num_layers: Number of layers to use. Each layer contains one mixer block and
one cost block. Default is 1.
num_iter: Number of optimization iterations to perform. Default is 100.
init_beta: Initial parameters for the mixer layer. Should be an 1-D ArrayLike
of length `num_layers` with values in the range $[0, \pi)$. If ``None``
(default) is provided the random values will be generated.
init_gamma: Initial parameters for the cost layer. Should be an 1-D ArrayLike
of length `num_layers` with values in range $[0, 2\pi)$. If ``None``
(default) is provided the random values will be generated.
training_device: Training device to use. Must be a ``DeviceConfig`` or a
mapping with the ``"name"`` and ``"options"`` keys. If ``None`` (default) is
provided, ``{"name": "default.qubit", "options": {}}`` is used.
evaluation_device: Evaluation device to use after the circuit has been
optimized. The chosen device must have the `shots` attribute. Must be a
``DeviceConfig`` or a mapping with the "name" and "options" keys. If
``None`` (default) is provided,
``{"name": "default.qubit","options": {"shots": 100}}`` is used.
optimizer: Optimizer to use. Must be an ``OptimizerConfig`` or a mapping with
the "name" and "options" keys. If ``None`` (default) is provided,
``{"name": "adagrad", "options": {}}`` is used.
diff_method: method used to compute the gradients during training. See the
documentation of :class:`pennylane.QNode` for more information of allowed
values.
verbose: If ``True`` print additional information. Default is ``False``.
Raises:
ValueError: If `num_layers` or `num_iter` is less than 1, or if `init_beta` or
`init_gamma` is not 1-dimensional of length `num_layers`.
TypeError: If `random_state`, `num_layers`, `num_iter`, `init_beta` or
`init_gamma` has invalid type.
"""
self._random_state = check_random_state(random_state, "random_state")
self._num_layers = check_int(num_layers, "num_layers", l_bound=1)
self.num_iter = num_iter
self.init_beta = (
init_beta
if init_beta is not None
else self._random_state.uniform(high=np.pi, size=self.num_layers)
)
self.init_gamma = (
init_gamma
if init_gamma is not None
else self._random_state.uniform(high=2.0 * np.pi, size=self.num_layers)
)
self.training_device = training_device
self.evaluation_device = evaluation_device
self.optimizer = optimizer
self._diff_method = (
None if diff_method is None else check_string(diff_method, "diff_method")
)
self.verbose = verbose
@property
def num_layers(self) -> int:
"""Number of layers to use."""
return self._num_layers
@property
def num_iter(self) -> int:
"""Number of optimization iterations to perform."""
return self._num_iter
@num_iter.setter
def num_iter(self, value: SupportsInt) -> None:
self._num_iter = check_int(value, "num_iter", l_bound=1)
@property
def init_beta(self) -> NDArray[np.float64]:
"""Initial parameters for the mixer layer."""
return self._init_beta
@init_beta.setter
def init_beta(self, value: ArrayLike) -> None:
self._init_beta = check_arraylike(
value, "init_beta", ndim=1, shape=(self.num_layers,)
).astype(np.float64)
@property
def init_gamma(self) -> NDArray[np.float64]:
"""Initial parameters for the cost layer."""
return self._init_gamma
@init_gamma.setter
def init_gamma(self, value: ArrayLike) -> None:
self._init_gamma = check_arraylike(
value, "init_gamma", ndim=1, shape=(self.num_layers,)
).astype(np.float64)
@property
def training_device(self) -> DeviceConfig:
"""Configuration for training device to use."""
return self._training_device
@training_device.setter
def training_device(self, value: DeviceConfig | Mapping[str, Any] | None) -> None:
self._training_device = get_default_training_device_if_none(value)
@property
def evaluation_device(self) -> DeviceConfig:
"""Configuration for evaluation device to use."""
return self._evaluation_device
@evaluation_device.setter
def evaluation_device(self, value: DeviceConfig | Mapping[str, Any] | None) -> None:
self._evaluation_device = get_default_evaluation_device_if_none(value)
@property
def optimizer(self) -> OptimizerConfig:
"""Configuration for optimizer to use."""
return self._optimizer
@optimizer.setter
def optimizer(self, value: OptimizerConfig | Mapping[str, Any] | None) -> None:
self._optimizer = get_default_optimizer_if_none(value)
@property
def diff_method(self) -> str | None:
"""Configuration for differantation method to use during training."""
return self._diff_method
@property
def verbose(self) -> bool:
"""Flag to indicate whether to print additional information."""
return self._verbose
@verbose.setter
def verbose(self, value: bool) -> None:
self._verbose = check_bool(value, "verbose")
def _solve(self, qubo: QUBO) -> QAOAResult:
"""Solve the given QUBO using the QAOA algorithm."""
properties = {
"num_iter": self.num_iter,
"init_beta": self._init_beta,
"init_gamma": self._init_gamma,
"training_device": self.training_device,
"evaluation_device": self.evaluation_device,
"optimizer": self.optimizer,
}
if qubo.size == 0:
properties["final_beta"] = properties["init_beta"]
properties["final_gamma"] = properties["init_gamma"]
properties["expval_history"] = []
raw_result: Mapping[str, int] = {"": 1}
return QAOAResult.from_result(qubo, raw_result, properties)
external_fields, interactions, _ = qubo.to_ising()
qnode_factory = QNodeFactory(external_fields, interactions)
qnode_cost = qnode_factory.make_qnode(
self.training_device, self.num_layers, "expval", self.diff_method
)
qnode_counts = qnode_factory.make_qnode(
self.evaluation_device, self.num_layers, "counts"
)
properties.update(
self._optimize_parameters(
qnode_cost=qnode_cost,
num_iter=self.num_iter,
init_beta=self._init_beta,
init_gamma=self._init_gamma,
optimizer=self.optimizer,
verbose=self.verbose,
)
)
raw_result = self._run_optimized_circuit(qnode_counts, properties)
return QAOAResult.from_result(qubo, raw_result, properties)
@staticmethod
def _optimize_parameters( # noqa: PLR0913
qnode_cost: qml.QNode,
num_iter: int,
init_beta: NDArray[np.float64],
init_gamma: NDArray[np.float64],
optimizer: OptimizerConfig,
*,
verbose: bool = False,
) -> dict[str, Any]:
"""Optimize the parameters for the cost circuit.
Args:
qnode_cost: QNode encoding the cost function to optimize.
num_layers: Number of layers to use.
num_iter: Number of optimization iterations to perform.
init_beta: Initial parameters for the mixer layer.
init_gamma: Initial parameters for the cost layer.
optimizer: Configuration for optimizer to use.
verbose: If ``True`` print additional information. Default is ``False``.
Returns:
Dictionary containing three keys: ``"final_gamma"``, ``"final_beta"`` and
``"expval_history"``.
"""
beta = torch.tensor(init_beta, requires_grad=True)
gamma = torch.tensor(init_gamma, requires_grad=True)
optimizer_instance: Optimizer = optimizer.get_instance([gamma, beta])
expval_history = np.zeros(num_iter)
# Optimize gamma and beta
for i in trange(num_iter) if verbose else range(num_iter):
loss = qnode_cost(gamma, beta)
expval_history[i] = loss.detach()
optimizer_instance.zero_grad()
loss.backward()
optimizer_instance.step()
final_gamma, final_beta = optimizer_instance.param_groups[0]["params"]
return {
"final_gamma": final_gamma.tolist(),
"final_beta": final_beta.tolist(),
"expval_history": expval_history.tolist(),
}
@classmethod
def _run_optimized_circuit(
cls, qnode_counts: qml.QNode, properties: dict[str, Any]
) -> Mapping[str, int]:
"""Optimize the parameters for the cost circuit.
Args:
qnode_counts: QNode of the QAOA circuit.
properties: Dictionary containing at least the keys ``"final_gamma"`` and
``"final_beta"``.
Returns:
A mapping with bitstrings as keys and frequencies as values.
"""
final_gamma = properties["final_gamma"]
final_beta = properties["final_beta"]
raw_result: Mapping[str, int] = qnode_counts(final_gamma, final_beta)
return raw_result
def get_default_training_device_if_none(
device: DeviceConfig | Mapping[str, Any] | None = None,
) -> DeviceConfig:
"""Set default training device if the one provided is ``None``.
Default training device ``{"name": "default.qubit", "options": {}}``.
Args:
device: Device configuration or ``None``.
Raises:
KeyError: If `device` does not contain key ``"name"``.
Returns:
Given device configuration or the default training device configuration.
"""
return DeviceConfig.from_mapping(
device if device is not None else {"name": "default.qubit"}
)
def get_default_evaluation_device_if_none(
device: DeviceConfig | Mapping[str, Any] | None = None,
) -> DeviceConfig:
"""Set default evaluation device if the one provided is ``None``.
Default evaluation device ``{"name": "default.qubit", "options": {"shots": 100}}``.
Args:
device: Evaluation configuration or ``None``.
Raises:
KeyError: If `device` does not contain key ``"name"``.
Returns:
Given device configuration or the default evaluation device configuration.
"""
return DeviceConfig.from_mapping(
device
if device is not None
else {"name": "default.qubit", "options": {"shots": 100}}
)
def get_default_optimizer_if_none(
optimizer: OptimizerConfig | Mapping[str, Any] | None = None,
) -> OptimizerConfig:
"""Set default optimizer if the one provided is ``None``.
Default optimizer ``{"name": "adagrad", "options": {}}``.
Args:
optimizer: optimizer configuration or ``None``.
Raises:
KeyError: If `optimizer` does not contain key ``"name"``.
Returns:
Given optimizer or the default optimizer.
"""
return OptimizerConfig.from_mapping(
optimizer if optimizer is not None else {"name": "adagrad"}
)