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

divi/qprog/workflows/_graph_partitioning.py

@@ -4,8 +4,7 @@
 
 import heapq
 import string
-from collections.abc import Callable, Sequence, Set
-from concurrent.futures import ProcessPoolExecutor
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from functools import partial
 from typing import Literal
@@ -21,7 +20,7 @@ from pymetis import part_graph
 from sklearn.cluster import SpectralClustering
 
 from divi.backends import CircuitRunner
-from divi.qprog import QAOA, ProgramBatch, QuantumProgram
+from divi.qprog import QAOA, ProgramBatch
 from divi.qprog.algorithms._qaoa import (
     _SUPPORTED_INITIAL_STATES_LITERAL,
     GraphProblem,
@@ -41,6 +40,43 @@ _MAXIMUM_AVAILABLE_QUBITS = 30
 
 @dataclass(frozen=True, eq=True)
 class PartitioningConfig:
+    """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"
@@ -335,7 +371,7 @@ def _node_partition_graph(
 
 def linear_aggregation(
     curr_solution: Sequence[Literal[0] | Literal[1]],
-    subproblem_solution: Set[int],
+    subproblem_solution: set[int],
     subproblem_reverse_index_map: dict[int, int],
 ):
     """Linearly combines a subproblem's solution into the main solution vector.
@@ -365,7 +401,7 @@ def linear_aggregation(
 
 def dominance_aggregation(
     curr_solution: Sequence[Literal[0] | Literal[1]],
-    subproblem_solution: Set[int],
+    subproblem_solution: set[int],
     subproblem_reverse_index_map: dict[int, int],
 ):
     for node in subproblem_solution:
@@ -387,14 +423,6 @@ def dominance_aggregation(
     return curr_solution
 
 
-def _run_and_compute_solution(program: QuantumProgram):
-    program.run()
-
-    final_sol_circuit_count, final_sol_run_time = program.compute_final_solution()
-
-    return final_sol_circuit_count, final_sol_run_time
-
-
 class GraphPartitioningQAOA(ProgramBatch):
     def __init__(
         self,
@@ -443,8 +471,6 @@ class GraphPartitioningQAOA(ProgramBatch):
         self.solution = None
         self.aggregate_fn = aggregate_fn
 
-        self._task_fn = _run_and_compute_solution
-
         self._constructor = partial(
             QAOA,
             initial_state=initial_state,
@@ -499,33 +525,12 @@ class GraphPartitioningQAOA(ProgramBatch):
                 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(
+                self._programs[prog_id] = self._constructor(
                     job_id=prog_id,
                     problem=_subgraph,
-                    losses=self._manager.list(),
-                    probs=self._manager.dict(),
-                    final_params=self._manager.list(),
-                    solution_nodes=self._manager.list(),
                     progress_queue=self._queue,
                 )
 
-    def compute_final_solutions(self):
-        if self._executor is not None:
-            self.join()
-
-        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):
         """
         Aggregates the results from all QAOA subprograms to form a global solution.
@@ -544,7 +549,7 @@ class GraphPartitioningQAOA(ProgramBatch):
         """
         super().aggregate_results()
 
-        if any(len(program.probs) == 0 for program in self.programs.values()):
+        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."
             )

divi/qprog/workflows/_qubo_partitioning.py

@@ -49,14 +49,6 @@ def _sanitize_problem_input(qubo: T) -> tuple[T, BinaryQuadraticModel]:
     raise ValueError(f"Got an unsupported QUBO input format: {type(qubo)}")
 
 
-def _run_and_compute_solution(program: QuantumProgram):
-    program.run()
-
-    final_sol_circuit_count, final_sol_run_time = program.compute_final_solution()
-
-    return final_sol_circuit_count, final_sol_run_time
-
-
 class QUBOPartitioningQAOA(ProgramBatch):
     def __init__(
         self,
@@ -94,10 +86,10 @@ class QUBOPartitioningQAOA(ProgramBatch):
         self._partitioning = hybrid.Unwind(decomposer)
         self._aggregating = hybrid.Reduce(hybrid.Lambda(_merge_substates)) | composer
 
-        self._task_fn = _run_and_compute_solution
-
         self.max_iterations = max_iterations
 
+        self.trivial_program_ids = set()
+
         self._constructor = partial(
             QAOA,
             optimizer=optimizer if optimizer is not None else MonteCarloOptimizer(),
@@ -140,6 +132,12 @@ class QUBOPartitioningQAOA(ProgramBatch):
                 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
@@ -155,33 +153,52 @@ class QUBOPartitioningQAOA(ProgramBatch):
                 ),
                 shape=(len(ldata), len(ldata)),
             )
-            self.prog_id_to_bqm_subproblem_states[prog_id] = partition
-            self.programs[prog_id] = self._constructor(
+
+            self._programs[prog_id] = self._constructor(
                 job_id=prog_id,
                 problem=coo_mat,
-                losses=self._manager.list(),
-                probs=self._manager.dict(),
-                final_params=self._manager.list(),
-                solution_bitstring=self._manager.list(),
                 progress_queue=self._queue,
             )
 
     def aggregate_results(self):
+        """
+        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 (np.ndarray): 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.probs) == 0 for program in self.programs.values()):
+        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, subproblem in self.programs.items():
-            bqm_subproblem_state = self.prog_id_to_bqm_subproblem_states[prog_id]
+        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)
+                )
 
-            curr_final_solution = subproblem.solution
-
-            var_to_val = dict(
-                zip(bqm_subproblem_state.subproblem.variables, curr_final_solution)
-            )
             sample_set = dimod.SampleSet.from_samples(
                 dimod.as_samples(var_to_val), "BINARY", 0
             )

divi/qprog/workflows/_vqe_sweep.py

@@ -127,6 +127,13 @@ def _zmatrix_to_cartesian(z_matrix: list[_ZMatrixEntry]) -> np.ndarray:
     if n_atoms == 0:
         return coords
 
+    # Validate bond lengths are positive
+    for i, entry in enumerate(z_matrix[1:], start=1):
+        if entry.bond_length is not None and entry.bond_length <= 0:
+            raise ValueError(
+                f"Bond length for atom {i} must be positive, got {entry.bond_length}"
+            )
+
     # --- First atom at origin ---
     coords[0] = np.array([0.0, 0.0, 0.0])
 
@@ -251,12 +258,14 @@ def _kabsch_align(P_in: np.ndarray, Q_in: np.ndarray, reference_atoms_idx=slice(
     H = P_centered.T @ Q_centered
     U, _, Vt = np.linalg.svd(H)
 
-    # Reflection check
-    d = np.sign(np.linalg.det(Vt.T @ U.T))
-    D = np.diag([1] * (P.shape[1] - 1) + [d])
+    # Compute rotation matrix
+    R = Vt.T @ U.T
 
-    # Optimal rotation and translation
-    R = Vt.T @ D @ U.T
+    # Ensure proper rotation (det = +1) by handling reflections
+    if np.linalg.det(R) < 0:
+        # Flip the last column of Vt to ensure proper rotation
+        Vt[-1, :] *= -1
+        R = Vt.T @ U.T
     t = Qc - Pc @ R
 
     # Apply transformation
@@ -270,27 +279,28 @@ def _kabsch_align(P_in: np.ndarray, Q_in: np.ndarray, reference_atoms_idx=slice(
 @dataclass(frozen=True, eq=True)
 class MoleculeTransformer:
     """
-    base_molecule: qml.qchem.Molecule
-        The reference molecule used as a template for generating variants.
-    bond_modifiers: Sequence[float]
-        A list of values used to adjust bond lengths. The class will generate
-        **one new molecule for each modifier** in this list. The modification
-        mode is detected automatically:
+    A class for transforming molecular structures by modifying bond lengths.
+
+    This class generates variants of a base molecule by adjusting bond lengths
+    according to specified modifiers. The modification mode is detected automatically.
+
+    Attributes:
+        base_molecule (qml.qchem.Molecule): The reference molecule used as a template for generating variants.
+        bond_modifiers (Sequence[float]): A list of values used to adjust bond lengths. The class will generate
+            **one new molecule for each modifier** in this list. The modification
+            mode is detected automatically:
             - **Scale mode**: If all values are positive, they are used as scaling
             factors (e.g., 1.1 for a 10% increase).
             - **Delta mode**: If any value is zero or negative, all values are
             treated as additive changes to the bond length, in Ångstroms.
-    atom_connectivity: Sequence[tuple[int, int]] | None
-        A sequence of atom index pairs specifying the bonds in the molecule.
-        If not provided, a chain structure will be assumed
-        e.g.: `[(0, 1), (1, 2), (2, 3), ...]`.
-    bonds_to_transform: Sequence[tuple[int, int]] | None
-        A subset of `atom_connectivity` that specifies the bonds to modify.
-        If None, all bonds will be transformed.
-    alignment_atoms: Sequence[int] | None
-        Indices of atoms onto which to align the orientation of the resulting
-        variants of the molecule. Only useful for visualization and debuggin.
-        If None, no alignment is carried out.
+        atom_connectivity (Sequence[tuple[int, int]] | None): A sequence of atom index pairs specifying the bonds in the molecule.
+            If not provided, a chain structure will be assumed
+            e.g.: `[(0, 1), (1, 2), (2, 3), ...]`.
+        bonds_to_transform (Sequence[tuple[int, int]] | None): A subset of `atom_connectivity` that specifies the bonds to modify.
+            If None, all bonds will be transformed.
+        alignment_atoms (Sequence[int] | None): Indices of atoms onto which to align the orientation of the resulting
+            variants of the molecule. Only useful for visualization and debugging.
+            If None, no alignment is carried out.
     """
 
     base_molecule: qml.qchem.Molecule
@@ -403,7 +413,7 @@ class VQEHyperparameterSweep(ProgramBatch):
 
         Parameters
         ----------
-        ansatze: Sequence[VQEAnsatz]
+        ansatze: Sequence[Ansatz]
             A sequence of ansatz circuits to test.
         molecule_transformer: MoleculeTransformer
             A `MoleculeTransformer` object defining the configuration for
@@ -430,6 +440,15 @@ class VQEHyperparameterSweep(ProgramBatch):
         )
 
     def create_programs(self):
+        """
+        Create VQE programs for all combinations of ansätze and molecule variants.
+
+        Generates molecule variants using the configured MoleculeTransformer, then
+        creates a VQE program for each (ansatz, molecule_variant) pair.
+
+        Note:
+            Program IDs are tuples of (ansatz_name, bond_modifier_value).
+        """
         super().create_programs()
 
         self.molecule_variants = self.molecule_transformer.generate()
@@ -437,23 +456,35 @@ class VQEHyperparameterSweep(ProgramBatch):
         for ansatz, (modifier, molecule) in product(
             self.ansatze, self.molecule_variants.items()
         ):
-            _job_id = (ansatz, modifier)
-            self.programs[_job_id] = self._constructor(
+            _job_id = (ansatz.name, modifier)
+            self._programs[_job_id] = self._constructor(
                 job_id=_job_id,
                 molecule=molecule,
                 ansatz=ansatz,
-                losses=self._manager.list(),
-                final_params=self._manager.list(),
                 progress_queue=self._queue,
             )
 
     def aggregate_results(self):
+        """
+        Find the best ansatz and bond configuration from all VQE runs.
+
+        Compares the final energies across all ansatz/molecule combinations
+        and returns the configuration that achieved the lowest ground state energy.
+
+        Returns:
+            tuple: A tuple containing:
+                - best_config (tuple): (ansatz_name, bond_modifier) of the best result.
+                - best_energy (float): The lowest energy achieved.
+
+        Raises:
+            RuntimeError: If programs haven't been run or have empty losses.
+        """
         super().aggregate_results()
 
-        all_energies = {key: prog.losses[-1] for key, prog in self.programs.items()}
+        all_energies = {key: prog.best_loss for key, prog in self.programs.items()}
 
-        smallest_key = min(all_energies, key=lambda k: min(all_energies[k].values()))
-        smallest_value = min(all_energies[smallest_key].values())
+        smallest_key = min(all_energies, key=lambda k: all_energies[k])
+        smallest_value = all_energies[smallest_key]
 
         return smallest_key, smallest_value
 
@@ -483,18 +514,17 @@ class VQEHyperparameterSweep(ProgramBatch):
             # Plot each ansatz's results as a separate series for clarity
             for ansatz in unique_ansatze:
                 modifiers = []
-                min_energies = []
+                energies = []
                 for modifier in self.molecule_transformer.bond_modifiers:
-                    program_key = (ansatz, modifier)
-                    if program_key in self.programs:
+                    program_key = (ansatz.name, modifier)
+                    if program_key in self._programs:
                         modifiers.append(modifier)
-                        curr_energies = self.programs[program_key].losses[-1]
-                        min_energies.append(min(curr_energies.values()))
+                        energies.append(self._programs[program_key].best_loss)
 
                 # Use the new .name property for the label and the color_map
                 plt.scatter(
                     modifiers,
-                    min_energies,
+                    energies,
                     color=color_map[ansatz],
                     label=ansatz.name,
                 )
@@ -503,9 +533,7 @@ class VQEHyperparameterSweep(ProgramBatch):
             for ansatz in unique_ansatze:
                 energies = []
                 for modifier in self.molecule_transformer.bond_modifiers:
-                    energies.append(
-                        min(self.programs[(ansatz, modifier)].losses[-1].values())
-                    )
+                    energies.append(self._programs[(ansatz.name, modifier)].best_loss)
 
                 plt.plot(
                     self.molecule_transformer.bond_modifiers,

divi/reporting/_pbar.py

@@ -9,11 +9,26 @@ from rich.progress import (
     ProgressColumn,
     SpinnerColumn,
     TextColumn,
+    TimeElapsedColumn,
 )
 from rich.text import Text
 
 
+class _UnfinishedTaskWrapper:
+    """Wrapper that forces a task to appear unfinished for spinner animation."""
+
+    def __init__(self, task):
+        self._task = task
+
+    def __getattr__(self, name):
+        if name == "finished":
+            return False
+        return getattr(self._task, name)
+
+
 class ConditionalSpinnerColumn(ProgressColumn):
+    _FINAL_STATUSES = ("Success", "Failed", "Cancelled", "Aborted")
+
     def __init__(self):
         super().__init__()
         self.spinner = SpinnerColumn("point")
@@ -21,10 +36,11 @@ class ConditionalSpinnerColumn(ProgressColumn):
     def render(self, task):
         status = task.fields.get("final_status")
 
-        if status in ("Success", "Failed"):
+        if status in self._FINAL_STATUSES:
             return Text("")
 
-        return self.spinner.render(task)
+        # Force the task to appear unfinished for spinner animation
+        return self.spinner.render(_UnfinishedTaskWrapper(task))
 
 
 class PhaseStatusColumn(ProgressColumn):
@@ -38,29 +54,55 @@ class PhaseStatusColumn(ProgressColumn):
             return Text("• Success! ✅", style="bold green")
         elif final_status == "Failed":
             return Text("• Failed! ❌", style="bold red")
+        elif final_status == "Cancelled":
+            return Text("• Cancelled ⏹️", style="bold yellow")
+        elif final_status == "Aborted":
+            return Text("• Aborted ⚠️", style="dim magenta")
 
         message = task.fields.get("message")
 
-        poll_attempt = task.fields.get("poll_attempt")
+        poll_attempt = task.fields.get("poll_attempt", 0)
         polling_str = ""
-        service_job_id = ""
-        if poll_attempt > 0:
-            max_retries = task.fields.get("max_retries")
-            service_job_id = task.fields.get("service_job_id").split("-")[0]
+        service_job_id = task.fields.get("service_job_id")
+
+        if service_job_id:
+            split_job_id = service_job_id.split("-")[0]
             job_status = task.fields.get("job_status")
-            polling_str = f" [Job {service_job_id} is {job_status}. Polling attempt {poll_attempt} / {max_retries}]"
+
+            if job_status == "COMPLETED":
+                polling_str = f" [Job {split_job_id} is complete.]"
+            elif poll_attempt > 0:
+                max_retries = task.fields.get("max_retries")
+                polling_str = f" [Job {split_job_id} is {job_status}. Polling attempt {poll_attempt} / {max_retries}]"
 
         final_text = Text(f"[{message}]{polling_str}")
-        final_text.highlight_words([service_job_id], "blue")
+        if service_job_id:
+            final_text.highlight_words([split_job_id], "blue")
 
         return final_text
 
 
 def make_progress_bar(is_jupyter: bool = False) -> Progress:
+    """
+    Create a customized Rich progress bar for tracking quantum program execution.
+
+    Builds a progress bar with custom columns including job name, completion status,
+    elapsed time, spinner, and phase status indicators. Automatically adapts refresh
+    behavior for Jupyter notebook environments.
+
+    Args:
+        is_jupyter (bool, optional): Whether the progress bar is being displayed in
+            a Jupyter notebook environment. Affects refresh behavior. Defaults to False.
+
+    Returns:
+        Progress: A configured Rich Progress instance with custom columns for
+            quantum program tracking.
+    """
     return Progress(
         TextColumn("[bold blue]{task.fields[job_name]}"),
         BarColumn(),
         MofNCompleteColumn(),
+        TimeElapsedColumn(),
         ConditionalSpinnerColumn(),
         PhaseStatusColumn(),
         # For jupyter notebooks, refresh manually instead

divi/reporting/_qlogger.py

@@ -30,6 +30,18 @@ def _is_jupyter():
         return False  # IPython is not installed
 
 
+class CustomFormatter(logging.Formatter):
+    """
+    A custom log formatter that removes '._reporter' from the logger name.
+    """
+
+    def format(self, record):
+        # Modify the record's name attribute in place
+        if record.name.endswith("._reporter"):
+            record.name = record.name.removesuffix("._reporter")
+        return super().format(record)
+
+
 class OverwriteStreamHandler(logging.StreamHandler):
     def __init__(self, stream=None):
         super().__init__(stream)
@@ -100,9 +112,24 @@ class OverwriteStreamHandler(logging.StreamHandler):
 
 
 def enable_logging(level=logging.INFO):
+    """
+    Enable logging for the divi package with custom formatting.
+
+    Sets up a custom logger with an OverwriteStreamHandler that supports
+    message overwriting (for progress updates) and removes the '._reporter'
+    suffix from logger names.
+
+    Args:
+        level (int, optional): Logging level to set (e.g., logging.INFO,
+            logging.DEBUG). Defaults to logging.INFO.
+
+    Note:
+        This function clears any existing handlers and sets up a new handler
+        with custom formatting.
+    """
     root_logger = logging.getLogger(__name__.split(".")[0])
 
-    formatter = logging.Formatter(
+    formatter = CustomFormatter(
         "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
         datefmt="%Y-%m-%d %H:%M:%S",
     )
@@ -116,6 +143,13 @@ def enable_logging(level=logging.INFO):
 
 
 def disable_logging():
+    """
+    Disable all logging for the divi package.
+
+    Removes all handlers and sets the logging level to above CRITICAL,
+    effectively suppressing all log messages. This is useful when using
+    progress bars that provide visual feedback.
+    """
     root_logger = logging.getLogger(__name__.split(".")[0])
     root_logger.handlers.clear()
     root_logger.setLevel(logging.CRITICAL + 1)

divi/reporting/_reporter.py

@@ -13,28 +13,22 @@ class ProgressReporter(ABC):
     """An abstract base class for reporting progress of a quantum program."""
 
     @abstractmethod
-    def update(self, **kwargs):
+    def update(self, **kwargs) -> None:
         """Provides a progress update."""
         pass
 
     @abstractmethod
-    def info(self, message: str, **kwargs):
-        """
-        Provides a simple informational message.
-        No changes to progress or state.
-        """
+    def info(self, message: str, **kwargs) -> None:
+        """Provides a simple informational message."""
         pass
 
 
 class QueueProgressReporter(ProgressReporter):
     """Reports progress by putting structured dictionaries onto a Queue."""
 
-    def __init__(
-        self, job_id: str, progress_queue: Queue, has_final_computation: bool = False
-    ):
+    def __init__(self, job_id: str, progress_queue: Queue):
         self._job_id = job_id
         self._queue = progress_queue
-        self.has_final_computation = has_final_computation
 
     def update(self, **kwargs):
         payload = {"job_id": self._job_id, "progress": 1}
@@ -43,20 +37,19 @@ class QueueProgressReporter(ProgressReporter):
     def info(self, message: str, **kwargs):
         payload = {"job_id": self._job_id, "progress": 0, "message": message}
 
-        # Determine if this message indicates the job is truly finished.
-        is_final_step = "Computed Final Solution" in message or (
-            "Finished Optimization" in message and not self.has_final_computation
-        )
-
-        if is_final_step:
+        if "Finished successfully!" in message:
             payload["final_status"] = "Success"
-        elif "poll_attempt" in kwargs:
+
+        if "poll_attempt" in kwargs:
             # For polling, remove the message key so the last message persists.
             del payload["message"]
             payload["poll_attempt"] = kwargs["poll_attempt"]
             payload["max_retries"] = kwargs["max_retries"]
             payload["service_job_id"] = kwargs["service_job_id"]
             payload["job_status"] = kwargs["job_status"]
+        else:
+            # For any other message, explicitly reset the polling attempt counter.
+            payload["poll_attempt"] = 0
 
         self._queue.put(payload)
 
@@ -82,9 +75,7 @@ class LoggingProgressReporter(ProgressReporter):
             return
 
         if "iteration" in kwargs:
-            logger.info(
-                f"Running Iteration #{kwargs['iteration'] + 1} circuits: {message}\r"
-            )
+            logger.info(f"Iteration #{kwargs['iteration'] + 1}: {message}\r")
             return
 
         logger.info(message)