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

divi/qprog/algorithms/_qaoa.py

@@ -21,9 +21,9 @@ from qiskit_optimization import QuadraticProgram
 from qiskit_optimization.converters import QuadraticProgramToQubo
 from qiskit_optimization.problems import VarType
 
-from divi.circuits import MetaCircuit
-from divi.qprog import QuantumProgram
+from divi.circuits import Circuit, MetaCircuit
 from divi.qprog.optimizers import MonteCarloOptimizer, Optimizer
+from divi.qprog.variational_quantum_algorithm import VariationalQuantumAlgorithm
 from divi.utils import convert_qubo_matrix_to_pennylane_ising
 
 logger = logging.getLogger(__name__)
@@ -32,7 +32,16 @@ GraphProblemTypes = nx.Graph | rx.PyGraph
 QUBOProblemTypes = list | np.ndarray | sps.spmatrix | QuadraticProgram
 
 
-def draw_graph_solution_nodes(main_graph, partition_nodes):
+def draw_graph_solution_nodes(main_graph: nx.Graph, partition_nodes):
+    """Visualize a graph with solution nodes highlighted.
+
+    Draws the graph with nodes colored to distinguish solution nodes (red) from
+    other nodes (light blue).
+
+    Args:
+        main_graph (nx.Graph): NetworkX graph to visualize.
+        partition_nodes: Collection of node indices that are part of the solution.
+    """
     # Create a dictionary for node colors
     node_colors = [
         "red" if node in partition_nodes else "lightblue" for node in main_graph.nodes()
@@ -55,6 +64,14 @@ def draw_graph_solution_nodes(main_graph, partition_nodes):
 
 
 class GraphProblem(Enum):
+    """Enumeration of supported graph problems for QAOA.
+
+    Each problem type defines:
+    - pl_string: The corresponding PennyLane function name
+    - constrained_initial_state: Recommended initial state for constrained problems
+    - unconstrained_initial_state: Recommended initial state for unconstrained problems
+    """
+
     MAX_CLIQUE = ("max_clique", "Zeros", "Superposition")
     MAX_INDEPENDENT_SET = ("max_independent_set", "Zeros", "Superposition")
     MAX_WEIGHT_CYCLE = ("max_weight_cycle", "Superposition", "Superposition")
@@ -70,6 +87,13 @@ class GraphProblem(Enum):
         constrained_initial_state: str,
         unconstrained_initial_state: str,
     ):
+        """Initialize the GraphProblem enum value.
+
+        Args:
+            pl_string (str): The corresponding PennyLane function name.
+            constrained_initial_state (str): Recommended initial state for constrained problems.
+            unconstrained_initial_state (str): Recommended initial state for unconstrained problems.
+        """
         self.pl_string = pl_string
 
         # Recommended initial state as per Pennylane's documentation.
@@ -84,6 +108,17 @@ _SUPPORTED_INITIAL_STATES_LITERAL = Literal[
 
 
 def _convert_quadratic_program_to_pennylane_ising(qp: QuadraticProgram):
+    """Convert a Qiskit QuadraticProgram to a PennyLane Ising Hamiltonian.
+
+    Args:
+        qp (QuadraticProgram): Qiskit QuadraticProgram to convert.
+
+    Returns:
+        tuple[qml.Hamiltonian, float, int]: (pennylane_ising, constant, n_qubits) where:
+            - pennylane_ising: The Ising Hamiltonian in PennyLane format
+            - constant: The constant term
+            - n_qubits: Number of qubits required
+    """
     qiskit_sparse_op, constant = qp.to_ising()
 
     pauli_list = qiskit_sparse_op.paulis
@@ -108,9 +143,16 @@ def _convert_quadratic_program_to_pennylane_ising(qp: QuadraticProgram):
 def _resolve_circuit_layers(
     initial_state, problem, graph_problem, **kwargs
 ) -> tuple[qml.operation.Operator, qml.operation.Operator, dict | None, str]:
-    """
-    Generates the cost and mixer hamiltonians for a given problem, in addition to
-    optional metadata returned by Pennylane if applicable
+    """Generate the cost and mixer Hamiltonians for a given problem.
+
+    Args:
+        initial_state (str): The initial state specification.
+        problem (GraphProblemTypes | QUBOProblemTypes): The problem to solve (graph or QUBO).
+        graph_problem (GraphProblem | None): The graph problem type (if applicable).
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        tuple[qml.operation.Operator, qml.operation.Operator, dict | None, str]: (cost_hamiltonian, mixer_hamiltonian, metadata, resolved_initial_state)
     """
 
     if isinstance(problem, GraphProblemTypes):
@@ -149,7 +191,36 @@ def _resolve_circuit_layers(
         )
 
 
-class QAOA(QuantumProgram):
+class QAOA(VariationalQuantumAlgorithm):
+    """Quantum Approximate Optimization Algorithm (QAOA) implementation.
+
+    QAOA is a hybrid quantum-classical algorithm designed to solve combinatorial
+    optimization problems. It alternates between applying a cost Hamiltonian
+    (encoding the problem) and a mixer Hamiltonian (enabling exploration).
+
+    The algorithm can solve:
+    - Graph problems (MaxCut, Max Clique, etc.)
+    - QUBO (Quadratic Unconstrained Binary Optimization) problems
+    - Quadratic programs (converted to QUBO)
+
+    Attributes:
+        problem (GraphProblemTypes | QUBOProblemTypes): The problem instance to solve.
+        graph_problem (GraphProblem | None): The graph problem type (if applicable).
+        n_layers (int): Number of QAOA layers.
+        n_qubits (int): Number of qubits required.
+        cost_hamiltonian (qml.Hamiltonian): The cost Hamiltonian encoding the problem.
+        mixer_hamiltonian (qml.Hamiltonian): The mixer Hamiltonian for exploration.
+        initial_state (str): The initial quantum state.
+        problem_metadata (dict | None): Additional metadata from problem setup.
+        loss_constant (float): Constant term from the problem.
+        optimizer (Optimizer): Classical optimizer for parameter updates.
+        max_iterations (int): Maximum number of optimization iterations.
+        current_iteration (int): Current optimization iteration.
+        _n_params (int): Number of parameters per layer (always 2 for QAOA).
+        _solution_nodes (list[int] | None): Solution nodes for graph problems.
+        _solution_bitstring (np.ndarray | None): Solution bitstring for QUBO problems.
+    """
+
     def __init__(
         self,
         problem: GraphProblemTypes | QUBOProblemTypes,
@@ -161,16 +232,18 @@ class QAOA(QuantumProgram):
         max_iterations: int = 10,
         **kwargs,
     ):
-        """
-        Initialize the QAOA problem.
+        """Initialize the QAOA problem.
 
         Args:
-            problem: The problem to solve, can either be a graph or a QUBO.
+            problem (GraphProblemTypes | QUBOProblemTypes): The problem to solve, can either be a graph or a QUBO.
                 For graph inputs, the graph problem to solve must be provided
                 through the `graph_problem` variable.
-            graph_problem (str): The graph problem to solve.
-            n_layers (int): number of QAOA layers
-            initial_state (str): The initial state of the circuit
+            graph_problem (GraphProblem | None): The graph problem to solve. Defaults to None.
+            n_layers (int): Number of QAOA layers. Defaults to 1.
+            initial_state (_SUPPORTED_INITIAL_STATES_LITERAL): The initial state of the circuit. Defaults to "Recommended".
+            optimizer (Optimizer | None): The optimizer to use. Defaults to MonteCarloOptimizer.
+            max_iterations (int): Maximum number of optimization iterations. Defaults to 10.
+            **kwargs: Additional keyword arguments passed to the parent class.
         """
 
         if isinstance(problem, QUBOProblemTypes):
@@ -225,17 +298,14 @@ class QAOA(QuantumProgram):
         self.max_iterations = max_iterations
         self.current_iteration = 0
         self._n_params = 2
-        self._is_compute_probabilites = False
         self.optimizer = optimizer if optimizer is not None else MonteCarloOptimizer()
 
-        # Shared Variables
-        self.probs = kwargs.pop("probs", {})
-        self._solution_nodes = kwargs.pop("solution_nodes", [])
-        self._solution_bitstring = kwargs.pop("solution_bitstring", [])
+        self._solution_nodes = []
+        self._solution_bitstring = []
 
         (
-            self.cost_hamiltonian,
-            self.mixer_hamiltonian,
+            self._cost_hamiltonian,
+            self._mixer_hamiltonian,
             *problem_metadata,
             self.initial_state,
         ) = _resolve_circuit_layers(
@@ -256,12 +326,28 @@ class QAOA(QuantumProgram):
             self.loss_constant = 0.0
 
         kwargs.pop("is_constrained", None)
-        super().__init__(has_final_computation=True, **kwargs)
+        super().__init__(**kwargs)
 
         self._meta_circuits = self._create_meta_circuits_dict()
 
+    @property
+    def cost_hamiltonian(self) -> qml.operation.Operator:
+        """The cost Hamiltonian for the QAOA problem."""
+        return self._cost_hamiltonian
+
+    @property
+    def mixer_hamiltonian(self) -> qml.operation.Operator:
+        """The mixer Hamiltonian for the QAOA problem."""
+        return self._mixer_hamiltonian
+
     @property
     def solution(self):
+        """Get the solution found by QAOA optimization.
+
+        Returns:
+            list[int] | np.ndarray: For graph problems, returns a list of selected node indices.
+                For QUBO problems, returns a list/array of binary values.
+        """
         return (
             self._solution_nodes
             if self.graph_problem is not None
@@ -269,11 +355,14 @@ class QAOA(QuantumProgram):
         )
 
     def _create_meta_circuits_dict(self) -> dict[str, MetaCircuit]:
-        """
-        Generate the meta circuits for the QAOA problem.
+        """Generate the meta circuits for the QAOA problem.
+
+        Creates both cost and measurement circuits for the QAOA algorithm.
+        The cost circuit is used during optimization, while the measurement
+        circuit is used for final solution extraction.
 
-        In this method, we generate the scaffolding for the circuits that will be
-        executed during optimization.
+        Returns:
+            dict[str, MetaCircuit]: Dictionary containing cost_circuit and meas_circuit.
         """
 
         betas = sp.symarray("β", self.n_layers)
@@ -283,14 +372,16 @@ class QAOA(QuantumProgram):
 
         def _qaoa_layer(params):
             gamma, beta = params
-            pqaoa.cost_layer(gamma, self.cost_hamiltonian)
-            pqaoa.mixer_layer(beta, self.mixer_hamiltonian)
+            pqaoa.cost_layer(gamma, self._cost_hamiltonian)
+            pqaoa.mixer_layer(beta, self._mixer_hamiltonian)
 
         def _prepare_circuit(hamiltonian, params, final_measurement):
-            """
-            Prepare the circuit for the QAOA problem.
+            """Prepare the circuit for the QAOA problem.
+
             Args:
-                hamiltonian (qml.Hamiltonian): The Hamiltonian term to measure
+                hamiltonian (qml.Hamiltonian): The Hamiltonian term to measure.
+                params (np.ndarray): The QAOA parameters (betas and gammas).
+                final_measurement (bool): Whether to perform final measurement.
             """
 
             # Note: could've been done as qml.[Insert Gate](wires=range(self.n_qubits))
@@ -313,121 +404,86 @@ class QAOA(QuantumProgram):
         return {
             "cost_circuit": self._meta_circuit_factory(
                 qml.tape.make_qscript(_prepare_circuit)(
-                    self.cost_hamiltonian, sym_params, final_measurement=False
+                    self._cost_hamiltonian, sym_params, final_measurement=False
                 ),
                 symbols=sym_params.flatten(),
             ),
             "meas_circuit": self._meta_circuit_factory(
                 qml.tape.make_qscript(_prepare_circuit)(
-                    self.cost_hamiltonian, sym_params, final_measurement=True
+                    self._cost_hamiltonian, sym_params, final_measurement=True
                 ),
                 symbols=sym_params.flatten(),
+                grouping_strategy="wires",
             ),
         }
 
-    def _generate_circuits(self):
-        """
-        Generate the circuits for the QAOA problem.
+    def _generate_circuits(self) -> list[Circuit]:
+        """Generate the circuits for the QAOA problem.
 
-        In this method, we generate bulk circuits based on the selected parameters.
-        """
+        Generates circuits for each parameter set in the current parameters.
+        The circuit type depends on whether we're computing probabilities
+        (for final solution extraction) or just expectation values (for optimization).
 
+        Returns:
+            list[Circuit]: List of Circuit objects for execution.
+        """
         circuit_type = (
             "cost_circuit" if not self._is_compute_probabilites else "meas_circuit"
         )
 
-        for p, params_group in enumerate(self._curr_params):
-            circuit = self._meta_circuits[circuit_type].initialize_circuit_from_params(
+        return [
+            self._meta_circuits[circuit_type].initialize_circuit_from_params(
                 params_group, tag_prefix=f"{p}"
             )
+            for p, params_group in enumerate(self._curr_params)
+        ]
 
-            self.circuits.append(circuit)
+    def _post_process_results(self, results, **kwargs):
+        """Post-process the results of the QAOA problem.
 
-    def _post_process_results(self, results):
-        """
-        Post-process the results of the QAOA problem.
+        Args:
+            results (dict[str, dict[str, int]]): Raw results from circuit execution.
+            **kwargs: Additional keyword arguments.
+                ham_ops (str): The Hamiltonian operators to measure, semicolon-separated.
+                    Only needed when the backend supports expval.
 
         Returns:
-            (dict) The losses for each parameter set grouping.
+            dict[str, dict[str, float]] | dict[int, float]: The losses for each parameter set grouping, or probability
+                distributions if computing probabilities.
         """
 
         if self._is_compute_probabilites:
-            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()
-            }
-
-        losses = super()._post_process_results(results)
+            return self._process_probability_results(results)
 
+        losses = super()._post_process_results(results, **kwargs)
         return losses
 
-    def _run_final_measurement(self):
-        self._is_compute_probabilites = True
-
-        self._curr_params = np.array(self.final_params)
-
-        self.circuits[:] = []
-
-        self._generate_circuits()
-
-        self.probs.update(self._dispatch_circuits_and_process_results())
+    def _perform_final_computation(self, **kwargs):
+        """Extract the optimal solution from the QAOA optimization process.
 
-        self._is_compute_probabilites = False
-
-    def compute_final_solution(self):
-        """
-        Computes and extracts the final solution from the QAOA optimization process.
         This method performs the following steps:
-            1. Identifies the best solution index based on the lowest loss value from the last optimization step.
-            2. Executes the final measurement circuit to obtain the probability distributions of solutions.
-            3. Retrieves the bitstring representing the best solution, correcting for endianness.
-            4. Depending on the problem type:
-            - For QUBO problems, stores the solution as a NumPy array of bits.
-            - For graph problems, stores the solution as a list of node indices corresponding to '1's in the bitstring.
-            5. Returns the total circuit count and total runtime for the optimization process.
+        1. Executes measurement circuits with the best parameters (those that achieved the lowest loss).
+        2. Retrieves the bitstring representing the best solution, correcting for endianness.
+        3. Depending on the problem type:
+           - For QUBO problems, stores the solution as a NumPy array of bits.
+           - For graph problems, stores the solution as a list of node indices corresponding to '1's in the bitstring.
 
         Returns:
-            tuple: A tuple containing:
-            - int: The total number of circuits executed.
-            - float: The total runtime of the optimization process.
+            tuple[int, float]: A tuple containing:
+                - int: The total number of circuits executed.
+                - float: The total runtime of the optimization process.
         """
 
-        self.reporter.info(message="🏁 Computing Final Solution 🏁")
-
-        # Convert losses dict to list to apply ordinal operations
-        final_losses_list = list(self.losses[-1].values())
-
-        # Get the index of the smallest loss in the last operation
-        best_solution_idx = min(
-            range(len(final_losses_list)),
-            key=lambda x: final_losses_list.__getitem__(x),
-        )
-
-        # Insert the measurement circuit here
-        self._run_final_measurement()
-
-        # Find the key matching the best_solution_idx with possible metadata in between
-        pattern = re.compile(rf"^{best_solution_idx}(?:_[^_]*)*_0$")
-        matching_keys = [k for k in self.probs.keys() if pattern.match(k)]
+        self.reporter.info(message="🏁 Computing Final Solution 🏁\r")
 
-        # Some minor sanity checks
-        if len(matching_keys) == 0:
-            raise RuntimeError("No matching key found.")
-        if len(matching_keys) > 1:
-            raise RuntimeError(f"More than one matching key found.")
+        self._run_solution_measurement()
 
-        best_solution_key = matching_keys[0]
-        # Retrieve the probability distribution dictionary of the best solution
-        best_solution_probs = self.probs[best_solution_key]
+        best_measurement_probs = next(iter(self._best_probs.values()))
 
-        # Retrieve the bitstring with the actual best solution
-        # Reverse to account for the endianness difference
-        best_solution_bitstring = max(best_solution_probs, key=best_solution_probs.get)[
-            ::-1
-        ]
+        # Endianness is corrected in _post_process_results
+        best_solution_bitstring = max(
+            best_measurement_probs, key=best_measurement_probs.get
+        )
 
         if isinstance(self.problem, QUBOProblemTypes):
             self._solution_bitstring[:] = np.fromiter(
@@ -439,17 +495,30 @@ class QAOA(QuantumProgram):
                 m.start() for m in re.finditer("1", best_solution_bitstring)
             ]
 
-        self.reporter.info(message="Computed Final Solution!")
+        self.reporter.info(message="🏁 Computed Final Solution! 🏁\r\n")
 
         return self._total_circuit_count, self._total_run_time
 
     def draw_solution(self):
+        """Visualize the solution found by QAOA for graph problems.
+
+        Draws the graph with solution nodes highlighted in red and other nodes
+        in light blue. If the solution hasn't been computed yet, it will be
+        calculated first.
+
+        Raises:
+            RuntimeError: If called on a QUBO problem instead of a graph problem.
+
+        Note:
+            This method only works for graph problems. For QUBO problems, access
+            the solution directly via the `solution` property.
+        """
         if self.graph_problem is None:
             raise RuntimeError(
                 "The problem is not a graph problem. Cannot draw solution."
             )
 
         if not self._solution_nodes:
-            self.compute_final_solution()
+            self._perform_final_computation()
 
         draw_graph_solution_nodes(self.problem, self._solution_nodes)