qoro-divi 0.2.0b1__py3-none-any.whl → 0.5.0__py3-none-any.whl

divi/qprog/workflows/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: 2025 Qoro Quantum Ltd <divi@qoroquantum.de>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from ._graph_partitioning import (
+    GraphPartitioningQAOA,
+    PartitioningConfig,
+)
+from ._qubo_partitioning import QUBOPartitioningQAOA
+from ._vqe_sweep import MoleculeTransformer, VQEHyperparameterSweep

divi/qprog/{_graph_partitioning.py → workflows/_graph_partitioning.py}

@@ -3,13 +3,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import heapq
-import re
 import string
 from collections.abc import Callable, Sequence
-from concurrent.futures import ProcessPoolExecutor
 from dataclasses import dataclass
 from functools import partial
-from typing import Literal, Optional
+from typing import Literal
 from warnings import warn
 
 import matplotlib.cm as cm
@@ -21,15 +19,15 @@ import scipy.sparse.linalg as spla
 from pymetis import part_graph
 from sklearn.cluster import SpectralClustering
 
-from divi.interfaces import CircuitRunner
+from divi.backends import CircuitRunner
 from divi.qprog import QAOA, ProgramBatch
-from divi.qprog._qaoa import (
+from divi.qprog.algorithms._qaoa import (
     _SUPPORTED_INITIAL_STATES_LITERAL,
     GraphProblem,
+    GraphProblemTypes,
     draw_graph_solution_nodes,
 )
-
-from .optimizers import Optimizer
+from divi.qprog.optimizers import MonteCarloOptimizer, Optimizer, copy_optimizer
 
 AggregateFn = Callable[
     [list[int], str, nx.Graph | rx.PyGraph, dict[int, int]], list[int]
@@ -42,8 +40,45 @@ _MAXIMUM_AVAILABLE_QUBITS = 30
 
 @dataclass(frozen=True, eq=True)
 class PartitioningConfig:
-    max_n_nodes_per_cluster: Optional[int] = None
-    minimum_n_clusters: Optional[int] = None
+    """Configuration for graph partitioning algorithms.
+
+    This class defines the parameters and constraints for partitioning large graphs
+    into smaller subgraphs for quantum algorithm execution. It supports multiple
+    partitioning algorithms and allows specification of size constraints.
+
+    Attributes:
+        max_n_nodes_per_cluster: Maximum number of nodes allowed in each cluster.
+            If None, no upper limit is enforced. Must be a positive integer.
+        minimum_n_clusters: Minimum number of clusters to create. If None, no
+            lower limit is enforced. Must be a positive integer.
+        partitioning_algorithm: Algorithm to use for partitioning. Options are:
+            - "spectral": Spectral partitioning using Fiedler vector (default)
+            - "metis": METIS graph partitioning library
+            - "kernighan_lin": Kernighan-Lin algorithm
+
+    Note:
+        At least one of `max_n_nodes_per_cluster` or `minimum_n_clusters` must be
+        specified. Both constraints cannot be None.
+
+    Examples:
+        >>> # Partition into clusters of at most 10 nodes
+        >>> config = PartitioningConfig(max_n_nodes_per_cluster=10)
+
+        >>> # Create at least 5 clusters using METIS
+        >>> config = PartitioningConfig(
+        ...     minimum_n_clusters=5,
+        ...     partitioning_algorithm="metis"
+        ... )
+
+        >>> # Both constraints: clusters of max 8 nodes, min 3 clusters
+        >>> config = PartitioningConfig(
+        ...     max_n_nodes_per_cluster=8,
+        ...     minimum_n_clusters=3
+        ... )
+    """
+
+    max_n_nodes_per_cluster: int | None = None
+    minimum_n_clusters: int | None = None
     partitioning_algorithm: Literal["spectral", "metis", "kernighan_lin"] = "spectral"
 
     def __post_init__(self):
@@ -334,19 +369,43 @@ def _node_partition_graph(
     return tuple(graph for (_, _, graph) in subgraphs)
 
 
-def linear_aggregation(curr_solution, solution_bitstring, graph, reverse_index_maps):
-    for node in graph.nodes():
-        solution_index = reverse_index_maps[node]
-        curr_solution[solution_index] = int(solution_bitstring[node])
+def linear_aggregation(
+    curr_solution: Sequence[Literal[0] | Literal[1]],
+    subproblem_solution: set[int],
+    subproblem_reverse_index_map: dict[int, int],
+):
+    """Linearly combines a subproblem's solution into the main solution vector.
+
+    This function iterates through each node of subproblem's solution. For each node,
+    it uses the reverse index map to find its original index in the main graph,
+    setting it to 1 in the current global solution, potentially overwriting any
+    previous states.
+
+    Args:
+        curr_solution (Sequence[Literal[0] | Literal[1]]): The main solution
+            vector being aggregated, represented as a sequence of 0s and 1s.
+        subproblem_solution (Set[int]): A set containing the original indices of
+            the nodes that form the solution for the subproblem.
+        subproblem_reverse_index_map (dict[int, int]): A mapping from the
+            subgraph's internal node labels back to their original indices in
+            the main solution vector.
+
+    Returns:
+        The updated main solution vector.
+    """
+    for node in subproblem_solution:
+        curr_solution[subproblem_reverse_index_map[node]] = 1
 
     return curr_solution
 
 
-def domninance_aggregation(
-    curr_solution, solution_bitstring, graph, reverse_index_maps
+def dominance_aggregation(
+    curr_solution: Sequence[Literal[0] | Literal[1]],
+    subproblem_solution: set[int],
+    subproblem_reverse_index_map: dict[int, int],
 ):
-    for node in graph.nodes():
-        solution_index = reverse_index_maps[node]
+    for node in subproblem_solution:
+        original_index = subproblem_reverse_index_map[node]
 
         # Use existing assignment if dominant in previous solutions
         # (e.g., more 0s than 1s or vice versa)
@@ -354,12 +413,12 @@ def domninance_aggregation(
         count_1 = curr_solution.count(1)
 
         if (
-            (count_0 > count_1 and curr_solution[node] == 0)
-            or (count_1 > count_0 and curr_solution[node] == 1)
+            (count_0 > count_1 and curr_solution[original_index] == 0)
+            or (count_1 > count_0 and curr_solution[original_index] == 1)
             or (count_0 == count_1)
         ):
             # Assign based on QAOA if tie
-            curr_solution[solution_index] = int(solution_bitstring[node])
+            curr_solution[original_index] = 1
 
     return curr_solution
 
@@ -367,14 +426,14 @@ def domninance_aggregation(
 class GraphPartitioningQAOA(ProgramBatch):
     def __init__(
         self,
-        graph: nx.Graph | rx.PyGraph,
+        graph: GraphProblemTypes,
         graph_problem: GraphProblem,
         n_layers: int,
         backend: CircuitRunner,
         partitioning_config: PartitioningConfig,
         initial_state: _SUPPORTED_INITIAL_STATES_LITERAL = "Recommended",
         aggregate_fn: AggregateFn = linear_aggregation,
-        optimizer=Optimizer.MONTE_CARLO,
+        optimizer: Optimizer | None = None,
         max_iterations=10,
         **kwargs,
     ):
@@ -409,15 +468,18 @@ class GraphPartitioningQAOA(ProgramBatch):
         self.partitioning_config = partitioning_config
         self.max_iterations = max_iterations
 
+        self.solution = None
         self.aggregate_fn = aggregate_fn
 
-        self._solution_nodes = None
+        # Store the optimizer template (will be copied for each program)
+        self._optimizer_template = (
+            optimizer if optimizer is not None else MonteCarloOptimizer()
+        )
 
         self._constructor = partial(
             QAOA,
             initial_state=initial_state,
             graph_problem=graph_problem,
-            optimizer=optimizer,
             max_iterations=self.max_iterations,
             backend=self.backend,
             n_layers=n_layers,
@@ -425,11 +487,23 @@ class GraphPartitioningQAOA(ProgramBatch):
         )
 
     def create_programs(self):
-        if len(self.programs) > 0:
-            raise RuntimeError(
-                "Some programs already exist. "
-                "Clear the program dictionary before creating new ones by using batch.reset()."
-            )
+        """
+        Creates and initializes QAOA programs for each partitioned subgraph.
+
+        The main graph is partitioned into node-based subgraphs
+        according to the specified partitioning configuration. Each subgraph is relabeled with
+        integer node labels for QAOA compatibility, and a reverse index map is stored for later
+        result aggregation.
+
+        Each program is assigned a unique program ID, which is a tuple of:
+            - An uppercase letter (A, B, C, ...) corresponding to the partition index.
+            - The number of nodes in the subgraph.
+
+        Example program ID: ('A', 5) for the first partition with 5 nodes.
+
+        The created QAOA programs are stored in the `self.programs` dictionary, keyed by their program IDs.
+
+        """
 
         super().create_programs()
 
@@ -449,84 +523,48 @@ class GraphPartitioningQAOA(ProgramBatch):
             self.reverse_index_maps = {}
 
             for i, subgraph in enumerate(subgraphs):
-                index_map = {node: idx for idx, node in enumerate(subgraph.nodes())}
-                self.reverse_index_maps[i] = {v: k for k, v in index_map.items()}
-                _subgraph = nx.relabel_nodes(subgraph, index_map)
-
                 prog_id = (string.ascii_uppercase[i], subgraph.number_of_nodes())
 
-                self.programs[prog_id] = self._constructor(
-                    job_id=prog_id,
+                index_map = {node: idx for idx, node in enumerate(subgraph.nodes())}
+                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(
+                    program_id=prog_id,
                     problem=_subgraph,
-                    losses=self._manager.list(),
-                    probs=self._manager.dict(),
-                    final_params=self._manager.list(),
+                    optimizer=copy_optimizer(self._optimizer_template),
                     progress_queue=self._queue,
                 )
 
-    def compute_final_solutions(self):
-        if self._executor is not None:
-            self.wait_for_all()
-
-        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):
-        if len(self.programs) == 0:
-            raise RuntimeError("No programs to aggregate. Run create_programs() first.")
+        """
+        Aggregates the results from all QAOA subprograms to form a global solution.
 
-        if self._executor is not None:
-            self.wait_for_all()
+        This method collects the final bitstring solutions from each partitioned subgraph's QAOA program,
+        using the aggregation function specified at initialization (e.g., linear or dominance aggregation).
+        It reconstructs the global solution by mapping each subgraph's solution back to the original node indices
+        using the stored reverse index maps.
 
-        if any(len(program.losses) == 0 for program in self.programs.values()):
-            raise RuntimeError(
-                "Some/All programs have empty losses. Did you call run()?"
-            )
+        The final solution is stored in `self.solution` as a list of node indices assigned to the selected partition.
 
-        if any(len(program.probs) == 0 for program in self.programs.values()):
+        Raises:
+            RuntimeError: If no programs exist, if programs have not been run, or if results are incomplete.
+        Returns:
+            list[int]: The list of node indices in the final aggregated solution.
+        """
+        super().aggregate_results()
+
+        if any(len(program.best_probs) == 0 for program in self.programs.values()):
             raise RuntimeError(
-                "Not all final probabilities computed yet. Please call `compute_final_solutions()` first."
+                "Not all final probabilities computed yet. Please call `run()` first."
             )
 
         # Extract the solutions from each program
-        for program, reverse_index_maps in zip(
-            self.programs.values(), self.reverse_index_maps.values()
-        ):
-            # Extract the final probabilities of the lowest energy
-            last_iteration_losses = program.losses[-1]
-            minimum_key = min(last_iteration_losses, key=last_iteration_losses.get)
-
-            # Find the key matching the best_solution_idx with possible metadata in between
-            pattern = re.compile(rf"^{minimum_key}(?:_[^_]*)*_0$")
-            matching_keys = [k for k in program.probs.keys() if pattern.match(k)]
-
-            if len(matching_keys) > 1:
-                raise RuntimeError(f"More than one matching key found.")
-
-            best_solution_key = matching_keys[0]
-
-            minimum_probabilities = program.probs[best_solution_key]
-
-            # The bitstring corresponding to the solution, with flip for correct endianness
-            max_prob_key = max(minimum_probabilities, key=minimum_probabilities.get)[
-                ::-1
-            ]
-
+        for prog_id, program in self.programs.items():
             self._bitstring_solution = self.aggregate_fn(
                 self._bitstring_solution,
-                max_prob_key,
-                program.problem,
-                reverse_index_maps,
+                program.solution,
+                self.reverse_index_maps[prog_id],
             )
 
         self.solution = list(np.where(self._bitstring_solution)[0])
@@ -559,9 +597,9 @@ class GraphPartitioningQAOA(ProgramBatch):
 
         # Convert partitions to node-to-partition mapping
         node_to_partition = {}
-        for partition_id, mapping in self.reverse_index_maps.items():
+        for (partition_id, _), mapping in self.reverse_index_maps.items():
             for node in mapping.values():
-                node_to_partition[node] = string.ascii_uppercase[partition_id]
+                node_to_partition[node] = partition_id
 
         # Get unique partition IDs and create color map
         unique_partitions = sorted(list(set(node_to_partition.values())))
@@ -613,7 +651,13 @@ class GraphPartitioningQAOA(ProgramBatch):
         plt.show()
 
     def draw_solution(self):
-        if self._solution_nodes is None:
-            self.aggregate_results()
+        """
+        Visualizes the main graph with nodes highlighted according to the final aggregated solution.
+
+        If the solution has not yet been computed, this method calls `aggregate_results()` to obtain it.
+        """
+
+        if self.solution is None:
+            self.solution = self.aggregate_results()
 
         draw_graph_solution_nodes(self.main_graph, self.solution)

divi/qprog/workflows/_qubo_partitioning.py

@@ -0,0 +1,220 @@
+# SPDX-FileCopyrightText: 2025 Qoro Quantum Ltd <divi@qoroquantum.de>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import string
+from functools import partial
+from typing import TypeVar
+
+import dimod
+import hybrid
+import numpy as np
+import numpy.typing as npt
+import scipy.sparse as sps
+from dimod import BinaryQuadraticModel
+
+from divi.backends import CircuitRunner
+from divi.qprog.algorithms import QAOA, QUBOProblemTypes
+from divi.qprog.batch import ProgramBatch
+from divi.qprog.optimizers import MonteCarloOptimizer, Optimizer, copy_optimizer
+
+
+# Helper function to merge subsamples in-place
+def _merge_substates(_, substates):
+    a, b = substates
+    return a.updated(subsamples=hybrid.hstack_samplesets(a.subsamples, b.subsamples))
+
+
+T = TypeVar("T", bound=QUBOProblemTypes | BinaryQuadraticModel)
+
+
+def _sanitize_problem_input(qubo: T) -> tuple[T, BinaryQuadraticModel]:
+    if isinstance(qubo, BinaryQuadraticModel):
+        return qubo, qubo
+
+    if isinstance(qubo, (np.ndarray, sps.spmatrix)):
+        x, y = qubo.shape
+        if x != y:
+            raise ValueError("Only matrices supported.")
+
+    if isinstance(qubo, np.ndarray):
+        return qubo, dimod.BinaryQuadraticModel(qubo, vartype=dimod.Vartype.BINARY)
+
+    if isinstance(qubo, sps.spmatrix):
+        return qubo, dimod.BinaryQuadraticModel(
+            {(row, col): data for row, col, data in zip(qubo.row, qubo.col, qubo.data)},
+            vartype=dimod.Vartype.BINARY,
+        )
+
+    raise ValueError(f"Got an unsupported QUBO input format: {type(qubo)}")
+
+
+class QUBOPartitioningQAOA(ProgramBatch):
+    def __init__(
+        self,
+        qubo: QUBOProblemTypes,
+        decomposer: hybrid.traits.ProblemDecomposer,
+        n_layers: int,
+        backend: CircuitRunner,
+        composer: hybrid.traits.SubsamplesComposer = hybrid.SplatComposer(),
+        optimizer: Optimizer | None = None,
+        max_iterations: int = 10,
+        **kwargs,
+    ):
+        """
+        Initialize a QUBOPartitioningQAOA instance for solving QUBO problems using partitioning and QAOA.
+
+        Args:
+            qubo (QUBOProblemTypes): The QUBO problem to solve, provided as a supported type.
+                Note: Variable types are assumed to be binary (not Spin).
+            decomposer (hybrid.traits.ProblemDecomposer): The decomposer used to partition the QUBO problem into subproblems.
+            n_layers (int): Number of QAOA layers to use for each subproblem.
+            backend (CircuitRunner): Backend responsible for running quantum circuits.
+            composer (hybrid.traits.SubsamplesComposer, optional): Composer to aggregate subsamples from subproblems.
+                Defaults to hybrid.SplatComposer().
+            optimizer (Optimizer, optional): Optimizer to use for QAOA.
+                Defaults to Optimizer.MONTE_CARLO.
+            max_iterations (int, optional): Maximum number of optimization iterations.
+                Defaults to 10.
+            **kwargs: Additional keyword arguments passed to the QAOA constructor.
+
+        """
+        super().__init__(backend=backend)
+
+        self.main_qubo, self._bqm = _sanitize_problem_input(qubo)
+
+        self._partitioning = hybrid.Unwind(decomposer)
+        self._aggregating = hybrid.Reduce(hybrid.Lambda(_merge_substates)) | composer
+
+        self.max_iterations = max_iterations
+
+        self.trivial_program_ids = set()
+
+        # Store the optimizer template (will be copied for each program)
+        self._optimizer_template = (
+            optimizer if optimizer is not None else MonteCarloOptimizer()
+        )
+
+        self._constructor = partial(
+            QAOA,
+            max_iterations=self.max_iterations,
+            backend=self.backend,
+            n_layers=n_layers,
+            **kwargs,
+        )
+        pass
+
+    def create_programs(self):
+        """
+        Partition the main QUBO problem and instantiate QAOA programs for each subproblem.
+
+        This implementation:
+        - Uses the configured decomposer to split the main QUBO into subproblems.
+        - For each subproblem, creates a QAOA program with the specified parameters.
+        - Stores each program in `self.programs` with a unique identifier.
+
+        Unique Identifier Format:
+            Each key in `self.programs` is a tuple of the form (letter, size), where:
+                - letter: An uppercase letter ('A', 'B', 'C', ...) indicating the partition index.
+                - size: The number of variables in the subproblem.
+
+            Example: ('A', 5) refers to the first partition with 5 variables.
+        """
+
+        super().create_programs()
+
+        self.prog_id_to_bqm_subproblem_states = {}
+
+        init_state = hybrid.State.from_problem(self._bqm)
+        _bqm_partitions = self._partitioning.run(init_state).result()
+
+        for i, partition in enumerate(_bqm_partitions):
+            if i > 0:
+                # We only need 'problem' on the first partition since
+                # it will propagate to the other partitions during
+                # aggregation, otherwise it's a waste of memory
+                del partition["problem"]
+
+            prog_id = (string.ascii_uppercase[i], len(partition.subproblem))
+            self.prog_id_to_bqm_subproblem_states[prog_id] = partition
+
+            if partition.subproblem.num_interactions == 0:
+                # Skip creating a full QAOA program for this trivial case.
+                self.trivial_program_ids.add(prog_id)
+                continue
+
+            ldata, (irow, icol, qdata), _ = partition.subproblem.to_numpy_vectors(
+                partition.subproblem.variables
+            )
+
+            coo_mat = sps.coo_matrix(
+                (
+                    np.r_[ldata, qdata],
+                    (
+                        np.r_[np.arange(len(ldata)), icol],
+                        np.r_[np.arange(len(ldata)), irow],
+                    ),
+                ),
+                shape=(len(ldata), len(ldata)),
+            )
+
+            self._programs[prog_id] = self._constructor(
+                program_id=prog_id,
+                problem=coo_mat,
+                optimizer=copy_optimizer(self._optimizer_template),
+                progress_queue=self._queue,
+            )
+
+    def aggregate_results(self) -> tuple[npt.NDArray[np.int32], float]:
+        """
+        Aggregate results from all QUBO subproblems into a global solution.
+
+        Collects solutions from each partitioned subproblem (both QAOA-optimized and
+        trivial ones) and uses the hybrid framework composer to combine them into
+        a final solution for the original QUBO problem.
+
+        Returns:
+            tuple: A tuple containing:
+                - solution (npt.NDArray[np.int32]): Binary solution vector for the QUBO problem.
+                - solution_energy (float): Energy/cost of the solution.
+
+        Raises:
+            RuntimeError: If programs haven't been run or if final probabilities
+                haven't been computed.
+        """
+        super().aggregate_results()
+
+        if any(len(program.best_probs) == 0 for program in self.programs.values()):
+            raise RuntimeError(
+                "Not all final probabilities computed yet. Please call `run()` first."
+            )
+
+        for (
+            prog_id,
+            bqm_subproblem_state,
+        ) in self.prog_id_to_bqm_subproblem_states.items():
+
+            if prog_id in self.trivial_program_ids:
+                # Case 1: Trivial problem. Solve classically.
+                # The solution is any bitstring (e.g., all zeros) with energy 0.
+                var_to_val = {v: 0 for v in bqm_subproblem_state.subproblem.variables}
+            else:
+                subproblem = self._programs[prog_id]
+                var_to_val = dict(
+                    zip(bqm_subproblem_state.subproblem.variables, subproblem.solution)
+                )
+
+            sample_set = dimod.SampleSet.from_samples(
+                dimod.as_samples(var_to_val), "BINARY", 0
+            )
+
+            self.prog_id_to_bqm_subproblem_states[prog_id] = (
+                bqm_subproblem_state.updated(subsamples=sample_set)
+            )
+
+        states = hybrid.States(*list(self.prog_id_to_bqm_subproblem_states.values()))
+        final_state = self._aggregating.run(states).result()
+
+        self.solution, self.solution_energy, _ = final_state.samples.record[0]
+
+        return self.solution, self.solution_energy