Add --scale-bar CLI option

grahamgower · May 3, 2023 · 09a904f · 09a904f
1 parent 5a5fcb4
commit 09a904f
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/demesdraw/__main__.py b/demesdraw/__main__.py
@@ -66,10 +66,23 @@ class TubesCommand(Command):
     def __init__(self, subparsers):
         super().__init__(subparsers, "tubes")
 
+        self.parser.add_argument(
+            "--scale-bar",
+            action="store_true",
+            default=False,
+            help="Draw a scale bar that indicates population size.",
+        )
+
     def __call__(self, args):
         graph = demes.load(args.input_file)
         fig, ax = demesdraw.utils.get_fig_axes(aspect=args.aspect, scale=args.scale)
-        demesdraw.tubes(graph, ax=ax, log_time=args.log_time, title=args.title)
+        demesdraw.tubes(
+            graph,
+            ax=ax,
+            log_time=args.log_time,
+            title=args.title,
+            scale_bar=args.scale_bar,
+        )
         if args.output_file is not None:
             fig.savefig(args.output_file)
         else:

diff --git a/demesdraw/tubes.py b/demesdraw/tubes.py
@@ -536,19 +536,18 @@ def random_migration_time(migration, log_scale: bool) -> float:
             ax.transData, ax.transAxes
         )
         x_offset = min(min(t.size1) for t in tubes.values())
-        # x_offset = max(max(t.size2) for t in tubes.values()) - width
-        ax_sb = ax.inset_axes([x_offset, -0.15, width, 0.01], transform=transform)
+        ax_sb = ax.inset_axes([x_offset, -0.16, width, 1e-6], transform=transform)
 
         ax_sb.yaxis.set_major_locator(matplotlib.ticker.NullLocator())
         ax_sb.spines[["left", "right", "top"]].set_visible(False)
-        ax_sb.xaxis.set_ticks_position("bottom")
+        ax_sb.xaxis.set_ticks_position("top")
         locator = matplotlib.ticker.AutoLocator()
         locator.set_params(integer=True)
         ax_sb.xaxis.set_major_locator(locator)
         ax_sb.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
         ax_sb.set_xlim(0, width)
         ax_sb.set_xlabel("deme size (individuals)")
-        ax_sb.xaxis.set_label_position("top")
+        ax_sb.xaxis.set_label_position("bottom")
 
     # Status bar text when in interactive mode.
     def format_coord(x, y):

diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -73,6 +73,9 @@ def test_output_file_is_created(self, subcommand, params):
             assert output_file.exists()
             assert output_file.stat().st_size > 0
 
+    def test_tubes_scale_bar(self):
+        self.test_output_file_is_created("tubes", "--scale-bar")
+
     @pytest.mark.parametrize("subcommand", ["tubes", "size_history"])
     def test_input_from_stdin(self, subcommand):
         input_file = tests.example_files()[0]