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

divi/qprog/quantum_program.py

@@ -5,23 +5,179 @@
 import logging
 import pickle
 from abc import ABC, abstractmethod
-from functools import partial
+from functools import lru_cache, partial
 from itertools import groupby
 from queue import Queue
+from threading import Event
 
 import numpy as np
-from pennylane.measurements import ExpectationMP
+import pennylane as qml
 from scipy.optimize import OptimizeResult
 
 from divi.backends import CircuitRunner, JobStatus, QoroService
 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.reporting import LoggingProgressReporter, QueueProgressReporter
 
 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: A PennyLane observable (e.g., qml.PauliZ(0), qml.PauliX(0) @ qml.PauliY(1)).
+
+    Returns:
+        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: A tuple of strings representing the observable's structure.
+
+    Returns:
+        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.
@@ -56,7 +212,6 @@ class QuantumProgram(ABC):
         backend: CircuitRunner,
         seed: int | None = None,
         progress_queue: Queue | None = None,
-        has_final_computation: bool = False,
         **kwargs,
     ):
         """
@@ -77,30 +232,22 @@ class QuantumProgram(ABC):
                 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.
-
         """
 
-        # Shared Variables
-        self.losses = kwargs.pop("losses", [])
-        self.final_params = kwargs.pop("final_params", [])
+        self._losses = []
+        self._final_params = []
 
-        self.circuits: list[Circuit] = []
+        self._circuits: list[Circuit] = []
 
         self._total_circuit_count = 0
         self._total_run_time = 0.0
-        self._curr_params = []
+        self._curr_params = None
 
         self._seed = seed
         self._rng = np.random.default_rng(self._seed)
@@ -114,9 +261,7 @@ class QuantumProgram(ABC):
         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
-            )
+            self.reporter = QueueProgressReporter(self.job_id, progress_queue)
         else:
             self.reporter = LoggingProgressReporter()
 
@@ -125,6 +270,8 @@ class QuantumProgram(ABC):
 
         self._qem_protocol = kwargs.pop("qem_protocol", None) or _NoMitigation()
 
+        self._cancellation_event = None
+
         self._meta_circuit_factory = partial(
             MetaCircuit,
             grouping_strategy=self._grouping_strategy,
@@ -133,20 +280,113 @@ class QuantumProgram(ABC):
 
     @property
     def total_circuit_count(self):
+        """
+        Get the total number of circuits executed so far.
+
+        Returns:
+            int: Cumulative count of circuits submitted for execution.
+        """
         return self._total_circuit_count
 
     @property
     def total_run_time(self):
+        """
+        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 circuits(self) -> list[Circuit]:
+        """
+        Get a copy of the generated circuits list.
+
+        Returns:
+            list[Circuit]: Copy of the circuits list. Modifications to this list
+                will not affect the internal state.
+        """
+        return self._circuits.copy()
+
+    @property
+    def losses(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.copy()
+
+    @property
+    def final_params(self) -> list:
+        """
+        Get a copy of the final optimized parameters.
+
+        Returns:
+            list: Copy of the final parameters. Modifications to this list
+                will not affect the internal state.
+        """
+        return self._final_params.copy()
+
+    @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
@@ -155,16 +395,60 @@ class QuantumProgram(ABC):
     def _generate_circuits(self, **kwargs):
         pass
 
+    def _set_cancellation_event(self, event: Event):
+        """
+        Set a cancellation event for graceful program termination.
+
+        This internal method is called by a batch runner to provide a mechanism
+        for stopping the optimization loop cleanly when requested.
+
+        Args:
+            event (Event): Threading Event object that signals cancellation when set.
+        """
+        self._cancellation_event = event
+
+    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):
-        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)
-            ]
+        """
+        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, store_data, data_file):
-        self.circuits[:] = []
+        self._circuits[:] = []
 
         self._generate_circuits()
 
@@ -177,7 +461,7 @@ class QuantumProgram(ABC):
     def _prepare_and_send_circuits(self):
         job_circuits = {}
 
-        for circuit in self.circuits:
+        for circuit in self._circuits:
             for tag, qasm_circuit in zip(circuit.tags, circuit.qasm_circuits):
                 job_circuits[tag] = qasm_circuit
 
@@ -261,6 +545,10 @@ class QuantumProgram(ABC):
             (dict) 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
@@ -291,18 +579,17 @@ class QuantumProgram(ABC):
                         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
-                    ]
+                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 obs_idx in range(len(curr_measurement_group)):
+                    intermediate_exp_values = expectation_matrix[obs_idx, :]
                     mitigated_exp_value = self._qem_protocol.postprocess_results(
                         intermediate_exp_values
                     )
-
                     curr_marginal_results.append(mitigated_exp_value)
 
                 marginal_results.append(
@@ -321,6 +608,20 @@ class QuantumProgram(ABC):
 
         return losses
 
+    def _perform_final_computation(self):
+        """
+        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.
+
+        Note:
+            The default implementation does nothing. Subclasses should override
+            this method if they need post-optimization processing.
+        """
+        pass
+
     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.
@@ -371,7 +672,8 @@ class QuantumProgram(ABC):
             return grads
 
         def _iteration_counter(intermediate_result: OptimizeResult):
-            self.losses.append(
+
+            self._losses.append(
                 dict(
                     zip(
                         range(len(intermediate_result.x)),
@@ -380,12 +682,13 @@ class QuantumProgram(ABC):
                 )
             )
 
-            self.final_params[:] = np.atleast_2d(intermediate_result.x)
-
             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.")
+
             if (
                 isinstance(self.optimizer, ScipyOptimizer)
                 and self.optimizer.method == ScipyMethod.COBYLA
@@ -396,26 +699,42 @@ class QuantumProgram(ABC):
         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!")
+        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[:] = np.atleast_2d(self._minimize_res.x)
+
+        self._perform_final_computation()
+
+        self.reporter.info(message="Finished successfully!")
 
         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.
+        Save the current state of the quantum program to a file.
+
+        Serializes the entire QuantumProgram instance including parameters,
+        losses, and circuit history using pickle.
 
         Args:
-            data_file (str): The file to save the iteration to.
+            data_file (str): Path to the file where the program state will be saved.
+
+        Note:
+            The file is written in binary mode and can be restored using
+            `import_iteration()`.
         """
 
         with open(data_file, "wb") as f:
@@ -424,10 +743,16 @@ class QuantumProgram(ABC):
     @staticmethod
     def import_iteration(data_file):
         """
-        Import an iteration of the program from a file.
+        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:

divi/qprog/workflows/_graph_partitioning.py

@@ -387,14 +387,6 @@ def dominance_aggregation(
     return curr_solution
 
 
-def _run_and_compute_solution(program: QuantumProgram):
-    program.run()
-
-    final_sol_circuit_count, final_sol_run_time = program.compute_final_solution()
-
-    return final_sol_circuit_count, final_sol_run_time
-
-
 class GraphPartitioningQAOA(ProgramBatch):
     def __init__(
         self,
@@ -443,8 +435,6 @@ class GraphPartitioningQAOA(ProgramBatch):
         self.solution = None
         self.aggregate_fn = aggregate_fn
 
-        self._task_fn = _run_and_compute_solution
-
         self._constructor = partial(
             QAOA,
             initial_state=initial_state,
@@ -499,33 +489,12 @@ class GraphPartitioningQAOA(ProgramBatch):
                 self.reverse_index_maps[prog_id] = {v: k for k, v in index_map.items()}
 
                 _subgraph = nx.relabel_nodes(subgraph, index_map)
-                self.programs[prog_id] = self._constructor(
+                self._programs[prog_id] = self._constructor(
                     job_id=prog_id,
                     problem=_subgraph,
-                    losses=self._manager.list(),
-                    probs=self._manager.dict(),
-                    final_params=self._manager.list(),
-                    solution_nodes=self._manager.list(),
                     progress_queue=self._queue,
                 )
 
-    def compute_final_solutions(self):
-        if self._executor is not None:
-            self.join()
-
-        if self._executor is not None:
-            raise RuntimeError("A batch is already being run.")
-
-        if len(self.programs) == 0:
-            raise RuntimeError("No programs to run.")
-
-        self._executor = ProcessPoolExecutor()
-
-        self.futures = [
-            self._executor.submit(program.compute_final_solution)
-            for program in self.programs.values()
-        ]
-
     def aggregate_results(self):
         """
         Aggregates the results from all QAOA subprograms to form a global solution.