Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draw layout and register on the same plot #772

Open
wants to merge 3 commits into
base: develop
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 44 additions & 9 deletions pulser-core/pulser/register/register.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

May I suggest to have the default color instead of None ? I also suggest we have the default color to be gray instead of red, to be consistent with the current color of a layout ?

Another idea I had, which might make the life of the user easier, would be to have the empty sites displayed as a trap in the layout, a black circle filled with white. What do you think of this alternative ? An advantage I see with this is that there will be no conflict with qubit_colors (in the current implementation, what happens if the user provide qubit_colors=empty_color ? The visualization will not be great...)

image

Suggested change
empty_color: Optional[str] = None,
empty_color: str = "gray",

) -> None:
"""Draws the entire register.

Expand Down Expand Up @@ -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'.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As in my comment above, I am wondering if empty_color is necessary, and I would perhaps suggest to have default to "gray" instead of "r".


Note:
When drawing half the blockade radius, we say there is a blockade
Expand All @@ -448,25 +452,56 @@ def draw(
draw_half_radius=draw_half_radius,
)

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
)
breakpoint()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should not leave this breakpoint() in the code :)

empty_qubit_colors : Mapping[Union[int, str], str] = {
trap: empty_color or "r" for trap in layout_ids
}
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],
)
super()._draw_2D(
custom_ax,
pos,
self._ids,
with_labels=with_labels,
draw_half_radius=draw_half_radius
)[1],
)

draw_kwargs = dict(
ax=custom_ax,
blockade_radius=blockade_radius,
draw_graph=draw_graph,
draw_half_radius=draw_half_radius,
)

if draw_empty_sites:
RegDrawer._draw_2D(empty_layout,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you can use super instead of RegDrawer :)

ids=empty_layout._ids,
pos=empty_pos,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Your implementation looks good. It enables you to implement the alternative I have presented above, where the empty spots are circles filled with white. To implement it fully, you might want to take out the positions of the atoms in the Register in empty_pos, because at the moment you are super-imposing two symbols for these filled positions and the symbol of the circle with white inside is larger than the symbol for the atoms (a disk of green by default, or a color defined in qubit_colors).

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:
Expand Down