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

divi/qprog/quantum_program.py

@@ -2,433 +2,256 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import logging
 import pickle
 from abc import ABC, abstractmethod
-from functools import partial
-from itertools import groupby
 from queue import Queue
+from threading import Event
+from typing import Any
 
-import numpy as np
-from pennylane.measurements import ExpectationMP
-from scipy.optimize import OptimizeResult
+from divi.backends import CircuitRunner, JobStatus
+from divi.circuits import Circuit
 
-from divi.backends import CircuitRunner, JobStatus, QoroService
-from divi.circuits import Circuit, MetaCircuit
-from divi.circuits.qem import _NoMitigation
-from divi.qprog.optimizers import ScipyMethod, ScipyOptimizer
-from divi.reporting import LoggingProgressReporter, QueueProgressReporter
 
-logger = logging.getLogger(__name__)
-
-
-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.
+class QuantumProgram(ABC):
+    """Abstract base class for quantum programs.
+
+    This class defines the interface and provides common functionality for quantum algorithms.
+    It handles circuit execution, result processing, and data persistence.
+
+    Subclasses must implement:
+        - run(): Execute the quantum algorithm
+        - _generate_circuits(): Generate quantum circuits for execution
+        - _post_process_results(): Process execution results
+
+    Attributes:
+        backend (CircuitRunner): The quantum circuit execution backend.
+        _seed (int | None): Random seed for reproducible results.
+        _progress_queue (Queue | None): Queue for progress reporting.
+        _circuits (list): List of circuits to be executed.
+        _curr_service_job_id: Current service job ID for QoroService backends.
     """
-    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 QuantumProgram(ABC):
     def __init__(
         self,
         backend: CircuitRunner,
         seed: int | None = None,
         progress_queue: Queue | None = None,
-        has_final_computation: bool = False,
         **kwargs,
     ):
-        """
-        Initializes the QuantumProgram class.
-
-        If a child class represents a hybrid quantum-classical algorithm,
-        the instance variables `n_layers` and `n_params` must be set, where:
-        - `n_layers` is the number of layers in the quantum circuit.
-        - `n_params` is the number of parameters per layer.
-
-        For exotic algorithms where these variables may not be applicable,
-        the `_initialize_params` method should be overridden to set the parameters.
+        """Initialize the QuantumProgram.
 
         Args:
-            backend (CircuitRunner): An instance of a CircuitRunner object, which
-                can either be ParallelSimulator or QoroService.
-            seed (int): A seed for numpy's random number generator, which will
-                be used for the parameter initialization.
-                Defaults to None.
-            progress_queue (Queue): a queue for progress bar updates.
-            has_final_computation (bool): Whether the program includes a final
-                computation step after optimization. This affects progress reporting.
-
-            **kwargs: Additional keyword arguments that influence behaviour.
-                - grouping_strategy (Literal["default", "wires", "qwc"]): A strategy for grouping operations, used in Pennylane's transforms.
-                    Defaults to None.
-                - qem_protocol (QEMProtocol, optional): the quantum error mitigation protocol to apply.
-                    Must be of type QEMProtocol. Defaults to None.
-
-                The following key values are reserved for internal use and should not be set by the user:
-                - losses (list, optional): A list to initialize the `losses` attribute. Defaults to an empty list.
-                - final_params (list, optional): A list to initialize the `final_params` attribute. Defaults to an empty list.
-
+            backend (CircuitRunner): Quantum circuit execution backend.
+            seed (int | None): Random seed for reproducible results. Defaults to None.
+            progress_queue (Queue | None): Queue for progress reporting. Defaults to None.
+            **kwargs: Additional keyword arguments for subclasses.
         """
-
-        # Shared Variables
-        self.losses = kwargs.pop("losses", [])
-        self.final_params = kwargs.pop("final_params", [])
-
-        self.circuits: list[Circuit] = []
-
+        self.backend = backend
+        self._seed = seed
+        self._progress_queue = progress_queue
         self._total_circuit_count = 0
         self._total_run_time = 0.0
-        self._curr_params = []
+        self._curr_circuits = []
+        self._curr_service_job_id = None
 
-        self._seed = seed
-        self._rng = np.random.default_rng(self._seed)
+    @abstractmethod
+    def run(self, data_file: str | None = None, **kwargs) -> tuple[int, float]:
+        """Execute the quantum algorithm.
 
-        # Lets child classes adapt their optimization
-        # step for grad calculation routine
-        self._grad_mode = False
+        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.
 
-        self.backend = backend
+        Returns:
+            tuple[int, float]: A tuple containing:
+                - int: Total number of circuits executed
+                - float: Total runtime in seconds
+        """
+        pass
 
-        self.job_id = kwargs.get("job_id", None)
-        self._progress_queue = progress_queue
-        if progress_queue and self.job_id:
-            self.reporter = QueueProgressReporter(
-                self.job_id, progress_queue, has_final_computation=has_final_computation
-            )
-        else:
-            self.reporter = LoggingProgressReporter()
+    @abstractmethod
+    def _generate_circuits(self, **kwargs) -> list[Circuit]:
+        """Generate quantum circuits for execution.
 
-        # Needed for Pennylane's transforms
-        self._grouping_strategy = kwargs.pop("grouping_strategy", None)
+        This method should generate and return a list of Circuit objects based on
+        the current algorithm state. The circuits will be executed by the backend.
 
-        self._qem_protocol = kwargs.pop("qem_protocol", None) or _NoMitigation()
+        Args:
+            **kwargs: Additional keyword arguments for circuit generation.
 
-        self._meta_circuit_factory = partial(
-            MetaCircuit,
-            grouping_strategy=self._grouping_strategy,
-            qem_protocol=self._qem_protocol,
-        )
+        Returns:
+            list[Circuit]: List of Circuit objects to be executed.
+        """
+        pass
 
-    @property
-    def total_circuit_count(self):
-        return self._total_circuit_count
+    @abstractmethod
+    def _post_process_results(self, results: dict, **kwargs) -> Any:
+        """Process execution results.
 
-    @property
-    def total_run_time(self):
-        return self._total_run_time
+        Args:
+            results (dict): Raw results from circuit execution.
 
-    @property
-    def meta_circuits(self):
-        return self._meta_circuits
+        Returns:
+            Any: Processed results specific to the algorithm.
+        """
+        pass
 
-    @property
-    def n_params(self):
-        return self._n_params
+    def _set_cancellation_event(self, event: Event):
+        """Set a cancellation event for graceful program termination.
 
-    @abstractmethod
-    def _create_meta_circuits_dict(self) -> dict[str, MetaCircuit]:
-        pass
+        This method is called by batch runners to provide a mechanism
+        for stopping the optimization loop cleanly when requested.
 
-    @abstractmethod
-    def _generate_circuits(self, **kwargs):
-        pass
+        Args:
+            event (Event): Threading Event object that signals cancellation when set.
+        """
+        self._cancellation_event = event
 
-    def _initialize_params(self):
-        self._curr_params = np.array(
-            [
-                self._rng.uniform(0, 2 * np.pi, self.n_layers * self.n_params)
-                for _ in range(self.optimizer.n_param_sets)
-            ]
-        )
+    @property
+    def total_circuit_count(self) -> int:
+        """Get the total number of circuits executed.
 
-    def _run_optimization_circuits(self, store_data, data_file):
-        self.circuits[:] = []
+        Returns:
+            int: Cumulative count of circuits submitted for execution.
+        """
+        return self._total_circuit_count
 
-        self._generate_circuits()
+    @property
+    def total_run_time(self) -> float:
+        """Get the total runtime across all circuit executions.
 
-        losses = self._dispatch_circuits_and_process_results(
-            store_data=store_data, data_file=data_file
-        )
+        Returns:
+            float: Cumulative execution time in seconds.
+        """
+        return self._total_run_time
 
-        return losses
+    def _prepare_and_send_circuits(self, **kwargs):
+        """Prepare circuits for execution and submit them to the backend.
 
-    def _prepare_and_send_circuits(self):
+        Returns:
+            Backend output from circuit submission.
+        """
         job_circuits = {}
 
-        for circuit in self.circuits:
+        for circuit in self._curr_circuits:
             for tag, qasm_circuit in zip(circuit.tags, circuit.qasm_circuits):
                 job_circuits[tag] = qasm_circuit
 
         self._total_circuit_count += len(job_circuits)
 
-        backend_output = self.backend.submit_circuits(job_circuits)
+        backend_output = self.backend.submit_circuits(job_circuits, **kwargs)
 
-        if isinstance(self.backend, QoroService):
+        if self.backend.is_async:
             self._curr_service_job_id = backend_output
 
         return backend_output
 
-    def _dispatch_circuits_and_process_results(self, store_data=False, data_file=None):
-        """
-        Run an iteration of the program. The outputs are stored in the Program object.
-        Optionally, the data can be stored in a file.
+    def _track_runtime(self, response):
+        """Extract and track runtime from a backend response.
 
         Args:
-            store_data (bool): Whether to store the data for the iteration
-            data_file (str): The file to store the data in
+            response: Backend response containing runtime information.
+                Can be a dict or a list of responses.
         """
+        if isinstance(response, dict):
+            self._total_run_time += float(response["run_time"])
+        elif isinstance(response, list):
+            self._total_run_time += sum(float(r.json()["run_time"]) for r in response)
 
-        results = self._prepare_and_send_circuits()
+    def _wait_for_qoro_job_completion(self, job_id: str) -> list[dict]:
+        """Wait for a QoroService job to complete and return results.
 
-        def add_run_time(response):
-            if isinstance(response, dict):
-                self._total_run_time += float(response["run_time"])
-            elif isinstance(response, list):
-                self._total_run_time += sum(
-                    float(r.json()["run_time"]) for r in response
-                )
+        Args:
+            job_id: The QoroService job identifier.
 
-        if isinstance(self.backend, QoroService):
+        Returns:
+            list[dict]: The job results from the backend.
+
+        Raises:
+            Exception: If job fails or doesn't complete.
+        """
+        # Build the poll callback if reporter is available
+        if hasattr(self, "reporter"):
             update_function = lambda n_polls, status: self.reporter.info(
                 message="",
                 poll_attempt=n_polls,
                 max_retries=self.backend.max_retries,
-                service_job_id=self._curr_service_job_id,
+                service_job_id=job_id,
                 job_status=status,
             )
+        else:
+            update_function = None
+
+        # Poll until complete
+        status = self.backend.poll_job_status(
+            job_id,
+            loop_until_complete=True,
+            on_complete=self._track_runtime,
+            verbose=False,  # Disable the default logger in QoroService
+            poll_callback=update_function,
+        )
 
-            status = self.backend.poll_job_status(
-                self._curr_service_job_id,
-                loop_until_complete=True,
-                on_complete=add_run_time,
-                verbose=False,  # Disable the default logger in QoroService
-                poll_callback=update_function,  # Use the new, more generic name
-            )
-
-            if status != JobStatus.COMPLETED:
-                raise Exception(
-                    "Job has not completed yet, cannot post-process results"
-                )
-
-            results = self.backend.get_job_results(self._curr_service_job_id)
-
-        results = {r["label"]: r["results"] for r in results}
-
-        result = self._post_process_results(results)
-
-        if store_data:
-            self.save_iteration(data_file)
+        if status != JobStatus.COMPLETED:
+            raise Exception("Job has not completed yet, cannot post-process results")
+        return self.backend.get_job_results(job_id)
 
-        return result
+    def _dispatch_circuits_and_process_results(
+        self, data_file: str | None = None, **kwargs
+    ):
+        """Run an iteration of the program.
 
-    def _post_process_results(
-        self, results: dict[str, dict[str, int]]
-    ) -> dict[int, float]:
-        """
-        Post-process the results of the quantum problem.
+        The outputs are stored in the Program object.
+        Optionally, the data can be stored in a file.
 
         Args:
-            results (dict): 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.
+            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 circuit submission and result processing.
 
         Returns:
-            (dict) The energies for each parameter set grouping, where the dict keys
-                correspond to the parameter indices.
+            Any: Processed results from _post_process_results.
         """
+        results = self._prepare_and_send_circuits(**kwargs)
 
-        losses = {}
-        measurement_groups = self._meta_circuits["cost_circuit"].measurement_groups
+        if self.backend.is_async:
+            results = self._wait_for_qoro_job_completion(self._curr_service_job_id)
 
-        # 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)
+        results = {r["label"]: r["results"] for r in results}
 
-            shots_by_qem_idx = zip(
-                *{
-                    gid: [value for _, value in group]
-                    for gid, group in groupby(param_group_iterator, key=get_qem_id)
-                }.values()
-            )
+        result = self._post_process_results(results, **kwargs)
 
-            marginal_results = []
-            for shots_dicts, curr_measurement_group in zip(
-                shots_by_qem_idx, measurement_groups
-            ):
-                if hasattr(self, "cost_hamiltonian"):
-                    wire_order = tuple(reversed(self.cost_hamiltonian.wires))
-                else:
-                    wire_order = tuple(
-                        reversed(range(len(next(iter(shots_dicts[0].keys())))))
-                    )
-
-                curr_marginal_results = []
-                for observable in curr_measurement_group:
-
-                    intermediate_exp_values = [
-                        ExpectationMP(observable).process_counts(shots_dict, wire_order)
-                        for shots_dict in shots_dicts
-                    ]
-
-                    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]
-                )
-
-            pl_loss = (
-                self._meta_circuits["cost_circuit"]
-                .postprocessing_fn(marginal_results)[0]
-                .item()
-            )
+        if data_file is not None:
+            self.save_iteration(data_file)
 
-            losses[p] = pl_loss + self.loss_constant
+        return result
 
-        return losses
+    def save_iteration(self, data_file: str):
+        """Save the current state of the quantum program to a file.
 
-    def run(self, store_data=False, data_file=None):
-        """
-        Run the QAOA problem. The outputs are stored in the QAOA object. Optionally, the data can be stored in a file.
+        Serializes the entire QuantumProgram instance including parameters,
+        losses, and circuit history using pickle.
 
         Args:
-            store_data (bool): Whether to store the data for the iteration
-            data_file (str): The file to store the data in
-        """
-
-        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(store_data, data_file)
+            data_file (str): Path to the file where the program state will be saved.
 
-            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(store_data, data_file)
-            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.append(
-                dict(
-                    zip(
-                        range(len(intermediate_result.x)),
-                        np.atleast_1d(intermediate_result.fun),
-                    )
-                )
-            )
-
-            self.final_params[:] = np.atleast_2d(intermediate_result.x)
-
-            self.current_iteration += 1
-
-            self.reporter.update(iteration=self.current_iteration)
-
-            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")
-
-        self._initialize_params()
-        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,
-        )
-        self.final_params[:] = np.atleast_2d(self._minimize_res.x)
-
-        self.reporter.info(message="Finished Optimization!")
-
-        return self._total_circuit_count, self._total_run_time
-
-    def save_iteration(self, data_file):
-        """
-        Save the current iteration of the program to a file.
-
-        Args:
-            data_file (str): The file to save the iteration to.
+        Note:
+            The file is written in binary mode and can be restored using
+            `import_iteration()`.
         """
-
         with open(data_file, "wb") as f:
             pickle.dump(self, f)
 
     @staticmethod
-    def import_iteration(data_file):
-        """
-        Import an iteration of the program from a file.
+    def import_iteration(data_file: str):
+        """Load a previously saved quantum program state from a file.
+
+        Deserializes a QuantumProgram instance that was saved using `save_iteration()`.
 
         Args:
-            data_file (str): The file to import the iteration from.
-        """
+            data_file (str): Path to the file containing the saved program state.
 
+        Returns:
+            QuantumProgram: The restored QuantumProgram instance with all its state,
+                including parameters, losses, and circuit history.
+        """
         with open(data_file, "rb") as f:
             return pickle.load(f)