qoro-divi 0.3.5__py3-none-any.whl → 0.4.0__py3-none-any.whl

divi/qprog/variational_quantum_algorithm.py

@@ -0,0 +1,786 @@
+# SPDX-FileCopyrightText: 2025 Qoro Quantum Ltd <divi@qoroquantum.de>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from abc import abstractmethod
+from functools import lru_cache, partial
+from itertools import chain, groupby
+from queue import Queue
+
+import numpy as np
+import pennylane as qml
+from scipy.optimize import OptimizeResult
+
+from divi.backends import CircuitRunner
+from divi.circuits import Circuit, MetaCircuit
+from divi.circuits.qem import _NoMitigation
+from divi.qprog.exceptions import _CancelledError
+from divi.qprog.optimizers import ScipyMethod, ScipyOptimizer
+from divi.qprog.quantum_program import QuantumProgram
+from divi.reporting import LoggingProgressReporter, QueueProgressReporter
+from divi.utils import hamiltonian_to_pauli_string, reverse_dict_endianness
+
+logger = logging.getLogger(__name__)
+
+
+def _get_structural_key(obs: qml.operation.Operation) -> tuple[str, ...]:
+    """Generates a hashable, wire-independent key from an observable's structure.
+
+    This function is used to create a canonical representation of an observable
+    based on its constituent Pauli operators, ignoring the wires they act on.
+    This key is ideal for caching computed eigenvalues, as observables with the
+    same structure (e.g., PauliX(0) @ PauliZ(1) and PauliX(2) @ PauliZ(3))
+    share the same eigenvalues. It maps PauliX and PauliY to PauliZ because
+    they are all isospectral (have eigenvalues [1, -1]).
+
+    Args:
+        obs (qml.operation.Operation): A PennyLane observable (e.g., qml.PauliZ(0), qml.PauliX(0) @ qml.PauliY(1)).
+
+    Returns:
+        tuple[str, ...]: A tuple of strings representing the structure of the observable,
+            e.g., ('PauliZ',) or ('PauliZ', 'PauliZ').
+    """
+
+    # Pennylane returns the same eigenvalues for PauliX and PauliY
+    # since it handles diagonalizing gates internally anyway
+    name_map = {
+        "PauliY": "PauliZ",
+        "PauliX": "PauliZ",
+        "PauliZ": "PauliZ",
+        "Identity": "Identity",
+    }
+
+    if isinstance(obs, qml.ops.Prod):
+        # Recursively build a tuple of operator names
+        return tuple(name_map[o.name] for o in obs.operands)
+
+    # For single operators, return a single-element tuple
+    return (name_map[obs.name],)
+
+
+@lru_cache(maxsize=512)
+def _get_eigvals_from_key(key: tuple[str, ...]) -> np.ndarray:
+    """Computes and caches eigenvalues based on a structural key.
+
+    This function takes a key generated by `_get_structural_key` and computes
+    the eigenvalues of the corresponding tensor product of operators. The results
+    are memoized using @lru_cache to avoid redundant calculations.
+
+    Args:
+        key (tuple[str, ...]): A tuple of strings representing the observable's structure.
+
+    Returns:
+        np.ndarray: A NumPy array containing the eigenvalues of the observable.
+    """
+
+    # Define a mapping from name to the base eigenvalue array
+    eigvals_map = {
+        "PauliZ": np.array([1, -1], dtype=np.int8),
+        "Identity": np.array([1, 1], dtype=np.int8),
+    }
+
+    # Start with the eigenvalues of the first operator in the key
+    final_eigvals = eigvals_map[key[0]]
+
+    # Iteratively compute the kronecker product for the rest
+    for op_name in key[1:]:
+        final_eigvals = np.kron(final_eigvals, eigvals_map[op_name])
+
+    return final_eigvals
+
+
+def _batched_expectation(shots_dicts, observables, wire_order):
+    """Efficiently calculates expectation values for multiple observables across multiple shot histograms.
+
+    This function is optimized to compute expectation values in a fully vectorized
+    manner, minimizing Python loops. It operates in four main steps:
+    1. Aggregates all unique bitstrings measured across all histograms.
+    2. Builds a "reduced" eigenvalue matrix corresponding only to the unique states.
+    3. Builds a "reduced" probability matrix from the shot counts for each histogram.
+    4. Computes all expectation values with a single matrix multiplication.
+
+    Args:
+        shots_dicts (list[dict[str, int]]): A list of shot dictionaries (histograms),
+            where each dictionary maps a measured bitstring to its count.
+        observables (list[qml.operation.Operation]): A list of PennyLane observables
+            for which to calculate expectation values.
+        wire_order (tuple[int, ...]): A tuple defining the order of wires, which maps
+            the bitstring to the qubits. Note: This is typically the reverse of the
+            qubit indices (e.g., (2, 1, 0) for a 3-qubit system).
+
+    Returns:
+        np.ndarray: A 2D NumPy array of shape (n_observables, n_shots) where
+            result[i, j] is the expectation value of observables[i] for the
+            histogram in shots_dicts[j].
+    """
+
+    n_histograms = len(shots_dicts)
+    n_total_wires = len(wire_order)
+    n_observables = len(observables)
+
+    # --- 1. Aggregate all unique measured states across all shots ---
+    all_measured_bitstrings = set()
+    for sd in shots_dicts:
+        all_measured_bitstrings.update(sd.keys())
+
+    unique_bitstrings = sorted(list(all_measured_bitstrings))
+    n_unique_states = len(unique_bitstrings)
+
+    bitstring_to_idx_map = {bs: i for i, bs in enumerate(unique_bitstrings)}
+
+    # --- 2. Build REDUCED Eigenvalue Matrix: (n_observables, n_unique_states) ---
+    unique_states_int = np.array(
+        [int(bs, 2) for bs in unique_bitstrings], dtype=np.uint64
+    )
+    reduced_eigvals_matrix = np.zeros((n_observables, n_unique_states))
+    wire_map = {w: i for i, w in enumerate(wire_order)}
+
+    powers_cache = {}
+
+    for obs_idx, observable in enumerate(observables):
+        obs_wires = observable.wires
+        n_obs_wires = len(obs_wires)
+
+        if n_obs_wires in powers_cache:
+            powers = powers_cache[n_obs_wires]
+        else:
+            powers = 2 ** np.arange(n_obs_wires - 1, -1, -1, dtype=np.intp)
+            powers_cache[n_obs_wires] = powers
+
+        obs_wire_indices = np.array([wire_map[w] for w in obs_wires], dtype=np.uint32)
+        eigvals = _get_eigvals_from_key(_get_structural_key(observable))
+
+        # Vectorized mapping, but on the *reduced* set of states
+        shifts = n_total_wires - 1 - obs_wire_indices
+        bits = ((unique_states_int[:, np.newaxis] >> shifts) & 1).astype(np.intp)
+        # powers = 2 ** np.arange(n_obs_wires - 1, -1, -1)
+
+        # obs_state_indices = (bits * powers).sum(axis=1).astype(np.intp)
+        obs_state_indices = np.dot(bits, powers)
+
+        reduced_eigvals_matrix[obs_idx, :] = eigvals[obs_state_indices]
+
+    # --- 3. Build REDUCED Probability Matrix: (n_shots, n_unique_states) ---
+    reduced_prob_matrix = np.zeros((n_histograms, n_unique_states), dtype=np.float32)
+    for i, shots_dict in enumerate(shots_dicts):
+        total = sum(shots_dict.values())
+
+        for bitstring, count in shots_dict.items():
+            col_idx = bitstring_to_idx_map[bitstring]
+            reduced_prob_matrix[i, col_idx] = count / total
+
+    # --- 4. Compute Final Expectation Values ---
+    # (n_shots, n_unique_states) @ (n_unique_states, n_observables)
+    result = reduced_prob_matrix @ reduced_eigvals_matrix.T
+
+    # Transpose to (n_observables, n_shots) as expected by the calling code
+    return result.T
+
+
+def _compute_parameter_shift_mask(n_params):
+    """
+    Generate a binary matrix mask for the parameter shift rule.
+    This mask is used to determine the shifts to apply to each parameter
+    when computing gradients via the parameter shift rule in quantum algorithms.
+
+    Args:
+        n_params (int): The number of parameters in the quantum circuit.
+
+    Returns:
+        np.ndarray: A (2 * n_params, n_params) matrix where each row encodes
+            the shift to apply to each parameter for a single evaluation.
+            The values are multiples of 0.5 * pi, with alternating signs.
+    """
+    mask_arr = np.arange(0, 2 * n_params, 2)
+    mask_arr[0] = 1
+
+    binary_matrix = ((mask_arr[:, np.newaxis] & (1 << np.arange(n_params))) > 0).astype(
+        np.float64
+    )
+
+    binary_matrix = binary_matrix.repeat(2, axis=0)
+    binary_matrix[1::2] *= -1
+    binary_matrix *= 0.5 * np.pi
+
+    return binary_matrix
+
+
+class VariationalQuantumAlgorithm(QuantumProgram):
+    """Base class for variational quantum algorithms.
+
+    This class provides the foundation for implementing variational quantum
+    algorithms in Divi. It handles circuit execution, parameter optimization,
+    and result management for algorithms that optimize parameterized quantum
+    circuits to minimize cost functions.
+
+    Variational algorithms work by:
+    1. Generating parameterized quantum circuits
+    2. Executing circuits on quantum hardware/simulators
+    3. Computing expectation values of cost Hamiltonians
+    4. Using classical optimizers to update parameters
+    5. Iterating until convergence
+
+    Attributes:
+        _losses_history (list[dict]): History of loss values during optimization.
+        _final_params (np.ndarray): Final optimized parameters.
+        _best_params (np.ndarray): Parameters that achieved the best loss.
+        _best_loss (float): Best loss achieved during optimization.
+        _circuits (list[Circuit]): Generated quantum circuits.
+        _total_circuit_count (int): Total number of circuits executed.
+        _total_run_time (float): Total execution time in seconds.
+        _curr_params (np.ndarray): Current parameter values.
+        _seed (int | None): Random seed for parameter initialization.
+        _rng (np.random.Generator): Random number generator.
+        _grad_mode (bool): Whether currently computing gradients.
+        _grouping_strategy (str): Strategy for grouping quantum operations.
+        _qem_protocol (QEMProtocol): Quantum error mitigation protocol.
+        _cancellation_event (Event | None): Event for graceful termination.
+        _meta_circuit_factory (callable): Factory for creating MetaCircuit instances.
+    """
+
+    def __init__(
+        self,
+        backend: CircuitRunner,
+        seed: int | None = None,
+        progress_queue: Queue | None = None,
+        **kwargs,
+    ):
+        """Initialize the VariationalQuantumAlgorithm.
+
+        This constructor is specifically designed for hybrid quantum-classical
+        variational algorithms. The instance variables `n_layers` and `n_params`
+        must be set by subclasses, where:
+        - `n_layers` is the number of layers in the quantum circuit.
+        - `n_params` is the number of parameters per layer.
+
+        For exotic variational algorithms where these variables may not be applicable,
+        the `_initialize_params` method should be overridden to set the parameters.
+
+        Args:
+            backend (CircuitRunner): Quantum circuit execution backend.
+            seed (int | None): Random seed for parameter initialization. Defaults to None.
+            progress_queue (Queue | None): Queue for progress reporting. Defaults to None.
+
+        Keyword Args:
+            grouping_strategy (str): Strategy for grouping operations in Pennylane transforms.
+                Options: "default", "wires", "qwc". Defaults to "qwc".
+            qem_protocol (QEMProtocol | None): Quantum error mitigation protocol to apply. Defaults to None.
+        """
+
+        self._losses_history = []
+
+        self._best_params = None
+        self._final_params = None
+        self._best_loss = float("inf")
+        self._best_probs = {}
+
+        self._curr_params = None
+
+        self._seed = seed
+        self._rng = np.random.default_rng(self._seed)
+
+        # Lets child classes adapt their optimization
+        # step for grad calculation routine
+        self._grad_mode = False
+        self._is_compute_probabilites = False
+
+        super().__init__(
+            backend=backend, seed=seed, progress_queue=progress_queue, **kwargs
+        )
+
+        self.job_id = kwargs.get("job_id", None)
+        if progress_queue and self.job_id:
+            self.reporter = QueueProgressReporter(self.job_id, progress_queue)
+        else:
+            self.reporter = LoggingProgressReporter()
+
+        # Needed for Pennylane's transforms
+        self._grouping_strategy = kwargs.pop("grouping_strategy", "qwc")
+
+        self._qem_protocol = kwargs.pop("qem_protocol", None) or _NoMitigation()
+
+        self._cancellation_event = None
+
+        self._meta_circuit_factory = partial(
+            MetaCircuit,
+            # No grouping strategy for expectation value measurements
+            grouping_strategy=(
+                "_backend_expval"
+                if self.backend and self.backend.supports_expval
+                else self._grouping_strategy
+            ),
+            qem_protocol=self._qem_protocol,
+        )
+
+    @property
+    @abstractmethod
+    def cost_hamiltonian(self) -> qml.operation.Operator:
+        """The cost Hamiltonian for the variational problem."""
+        pass
+
+    @property
+    def total_circuit_count(self) -> int:
+        """Get the total number of circuits executed.
+
+        Returns:
+            int: Cumulative count of circuits submitted for execution.
+        """
+        return self._total_circuit_count
+
+    @property
+    def total_run_time(self) -> float:
+        """Get the total runtime across all circuit executions.
+
+        Returns:
+            float: Cumulative execution time in seconds.
+        """
+        return self._total_run_time
+
+    @property
+    def meta_circuits(self):
+        """Get the meta-circuit templates used by this program.
+
+        Returns:
+            dict[str, MetaCircuit]: Dictionary mapping circuit names to their
+                MetaCircuit templates.
+        """
+        return self._meta_circuits
+
+    @property
+    def n_params(self):
+        """Get the total number of parameters in the quantum circuit.
+
+        Returns:
+            int: Total number of trainable parameters (n_layers * n_params_per_layer).
+        """
+        return self._n_params
+
+    @property
+    def losses_history(self) -> list[dict]:
+        """Get a copy of the optimization loss history.
+
+        Each entry is a dictionary mapping parameter indices to loss values.
+
+        Returns:
+            list[dict]: Copy of the loss history. Modifications to this list
+                will not affect the internal state.
+        """
+        return self._losses_history.copy()
+
+    @property
+    def final_params(self) -> np.ndarray:
+        """Get a copy of the final optimized parameters.
+
+        Returns:
+            np.ndarray: Copy of the final parameters. Modifications to this array
+                will not affect the internal state.
+        """
+        return self._final_params.copy()
+
+    @property
+    def best_params(self) -> np.ndarray:
+        """Get a copy of the parameters that achieved the best (lowest) loss.
+
+        Returns:
+            np.ndarray: Copy of the best parameters. Modifications to this array
+                will not affect the internal state.
+        """
+        return self._best_params.copy()
+
+    @property
+    def best_loss(self) -> float:
+        """Get the best loss achieved so far.
+
+        Returns:
+            float: The best loss achieved so far.
+        """
+        return self._best_loss
+
+    @property
+    def best_probs(self):
+        """Get a copy of the probability distribution for the best parameters.
+
+        Returns:
+            dict: A copy of the best probability distribution.
+        """
+        return self._best_probs.copy()
+
+    def _convert_counts_to_probs(self, results):
+        """Convert raw counts to probability distributions."""
+        return {
+            outer_k: {
+                inner_k: inner_v / self.backend.shots
+                for inner_k, inner_v in outer_v.items()
+            }
+            for outer_k, outer_v in results.items()
+        }
+
+    def _process_probability_results(self, results):
+        """Convert raw counts to probabilities and fix endianness."""
+        probs = self._convert_counts_to_probs(results)
+        return reverse_dict_endianness(probs)
+
+    @property
+    def initial_params(self) -> np.ndarray:
+        """Get the current initial parameters.
+
+        Returns:
+            np.ndarray: Current initial parameters. If not yet initialized,
+                they will be generated automatically.
+        """
+        if self._curr_params is None:
+            self._initialize_params()
+        return self._curr_params.copy()
+
+    @initial_params.setter
+    def initial_params(self, value: np.ndarray | None):
+        """
+        Set initial parameters.
+
+        Args:
+            value (np.ndarray | None): Initial parameters with shape
+                (n_param_sets, n_layers * n_params), or None to reset
+                to uninitialized state.
+
+        Raises:
+            ValueError: If parameters have incorrect shape.
+        """
+        if value is not None:
+            self._validate_initial_params(value)
+            self._curr_params = value.copy()
+        else:
+            # Reset to uninitialized state
+            self._curr_params = None
+
+    @abstractmethod
+    def _create_meta_circuits_dict(self) -> dict[str, MetaCircuit]:
+        pass
+
+    @abstractmethod
+    def _generate_circuits(self, **kwargs) -> list[Circuit]:
+        """Generate quantum circuits for execution.
+
+        This method should generate and return a list of Circuit objects based on
+        the current algorithm state and parameters. The circuits will be executed
+        by the backend.
+
+        Args:
+            **kwargs: Additional keyword arguments for circuit generation.
+
+        Returns:
+            list[Circuit]: List of Circuit objects to be executed.
+        """
+        pass
+
+    def get_expected_param_shape(self) -> tuple[int, int]:
+        """
+        Get the expected shape for initial parameters.
+
+        Returns:
+            tuple[int, int]: Shape (n_param_sets, n_layers * n_params) that
+                initial parameters should have for this quantum program.
+        """
+        return (self.optimizer.n_param_sets, self.n_layers * self.n_params)
+
+    def _validate_initial_params(self, params: np.ndarray):
+        """
+        Validate user-provided initial parameters.
+
+        Args:
+            params (np.ndarray): Parameters to validate.
+
+        Raises:
+            ValueError: If parameters have incorrect shape.
+        """
+        expected_shape = self.get_expected_param_shape()
+
+        if params.shape != expected_shape:
+            raise ValueError(
+                f"Initial parameters must have shape {expected_shape}, "
+                f"got {params.shape}"
+            )
+
+    def _initialize_params(self):
+        """
+        Initialize the circuit parameters randomly.
+
+        Generates random parameters with values uniformly distributed between
+        0 and 2π. The number of parameter sets depends on the optimizer being used.
+        """
+        total_params = self.n_layers * self.n_params
+        self._curr_params = self._rng.uniform(
+            0, 2 * np.pi, (self.optimizer.n_param_sets, total_params)
+        )
+
+    def _run_optimization_circuits(self, data_file, **kwargs):
+        self._curr_circuits = self._generate_circuits(**kwargs)
+
+        if self.backend.supports_expval:
+            kwargs["ham_ops"] = hamiltonian_to_pauli_string(
+                self.cost_hamiltonian, self.n_qubits
+            )
+
+        losses = self._dispatch_circuits_and_process_results(
+            data_file=data_file, **kwargs
+        )
+
+        return losses
+
+    def _post_process_results(
+        self, results: dict[str, dict[str, int]], **kwargs
+    ) -> dict[int, float]:
+        """
+        Post-process the results of the quantum problem.
+
+        Args:
+            results (dict[str, dict[str, int]]): The shot histograms of the quantum execution step.
+                The keys should be strings of format {param_id}_*_{measurement_group_id}.
+                i.e. an underscore-separated bunch of metadata, starting always with
+                the index of some parameter and ending with the index of some measurement group.
+                Any extra piece of metadata that might be relevant to the specific
+                application can be kept in the middle.
+
+        Returns:
+            dict[int, float]: The energies for each parameter set grouping, where the dict keys
+                correspond to the parameter indices.
+        """
+        if not (self._cancellation_event and self._cancellation_event.is_set()):
+            self.reporter.info(
+                message="Post-processing output", iteration=self.current_iteration
+            )
+
+        losses = {}
+        measurement_groups = self._meta_circuits["cost_circuit"].measurement_groups
+
+        # Flatten measurement groups for expectation value backends
+        if self.backend.supports_expval:
+            flattened_measurement_groups = tuple(
+                chain.from_iterable(measurement_groups)
+            )
+        else:
+            flattened_measurement_groups = measurement_groups
+
+        # Define key functions for both levels of grouping
+        get_param_id = lambda item: int(item[0].split("_")[0])
+        get_qem_id = lambda item: int(item[0].split("_")[1].split(":")[1])
+
+        # Group the pre-sorted results by parameter ID.
+        for p, param_group_iterator in groupby(results.items(), key=get_param_id):
+            param_group_iterator = list(param_group_iterator)
+
+            shots_by_qem_idx = zip(
+                *{
+                    gid: [value for _, value in group]
+                    for gid, group in groupby(param_group_iterator, key=get_qem_id)
+                }.values()
+            )
+
+            marginal_results = []
+            for shots_dicts, curr_measurement_group in zip(
+                shots_by_qem_idx, flattened_measurement_groups
+            ):
+                if self.backend.supports_expval:
+                    ham_ops = kwargs.get("ham_ops")
+                    if ham_ops is None:
+                        # Internal consistency check: ham_ops should be set by _run_optimization_circuits
+                        # when backend supports expectation values
+                        raise ValueError(
+                            "Hamiltonian operators (ham_ops) are required when using a backend "
+                            "that supports expectation values, but were not provided."
+                        )
+
+                    expectation_matrix = np.array(
+                        [
+                            [shot_dict[op_name] for op_name in ham_ops.split(";")]
+                            for shot_dict in shots_dicts
+                        ]
+                    ).T
+                else:
+                    wire_order = tuple(reversed(self.cost_hamiltonian.wires))
+
+                    expectation_matrix = _batched_expectation(
+                        shots_dicts, curr_measurement_group, wire_order
+                    )
+
+                # expectation_matrix[i, j] = expectation value for observable i, histogram j
+                curr_marginal_results = []
+                for intermediate_exp_values in expectation_matrix:
+                    mitigated_exp_value = self._qem_protocol.postprocess_results(
+                        intermediate_exp_values
+                    )
+                    curr_marginal_results.append(mitigated_exp_value)
+
+                marginal_results.append(
+                    curr_marginal_results
+                    if len(curr_marginal_results) > 1
+                    else curr_marginal_results[0]
+                )
+
+            if self.backend.supports_expval:
+                marginal_results = marginal_results[0]
+
+            pl_loss = (
+                self._meta_circuits["cost_circuit"]
+                .postprocessing_fn(marginal_results)[0]
+                .item()
+            )
+
+            losses[p] = pl_loss + self.loss_constant
+
+        return losses
+
+    def _perform_final_computation(self, **kwargs):
+        """
+        Perform final computations after optimization completes.
+
+        This is an optional hook method that subclasses can override to perform
+        any post-optimization processing, such as extracting solutions, running
+        final measurements, or computing additional metrics.
+
+        Args:
+            **kwargs: Additional keyword arguments for subclasses.
+
+        Note:
+            The default implementation does nothing. Subclasses should override
+            this method if they need post-optimization processing.
+        """
+        pass
+
+    def run(self, data_file=None, **kwargs):
+        """Run the variational quantum algorithm.
+
+        The outputs are stored in the algorithm object.
+        Optionally, the data can be stored in a file.
+
+        Args:
+            data_file (str | None): The file to store the data in. If None, no data is stored. Defaults to None.
+            **kwargs: Additional keyword arguments for subclasses.
+        """
+
+        def cost_fn(params):
+            self.reporter.info(
+                message="💸 Computing Cost 💸", iteration=self.current_iteration
+            )
+
+            self._curr_params = np.atleast_2d(params)
+
+            losses = self._run_optimization_circuits(data_file, **kwargs)
+
+            losses = np.fromiter(losses.values(), dtype=np.float64)
+
+            if params.ndim > 1:
+                return losses
+            else:
+                return losses.item()
+
+        self._grad_shift_mask = _compute_parameter_shift_mask(
+            self.n_layers * self.n_params
+        )
+
+        def grad_fn(params):
+            self._grad_mode = True
+
+            self.reporter.info(
+                message="📈 Computing Gradients 📈", iteration=self.current_iteration
+            )
+
+            self._curr_params = self._grad_shift_mask + params
+
+            exp_vals = self._run_optimization_circuits(data_file, **kwargs)
+            exp_vals_arr = np.fromiter(exp_vals.values(), dtype=np.float64)
+
+            pos_shifts = exp_vals_arr[::2]
+            neg_shifts = exp_vals_arr[1::2]
+            grads = 0.5 * (pos_shifts - neg_shifts)
+
+            self._grad_mode = False
+
+            return grads
+
+        def _iteration_counter(intermediate_result: OptimizeResult):
+
+            self._losses_history.append(
+                dict(
+                    zip(
+                        range(len(intermediate_result.x)),
+                        intermediate_result.fun,
+                    )
+                )
+            )
+
+            current_loss = np.min(intermediate_result.fun)
+            if current_loss < self._best_loss:
+                self._best_loss = current_loss
+                best_idx = np.argmin(intermediate_result.fun)
+
+                self._best_params = intermediate_result.x[best_idx].copy()
+
+            self.current_iteration += 1
+
+            self.reporter.update(iteration=self.current_iteration)
+
+            if self._cancellation_event and self._cancellation_event.is_set():
+                raise _CancelledError("Cancellation requested by batch.")
+
+            # The scipy implementation of COBYLA interprets the `maxiter` option
+            # as the maximum number of function evaluations, not iterations.
+            # To provide a consistent user experience, we disable `scipy`'s
+            # `maxiter` and manually stop the optimization from the callback
+            # when the desired number of iterations is reached.
+            if (
+                isinstance(self.optimizer, ScipyOptimizer)
+                and self.optimizer.method == ScipyMethod.COBYLA
+                and intermediate_result.nit + 1 == self.max_iterations
+            ):
+                raise StopIteration
+
+        self.reporter.info(message="Finished Setup")
+
+        # Only initialize if user hasn't already set initial_params
+        if self._curr_params is None:
+            self._initialize_params()
+
+        try:
+            self._minimize_res = self.optimizer.optimize(
+                cost_fn=cost_fn,
+                initial_params=self._curr_params,
+                callback_fn=_iteration_counter,
+                jac=grad_fn,
+                maxiter=self.max_iterations,
+                rng=self._rng,
+            )
+        except _CancelledError:
+            # The optimizer was stopped by our callback. This is not a real
+            # error, just a signal to exit this task cleanly.
+            return self._total_circuit_count, self._total_run_time
+
+        self._final_params = self._minimize_res.x
+
+        self._perform_final_computation(**kwargs)
+
+        self.reporter.info(message="Finished successfully!")
+
+        return self.total_circuit_count, self.total_run_time
+
+    def _run_solution_measurement(self):
+        """Execute measurement circuits to obtain probability distributions for solution extraction."""
+        if self._best_params is None:
+            raise RuntimeError(
+                "Optimization has not been run, no best parameters available."
+            )
+
+        if "meas_circuit" not in self._meta_circuits:
+            raise NotImplementedError(
+                f"{type(self).__name__} does not implement a 'meas_circuit'."
+            )
+
+        self._is_compute_probabilites = True
+
+        # Compute probabilities for best parameters (the ones that achieved best loss)
+        self._curr_params = np.atleast_2d(self._best_params)
+        self._curr_circuits = self._generate_circuits()
+        best_probs = self._dispatch_circuits_and_process_results()
+        self._best_probs.update(best_probs)
+
+        self._is_compute_probabilites = False