feat(progress): add Rich progress bars to workflow and expand tests:

- Add ResultsReporter progress context manager
- Propagate optional progress sink through workflow orchestration
- Add progress reporting for:
  - Conformational state construction (per group)
  - Frame processing stage (per frame)
- Keep entropy graph execution silent due to fast runtime
- Update runtime tests to reflect wrapped RuntimeError behavior

Author: harryswift01

CodeEntropy/entropy/graph.py

@@ -80,23 +80,44 @@ def build(self) -> "EntropyGraph":
 
         return self
 
-    def execute(self, shared_data: SharedData) -> Dict[str, Any]:
+    def execute(
+        self, shared_data: SharedData, *, progress: object | None = None
+    ) -> Dict[str, Any]:
         """Execute the entropy graph in topological order.
 
+        Nodes are executed in dependency order (topological sort). Each node reads
+        from and may mutate `shared_data`. Dict-like outputs returned by nodes are
+        merged into a single results dictionary.
+
+        This method intentionally does *not* create a progress bar/task for the
+        entropy graph itself because the graph is typically very fast. If a progress
+        sink is provided, it is forwarded to nodes that accept it.
+
         Args:
             shared_data: Mutable shared data dictionary passed to each node.
+            progress: Optional progress sink (e.g., from ResultsReporter.progress()).
+                Forwarded to node `run()` methods that accept a `progress` keyword.
 
         Returns:
-            Dictionary containing the merged outputs of all nodes (only including
-            outputs that are dict-like).
+            Dictionary containing merged dict outputs produced by nodes. On key
+            collision, later nodes overwrite earlier keys.
 
         Raises:
             KeyError: If a node name is missing from the internal node registry.
         """
         results: Dict[str, Any] = {}
+
         for node_name in nx.topological_sort(self._graph):
             node = self._nodes[node_name]
-            out = node.run(shared_data)
+
+            if progress is not None:
+                try:
+                    out = node.run(shared_data, progress=progress)
+                except TypeError:
+                    out = node.run(shared_data)
+            else:
+                out = node.run(shared_data)
+
             if isinstance(out, dict):
                 results.update(out)
         return results

CodeEntropy/entropy/workflow.py

@@ -124,8 +124,9 @@ def execute(self) -> None:
             traj=traj,
         )
 
-        self._run_level_dag(shared_data)
-        self._run_entropy_graph(shared_data)
+        with self._reporter.progress(transient=False) as p:
+            self._run_level_dag(shared_data, progress=p)
+            self._run_entropy_graph(shared_data, progress=p)
 
         self._finalize_molecule_results()
         self._reporter.log_tables()
@@ -164,21 +165,29 @@ def _build_shared_data(
         }
         return shared_data
 
-    def _run_level_dag(self, shared_data: SharedData) -> None:
+    def _run_level_dag(
+        self, shared_data: SharedData, *, progress: object | None = None
+    ) -> None:
         """Execute the structural/level DAG.
 
         Args:
             shared_data: Shared data dict that will be mutated by the DAG.
+            progress: Optional progress sink provided by ResultsReporter.progress().
         """
-        LevelDAG(self._universe_operations).build().execute(shared_data)
+        LevelDAG(self._universe_operations).build().execute(
+            shared_data, progress=progress
+        )
 
-    def _run_entropy_graph(self, shared_data: SharedData) -> None:
+    def _run_entropy_graph(
+        self, shared_data: SharedData, *, progress: object | None = None
+    ) -> None:
         """Execute the entropy calculation graph and merge results into shared_data.
 
         Args:
             shared_data: Shared data dict that will be mutated by the graph.
+            progress: Optional progress sink provided by ResultsReporter.progress().
         """
-        entropy_results = EntropyGraph().build().execute(shared_data)
+        entropy_results = EntropyGraph().build().execute(shared_data, progress=progress)
         shared_data.update(entropy_results)
 
     def _build_trajectory_slice(self) -> TrajectorySlice:

CodeEntropy/levels/dihedrals.py

@@ -45,79 +45,114 @@ def build_conformational_states(
         end: int,
         step: int,
         bin_width: float,
+        progress: object | None = None,
     ):
-        """Build conformational state labels for UA and residue levels.
+        """Build conformational state labels from trajectory dihedrals.
+
+        This method constructs discrete conformational state descriptors used in
+        configurational entropy calculations. It supports united-atom (UA) and
+        residue-level state generation depending on which hierarchy levels are
+        enabled per molecule.
+
+        Progress reporting is optional and UI-agnostic: if a progress sink is
+        provided, the method will create a single task and advance it once per
+        molecule group.
 
         Args:
-            data_container: MDAnalysis universe containing the system.
-            levels: Mapping of molecule_id -> list of enabled levels.
+            data_container: MDAnalysis Universe (or compatible container) used to
+                extract fragments and compute dihedral time series.
+            levels: Mapping of molecule_id -> iterable of enabled level names
+                (e.g., ["united_atom", "residue"]).
             groups: Mapping of group_id -> list of molecule_ids.
-            start: Start frame index (currently not applied in legacy sampling).
-            end: End frame index (currently not applied in legacy sampling).
-            step: Step size (currently not applied in legacy sampling).
-            bin_width: Histogram bin width (degrees).
+            start: Inclusive start frame index.
+            end: Exclusive end frame index.
+            step: Frame stride.
+            bin_width: Histogram bin width in degrees used when identifying peak
+                dihedral populations.
+            progress: Optional progress sink (e.g., from ResultsReporter.progress()).
+                Must expose add_task(), update(), and advance().
 
         Returns:
             Tuple of:
-                states_ua: Dict[(group_id, res_id)] -> list of state labels.
-                states_res: List indexed by group_id -> list of state labels.
+                states_ua: Dict mapping (group_id, local_residue_id) -> list of state
+                    labels (strings) across the analyzed trajectory.
+                states_res: List-like structure indexed by group_id (or equivalent)
+                    containing residue-level state labels (strings) across the
+                    analyzed trajectory.
+
+        Notes:
+            - This function advances progress once per group_id.
+            - Frame slicing arguments (start/end/step) are forwarded to downstream
+            helpers as implemented in this module.
         """
         number_groups = len(groups)
         states_ua: Dict[UAKey, List[str]] = {}
         states_res: List[List[str]] = [None] * number_groups
 
-        total_items = self._count_total_items(levels=levels, groups=groups)
-
-        with self._progress_bar(total_items) as progress:
+        task = None
+        if progress is not None:
+            total = max(1, len(groups))
             task = progress.add_task(
-                "[green]Building Conformational States...",
-                total=total_items,
-                title="Starting...",
+                "[green]Conformational states",
+                total=total,
+                title="Initializing",
             )
 
-            for group_id in groups.keys():
-                molecules = groups[group_id]
-                if not molecules:
+        if not groups:
+            if task is not None:
+                progress.update(task, title="No groups")
+                progress.advance(task)
+            return states_ua, states_res
+
+        for group_id in groups.keys():
+            molecules = groups[group_id]
+            if not molecules:
+                if task is not None:
+                    progress.update(task, title=f"Group {group_id} (empty)")
                     progress.advance(task)
-                    continue
+                continue
 
-                mol = self._universe_operations.extract_fragment(
-                    data_container, molecules[0]
-                )
+            if task is not None:
+                progress.update(task, title=f"Group {group_id}")
 
-                dihedrals_ua, dihedrals_res = self._collect_dihedrals_for_group(
-                    mol=mol,
-                    level_list=levels[molecules[0]],
-                )
+            mol = self._universe_operations.extract_fragment(
+                data_container, molecules[0]
+            )
 
-                peaks_ua, peaks_res = self._collect_peaks_for_group(
-                    data_container=data_container,
-                    molecules=molecules,
-                    dihedrals_ua=dihedrals_ua,
-                    dihedrals_res=dihedrals_res,
-                    bin_width=bin_width,
-                    start=start,
-                    end=end,
-                    step=step,
-                    level_list=levels[molecules[0]],
-                )
+            dihedrals_ua, dihedrals_res = self._collect_dihedrals_for_group(
+                mol=mol,
+                level_list=levels[molecules[0]],
+            )
 
-                self._assign_states_for_group(
-                    data_container=data_container,
-                    group_id=group_id,
-                    molecules=molecules,
-                    dihedrals_ua=dihedrals_ua,
-                    peaks_ua=peaks_ua,
-                    dihedrals_res=dihedrals_res,
-                    peaks_res=peaks_res,
-                    start=start,
-                    end=end,
-                    step=step,
-                    level_list=levels[molecules[0]],
-                    states_ua=states_ua,
-                    states_res=states_res,
-                )
+            peaks_ua, peaks_res = self._collect_peaks_for_group(
+                data_container=data_container,
+                molecules=molecules,
+                dihedrals_ua=dihedrals_ua,
+                dihedrals_res=dihedrals_res,
+                bin_width=bin_width,
+                start=start,
+                end=end,
+                step=step,
+                level_list=levels[molecules[0]],
+            )
+
+            self._assign_states_for_group(
+                data_container=data_container,
+                group_id=group_id,
+                molecules=molecules,
+                dihedrals_ua=dihedrals_ua,
+                peaks_ua=peaks_ua,
+                dihedrals_res=dihedrals_res,
+                peaks_res=peaks_res,
+                start=start,
+                end=end,
+                step=step,
+                level_list=levels[molecules[0]],
+                states_ua=states_ua,
+                states_res=states_res,
+            )
 
+            if task is not None:
                 progress.advance(task)
 
         return states_ua, states_res

CodeEntropy/levels/level_dag.py

@@ -82,24 +82,28 @@ def build(self) -> "LevelDAG":
         self._frame_dag.build()
         return self
 
-    def execute(self, shared_data: Dict[str, Any]) -> Dict[str, Any]:
-        """Execute the full hierarchy workflow and mutate shared_data.
-
-        Args:
-            shared_data: Shared workflow data dict that will be mutated by the DAG.
-
-        Returns:
-            The mutated shared_data dict.
-        """
+    def execute(
+        self, shared_data: Dict[str, Any], *, progress: object | None = None
+    ) -> Dict[str, Any]:
+        """Execute the full hierarchy workflow and mutate shared_data."""
         shared_data.setdefault("axes_manager", AxesCalculator())
-        self._run_static_stage(shared_data)
-        self._run_frame_stage(shared_data)
+        self._run_static_stage(shared_data, progress=progress)
+        self._run_frame_stage(shared_data, progress=progress)
         return shared_data
 
-    def _run_static_stage(self, shared_data: Dict[str, Any]) -> None:
+    def _run_static_stage(
+        self, shared_data: Dict[str, Any], *, progress: object | None = None
+    ) -> None:
         """Run all static nodes in dependency order."""
         for node_name in nx.topological_sort(self._static_graph):
-            self._static_nodes[node_name].run(shared_data)
+            node = self._static_nodes[node_name]
+            if progress is not None:
+                try:
+                    node.run(shared_data, progress=progress)
+                    continue
+                except TypeError:
+                    pass
+            node.run(shared_data)
 
     def _add_static(
         self, name: str, node: Any, deps: Optional[list[str]] = None
@@ -110,16 +114,74 @@ def _add_static(
         for dep in deps or []:
             self._static_graph.add_edge(dep, name)
 
-    def _run_frame_stage(self, shared_data: Dict[str, Any]) -> None:
-        """Run the frame DAG for each selected trajectory frame and reduce outputs."""
+    def _run_frame_stage(
+        self, shared_data: Dict[str, Any], *, progress: object | None = None
+    ) -> None:
+        """Execute the per-frame DAG stage and reduce frame outputs.
+
+        This method iterates over the selected trajectory frames, executes the
+        frame-local DAG for each frame, and reduces the resulting outputs into the
+        shared accumulators stored in `shared_data`.
+
+        Progress reporting is optional. If a progress sink is provided, a task is
+        always created. When the total number of frames cannot be determined, the
+        task is created with total=None (indeterminate).
+
+        Args:
+            shared_data: Shared data dictionary. Must contain:
+                - "reduced_universe": MDAnalysis Universe providing the trajectory.
+                - "start", "end", "step": frame slicing parameters.
+                - any additional keys required by the frame DAG and reducer.
+            progress: Optional progress sink (e.g., from ResultsReporter.progress()).
+                Must expose add_task(), update(), and advance().
+
+        Returns:
+            None. Mutates `shared_data` in-place via reduction.
+
+        Notes:
+            The task title shows the current frame index being processed.
+        """
         u = shared_data["reduced_universe"]
         start, end, step = shared_data["start"], shared_data["end"], shared_data["step"]
 
+        task = None
+        total_frames = None
+
+        if progress is not None:
+            try:
+                n_frames = len(u.trajectory)
+
+                s = 0 if start is None else int(start)
+                e = n_frames if end is None else int(end)
+
+                if e < 0:
+                    e = n_frames + e
+
+                e = max(0, min(e, n_frames))
+                s = max(0, min(s, e))
+
+                st = 1 if step is None else int(step)
+                if st > 0:
+                    total_frames = max(0, (e - s + st - 1) // st)
+            except Exception:
+                total_frames = None
+
+            task = progress.add_task(
+                "[green]Frame processing",
+                total=total_frames,
+                title="Initializing",
+            )
+
         for ts in u.trajectory[start:end:step]:
-            frame_index = ts.frame
-            frame_out = self._frame_dag.execute_frame(shared_data, frame_index)
+            if task is not None:
+                progress.update(task, title=f"Frame {ts.frame}")
+
+            frame_out = self._frame_dag.execute_frame(shared_data, ts.frame)
             self._reduce_one_frame(shared_data, frame_out)
 
+            if task is not None:
+                progress.advance(task)
+
     @staticmethod
     def _incremental_mean(old: Any, new: Any, n: int) -> Any:
         """Compute an incremental mean.