diff --git a/qokit/fur/c/qaoa_simulator.py b/qokit/fur/c/qaoa_simulator.py index 985acf4..8035c99 100644 --- a/qokit/fur/c/qaoa_simulator.py +++ b/qokit/fur/c/qaoa_simulator.py @@ -3,14 +3,15 @@ # // Copyright : JP Morgan Chase & Co ############################################################################### from __future__ import annotations + from collections.abc import Sequence + import numpy as np -from ..qaoa_simulator_base import QAOAFastSimulatorBase, CostsType, TermsType, ParamType from ..diagonal_precomputation import precompute_vectorized_cpu_parallel - -from .gates import ComplexArray +from ..qaoa_simulator_base import CostsType, ParamType, QAOAFastSimulatorBase, TermsType from . import csim +from .gates import ComplexArray class QAOAFastSimulatorCBase(QAOAFastSimulatorBase): @@ -59,6 +60,12 @@ def get_expectation(self, result: ComplexArray, costs: np.ndarray | None = None, costs = -1 * np.asarray(costs) return np.dot(costs, self.get_probabilities(result, **kwargs)) + def get_std(self, result: ComplexArray, costs: np.ndarray | None = None, **kwargs) -> float: + if costs is None: + costs = self._hc_diag + probs = self.get_probabilities(result) + return np.sqrt(max(np.dot(costs**2, probs) - np.dot(costs, probs) ** 2, 0)) + def get_overlap( self, result: ComplexArray, costs: CostsType | None = None, indices: np.ndarray | Sequence[int] | None = None, optimization_type="min", **kwargs ) -> float: diff --git a/qokit/fur/lazy_import.py b/qokit/fur/lazy_import.py index 2fb2ea8..3e1f049 100644 --- a/qokit/fur/lazy_import.py +++ b/qokit/fur/lazy_import.py @@ -25,3 +25,4 @@ def __getattr__(self, attr): MPI = LasyModule("mpi4py.MPI") +pkl5 = LasyModule("mpi4py.util.pkl5") diff --git a/qokit/fur/mpi_nbcuda/qaoa_simulator.py b/qokit/fur/mpi_nbcuda/qaoa_simulator.py index 8f5b971..0e22ca9 100644 --- a/qokit/fur/mpi_nbcuda/qaoa_simulator.py +++ b/qokit/fur/mpi_nbcuda/qaoa_simulator.py @@ -3,20 +3,22 @@ # // Copyright : JP Morgan Chase & Co ############################################################################### from __future__ import annotations + +import math import typing from collections.abc import Sequence -import numpy as np -import math + import numba.cuda +import numpy as np from qokit.fur.nbcuda.qaoa_simulator import DeviceArray -from ..lazy_import import MPI -from ..nbcuda.qaoa_simulator import QAOAFastSimulatorGPUBase, CostsType, TermsType, ParamType -from ..nbcuda.utils import norm_squared, initialize_uniform, multiply, sum_reduce, copy from ..diagonal_precomputation import precompute_gpu -from .qaoa_fur import apply_qaoa_furx # , apply_qaoa_furxy_complete, apply_qaoa_furxy_ring +from ..lazy_import import MPI, pkl5 +from ..nbcuda.qaoa_simulator import CostsType, ParamType, QAOAFastSimulatorGPUBase, TermsType +from ..nbcuda.utils import copy, initialize_uniform, multiply, norm_squared, sum_reduce from .compute_costs import compute_costs, zero_init +from .qaoa_fur import apply_qaoa_furx # , apply_qaoa_furxy_complete, apply_qaoa_furxy_ring def mpi_available(allow_single_process=False): @@ -66,7 +68,7 @@ def get_costs(terms, N, rank=0): class QAOAFastSimulatorGPUMPIBase(QAOAFastSimulatorGPUBase): def __init__(self, n_qubits: int, costs: CostsType | None = None, terms: TermsType | None = None) -> None: - self._comm = MPI.COMM_WORLD + self._comm = pkl5.Intracomm(MPI.COMM_WORLD) self._rank = self._comm.Get_rank() self._size = self._comm.Get_size() _set_gpu_device_roundrobin() @@ -193,6 +195,34 @@ def get_expectation(self, result: DeviceArray, costs: CostsType | None = None, o else: return global_sum + def get_std(self, result: DeviceArray, costs: CostsType | None = None, **kwargs) -> float: + if costs is None: + preserve_costs = kwargs.get("preserve_costs", True) + if preserve_costs: + costs = numba.cuda.device_array_like(self._hc_diag) + copy(costs, self._hc_diag) + else: + costs = self._hc_diag + else: + costs = self._diag_from_costs(costs) + preserve_state = kwargs.get("preserve_state", True) + if preserve_state: + result_orig = result + result = numba.cuda.device_array_like(result_orig) + copy(result, result_orig) + norm_squared(result) + probs = numba.cuda.device_array_like(result) + copy(probs, result) + multiply(probs, costs) + local_sum = sum_reduce(probs).real # type: ignore + exp = self._comm.allreduce(local_sum, op=MPI.SUM) + del probs + norm_squared(costs) + multiply(result, costs) + local_sum = sum_reduce(result).real # type: ignore + global_sum = self._comm.allreduce(local_sum, op=MPI.SUM) + return np.sqrt(max(global_sum - exp**2, 0)) + def get_overlap( self, result: DeviceArray, costs: CostsType | None = None, indices: np.ndarray | Sequence[int] | None = None, optimization_type="min", **kwargs ) -> float: diff --git a/qokit/fur/nbcuda/qaoa_simulator.py b/qokit/fur/nbcuda/qaoa_simulator.py index d3d413a..2bb98cd 100644 --- a/qokit/fur/nbcuda/qaoa_simulator.py +++ b/qokit/fur/nbcuda/qaoa_simulator.py @@ -3,18 +3,19 @@ # // Copyright : JP Morgan Chase & Co ############################################################################### from __future__ import annotations + +import warnings from collections.abc import Sequence -import numpy as np + import numba.cuda -import warnings +import numpy as np from qokit.fur.qaoa_simulator_base import TermsType -from ..qaoa_simulator_base import QAOAFastSimulatorBase, ParamType, CostsType, TermsType -from .qaoa_fur import apply_qaoa_furx, apply_qaoa_furxy_complete, apply_qaoa_furxy_ring from ..diagonal_precomputation import precompute_gpu -from .utils import norm_squared, initialize_uniform, multiply, sum_reduce, copy - +from ..qaoa_simulator_base import CostsType, ParamType, QAOAFastSimulatorBase, TermsType +from .qaoa_fur import apply_qaoa_furx, apply_qaoa_furxy_complete, apply_qaoa_furxy_ring +from .utils import copy, initialize_uniform, multiply, norm_squared, sum_reduce DeviceArray = numba.cuda.devicearray.DeviceNDArray @@ -90,6 +91,31 @@ def get_expectation(self, result: DeviceArray, costs: DeviceArray | np.ndarray | else: return sum_reduce(result).real + def get_std(self, result: DeviceArray, costs: DeviceArray | np.ndarray | None = None, **kwargs) -> float: + if costs is None: + costs = self._hc_diag + else: + costs = self._diag_from_costs(costs) + preserve_costs = kwargs.get("preserve_costs", True) + if preserve_costs: + costs_orig = costs + costs = numba.cuda.device_array_like(self._hc_diag) + copy(costs, costs_orig) + preserve_state = kwargs.get("preserve_state", True) + if preserve_state: + result_orig = result + result = numba.cuda.device_array_like(result_orig) + copy(result, result_orig) + norm_squared(result) + probs = numba.cuda.device_array_like(result) + copy(probs, result) + multiply(probs, costs) + exp = sum_reduce(probs).real # type: ignore + del probs + norm_squared(costs) + multiply(result, costs) + return np.sqrt(max(sum_reduce(result).real - exp**2, 0)) # type: ignore + def get_overlap( self, result: DeviceArray, costs: CostsType | None = None, indices: np.ndarray | Sequence[int] | None = None, optimization_type="min", **kwargs ) -> float: diff --git a/qokit/fur/python/qaoa_simulator.py b/qokit/fur/python/qaoa_simulator.py index c723623..38728be 100644 --- a/qokit/fur/python/qaoa_simulator.py +++ b/qokit/fur/python/qaoa_simulator.py @@ -3,10 +3,13 @@ # // Copyright : JP Morgan Chase & Co ############################################################################### from __future__ import annotations + from collections.abc import Sequence + import numpy as np -from ..qaoa_simulator_base import QAOAFastSimulatorBase, CostsType, TermsType, ParamType + from ..diagonal_precomputation import precompute_vectorized_cpu_parallel +from ..qaoa_simulator_base import CostsType, ParamType, QAOAFastSimulatorBase, TermsType from .qaoa_fur import apply_qaoa_furx, apply_qaoa_furxy_complete, apply_qaoa_furxy_ring @@ -73,6 +76,12 @@ def get_expectation(self, result: np.ndarray, costs: np.ndarray | None = None, o return -1 * np.dot(costs, np.abs(result) ** 2) return np.dot(costs, np.abs(result) ** 2) + def get_std(self, result: np.ndarray, costs: np.ndarray | None = None, **kwargs) -> float: + if costs is None: + costs = self._hc_diag + probs = np.abs(result) ** 2 + return np.sqrt(max(np.dot(costs**2, probs) - np.dot(costs, probs) ** 2, 0)) + def get_overlap( self, result: np.ndarray, costs: CostsType | None = None, indices: np.ndarray | Sequence[int] | None = None, optimization_type="min", **kwargs ) -> float: diff --git a/qokit/fur/qaoa_simulator_base.py b/qokit/fur/qaoa_simulator_base.py index a318260..5faca36 100644 --- a/qokit/fur/qaoa_simulator_base.py +++ b/qokit/fur/qaoa_simulator_base.py @@ -3,10 +3,12 @@ # // Copyright : JP Morgan Chase & Co ############################################################################### from __future__ import annotations + +import sys import typing -import numpy as np from abc import ABC, abstractmethod -import sys + +import numpy as np # Terms are a list of tuples (coeff, [qubit indices]) # Run this only for python> 3.9 @@ -139,6 +141,21 @@ def get_expectation(self, result, costs: typing.Any = None, optimization_type="m result: obtained from `sim.simulate_qaoa` costs: (optional) array containing values of the cost function at each computational basis state. Accepted types depend on the implementation. + **kwargs : additional arguments for the simulator depending on the implementation + """ + ... + + @abstractmethod + def get_std(self, result, costs: typing.Any = None, **kwargs) -> float: + """ + Return the standard deviation of the expectation of the cost Hamiltonian + + Parameters + ---------- + result: obtained from `sim.simulate_qaoa` + costs: (optional) array containing values of the cost function at + each computational basis state. Accepted types depend on the implementation. + **kwargs : additional arguments for the simulator depending on the implementation """ ... @@ -166,6 +183,7 @@ def get_statevector(self, result: typing.Any, **kwargs) -> np.ndarray: Parameters ---------- result: obtained from `sim.simulate_qaoa` + **kwargs : additional arguments for the simulator depending on the implementation """ ... @@ -177,5 +195,6 @@ def get_probabilities(self, result: typing.Any, **kwargs) -> np.ndarray: Parameters ---------- result: obtained from `sim.simulate_qaoa` + **kwargs : additional arguments for the simulator depending on the implementation """ ... diff --git a/qokit/qaoa_objective.py b/qokit/qaoa_objective.py index e10de81..e77744d 100644 --- a/qokit/qaoa_objective.py +++ b/qokit/qaoa_objective.py @@ -3,22 +3,24 @@ # // Copyright : JP Morgan Chase & Co ############################################################################### from __future__ import annotations + +import typing from calendar import c -import numpy as np -from qiskit import Aer, execute +from collections.abc import Sequence from functools import reduce -import numba.cuda -from .fur import choose_simulator, choose_simulator_xyring, QAOAFastSimulatorBase -import typing +import numba.cuda +import numpy as np +from qiskit import Aer, execute -from .parameter_utils import from_fourier_basis import qokit.parameter_utils from qokit.parameter_utils import QAOAParameterization -from .qaoa_circuit_portfolio import measure_circuit -from .utils import reverse_array_index_bit_order +from .fur import QAOAFastSimulatorBase, choose_simulator, choose_simulator_xyring from .fur.diagonal_precomputation import precompute_vectorized_cpu_parallel +from .parameter_utils import from_fourier_basis +from .qaoa_circuit_portfolio import measure_circuit +from .utils import reverse_array_index_bit_order def _get_qiskit_objective( @@ -104,7 +106,7 @@ def get_qaoa_objective( terms=None, precomputed_optimal_bitstrings=None, parameterization: str | QAOAParameterization = "theta", - objective: str = "expectation", + objective: str | Sequence[str] = "expectation", parameterized_circuit=None, simulator: str = "auto", mixer: str = "x", @@ -126,9 +128,10 @@ def get_qaoa_objective( For below Fourier parameters, q=p If parameterization == 'freq', then f takes one parameter (fourier parameters u and v concatenated) If parameterization == 'u v', then f takes two parameters (fourier parameters u and v) - objective : str + objective : str | Sequence[str] If objective == 'expectation', then returns f(theta) = - < theta | C_{LABS} | theta > (minus for minimization) If objective == 'overlap', then returns f(theta) = 1 - Overlap ||^2 (1-overlap for minimization) + If objective == "std", then returns the standard deviation of the expectation value calculated with the probabilities simulator : str If simulator == 'auto', implementation is chosen automatically (either the fastest CPU simulator or a GPU simulator if CUDA is available) @@ -178,19 +181,32 @@ def fq(*args): if precomputed_costs is None: precomputed_costs = sim.get_cost_diagonal() + objective = [objective] if isinstance(objective, str) else objective + bitstring_loc = None - if precomputed_optimal_bitstrings is not None and objective != "expectation": + if precomputed_optimal_bitstrings is not None and "overlap" in objective: bitstring_loc = np.array([reduce(lambda a, b: 2 * a + b, x) for x in precomputed_optimal_bitstrings]) # -- Final function def f(*args): gamma, beta = qokit.parameter_utils.convert_to_gamma_beta(*args, parameterization=parameterization) result = sim.simulate_qaoa(gamma, beta, initial_state, n_trotters=n_trotters) - if objective == "expectation": - return sim.get_expectation(result, costs=precomputed_costs, preserve_state=False, optimization_type=optimization_type) - elif objective == "overlap": - overlap = sim.get_overlap(result, costs=precomputed_costs, indices=bitstring_loc, preserve_state=False, optimization_type=optimization_type) - return 1 - overlap + ret = [] + if "expectation" in objective: + ret.append(sim.get_expectation(result, costs=precomputed_costs, preserve_state=len(objective) != 1, optimization_type=optimization_type)) + if "overlap" in objective: + ret.append( + 1 + - sim.get_overlap( + result, costs=precomputed_costs, indices=bitstring_loc, preserve_state="std" in objective, optimization_type=optimization_type + ) + ) + if "std" in objective: + ret.append(sim.get_std(result, costs=precomputed_costs, preserve_state=False, preserve_costs=False)) + if len(ret) == 1: + return ret[0] + order = {"expectation": 0, "overlap": int("expectation" in objective), "std": -1} + return tuple([ret[order[obj]] for obj in objective]) return f