fixing bug with display module. (#586)

* fixing bug with display module.

* add tests for builder

* add codecov setting, 4%

---------

Co-authored-by: kaihsin <kaihsinwu@gmail.com>

codecov.yml

@@ -0,0 +1,10 @@
+coverage:
+  status:
+    project:
+      default:
+        target: auto
+        threshold: 4%
+    patch:
+      default:
+        target: auto
+        threshold: 4%

src/bloqade/visualization/__init__.py

@@ -1,7 +1,15 @@
 Use_bokeh = True
 
 if Use_bokeh:
-    from .display import display_ir, display_report, display_task_ir, display_builder
+    from .display import (
+        display_ir,
+        display_report,
+        display_task_ir,
+        display_builder,
+        figure_ir,
+        builder_figure,
+        report_figure,
+    )
     from .task_visualize import get_task_ir_figure
     from .atom_arragement_visualize import get_atom_arrangement_figure
     from .ir_visualize import get_ir_figure, get_field_figure, get_pulse_figure
@@ -36,3 +44,13 @@ def get_field_figure(ir_field, title, indicator, **assignments):
 
     def get_pulse_figure(ir_pulse, title: str = None, **assginments):
         raise Warning("Bokeh not installed", UserWarning)
+
+    # figure:
+    def figure_ir(obj, assignemnts):
+        raise Warning("Bokeh not installed", UserWarning)
+
+    def builder_figure(builder, batch_id):
+        raise Warning("Bokeh not installed", UserWarning)
+
+    def report_figure(report):
+        raise Warning("Bokeh not installed", UserWarning)

src/bloqade/visualization/display.py

@@ -36,11 +36,37 @@ def display_ir(obj, assignments):
         raise NotImplementedError(f"not supported IR for display, got {type(obj)}")
 
 
+def figure_ir(obj, assignments):
+    from bloqade.ir.analog_circuit import AnalogCircuit
+    from bloqade.ir.control.pulse import PulseExpr
+    from bloqade.ir.control.sequence import SequenceExpr
+    from bloqade.ir.control.field import Field, SpatialModulation
+    from bloqade.ir.control.waveform import Waveform
+    from bloqade.ir.location import AtomArrangement
+
+    if isinstance(obj, AnalogCircuit):
+        return analog_circuit_figure(obj, assignments)
+    elif isinstance(obj, PulseExpr):
+        return obj.figure(**assignments)
+    elif isinstance(obj, SequenceExpr):
+        return obj.figure(**assignments)
+    elif isinstance(obj, Field):
+        return obj.figure(**assignments)
+    elif isinstance(obj, SpatialModulation):
+        return obj.figure(**assignments)
+    elif isinstance(obj, Waveform):
+        return obj.figure(**assignments)
+    elif isinstance(obj, AtomArrangement):
+        return atom_arrangement_figure(obj, assignments)
+    else:
+        raise NotImplementedError(f"not supported IR for display, got {type(obj)}")
+
+
 def liner(txt):
     return f"<p>{txt}</p>"
 
 
-def display_builder(builder, batch_id):
+def builder_figure(builder, batch_id):
     from bloqade.builder.parse.builder import Parser
 
     analog_circ, metas = Parser().parse_source(builder)
@@ -52,18 +78,25 @@ def display_builder(builder, batch_id):
     for key, val in kwargs.items():
         out += liner(f" :: {key} = {val}")
 
-    rowobj = analog_circ.figure(**kwargs)
+    field, reg = analog_circ.figure(**kwargs)
 
-    field = rowobj.children[0]
-    reg = rowobj.children[1]
     div = Div(text=out, width=200, height=200)
+    return row(field, column(reg, div))
 
-    show(row(field, column(reg, div)))
+
+def display_builder(builder, batch_id):
+    fig = builder_figure(builder, batch_id)
+    show(fig)
 
 
 ## ir
+def analog_circuit_figure(analog_circuit, assignments):
+    return row(*analog_circuit.figure(**assignments))
+
+
 def display_analog_circuit(analog_circuit, assignments):
-    show(row(*analog_circuit.figure(**assignments)))
+    fig = analog_circuit_figure(analog_circuit, assignments)
+    show(fig)
 
 
 def display_pulse(pulse, assignments):
@@ -86,7 +119,7 @@ def display_waveform(wvfm_ir, assignments):
     show(wvfm_ir.figure(**assignments))
 
 
-def display_atom_arrangement(atom_arrangement, assignments):
+def atom_arrangement_figure(atom_arrangement, assignments):
     """show the register."""
     p = atom_arrangement.figure(None, **assignments)
 
@@ -113,14 +146,26 @@ def display_atom_arrangement(atom_arrangement, assignments):
 
     full = column(p, row(Brad_input, toggle_button))
     # full.sizing_mode="scale_both"
-    show(full)
+    return full
+
+
+def display_atom_arrangement(atom_arrangement, assignments):
+    fig = atom_arrangement_figure(atom_arrangement, assignments)
+    show(fig)
 
 
 ## report
-def display_report(report):
+
+
+def report_figure(report):
     dat = report_visualize.format_report_data(report)
     p = report_visualize.report_visual(*dat)
-    show(p)
+    return p
+
+
+def display_report(report):
+    fig = report_figure(report)
+    show(fig)
 
 
 ## task_ir

tests/test_builder_visual.py

@@ -0,0 +1,20 @@
+from bloqade.ir.location import Square
+import numpy as np
+from decimal import Decimal
+from bloqade.visualization import builder_figure
+from bokeh.io import save
+
+
+def test_builder_vis():
+    program = (
+        Square(3, lattice_spacing=1 / np.sqrt(Decimal(2)))
+        .apply_defect_count(4)
+        .scale(6.1)
+        .rydberg.detuning.uniform.piecewise_linear(
+            [0.1, 0.5, 0.1], [-10, -10, "final_detuning", "final_detuning"]
+        )
+        .amplitude.uniform.piecewise_linear([0.1, 0.5, 0.1], [0, 15, 15, 0])
+        .batch_assign(final_detuning=np.linspace(0, 50, 51).tolist())
+    )
+    fig = builder_figure(program, 0)
+    save(fig)