add line_width option

ihincks · Oct 29, 2024 · ab9b77f · ab9b77f
1 parent ca0d723
commit ab9b77f
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 10 deletions.
diff --git a/qiskit_ibm_runtime/execution_span/execution_spans.py b/qiskit_ibm_runtime/execution_span/execution_spans.py
@@ -117,7 +117,7 @@ def sort(self, inplace: bool = True) -> "ExecutionSpans":
         obj._spans.sort()
         return obj
 
-    def draw(self, normalize_y: bool = False) -> "PlotlyFigure":
+    def draw(self, normalize_y: bool = False, line_width: int = 4) -> "PlotlyFigure":
         """Draw these execution spans.
 
         .. note::
@@ -127,11 +127,12 @@ def draw(self, normalize_y: bool = False) -> "PlotlyFigure":
         Args:
             normalize_y: Whether to display the y-axis units as a percentage of work
                 complete, rather than cumulative shots completed.
+            line_width: The thickness of line segments.
 
         Returns:
             A plotly figure.
         """
         # pylint: disable=import-outside-toplevel, cyclic-import
         from ..visualization import draw_execution_spans
 
-        return draw_execution_spans(self, normalize_y=normalize_y)
+        return draw_execution_spans(self, normalize_y=normalize_y, line_width=line_width)
diff --git a/qiskit_ibm_runtime/visualization/draw_execution_spans.py b/qiskit_ibm_runtime/visualization/draw_execution_spans.py
@@ -49,7 +49,10 @@ def _get_id(span: ExecutionSpan, multiple: bool) -> str:
 
 
 def draw_execution_spans(
-    *spans: ExecutionSpans, common_start: bool = False, normalize_y: bool = False
+    *spans: ExecutionSpans,
+    common_start: bool = False,
+    normalize_y: bool = False,
+    line_width: int = 4,
 ) -> PlotlyFigure:
     """Draw one or more :class:`~.ExecutionSpans` on a bar plot.
 
@@ -59,6 +62,7 @@ def draw_execution_spans(
             at :math:`t=0`.
         normalize_y: Whether to display the y-axis units as a percentage of work complete, rather
             than cumulative shots completed.
+        line_width: The thickness of line segments.
 
     Returns:
         A plotly figure.
@@ -92,7 +96,7 @@ def draw_execution_spans(
                     x=[span.start - offset, span.stop - offset],
                     y=[y_value, y_value],
                     mode="lines",
-                    line={"width": 4, "color": color},
+                    line={"width": line_width, "color": color},
                     text=text,
                     hoverinfo="text",
                 )

diff --git a/test/unit/test_execution_span.py b/test/unit/test_execution_span.py
@@ -292,8 +292,9 @@ def test_sort(self):
         self.assertLess(spans[1], spans[0])
         self.assertLess(new_sort[0], new_sort[1])
 
-    @ddt.data(False, True)
-    def test_draw(self, normalize_y):
+    @ddt.data((False, 4), (True, 6))
+    @ddt.unpack
+    def test_draw(self, normalize_y, width):
         """Test the draw method."""
         spans = ExecutionSpans([self.span2, self.span1])
-        self.save_plotly_artifact(spans.draw(normalize_y=normalize_y))
+        self.save_plotly_artifact(spans.draw(normalize_y=normalize_y, line_width=width))
diff --git a/test/unit/test_visualization.py b/test/unit/test_visualization.py
@@ -75,11 +75,15 @@ def test_one_spans(self, normalize_y):
         fig = draw_execution_spans(self.spans0, normalize_y=normalize_y)
         self.save_plotly_artifact(fig)
 
-    @ddt.data((False, False), (True, True))
+    @ddt.data((False, False, 4), (True, True, 8))
     @ddt.unpack
-    def test_two_spans(self, normalize_y, common_start):
+    def test_two_spans(self, normalize_y, common_start, width):
         """Test with two sets of spans."""
         fig = draw_execution_spans(
-            self.spans0, self.spans1, normalize_y=normalize_y, common_start=common_start
+            self.spans0,
+            self.spans1,
+            normalize_y=normalize_y,
+            common_start=common_start,
+            line_width=width,
         )
         self.save_plotly_artifact(fig)