Repository: tskit-dev/tsdate

Issue:
Nicer EP iteration progress bars
At the moment we spew out a lot of separate progress bars to screen, 2 for each EP iteration. We should use nested progress bars instead (https://github.com/tqdm/tqdm#nested-progress-bars):

<img width="699" alt="Screenshot 2023-07-14 at 14 14 19" src="https://github.com/tskit-dev/tsdate/assets/4699014/a332b903-0534-49de-8d40-a98e078bae4d">


Discussion:

Patch:
diff --git a/tsdate/core.py b/tsdate/core.py
index b1226f0..e0025b7 100644
--- a/tsdate/core.py
+++ b/tsdate/core.py
@@ -996,7 +996,9 @@ class ExpectationPropagation(InOutAlgorithms):
         if progress is None:
             progress = self.progress
         # TODO: this will still converge if parallelized (potentially slower)
-        for edge in tqdm(edges, desc, total=self.ts.num_edges, disable=not progress):
+        for edge in tqdm(
+            edges, desc, total=self.ts.num_edges, disable=not progress, leave=False
+        ):
             if edge.child in self.fixednodes:
                 continue
             if edge.parent in self.fixednodes:
@@ -1039,17 +1041,17 @@ class ExpectationPropagation(InOutAlgorithms):
         Update edge factors from leaves to root then from root to leaves,
         and return approximate log marginal likelihood
         """
-        desc = "Expectation Propagation"
-        if iter_num:  # Show iteration number if not first iteration
-            desc = f"EP (iter {iter_num + 1:>2}, rootwards)"
-        self.propagate(
-            edges=self.edges_by_parent_asc(grouped=False), desc=desc, progress=progress
-        )
-        if iter_num:
-            desc = f"EP (iter {iter_num + 1:>2}, leafwards)"
-        self.propagate(
-            edges=self.edges_by_child_desc(grouped=False), desc=desc, progress=progress
-        )
+        if progress is None:
+            progress = self.progress
+        it = iter_num + 1  # For display purposes: show 1-based iteration
+        tasks = {
+            "Rootwards": self.edges_by_parent_asc,
+            "Leafwards": self.edges_by_child_desc,
+        }
+        for desc, func in tqdm(
+            tasks.items(), f"Iteration {it}", disable=not progress, leave=False
+        ):
+            self.propagate(edges=func(grouped=False), desc=desc, progress=progress)
         # TODO
         # marginal_lik = np.sum(self.factor_norm)
         # return marginal_lik
@@ -1538,7 +1540,10 @@ def variational_dates(
     )
 
     dynamic_prog = ExpectationPropagation(priors, liklhd, progress=progress)
-    for it in range(max_iterations):
+    for it in tqdm(
+        np.arange(max_iterations),
+        desc="Expectation Propagation",
+    ):
         dynamic_prog.iterate(iter_num=it)
     posterior = dynamic_prog.posterior
     tree_sequence, mn_post, _ = variational_mean_var(


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index e55910d..7633da8 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -271,9 +271,12 @@ class TestOutput(RunCLI):
         (out, err) = capfd.readouterr()
         assert out == ""
         # run_tsdate_cli print logging to stderr
-        assert err.count("Expectation Propagation: 100%") == 2
-        assert err.count("EP (iter  2, rootwards): 100%") == 1
-        assert err.count("rootwards): 100%") == err.count("leafwards): 100%")
+        assert err.count("Expectation Propagation: 100%") == 1
+        # The capfd fixture doesn't end up capturing progress bars with
+        # leave=False (they get deleted) so we can't see these in the output
+        # assert err.count("Iteration 1: 100%") == 1
+        # assert err.count("Rootwards: 100%") > 1
+        # assert err.count("Rootwards: 100%") == err.count("Leafwards: 100%")
 
 
 class TestEndToEnd(RunCLI):