qoro-divi 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

divi/qprog/algorithms/_custom_vqa.py

@@ -0,0 +1,263 @@
+# SPDX-FileCopyrightText: 2025-2026 Qoro Quantum Ltd <divi@qoroquantum.de>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+from warnings import warn
+
+import numpy as np
+import pennylane as qml
+import sympy as sp
+from qiskit import QuantumCircuit
+
+from divi.circuits import CircuitBundle, MetaCircuit
+from divi.qprog._hamiltonians import _clean_hamiltonian
+from divi.qprog.variational_quantum_algorithm import VariationalQuantumAlgorithm
+
+
+class CustomVQA(VariationalQuantumAlgorithm):
+    """Custom variational algorithm for a parameterized QuantumScript.
+
+    This implementation wraps a PennyLane QuantumScript (or converts a Qiskit
+    QuantumCircuit into one) and optimizes its trainable parameters to minimize
+    a single expectation-value measurement. Qiskit measurements are converted
+    into a PauliZ expectation on the measured wires. Parameters are bound to sympy
+    symbols to enable QASM substitution and reuse of MetaCircuit templates
+    during optimization.
+
+    Attributes:
+        qscript (qml.tape.QuantumScript): The parameterized QuantumScript.
+        param_shape (tuple[int, ...]): Shape of a single parameter set.
+        n_qubits (int): Number of qubits in the script.
+        n_layers (int): Layer count (fixed to 1 for this wrapper).
+        cost_hamiltonian (qml.operation.Operator): Observable being minimized.
+        loss_constant (float): Constant term extracted from the observable.
+        optimizer (Optimizer): Classical optimizer for parameter updates.
+        max_iterations (int): Maximum number of optimization iterations.
+        current_iteration (int): Current optimization iteration.
+    """
+
+    def __init__(
+        self,
+        qscript: qml.tape.QuantumScript | QuantumCircuit,
+        *,
+        param_shape: tuple[int, ...] | int | None = None,
+        max_iterations: int = 10,
+        **kwargs,
+    ) -> None:
+        """Initialize a CustomVQA instance.
+
+        Args:
+            qscript (qml.tape.QuantumScript | QuantumCircuit): A parameterized QuantumScript with a
+                single expectation-value measurement, or a Qiskit QuantumCircuit with
+                computational basis measurements.
+            param_shape (tuple[int, ...] | int | None): Shape of a single parameter
+                set. If None, uses a flat shape inferred from trainable parameters.
+            max_iterations (int): Maximum number of optimization iterations.
+            **kwargs: Additional keyword arguments passed to the parent class, including
+                backend and optimizer.
+
+        Raises:
+            TypeError: If qscript is not a supported PennyLane QuantumScript or Qiskit QuantumCircuit.
+            ValueError: If the script has an invalid measurement or no trainable parameters.
+        """
+        super().__init__(**kwargs)
+
+        self._qiskit_param_names = (
+            [param.name for param in qscript.parameters]
+            if isinstance(qscript, QuantumCircuit)
+            else None
+        )
+        self.qscript = self._coerce_to_quantum_script(qscript)
+
+        if len(self.qscript.measurements) != 1:
+            raise ValueError(
+                "QuantumScript must contain exactly one measurement for optimization."
+            )
+
+        measurement = self.qscript.measurements[0]
+        if not hasattr(measurement, "obs") or measurement.obs is None:
+            raise ValueError(
+                "QuantumScript must contain a single expectation-value measurement."
+            )
+
+        self._cost_hamiltonian, self.loss_constant = _clean_hamiltonian(measurement.obs)
+        if (
+            isinstance(self._cost_hamiltonian, qml.Hamiltonian)
+            and not self._cost_hamiltonian.operands
+        ):
+            raise ValueError("Hamiltonian contains only constant terms.")
+
+        self.n_qubits = self.qscript.num_wires
+        self.n_layers = 1
+        self.max_iterations = max_iterations
+        self.current_iteration = 0
+
+        trainable_param_indices = (
+            list(self.qscript.trainable_params)
+            if self.qscript.trainable_params
+            else list(range(len(self.qscript.get_parameters())))
+        )
+        if not trainable_param_indices:
+            raise ValueError("QuantumScript does not contain any trainable parameters.")
+
+        self._param_shape = self._resolve_param_shape(
+            param_shape, len(trainable_param_indices)
+        )
+        self._n_params = int(np.prod(self._param_shape))
+
+        self._trainable_param_indices = trainable_param_indices
+        self._param_symbols = (
+            np.array(
+                [sp.Symbol(name) for name in self._qiskit_param_names], dtype=object
+            ).reshape(self._param_shape)
+            if self._qiskit_param_names is not None
+            else sp.symarray("p", self._param_shape)
+        )
+
+        flat_symbols = self._param_symbols.flatten().tolist()
+        self._qscript = self.qscript.bind_new_parameters(
+            flat_symbols, self._trainable_param_indices
+        )
+
+    @property
+    def cost_hamiltonian(self) -> qml.operation.Operator:
+        """The cost Hamiltonian for the QuantumScript optimization."""
+        return self._cost_hamiltonian
+
+    @property
+    def param_shape(self) -> tuple[int, ...]:
+        """Shape of a single parameter set."""
+        return self._param_shape
+
+    def _resolve_param_shape(
+        self, param_shape: tuple[int, ...] | int | None, n_params: int
+    ) -> tuple[int, ...]:
+        """Validate and normalize the parameter shape.
+
+        Args:
+            param_shape (tuple[int, ...] | int | None): User-provided parameter shape.
+            n_params (int): Number of trainable parameters in the script.
+
+        Returns:
+            tuple[int, ...]: Normalized parameter shape.
+
+        Raises:
+            ValueError: If the shape is invalid or does not match n_params.
+        """
+        if param_shape is None:
+            return (n_params,)
+
+        param_shape = (param_shape,) if isinstance(param_shape, int) else param_shape
+
+        if any(dim <= 0 for dim in param_shape):
+            raise ValueError(
+                f"param_shape entries must be positive, got {param_shape}."
+            )
+
+        if int(np.prod(param_shape)) != n_params:
+            raise ValueError(
+                f"param_shape does not match the number of trainable parameters. "
+                f"Expected product {n_params}, got {int(np.prod(param_shape))}."
+            )
+
+        return tuple(param_shape)
+
+    def _coerce_to_quantum_script(
+        self,
+        qscript: qml.tape.QuantumScript | QuantumCircuit,
+    ) -> qml.tape.QuantumScript:
+        """Convert supported inputs into a PennyLane QuantumScript.
+
+        Args:
+            qscript (qml.tape.QuantumScript): Input QuantumScript or Qiskit QuantumCircuit.
+
+        Returns:
+            qml.tape.QuantumScript: The converted QuantumScript.
+
+        Raises:
+            TypeError: If the input type is unsupported.
+        """
+        if isinstance(qscript, qml.tape.QuantumScript):
+            return qscript
+
+        if isinstance(qscript, QuantumCircuit):
+            measured_wires = sorted(
+                {
+                    qscript.qubits.index(qubit)
+                    for instruction in qscript.data
+                    if instruction.operation.name == "measure"
+                    for qubit in instruction.qubits
+                }
+            )
+            if not measured_wires:
+                warn(
+                    "Provided QuantumCircuit has no measurement operations. "
+                    "Defaulting to measuring all wires with PauliZ.",
+                    UserWarning,
+                )
+                measured_wires = list(range(len(qscript.qubits)))
+
+            obs = (
+                qml.Z(measured_wires[0])
+                if len(measured_wires) == 1
+                else qml.sum(*(qml.Z(wire) for wire in measured_wires))
+            )
+            # Remove measurements before conversion to avoid MidMeasureMP issues
+            qc_no_measure = QuantumCircuit(qscript.num_qubits)
+            for instruction in qscript.data:
+                if instruction.operation.name != "measure":
+                    qc_no_measure.append(
+                        instruction.operation, instruction.qubits, instruction.clbits
+                    )
+            qfunc = qml.from_qiskit(qc_no_measure)
+            qiskit_params = [
+                qml.numpy.array(0.0, requires_grad=True) for _ in qscript.parameters
+            ]
+
+            def qfunc_with_measurement(*params):
+                qfunc(*params)
+                return qml.expval(obs)
+
+            return qml.tape.make_qscript(qfunc_with_measurement)(*qiskit_params)
+
+        raise TypeError(
+            "qscript must be a PennyLane QuantumScript or a Qiskit QuantumCircuit."
+        )
+
+    def _create_meta_circuits_dict(self) -> dict[str, MetaCircuit]:
+        """Create the meta-circuit dictionary for CustomVQA.
+
+        Returns:
+            dict[str, MetaCircuit]: Dictionary containing the cost circuit template.
+        """
+        return {
+            "cost_circuit": self._meta_circuit_factory(
+                self._qscript, symbols=self._param_symbols.flatten()
+            )
+        }
+
+    def _generate_circuits(self) -> list[CircuitBundle]:
+        """Generate circuits for the current parameter sets.
+
+        Returns:
+            list[CircuitBundle]: Circuit bundles tagged by parameter index.
+        """
+        return [
+            self.meta_circuits["cost_circuit"].initialize_circuit_from_params(
+                params_group, param_idx=p
+            )
+            for p, params_group in enumerate(self._curr_params)
+        ]
+
+    def _perform_final_computation(self, **kwargs) -> None:
+        """No-op by default for custom QuantumScript optimization."""
+        pass
+
+    def _save_subclass_state(self) -> dict[str, Any]:
+        """Save subclass-specific state for checkpointing."""
+        return {}
+
+    def _load_subclass_state(self, state: dict[str, Any]) -> None:
+        """Load subclass-specific state from a checkpoint."""
+        pass

divi/qprog/algorithms/_pce.py

@@ -0,0 +1,262 @@
+# SPDX-FileCopyrightText: 2026 Qoro Quantum Ltd <divi@qoroquantum.de>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from warnings import warn
+
+import numpy as np
+import numpy.typing as npt
+import pennylane as qml
+import sympy as sp
+
+from divi.circuits import MetaCircuit
+from divi.qprog.typing import QUBOProblemTypes, qubo_to_matrix
+
+from ._vqe import VQE
+
+# Pre-computed 8-bit popcount table for O(1) lookups
+_POPCOUNT_TABLE_8BIT = np.array([bin(i).count("1") for i in range(256)], dtype=np.uint8)
+
+
+def _fast_popcount_parity(arr_input: npt.NDArray[np.integer]) -> npt.NDArray[np.uint8]:
+    """
+    Vectorized calculation of (popcount % 2) for an array of integers.
+    Uses numpy view casting for extreme speed over large arrays.
+    """
+    # 1. Ensure array is uint64
+    arr_u64 = arr_input.astype(np.uint64)
+
+    # 2. View as bytes to use 8-bit lookup table
+    arr_bytes = arr_u64.view(np.uint8).reshape(arr_input.shape + (8,))
+
+    # 3. Lookup and sum bits
+    total_bits = _POPCOUNT_TABLE_8BIT[arr_bytes].sum(axis=-1)
+
+    # 4. Return Parity (0 or 1)
+    return total_bits % 2
+
+
+def _aggregate_param_group(
+    param_group: list[tuple[str, dict[str, int]]],
+    merge_counts_fn,
+) -> tuple[list[str], npt.NDArray[np.float64], float]:
+    """Aggregate a parameter group into states, counts, and total shots."""
+    shots_dict = merge_counts_fn(param_group)
+    state_strings = list(shots_dict.keys())
+    counts = np.array(list(shots_dict.values()), dtype=float)
+    total_shots = counts.sum()
+    return state_strings, counts, float(total_shots)
+
+
+def _decode_parities(
+    state_strings: list[str], variable_masks_u64: npt.NDArray[np.uint64]
+) -> npt.NDArray[np.uint8]:
+    """Decode bitstring parities using the precomputed variable masks."""
+    states = np.array([int(s, 2) for s in state_strings], dtype=np.uint64)
+    overlaps = variable_masks_u64[:, None] & states[None, :]
+    return _fast_popcount_parity(overlaps)
+
+
+def _compute_soft_energy(
+    parities: npt.NDArray[np.uint8],
+    probs: npt.NDArray[np.float64],
+    alpha: float,
+    qubo_matrix: npt.NDArray[np.float64] | np.ndarray,
+) -> float:
+    """Compute the relaxed (soft) QUBO energy from parity expectations."""
+    mean_parities = parities.dot(probs)
+    z_expectations = 1.0 - (2.0 * mean_parities)
+    x_soft = 0.5 * (1.0 + np.tanh(alpha * z_expectations))
+    Qx = qubo_matrix @ x_soft
+    return float(np.dot(x_soft, Qx))
+
+
+def _compute_hard_cvar_energy(
+    parities: npt.NDArray[np.uint8],
+    counts: npt.NDArray[np.float64],
+    total_shots: float,
+    qubo_matrix: npt.NDArray[np.float64] | np.ndarray,
+    alpha_cvar: float = 0.25,
+) -> float:
+    """Compute CVaR energy from sampled hard assignments."""
+    x_vals = 1.0 - parities.astype(float)
+    Qx = qubo_matrix @ x_vals
+    energies = np.einsum("ij,ij->j", x_vals, Qx)
+
+    sorted_indices = np.argsort(energies)
+    sorted_energies = energies[sorted_indices]
+    sorted_counts = counts[sorted_indices]
+
+    cutoff_count = int(np.ceil(alpha_cvar * total_shots))
+    accumulated_counts = np.cumsum(sorted_counts)
+    limit_idx = np.searchsorted(accumulated_counts, cutoff_count)
+
+    cvar_energy = 0.0
+    count_sum = 0
+    if limit_idx > 0:
+        cvar_energy += np.sum(sorted_energies[:limit_idx] * sorted_counts[:limit_idx])
+        count_sum += np.sum(sorted_counts[:limit_idx])
+
+    remaining = cutoff_count - count_sum
+    cvar_energy += sorted_energies[limit_idx] * remaining
+    return float(cvar_energy / cutoff_count)
+
+
+class PCE(VQE):
+    """
+    Generalized Pauli Correlation Encoding (PCE) VQE.
+
+    Encodes an N-variable QUBO into O(log2(N)) qubits by mapping each variable
+    to a parity (Pauli-Z correlation) of the measured bitstring. The algorithm
+    uses the measurement distribution to estimate these parities, applies a
+    smooth relaxation when `alpha` is small, and evaluates the classical QUBO
+    objective: E = x.T @ Q @ x. For larger `alpha`, it switches to a discrete
+    objective (CVaR over sampled energies) for harder convergence.
+    """
+
+    def __init__(
+        self,
+        qubo_matrix: QUBOProblemTypes,
+        n_qubits: int | None = None,
+        alpha: float = 2.0,
+        **kwargs,
+    ):
+        """
+        Args:
+            qubo_matrix (QUBOProblemTypes): The N x N matrix to minimize. Accepts
+                a dense array, sparse matrix, list, or BinaryQuadraticModel.
+            n_qubits (int | None): Optional override. Must be >= ceil(log2(N)).
+                Larger values increase circuit size without adding representational power.
+            alpha (float): Scaling factor for the tanh() activation. Higher = harder
+                binary constraints, Lower = smoother gradient.
+        """
+
+        self.qubo_matrix = qubo_to_matrix(qubo_matrix)
+        self.n_vars = self.qubo_matrix.shape[0]
+        self.alpha = alpha
+        self._use_soft_objective = self.alpha < 5.0
+        self._final_vector: npt.NDArray[np.integer] | None = None
+
+        if kwargs.get("qem_protocol") is not None:
+            raise ValueError("PCE does not currently support qem_protocol.")
+
+        # Calculate required qubits (Logarithmic Scaling)
+        min_qubits = int(np.ceil(np.log2(self.n_vars + 1)))
+        if n_qubits is not None and n_qubits < min_qubits:
+            raise ValueError(
+                "n_qubits must be >= ceil(log2(N + 1)) to represent all variables. "
+                f"Got n_qubits={n_qubits}, minimum={min_qubits}."
+            )
+        if n_qubits is not None and n_qubits > min_qubits:
+            warn(
+                "n_qubits exceeds the minimum required; extra qubits increase circuit "
+                "size and can add noise without representing more variables.",
+                UserWarning,
+            )
+        self.n_qubits = n_qubits if n_qubits is not None else min_qubits
+
+        # Pre-compute U64 masks for the fast broadcasting step later
+        self._variable_masks_u64 = np.arange(1, self.n_vars + 1, dtype=np.uint64)
+
+        # Placeholder Hamiltonian required by VQE; we care about the measurement
+        # probability distribution, and Z-basis measurements provide it.
+        placeholder_hamiltonian = qml.Hamiltonian(
+            [1.0] * self.n_qubits, [qml.PauliZ(i) for i in range(self.n_qubits)]
+        )
+        super().__init__(hamiltonian=placeholder_hamiltonian, **kwargs)
+
+    def _create_meta_circuits_dict(self) -> dict[str, MetaCircuit]:
+        """Create meta circuits, handling the edge case of zero parameters."""
+        n_params = self.ansatz.n_params_per_layer(
+            self.n_qubits, n_electrons=self.n_electrons
+        )
+
+        weights_syms = sp.symarray("w", (self.n_layers, n_params))
+
+        ops = self.ansatz.build(
+            weights_syms,
+            n_qubits=self.n_qubits,
+            n_layers=self.n_layers,
+            n_electrons=self.n_electrons,
+        )
+
+        return {
+            "cost_circuit": self._meta_circuit_factory(
+                qml.tape.QuantumScript(
+                    ops=ops, measurements=[qml.expval(self._cost_hamiltonian)]
+                ),
+                symbols=weights_syms.flatten(),
+            ),
+            "meas_circuit": self._meta_circuit_factory(
+                qml.tape.QuantumScript(ops=ops, measurements=[qml.probs()]),
+                symbols=weights_syms.flatten(),
+                grouping_strategy="wires",
+            ),
+        }
+
+    def _post_process_results(
+        self, results: dict[str, dict[str, int]]
+    ) -> dict[int, float]:
+        """
+        Calculates loss.
+        If self.alpha < 5.0, computes 'Soft Energy' (Relaxed VQE) for smooth gradients.
+        If self.alpha >= 5.0, computes 'Hard CVaR Energy' for final convergence.
+        """
+
+        # Return raw probabilities if requested (skip processing)
+        if getattr(self, "_is_compute_probabilities", False):
+            return super()._post_process_results(results)
+
+        losses = {}
+
+        for p_idx, qem_groups in self._group_results(results).items():
+            # PCE ignores QEM ids; aggregate all shots for this parameter set.
+            param_group = [
+                ("0", shots)
+                for shots_list in qem_groups.values()
+                for shots in shots_list
+            ]
+
+            state_strings, counts, total_shots = _aggregate_param_group(
+                param_group, self._merge_param_group_counts
+            )
+
+            parities = _decode_parities(state_strings, self._variable_masks_u64)
+            if self._use_soft_objective:
+                probs = counts / total_shots
+                losses[p_idx] = _compute_soft_energy(
+                    parities, probs, self.alpha, self.qubo_matrix
+                )
+            else:
+                losses[p_idx] = _compute_hard_cvar_energy(
+                    parities, counts, total_shots, self.qubo_matrix
+                )
+
+        return losses
+
+    def _perform_final_computation(self, **kwargs) -> None:
+        """Compute the final eigenstate and decode it into a PCE vector."""
+        super()._perform_final_computation(**kwargs)
+
+        if self._eigenstate is None:
+            self._final_vector = None
+            return
+
+        best_bitstring = "".join(str(x) for x in self._eigenstate)
+        state_int = int(best_bitstring, 2)
+        state_u64 = np.array([state_int], dtype=np.uint64)
+
+        overlaps = self._variable_masks_u64[:, None] & state_u64[None, :]
+        parities = _fast_popcount_parity(overlaps).flatten()
+        self._final_vector = 1 - parities
+
+    @property
+    def solution(self) -> npt.NDArray[np.integer]:
+        """
+        Returns the final optimized vector (hard binary 0/1) based on the best parameters found.
+        You must run .run() before calling this.
+        """
+        if self._final_vector is None:
+            raise RuntimeError("Run the VQE optimization first.")
+
+        return self._final_vector

divi/qprog/algorithms/_qaoa.py

@@ -1,8 +1,9 @@
-# SPDX-FileCopyrightText: 2025 Qoro Quantum Ltd <divi@qoroquantum.de>
+# SPDX-FileCopyrightText: 2025-2026 Qoro Quantum Ltd <divi@qoroquantum.de>
 #
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
+from collections.abc import Callable
 from enum import Enum
 from typing import Any, Literal, get_args
 from warnings import warn
@@ -13,8 +14,6 @@ import networkx as nx
 import numpy as np
 import pennylane as qml
 import pennylane.qaoa as pqaoa
-import rustworkx as rx
-import scipy.sparse as sps
 import sympy as sp
 
 from divi.circuits import CircuitBundle, MetaCircuit
@@ -22,13 +21,11 @@ from divi.qprog._hamiltonians import (
     _clean_hamiltonian,
     convert_qubo_matrix_to_pennylane_ising,
 )
+from divi.qprog.typing import GraphProblemTypes, QUBOProblemTypes, qubo_to_matrix
 from divi.qprog.variational_quantum_algorithm import VariationalQuantumAlgorithm
 
 logger = logging.getLogger(__name__)
 
-GraphProblemTypes = nx.Graph | rx.PyGraph
-QUBOProblemTypes = list | np.ndarray | sps.spmatrix | dimod.BinaryQuadraticModel
-
 
 def _extract_loss_constant(
     problem_metadata: dict, constant_from_hamiltonian: float
@@ -163,17 +160,7 @@ def _resolve_circuit_layers(
 
         return *getattr(pqaoa, graph_problem.pl_string)(*params), resolved_initial_state
     else:
-        # Convert BinaryQuadraticModel to matrix if needed
-        if isinstance(problem, dimod.BinaryQuadraticModel):
-            # Manual conversion from BQM to matrix (replacing deprecated to_numpy_matrix)
-            variables = list(problem.variables)
-            var_to_idx = {v: i for i, v in enumerate(variables)}
-            qubo_matrix = np.diag([problem.linear.get(v, 0) for v in variables])
-            for (u, v), coeff in problem.quadratic.items():
-                i, j = var_to_idx[u], var_to_idx[v]
-                qubo_matrix[i, j] = qubo_matrix[j, i] = coeff
-        else:
-            qubo_matrix = problem
+        qubo_matrix = qubo_to_matrix(problem)
 
         cost_hamiltonian, constant = convert_qubo_matrix_to_pennylane_ising(qubo_matrix)
 
@@ -225,6 +212,7 @@ class QAOA(VariationalQuantumAlgorithm):
         n_layers: int = 1,
         initial_state: _SUPPORTED_INITIAL_STATES_LITERAL = "Recommended",
         max_iterations: int = 10,
+        decode_solution_fn: Callable[[str], Any] | None = None,
         **kwargs,
     ):
         """Initialize the QAOA problem.
@@ -237,15 +225,21 @@ class QAOA(VariationalQuantumAlgorithm):
             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".
             max_iterations (int): Maximum number of optimization iterations. Defaults to 10.
+            decode_solution_fn (callable[[str], Any] | None): Optional decoder for bitstrings.
+                If not provided, a default decoder is selected based on problem type.
             **kwargs: Additional keyword arguments passed to the parent class, including `optimizer`.
         """
-        super().__init__(**kwargs)
-
         self.graph_problem = graph_problem
 
-        # Validate and process problem
+        # Validate and process problem (needed to determine decode function)
+        # This sets n_qubits which is needed before parent init
         self.problem = self._validate_and_set_problem(problem, graph_problem)
 
+        if decode_solution_fn is not None:
+            kwargs["decode_solution_fn"] = decode_solution_fn
+
+        super().__init__(**kwargs)
+
         # Validate initial state
         if initial_state not in get_args(_SUPPORTED_INITIAL_STATES_LITERAL):
             raise ValueError(
@@ -286,6 +280,22 @@ class QAOA(VariationalQuantumAlgorithm):
         # Extract wire labels from the cost Hamiltonian to ensure consistency
         self._circuit_wires = tuple(self._cost_hamiltonian.wires)
 
+        # Set up decode function based on problem type if user didn't provide one
+        if decode_solution_fn is None:
+            if isinstance(self.problem, QUBOProblemTypes):
+                # For QUBO: convert bitstring to numpy array of int32
+                self._decode_solution_fn = lambda bitstring: np.fromiter(
+                    bitstring, dtype=np.int32
+                )
+            elif isinstance(self.problem, GraphProblemTypes):
+                # For Graph: map bitstring positions to graph node labels
+                circuit_wires = self._circuit_wires  # Capture for closure
+                self._decode_solution_fn = lambda bitstring: [
+                    circuit_wires[idx]
+                    for idx, bit in enumerate(bitstring)
+                    if bit == "1" and idx < len(circuit_wires)
+                ]
+
     def _save_subclass_state(self) -> dict[str, Any]:
         """Save QAOA-specific runtime state."""
         return {
@@ -497,7 +507,7 @@ class QAOA(VariationalQuantumAlgorithm):
 
         return [
             self.meta_circuits[circuit_type].initialize_circuit_from_params(
-                params_group, tag_prefix=f"{p}"
+                params_group, param_idx=p
             )
             for p, params_group in enumerate(self._curr_params)
         ]
@@ -508,9 +518,11 @@ class QAOA(VariationalQuantumAlgorithm):
         This method performs the following steps:
         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.
+        3. Uses the `decode_solution_fn` (configured in constructor based on problem type) to decode
+           the bitstring into the appropriate format:
+           - For QUBO problems: NumPy array of bits (int32).
+           - For graph problems: List of node indices corresponding to '1's in the bitstring.
+        4. Stores the decoded solution in the appropriate attribute.
 
         Returns:
             tuple[int, float]: A tuple containing:
@@ -529,19 +541,14 @@ class QAOA(VariationalQuantumAlgorithm):
             best_measurement_probs, key=best_measurement_probs.get
         )
 
-        if isinstance(self.problem, QUBOProblemTypes):
-            self._solution_bitstring[:] = np.fromiter(
-                best_solution_bitstring, dtype=np.int32
-            )
+        # Use decode function to get the decoded solution
+        decoded_solution = self._decode_solution_fn(best_solution_bitstring)
 
-        if isinstance(self.problem, GraphProblemTypes):
-            # Map bitstring positions to actual graph node labels
-            # Bitstring is already endianness-corrected, so positions map directly to circuit_wires
-            self._solution_nodes[:] = [
-                self._circuit_wires[idx]
-                for idx, bit in enumerate(best_solution_bitstring)
-                if bit == "1" and idx < len(self._circuit_wires)
-            ]
+        # Store in appropriate attribute based on problem type
+        if isinstance(self.problem, QUBOProblemTypes):
+            self._solution_bitstring[:] = decoded_solution
+        elif isinstance(self.problem, GraphProblemTypes):
+            self._solution_nodes[:] = decoded_solution
 
         self.reporter.info(message="🏁 Computed Final Solution! 🏁")