From abb714b039dcd887cc13ae9fabf45501bb37863f Mon Sep 17 00:00:00 2001 From: Briac Flesselles Date: Sat, 26 Oct 2024 15:45:12 +0200 Subject: [PATCH 1/3] merged upstream --- pulser-core/pulser/register/register.py | 44 ++++++++++++++++++++----- 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/pulser-core/pulser/register/register.py b/pulser-core/pulser/register/register.py index b110ca26c..7951282d3 100644 --- a/pulser-core/pulser/register/register.py +++ b/pulser-core/pulser/register/register.py @@ -404,6 +404,8 @@ def draw( kwargs_savefig: dict = {}, custom_ax: Optional[Axes] = None, show: bool = True, + draw_empty_sites: bool = False, + empty_color: Optional[str] = None, ) -> None: """Draws the entire register. @@ -434,6 +436,8 @@ def draw( show: Whether or not to call `plt.show()` before returning. When combining this plot with other ones in a single figure, one may need to set this flag to False. + draw_empty_sites: If True, draws the empty sites as well. + empty_color: The color of the empty sites. Default is 'r'. Note: When drawing half the blockade radius, we say there is a blockade @@ -441,6 +445,7 @@ def draw( This representation is preferred over drawing the full Rydberg radius because it helps in seeing the interactions between atoms. """ + super()._draw_checks( len(self._ids), blockade_radius=blockade_radius, @@ -448,25 +453,48 @@ def draw( draw_half_radius=draw_half_radius, ) + if draw_empty_sites: + layout_ids = list(self.layout.traps_dict.keys()) + empty_layout = self.layout.define_register(*layout_ids, qubit_ids=layout_ids) + empty_qubit_colors={trap: empty_color or "r" for trap in layout_ids} + breakpoint() + empty_pos = empty_layout._coords_arr.as_array(detach=True) + pos = self._coords_arr.as_array(detach=True) if custom_ax is None: custom_ax = cast( plt.Axes, self._initialize_fig_axes( - pos, + empty_pos if draw_empty_sites else pos, blockade_radius=blockade_radius, - draw_half_radius=draw_half_radius, - )[1], + draw_half_radius=draw_half_radius + )[1], ) - super()._draw_2D( - custom_ax, - pos, - self._ids, - with_labels=with_labels, + + + draw_kwargs = dict( + ax=custom_ax, blockade_radius=blockade_radius, draw_graph=draw_graph, draw_half_radius=draw_half_radius, + ) + + if draw_empty_sites: + empty_layout._draw_2D( + ids=empty_layout._ids, + pos=empty_pos, + qubit_colors=empty_qubit_colors, + with_labels=False, + label_name="empty", + **draw_kwargs, + ) + + super()._draw_2D( + ids=self._ids, + pos=pos, qubit_colors=qubit_colors, + with_labels=with_labels, + **draw_kwargs, ) if fig_name is not None: From 1b1df5ffb28157587243415db284539a5799a1fc Mon Sep 17 00:00:00 2001 From: Briac Flesselles Date: Sat, 26 Oct 2024 15:49:09 +0200 Subject: [PATCH 2/3] added option do draw empty sites\nformatting --- pulser-core/pulser/register/register.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/pulser-core/pulser/register/register.py b/pulser-core/pulser/register/register.py index 7951282d3..a983d4ba9 100644 --- a/pulser-core/pulser/register/register.py +++ b/pulser-core/pulser/register/register.py @@ -455,8 +455,12 @@ def draw( if draw_empty_sites: layout_ids = list(self.layout.traps_dict.keys()) - empty_layout = self.layout.define_register(*layout_ids, qubit_ids=layout_ids) - empty_qubit_colors={trap: empty_color or "r" for trap in layout_ids} + empty_layout = self.layout.define_register( + *layout_ids, qubit_ids=layout_ids + ) + empty_qubit_colors = { + trap: empty_color or "r" for trap in layout_ids + } breakpoint() empty_pos = empty_layout._coords_arr.as_array(detach=True) @@ -469,9 +473,8 @@ def draw( blockade_radius=blockade_radius, draw_half_radius=draw_half_radius )[1], - ) - - + ) + draw_kwargs = dict( ax=custom_ax, blockade_radius=blockade_radius, @@ -485,10 +488,10 @@ def draw( pos=empty_pos, qubit_colors=empty_qubit_colors, with_labels=False, - label_name="empty", + label_name="empty", **draw_kwargs, ) - + super()._draw_2D( ids=self._ids, pos=pos, From f392db65b08dfc86895273367010523fa232f525 Mon Sep 17 00:00:00 2001 From: Briac Flesselles Date: Sat, 26 Oct 2024 15:38:30 +0200 Subject: [PATCH 3/3] typing hints --- pulser-core/pulser/register/register.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/pulser-core/pulser/register/register.py b/pulser-core/pulser/register/register.py index a983d4ba9..f1e67947f 100644 --- a/pulser-core/pulser/register/register.py +++ b/pulser-core/pulser/register/register.py @@ -445,7 +445,6 @@ def draw( This representation is preferred over drawing the full Rydberg radius because it helps in seeing the interactions between atoms. """ - super()._draw_checks( len(self._ids), blockade_radius=blockade_radius, @@ -454,14 +453,19 @@ def draw( ) if draw_empty_sites: + if self.layout is None: + raise ValueError( + "The register must have an associated RegisterLayout " + "to draw the empty sites." + ) layout_ids = list(self.layout.traps_dict.keys()) empty_layout = self.layout.define_register( *layout_ids, qubit_ids=layout_ids ) - empty_qubit_colors = { + breakpoint() + empty_qubit_colors : Mapping[Union[int, str], str] = { trap: empty_color or "r" for trap in layout_ids } - breakpoint() empty_pos = empty_layout._coords_arr.as_array(detach=True) pos = self._coords_arr.as_array(detach=True) @@ -483,7 +487,7 @@ def draw( ) if draw_empty_sites: - empty_layout._draw_2D( + RegDrawer._draw_2D(empty_layout, ids=empty_layout._ids, pos=empty_pos, qubit_colors=empty_qubit_colors,