Added BaseSamplerV2 support for Amplitude Estimation algorithms with StatevectorSampler

qiskit_algorithms/amplitude_estimators/ae.py

@@ -20,7 +20,7 @@
 from scipy.optimize import bisect
 
 from qiskit import QuantumCircuit, ClassicalRegister
-from qiskit.primitives import BaseSampler, Sampler
+from qiskit.primitives import BaseSamplerV2, StatevectorSampler
 from .amplitude_estimator import AmplitudeEstimator, AmplitudeEstimatorResult
 from .ae_utils import pdf_a, derivative_log_pdf_a, bisect_max
 from .estimation_problem import EstimationProblem
@@ -66,7 +66,7 @@ def __init__(
         num_eval_qubits: int,
         phase_estimation_circuit: QuantumCircuit | None = None,
         iqft: QuantumCircuit | None = None,
-        sampler: BaseSampler | None = None,
+        sampler: BaseSamplerV2 | None = None,
     ) -> None:
         r"""
         Args:
@@ -95,7 +95,7 @@ def __init__(
         self._sampler = sampler
 
     @property
-    def sampler(self) -> BaseSampler | None:
+    def sampler(self) -> BaseSamplerV2 | None:
         """Get the sampler primitive.
 
         Returns:
@@ -104,7 +104,7 @@ def sampler(self) -> BaseSampler | None:
         return self._sampler
 
     @sampler.setter
-    def sampler(self, sampler: BaseSampler) -> None:
+    def sampler(self, sampler: BaseSamplerV2) -> None:
         """Set sampler primitive.
 
         Args:
@@ -145,7 +145,7 @@ def construct_circuit(
 
         # add measurements if necessary
         if measurement:
-            cr = ClassicalRegister(self._m)
+            cr = ClassicalRegister(self._m, name="meas")
             circuit.add_register(cr)
             circuit.measure(list(range(self._m)), list(range(self._m)))
 
@@ -288,8 +288,10 @@ def estimate(self, estimation_problem: EstimationProblem) -> "AmplitudeEstimatio
                 "The state_preparation property of the estimation problem must be set."
             )
         if self._sampler is None:
-            warnings.warn("No sampler provided, defaulting to Sampler from qiskit.primitives")
-            self._sampler = Sampler()
+            warnings.warn(
+                "No sampler provided, defaulting to StatevectorSampler from qiskit.primitives"
+            )
+            self._sampler = StatevectorSampler()
 
         if estimation_problem.objective_qubits is None:
             raise ValueError("The objective_qubits property of the estimation problem must be set.")
@@ -307,29 +309,32 @@ def estimate(self, estimation_problem: EstimationProblem) -> "AmplitudeEstimatio
         result.post_processing = estimation_problem.post_processing  # type: ignore[assignment]
 
         circuit = self.construct_circuit(estimation_problem, measurement=True)
+
         try:
-            job = self._sampler.run([circuit])
-            ret = job.result()
-        except Exception as exc:
-            raise AlgorithmError("The job was not completed successfully. ") from exc
+            # V2 requires PUB format with shots handling
+            job = self._sampler.run([(circuit,)])
+            pub_result = job.result()[0]
+
+            if pub_result.metadata["shots"]:  # Shot-based results
+                counts = pub_result.data.meas.get_counts()
+                total_shots = sum(counts.values())
+                result.circuit_results = {
+                    bitstr: count / total_shots for bitstr, count in counts.items()
+                }
+            else:  # Statevector results
+                result.circuit_results = {
+                    bitstr: prob for bitstr, prob in pub_result.data.meas.items()
+                }
+
+            exact = True
+            # Store shots from metadata
+            result.shots = pub_result.metadata["shots"] or 1  # Handle statevector case
 
-        shots = ret.metadata[0].get("shots")
-        exact = True
+        except Exception as exc:
+            raise AlgorithmError("Job failed") from exc
 
-        if shots is None:
-            result.circuit_results = ret.quasi_dists[0].binary_probabilities()
-            shots = 1
-        else:
-            result.circuit_results = {
-                k: round(v * shots) for k, v in ret.quasi_dists[0].binary_probabilities().items()
-            }
-            exact = False
-
-        # store shots
-        result.shots = shots
-        samples, measurements = self.evaluate_measurements(
-            result.circuit_results  # type: ignore[arg-type]
-        )
+        # Update measurement processing for V2 bitstring format
+        samples, measurements = self.evaluate_measurements(result.circuit_results)
 
         result.samples = samples
         result.samples_processed = {

qiskit_algorithms/amplitude_estimators/fae.py

@@ -18,7 +18,7 @@
 import numpy as np
 
 from qiskit.circuit import QuantumCircuit, ClassicalRegister
-from qiskit.primitives import BaseSampler, Sampler
+from qiskit.primitives import BaseSamplerV2, StatevectorSampler
 from qiskit_algorithms.exceptions import AlgorithmError
 
 from .amplitude_estimator import AmplitudeEstimator, AmplitudeEstimatorResult
@@ -51,7 +51,7 @@ def __init__(
         delta: float,
         maxiter: int,
         rescale: bool = True,
-        sampler: BaseSampler | None = None,
+        sampler: BaseSamplerV2 | None = None,
     ) -> None:
         r"""
         Args:
@@ -70,7 +70,7 @@ def __init__(
         self._sampler = sampler
 
     @property
-    def sampler(self) -> BaseSampler | None:
+    def sampler(self) -> BaseSamplerV2 | None:
         """Get the sampler primitive.
 
         Returns:
@@ -79,7 +79,7 @@ def sampler(self) -> BaseSampler | None:
         return self._sampler
 
     @sampler.setter
-    def sampler(self, sampler: BaseSampler) -> None:
+    def sampler(self, sampler: BaseSamplerV2) -> None:
         """Set sampler primitive.
 
         Args:
@@ -91,13 +91,13 @@ def _cos_estimate(self, estimation_problem, k, shots):
 
         if self._sampler is None:
             warnings.warn("No sampler provided, defaulting to Sampler from qiskit.primitives")
-            self._sampler = Sampler()
+            self._sampler = StatevectorSampler()
 
         circuit = self.construct_circuit(estimation_problem, k, measurement=True)
 
         try:
-            job = self._sampler.run([circuit], shots=shots)
-            result = job.result()
+            job = self._sampler.run([(circuit,)], shots=shots)  # PUB format
+            pub_result = job.result()[0]
         except Exception as exc:
             raise AlgorithmError("The job was not completed successfully. ") from exc
 
@@ -106,11 +106,15 @@ def _cos_estimate(self, estimation_problem, k, shots):
         self._num_oracle_calls += (2 * k + 1) * shots
 
         # sum over all probabilities where the objective qubits are 1
-        prob = 0
-        for bit, probabilities in result.quasi_dists[0].binary_probabilities().items():
-            # check if it is a good state
-            if estimation_problem.is_good_state(bit):
-                prob += probabilities
+        counts = pub_result.data.meas.get_counts()
+        prob = (
+            sum(
+                count
+                for bitstr, count in counts.items()
+                if estimation_problem.is_good_state(bitstr)
+            )
+            / shots
+        )
 
         cos_estimate = 1 - 2 * prob
 
@@ -146,7 +150,7 @@ def construct_circuit(
 
         # add classical register if needed
         if measurement:
-            c = ClassicalRegister(len(estimation_problem.objective_qubits))
+            c = ClassicalRegister(len(estimation_problem.objective_qubits), name="meas")
             circuit.add_register(c)
 
         # add A operator
@@ -179,7 +183,7 @@ def estimate(self, estimation_problem: EstimationProblem) -> "FasterAmplitudeEst
         """
         if self._sampler is None:
             warnings.warn("No sampler provided, defaulting to Sampler from qiskit.primitives")
-            self._sampler = Sampler()
+            self._sampler = StatevectorSampler()
 
         self._num_oracle_calls = 0
 

qiskit_algorithms/amplitude_estimators/iae.py

@@ -19,7 +19,7 @@
 from scipy.stats import beta
 
 from qiskit import ClassicalRegister, QuantumCircuit
-from qiskit.primitives import BaseSampler, Sampler
+from qiskit.primitives import BaseSamplerV2, StatevectorSampler
 
 from .amplitude_estimator import AmplitudeEstimator, AmplitudeEstimatorResult
 from .estimation_problem import EstimationProblem
@@ -54,7 +54,7 @@ def __init__(
         alpha: float,
         confint_method: str = "beta",
         min_ratio: float = 2,
-        sampler: BaseSampler | None = None,
+        sampler: BaseSamplerV2 | None = None,
     ) -> None:
         r"""
         The output of the algorithm is an estimate for the amplitude `a`, that with at least
@@ -98,7 +98,7 @@ def __init__(
         self._sampler = sampler
 
     @property
-    def sampler(self) -> BaseSampler | None:
+    def sampler(self) -> BaseSamplerV2 | None:
         """Get the sampler primitive.
 
         Returns:
@@ -107,7 +107,7 @@ def sampler(self) -> BaseSampler | None:
         return self._sampler
 
     @sampler.setter
-    def sampler(self, sampler: BaseSampler) -> None:
+    def sampler(self, sampler: BaseSamplerV2) -> None:
         """Set sampler primitive.
 
         Args:
@@ -214,7 +214,7 @@ def construct_circuit(
 
         # add classical register if needed
         if measurement:
-            c = ClassicalRegister(len(estimation_problem.objective_qubits))
+            c = ClassicalRegister(len(estimation_problem.objective_qubits), name="meas")
             circuit.add_register(c)
 
         # add A operator
@@ -272,7 +272,7 @@ def estimate(
         """
         if self._sampler is None:
             warnings.warn("No sampler provided, defaulting to Sampler from qiskit.primitives")
-            self._sampler = Sampler()
+            self._sampler = StatevectorSampler()
 
         # initialize memory variables
         powers = [0]  # list of powers k: Q^k, (called 'k' in paper)
@@ -310,23 +310,23 @@ def estimate(
             counts = {}
 
             try:
-                job = self._sampler.run([circuit])
-                ret = job.result()
+                job = self._sampler.run([(circuit,)])  # PUB format
+                pub_result = job.result()[0]
             except Exception as exc:
                 raise AlgorithmError("The job was not completed successfully. ") from exc
 
-            shots = ret.metadata[0].get("shots")
+            shots = pub_result.metadata["shots"]
             if shots is None:
                 circuit = self.construct_circuit(estimation_problem, k=0, measurement=True)
                 try:
-                    job = self._sampler.run([circuit])
-                    ret = job.result()
+                    job = self._sampler.run([(circuit,)])  # PUB format
+                    pub_result = job.result()[0]
                 except Exception as exc:
                     raise AlgorithmError("The job was not completed successfully. ") from exc
 
                 # calculate the probability of measuring '1'
                 prob = 0.0
-                for bit, probabilities in ret.quasi_dists[0].binary_probabilities().items():
+                for bit, probabilities in pub_result.data.meas.items():
                     # check if it is a good state
                     if estimation_problem.is_good_state(bit):
                         prob += probabilities
@@ -341,13 +341,11 @@ def estimate(
                 num_oracle_queries = 0  # no Q-oracle call, only a single one to A
                 break
 
-            counts = {
-                k: round(v * shots) for k, v in ret.quasi_dists[0].binary_probabilities().items()
-            }
+            # Handle shot-based case with StateVector
+            counts = pub_result.data.meas.get_counts()
 
             # calculate the probability of measuring '1', 'prob' is a_i in the paper
             one_counts, prob = self._good_state_probability(estimation_problem, counts)
-
             num_one_shots.append(one_counts)
 
             # track number of Q-oracle calls

qiskit_algorithms/amplitude_estimators/mlae.py

@@ -22,7 +22,7 @@
 from scipy.stats import norm, chi2
 
 from qiskit import ClassicalRegister, QuantumRegister, QuantumCircuit
-from qiskit.primitives import BaseSampler, Sampler
+from qiskit.primitives import BaseSamplerV2, StatevectorSampler
 
 from .amplitude_estimator import AmplitudeEstimator, AmplitudeEstimatorResult
 from .estimation_problem import EstimationProblem
@@ -54,7 +54,7 @@ def __init__(
         self,
         evaluation_schedule: list[int] | int,
         minimizer: MINIMIZER | None = None,
-        sampler: BaseSampler | None = None,
+        sampler: BaseSamplerV2 | None = None,
     ) -> None:
         r"""
         Args:
@@ -101,7 +101,7 @@ def default_minimizer(objective_fn, bounds):
         self._sampler = sampler
 
     @property
-    def sampler(self) -> BaseSampler | None:
+    def sampler(self) -> BaseSamplerV2 | None:
         """Get the sampler primitive.
 
         Returns:
@@ -110,7 +110,7 @@ def sampler(self) -> BaseSampler | None:
         return self._sampler
 
     @sampler.setter
-    def sampler(self, sampler: BaseSampler) -> None:
+    def sampler(self, sampler: BaseSamplerV2) -> None:
         """Set sampler primitive.
 
         Args:
@@ -142,7 +142,7 @@ def construct_circuits(
 
         # add classical register if needed
         if measurement:
-            c = ClassicalRegister(len(estimation_problem.objective_qubits))
+            c = ClassicalRegister(len(estimation_problem.objective_qubits), name="meas")
             qc_0.add_register(c)
 
         qc_0.compose(estimation_problem.state_preparation, inplace=True)
@@ -276,7 +276,7 @@ def estimate(
         """
         if self._sampler is None:
             warnings.warn("No sampler provided, defaulting to Sampler from qiskit.primitives")
-            self._sampler = Sampler()
+            self._sampler = StatevectorSampler()
 
         if estimation_problem.state_preparation is None:
             raise AlgorithmError(
@@ -291,23 +291,23 @@ def estimate(
         circuits = self.construct_circuits(estimation_problem, measurement=True)
 
         try:
-            job = self._sampler.run(circuits)
-            ret = job.result()
+            pubs = [(circuit,) for circuit in circuits]
+            job = self._sampler.run(pubs)
+            pub_results = job.result()
         except Exception as exc:
             raise AlgorithmError("The job was not completed successfully. ") from exc
 
         circuit_results = []
-        shots = ret.metadata[0].get("shots")
+        shots = pub_results[0].metadata["shots"]
         exact = True
-        if shots is None:
-            for quasi_dist in ret.quasi_dists:
-                circuit_result = quasi_dist.binary_probabilities()
+        if shots is None:  # Statevector simulation
+            for pub_result in pub_results:
+                circuit_result = {bitstr: prob for bitstr, prob in pub_result.data.meas.items()}
                 circuit_results.append(circuit_result)
             shots = 1
-        else:
-            # get counts and construct MLE input
-            for quasi_dist in ret.quasi_dists:
-                counts = {k: round(v * shots) for k, v in quasi_dist.binary_probabilities().items()}
+        else:  # Shot-based results
+            for pub_result in pub_results:
+                counts = pub_result.data.meas.get_counts()
                 circuit_results.append(counts)
             exact = False