One progress bar per chain when samplings

Author: jessegrabowski

pymc/sampling/parallel.py

@@ -28,7 +28,9 @@
 import numpy as np
 
 from rich.console import Console
-from rich.progress import BarColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn
+from rich.progress import TextColumn
+from rich.style import Style
+from rich.table import Column
 from rich.theme import Theme
 from threadpoolctl import threadpool_limits
 
@@ -37,6 +39,7 @@
 from pymc.exceptions import SamplingError
 from pymc.util import (
     CustomProgress,
+    DivergenceBarColumn,
     RandomGeneratorState,
     default_progress_theme,
     get_state_from_generator,
@@ -487,20 +490,35 @@ def __init__(
         self._in_context = False
 
         self._progress = CustomProgress(
-            "[progress.description]{task.description}",
-            BarColumn(),
-            "[progress.percentage]{task.percentage:>3.0f}%",
-            TimeRemainingColumn(),
-            TextColumn("/"),
-            TimeElapsedColumn(),
+            DivergenceBarColumn(
+                table_column=Column("Progress", ratio=2),
+                diverging_color="tab:red",
+                diverging_finished_color="tab:purple",
+                complete_style=Style.parse("rgb(31,119,180)"),  # tab:blue
+                finished_style=Style.parse("rgb(44,160,44)"),  # tab:green
+            ),
+            TextColumn("{task.fields[draws]:,d}", table_column=Column("Draws", ratio=1)),
+            TextColumn(
+                "{task.fields[divergences]:,d}", table_column=Column("Divergences", ratio=1)
+            ),
+            TextColumn("{task.fields[step_size]:0.2f}", table_column=Column("Step size", ratio=1)),
+            TextColumn("{task.fields[tree_depth]:,d}", table_column=Column("Tree depth", ratio=1)),
+            TextColumn(
+                "{task.fields[sampling_speed]:0.2f} {task.fields[speed_unit]}",
+                table_column=Column("Sampling Speed", ratio=1),
+            ),
             console=Console(theme=progressbar_theme),
             disable=not progressbar,
+            include_headers=True,
         )
+
         self._show_progress = progressbar
         self._divergences = 0
+        self._divergences_by_chain = [0] * chains
         self._completed_draws = 0
-        self._total_draws = chains * (draws + tune)
-        self._desc = "Sampling {0._chains:d} chains, {0._divergences:,d} divergences"
+        self._completed_draws_by_chain = [0] * chains
+        self._total_draws = draws + tune
+        self._desc = "Sampling chain"
         self._chains = chains
 
     def _make_active(self):
@@ -517,31 +535,71 @@ def __iter__(self):
         self._make_active()
 
         with self._progress as progress:
-            task = progress.add_task(
-                self._desc.format(self),
-                completed=self._completed_draws,
-                total=self._total_draws,
-            )
+            tasks = [
+                progress.add_task(
+                    self._desc.format(self),
+                    completed=self._completed_draws,
+                    total=self._total_draws,
+                    chain_idx=chain_idx,
+                    draws=0,
+                    divergences=0,
+                    step_size=0.0,
+                    tree_depth=0,
+                    sampling_speed=0,
+                    speed_unit="draws/s",
+                )
+                for chain_idx in range(self._chains)
+            ]
 
             while self._active:
                 draw = ProcessAdapter.recv_draw(self._active)
                 proc, is_last, draw, tuning, stats = draw
+                speed = 0
+                unit = "draws/s"
+
                 self._completed_draws += 1
+                self._completed_draws_by_chain[proc.chain] += 1
+
                 if not tuning and stats and stats[0].get("diverging"):
                     self._divergences += 1
+                    self._divergences_by_chain[proc.chain] += 1
+
+                if self._show_progress:
+                    elapsed = progress._tasks[proc.chain].elapsed
+                    speed = self._completed_draws_by_chain[proc.chain] / elapsed
+
+                    if speed > 1:
+                        unit = "draws/s"
+                    else:
+                        unit = "s/draws"
+                        speed = 1 / speed
+
                 progress.update(
-                    task,
-                    completed=self._completed_draws,
-                    total=self._total_draws,
-                    description=self._desc.format(self),
+                    tasks[proc.chain],
+                    completed=self._completed_draws_by_chain[proc.chain],
+                    draws=draw,
+                    divergences=self._divergences_by_chain[proc.chain],
+                    step_size=stats[0].get("step_size", 0),
+                    tree_depth=stats[0].get("tree_size", 0),
+                    sampling_speed=speed,
+                    speed_unit=unit,
                 )
 
                 if is_last:
+                    self._completed_draws_by_chain[proc.chain] += 1
+
                     proc.join()
                     self._active.remove(proc)
                     self._finished.append(proc)
                     self._make_active()
-                    progress.update(task, description=self._desc.format(self), refresh=True)
+                    progress.update(
+                        tasks[proc.chain],
+                        draws=draw + 1,
+                        divergences=self._divergences_by_chain[proc.chain],
+                        step_size=stats[0].get("step_size", 0),
+                        tree_depth=stats[0].get("tree_size", 0),
+                        refresh=True,
+                    )
 
                 # We could also yield proc.shared_point_view directly,
                 # and only call proc.write_next() after the yield returns.

pymc/util.py

@@ -17,7 +17,7 @@
 import warnings
 
 from collections import namedtuple
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from copy import deepcopy
 from typing import NewType, cast
 
@@ -30,7 +30,10 @@
 from pytensor import Variable
 from pytensor.compile import SharedVariable
 from pytensor.graph.utils import ValidatingScratchpad
-from rich.progress import Progress
+from rich.box import SIMPLE_HEAD
+from rich.progress import BarColumn, Progress, Task
+from rich.style import Style
+from rich.table import Column, Table
 from rich.theme import Theme
 
 from pymc.exceptions import BlockModelAccessError
@@ -556,8 +559,10 @@ class CustomProgress(Progress):
     it's `True`.
     """
 
-    def __init__(self, *args, **kwargs):
-        self.is_enabled = kwargs.get("disable", None) is not True
+    def __init__(self, *args, disable=False, include_headers=False, **kwargs):
+        self.is_enabled = not disable
+        self.include_headers = include_headers
+
         if self.is_enabled:
             super().__init__(*args, **kwargs)
 
@@ -607,6 +612,88 @@ def update(
             )
         return None
 
+    def make_tasks_table(self, tasks: Iterable[Task]) -> Table:
+        """Get a table to render the Progress display.
+
+        Unlike the parent method, this one returns a full table (not a grid), allowing for column headings.
+
+        Parameters
+        ----------
+        tasks: Iterable[Task]
+            An iterable of Task instances, one per row of the table.
+
+        Returns
+        -------
+        table: Table
+            A table instance.
+        """
+
+        def call_column(column, task):
+            if hasattr(column, "callbacks"):
+                column.callbacks(task)
+
+            return column(task)
+
+        table_columns = (
+            (
+                Column(no_wrap=True)
+                if isinstance(_column, str)
+                else _column.get_table_column().copy()
+            )
+            for _column in self.columns
+        )
+        if self.include_headers:
+            table = Table(
+                *table_columns,
+                padding=(0, 1),
+                expand=self.expand,
+                show_header=True,
+                show_edge=True,
+                box=SIMPLE_HEAD,
+            )
+        else:
+            table = Table.grid(*table_columns, padding=(0, 1), expand=self.expand)
+
+        for task in tasks:
+            if task.visible:
+                table.add_row(
+                    *(
+                        (
+                            column.format(task=task)
+                            if isinstance(column, str)
+                            else call_column(column, task)
+                        )
+                        for column in self.columns
+                    )
+                )
+
+        return table
+
+
+class DivergenceBarColumn(BarColumn):
+    def __init__(self, *args, diverging_color="red", diverging_finished_color="purple", **kwargs):
+        from matplotlib.colors import to_rgb
+
+        self.diverging_color = diverging_color
+        self.diverging_rgb = [int(x * 255) for x in to_rgb(self.diverging_color)]
+
+        self.diverging_finished_color = diverging_finished_color
+        self.diverging_finished_rgb = [int(x * 255) for x in to_rgb(self.diverging_finished_color)]
+
+        super().__init__(*args, **kwargs)
+
+        self.non_diverging_style = self.complete_style
+        self.non_diverging_finished_style = self.finished_style
+
+    def callbacks(self, task: "Task"):
+        divergences = task.fields.get("divergences", 0)
+        if divergences > 0:
+            self.complete_style = Style.parse("rgb({},{},{})".format(*self.diverging_rgb))
+            self.finished_style = Style.parse("rgb({},{},{})".format(*self.diverging_finished_rgb))
+        else:
+            self.complete_style = self.non_diverging_style
+            self.finished_style = self.non_diverging_finished_style
+
 
 RandomGeneratorState = namedtuple("RandomGeneratorState", ["bit_generator_state", "seed_seq_state"])