From 4eb41d12c41e330a7111e530befb02a9da60b8f6 Mon Sep 17 00:00:00 2001 From: Tianyi Hao Date: Tue, 20 Feb 2024 00:59:36 -0600 Subject: [PATCH 1/4] add get_std() to simulators --- qokit/fur/c/qaoa_simulator.py | 14 ++++++-- qokit/fur/mpi_nbcuda/qaoa_simulator.py | 45 ++++++++++++++++++++++---- qokit/fur/nbcuda/qaoa_simulator.py | 40 +++++++++++++++++++---- qokit/fur/python/qaoa_simulator.py | 15 +++++++-- qokit/fur/qaoa_simulator_base.py | 23 +++++++++++-- 5 files changed, 118 insertions(+), 19 deletions(-) diff --git a/qokit/fur/c/qaoa_simulator.py b/qokit/fur/c/qaoa_simulator.py index 985acf48..4e590530 100644 --- a/qokit/fur/c/qaoa_simulator.py +++ b/qokit/fur/c/qaoa_simulator.py @@ -3,14 +3,16 @@ # // 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 +61,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/mpi_nbcuda/qaoa_simulator.py b/qokit/fur/mpi_nbcuda/qaoa_simulator.py index 8f5b971c..6ec42db2 100644 --- a/qokit/fur/mpi_nbcuda/qaoa_simulator.py +++ b/qokit/fur/mpi_nbcuda/qaoa_simulator.py @@ -3,20 +3,25 @@ # // 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 +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): @@ -193,6 +198,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 d3d413af..412dd939 100644 --- a/qokit/fur/nbcuda/qaoa_simulator.py +++ b/qokit/fur/nbcuda/qaoa_simulator.py @@ -3,18 +3,21 @@ # // 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 +93,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 c7236233..a76dfe94 100644 --- a/qokit/fur/python/qaoa_simulator.py +++ b/qokit/fur/python/qaoa_simulator.py @@ -3,11 +3,16 @@ # // 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_fur import apply_qaoa_furx, apply_qaoa_furxy_complete, apply_qaoa_furxy_ring +from ..qaoa_simulator_base import (CostsType, ParamType, QAOAFastSimulatorBase, + TermsType) +from .qaoa_fur import (apply_qaoa_furx, apply_qaoa_furxy_complete, + apply_qaoa_furxy_ring) def little_to_big_endian(arr): @@ -73,6 +78,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 a3182603..5faca360 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 """ ... From 02978244bc442aca8325d53292b6a9a435977114 Mon Sep 17 00:00:00 2001 From: Tianyi Hao Date: Tue, 20 Feb 2024 01:00:09 -0600 Subject: [PATCH 2/4] let get_qaoa_objective() accept sequence of strings --- qokit/qaoa_objective.py | 49 +++++++++++++++++++++++++++-------------- 1 file changed, 33 insertions(+), 16 deletions(-) diff --git a/qokit/qaoa_objective.py b/qokit/qaoa_objective.py index e10de81b..0090d200 100644 --- a/qokit/qaoa_objective.py +++ b/qokit/qaoa_objective.py @@ -3,22 +3,25 @@ # // 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 +107,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 +129,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 +182,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 From b7500ae447c8369eb49cf351803d1fce76a89de7 Mon Sep 17 00:00:00 2001 From: Tianyi Hao Date: Tue, 20 Feb 2024 01:05:01 -0600 Subject: [PATCH 3/4] use pkl5 for mpi --- qokit/fur/lazy_import.py | 1 + qokit/fur/mpi_nbcuda/qaoa_simulator.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/qokit/fur/lazy_import.py b/qokit/fur/lazy_import.py index 2fb2ea88..3e1f049d 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 6ec42db2..3ecfcb59 100644 --- a/qokit/fur/mpi_nbcuda/qaoa_simulator.py +++ b/qokit/fur/mpi_nbcuda/qaoa_simulator.py @@ -14,7 +14,7 @@ from qokit.fur.nbcuda.qaoa_simulator import DeviceArray from ..diagonal_precomputation import precompute_gpu -from ..lazy_import import MPI +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, @@ -71,7 +71,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() From b38a32e92f6214a095dd76389f970df1734849ec Mon Sep 17 00:00:00 2001 From: Tianyi Hao Date: Tue, 20 Feb 2024 01:07:37 -0600 Subject: [PATCH 4/4] black --- qokit/fur/c/qaoa_simulator.py | 3 +-- qokit/fur/mpi_nbcuda/qaoa_simulator.py | 9 +++------ qokit/fur/nbcuda/qaoa_simulator.py | 6 ++---- qokit/fur/python/qaoa_simulator.py | 6 ++---- qokit/qaoa_objective.py | 3 +-- 5 files changed, 9 insertions(+), 18 deletions(-) diff --git a/qokit/fur/c/qaoa_simulator.py b/qokit/fur/c/qaoa_simulator.py index 4e590530..8035c997 100644 --- a/qokit/fur/c/qaoa_simulator.py +++ b/qokit/fur/c/qaoa_simulator.py @@ -9,8 +9,7 @@ import numpy as np from ..diagonal_precomputation import precompute_vectorized_cpu_parallel -from ..qaoa_simulator_base import (CostsType, ParamType, QAOAFastSimulatorBase, - TermsType) +from ..qaoa_simulator_base import CostsType, ParamType, QAOAFastSimulatorBase, TermsType from . import csim from .gates import ComplexArray diff --git a/qokit/fur/mpi_nbcuda/qaoa_simulator.py b/qokit/fur/mpi_nbcuda/qaoa_simulator.py index 3ecfcb59..0e22ca9d 100644 --- a/qokit/fur/mpi_nbcuda/qaoa_simulator.py +++ b/qokit/fur/mpi_nbcuda/qaoa_simulator.py @@ -15,13 +15,10 @@ from ..diagonal_precomputation import precompute_gpu 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 ..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 +from .qaoa_fur import apply_qaoa_furx # , apply_qaoa_furxy_complete, apply_qaoa_furxy_ring def mpi_available(allow_single_process=False): diff --git a/qokit/fur/nbcuda/qaoa_simulator.py b/qokit/fur/nbcuda/qaoa_simulator.py index 412dd939..2bb98cd8 100644 --- a/qokit/fur/nbcuda/qaoa_simulator.py +++ b/qokit/fur/nbcuda/qaoa_simulator.py @@ -13,10 +13,8 @@ from qokit.fur.qaoa_simulator_base import TermsType from ..diagonal_precomputation import precompute_gpu -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 ..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 diff --git a/qokit/fur/python/qaoa_simulator.py b/qokit/fur/python/qaoa_simulator.py index a76dfe94..38728be7 100644 --- a/qokit/fur/python/qaoa_simulator.py +++ b/qokit/fur/python/qaoa_simulator.py @@ -9,10 +9,8 @@ import numpy as np 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) +from ..qaoa_simulator_base import CostsType, ParamType, QAOAFastSimulatorBase, TermsType +from .qaoa_fur import apply_qaoa_furx, apply_qaoa_furxy_complete, apply_qaoa_furxy_ring def little_to_big_endian(arr): diff --git a/qokit/qaoa_objective.py b/qokit/qaoa_objective.py index 0090d200..e77744d7 100644 --- a/qokit/qaoa_objective.py +++ b/qokit/qaoa_objective.py @@ -16,8 +16,7 @@ import qokit.parameter_utils from qokit.parameter_utils import QAOAParameterization -from .fur import (QAOAFastSimulatorBase, choose_simulator, - choose_simulator_xyring) +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