From 1b088e37680a4bd1b7107f69bab32bfc66337a34 Mon Sep 17 00:00:00 2001
From: Kai-Hsin Wu <khwu@KHWus-MacBook-Pro.local>
Date: Thu, 3 Aug 2023 10:39:51 -0400
Subject: [PATCH] add labels for registers

---
 src/bloqade/ir/location/base.py                 | 6 +++++-
 src/bloqade/submission/ir/task_specification.py | 8 ++++----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/bloqade/ir/location/base.py b/src/bloqade/ir/location/base.py
index cf41596c7..9bf54d39c 100644
--- a/src/bloqade/ir/location/base.py
+++ b/src/bloqade/ir/location/base.py
@@ -50,7 +50,7 @@ def enumerate(self) -> Generator[LocationInfo, None, None]:
         """enumerate all locations in the register."""
         raise NotImplementedError
 
-    def figure(self, **assignments):
+    def figure(self, fig_kwargs=None, **assignments):
         """obtain a figure object from the atom arrangement."""
         xs_filled, ys_filled, labels_filled = [], [], []
         xs_vacant, ys_vacant, labels_vacant = [], [], []
@@ -103,7 +103,11 @@ def figure(self, **assignments):
         ]
 
         ## remove box_zoom since we don't want to change the scale
+        if fig_kwargs is None:
+            fig_kwargs = {}
+
         p = figure(
+            **fig_kwargs,
             width=400,
             height=400,
             tools="wheel_zoom,reset, undo, redo, pan",
diff --git a/src/bloqade/submission/ir/task_specification.py b/src/bloqade/submission/ir/task_specification.py
index 4c578a399..b5d72d406 100644
--- a/src/bloqade/submission/ir/task_specification.py
+++ b/src/bloqade/submission/ir/task_specification.py
@@ -329,14 +329,14 @@ def discretize(self, task_capabilities: QuEraCapabilities):
             filling=self.filling,
         )
 
-    def figure(self):
+    def figure(self, **fig_kwargs):
         ## use ir.Atom_oarrangement's plotting:
         ## covert unit to m -> um
         sites_um = list(
             map(lambda cord: (float(cord[0]) * 1e6, float(cord[1]) * 1e6), self.sites)
         )
         reg = ListOfLocations().add_positions(sites_um, self.filling)
-        fig_reg = reg.figure()  # ignore the B-rad widget
+        fig_reg = reg.figure(fig_kwargs=fig_kwargs)  # ignore the B-rad widget
         return fig_reg
 
     def show(self):
@@ -386,7 +386,7 @@ def figure(self):
         # grab global figures
         rabi_amplitude = (
             self.effective_hamiltonian.rydberg.rabi_frequency_amplitude.figure(
-                rabi_amp_src, tools="hover,wheel_zoom,reset, undo, redo, pan"
+                rabi_amp_src, tools="wheel_zoom,reset, undo, redo, pan"
             )
         )
 
@@ -402,7 +402,7 @@ def figure(self):
         )
 
         # lattice:
-        register = self.lattice.figure()
+        register = self.lattice.figure(x_axis_label="x (um)", y_axis_label="y (um)")
 
         col_plt = gridplot(
             [[rabi_amplitude], [global_detuning], [rabi_phase]],