add get_std() to simulators

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:

qokit/fur/lazy_import.py

@@ -25,3 +25,4 @@ def __getattr__(self, attr):
 
 
 MPI = LasyModule("mpi4py.MPI")
+pkl5 = LasyModule("mpi4py.util.pkl5")

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:

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:

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:

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
         """
         ...

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 |<theta|optimal_bitstring>|^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