Source code for divi.qprog.problems._graphs

# SPDX-FileCopyrightText: 2025-2026 Qoro Quantum Ltd <divi@qoroquantum.de>
#
# SPDX-License-Identifier: Apache-2.0

"""Graph problem classes for QAOA."""

from collections.abc import Callable, Hashable
from functools import cached_property
from typing import Any
from warnings import warn

import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from qiskit.quantum_info import SparsePauliOp

from divi.hamiltonians._term_ops import _clean_hamiltonian_spo
from divi.qprog import GraphProblemTypes
from divi.qprog.algorithms import (
    InitialState,
    OnesState,
    SuperpositionState,
    ZerosState,
)
from divi.qprog.problems import GraphPartitioningConfig, QAOAProblem
from divi.qprog.problems._graph_hamiltonians import (
    max_clique_hamiltonians,
    max_independent_set_hamiltonians,
    max_weight_cycle_hamiltonians,
    maxcut_hamiltonians,
    min_vertex_cover_hamiltonians,
)
from divi.qprog.problems._graph_partitioning_utils import _node_partition_graph


class _GraphProblemBase(QAOAProblem):
    """Shared logic for graph problems built directly from ``SparsePauliOp``.

    Subclasses set ``_resolver`` (a function returning ``(cost_spo, mixer_spo)``
    or ``(cost_spo, mixer_spo, metadata)``) and the two ``_*_state_cls`` class
    attributes, then call ``super().__init__``.
    """

    _resolver: staticmethod
    _constrained_state_cls: type[InitialState]
    _unconstrained_state_cls: type[InitialState]

    def __init__(
        self,
        graph: GraphProblemTypes,
        *,
        is_constrained: bool = True,
        config: GraphPartitioningConfig | None = None,
    ):
        self._graph = graph
        self._is_constrained = is_constrained

        cost_spo, self._mixer_hamiltonian, *self._metadata = self._resolve(
            graph, is_constrained
        )

        cleaned, ham_constant = _clean_hamiltonian_spo(cost_spo, raise_on_constant=True)

        self._cost_hamiltonian = cleaned
        self._loss_constant = ham_constant
        self._wire_labels = self._compute_wire_labels(graph)
        self._initial_state = (
            self._constrained_state_cls
            if is_constrained
            else self._unconstrained_state_cls
        )()
        self._config = config
        self._reverse_index_maps = {}

    @classmethod
    def _resolve(cls, graph, is_constrained):
        """Build cost/mixer SPOs for this problem type."""
        try:
            return cls._resolver(graph, constrained=is_constrained)
        except TypeError:
            return cls._resolver(graph)

    @staticmethod
    def _compute_wire_labels(graph: GraphProblemTypes) -> tuple:
        """Map qubit positions back to original node values in node-iteration order."""
        if isinstance(graph, nx.Graph):
            return tuple(graph.nodes())
        # rustworkx graph: edge_list() / node values; mirror the relabeling done
        # inside the SPO builders.
        return tuple(graph.nodes())

    @property
    def graph(self) -> GraphProblemTypes:
        """The underlying graph.

        Treat as read-only: the cost Hamiltonian is fixed at construction, and
        ``evaluate_global_solution`` caches the Pauli term list on first call.
        Mutating this graph (adding nodes/edges, changing weights) regenerates
        neither, so scores would go stale. Build a new problem instead.
        """
        return self._graph

    @property
    def cost_hamiltonian(self) -> SparsePauliOp:
        return self._cost_hamiltonian

    @property
    def mixer_hamiltonian(self) -> SparsePauliOp:
        return self._mixer_hamiltonian

    @property
    def wire_labels(self) -> tuple:
        return self._wire_labels

    @property
    def loss_constant(self) -> float:
        return self._loss_constant

    @property
    def recommended_initial_state(self) -> InitialState:
        return self._initial_state

    @property
    def decode_fn(self) -> Callable[[str], Any]:
        wires = self._wire_labels

        def _decode(bitstring: str) -> list:
            return [
                wires[idx]
                for idx, bit in enumerate(bitstring)
                if bit == "1" and idx < len(wires)
            ]

        return _decode

    @property
    def metadata(self) -> dict[str, Any]:
        return self._metadata[0] if self._metadata else {}

    def decompose(self) -> dict[Hashable, QAOAProblem]:
        if self._config is None:
            raise ValueError(
                "Cannot decompose: no config was provided at construction."
            )

        # Warn if this problem type has known partitioning risks
        tier = _PARTITIONING_COMPATIBILITY_TIERS.get(type(self))
        if tier is not None:
            risk_level, rationale = tier
            prefix = "High-risk" if risk_level == "high-risk" else "Heuristic-risk"
            detail = (
                "Aggregation is heuristic and may miss globally valid/high-quality "
                f"solutions because {rationale}"
                if risk_level == "high-risk"
                else "Results may be sensitive to partition boundaries because "
                f"{rationale}"
            )
            warn(
                f"{prefix} graph partitioning objective: "
                f"{type(self).__name__}. {detail}",
                UserWarning,
                stacklevel=2,
            )

        subgraphs = _node_partition_graph(
            self.graph,
            partitioning_config=self._config,
        )

        self._reverse_index_maps = {}
        sub_problems: dict[Hashable, QAOAProblem] = {}

        for i, (subgraph, cluster_ids) in enumerate(subgraphs):
            prog_id = (f"P{i}", len(subgraph))
            # ``cluster_ids[local_idx] == original_node_id``; the partitioner
            # has already relabeled each subgraph to ``0..M-1``.
            self._reverse_index_maps[prog_id] = dict(enumerate(cluster_ids))
            sub_problems[prog_id] = type(self)(
                subgraph, is_constrained=self._is_constrained
            )

        return sub_problems

    def initial_solution_size(self) -> int:
        return len(self.graph)

    def extend_solution(
        self,
        current_solution: list[int],
        prog_id: Hashable,
        candidate_decoded: list[int],
    ) -> list[int]:
        extended = list(current_solution)
        reverse_map = self._reverse_index_maps[prog_id]

        # Reset all positions belonging to this partition to 0
        for global_idx in reverse_map.values():
            extended[global_idx] = 0

        # Set positions for nodes in the candidate's decoded solution to 1
        for local_node in candidate_decoded:
            global_idx = reverse_map[local_node]
            extended[global_idx] = 1

        return extended

    @cached_property
    def _diagonal_terms(self) -> list[tuple[float, tuple[int, ...]]]:
        """``(coeff, z_qubit_indices)`` per cost-Hamiltonian term.

        Computed once from the (immutable) cost Hamiltonian so
        :meth:`evaluate_global_solution` need not rebuild Pauli labels on every
        call. Validates the Hamiltonian is diagonal (Z/I only) up front.
        """
        spo: SparsePauliOp = self.cost_hamiltonian
        terms: list[tuple[float, tuple[int, ...]]] = []
        for label, coeff in zip(spo.paulis.to_labels(), spo.coeffs):
            z_qubits = []
            for qubit, char in enumerate(reversed(label)):
                if char == "I":
                    continue
                if char != "Z":
                    raise ValueError(
                        f"Cost Hamiltonian contains non-diagonal term {label!r}; "
                        f"evaluate_global_solution requires Z-only operators."
                    )
                z_qubits.append(qubit)
            terms.append((float(np.real(coeff)), tuple(z_qubits)))
        return terms

    def evaluate_global_solution(self, solution: list[int]) -> float:
        energy = self.loss_constant
        for coeff, z_qubits in self._diagonal_terms:
            eigenvalue = 1.0
            for qubit in z_qubits:
                eigenvalue *= 1 - 2 * solution[qubit]
            energy += coeff * eigenvalue
        return energy

    def postprocess_candidates(
        self, candidates: list[tuple[float, list[int]]], *, strict: bool = False
    ) -> list[tuple[list[int], float]]:
        return [(list(np.where(solution)[0]), score) for score, solution in candidates]


[docs] class MaxCutProblem(_GraphProblemBase): """MaxCut problem on a graph. Args: graph: NetworkX or RustworkX graph. """ _resolver = staticmethod(maxcut_hamiltonians) # type: ignore[assignment, bad-override] _constrained_state_cls = SuperpositionState _unconstrained_state_cls = SuperpositionState
[docs] class MaxCliqueProblem(_GraphProblemBase): """Max clique problem on a graph. Args: graph: NetworkX or RustworkX graph. is_constrained: Use constrained mixer. Defaults to True. """ _resolver = staticmethod(max_clique_hamiltonians) # type: ignore[assignment, bad-override] _constrained_state_cls = ZerosState _unconstrained_state_cls = SuperpositionState
[docs] class MaxIndependentSetProblem(_GraphProblemBase): """Max independent set problem on a graph. Args: graph: NetworkX or RustworkX graph. is_constrained: Use constrained mixer. Defaults to True. """ _resolver = staticmethod(max_independent_set_hamiltonians) # type: ignore[assignment, bad-override] _constrained_state_cls = ZerosState _unconstrained_state_cls = SuperpositionState
[docs] class MinVertexCoverProblem(_GraphProblemBase): """Min vertex cover problem on a graph. Args: graph: NetworkX or RustworkX graph. is_constrained: Use constrained mixer. Defaults to True. """ _resolver = staticmethod(min_vertex_cover_hamiltonians) # type: ignore[assignment, bad-override] _constrained_state_cls = OnesState _unconstrained_state_cls = SuperpositionState
[docs] class MaxWeightCycleProblem(_GraphProblemBase): """Max weight cycle problem on a directed graph. Args: graph: NetworkX DiGraph or RustworkX PyDiGraph with weighted edges. is_constrained: Use cycle-mixer (preserves valid cycles). Defaults to True. """ _resolver = staticmethod(max_weight_cycle_hamiltonians) # type: ignore[assignment, bad-override] _constrained_state_cls = SuperpositionState _unconstrained_state_cls = SuperpositionState @staticmethod def _compute_wire_labels(graph: GraphProblemTypes) -> tuple: # Cycle problems use edge variables; wires are 0-indexed by edge count. if hasattr(graph, "number_of_edges"): return tuple(range(graph.number_of_edges())) return tuple(range(len(graph.edge_list()))) # type: ignore[attr-defined]
# Partitioning is most robust for cut-style objectives (e.g. MaxCut). # Structure-dependent objectives may lose cross-partition constraints. _PARTITIONING_COMPATIBILITY_TIERS = { MaxWeightCycleProblem: ( "high-risk", "partitioning can break cycles across cluster boundaries.", ), MaxCliqueProblem: ( "heuristic-risk", "partitioning can hide cross-partition adjacency needed for global cliques.", ), MaxIndependentSetProblem: ( "heuristic-risk", "partitioning can hide cross-partition conflicts between selected vertices.", ), MinVertexCoverProblem: ( "heuristic-risk", "partitioning can hide cross-partition edges that must be covered globally.", ), }
[docs] 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. """ node_colors = [ "red" if node in partition_nodes else "lightblue" for node in main_graph.nodes() ] plt.figure(figsize=(10, 8)) pos = nx.spring_layout(main_graph) nx.draw_networkx_nodes(main_graph, pos, node_color=node_colors, node_size=500) nx.draw_networkx_edges(main_graph, pos) nx.draw_networkx_labels(main_graph, pos, font_size=10, font_weight="bold") plt.axis("off") plt.tight_layout() plt.show()