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

divi/qprog/{_qaoa.py → algorithms/_qaoa.py}

@@ -33,6 +33,16 @@ QUBOProblemTypes = list | np.ndarray | sps.spmatrix | QuadraticProgram
 
 
 def draw_graph_solution_nodes(main_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: 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()
@@ -224,14 +234,13 @@ class QAOA(QuantumProgram):
         self.n_layers = n_layers
         self.max_iterations = max_iterations
         self.current_iteration = 0
-        self.n_params = 2
+        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._probs = {}
+        self._solution_nodes = []
+        self._solution_bitstring = []
 
         (
             self.cost_hamiltonian,
@@ -256,18 +265,41 @@ 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 solution(self):
+        """
+        Get the solution found by QAOA optimization.
+
+        Returns:
+            list: 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
             else self._solution_bitstring
         )
 
+    @property
+    def probs(self) -> dict:
+        """
+        Get a copy of the probability distributions from final measurements.
+
+        Returns:
+            dict: Copy of the probability distributions dictionary. Keys are
+                circuit tags, values are probability distributions.
+        """
+        return self._probs.copy()
+
+    @probs.setter
+    def probs(self, value: dict):
+        """Set the probability distributions."""
+        self._probs = value
+
     def _create_meta_circuits_dict(self) -> dict[str, MetaCircuit]:
         """
         Generate the meta circuits for the QAOA problem.
@@ -341,7 +373,7 @@ class QAOA(QuantumProgram):
                 params_group, tag_prefix=f"{p}"
             )
 
-            self.circuits.append(circuit)
+            self._circuits.append(circuit)
 
     def _post_process_results(self, results):
         """
@@ -365,19 +397,26 @@ class QAOA(QuantumProgram):
         return losses
 
     def _run_final_measurement(self):
+        """
+        Execute final measurement circuits to obtain probability distributions.
+
+        Runs the optimized circuit with final parameters to get the full probability
+        distribution over all measurement outcomes, which is used to extract the
+        solution bitstring.
+        """
         self._is_compute_probabilites = True
 
-        self._curr_params = np.array(self.final_params)
+        self._curr_params = np.array(self._final_params)
 
-        self.circuits[:] = []
+        self._circuits[:] = []
 
         self._generate_circuits()
 
-        self.probs.update(self._dispatch_circuits_and_process_results())
+        self._probs.update(self._dispatch_circuits_and_process_results())
 
         self._is_compute_probabilites = False
 
-    def compute_final_solution(self):
+    def _perform_final_computation(self):
         """
         Computes and extracts the final solution from the QAOA optimization process.
         This method performs the following steps:
@@ -397,37 +436,14 @@ class QAOA(QuantumProgram):
 
         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)]
-
-        # 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.")
+        final_measurement_probs = next(iter(self._probs.values()))
 
-        best_solution_key = matching_keys[0]
-        # Retrieve the probability distribution dictionary of the best solution
-        best_solution_probs = self.probs[best_solution_key]
-
-        # 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
-        ]
+        best_solution_bitstring = max(
+            final_measurement_probs, key=final_measurement_probs.get
+        )[::-1]
 
         if isinstance(self.problem, QUBOProblemTypes):
             self._solution_bitstring[:] = np.fromiter(
@@ -439,17 +455,29 @@ class QAOA(QuantumProgram):
                 m.start() for m in re.finditer("1", best_solution_bitstring)
             ]
 
-        self.reporter.info(message="Computed Final Solution!")
-
         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)

divi/qprog/{_vqe.py → algorithms/_vqe.py}

@@ -2,7 +2,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from enum import Enum
 from warnings import warn
 
 import pennylane as qml
@@ -10,38 +9,10 @@ import sympy as sp
 
 from divi.circuits import MetaCircuit
 from divi.qprog import QuantumProgram
+from divi.qprog.algorithms._ansatze import Ansatz, HartreeFockAnsatz
 from divi.qprog.optimizers import MonteCarloOptimizer, Optimizer
 
 
-class VQEAnsatz(Enum):
-    UCCSD = "UCCSD"
-    RY = "RY"
-    RYRZ = "RYRZ"
-    HW_EFFICIENT = "HW_EFFICIENT"
-    QAOA = "QAOA"
-    HARTREE_FOCK = "HF"
-
-    def describe(self):
-        return self.name, self.value
-
-    def n_params(self, n_qubits, **kwargs):
-        if self in (VQEAnsatz.UCCSD, VQEAnsatz.HARTREE_FOCK):
-            singles, doubles = qml.qchem.excitations(
-                kwargs.pop("n_electrons"), n_qubits
-            )
-            s_wires, d_wires = qml.qchem.excitations_to_wires(singles, doubles)
-
-            return len(s_wires) + len(d_wires)
-        elif self == VQEAnsatz.RY:
-            return n_qubits
-        elif self == VQEAnsatz.RYRZ:
-            return 2 * n_qubits
-        elif self == VQEAnsatz.HW_EFFICIENT:
-            raise NotImplementedError
-        elif self == VQEAnsatz.QAOA:
-            return qml.QAOAEmbedding.shape(n_layers=1, n_wires=n_qubits)[1]
-
-
 class VQE(QuantumProgram):
     def __init__(
         self,
@@ -49,7 +20,7 @@ class VQE(QuantumProgram):
         molecule: qml.qchem.Molecule | None = None,
         n_electrons: int | None = None,
         n_layers: int = 1,
-        ansatz=VQEAnsatz.HARTREE_FOCK,
+        ansatz: Ansatz | None = None,
         optimizer: Optimizer | None = None,
         max_iterations=10,
         **kwargs,
@@ -58,19 +29,26 @@ class VQE(QuantumProgram):
         Initialize the VQE problem.
 
         Args:
-            hamiltonain (pennylane.operation.Operator, optional): A Hamiltonian representing the problem.
-            molecule (pennylane.qchem.Molecule, optional): The molecule representing the problem.
-            n_electrons (int, optional): Number of electrons associated with the Hamiltonian.
-                Only needs to be provided when a Hamiltonian is given.
-            ansatz (VQEAnsatz): The ansatz to use for the VQE problem
-            optimizer (Optimizers): The optimizer to use.
-            max_iterations (int): Maximum number of iteration optimizers.
+            hamiltonian (pennylane.operation.Operator, optional): A Hamiltonian
+                representing the problem.
+            molecule (pennylane.qchem.Molecule, optional): The molecule representing
+                the problem.
+            n_electrons (int, optional): Number of electrons associated with the
+                Hamiltonian. Only needs to be provided when a Hamiltonian is given.
+            n_layers (int, optional): Number of ansatz layers. Defaults to 1.
+            ansatz (Ansatz, optional): The ansatz to use for the VQE problem.
+                Defaults to HartreeFockAnsatz.
+            optimizer (Optimizer, optional): The optimizer to use. Defaults to
+                MonteCarloOptimizer.
+            max_iterations (int, optional): Maximum number of optimization iterations.
+                Defaults to 10.
+            **kwargs: Additional keyword arguments passed to the parent QuantumProgram.
         """
 
         # Local Variables
+        self.ansatz = HartreeFockAnsatz() if ansatz is None else ansatz
         self.n_layers = n_layers
         self.results = {}
-        self.ansatz = ansatz
         self.max_iterations = max_iterations
         self.current_iteration = 0
 
@@ -84,20 +62,43 @@ class VQE(QuantumProgram):
 
         self._meta_circuits = self._create_meta_circuits_dict()
 
+    @property
+    def n_params(self):
+        """
+        Get the total number of parameters for the VQE ansatz.
+
+        Returns:
+            int: Total number of parameters (n_params_per_layer * n_layers).
+        """
+        return (
+            self.ansatz.n_params_per_layer(self.n_qubits, n_electrons=self.n_electrons)
+            * self.n_layers
+        )
+
     def _process_problem_input(self, hamiltonian, molecule, n_electrons):
+        """
+        Process and validate the VQE problem input.
+
+        Handles both Hamiltonian-based and molecule-based problem specifications,
+        extracting the necessary information (n_qubits, n_electrons, hamiltonian).
+
+        Args:
+            hamiltonian: PennyLane Hamiltonian operator or None.
+            molecule: PennyLane Molecule object or None.
+            n_electrons: Number of electrons or None.
+
+        Raises:
+            ValueError: If neither hamiltonian nor molecule is provided.
+            UserWarning: If n_electrons conflicts with the molecule's electron count.
+        """
         if hamiltonian is None and molecule is None:
             raise ValueError(
                 "Either one of `molecule` and `hamiltonian` must be provided."
             )
 
         if hamiltonian is not None:
-            if not isinstance(n_electrons, int) or n_electrons < 0:
-                raise ValueError(
-                    f"`n_electrons` is expected to be a non-negative integer. Got {n_electrons}."
-                )
-
-            self.n_electrons = n_electrons
             self.n_qubits = len(hamiltonian.wires)
+            self.n_electrons = n_electrons
 
         if molecule is not None:
             self.molecule = molecule
@@ -112,10 +113,6 @@ class VQE(QuantumProgram):
                     UserWarning,
                 )
 
-        self.n_params = self.ansatz.n_params(
-            self.n_qubits, n_electrons=self.n_electrons
-        )
-
         self.cost_hamiltonian = self._clean_hamiltonian(hamiltonian)
 
     def _clean_hamiltonian(
@@ -148,9 +145,17 @@ class VQE(QuantumProgram):
         return hamiltonian.simplify()
 
     def _create_meta_circuits_dict(self) -> dict[str, MetaCircuit]:
-        weights_syms = sp.symarray("w", (self.n_layers, self.n_params))
+        weights_syms = sp.symarray(
+            "w",
+            (
+                self.n_layers,
+                self.ansatz.n_params_per_layer(
+                    self.n_qubits, n_electrons=self.n_electrons
+                ),
+            ),
+        )
 
-        def _prepare_circuit(ansatz, hamiltonian, params):
+        def _prepare_circuit(hamiltonian, params):
             """
             Prepare the circuit for the VQE problem.
             Args:
@@ -158,7 +163,12 @@ class VQE(QuantumProgram):
                 hamiltonian (qml.Hamiltonian): The Hamiltonian to use
                 params (list): The parameters to use for the ansatz
             """
-            self._set_ansatz(ansatz, params)
+            self.ansatz.build(
+                params,
+                n_qubits=self.n_qubits,
+                n_layers=self.n_layers,
+                n_electrons=self.n_electrons,
+            )
 
             # Even though in principle we want to sample from a state,
             # we are applying an `expval` operation here to make it compatible
@@ -169,93 +179,12 @@ class VQE(QuantumProgram):
         return {
             "cost_circuit": self._meta_circuit_factory(
                 qml.tape.make_qscript(_prepare_circuit)(
-                    self.ansatz, self.cost_hamiltonian, weights_syms
+                    self.cost_hamiltonian, weights_syms
                 ),
                 symbols=weights_syms.flatten(),
             )
         }
 
-    def _set_ansatz(self, ansatz: VQEAnsatz, params):
-        """
-        Set the ansatz for the VQE problem.
-        Args:
-            ansatz (Ansatze): The ansatz to use
-            params (list): The parameters to use for the ansatz
-            n_layers (int): The number of layers to use for the ansatz
-        """
-
-        def _add_hw_efficient_ansatz(params):
-            raise NotImplementedError
-
-        def _add_qaoa_ansatz(params):
-            # This infers layers automatically from the parameters shape
-            qml.QAOAEmbedding(
-                features=[],
-                weights=params.reshape(self.n_layers, -1),
-                wires=range(self.n_qubits),
-            )
-
-        def _add_ry_ansatz(params):
-            qml.layer(
-                qml.AngleEmbedding,
-                self.n_layers,
-                params.reshape(self.n_layers, -1),
-                wires=range(self.n_qubits),
-                rotation="Y",
-            )
-
-        def _add_ryrz_ansatz(params):
-            def _ryrz(params, wires):
-                ry_rots, rz_rots = params.reshape(2, -1)
-                qml.AngleEmbedding(ry_rots, wires=wires, rotation="Y")
-                qml.AngleEmbedding(rz_rots, wires=wires, rotation="Z")
-
-            qml.layer(
-                _ryrz,
-                self.n_layers,
-                params.reshape(self.n_layers, -1),
-                wires=range(self.n_qubits),
-            )
-
-        def _add_uccsd_ansatz(params):
-            hf_state = qml.qchem.hf_state(self.n_electrons, self.n_qubits)
-
-            singles, doubles = qml.qchem.excitations(self.n_electrons, self.n_qubits)
-            s_wires, d_wires = qml.qchem.excitations_to_wires(singles, doubles)
-
-            qml.UCCSD(
-                params.reshape(self.n_layers, -1),
-                wires=range(self.n_qubits),
-                s_wires=s_wires,
-                d_wires=d_wires,
-                init_state=hf_state,
-                n_repeats=self.n_layers,
-            )
-
-        def _add_hartree_fock_ansatz(params):
-            singles, doubles = qml.qchem.excitations(self.n_electrons, self.n_qubits)
-            hf_state = qml.qchem.hf_state(self.n_electrons, self.n_qubits)
-
-            qml.layer(
-                qml.AllSinglesDoubles,
-                self.n_layers,
-                params.reshape(self.n_layers, -1),
-                wires=range(self.n_qubits),
-                hf_state=hf_state,
-                singles=singles,
-                doubles=doubles,
-            )
-
-            # Reset the BasisState operations after the first layer
-            # for behaviour similar to UCCSD ansatz
-            for op in qml.QueuingManager.active_context().queue[1:]:
-                op._hyperparameters["hf_state"] = 0
-
-        if ansatz in VQEAnsatz:
-            locals()[f"_add_{ansatz.name.lower()}_ansatz"](params)
-        else:
-            raise ValueError(f"Invalid Ansatz Value. Got {ansatz}.")
-
     def _generate_circuits(self):
         """
         Generate the circuits for the VQE problem.
@@ -274,9 +203,24 @@ class VQE(QuantumProgram):
                 "cost_circuit"
             ].initialize_circuit_from_params(params_group, tag_prefix=f"{p}")
 
-            self.circuits.append(circuit)
+            self._circuits.append(circuit)
 
     def _run_optimization_circuits(self, store_data, data_file):
+        """
+        Execute the circuits for the current optimization iteration.
+
+        Validates that the Hamiltonian is properly set before running circuits.
+
+        Args:
+            store_data (bool): Whether to save iteration data.
+            data_file (str): Path to file for saving data.
+
+        Returns:
+            dict: Loss values for each parameter set.
+
+        Raises:
+            RuntimeError: If the cost Hamiltonian is not set or empty.
+        """
         if self.cost_hamiltonian is None or len(self.cost_hamiltonian) == 0:
             raise RuntimeError(
                 "Hamiltonian operators must be generated before running the VQE"