diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 4ad284c..480073c 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -1,51 +1,7 @@ -// For format details, see https://aka.ms/devcontainer.json. For config options, see the -// README at: https://github.com/devcontainers/templates/tree/main/src/debian +// For format details, see https://aka.ms/devcontainer.json. { - "name": "Ocean Development Environment", - - // python 3.11 on debian, with latest Ocean and optional packages - // source repo: https://github.com/dwavesystems/ocean-dev-docker - "image": "docker.io/dwavesys/ocean-dev:latest", - - // install repo pip requirements (only if present) on content update - "updateContentCommand": "[ ! -r requirements.txt ] || pip install -r requirements.txt", - - // forward/expose container services (relevant only when run locally) - "forwardPorts": [ - // dwave-inspector web app - 18000, 18001, 18002, 18003, 18004, - // OAuth connect redirect URIs - 36000, 36001, 36002, 36003, 36004 - ], - - "portsAttributes": { - "18000-18004": { - "label": "D-Wave Problem Inspector", - "requireLocalPort": true - }, - "36000-36004": { - "label": "OAuth 2.0 authorization code redirect URI", - "requireLocalPort": true - } - }, - - // Configure tool-specific properties. - "customizations": { - // Configure properties specific to VS Code. - "vscode": { - // Set *default* container specific settings.json values on container create. - "settings": { - "workbench": { - "editorAssociations": { - "*.md": "vscode.markdown.preview.editor" - }, - "startupEditor": "readme" - } - }, - "extensions": [ - "ms-python.python", - "ms-toolsai.jupyter" - ] - } - } + // Debian stable with the second-latest Python, latest Ocean, and optional dev packages. + // Docker image source: https://github.com/dwavesystems/ocean-dev-docker. + // Devcontainer config: https://github.com/dwavesystems/ocean-devcontainer. + "image": "docker.io/dwavesys/ocean-dev:latest" } diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000..ac3e9dc --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,24 @@ + + +### Associated GitHub Issue + + +### Feature Implemented/Bugs Fixed + + +### Additional Information + + +### Accessibility Score + + +### AI Generation Disclosure + diff --git a/requirements.txt b/requirements.txt index cd4504d..e74f4c8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ dash[diskcache]~=3.2 dash-bootstrap-components~=2.0 +dwave-experimental==2026.5.27 dwave-ocean-sdk~=9.0 dwave-pytorch-plugin~=0.3 einops~=0.8 diff --git a/src/model_wrapper.py b/src/model_wrapper.py index 19596a1..7699e49 100755 --- a/src/model_wrapper.py +++ b/src/model_wrapper.py @@ -18,6 +18,7 @@ from typing import Optional import numpy as np +from dimod.exceptions import BinaryQuadraticModelStructureError import plotly.express as px import torch import yaml @@ -306,14 +307,25 @@ def step(self, batch: tuple[torch.Tensor, torch.Tensor], epoch: int) -> torch.Te self.losses["mse_losses"].append(mse_loss.item()) with torch.no_grad(): - samples = self._grbm.sample( # type: ignore - sampler=self.sampler, - prefactor=self.PREFACTOR, - linear_range=self.linear_range, - quadratic_range=self.quadratic_range, - device=spins.device, - sample_params=self.sampler_kwargs, - ) + try: + samples = self._grbm.sample( # type: ignore + sampler=self.sampler, + prefactor=self.PREFACTOR, + linear_range=self.linear_range, + quadratic_range=self.quadratic_range, + device=spins.device, + sample_params=self.sampler_kwargs, + ) + except BinaryQuadraticModelStructureError: + self._rebuild_sampler() + samples = self._grbm.sample( # type: ignore + sampler=self.sampler, + prefactor=self.PREFACTOR, + linear_range=self.linear_range, + quadratic_range=self.quadratic_range, + device=spins.device, + sample_params=self.sampler_kwargs, + ) spins = spins.reshape(-1, spins.shape[-1]) @@ -329,17 +341,31 @@ def step(self, batch: tuple[torch.Tensor, torch.Tensor], epoch: int) -> torch.Te # train boltzmann machine if train_grbm(self._tpar["opt_step"], epoch): self._grbm_optimizer.zero_grad() - grbm_loss, self._tpar["sample_set"] = nll_loss( - spins=spins.detach(), - grbm=self._grbm, - sampler=self.sampler, - sampler_kwargs=self.sampler_kwargs, - linear_range=self.linear_range, - quadratic_range=self.quadratic_range, - prefactor=self.PREFACTOR, - persistent_qpu_sample_helper=self._tpar["persistent_qpu_sample_helper"], - sample_set=self._tpar["sample_set"], - ) + try: + grbm_loss, self._tpar["sample_set"] = nll_loss( + spins=spins.detach(), + grbm=self._grbm, + sampler=self.sampler, + sampler_kwargs=self.sampler_kwargs, + linear_range=self.linear_range, + quadratic_range=self.quadratic_range, + prefactor=self.PREFACTOR, + persistent_qpu_sample_helper=self._tpar["persistent_qpu_sample_helper"], + sample_set=self._tpar["sample_set"], + ) + except BinaryQuadraticModelStructureError: + self._rebuild_sampler() + grbm_loss, self._tpar["sample_set"] = nll_loss( + spins=spins.detach(), + grbm=self._grbm, + sampler=self.sampler, + sampler_kwargs=self.sampler_kwargs, + linear_range=self.linear_range, + quadratic_range=self.quadratic_range, + prefactor=self.PREFACTOR, + persistent_qpu_sample_helper=self._tpar["persistent_qpu_sample_helper"], + sample_set=self._tpar["sample_set"], + ) grbm_loss.backward() self._grbm_optimizer.step() @@ -352,6 +378,34 @@ def step(self, batch: tuple[torch.Tensor, torch.Tensor], epoch: int) -> torch.Te return mse_loss + def _rebuild_sampler(self) -> None: + """Reconnect to the QPU and rebuild the sampler after a structure error. + + Used as a Zephyr-only fallback when the cached embedding references qubits + or couplers that are no longer available on the QPU. ``EmbeddingComposite`` + re-embeds the GRBM's logical graph dynamically onto the current working graph, + so the GRBM weights remain valid and no retraining is needed. + """ + from dwave.system import DWaveSampler, EmbeddingComposite + + print( + f"Sampler structure error detected; reconnecting to QPU '{self.qpu}' and rebuilding sampler." + ) + qpu_sampler = DWaveSampler(solver=self.qpu) + qpu_topology = qpu_sampler.properties["topology"]["type"] + + if qpu_topology != "zephyr": + raise RuntimeError( + f"Sampler rebuild is only supported for Zephyr QPUs (got '{qpu_topology}')." + ) + + # EmbeddingComposite finds a valid minor embedding of the GRBM's logical + # graph on the current working QPU topology without requiring retraining. + self.sampler = EmbeddingComposite(qpu_sampler) + self.linear_range = qpu_sampler.properties["h_range"] + self.quadratic_range = qpu_sampler.properties["j_range"] + print(f"Sampler rebuilt with EmbeddingComposite for Zephyr QPU '{self.qpu}'.") + def generate_output(self, latent_qpu_file: str, sharpen: bool = False, save_to_file: str = "") -> go.Figure: """Generate output images from trained model. @@ -366,14 +420,25 @@ def generate_output(self, latent_qpu_file: str, sharpen: bool = False, save_to_f self._grbm.eval() with torch.no_grad(): - samples = self._grbm.sample( - self.sampler, - prefactor=self.PREFACTOR, - device=self._device, - linear_range=self.linear_range, - quadratic_range=self.quadratic_range, - sample_params=self.sampler_kwargs, - ) + try: + samples = self._grbm.sample( + self.sampler, + prefactor=self.PREFACTOR, + device=self._device, + linear_range=self.linear_range, + quadratic_range=self.quadratic_range, + sample_params=self.sampler_kwargs, + ) + except BinaryQuadraticModelStructureError: + self._rebuild_sampler() + samples = self._grbm.sample( + self.sampler, + prefactor=self.PREFACTOR, + device=self._device, + linear_range=self.linear_range, + quadratic_range=self.quadratic_range, + sample_params=self.sampler_kwargs, + ) with open(latent_qpu_file, "w") as f: json.dump(samples[0].tolist(), f) diff --git a/src/utils/callback_helpers.py b/src/utils/callback_helpers.py index 5686cc0..4c5aef1 100644 --- a/src/utils/callback_helpers.py +++ b/src/utils/callback_helpers.py @@ -30,7 +30,7 @@ from torchvision.utils import save_image from demo_configs import GENERATE_NEW_MODEL_DIAGRAM, GRAPH_COLORS, SHARPEN_OUTPUT, THEME_COLOR_SECONDARY -from src.utils.common import get_graph_mapping, greedy_get_subgraph +from src.utils.common import get_graph_mapping, greedy_get_subgraph, _try_zephyr_sublattice_fallback MODEL_PATH = Path("models") JSON_FILE_DIR = "generated_json" @@ -358,13 +358,24 @@ def generate_model_fig( """ qpu = DWaveSampler(solver=qpu) qpu_graph = qpu.to_networkx_graph() - subgraph = greedy_get_subgraph(n_nodes=n_latents, random_seed=random_seed, graph=qpu_graph) - _, mapping = get_graph_mapping(subgraph) + qpu_topology = qpu.properties["topology"]["type"] - latent_mapping = [mapping[node] for node in subgraph.nodes()] + try: + subgraph = greedy_get_subgraph(n_nodes=n_latents, random_seed=random_seed, graph=qpu_graph) + _, mapping = get_graph_mapping(subgraph) + latent_mapping = [mapping[node] for node in subgraph.nodes()] + except Exception: + if qpu_topology != "zephyr": + raise + fallback = _try_zephyr_sublattice_fallback( + n_nodes=n_latents, qpu_graph=qpu_graph, random_seed=random_seed + ) + if fallback is None: + raise + subgraph, _ = fallback + latent_mapping = list(subgraph.nodes()) qpu_shape = qpu.properties["topology"]["shape"][0] - qpu_topology = qpu.properties["topology"]["type"] if qpu_topology == "pegasus": node_coords = dnx.drawing.pegasus_layout(dnx.pegasus_graph(qpu_shape), crosses=True) diff --git a/src/utils/common.py b/src/utils/common.py index b1b6d13..616ae0d 100644 --- a/src/utils/common.py +++ b/src/utils/common.py @@ -14,9 +14,13 @@ import random from typing import Callable, Literal, Optional +import dwave_networkx as dnx import networkx as nx import torch +from dwave.experimental.embedding_methods import zephyr_quotient_search from dwave.system import DWaveSampler, FixedEmbeddingComposite +from minorminer import find_embedding +from minorminer.utils.parallel_embeddings import find_sublattice_embeddings def greedy_get_subgraph( @@ -83,6 +87,112 @@ def greedy_get_subgraph( return subgraph + +def _try_zephyr_sublattice_fallback( + n_nodes: int, + qpu_graph: nx.Graph, + random_seed: Optional[int], +) -> Optional[tuple[nx.Graph, dict]]: + """Attempt to find a defect-free Zephyr sublattice embedding as a backup. + + Follows the approach demonstrated in + https://github.com/dwavesystems/dwave-experimental/blob/main/examples/fully_yielded_zephyr_subgraph.py: + 1. Locate a complete Zephyr sublattice in the defective QPU graph via + ``find_sublattice_embeddings``. + 2. Relabel it to canonical coordinates for ``zephyr_quotient_search``. + 3. Embed a source Zephyr graph (with ``n_nodes`` nodes) into the sublattice. + 4. Refine with ``find_embedding`` if quotient search did not reach full yield. + 5. Map the resulting chain embedding back to original QPU qubit labels. + + Args: + n_nodes: Required number of logical nodes (equal to the latent space size). + qpu_graph: The defective QPU graph as returned by ``qpu.to_networkx_graph()``. + random_seed: Random seed forwarded to sublattice search. + + Returns: + ``(logical_graph, chain_embedding)`` where ``logical_graph`` has integer + nodes ``0..n_nodes-1`` and ``chain_embedding`` maps each integer node to a + list of physical QPU qubits suitable for ``FixedEmbeddingComposite``. + Returns ``None`` when no suitable embedding can be found. + """ + qpu_rows = qpu_graph.graph.get("rows") + qpu_tile = qpu_graph.graph.get("tile", 4) + labels = qpu_graph.graph.get("labels") + + if not isinstance(qpu_rows, int) or qpu_rows <= 0: + return None + + use_coords = labels == "coordinate" + zephyr_tile = dnx.zephyr_graph(qpu_rows, qpu_tile, coordinates=use_coords) + + tile_embeddings = find_sublattice_embeddings( + S=zephyr_tile, + T=qpu_graph, + max_num_emb=1, + one_to_iterable=False, + seed=random_seed, + ) + if not tile_embeddings: + print("Zephyr fallback: no complete sublattice found in defective QPU graph.") + return None + + tile_embedding = tile_embeddings[0] # {tile_node: physical_qubit} + + sublattice_nodes = set(tile_embedding.values()) + target_sub = qpu_graph.subgraph(sublattice_nodes).copy() + inv_map = {phys: tile_node for tile_node, phys in tile_embedding.items()} + target_sub = nx.relabel_nodes(target_sub, inv_map, copy=True) + target_sub.graph.update( + family="zephyr", + rows=qpu_rows, + tile=qpu_tile, + labels="coordinate" if use_coords else "int", + ) + + source = None + for t_s in range(qpu_tile, 0, -1): + candidate = dnx.zephyr_graph(qpu_rows, t_s, coordinates=use_coords) + if candidate.number_of_nodes() == n_nodes: + source = candidate + break + + if source is None: + print( + f"Zephyr fallback: no source Zephyr graph with {n_nodes} nodes found " + f"for QPU rows={qpu_rows} tile={qpu_tile}." + ) + return None + + emb, metadata = zephyr_quotient_search(source, target_sub, yield_type="edge") + best_emb = emb + + if metadata.final_num_yielded < metadata.max_num_yielded: + print( + f"Zephyr fallback: quotient search placed {metadata.final_num_yielded}/" + f"{metadata.max_num_yielded} edges; refining with find_embedding." + ) + refined = find_embedding(S=source, T=target_sub, initial_chains=emb, timeout=50) + if refined: + best_emb = refined + + if not best_emb: + print("Zephyr fallback: embedding search returned no result.") + return None + + chain_embedding_in_qpu = { + s_node: [tile_embedding[v] for v in chain] + for s_node, chain in best_emb.items() + } + + source_to_int = {node: i for i, node in enumerate(source.nodes())} + logical_graph = nx.relabel_nodes(source, source_to_int) + int_chain_embedding = { + source_to_int[s]: chains for s, chains in chain_embedding_in_qpu.items() + } + + return logical_graph, int_chain_embedding + + def get_graph_mapping(graph: Optional[nx.Graph]) -> tuple[nx.Graph, dict]: """Maps a graph of the QPU to the encoded latent data. @@ -122,10 +232,23 @@ def get_sampler_and_sampler_kwargs( qpu = DWaveSampler(solver=qpu) qpu_graph = qpu.to_networkx_graph() - subgraph = greedy_get_subgraph(n_nodes=n_latents, random_seed=random_seed, graph=qpu_graph) - mapped_graph, mapping = get_graph_mapping(subgraph) - sampler = FixedEmbeddingComposite(qpu, {l_: [p] for p, l_ in mapping.items()}) + try: + subgraph = greedy_get_subgraph(n_nodes=n_latents, random_seed=random_seed, graph=qpu_graph) + mapped_graph, mapping = get_graph_mapping(subgraph) + sampler = FixedEmbeddingComposite(qpu, {l_: [p] for p, l_ in mapping.items()}) + except Exception: + qpu_topology = qpu.properties["topology"]["type"] + if qpu_topology != "zephyr": + raise + fallback = _try_zephyr_sublattice_fallback( + n_nodes=n_latents, qpu_graph=qpu_graph, random_seed=random_seed + ) + if fallback is None: + raise + mapped_graph, chain_embedding = fallback + sampler = FixedEmbeddingComposite(qpu, chain_embedding) + linear_range, quadratic_range = qpu.properties["h_range"], qpu.properties["j_range"] sampler_kwargs = dict( num_reads=num_reads,