Skip to content

Commit

Permalink
[Cleanup] Getting rid of magic constants
Browse files Browse the repository at this point in the history 
The notebook uses a magic constant RADIUS=5.001um, without explaining
where it comes from. We replace this by a value computed from the device,
which makes more sense. This simplifies both the code and the notebook.
  • Loading branch information
@Yoric
Yoric committed Jan 8, 2025
1 parent f8ee2f9 commit d70206a
Show file tree
Hide file tree
Showing 3 changed files with 99 additions and 74 deletions.
58 changes: 32 additions & 26 deletions examples/pipeline.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,18 @@
},
{
"cell_type": "code",
"execution_count": 13,
"execution_count": 1,
"metadata": {},
"outputs": [],
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/tmp/ipykernel_1614903/2533359462.py:8: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
" from tqdm.autonotebook import tqdm\n"
]
}
],
"source": [
"from __future__ import annotations\n",
"\n",
Expand All @@ -41,7 +50,7 @@
},
{
"cell_type": "code",
"execution_count": 14,
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -51,7 +60,7 @@
},
{
"cell_type": "code",
"execution_count": 15,
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -66,24 +75,22 @@
"\n",
"QEK lets researchers embed _graphs_ on Quantum Devices. To do this, we need to give these graphs a geometry (positions in\n",
"space) and to confirm that the geometry is compatible with a Quantum Device. Here, our dataset consists in molecules (represented\n",
"as graphs). To simplify things, QEK comes with a dedicated class `qek_datatools.MoleculeGraph` that adds a geometry to the graphs.\n",
"as graphs). To simplify things, QEK comes with a dedicated class `qek_datatools.MoleculeGraph` that use physical tools to compute\n",
"a reasonable geometry from molecular data for a specific Quantum Device.\n",
"\n",
"One of the core ideas behind QEK is that each nodes (aka atoms) in a graph (aka molecule) from the dataset is represented by one\n",
"cold atom on the Device and if two nodes are joined by an edge, their cold atoms must be close to each other. In geometrical terms,\n",
"this means that the `MoleculeGraph` must be a _disk graph_, with a radius of 5.001 $\\mu m$. In this notebook, for the sake of\n",
"simplicity, we simply discard graphs that are not disk graphs.\n",
" "
"As the geometry depends on the Quantum Device, we need to specify a device to use. We'll use Pulser's `AnalogDevice`, which is\n",
"a reasonable default device. If you need to adapt your code to a different device, you'll need to pick another one."
]
},
{
"cell_type": "code",
"execution_count": 16,
"execution_count": 4,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "9deae7c65fae4de9b557c916e3df32bc",
"model_id": "7befc14e41dd443ba608c19230fc5030",
"version_major": 2,
"version_minor": 0
},
Expand All @@ -97,12 +104,9 @@
],
"source": [
"list_of_graphs = []\n",
"RADIUS = 5.001\n",
"EPS = 0.01\n",
"for data in tqdm(og_ptcfm):\n",
" graph = qek_datatools.MoleculeGraph(data=data, blockade_radius=RADIUS)\n",
" if graph.is_disk_graph(radius=RADIUS+EPS):\n",
" list_of_graphs.append((graph, graph.pyg.y.item()))"
" graph = qek_datatools.MoleculeGraph(data=data, device=pl.AnalogDevice)\n",
" list_of_graphs.append((graph, graph.pyg.y.item()))"
]
},
{
Expand All @@ -111,23 +115,25 @@
"source": [
"## Create a Pulser sequence\n",
"\n",
"Once the embedding is found, we create a Pulser Sequence that can be interpreted by a Quantum Device. A Sequence consists of a **register** (i.e. a geometry of cold atoms on the device) and **pulse**s. Sequences need to be designed for a specific device, so our graph object offers a method `compute_sequence` that does exactly that."
"Once the embedding is found, we create a Pulser Sequence that can be interpreted by a Quantum Device. A Sequence consists of a **register** (i.e. a geometry of cold atoms on the device) and **pulse**s. Sequences need to be designed for a specific device, so our graph object offers a method `compute_sequence` that does exactly that.\n",
"\n",
"Not all geometries can be embedded on a given device. In this notebook, for the sake of simplicity, we simply discard graphs that cannot be trivially embedded."
]
},
{
"cell_type": "code",
"execution_count": 17,
"execution_count": 5,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "f3e5e9dc5b864843ba85a4281951072b",
"model_id": "4296d31619e64d8a904660b67d68d00f",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
" 0%| | 0/294 [00:00<?, ?it/s]"
" 0%| | 0/349 [00:00<?, ?it/s]"
]
},
"metadata": {},
Expand All @@ -138,9 +144,9 @@
"dataset_sequence = []\n",
"\n",
"for graph, target in tqdm(list_of_graphs):\n",
" # Some graph are not compatible with AnalogDevice\n",
" if graph.is_embeddable(device=pl.AnalogDevice):\n",
" dataset_sequence.append((graph.compute_sequence(device=pl.AnalogDevice), target))\n"
" # Some graph are not compatible with AnalogDevice, just skip them.\n",
" if graph.is_embeddable():\n",
" dataset_sequence.append((graph.compute_sequence(), target))\n"
]
},
{
Expand All @@ -152,13 +158,13 @@
},
{
"cell_type": "code",
"execution_count": 18,
"execution_count": 6,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "36483a40288c431eb0482e66325b375a",
"model_id": "acf47b6f06f64c9082b8bc4638556187",
"version_major": 2,
"version_minor": 0
},
Expand Down
86 changes: 50 additions & 36 deletions qek/data/datatools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,22 +92,29 @@ def load_dataset(file_path: str) -> list[ProcessedData]:
]


EPSILON_DISTANCE_UM = 0.01


class BaseGraph:
"""
A graph being prepared for embedding on a quantum device.
"""

def __init__(self, data: pyg_data.Data):
def __init__(self, data: pyg_data.Data, device: pl.devices.Device):
"""
Create a graph from geometric data.
Args:
data: A homogeneous graph, in PyTorch Geometric format. Unchanged.
It MUST have attributes 'pos'
device: The device for which the graph is prepared.
"""
if not hasattr(data, "pos"):
raise AttributeError("The graph should have an attribute 'pos'.")

# The device for which the graph is prepared.
self.device = device

# The graph in torch geometric format.
self.pyg = data.clone()

Expand Down Expand Up @@ -135,66 +142,69 @@ def is_disk_graph(self, radius: float) -> bool:
`True` if the graph is a disk graph with the specified radius,
`False` otherwise.
"""

if self.pyg.num_nodes == 0 or self.pyg.num_nodes is None:
return False

# Check if the graph is connected.
if len(self.nx_graph) == 0 or not nx.is_connected(self.nx_graph):
return False

# Check the distances between all pairs of nodes.
pos = self.pyg.pos
for u, v in nx.non_edges(self.nx_graph):
distance = np.linalg.norm(np.array(pos[u]) - np.array(pos[v]))
if distance <= radius:
distance_um = np.linalg.norm(np.array(pos[u]) - np.array(pos[v]))
if distance_um <= radius:
# These disjointed nodes would interact with each other, so
# this is not an embeddable graph.
return False

for u, v in self.nx_graph.edges():
distance = np.linalg.norm(np.array(pos[u]) - np.array(pos[v]))
if distance > radius:
distance_um = np.linalg.norm(np.array(pos[u]) - np.array(pos[v]))
if distance_um > radius:
# These joined nodes would not interact with each other, so
# this is not an embeddable graph.
return False

return True

def is_embeddable(self, device: pl.devices.Device) -> bool:
def is_embeddable(self) -> bool:
"""
A predicate to check if the graph can be embedded in the
quantum device.
For a graph to be embeddable on a device, all the following
criteria must be fulfilled:
- the graph must be non-empty;
- the device must have at least as many atoms as the graph has
nodes;
- the device must be physically large enough to place all the
nodes (device.max_radial_distance);
- the nodes must be distant enough that quantum interactions
may take place (device.min_atom_distance)
Args:
device (pulser.devices.Device): the device
Returns:
bool: True if possible, False if not
"""

# Check the number of atoms
if self.pyg.num_nodes > device.max_atom_num:
# Reject empty graphs.
if self.pyg.num_nodes == 0 or self.pyg.num_nodes is None:
return False

# Reject graphs that have more nodes than can be represented
# on the device.
if self.pyg.num_nodes > self.device.max_atom_num:
return False

# Check the distance from the center
pos_graph = self.pyg.pos
distance_from_center = np.linalg.norm(pos_graph, ord=2, axis=-1)
if any(distance_from_center > device.max_radial_distance):
pos = self.pyg.pos
distance_from_center = np.linalg.norm(pos, ord=2, axis=-1)
if any(distance_from_center > self.device.max_radial_distance):
return False

# Check the distance between nodes.
nodes = list(self.nx_graph.nodes)
for i in range(0, len(nodes)):
for j in range(i + 1, len(nodes)):
dist = np.linalg.norm(pos_graph[i] - pos_graph[j], ord=2)
if dist < device.min_atom_distance:
return False
# Check distance between nodes
if not self.is_disk_graph(self.device.min_atom_distance + EPSILON_DISTANCE_UM):
return False

for u, v in self.nx_graph.edges():
distance_um = np.linalg.norm(np.array(pos[u]) - np.array(pos[v]))
if distance_um < self.device.min_atom_distance:
# These nodes are too close to each other, preventing quantum
# interactions on the device.
return False

return True

Expand All @@ -206,17 +216,17 @@ def compute_register(self) -> pl.Register:
"""
return pl.Register.from_coordinates(coords=self.pyg.pos)

def compute_sequence(self, device: pl.devices.Device) -> pl.Sequence:
def compute_sequence(self) -> pl.Sequence:
"""
Compile a Quantum Sequence from a graph for a specific device.
Raises:
ValueError if the graph cannot be embedded on the given device.
"""
if not self.is_embeddable(device):
raise ValueError(f"The graph is not compatible with {device}")
if not self.is_embeddable():
raise ValueError(f"The graph is not compatible with {self.device}")
reg = self.compute_register()
seq = pl.Sequence(register=reg, device=device)
seq = pl.Sequence(register=reg, device=self.device)

# See the companion paper for an explanation on these constants.
Omega_max = 1.0 * 2 * np.pi
Expand Down Expand Up @@ -272,7 +282,7 @@ class MoleculeGraph(BaseGraph):
def __init__(
self,
data: pyg_data.Data,
blockade_radius: float,
device: pl.devices.Device,
node_mapping: dict[int, str] = PTCFM_ATOM_NAMES,
edge_mapping: dict[int, Chem.BondType] = PTCFM_BOND_TYPES,
):
Expand All @@ -294,7 +304,7 @@ def __init__(
"""
pyg = data.clone()
pyg.pos = None # Placeholder
super().__init__(pyg)
super().__init__(pyg, device)

# Reconstruct the molecule.
tmp_mol = graph_to_mol(
Expand All @@ -306,10 +316,14 @@ def __init__(
# Extract the geometry.
AllChem.Compute2DCoords(tmp_mol, useRingTemplates=True)
pos = tmp_mol.GetConformer().GetPositions()[..., :2] # Convert to 2D

# Scale the geometry so that the longest edge is as long as
# `device.min_atom_distance`.
dist_list = []
for start, end in self.nx_graph.edges():
dist_list.append(np.linalg.norm(pos[start] - pos[end]))
norm_factor = np.max(dist_list)
pos = pos * device.min_atom_distance / norm_factor

# Finally, store the geometry.
self.pyg.pos = pos * blockade_radius / norm_factor
# Finally, store the position.
self.pyg.pos = pos
Loading

0 comments on commit d70206a

Please sign in to comment.