Bokeh test and remove redundant __main__ in visualization/ (#450)

* explicit import

* fix testing import

* change show to save for tests

---------

Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MBP.hsd1.ma.comcast.net>
Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MacBook-Pro.local>
Co-authored-by: Phillip Weinberg <weinbe58@gmail.com>

src/bloqade/visualization/ir_visualize.py

@@ -1,4 +1,4 @@
-from bokeh.plotting import figure, show
+from bokeh.plotting import figure
 from bokeh.layouts import row, column
 from bokeh.models import CustomJS, MultiChoice
 from bokeh.models import (
@@ -294,79 +294,3 @@ def get_sequence_figure(ir_seq, **assignments):
 
     Seq = assemble_sequences(pulse_all)
     return Seq
-
-
-if __name__ == "__main__":
-    shared_indicator = Span(dimension="height")
-
-    ## Rydberg:
-    dats, names, spinfo = mock_data(10)
-    fig = Field_wvfm(
-        colors=itertools.cycle(Dark2_5),
-        data_sources=dats,
-        ch_names=names,
-        crx_hair_overlay=shared_indicator,
-    )
-    cube = SpacialMod(spinfo)
-    p1 = assemble_field(cube, fig, "Detuning Fields")
-
-    dats, names, spinfo = mock_data(10)
-    fig = Field_wvfm(
-        colors=itertools.cycle(Dark2_5),
-        data_sources=dats,
-        ch_names=names,
-        crx_hair_overlay=shared_indicator,
-    )
-    cube = SpacialMod(spinfo)
-    p2 = assemble_field(cube, fig, "Rabi amp Fields")
-
-    dats, names, spinfo = mock_data(10)
-    fig = Field_wvfm(
-        colors=itertools.cycle(Dark2_5),
-        data_sources=dats,
-        ch_names=names,
-        crx_hair_overlay=shared_indicator,
-    )
-    cube = SpacialMod(spinfo)
-    p3 = assemble_field(cube, fig, "Rabi phase Fields")
-
-    Panel_Pulse1 = assemble_pulse_panel([p1, p2, p3], "Rydberg")
-
-    shared_indicator = Span(dimension="height")
-
-    ## Hyperfine:
-    dats, names, spinfo = mock_data(10)
-    fig = Field_wvfm(
-        colors=itertools.cycle(Dark2_5),
-        data_sources=dats,
-        ch_names=names,
-        crx_hair_overlay=shared_indicator,
-    )
-    cube = SpacialMod(spinfo)
-    p1 = assemble_field(cube, fig, "Detuning Fields")
-
-    dats, names, spinfo = mock_data(10)
-    fig = Field_wvfm(
-        colors=itertools.cycle(Dark2_5),
-        data_sources=dats,
-        ch_names=names,
-        crx_hair_overlay=shared_indicator,
-    )
-    cube = SpacialMod(spinfo)
-    p2 = assemble_field(cube, fig, "Rabi amp Fields")
-
-    dats, names, spinfo = mock_data(10)
-    fig = Field_wvfm(
-        colors=itertools.cycle(Dark2_5),
-        data_sources=dats,
-        ch_names=names,
-        crx_hair_overlay=shared_indicator,
-    )
-    cube = SpacialMod(spinfo)
-    p3 = assemble_field(cube, fig, "Rabi phase Fields")
-
-    Panel_Pulse2 = assemble_pulse_panel([p1, p2, p3], "Hyperfine")
-
-    Seq = assemble_sequences([Panel_Pulse1, Panel_Pulse2])
-
-    show(row(Seq, mock_register()))

src/bloqade/visualization/report_visualize.py

@@ -1,4 +1,4 @@
-from bokeh.plotting import figure, show
+from bokeh.plotting import figure
 from bokeh.layouts import column, row
 from bokeh.models import CustomJS, MultiChoice, Div, HoverTool, Range1d, ColorBar
 from bokeh.models import ColumnDataSource, LinearColorMapper
@@ -421,17 +421,3 @@ def bloqadeICON():
         lines = f.readlines()
 
     return lines
-
-
-if __name__ == "__main__":
-    dat = mock_data()
-
-    fig = report_visual(*dat)
-
-    show(fig)
-    # from bokeh.models import SVGIcon
-
-    # p = figure(width=200, height=100, toolbar_location=None)
-    # p.image_url(url="file:///./logo.png")
-    # button = Button(label="", icon=SVGIcon(svg=bloqadeICON(), size=50))
-    # show(button)

tests/test_ir_visual.py

@@ -0,0 +1,90 @@
+from bloqade.visualization.ir_visualize import (
+    mock_data,
+    Field_wvfm,
+    SpacialMod,
+    assemble_field,
+    assemble_pulse_panel,
+    assemble_sequences,
+    mock_register,
+)
+import itertools
+from bokeh.palettes import Dark2_5
+from bokeh.models import Span
+from bokeh.io import save
+from bokeh.layouts import row
+
+
+def test_mock_data():
+    shared_indicator = Span(dimension="height")
+
+    ## Rydberg:
+    dats, names, spinfo = mock_data(10)
+    fig = Field_wvfm(
+        colors=itertools.cycle(Dark2_5),
+        data_sources=dats,
+        ch_names=names,
+        crx_hair_overlay=shared_indicator,
+    )
+    cube = SpacialMod(spinfo)
+    p1 = assemble_field(cube, fig, "Detuning Fields")
+
+    dats, names, spinfo = mock_data(10)
+    fig = Field_wvfm(
+        colors=itertools.cycle(Dark2_5),
+        data_sources=dats,
+        ch_names=names,
+        crx_hair_overlay=shared_indicator,
+    )
+    cube = SpacialMod(spinfo)
+    p2 = assemble_field(cube, fig, "Rabi amp Fields")
+
+    dats, names, spinfo = mock_data(10)
+    fig = Field_wvfm(
+        colors=itertools.cycle(Dark2_5),
+        data_sources=dats,
+        ch_names=names,
+        crx_hair_overlay=shared_indicator,
+    )
+    cube = SpacialMod(spinfo)
+    p3 = assemble_field(cube, fig, "Rabi phase Fields")
+
+    Panel_Pulse1 = assemble_pulse_panel([p1, p2, p3], "Rydberg")
+
+    shared_indicator = Span(dimension="height")
+
+    ## Hyperfine:
+    dats, names, spinfo = mock_data(10)
+    fig = Field_wvfm(
+        colors=itertools.cycle(Dark2_5),
+        data_sources=dats,
+        ch_names=names,
+        crx_hair_overlay=shared_indicator,
+    )
+    cube = SpacialMod(spinfo)
+    p1 = assemble_field(cube, fig, "Detuning Fields")
+
+    dats, names, spinfo = mock_data(10)
+    fig = Field_wvfm(
+        colors=itertools.cycle(Dark2_5),
+        data_sources=dats,
+        ch_names=names,
+        crx_hair_overlay=shared_indicator,
+    )
+    cube = SpacialMod(spinfo)
+    p2 = assemble_field(cube, fig, "Rabi amp Fields")
+
+    dats, names, spinfo = mock_data(10)
+    fig = Field_wvfm(
+        colors=itertools.cycle(Dark2_5),
+        data_sources=dats,
+        ch_names=names,
+        crx_hair_overlay=shared_indicator,
+    )
+    cube = SpacialMod(spinfo)
+    p3 = assemble_field(cube, fig, "Rabi phase Fields")
+
+    Panel_Pulse2 = assemble_pulse_panel([p1, p2, p3], "Hyperfine")
+
+    Seq = assemble_sequences([Panel_Pulse1, Panel_Pulse2])
+
+    save(row(Seq, mock_register()))

tests/test_report_visual.py

@@ -0,0 +1,16 @@
+from bloqade.visualization.report_visualize import mock_data, report_visual
+from bokeh.io import save
+
+
+def test_report_vis_mock():
+    dat = mock_data()
+
+    fig = report_visual(*dat)
+
+    save(fig)
+    # from bokeh.models import SVGIcon
+
+    # p = figure(width=200, height=100, toolbar_location=None)
+    # p.image_url(url="file:///./logo.png")
+    # button = Button(label="", icon=SVGIcon(svg=bloqadeICON(), size=50))
+    # show(button)