diff --git a/dwave/experimental/lattice_utils/__init__.py b/dwave/experimental/lattice_utils/__init__.py new file mode 100644 index 0000000..bd91f26 --- /dev/null +++ b/dwave/experimental/lattice_utils/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dwave.experimental.lattice_utils import experiment, lattice, observable + +__all__ = ["experiment", "lattice", "observable"] diff --git a/dwave/experimental/lattice_utils/experiment/__init__.py b/dwave/experimental/lattice_utils/experiment/__init__.py new file mode 100644 index 0000000..667811a --- /dev/null +++ b/dwave/experimental/lattice_utils/experiment/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dwave.experimental.lattice_utils.experiment.experiment import * +from dwave.experimental.lattice_utils.experiment.fast_anneal_experiment import * +from dwave.experimental.lattice_utils.experiment.samplercall import * diff --git a/dwave/experimental/lattice_utils/experiment/experiment.py b/dwave/experimental/lattice_utils/experiment/experiment.py new file mode 100644 index 0000000..524c66a --- /dev/null +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -0,0 +1,850 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import copy +import lzma +import os +import pickle +import tempfile +import time +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, TYPE_CHECKING + +import dimod +import numpy as np + +try: + from tqdm.auto import tqdm +except ImportError: + tqdm = None + +if TYPE_CHECKING: + from tqdm.std import tqdm as tqdm_bar + +from dwave.experimental.lattice_utils.experiment.samplercall import SamplerCall +from dwave.experimental.lattice_utils.lattice.lattice import Lattice +from dwave.experimental.lattice_utils.observable import ( + BitpackedSpins, + CouplerCorrelation, + CouplerFrustration, + QubitMagnetization, + ReferenceEnergy, + SampleEnergy, +) + +__all__ = ['Experiment', 'ExperimentConfig'] + +DW_TEAL = "#17bebb" +DW_BLUE = "#2a7de1" +DW_ORANGE = "#f37820" + + +@dataclass +class ExperimentConfig: + """Container for the parameters that define an experiment.""" + + signed_energy_scale: float = 1.0 + num_reads: int = 100 + anneal_time: float = 1.0 + num_random_instances: int | None = 1 + readout_thermalization: int = 100 + flux_bias_shim_step: float = 0.0 + coupler_shim_step: float = 0.0 + anneal_offset_shim_step: float = 0.0 + target_magnetization: float = 0.0 + + +class Experiment: + """Base class for running experiments on lattice instances. + + Includes common functionality for managing parameters, running iterations, + parsing results, and saving data. + + Args: + inst: The lattice instance to run the experiment on. + sampler: The dimod sampler to use for sampling. + max_iterations: The maximum number of iterations to run the experiment for. + config: An ExperimentConfig object containing experiment parameters. + """ + + def __init__( + self, + *, + inst: Lattice, + sampler: dimod.Sampler, + reference_energy_sampler: dimod.Sampler | None = None, + reference_energy_sampler_kwargs: dict[str, Any] | None = None, + max_iterations: int | None = None, + config: ExperimentConfig | None = None, + ): + if config is None: + config = ExperimentConfig() + + self.inst = inst + self.sampler = sampler + self.reference_energy_sampler = reference_energy_sampler + self.reference_energy_sampler_kwargs = reference_energy_sampler_kwargs + self.param = vars(config).copy() + self.experiment_results_root = inst.data_root / "results" + self.data_path = None + self.run_index = 0 + self.config = config + self.max_iterations = max_iterations + self.already_initialized: bool = False + self.observables_to_collect = [ + QubitMagnetization(), + CouplerCorrelation(), + CouplerFrustration(), + SampleEnergy(), + BitpackedSpins(), + ReferenceEnergy(), + ] + + def load_results( + self, + num_iterations: int = 100, + start_iteration: int | None = None, + result_fields: list[str] | None = None, + ignore_shim: bool = False, + ) -> list[dict[str, Any]]: + """Load results from the highest-numbered iterations of the experiment. + + Args: + num_iterations: Maximum number of iterations to load. + start_iteration: If provided, load results starting from this + iteration index. Otherwise the most recent ``num_iterations`` + results are loaded. + result_fields: Subset of fields to extract from each result file. If + ``None``, all fields present in the first result file are used. + ignore_shim: If true, the ``shim_data`` field is removed from the + returned results. + + Returns: + A list of dictionaries containing the results for each iteration. + """ + fnlist = self._get_sorted_results_file_list() + if start_iteration is not None: + fnlist = fnlist[max(start_iteration, 0) : max(start_iteration + num_iterations, 0)] + else: + fnlist = fnlist[-num_iterations:] + + results = [] + for filename in fnlist: + + try: + with lzma.open(filename, "rb") as f: + data = pickle.load(f) + except lzma.LZMAError as e: + raise lzma.LZMAError(f"Failing to load {filename}", e) + + if result_fields is None: + result_fields = list(data.keys()) + if ignore_shim: + result_fields.remove("shim_data") + + results.append({k: data[k] for k in result_fields}) + + return results + + def apply_param(self, param: dict[str, float]) -> None: + """Apply a parameter configuration to the experiment. + + Parameters are formatted to ensure filename consistency, which can be + important for loading data. + + Args: + param: Dictionary of parameter values to apply to the experiment. + """ + param = self._format_parameter_list([param])[0] + + # anneal_time and anneal_schedule are mutually exclusive + if "anneal_schedule" in param: + self.param.pop("anneal_time", None) + elif "anneal_time" in param: + self.param.pop("anneal_schedule", None) + + self.param.update(param) + + self.data_path = self.experiment_results_root / self._get_relative_data_path() + self.already_initialized = self._prepare_run_index() + + def run_iteration( + self, + parameter_list: list, + progress: bool = False, + scaling_factor: float = 1.0, + ) -> bool: + """Run one experiment iteration for each parameter set in ``parameter_list``. + + For each parametrization, this method applies the parameters, builds the + sampler call, submits the sampling job, waits for completion, parses the + returned results, updates the shim, and saves the results. + + Args: + parameter_list: List of parameter dictionaries to run. + progress: If true, displays a progress bar for waiting on results. + scaling_factor: A multiplicative factor to apply to the BQM before sampling. + + Returns: + A boolean value corresponding to whether or not the experiment is + finished. + """ + if progress and tqdm is None: + raise ImportError("Progress reporting requires the optional 'tqdm' dependency.") + + try: + self.inst._load_embeddings(self.sampler) + except FileNotFoundError as e: + raise FileNotFoundError("No Embedding Found: ", e) from e + + if progress: + tqdm.write( + f"\n{type(self.inst).__name__}={self.inst.dimensions}, " + f"J={self.param['signed_energy_scale']}, " + f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')} " + f"({self.inst._get_instance_pathstring()}/{self._get_solver_pathstring()})" + ) + + parameter_list = self._format_parameter_list(parameter_list) + response_dict = {} + call_dict = {} + + create_bar = ( + self._make_progress_bar( + total=len(parameter_list), + desc="Creating sampler calls", + colour=DW_BLUE, + ) + if progress + else None + ) + + for index, param in enumerate(parameter_list): + self.apply_param(param) + call_dict[index] = self._build_sampler_call() + if call_dict[index] is None: + call_dict.pop(index) + else: + response_dict[index] = self.sampler.sample( + call_dict[index].bqm * scaling_factor, + **call_dict[index].sampler_params, + ) + if create_bar is not None: + create_bar.update() + + if create_bar is not None: + create_bar.close() + + if not call_dict: + if progress: + tqdm.write( + f"***\n***\nFINISHED for all {len(parameter_list)} parameterizations.\n***\n***" + ) + return True + + wait_bar = ( + self._make_progress_bar( + total=len(call_dict), + desc=" Awaiting/parsing data", + colour=DW_TEAL, + ) + if progress + else None + ) + + # Get and manage all the results + while response_dict: + made_progress = False + + for index, val in response_dict.items(): + + if val.done(): + self.apply_param(parameter_list[index]) + results = self.parse_results(call_dict[index], val) + self._update_shim(call_dict[index], results) + savedata = self._generate_data_to_save(call_dict[index], results) + self._save_results(savedata) + if wait_bar is not None: + wait_bar.update() + del response_dict[index] + made_progress = True + break + + if not made_progress: + time.sleep(0.1) # Waiting for results to come in + + if wait_bar is not None: + wait_bar.close() + + if progress: + self._print_iteration_status(call_dict, len(parameter_list)) + + return False + + def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[str, Any]: + """Parse a sampler response into per-embedding observable results. + + Args: + call: Sampler call metadata, inluding the logical BQMs + response: Raw sample set returned by the sampler. + + Returns: + Dictionary mapping observable names to their evaluated results across + embeddings. + """ + if hasattr(self.inst, "embedding_list"): + embedding_list = self.inst.embedding_list + myarr = response.samples(sorted_by=None) + sample_arrays = [myarr[:, emb].copy() for emb in embedding_list] + else: + sample_arrays = [response.samples(sorted_by=None)[:, np.arange(self.inst.num_spins)]] + + sample_set = {} + for iemb, sample_array in enumerate(sample_arrays): + sample_set[iemb] = dimod.SampleSet.from_samples_bqm( + sample_array, call.logical_bqms[iemb] + ) + + results = {} + for observable in self.observables_to_collect: + results[observable.name] = [] + for iemb, sample_array in enumerate(sample_arrays): + bqm = call.logical_bqms[iemb] + obs_result = observable.evaluate(self, bqm, sample_set[iemb]) + results[observable.name].append(obs_result) + + if isinstance(results[observable.name][0], np.ndarray): + results[observable.name] = np.asarray(results[observable.name]) + + return results + + def _make_progress_bar( + self, + *, + total: int, + desc: str, + colour: str, + bar_format: str | None = None, + initial: int | float = 0, + ) -> tqdm_bar: + """Create a tqdm progress bar with consistent formatting. + + The optional ``tqdm`` dependency is required to call this method. If + ``tqdm`` is not installed, an ImportError will be raised. + """ + if tqdm is None: + raise ImportError("Progress reporting requires the optional 'tqdm' dependency.") + + if bar_format is None: + bar_width = min(100, max(total, 20)) + bar_format = f"{{desc}}: |{{bar:{bar_width}}}{{r_bar}}{{bar:-{bar_width}b}}" + + return tqdm( + total=total, + initial=initial, + desc=desc, + bar_format=bar_format, + colour=colour, + ) + + def _print_iteration_status( + self, + call_dict: dict[int, SamplerCall], + num_params: int, + ) -> None: + """Print a summary of the iteration status, including progress and iteration ranges. + + The optional ``tqdm`` dependency is required to call this method. If + ``tqdm`` is not installed, an ImportError will be raised. + """ + if tqdm is None: + raise ImportError("Progress reporting requires the optional 'tqdm' dependency.") + + iteration_range = ( + f"Iteration range " + f"{min(call.shim_data['total_iterations'] for call in call_dict.values())}-" + f"{max(call.shim_data['total_iterations'] for call in call_dict.values())} " + ) + if self.max_iterations is None: + tqdm.write(" Total progress: " + iteration_range) + return + + total = num_params * self.max_iterations + progress_value = ( + sum(call.shim_data["total_iterations"] for call in call_dict.values()) + + (num_params - len(call_dict)) * self.max_iterations + ) + + progress_string = ( + f"{progress_value / total * 100:.1f}% " + f"Iteration range " + f"{min(call.shim_data['total_iterations'] for call in call_dict.values())}-" + f"{max(call.shim_data['total_iterations'] for call in call_dict.values())} " + f"of {self.max_iterations} " + f"({num_params - len(call_dict)} of {num_params} parameters finished)" + ) + + bar_width = min(100, max(num_params, 20)) + bar_format = f"{{desc}}: |{{bar:{bar_width}}}| {progress_string}" + + total_bar = self._make_progress_bar( + total=total, + desc=" Total progress", + bar_format=bar_format, + colour=DW_ORANGE, + initial=progress_value, + ) + total_bar.close() + + def _save_results( + self, + data_dict: dict[str, Any], + run_index: int | None = None, + filename: str | None = None, + ) -> None: + """Save results to disk using LZMA-compressed pickle.""" + if filename is None: + if run_index is None: + run_index = self.run_index + filename = f"iter{run_index:05d}.pkl.lzma" + else: + if run_index is not None: + raise ValueError("Cannot specify both filename and run_index.") + + # Write to a temp directory first to reduce disk write errors from killed jobs. + with tempfile.TemporaryDirectory(dir=self.data_path) as tmp: + temp_filename = Path(tmp) / filename + with lzma.open(temp_filename, "wb") as f: + pickle.dump(data_dict, f) + os.rename(temp_filename, self.data_path / filename) + + def _get_sorted_results_file_list(self) -> list[str]: + """Return result filenames sorted lexicographically.""" + return [str(fn) for fn in sorted(self.data_path.glob("iter*.pkl.lzma"))] + + def _get_next_run_index(self) -> tuple[int, bool]: + """Get the next run index based on the existing files in the data path.""" + if not self.data_path.exists(): + return 0, False + + fnlist = list(self.data_path.glob("iter*.pkl.lzma")) + if not fnlist: + return 0, False + + latest_file_iter = max(int(fn.stem.split(".")[0][4:]) for fn in fnlist) + return latest_file_iter + 1, True + + def _prepare_run_index(self) -> bool: + """Prepare the run index for the next iteration, creating the data path if needed.""" + if self.data_path is None: + raise RuntimeError("No parameterization selected. Call apply_param() first.") + + self.data_path.mkdir(parents=True, exist_ok=True) + self.run_index, already_initialized = self._get_next_run_index() + return already_initialized + + def _get_solver_pathstring(self) -> str: + """Construct a pathstring for the solver. + + Structured to support additional sampler types in the future. + """ + pathstring = None + rules = [ + (lambda s: s == "DWaveSampler", "qpu"), + ] + for check, label in rules: + if check(type(self.sampler).__name__): + pathstring = label + if pathstring is None: + raise TypeError("Sampler type not compatible with known possibilities") + + if pathstring in ["qpu"]: + pathstring += f"/{self.sampler.solver.name}" + + return pathstring + + def _get_parameter_pathstring(self) -> str: + """Construct a pathstring for the experimental parameters. + + Assumes a forward anneal. Annealing time format is in microseconds (up + to 999.9999us), with six decimal places (picosecond resolution). + """ + signed_energy_scale = self.param["signed_energy_scale"] + + if "anneal_time" in self.param: + pathstring = ( + f'energyscale{signed_energy_scale:0.3}/atime{self.param["anneal_time"]:010.6f}us' + ) + elif "anneal_schedule" in self.param: + pathstring = ( + f'energyscale{signed_energy_scale:0.3}/asched{self.param["anneal_schedule"]}' + ) + else: + raise ValueError( + "Parameter list must contain either 'anneal_time' or 'anneal_schedule'." + ) + + # Strip spaces and replace other unswanted symbols with underscores. + pathstring = pathstring.replace(" ", "_") + for bad_symbol in ":;,": + pathstring = pathstring.replace(bad_symbol, "") + + return pathstring + + def _get_relative_data_path(self) -> str: + """Make a subdirectory name for a sampler call's data.""" + return "/".join( + [ + self.inst._get_instance_pathstring(), + self._get_solver_pathstring(), + self._get_parameter_pathstring(), + ] + ) + + def _make_logical_bqms(self) -> list[dimod.BQM]: + """Make logical BQMs for the experiment.""" + logical_bqm = self.inst.make_bqm() + + if not hasattr(self.inst, "embedding_list"): + return [logical_bqm] + + return [logical_bqm] * len(self.inst.embedding_list) + + def _build_sampler_call(self) -> None | SamplerCall: + """Build the sampler call using attributes of the experiment and instance. + + Returns a SamplerCall. + """ + sampler_call = SamplerCall(run_index=self.run_index) + sampler_call.logical_bqms = self._make_logical_bqms() + sampler_call.shim_data = self._get_shim_data() + + # Here we can find out that we're finished. + if ( + self.max_iterations is not None + and sampler_call.shim_data["total_iterations"] >= self.max_iterations + ): + return None + + sampler_call.bqm = self._make_bqm(sampler_call) + sampler_call.sampler_params = self._make_sampler_params(shim_data=sampler_call.shim_data) + + return sampler_call + + def _format_parameter_list( + self, + parameter_list: list[dict[str, float]], + ) -> list[dict[str, float]]: + """Deduplicate and format the parameter list for filename consistency. + + Some parameters can cause bugs if they are not appropriately formatted, + rounded, etc. in accordance with filenames. + """ + ret_unique = [] + ret = copy.deepcopy(parameter_list) + for entry in ret: + if "anneal_time" in entry: + entry["anneal_time"] = np.round(entry["anneal_time"], 6) + if "anneal_schedule" in entry: + entry["anneal_schedule"] = [tuple(np.round(p, 6)) for p in entry["anneal_schedule"]] + + if entry not in ret_unique: + ret_unique.append(entry) + + return ret_unique + + def _generate_data_to_save( + self, + sampler_call: SamplerCall, + results: dict[str, Any], + ) -> dict[str, Any]: + """Construct a single dictionary containing results and shim data for saving.""" + savedata = {} + for key in results: + if isinstance(results[key], np.ndarray): + if results[key].dtype == "complex128": + savedata[key] = results[key].astype(np.complex64) + elif results[key].dtype == "float64": + savedata[key] = results[key].astype(np.float32) + else: + savedata[key] = results[key] + else: + savedata[key] = results[key].copy() + + savedata["shim_data"] = {} + for key in sampler_call.shim_data: + if isinstance(sampler_call.shim_data[key], np.ndarray): + savedata["shim_data"][key] = sampler_call.shim_data[key].astype(np.float32) + elif isinstance(sampler_call.shim_data[key], int): + savedata["shim_data"][key] = sampler_call.shim_data[key] + else: + savedata["shim_data"][key] = sampler_call.shim_data[key].copy() + + return savedata + + def _make_sampler_params(self, **kwargs) -> dict[str, Any]: + """Construct a dictionary containing sampler parameters.""" + ret = { + "answer_mode": "raw", + "auto_scale": False, + "flux_drift_compensation": False, + "readout_thermalization": int(self.param["readout_thermalization"]), + "num_reads": self.param["num_reads"], + "label": os.path.join(self._get_relative_data_path(), f"iter{self.run_index:05d}"), + } + + if "shim_data" in kwargs: + if "flux_biases" in kwargs["shim_data"]: + ret["flux_biases"] = list(kwargs["shim_data"]["flux_biases"]) + if "anneal_offsets" in kwargs["shim_data"]: + ret["anneal_offsets"] = list(kwargs["shim_data"]["anneal_offsets"]) + + if self.param.get("fast_anneal", False): + ret["fast_anneal"] = True + + if "anneal_schedule" in self.param: + ret["anneal_schedule"] = self.param["anneal_schedule"] + elif "anneal_time" in self.param: + ret["annealing_time"] = self.param["anneal_time"] + + return ret + + def _get_shim_data(self) -> dict[str, Any]: + """Load shim data if possible, otherwise make an initial shim.""" + if self.already_initialized: + return self._load_shim() + return self._make_initial_shim() + + def _make_initial_shim(self) -> dict[str, Any]: + """Create the initial shim and dictate what shim will be saved and modified.""" + shim_data = {"total_iterations": 0} + if hasattr(self.inst, "embedding_list"): + num_embeddings = len(self.inst.embedding_list) + shim_data["flux_biases"] = np.zeros(self.sampler.properties["num_qubits"]) + shim_data["anneal_offsets"] = np.zeros(self.sampler.properties["num_qubits"]) + shim_data["relative_coupler_strength"] = np.ones((num_embeddings, self.inst.num_edges)) + + if self.param.get("flux_biases", None) is not None: + shim_data["flux_biases"] = self.param.get("flux_biases") + + return shim_data + + def _get_latest_iteration_filename(self) -> Path: + """Return the filename of the most recently completed iteration.""" + return self.data_path / f"iter{self.run_index - 1:05d}.pkl.lzma" + + def _load_shim(self) -> dict[str, Any]: + """Load shim data from the most recently completed iteration.""" + filename = self._get_latest_iteration_filename() + + if os.path.getsize(filename) == 0: + os.remove(filename) + raise FileNotFoundError(f"{filename} does not exist") + + try: + with lzma.open(filename, "rb") as f: + data = pickle.load(f) + shim_data = data["shim_data"] + return shim_data + except FileNotFoundError as e: + raise FileNotFoundError(f"{filename} does not exist") from e + except Exception as e: + raise OSError("Failed to open file") from e + + def _update_shim(self, sampler_call: SamplerCall, results: dict[str, Any]) -> None: + """Update shim parameters according to shim data and parameters.""" + if "flux_biases" in sampler_call.shim_data and self.param.get("flux_bias_shim_step", 0) != 0: + self._update_flux_bias_shim(sampler_call, results) + if ( + "relative_coupler_strength" in sampler_call.shim_data + and self.param.get("coupler_shim_step", 0) != 0 + ): + self._update_coupler_shim(sampler_call, results) + + sampler_call.shim_data["total_iterations"] += 1 + + def _update_flux_bias_shim(self, sampler_call: SamplerCall, results: dict[str, Any]) -> None: + """Update flux-bias shim values based on qubit magnetization.""" + target_magnetization = self.param["target_magnetization"] + qubit_magnetization = results["QubitMagnetization"] + flux_biases = sampler_call.shim_data["flux_biases"] + shim_step = self.param["flux_bias_shim_step"] + + steps = shim_step * (qubit_magnetization.ravel() - target_magnetization) + flux_biases[self.inst.embedding_list.ravel()] -= steps + mean_magnetization = np.mean(qubit_magnetization) + + if target_magnetization > 0: + if mean_magnetization < target_magnetization - 0.001: + flux_biases *= 1.01 + elif mean_magnetization > target_magnetization + 0.001: + flux_biases /= 1.01 + + elif target_magnetization < 0: + if mean_magnetization > target_magnetization + 0.001: + flux_biases *= 1.01 + elif mean_magnetization < target_magnetization - 0.001: + flux_biases /= 1.01 + + def _update_coupler_shim( + self, + sampler_call: SamplerCall, + results: dict[str, Any], + step_size: float | None = None, + ) -> None: + """Update relative coupler strength based on measured frustration.""" + orbits = self.inst.coupler_orbits + signed_energy_scale = self.param["signed_energy_scale"] + relative_coupler_strength = sampler_call.shim_data["relative_coupler_strength"] + + # Allow for zero step size, which will just truncate the shim. + if step_size is None: + step_size = self.param["coupler_shim_step"] + if step_size == 0: + return + + # Get the set over which we normalize. + normalization_basis = np.ones_like(orbits, dtype=bool) + + # Assume we have multiple embeddings of the same BQM. + bqms = sampler_call.logical_bqms + if len(bqms) > 1 and any(bqm != bqms[0] for bqm in bqms[1:]): + raise NotImplementedError("Case for distinct embedded BQMs not implemented yet.") + + bqm = bqms[0] + logical_values = np.array([bqm.quadratic[edge] for edge in self.inst.edge_list]) + coupler_signs = np.sign(logical_values) + for orbit_bin in range(max(orbits) + 1): + bin_edges = np.argwhere(orbits == orbit_bin).ravel() + if step_size != 0: + frust = results["CouplerFrustration"][:, bin_edges] + meanfrust = np.mean(frust) + relative_coupler_strength[:, bin_edges] += step_size * (frust - meanfrust) + + # Damp the couplers (push toward default value) + if "coupler_damp" in self.param and self.param["coupler_damp"] > 0: + excess = relative_coupler_strength[:, bin_edges] - np.mean( + relative_coupler_strength[:, bin_edges] + ) + relative_coupler_strength[:, bin_edges] -= ( + np.multiply(coupler_signs[bin_edges], excess) * self.param["coupler_damp"] + ) + + # New truncation method... previous is buggy when we mix signs of logical values. + # Let's try being more explicit. + for iemb in range(len(relative_coupler_strength)): + violators = ( + relative_coupler_strength[iemb, bin_edges] + * logical_values[bin_edges] + * signed_energy_scale + > 1 + ) + relative_coupler_strength[iemb, bin_edges[violators]] = ( + 0.99999 / logical_values[bin_edges[violators]] / signed_energy_scale + ) + + violators = ( + relative_coupler_strength[iemb, bin_edges] + * logical_values[bin_edges] + * signed_energy_scale + < -2 + ) + relative_coupler_strength[iemb, bin_edges[violators]] = ( + -1.99999 / logical_values[bin_edges[violators]] / signed_energy_scale + ) + + # Renormalize each orbit after truncation + for orbit_bin in range(np.max(orbits) + 1): + bin_edges = orbits == orbit_bin + mean_relative = np.mean( + np.abs(relative_coupler_strength[:, bin_edges * normalization_basis]) + ) + relative_coupler_strength[:, bin_edges] /= mean_relative + + # And truncate again + for orbit_bin in range(np.max(orbits) + 1): + bin_edges = np.argwhere(orbits == orbit_bin).ravel() + + # New truncation method... previous is buggy when we mix signs of logical values. + # Let's try being more explicit. + for iemb in range(len(relative_coupler_strength)): + violators = ( + relative_coupler_strength[iemb, bin_edges] + * logical_values[bin_edges] + * signed_energy_scale + > 1 + ) + relative_coupler_strength[iemb, bin_edges[violators]] = ( + 0.99999 / logical_values[bin_edges[violators]] / signed_energy_scale + ) + + violators = ( + relative_coupler_strength[iemb, bin_edges] + * logical_values[bin_edges] + * signed_energy_scale + < -2 + ) + relative_coupler_strength[iemb, bin_edges[violators]] = ( + -1.99999 / logical_values[bin_edges[violators]] / signed_energy_scale + ) + + Q = logical_values * relative_coupler_strength * signed_energy_scale + Q_max = np.max(Q) + Q_min = np.min(Q) + if Q_max > 1 or Q_min < -2: + raise ValueError( + "Effective coupler strengths violate hardware bounds: " + f"min={Q_min:.6f}, max={Q_max:.6f}" + ) + + def _make_bqm(self, sampler_call: SamplerCall) -> dimod.BQM: + """Construct a BQM for the current sampler call.""" + signed_energy_scale = self.param["signed_energy_scale"] + bqm = dimod.BQM(vartype="SPIN") + if not hasattr(self.inst, "embedding_list"): + logical_bqm = sampler_call.logical_bqms[0] + + for v in range(self.inst.num_spins): + # Make sure variables appear in the correct order when dealing with software solvers + bqm.add_variable(v) + if v in logical_bqm.variables: + bqm.add_linear(v, logical_bqm.linear[v]) + + for u, v in self.inst.edge_list: + bqm.add_quadratic(u, v, logical_bqm.quadratic[u, v] * signed_energy_scale) + + return bqm + + relative_coupler_strength = sampler_call.shim_data["relative_coupler_strength"] + for iemb, emb in enumerate(self.inst.embedding_list): + logical_bqm = sampler_call.logical_bqms[iemb].copy() + + for v in range(self.inst.num_spins): + # Don't touch degree-zero spins. Relevant to partial yield. + if logical_bqm.degree(v) > 0: + bqm.add_linear(emb[v], logical_bqm.linear[v]) + + for iedge, edge in enumerate(self.inst.edge_list): + bias = ( + logical_bqm.quadratic[tuple(edge)] + * relative_coupler_strength[iemb, iedge] + * signed_energy_scale + ) + bqm.add_quadratic(emb[edge[0]], emb[edge[1]], bias) + + return bqm diff --git a/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py new file mode 100644 index 0000000..a34860d --- /dev/null +++ b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py @@ -0,0 +1,31 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass + +from dwave.experimental.lattice_utils.experiment import ExperimentConfig + +__all__ = ['FastAnnealExperimentConfig'] + + +@dataclass +class FastAnnealExperimentConfig(ExperimentConfig): + """Configuration class for Fast Anneal Experiments.""" + + fast_anneal: bool = True + automorph_embeddings: bool = False + coupler_damp: float = 0.0 + anneal_offset_damp: float = 0.0 + individual_qubit_anneal_offsets: list[float] | None = None + logical_software: bool = False diff --git a/dwave/experimental/lattice_utils/experiment/samplercall.py b/dwave/experimental/lattice_utils/experiment/samplercall.py new file mode 100644 index 0000000..b8cffb9 --- /dev/null +++ b/dwave/experimental/lattice_utils/experiment/samplercall.py @@ -0,0 +1,30 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field + +import dimod + +__all__ = ['SamplerCall'] + + +@dataclass +class SamplerCall: + """Data class for managing asynchronous sampler calls.""" + + run_index: int + bqm: dimod.BQM | None = None + shim_data: dict = field(default_factory=dict) + logical_bqms: list = field(default_factory=list) + sampler_params: dict = field(default_factory=dict) diff --git a/dwave/experimental/lattice_utils/lattice/__init__.py b/dwave/experimental/lattice_utils/lattice/__init__.py new file mode 100644 index 0000000..1ff81e5 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/__init__.py @@ -0,0 +1,20 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dwave.experimental.lattice_utils.lattice.chain import * +from dwave.experimental.lattice_utils.lattice.embedded_lattice import * +from dwave.experimental.lattice_utils.lattice.lattice import * +from dwave.experimental.lattice_utils.lattice.optimize import * +from dwave.experimental.lattice_utils.lattice.orbits import * +from dwave.experimental.lattice_utils.lattice.triangular import * diff --git a/dwave/experimental/lattice_utils/lattice/chain.py b/dwave/experimental/lattice_utils/lattice/chain.py new file mode 100644 index 0000000..6636a78 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/chain.py @@ -0,0 +1,86 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Generator +from pathlib import Path +from typing import Any + +import dimod +from numpy.typing import NDArray + +from dwave.experimental.lattice_utils.lattice.lattice import Lattice + +__all__ = ['Chain'] + + +class Chain(Lattice): + """One-dimensional chain lattice. + + This class represents a 1D chain of spins, where each spin is connected to + its nearest neighbors. The chain can be periodic (forming a ring) or + non-periodic (open chain) based on the ``periodic`` parameter. + + Args: + dimensions: One-element tuple giving the number of spins in the chain. + periodic: One-element tuple indicating whether the chain is periodic. + data_root: A string or Path to the root directory for storing lattice data. + orbit_type: A string specifying the type of orbits to compute for the + lattice. + qubit_orbits: Explicit qubit orbit labels, used only when ``orbit_type == "explicit"``. + Must have length equal to the number of spins in the lattice. + coupler_orbits: Explicit coupler orbit labels, used only when ``orbit_type == "explicit"``. + Must have length equal to the number of edges in the lattice. + """ + + def __init__( + self, + *, + dimensions: tuple[int], + data_root: str | Path, + periodic: tuple[bool] = (True,), + orbit_type: str = "singleton", + qubit_orbits: NDArray | None = None, + coupler_orbits: NDArray | None = None, + reference_energy_sampler: dimod.Sampler | None = None, + reference_energy_sampler_kwargs: dict[str, Any] | None = None, + ): + self.geometry_name = "Chain" + self.num_spins = dimensions[0] + if len(dimensions) != 1: + raise ValueError(f"Chain requires dimensions of length 1, got {len(dimensions)}.") + + super().__init__( + dimensions=dimensions, + periodic=periodic, + data_root=data_root, + orbit_type=orbit_type, + qubit_orbits=qubit_orbits, + coupler_orbits=coupler_orbits, + reference_energy_sampler=reference_energy_sampler, + reference_energy_sampler_kwargs=reference_energy_sampler_kwargs, + ) + + def generate_edges(self) -> Generator[tuple[int, int]]: + """Yield edges for a 1D chain lattice. + + Returns: + A generator of tuples, where each tuple represents an edge between + two spins in the chain. + """ + n = self.dimensions[0] + for i in range(n - 1): + yield (i, i + 1) + + if self.periodic[0] and n > 2: + yield (n - 1, 0) diff --git a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py new file mode 100644 index 0000000..9ad8a71 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -0,0 +1,288 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Generator, Hashable +from itertools import combinations, product +from numbers import Integral +from pathlib import Path + +import dimod +import numpy as np +from numpy.typing import NDArray + +from dwave.experimental.lattice_utils.lattice.lattice import Lattice + +__all__ = ['EmbeddedLattice'] + + +class EmbeddedLattice(Lattice): + """Embed a logical lattice onto a physical lattice using chains. + + Logical nodes are represented by chains of physical spins. Subclasses can + specialize ``get_chain_connectivity`` to describe how spins within a chain, + and between neighboring logical chains, should be connected. + + For example, a dimer-style embedding might map different logical couplings to + different physical index pairs. In a 3D dimer class, x-, y-, and z-couplings + could return ``((1, 1),)``, ``((0, 0),)``, and ``((0, 1), (1, 0))``, + respectively. A chain coupling, where the logical edge is ``(u, u)``, could + return ``((0, 1),)``. + + Args: + logical_lattice: The logical lattice instance to embed. + chain_nodes: Mapping from logical nodes to their physical chains. + """ + + def __init__( + self, + *, + logical_lattice: Lattice, + chain_nodes: dict[int, tuple[int, Integral]], + dimensions: tuple[int, ...], + data_root: str | Path | None = None, + periodic: tuple[bool, ...] | None = None, + orbit_type: str = "singleton", + qubit_orbits: NDArray | None = None, + coupler_orbits: NDArray | None = None, + chain_strength: float = 2, + ): + if not isinstance(logical_lattice, Lattice): + raise TypeError("logical_lattice must be a Lattice instance.") + + self.logical_lattice = logical_lattice + if hasattr(self.logical_lattice, "logical_lattice"): + raise NotImplementedError("Nested embedded lattices not supported.") + + if data_root is None: + data_root = logical_lattice.data_root + + self.chain_nodes = chain_nodes + self.chain_coupling = -chain_strength + + if not hasattr(self, "num_spins"): + self.num_spins = sum(len(c) for c in chain_nodes.values()) + + if periodic is None: + periodic = self.logical_lattice.periodic + + super().__init__( + dimensions=dimensions, + data_root=data_root, + periodic=periodic, + orbit_type=orbit_type, + qubit_orbits=qubit_orbits, + coupler_orbits=coupler_orbits, + ) + + def get_chain_connectivity( + self, + u: Hashable, + v: Hashable | None = None, + ) -> tuple[tuple[int, int], ...]: + """Get the connectivity for a given edge in the logical lattice. + + Args: + u: The first node in the logical edge. + v: The second node in the logical edge. If None, this is treated as + a chain edge (u == v). + + Returns: + A tuple of tuples, where each inner tuple represents a pair of indices + in the chains corresponding to u and v that should be connected. For + a chain edge (u == v or v is None), this will return pairs of indices + within the same chain. For a logical edge (u != v), this will return + pairs of indices between the two chains. + """ + # Interior chain connectivity. Generic version: add all possible edges. + if u == v or v is None: + return tuple(combinations(range(len(self.chain_nodes[u])), 2)) + + # Connectivity between two edges. Generic version: add all possible edges. + return tuple(product(range(len(self.chain_nodes[u])), range(len(self.chain_nodes[v])))) + + def generate_edges(self) -> Generator[tuple[Hashable, Hashable]]: + """Yield physical edges for the embedded lattice. + + Returns: + A generator of tuples, where each tuple represents an edge between + two spins in the physical lattice. + """ + logical_bqm = self.logical_lattice.make_bqm() + + # Now embed it. First make embedded spins and connect the chains. + for v in logical_bqm.variables: + for edge in self.get_chain_connectivity(v): + yield self.chain_nodes[v][edge[0]], self.chain_nodes[v][edge[1]] + + # Next, connect the chains together + for u, v in self.logical_lattice.edge_list: + u_chain = self.chain_nodes[u] + v_chain = self.chain_nodes[v] + for edge in self.get_chain_connectivity(u, v): + yield u_chain[edge[0]], v_chain[edge[1]] + + def make_bqm(self, **kwargs) -> dimod.BQM: + """Construct the physical BQM for this embedded lattice. + + Overrides the base class ``make_bqm`` for the + + Args: + kwargs: Keyword arguments to pass to the logical lattice's + `make_bqm` method. + + Returns: + A dimod.BQM representing the embedded logical BQM. + """ + if hasattr(self, "fixed_seed"): + self.logical_lattice.fixed_seed = self.fixed_seed + kwargs.pop("seed", None) + + return self.embed_bqm(self.logical_lattice.make_bqm(**kwargs)) + + def embed_bqm(self, logical_bqm: dimod.BQM) -> dimod.BQM: + """Embed a logical BQM onto the physical lattice. + + This is a lattice-aware alternative to ``dwave.embedding.embed_bqm``. + The standard implementation treats each chain as an unodered set of + physical qubits and routes interactions across whatever target edges + happen to be available. Here, chains are ordered tuples and the + physical edges used for each logical interaction are chosen + deterministically by ``get_chain_connectivity``, so that the position of + a qubit within its chain carries geometric meaining (e.g. in dimerized + lattices, index 0 vs. 1 corresponds to a specific sublattice). This + allows for more structured embeddings that can be tailored to the + geometry of the logical lattice and the physics of the problem. + + Chain couplings are fixed at ``self.chain_coupling`` rather than + computed by a chain-strength heuristic. + + Args: + logical_bqm: A dimod.BQM representing the BQM defined on the logical + variable space of the embedded lattice. + + Returns: + A dimod.BQM representing the embedded BQM defined on the physical + variable space of the embedded lattice. + """ + # First make embedded spins and connect the chains. + embedded_bqm = dimod.BQM(vartype="SPIN") + embedded_variables = np.concatenate(list(self.chain_nodes.values())) + embedded_variables.sort() + + for v in embedded_variables: + embedded_bqm.add_variable(v) + for v in logical_bqm.variables: + if logical_bqm.degree(v) > 0: # If the degree is zero we won't add any chain couplings. + for embedded_v in self.chain_nodes[v]: + embedded_bqm.add_linear( + embedded_v, + logical_bqm.linear[v] / len(self.chain_nodes[v]), + ) + for edge in self.get_chain_connectivity(v): + embedded_bqm.add_quadratic( + self.chain_nodes[v][edge[0]], + self.chain_nodes[v][edge[1]], + self.chain_coupling, + ) + + # Next, connect the chains together + for u, v in self.logical_lattice.edge_list: + u_chain = self.chain_nodes[u] + v_chain = self.chain_nodes[v] + bias_uv = logical_bqm.quadratic[u, v] + edges = self.get_chain_connectivity(u, v) + for x, y in edges: + embedded_bqm.add_quadratic(u_chain[x], v_chain[y], bias_uv / len(edges)) + + return embedded_bqm + + def unembed_bqm(self, embedded_bqm: dimod.BQM) -> dimod.BQM: + """Unembed an embedded BQM back onto the logical variable space. + + Args: + embedded_bqm: A dimod.BQM representing the BQM defined on the physical + variable space of the embedded lattice. + + Returns: + A dimod.BQM representing the unembedded logical BQM. + """ + logical_bqm = dimod.BQM(vartype="SPIN") + for v in range(self.logical_lattice.num_spins): + logical_bqm.add_variable(v) + + which_spin = np.zeros(self.num_spins).astype(int) + for spin, chain in self.chain_nodes.items(): + which_spin[np.array(chain)] = spin + + for v in embedded_bqm.variables: + logical_bqm.add_linear(which_spin[v], embedded_bqm.linear[v]) + + for u, v in embedded_bqm.quadratic: + if which_spin[u] != which_spin[v]: + bias_uv = embedded_bqm.quadratic[u, v] + logical_bqm.add_quadratic(which_spin[u], which_spin[v], bias_uv) + + return logical_bqm + + def unembed_sampleset(self, sampleset: dimod.SampleSet) -> dimod.SampleSet: + """Unembed a SampleSet using majority vote with random tie-breaking. + + Args: + sampleset: A dimod.SampleSet representing samples in the physical + variable space. + + Returns: + A dimod.SampleSet representing the unembedded logical samples. + """ + sample_array = dimod.as_samples(sampleset)[0].T + + voted_samples = np.asarray( + [ + np.sum(sample_array[self.chain_nodes[v], :], axis=0) + for v in range(len(self.chain_nodes)) + ] + ) + voted_samples = np.sign(voted_samples + np.random.rand(*voted_samples.shape) - 0.5).T + + return dimod.SampleSet.from_samples(voted_samples, vartype=dimod.SPIN, energy=0) + + def embed_sample(self, sample: NDArray) -> NDArray: + """Embed a logical sample onto the physical lattice. + + Args: + sample: A NumPy array representing a sample in the logical variable space. + + Returns + A NumPy array representing the embedded physical sample. + """ + ret = np.zeros(self.num_spins) + for spin, chain in self.chain_nodes.items(): + ret[np.array(chain)] = sample[spin] + + return ret + + def unembed_sample(self, sample: NDArray) -> NDArray: + """Unembed a physical sample using majority vote with random tie-breaking. + + Args: + sample: A NumPy array representing a sample in the physical variable space. + + Returns: + A NumPy array representing the unembedded logical sample. + """ + ret = np.zeros(self.logical_lattice.num_spins) + for spin, chain in self.chain_nodes.items(): + ret[spin] = np.sign(np.sum(sample[np.array(chain)]) + np.random.rand() - 0.5) + + return ret diff --git a/dwave/experimental/lattice_utils/lattice/lattice.py b/dwave/experimental/lattice_utils/lattice/lattice.py new file mode 100644 index 0000000..e6784eb --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -0,0 +1,330 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import warnings +from abc import ABC, abstractmethod +from collections.abc import Generator, Hashable +from pathlib import Path +from typing import Any + +import dimod +import networkx as nx +import numpy as np +from minorminer.utils.parallel_embeddings import find_multiple_embeddings +from numpy.typing import NDArray + +from dwave.experimental.lattice_utils.lattice.optimize import optimize +from dwave.experimental.lattice_utils.lattice.orbits import get_orbits + +__all__ = ['Lattice'] + + +class Lattice(ABC): + """An abstract base class for representing lattice geometries used in lattice-utils experiments. + + Subclasses are responsible for defining the lattice geometry itself. In particular, + a subclass must: + + - Implement the ``generate_edges`` method, which yields the edges of the lattice as pairs + - Initialize the ``self.num_spins`` attribute in the constructor, which is used by the base class + - set any geometry-specific identifiers such as ``self.geometry_name`` + + Args: + dimensions: Tuple specifying the size of the lattice in each dimension. + data_root: Root directory for loading and saving lattice data such as embeddings and orbits. + periodic: Tuple indicating whether each dimension is periodic (True) or open (False). + orbit_type: Method for determining qubit and coupler orbits. Must be one of "global", + "standard", "singleton", or "explicit". See ``initialize_orbits`` for details. + qubit_orbits: Explicit qubit orbit labels, used only when ``orbit_type == "explicit"``. + Must have length equal to the number of spins in the lattice. + coupler_orbits: Explicit coupler orbit labels, used only when ``orbit_type == "explicit"``. + Must have length equal to the number of edges in the lattice. + """ + + def __init__( + self, + *, + dimensions: tuple[int, ...], + data_root: str | Path, + periodic: tuple[bool, ...] | None = None, + orbit_type: str = "singleton", + qubit_orbits: NDArray | None = None, + coupler_orbits: NDArray | None = None, + reference_energy_sampler: dimod.Sampler | None = None, + reference_energy_sampler_kwargs: dict[str, Any] | None = None, + ): + self.dimensions = dimensions + self.data_root = Path(data_root) + self.reference_energy_sampler = reference_energy_sampler + self.reference_energy_sampler_kwargs = reference_energy_sampler_kwargs + + self.periodic = periodic if periodic is not None else tuple(False for _ in dimensions) + if len(self.periodic) != len(self.dimensions): + raise ValueError( + f"periodic and dimensions must have the same length: " + f"got {len(self.periodic)} and {len(self.dimensions)}." + ) + if not hasattr(self, "num_spins"): + raise AttributeError(f"{type(self).__name__} subclass must initialize self.num_spins") + + self.edge_list: list[tuple[Hashable, Hashable]] = list(self.generate_edges()) + self.num_edges: int = len(self.edge_list) + self.orbit_type: str = orbit_type + self.initialize_orbits(qubit_orbits, coupler_orbits) + + @abstractmethod + def generate_edges(self) -> Generator[tuple[Hashable, Hashable]]: + """Yield the edges for this lattice.""" + + def embed_lattice( + self, + sampler: dimod.Sampler, + try_to_load: bool = True, + timeout: int = 10, + max_number_of_embeddings: int | None = None, + min_number_of_embeddings: int = 1, + exclude_qubits: list | None = None, + **kwargs, + ) -> None: + """Find or load embeddings onto the sampler graph. + + Args: + sampler: Sampler whose hardware graph is used as the target for embedding. + try_to_load: If True, attempt to load embeddings from disk before + trying to find them. + timeout: Time limit for the embedding search, in seconds. + max_number_of_embeddings: Maximum number of embeddings to search for. + min_number_of_embeddings: Minimum number of embeddings required to save. + exclude_qubits: Qubits to remove from the sampler graph before searching + for embeddings. + """ + if exclude_qubits is None: + exclude_qubits = [] + + graph_bqm = dimod.to_networkx_graph(self.make_bqm()) + graph_sampler = sampler.to_networkx_graph() + graph_sampler.remove_nodes_from(exclude_qubits) + + if try_to_load: + try: + self._load_embeddings(sampler) + return + except FileNotFoundError: + warnings.warn("No cached embedding file found, computing new embeddings.") + + embedding_dicts = find_multiple_embeddings( + graph_bqm, + graph_sampler, + max_num_emb=max_number_of_embeddings, + embedder_kwargs={'timeout': timeout}, + ) + if not embedding_dicts: + raise ValueError( + f"No embeddings found for {type(self).__name__}" + f"(dimensions={self.dimensions}) on " + f"{type(sampler).__name__}" + f" (target graph: {graph_sampler.number_of_nodes()} nodes, " + f"{graph_sampler.number_of_edges()} edges; timeout={timeout}s). " + f"Try increasing timeout, reducing dimensions, or relaxing exclude_qubits." + ) + + embeddings = np.stack([list(emb.values()) for emb in embedding_dicts]) + if len(embeddings) >= min_number_of_embeddings and np.prod(embeddings.shape): + self._save_embeddings(sampler, embeddings) + + def make_bqm(self) -> dimod.BQM: + """Construct a default with BQM coupling strength values set to +1. + + Returns: + A binary quadratic model representing the lattice with uniform + coupling strength. + """ + bqm = dimod.BQM(vartype="SPIN") + for v in range(self.num_spins): + bqm.add_variable(v) + for u, v in self.edge_list: + bqm.add_quadratic(u, v, 1.0) + + return bqm + + def initialize_orbits( + self, + qubit_orbits: NDArray | None = None, + coupler_orbits: NDArray | None = None, + ) -> None: + """Initialize qubit and coupler orbits. + + Orbit assignments are determined according to ``self.orbit_type``: + + -``global``: Put all the couplers in one orbit and all the qubits in one + orbit. Exception: for embedded lattices, put all logical couplers in one + orbit and all chain couplers in another. + -``standard``: Load previously computed automorphism-based orbits, or + compute them and save them if unavailable. + -``singleton``: Put each qubit and coupler in its own orbit. + -``explicit``: Use the orbit assignments provided via ``qubit_orbits`` + and ``coupler_orbits``. + + Args: + qubit_orbits: Explicit qubit orbit labels, used only when + ``self.orbit_type == "explicit"``. Must have length ``self.num_spins``. + coupler_orbits: Explicit coupler orbit labels. Used only when + ``self.orbit_type == "explicit"``. Must have length ``self.num_edges``. + """ + if self.orbit_type == "global": + self.qubit_orbits = np.zeros(self.num_spins, dtype=int) + + if hasattr(self, "logical_lattice"): + which_chain = {v: key for key, val in self.chain_nodes.items() for v in val} + self.coupler_orbits = np.zeros(self.num_edges, dtype=int) + + for i, (u, v) in enumerate(self.edge_list): + if which_chain[u] == which_chain[v]: + self.coupler_orbits[i] = 1 + else: + self.coupler_orbits = np.zeros(self.num_edges, dtype=int) + + elif self.orbit_type == "standard": + try: + self._load_orbits() + except FileNotFoundError: + # calculating orbits + bqm = self.make_bqm() + self.qubit_orbits, self.coupler_orbits = get_orbits(bqm, self.edge_list) + self._save_orbits() + + elif self.orbit_type == "singleton": + self.qubit_orbits = np.arange(self.num_spins) + self.coupler_orbits = np.arange(self.num_edges) + + elif self.orbit_type == "explicit": + if qubit_orbits is None or coupler_orbits is None: + raise ValueError( + 'orbit_type "explicit" requires both qubit_orbits and coupler_orbits.' + ) + if len(qubit_orbits) != self.num_spins: + raise ValueError( + f"qubit_orbits must have length {self.num_spins}, got {len(qubit_orbits)}." + ) + if len(coupler_orbits) != self.num_edges: + raise ValueError( + f"coupler_orbits must have length {self.num_edges}, " + f"got {len(coupler_orbits)}." + ) + self.qubit_orbits = qubit_orbits + self.coupler_orbits = coupler_orbits + else: + raise ValueError( + f'Unknown orbit type {self.orbit_type}. ' + 'Must be "global", "standard", "singleton", or "explicit".' + ) + + def make_networkx_graph(self) -> nx.Graph: + """Construct a NetworkX graph representation of the lattice. + + Returns: + A NetworkX graph where nodes correspond to spins and edges correspond to couplers. + """ + graph = nx.Graph() + for v in range(self.num_spins): + graph.add_node(v) + for u, v in self.edge_list: + graph.add_edge(u, v) + + return graph + + def optimize(self, bqm: dimod.BQM) -> tuple[float, NDArray, str]: + """Return the lowest energy sample by optimizing the BQM using the reference sampler. + + Returns: + A tuple containing the best energy found, the corresponding sample as a + NumPy array, and a string indicating the optimization method used. + """ + return optimize( + lattice=self, + bqm=bqm, + sampler=self.reference_energy_sampler, + sampler_kwargs=self.reference_energy_sampler_kwargs, + ) + + def _get_path( + self, + kind: str, + sampler_name: str | None = None, + extra_subdir: str | Path | None = None, + ) -> Path: + """Construct a standardized file path for embedding or orbit data.""" + if kind not in {"embedding", "orbits"}: + raise ValueError("kind must be provided as either `embedding` or `orbits`") + + class_subdir = Path(self.geometry_name) + if extra_subdir is not None: + class_subdir = class_subdir / extra_subdir + + base_dir = self.data_root / "lattice_data" / kind / class_subdir + if sampler_name is not None: + base_dir = base_dir / sampler_name + + filename = f"{self._get_size_pathstring()}.txt" + return base_dir / filename + + def _make_filename(self, kind: str, sampler: dimod.Sampler | None = None) -> Path: + """Construct a data filename for the specified sampler and data type.""" + if sampler is None: + return self._get_path(kind) + + if type(sampler).__name__ == "MockDWaveSampler": + return self._get_path(kind, sampler_name="MockDWaveSampler") + return self._get_path(kind, sampler_name=sampler.solver.name) + + def _save_embeddings(self, sampler: dimod.Sampler, embeddings: NDArray) -> None: + """Save embedding data to disk.""" + cache_filename = self._make_filename("embedding", sampler=sampler) + os.makedirs(cache_filename.parent, exist_ok=True) + np.savetxt(cache_filename, embeddings, fmt="%d") + + def _load_embeddings(self, sampler: dimod.Sampler) -> None: + """Load embedding data.""" + filename = self._make_filename("embedding", sampler=sampler) + self.embedding_list = np.atleast_2d(np.loadtxt(filename, dtype=int)) + + def _save_orbits(self) -> None: + """Save qubit and coupler orbits to disk.""" + cache_filename = self._make_filename("orbits") + cache_dir = cache_filename.parent / cache_filename.stem + os.makedirs(cache_dir, exist_ok=True) + np.savetxt(cache_dir / "qubit_orbits.txt", self.qubit_orbits, fmt="%d") + np.savetxt(cache_dir / "coupler_orbits.txt", self.coupler_orbits, fmt="%d") + + def _load_orbits(self) -> None: + """Load qubit and coupler orbits.""" + cache_filename = self._make_filename("orbits") + cache_dir = cache_filename.parent / cache_filename.stem + + self.qubit_orbits = np.loadtxt(cache_dir / "qubit_orbits.txt", dtype=int) + self.coupler_orbits = np.loadtxt(cache_dir / "coupler_orbits.txt", dtype=int) + + def _get_instance_pathstring(self) -> str: + """Construct an instance-specific pathstring. + + Generic version. Let more complex classes, including inputs that are + processor-dependent, redefine their pathstrings. This will incorporate + periodic dimensions, if available. + """ + return type(self).__name__ + "/" + self._get_size_pathstring() + + def _get_size_pathstring(self) -> str: + """Construct a size-specific pathstring including dimensions and periodicity.""" + return "size" + "x".join(f"{dim}{'p'*p}" for dim, p in zip(self.dimensions, self.periodic)) diff --git a/dwave/experimental/lattice_utils/lattice/optimize.py b/dwave/experimental/lattice_utils/lattice/optimize.py new file mode 100644 index 0000000..574ce7a --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/optimize.py @@ -0,0 +1,135 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import dimod +import numpy as np +from dwave.samplers import SimulatedAnnealingSampler +from numpy.typing import NDArray + +if TYPE_CHECKING: + from dwave.experimental.lattice_utils.experiment.experiment import Lattice + +__all__ = ['optimize', 'ExponentialBackoffSimulatedAnnealingSampler'] + + +class ExponentialBackoffSimulatedAnnealingSampler(dimod.Sampler): + """SA sampler that doubles num_sweeps until energy stops improving or a cap is hit. + + Starts at ``min_num_sweeps`` and doubles after each round that improves the + best energy, stopping when no improvement is found or ``max_num_sweeps`` is + exceeded. + + Args: + max_num_sweeps: Upper bound on the number of sweeps per SA call. Once + ``num_sweeps`` exceeds this value, the backoff loop terminates. + min_num_sweeps: Initial number of sweeps for the first SA call. + """ + + properties = None + parameters = None + + def __init__(self, max_num_sweeps=1024, min_num_sweeps=256): + self.sampler = SimulatedAnnealingSampler() + self.max_num_sweeps = max_num_sweeps + self.min_num_sweeps = min_num_sweeps + self.properties = self.sampler.properties.copy() + self.parameters = self.sampler.parameters.copy() + + def sample(self, bqm, **parameters): + num_sweeps = parameters.pop("num_sweeps", self.min_num_sweeps) + num_reads = parameters.pop("num_reads", 256) + + best_energy = np.inf + best_sampleset = None + + while num_sweeps <= self.max_num_sweeps: + ss = self.sampler.sample(bqm, num_sweeps=num_sweeps, num_reads=num_reads, **parameters) + energy = ss.first.energy + + if energy < best_energy: + best_energy = energy + best_sampleset = ss + num_sweeps *= 2 + else: + break + + best_sampleset.info["num_sweeps_exit"] = num_sweeps + return best_sampleset + + +def optimize( + lattice: Lattice, + bqm: dimod.BQM, + sampler: dimod.Sampler | None = None, + sampler_kwargs: dict[str, Any] | None = None, +) -> tuple[float, NDArray, str]: + """Return the best sample found by optimizing the BQM using simulated annealing. + + For ordinary lattices, this function applies simulated annealing directly to + the BQM. + + For embedded lattices, this function first unembeds the BQM to get the logical + BQM, optimizes the logical BQM, and then embeds the resulting sample back into + the physical lattice. The energy of the embedded sample is then optimized using + simulated annealing. + + Args: + lattice: Lattice instance defining how the optimization should be performed. + If the lattice is an EmbeddedLattice, the logical lattice will be + optimized and the resulting sample will be embedded back into the + physical lattice. + bqm: The binary quadratic model to optimize. + sampler: A dimod Sampler to use for optimization of the reference energy. + If None, a default ExponentialBackoffSimulatedAnnealingSampler will + be used. + sampler_kwargs: Optional keyword arguments to pass to the provided + sampler, such as ``num_reads`` and ``num_sweeps`` in the case of a + SA sampler. + + Returns: + A tuple containing the best energy found, the corresponding sample as a + NumPy array, and a string indicating the optimization method used. + """ + if sampler_kwargs is None: + sampler_kwargs = {} + + if sampler is None: + sampler = ExponentialBackoffSimulatedAnnealingSampler() + + reference_energy = np.inf + reference_sample = None + + # If the lattice is embedded, we should optimize the logical lattice + if hasattr(lattice, "logical_lattice"): + _, logical_sample, _ = optimize( + lattice.logical_lattice, + lattice.unembed_bqm(bqm), + sampler=sampler, + sampler_kwargs=sampler_kwargs, + ) + reference_sample = lattice.embed_sample(logical_sample) + reference_energy = bqm.energy(reference_sample) + + sampleset = sampler.sample(bqm, **sampler_kwargs) + best = sampleset.first + + if best.energy < reference_energy: + sample = np.array([best.sample[v] for v in bqm.variables]) + return best.energy, sample, type(sampler).__name__ + + return reference_energy, reference_sample, type(sampler).__name__ diff --git a/dwave/experimental/lattice_utils/lattice/orbits.py b/dwave/experimental/lattice_utils/lattice/orbits.py new file mode 100644 index 0000000..6fb9af9 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/orbits.py @@ -0,0 +1,241 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections import defaultdict +from collections.abc import Hashable + +import dimod +import networkx as nx +import numpy as np +from numpy.typing import NDArray + +from dwave.experimental.automorphism import schreier_rep + +__all__ = [ + 'reindex', + 'make_signed_bqm', + 'get_bqm_orbits', + 'get_unsigned_bqm_orbits', + 'get_orbits', +] + + +def reindex(mapping: dict[Hashable, int]) -> dict[Hashable, int]: + """Reindex dictionary values to consecutive integers starting at zero. + + Args: + mapping: Dictionary whose values represent indices or labels. + + Returns: + A new dictionary with the same keys as ``mapping`` but with values reindexed + to consecutive integers starting at zero. + """ + value_mapping = {v: i for i, v in enumerate(dict.fromkeys(mapping.values()))} + return {k: value_mapping[v] for k, v in mapping.items()} + + +def make_signed_bqm(bqm: dimod.BQM) -> dimod.BQM: + """Construct a signed expansion of a BQM. + + Takes a bqm and duplicates every spin s into two copies corresponding to + s and -s. Each field h gets mapped to two opposing fields: + h(s1) = -h(s2) + Each coupler gets mapped to four couplers: + J(s1,s2) = J(-s1,-s2) = -J(s1,-s2) = -J(-s1,s2) + + Args: + bqm: Input binary quadratic model. + + Returns: + A new BQM with duplicated variables representing both signs of each spin. + """ + ret = dimod.BinaryQuadraticModel(vartype="SPIN") + for var in bqm.variables: + ret.add_variable(f"p{var}", bqm.linear[var]) + ret.add_variable(f"m{var}", -bqm.linear[var]) + + for u, v in bqm.quadratic: + ret.add_quadratic(f"p{u}", f"p{v}", bqm.quadratic[(u, v)]) + ret.add_quadratic(f"m{u}", f"m{v}", bqm.quadratic[(u, v)]) + ret.add_quadratic(f"p{u}", f"m{v}", -bqm.quadratic[(u, v)]) + ret.add_quadratic(f"m{u}", f"p{v}", -bqm.quadratic[(u, v)]) + + return ret + + +def get_bqm_orbits( + bqm: dimod.BQM, +) -> tuple[dict[Hashable, int], dict[tuple[Hashable, Hashable], int]]: + """Take a bqm, perhaps a "signed" bqm from make_signed_bqm, and convert it + into a vertex-colored graph as needed. + + Since the automorphism module only takes edge colorings, the couplings + (J terms) need to be specified using auxiliary vertices. Thus for every + edge (u,v) of the BQM graph, we add a new vertex w(u,v) and give it the color + corresponding to J(u,v) in the BQM. + + To avoid ambiguity, we add a pendant (degree 1) vertex corresponding to each + original vertex. + + Args: + bqm: Input binary quadratic model. + + Returns: + A tuple ``(qubit_orbits, coupler_orbits)`` where ``qubit_orbits`` maps + each node to an integer orbit label and ``coupler_orbits`` maps each + edge to an integer orbit label. + """ + # The function first adds auxiliary elements to a BQM + graph = nx.Graph() + + for v in bqm.variables: + graph.add_node(f"hnode_{v}") + graph.add_node(v) + graph.add_edge(v, f"hnode_{v}") + + for u, v in bqm.quadratic: + graph.add_edge(u, v) + graph.add_node(f"Jnode_{u}_{v}") + graph.add_edge(u, f"Jnode_{u}_{v}") + graph.add_edge(v, f"Jnode_{u}_{v}") + + node_labels = list(graph.nodes) + num_nodes = graph.number_of_nodes() + node_to_idx = {node: i for i, node in enumerate(graph.nodes())} + + mapping_h = defaultdict(list) + mapping_mp = defaultdict(list) + mapping_J = defaultdict(list) + + for p, q in bqm.linear.items(): + mapping_h[q].append(f"hnode_{p}") + mapping_mp[q].append(p) + + for p, q in bqm.quadratic.items(): + mapping_J[q].append(f"Jnode_{p[0]}_{p[1]}") + + # Make color classes + coloring = [] + for nodes_h in mapping_h.values(): + coloring.append({node_to_idx[v] for v in nodes_h}) + for nodes_J in mapping_J.values(): + coloring.append({node_to_idx[e] for e in nodes_J}) + for nodes_mp in mapping_mp.values(): + coloring.append({node_to_idx[v] for v in nodes_mp}) + + graph_coloring = {} + node_colors = np.zeros(num_nodes) + for i, color in enumerate(coloring): + node_colors[list(color)] = i + for node in color: + graph_coloring[node_labels[node]] = i + + result = schreier_rep(graph, graph_coloring=graph_coloring) + + vertex_orbits = result.vertex_orbits_original_labels + vertex_orbit_array = np.zeros(num_nodes, dtype=int) + for i in range(len(vertex_orbits)): + vertex_orbit_array[[node_to_idx[x] for x in vertex_orbits[i]]] = i + qubit_orbits = { + spin: vertex_orbit_array[node_to_idx[f"hnode_{spin}"]] for spin in bqm.variables + } + + edge_orbits = result.edge_orbits_original_labels + edge_orbit_array = np.zeros(num_nodes, dtype=int) + for i in range(len(edge_orbits)): + edge_orbit_array[[(node_to_idx[x], node_to_idx[y]) for x, y in edge_orbits[i]]] = i + coupler_orbits = { + (u, v): edge_orbit_array[node_to_idx[f"Jnode_{u}_{v}"]] for u, v in bqm.quadratic + } + + return reindex(qubit_orbits), reindex(coupler_orbits) + + +def get_unsigned_bqm_orbits( + signed_qubit_orbits: dict[Hashable, int], + signed_coupler_orbits: dict[tuple[Hashable, Hashable], int], + bqm: dimod.BQM, +) -> tuple[dict[Hashable, int], dict[tuple[Hashable, Hashable], int]]: + """Convert orbits for a signed BQM into orbits for the corresponding unsigned BQM. + + Assumes that orbits are given for a signed BQM, and turns them into signed + orbits for an unsigned BQM. + + Coupler orbits are combined so that the orbit index of (p1,p2) is the same as + the orbit index of (m1,m2) and the orbit index of (p1,m2) is the same as the + orbit of index (m1,p2). This is because these pairs are related by a symmetry + of the unsigned BQM that flips both spins, and thus should be in the same orbit. + + Args: + signed_qubit_orbits: Mapping from signed variable labels to orbit indices. + signed_coupler_orbits: Mapping from signed coupler pairs to orbit indices. + bqm: Original unsigned BQM. + + Returns: + A tuple ``(qubit_orbits, coupler_orbits)`` where ``qubit_orbits`` maps + each original variable to its orbit index and ``coupler_orbits`` maps + each coupling to its orbit index. + """ + coupler_orbits = {} + for u, v in bqm.quadratic: + signed_coupler_orbits[(f"p{u}", f"p{v}")] = min( + signed_coupler_orbits[(f"p{u}", f"p{v}")], + signed_coupler_orbits[(f"m{u}", f"m{v}")], + ) + signed_coupler_orbits[(f"m{u}", f"m{v}")] = signed_coupler_orbits[(f"p{u}", f"p{v}")] + + signed_coupler_orbits[(f"p{u}", f"m{v}")] = min( + signed_coupler_orbits[(f"p{u}", f"m{v}")], + signed_coupler_orbits[(f"m{u}", f"p{v}")], + ) + signed_coupler_orbits[(f"m{u}", f"p{v}")] = signed_coupler_orbits[(f"p{u}", f"m{v}")] + + coupler_orbits[(u, v)] = signed_coupler_orbits[(f"p{u}", f"p{v}")] + + qubit_orbits = {} + for v in bqm.linear: + qubit_orbits[v] = signed_qubit_orbits[(f"p{v}")] + + return reindex(qubit_orbits), reindex(coupler_orbits) + + +def get_orbits(bqm: dimod.BQM, edge_list: list[int, int]) -> tuple[NDArray, NDArray]: + """Provide a bqm and receive a set of usable orbits derived from the signed BQM. + + Args: + bqm: Ising model to analyze. + edge_list: List of edges from the original unsigned BQM. + + Returns: + A tuple ``(qubit_orbits_array, coupler_orbits_array)`` where + ``qubit_orbits_array`` is a 1-D array of length ``num_spins`` mapping + each variable index to an orbit index, and ``coupler_orbits_array`` is a + 1-D array of length ``len(edge_list)`` mapping each entry of ``edge_list`` + to an orbit index. + """ + signed_bqm = make_signed_bqm(bqm) + signed_qubit_orbits, signed_coupler_orbits = get_bqm_orbits(signed_bqm) + qubit_orbits, coupler_orbits = get_unsigned_bqm_orbits( + signed_qubit_orbits, + signed_coupler_orbits, + bqm, + ) + + qubit_orbits_array = np.array([qubit_orbits[q] for q in range(len(qubit_orbits))]).astype(int) + coupler_orbit_dict = {tuple(sorted(list(key))): val for key, val in coupler_orbits.items()} + coupler_orbits_array = np.array( + [coupler_orbit_dict[tuple(sorted(list(c)))] for c in edge_list] + ).astype(int) + + return qubit_orbits_array, coupler_orbits_array diff --git a/dwave/experimental/lattice_utils/lattice/triangular.py b/dwave/experimental/lattice_utils/lattice/triangular.py new file mode 100644 index 0000000..1046341 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -0,0 +1,293 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Generator, Hashable +from pathlib import Path + +import dimod +import networkx as nx +import numpy as np +from numpy.typing import NDArray + +from dwave.experimental.lattice_utils.lattice.embedded_lattice import EmbeddedLattice +from dwave.experimental.lattice_utils.lattice.lattice import Lattice + +__all__ = ['Triangular', 'DimerizedTriangular'] + + +class Triangular(Lattice): + """Triangular lattice class. + + This class represents a 2D triangular lattice, where each node is connected + to its six nearest neighbors (except at boundaries, if not periodic). + + Args: + dimensions: Two-element tuple giving the number of spins in the y and x + dimensions. + periodic: Two-element tuple indicating whether the lattice is periodic + in the y and x dimensions. + data_root: A string or Path to the root directory for storing lattice + data. orbit_type: Method for determining qubit and coupler orbits. + Must be one of "global", "standard", "singleton", or "explicit". See + ``initialize_orbits`` for details. + qubit_orbits: Explicit qubit orbit labels, used only when + ``orbit_type == "explicit"``. Must have length equal to the number + of spins in the lattice. + coupler_orbits: Explicit coupler orbit labels, used only when + ``orbit_type == "explicit"``. Must have length equal to the number + of edges in the lattice. + halve_boundary_couplers: A boolean indicating whether to assign half the + coupling strength to boundary couplers. + """ + + def __init__( + self, + *, + dimensions: tuple[int, int], + periodic: tuple[bool, bool] = (True, False), + data_root: str | Path, + orbit_type: str = "singleton", + qubit_orbits: NDArray | None = None, + coupler_orbits: NDArray | None = None, + halve_boundary_couplers: bool = False, + ): + if len(dimensions) != 2: + raise ValueError(f"Triangular requires dimensions of length 2, got {len(dimensions)}.") + + self.geometry_name: str = "Triangular" + self.halve_boundary_couplers: bool = halve_boundary_couplers + self.num_spins = dimensions[0] * dimensions[1] + self.sublattice: NDArray | None = None + self.integer_coords: list[tuple[int, int]] | None = None + self.xy_coords: list[tuple[float, float]] | None = None + self.xy_size: tuple[float, float] | None = None + super().__init__( + dimensions=dimensions, + periodic=periodic, + data_root=data_root, + orbit_type=orbit_type, + qubit_orbits=qubit_orbits, + coupler_orbits=coupler_orbits, + ) + if self.periodic[0] and self.dimensions[0] % 3 != 0: + raise ValueError( + "For Triangular with periodic[0]=True, dimensions[0] must be divisible by 3." + ) + if self.periodic[1] and self.dimensions[1] % 3 != 0: + raise ValueError( + "For Triangular with periodic[1]=True, dimensions[1] must be divisible by 3." + ) + + def coordinates(self, node: int) -> tuple[int, int]: + """Return the coordinates of a node in the lattice given its index. + + Node indices are ordered by traversing the y direction first. + + Args: + node: The index of the node for which to return coordinates. + + Returns: + A tuple (y, x) representing the coordinates of the node in the lattice. + """ + length_y = self.dimensions[0] + return node % length_y, node // length_y + + def make_bqm(self) -> dimod.BQM: + """Construct the nominal triangular lattice BQM. + + If ``halve_boundary_couplers`` is True, couplers that are on the boundary + of the lattice are assigned a coupling strength of 0.5 instead of 1.0. + + Returns: + A dimod.BQM representing the nominal triangular lattice. + """ + graph = self.make_networkx_graph() + bqm = dimod.BQM(vartype="SPIN") + + for v in range(self.num_spins): + bqm.add_variable(v) + for u, v in self.edge_list: + if not self.halve_boundary_couplers or graph.degree[u] == 6 or graph.degree[v] == 6: + bqm.add_quadratic(u, v, 1.0) + else: + bqm.add_quadratic(u, v, 0.5) + + return bqm + + def generate_edges(self) -> Generator[tuple[int, int]]: + """Yield edges for the triangular lattice and initialize coordinate attributes. + + y is the first dimension, x is the second. Edges are straight along + the y dimension, so boundary must be staggered in the x dimension, if + not periodic. + + Returns: + A generator of tuples, where each tuple represents an edge between + two spins in the lattice. + """ + length_y, length_x = self.dimensions + + graph = nx.Graph() + for x in range(length_x): + for y in range(length_y): + graph.add_node((y, x)) + + for x in range(length_x): + for y in range(length_y): + # Do y couplers + if y < length_y - 1 or self.periodic[0]: + graph.add_edge((y, x), ((y + 1) % length_y, x)) + + if x < length_x - 1 or self.periodic[1]: + + # Do up-up couplers + graph.add_edge((y, x), (y, (x + 1) % length_x)) + # Do up-down couplers + if y > 0 or self.periodic[0]: + graph.add_edge((y, x), ((y - 1) % length_y, (x + 1) % length_x)) + + num_nodes = graph.number_of_nodes() + relabeling = {self.coordinates(v): v for v in range(num_nodes)} + graph = nx.relabel_nodes(graph, relabeling) + + self.sublattice = np.array([(v - (v // length_y)) % 3 for v in range(num_nodes)]) + + self.integer_coords = [ + (self.coordinates(v)[1], (self.coordinates(v)[0])) for v in range(num_nodes) + ] + self.xy_coords = [ + ( + self.integer_coords[v][0] * 3**0.5 / 2, + self.integer_coords[v][0] / 2 + self.integer_coords[v][1], + ) + for v in range(num_nodes) + ] + self.xy_size = (length_y, length_x * 3**0.5 / 2) # Size as though periodic. + + yield from sorted([tuple(sorted(e)) for e in graph.edges]) + + +class DimerizedTriangular(EmbeddedLattice): + """Dimerized triangular lattice class. + + This class represents a dimerized version of the 2D triangular lattice, + where each node in the logical lattice is represented by a chain of two spins + in the physical lattice. + + Args: + dimensions: Two-element tuple giving the number of spins in the y and x + dimensions. + periodic: Two-element tuple indicating whether the lattice is periodic + in the y and x dimensions. + data_root: A string or Path to the root directory for storing lattice data. + orbit_type: Method for determining qubit and coupler orbits. Must be one of "global", + "standard", "singleton", or "explicit". See ``initialize_orbits`` for details. + qubit_orbits: Explicit qubit orbit labels, used only when ``orbit_type == "explicit"``. + Must have length equal to the number of spins in the lattice. + coupler_orbits: Explicit coupler orbit labels, used only when ``orbit_type == "explicit"``. + Must have length equal to the number of edges in the lattice. + halve_boundary_couplers: A boolean indicating whether to assign half the + coupling strength to boundary couplers in the logical lattice. + chain_strength: The strength of the couplings within each chain. + logical_lattice: Optional logical lattice instance to embed. If not + provided, a ``Triangular`` lattice is constructed from the other + initialization arguments. + """ + + def __init__( + self, + *, + dimensions: tuple[int, int], + periodic: tuple[bool, bool] = (True, False), + data_root: str | Path, + orbit_type: str = "singleton", + qubit_orbits: NDArray | None = None, + coupler_orbits: NDArray | None = None, + halve_boundary_couplers: bool = False, + chain_strength: float = 2, + logical_lattice: Lattice | None = None, + ): + if len(dimensions) != 2: + raise ValueError( + f"DimerizedTriangular requires dimensions of length 2, got {len(dimensions)}." + ) + chain_nodes = {v: (v, v + np.prod(dimensions)) for v in range(np.prod(dimensions))} + self.geometry_name: str = "DimerizedTriangular" + self.num_spins = 2 * int(np.prod(dimensions)) + if logical_lattice is None: + logical_lattice = Triangular( + dimensions=dimensions, + periodic=periodic, + data_root=data_root, + orbit_type=orbit_type, + qubit_orbits=qubit_orbits, + coupler_orbits=coupler_orbits, + halve_boundary_couplers=halve_boundary_couplers, + ) + + super().__init__( + logical_lattice=logical_lattice, + chain_nodes=chain_nodes, + dimensions=dimensions, + periodic=periodic, + data_root=data_root, + orbit_type=orbit_type, + qubit_orbits=qubit_orbits, + coupler_orbits=coupler_orbits, + chain_strength=chain_strength, + ) + self.halve_boundary_couplers: bool = self.logical_lattice.halve_boundary_couplers + + def get_chain_connectivity( + self, + u: Hashable, + v: Hashable | None = None, + ) -> tuple[tuple[int, int]]: + """Return the connectivity for a chain or edge in the logical lattice. + + Args: + u: The first node in the logical edge. + v: The second node in the logical edge. If None, this is treated as + a chain edge (u == v). + Returns: + A tuple of tuples, where each inner tuple represents a pair of indices + in the chains corresponding to u and v that should be connected. For + a chain edge (u == v or v is None), this will return pairs of indices + within the same chain. For a logical edge (u != v), this will return + pairs of indices between the two chains. + """ + if u == v or v is None: + # Interior chain connectivity. + # Generic version: add all possible edges. + return ((0, 1),) + + # Connectivity between two edges. + # Triangular version + uy, ux = self.logical_lattice.coordinates(u) + vy, vx = self.logical_lattice.coordinates(v) + + if ux == vx: # straight up. + if uy > vy or (uy == 0 and vy == self.dimensions[0] - 1 and self.periodic[0]): + return ((0, 1),) + return ((1, 0),) + + # x-edge, i.e. tilted. + if ux == vx - 1 or vx == 0: # (ux == self.dimensions[1] - 1 and self.periodic[1]): + if uy == vy: + return ((1, 0),) + return ((0, 1),) + + if uy == vy: + return ((0, 1),) + return ((1, 0),) diff --git a/dwave/experimental/lattice_utils/observable/__init__.py b/dwave/experimental/lattice_utils/observable/__init__.py new file mode 100644 index 0000000..8f32897 --- /dev/null +++ b/dwave/experimental/lattice_utils/observable/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dwave.experimental.lattice_utils.observable.kinks import * +from dwave.experimental.lattice_utils.observable.observable import * +from dwave.experimental.lattice_utils.observable.triangular import * diff --git a/dwave/experimental/lattice_utils/observable/kinks.py b/dwave/experimental/lattice_utils/observable/kinks.py new file mode 100644 index 0000000..3f55717 --- /dev/null +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -0,0 +1,63 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import dimod +import numpy as np +from dimod import BQM, SampleSet +from numpy.typing import NDArray + +from dwave.experimental.lattice_utils.observable.observable import Observable + +if TYPE_CHECKING: + from dwave.experimental.lattice_utils.experiment.experiment import Experiment + +__all__ = ['KinkKinkCorrelator'] + + +class KinkKinkCorrelator(Observable): + """A class for computing the kink-kink correlator for 1D spin chains.""" + + def evaluate(self, experiment: Experiment, bqm: BQM, sample_set: SampleSet) -> NDArray: + """Compute the kink-kink correlator for 1D spin chains. + + Args: + experiment: The experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: The samples on which to compute the kink-kink correlator. + + Returns: + A numpy array containing the kink-kink correlator values for each sample. + """ + samples = dimod.as_samples(sample_set)[0] + + shifted_samples = np.roll(samples, 1, axis=1) + kink_mask = shifted_samples * samples == np.sign(experiment.param["signed_energy_scale"]) + chain_length = kink_mask.shape[-1] + kink_mask = np.reshape(kink_mask, (-1, chain_length)) + kink_density = np.mean(kink_mask) + + kink_kink_correlator = np.zeros((kink_mask.shape[-1],)) + + mean_kink = np.mean(kink_mask) + for distance in range(1, chain_length): + shifted_kink_mask = np.roll(kink_mask, distance, axis=1) + kink_kink_correlator[distance] = np.mean(kink_mask * shifted_kink_mask) - mean_kink**2 + + kink_kink_correlator /= kink_density**2 + + return kink_kink_correlator diff --git a/dwave/experimental/lattice_utils/observable/observable.py b/dwave/experimental/lattice_utils/observable/observable.py new file mode 100644 index 0000000..fed8475 --- /dev/null +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -0,0 +1,388 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Any, TypeAlias, TYPE_CHECKING + +import dimod +import numpy as np +from numpy.typing import NDArray + +from dwave.experimental.lattice_utils.lattice.lattice import Lattice + +if TYPE_CHECKING: + from dwave.experimental.lattice_utils.experiment.experiment import Experiment + +__all__ = [ + 'Observable', + 'QubitMagnetization', + 'CouplerCorrelation', + 'CouplerFrustration', + 'SampleEnergy', + 'BitpackedSpins', + 'ReferenceEnergy', +] + +ObservableResult: TypeAlias = NDArray | float | int | tuple[NDArray, tuple[int, int]] + + +class Observable(ABC): + """Abstract base class for observables in lattice experiments. + + Each observable should inherit from this class and implement the 'evaluate' + method, which computes the observable from a given sample set. + """ + + def __init__(self): + self.name: str = type(self).__name__ + + @abstractmethod + def evaluate( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample_set: dimod.SampleSet, + ) -> ObservableResult: + """Compute the observable from the provided sample set. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: The samples used to compute the observable. + """ + + +class QubitMagnetization(Observable): + """Compute the mean magnetization of each qubit.""" + + def evaluate( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample_set: dimod.SampleSet, + ) -> NDArray: + """Return per-qubit mean spin values over the provided samples. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples used to compute the magnetization. + + Returns: + A numpy array containing the mean magnetization for each qubit. + """ + sample_array = dimod.as_samples(sample_set)[0].astype(float) + return np.mean(sample_array, axis=0) + + +class CouplerCorrelation(Observable): + """Compute pairwise spin correlations for each coupler.""" + + def evaluate( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample_set: dimod.SampleSet, + ) -> NDArray: + """Return per-coupler pairwise spin correlations over the provided samples. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples used to compute the coupler correlations. + + Returns: + A numpy array containing the pairwise spin correlations for each coupler. + """ + sample_array = dimod.as_samples(sample_set)[0].astype(float) + if not experiment.inst.edge_list: + return np.empty(0, dtype=float) + + row, col = np.asarray(experiment.inst.edge_list).T + spin_product = np.matmul(sample_array.T, sample_array)[row, col] / len(sample_array) + return spin_product + + +class CouplerFrustration(Observable): + """Compute the mean coupler frustration for each edge.""" + + def evaluate( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample_set: dimod.SampleSet, + ) -> NDArray: + """Return the mean coupler frustration over the provided samples. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples used to compute the mean coupler frustration. + + Returns: + A numpy array containing the mean coupler frustration for each edge. + """ + sample_array = dimod.as_samples(sample_set)[0].astype(float) + if not experiment.inst.edge_list: + return np.empty(0, dtype=float) + + row, col = np.asarray(experiment.inst.edge_list).T + spin_product = np.matmul(sample_array.T, sample_array)[row, col] / len(sample_array) + coupler_signs = np.sign( + [bqm.quadratic[edge] for edge in experiment.inst.edge_list] + ) * np.sign(experiment.param["signed_energy_scale"]) + + return spin_product * coupler_signs / 2 + 1 / 2 + + +class SampleEnergy(Observable): + """Compute sample energies with respect to the nominal BQM. + + Energies exclude the magnitude of ``signed_energy_scale`` but include its sign. + """ + + def evaluate( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample_set: dimod.SampleSet, + ) -> NDArray: + """Return signed sample energies from the sample set. + + Args: + experiment: Experiment context providing ``signed_energy_scale``. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples containing energy data. + + Returns: + A numpy array containing the sample energies multiplied by the sign + of ``signed_energy_scale``. + """ + return sample_set.data_vectors["energy"] * np.sign(experiment.param["signed_energy_scale"]) + + +class BitpackedSpins(Observable): + """Compute bitpacked spins.""" + + def evaluate( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample_set: dimod.SampleSet, + ) -> tuple[NDArray, tuple[int, int]]: + """Return bitpacked spin samples and their original array shape. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples containing the spin values to unpack. + + Returns: + A tuple containing the bitpacked spin array and the original array + shape. + """ + sample_array = dimod.as_samples(sample_set)[0] + + # Bitpack solutions + results_bool = np.equal(sample_array, 1) + results_bitpacked = np.packbits(results_bool) + results_shape = sample_array.shape + + return results_bitpacked, results_shape + + +class ReferenceEnergy(Observable): + """Return a cached reference energy, computing it and saving it if needed.""" + + def evaluate( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample_set: dimod.SampleSet, + path: str | Path | None = None, + inst: Lattice | None = None, + ) -> float: + """Get the reference energy for the given BQM, computing and caching it + if needed. + + Args: + experiment: The experiment for which to get the reference energy. Used + to determine the path for caching and loading the reference energy. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: The sample set is not used in this observable, but is + included in the signature for consistency with other observables. + path: Optional path to load/save the reference energy. If not provided, + a default path will be generated based on the experiment and BQM. + + Returns: + The reference energy for the given BQM. + """ + if path is not None: + path = Path(path) + else: + path = get_reference_energy_path(bqm, experiment) + + if path.exists(): + energy, sample, method_string = self.load(experiment, bqm, path) + return energy + + # And if we can't load, we generate a reference sample. + if experiment is not None: + energy, sample, method_string = experiment.inst.optimize(bqm) + elif inst is not None: + energy, sample, method_string = inst.optimize(bqm) + else: + raise ValueError( + "Must provide either an experiment or a lattice to compute reference energy." + ) + + self.save(path, energy, sample, method_string) + + return energy + + def load( + self, + experiment: Experiment, + bqm: dimod.BQM, + path: str | Path | None = None, + ) -> tuple[float, NDArray, str]: + """Load and get the full data tuple, not just the energy. + + Args: + experiment: The experiment for which to load the reference energy. + bqm: The binary quadratic model corresponding to the problem instance. + path: Optional path to load the reference energy. If not provided, + a default path will be generated based on the experiment and BQM. + + Returns: + A tuple containing the reference energy, the corresponding sample + as a NumPy array, and a string indicating the optimization method used. + """ + if path is not None: + path = Path(path) + else: + path = get_reference_energy_path(bqm, experiment) + + with open(path, "r") as f: + method_string = f.readline().strip() + energy = float(f.readline().strip()) + + sample = np.loadtxt(path, skiprows=2) + + return energy, sample, method_string + + def save(self, path: str | Path, energy: float, sample: NDArray, method_string: str) -> None: + """Save the reference energy to disk. + + Args: + path: Path to save the reference energy file. + energy: The reference energy to save. + sample: The corresponding sample to save. + method_string: A string indicating the optimization method used to + obtain the reference energy. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + np.savetxt(path, sample, fmt="%d", header=f"{method_string}\n{energy}", comments="") + + def update( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample: NDArray, + path: str | Path | None = None, + ) -> None: + """Update the cached reference energy if the provided sample improves it. + + Use this when you get an energy that is lower than the reference energy. + We want to keep the old method string unless it is specified. + + Args: + experiment: The experiment for which to update the reference energy. + bqm: The binary quadratic model corresponding to the problem instance. + sample: The new sample that may improve the reference energy. + path: Optional path to load/save the reference energy. If not provided, + a default path will be generated based on the experiment and BQM. + """ + if path is not None: + path = Path(path) + + reference_energy, _, reference_method_string = self.load(experiment, bqm, path) + new_energy = bqm.energy(sample) + + if new_energy >= reference_energy: + raise ValueError("New energy is not better than reference energy, not updating.") + + if path is None: + path = get_reference_energy_path(bqm, experiment) + self.save(path, new_energy, sample, reference_method_string) + + +def get_reference_energy_path( + bqm: dimod.BQM, + experiment: Experiment | None = None, + root: str | Path | None = None, + dummy_experiment_data_dict: dict[str, Any] | None = None, +) -> Path: + """Return the path to the reference energy file for the given experiment and BQM. + + This should be revised if relevant factors are not captured in the instance + pathstring, for example when ground-state energies depend on the specific chip. + + Args: + bqm: The binary quadratic model for which to get the reference energy path. + experiment: The experiment for which to get the reference energy path. + root: Optional root directory to use instead of the experiment's data root. + dummy_experiment_data_dict: A dictionary containing the keys ``run_index``, + ``num_random_instances``, and ``inst`` to use when no experiment is + provided. This allows for generation of dummy experiment data without + all the overhead, for running without an actual experiment. + + Returns: + The path to the reference energy file. + """ + if experiment is None: + if dummy_experiment_data_dict is None: + raise ValueError("Provide either 'experiment' or 'dummy_experiment_data_dict'.") + experiment_data_dict = dummy_experiment_data_dict + else: + experiment_data_dict = { + "run_index": experiment.run_index, + "num_random_instances": experiment.param["num_random_instances"], + "inst": experiment.inst, + } + + if root is None: + root = experiment_data_dict["inst"].data_root + else: + root = Path(root) + + path = ( + root + / "lattice_data" + / "reference_energies" + / experiment_data_dict["inst"]._get_instance_pathstring() + ) + + # Use hash. BQM is not hashable so use the experiment.inst data to generate a tuple. + bqm_as_tuple = tuple(bqm.linear[v] for v in sorted(bqm.variables)) + tuple( + bqm.quadratic[e] for e in experiment_data_dict["inst"].edge_list + ) + bqm_hash = hash(bqm_as_tuple) + path = path / str(bqm_hash) + + return path.with_suffix('.txt') diff --git a/dwave/experimental/lattice_utils/observable/triangular.py b/dwave/experimental/lattice_utils/observable/triangular.py new file mode 100644 index 0000000..a24f7b4 --- /dev/null +++ b/dwave/experimental/lattice_utils/observable/triangular.py @@ -0,0 +1,83 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import dimod +import numpy as np +from dimod import BQM +from numpy.typing import NDArray + +from dwave.experimental.lattice_utils.observable.observable import Observable + +if TYPE_CHECKING: + from dwave.experimental.lattice_utils.experiment.experiment import Experiment + +__all__ = ['TriangularOP'] + + +class TriangularOP(Observable): + """A class for calculating the order parameter of triangular lattices.""" + + def evaluate( + self, + experiment: Experiment, + bqm: BQM, + sample_set: dimod.SampleSet, + ) -> NDArray: + """Calculate the triangular lattice order parameter. + + This observable uses the three-sublattice complex order parameter described in + `King et al. (2023) _`. + + Args: + experiment: The experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: The samples on which to compute the order parameter. + + Returns: + A numpy array containing the order parameter values for each sample. + """ + + # If the lattice is an embedded lattice then the BQM and sampleset must be unembedded. + if hasattr(experiment.inst, "logical_lattice"): + lbqm = experiment.inst.unembed_bqm(bqm) + + lss = experiment.inst.unembed_sampleset(sample_set) + triangular_sublattice = experiment.inst.logical_lattice.sublattice + else: + lbqm, lss = bqm, sample_set + triangular_sublattice = experiment.inst.sublattice + + sample_array = dimod.as_samples(lss)[0] + + for u, v in lbqm.quadratic: + if triangular_sublattice[u] == triangular_sublattice[v]: + raise ValueError( + "Invalid triangular sublattice assignment: edge " + f"({u}, {v}) connects nodes in the same sublattice" + ) + + sublattice_mags = np.zeros((sample_array.shape[0], 3), dtype=float) + for sublattice in range(3): + sublattice_mags[:, sublattice] = np.mean( + sample_array[:, triangular_sublattice == sublattice], axis=1 + ) + + angles = np.array(np.exp([0.0, 1.0j * 4 * np.pi / 3, 1.0j * 2 * np.pi / 3])).T + order_parameter = np.matmul(sublattice_mags, angles).ravel() / np.sqrt(3) + + return order_parameter diff --git a/dwave/experimental/lattice_utils/utils.py b/dwave/experimental/lattice_utils/utils.py new file mode 100644 index 0000000..b191cd7 --- /dev/null +++ b/dwave/experimental/lattice_utils/utils.py @@ -0,0 +1,113 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Callable, Iterator + +import numpy as np +from numpy.typing import NDArray + + +def bootstrap( + array: NDArray, + rng: np.random.Generator, + repetitions: int = 200, + bootstrap_function: Callable[[NDArray], float] = np.nanmedian, + skipnan: bool = True, +) -> list[float]: + """Estimate a statistic by bootstrap resampling. + + The input is flattened to one dimension before resampling. For each + bootstrap sample, the statistic is computed by applying `bootstrap_function` + to the resampled array. + + If `skipnan` is True, then NaN values are removed from the input before + resampling. If all values are NaN, then the output is a list of NaN values. + + Args: + array: The input data to resample. + rng: A random number generator used for resampling. + repetitions: The number of bootstrap samples to generate. + bootstrap_function: A function that takes an array and returns a statistic. + skipnan: Whether to ignore NaN values in the input. + + Returns: + A list of bootstrap estimates of the statistic. + """ + array = np.asarray(np.atleast_1d(array)).ravel() + if skipnan: + array = array[~np.isnan(array)] + if len(array) == 0: + return [np.nan] * repetitions + + output = [] + if len(array) > 0: + for inds in generate_bootstrap_indices(array.size, repetitions, rng): + output.append(bootstrap_function(array[inds])) + + return output + + +def generate_bootstrap_indices( + size: int, + repetitions: int, + rng: np.random.Generator, +) -> Iterator[NDArray]: + """Yield bootstrap index arrays of a given size and number of repetitions. + + Each index array is generated by sampling with replacement from the range of + indices corresponding to the input size. + + Args: + size: The size of the array for which to generate bootstrap indices. + repetitions: The number of bootstrap index arrays to generate. + rng: A random number generator used for sampling. + + Yields: + An array of indices for a bootstrap sample. + """ + for _ in range(repetitions): + inds = rng.choice(range(size), replace=True, size=size) + yield inds + + +def confidence_interval(array: NDArray, width: float = 0.95) -> tuple[float, float, float]: + """Calculate a confidence interval for a statistic using quantiles. + + The input is flattened to one dimension before calculating quantiles. The + confidence interval is calculated by taking the quantiles corresponding to + the specified width. + + Args: + array: The input data from which to calculate the confidence interval. + width: The width of the confidence interval (e.g., 0.95 for a 95% + confidence interval). + + Returns: + A tuple of the form (median, lower_error, upper_error) where: + - median is the median of the input array. + - lower_error is the distance from the median to the lower bound of + the confidence interval. + - upper_error is the distance from the median to the upper bound of + the confidence interval. + """ + x = np.asarray(array).ravel() + if len(x) == 0: + return np.nan, np.nan, np.nan + + x.sort() + low = x[int(np.floor((1 - width) / 2 * x.size))] + high = x[int(np.floor((1 - (1 - width) / 2) * x.size))] + med = np.median(x) + + return med, med - low, high - med diff --git a/examples/lattice_utils/1D_Ising_chain_shim.py b/examples/lattice_utils/1D_Ising_chain_shim.py new file mode 100644 index 0000000..95fd885 --- /dev/null +++ b/examples/lattice_utils/1D_Ising_chain_shim.py @@ -0,0 +1,184 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shimming example for 1D Ising chain.""" + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +from dwave.system import DWaveSampler + +from dwave.experimental.lattice_utils import lattice, experiment + +# Set up the parameters + +# Get a QPU sampler with Zephyr topology. +sampler = DWaveSampler(solver=dict(topology__type="zephyr")) + +NUM_SPINS = 256 +SIGNED_ENERGY_SCALE = -1.6 + +# We will simulate four orders of magnitude in anneal time. +# File format rounds to the nearest picosecond, so we will do so explicitly here. +ANNEAL_TIMES = np.round(np.geomspace(0.01, 100, 3), 6) + +errorbar_style = {"marker": "", "linestyle": "", "capsize": 2} +point_style = {"linestyle": "", "markersize": 5} +cm = plt.get_cmap("tab10") + +data_root = Path(__file__).resolve().parents[1] + +# Make a lattice instance for a periodic 256-spin chain, so we can embed it. +inst = lattice.Chain( + dimensions=(NUM_SPINS,), + data_root=data_root, + periodic=(True,), + orbit_type="standard", +) + +# Find parallel embeddings of the lattice heuristically. The embed_lattice +# function is heuristic and is run here with a default timeout (10s) and no +# tuning of any parameters. Larger and more complex lattices can take longer +# to embed. +inst.embed_lattice(sampler) + +# Time to make an experiment. The aim in this example is to demonstrate a coupler +# and flux-bias shim on a chain, at fixed energy scale and varying anneal time. +# Since the susceptibility to changes in the shim changes as a function of anneal +# time, we will run each anneal time using separate shim steps. +# We will also set the orbit_type to 'standard', allowing the use of graph +# automorphisms to determine symmetries in the system that can be exploited by +# shimming. In this case, all couplers are equivalent (they go in the same orbit) +# so the coupler shim will compel them all to have the same spin-spin correlation +# for a given parameterization. +flux_bias_shim_step = { + 0.01: 5e-6, + 1.0: 2e-6, + 100: 0.5e-6, +} +coupler_shim_step = { + 0.01: 0.05, + 1.0: 0.2, + 100: 1.0, +} +for anneal_time in ANNEAL_TIMES: + + config = experiment.FastAnnealExperimentConfig( + signed_energy_scale=SIGNED_ENERGY_SCALE, + coupler_shim_step=coupler_shim_step[anneal_time], + flux_bias_shim_step=flux_bias_shim_step[anneal_time], + ) + exp = experiment.Experiment(inst=inst, sampler=sampler, max_iterations=100, config=config) + + # Make parameter list. We will only vary anneal time. + for _ in range(120): + done = exp.run_iteration([{"anneal_time": anneal_time}]) + if done: + break + + +# We will make a dict for the data we want to analyze. For each anneal time +# we will load all iterations into a corresponding dict entry. +# Disjoint embeddings are given along a separate axis in the results, except +# flux biases and anneal offsets, which are given as a single +# array since they are indexed by physical qubit. + +mag = {} # average qubit magnetization +frust = {} # average coupler frustration (kink density) +cshim = {} # coupler shim +fbshim = {} # flux bias shim +for anneal_time in ANNEAL_TIMES: + exp.apply_param({"anneal_time": anneal_time}) + res = exp.load_results() + mag[anneal_time] = np.array([it["QubitMagnetization"] for it in res]) + frust[anneal_time] = np.array([it["CouplerFrustration"] for it in res]) + cshim[anneal_time] = np.array( + [it["shim_data"]["relative_coupler_strength"].ravel() for it in res] + ) + fbshim[anneal_time] = np.array([it["shim_data"]["flux_biases"] for it in res]) + +title = ( + f"1D chain shim, " + f"{'x'.join([str(dim) for dim in inst.dimensions])}, " + f"J={exp.param['signed_energy_scale']}, " + f"{sampler.solver.name}" +) +fig, axes = plt.subplots(3, 6, figsize=(16, 8), sharex="col", sharey="col") +fig.suptitle(title, fontsize=16) + +for iat, anneal_time in enumerate(ANNEAL_TIMES): + + # Plot std of qubit magnetization + ax = axes[iat, 0] + ax.plot(mag[anneal_time].std(axis=(1, 2)), label=r"Mag std") + ax.set_ylabel("Magnetization std") + + # Plot histograms of first and last iterations + ax = axes[iat, 1] + ax.hist(mag[anneal_time][:5].ravel(), label="First 5 iterations", alpha=0.5) + ax.hist(mag[anneal_time][-5:].ravel(), label="Last 5 iterations", alpha=0.5) + ax.set_yticks([]) + ax.set_ylabel("Frequency") + + # Plot std of coupler frustration + ax = axes[iat, 2] + ax.plot(frust[anneal_time].std(axis=(1, 2)), label=r"Frust std") + ax.set_ylabel("Frustration std") + + # Plot histograms of first and last iterations + ax = axes[iat, 3] + ax.hist(frust[anneal_time][:5].ravel(), label="First iteration", alpha=0.5) + ax.hist(frust[anneal_time][-5:].ravel(), label="Last iteration", alpha=0.5) + ax.set_yticks([]) + ax.set_ylabel("Frequency") + + # Plot flux biases + ax = axes[iat, 4] + ax.plot(fbshim[anneal_time], alpha=0.2) + ax.set_ylabel("Flux bias") + + # Plot coupler shim + ax = axes[iat, 5] + ax.plot(cshim[anneal_time], alpha=0.2) + ax.set_ylabel("Rel. cplr. strength") + + +axes[iat, 0].set_xlabel("Iteration") +axes[iat, 1].set_xlabel("Magnetization") +axes[iat, 1].legend() +axes[iat, 2].set_xlabel("Iteration") +axes[iat, 3].set_xlabel("Frustration") +axes[iat, 3].legend() +axes[iat, 4].set_xlabel("Iteration") +axes[iat, 5].set_xlabel("Iteration") + +axes[0, 0].set_title("Magnetization std") +axes[0, 1].set_title("Magnetization") +axes[0, 2].set_title("Frustration std") +axes[0, 3].set_title("Frustration") +axes[0, 4].set_title("Flux offset shim") +axes[0, 5].set_title("Coupler shim") + +fig.tight_layout() + +filename = title +for bad_symbol in "/: ;,": + filename = filename.replace(bad_symbol, "_") + +# Create a folder to save figures in if it doesn't already exist. +(data_root / "figures").mkdir(exist_ok=True) + +fig.savefig(data_root / "figures" / f"{filename}.png") +plt.show() diff --git a/examples/mca_shim_AO_FB.py b/examples/mca_shim_AO_FB.py index 72a494a..547f4f3 100644 --- a/examples/mca_shim_AO_FB.py +++ b/examples/mca_shim_AO_FB.py @@ -25,6 +25,12 @@ import numpy as np from tqdm import tqdm +from pathlib import Path +import sys +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + import dimod from dwave.system import DWaveSampler from dwave.system.composites import ParallelEmbeddingComposite @@ -38,7 +44,6 @@ ) from dwave.experimental.shimming import shim_flux_biases - def _make_anneal_schedules( exp_feature_info: list, target_c: float = 0.37, diff --git a/pyproject.toml b/pyproject.toml index 57a587a..1b8a959 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,9 @@ dependencies = [ # https://github.com/python/cpython/issues/142214 requires-python = ">=3.10, !=3.14.1" +[project.optional-dependencies] +progress = ["tqdm>=4,<5"] + [project.urls] Issues = "https://github.com/dwavesystems/dwave-experimental/issues" Repository = "https://github.com/dwavesystems/dwave-experimental.git" diff --git a/releasenotes/notes/add-lattice-utils-49e5efd79268e69d.yaml b/releasenotes/notes/add-lattice-utils-49e5efd79268e69d.yaml new file mode 100644 index 0000000..71d728b --- /dev/null +++ b/releasenotes/notes/add-lattice-utils-49e5efd79268e69d.yaml @@ -0,0 +1,20 @@ +--- +features: + - | + Add ``dwave.experimental.lattice_utils`` submodule with utilities for + constructing lattice graphs, evaluating physics + observables on sample sets, and running shimmed Ising experiments on + QPU samplers. + - | + Add lattice graph constructors ``Chain``, ``Triangular``, and + ``DimerizedTriangular`` for building common lattice geometries. + - | + Add the ``EmbeddedLattice`` class for embedding logical lattices onto physical + hardware graphs using ordered chains. + - | + Add the ``Observable`` base class and several observables (e.g. qubit + magnetization, coupler correlation, and coupler frustration) for evaluating + physics quantities on sample sets. + - | + Add the ``Experiment`` class for running shimmed Ising experiments on QPU + samplers and collecting observables from the resulting sample sets. diff --git a/tests/__init__.py b/tests/__init__.py index de79690..0002aae 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/test_lattice_utils/__init__.py b/tests/test_lattice_utils/__init__.py new file mode 100644 index 0000000..0002aae --- /dev/null +++ b/tests/test_lattice_utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/test_lattice_utils/_helpers.py b/tests/test_lattice_utils/_helpers.py new file mode 100644 index 0000000..680efbf --- /dev/null +++ b/tests/test_lattice_utils/_helpers.py @@ -0,0 +1,103 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared test fixtures for the ``lattice_utils`` test package.""" + +from unittest import mock + +import dimod +import numpy as np + +from dwave.experimental.lattice_utils import lattice + +__all__ = [ + "_make_triangular", + "_make_mock_sampler", + "_make_mock_experiment", + "_make_embedded_chain", +] + + +def _make_triangular( + data_root, + ly=3, + lx=3, + periodic=(True, False), + orbit_type="singleton", + halve_boundary_couplers=False, +): + return lattice.Triangular( + dimensions=(ly, lx), + periodic=periodic, + data_root=data_root, + orbit_type=orbit_type, + halve_boundary_couplers=halve_boundary_couplers, + ) + + +def _make_mock_sampler( + num_qubits=128, + nodelist=None, + solver_name="TestSolver", + type_name="DWaveSampler", + *, + sync_response=False, +): + """Build a mock sampler resembling DWaveSampler. + + Production code detects the sampler via ``type(sampler).__name__`` and reads + ``nodelist``, ``properties["num_qubits"]``, and ``solver.name``. When + ``sync_response`` is True, ``sampler.sample(...)`` returns a mock response + mimicking the DWaveSampler async interface (``.done()`` -> True, + ``.samples()`` -> all-ones ndarray) used by ``run_iteration`` tests. + """ + sampler = mock.MagicMock(spec=dimod.Sampler) + type(sampler).__name__ = type_name + if nodelist is None: + nodelist = list(range(num_qubits)) + sampler.nodelist = nodelist + sampler.properties = {"num_qubits": num_qubits} + sampler.solver = mock.MagicMock() + sampler.solver.name = solver_name + if sync_response: + response = mock.MagicMock() + response.done.return_value = True + response.samples.return_value = np.ones((10, num_qubits), dtype=float) + sampler.sample.return_value = response + return sampler + + +def _make_mock_experiment(inst, signed_energy_scale=1.0, run_index=0, num_random_instances=1): + """Return a lightweight mock Experiment with .inst and .param.""" + exp = mock.MagicMock() + exp.inst = inst + exp.param = { + "signed_energy_scale": signed_energy_scale, + "num_random_instances": num_random_instances, + } + exp.run_index = run_index + return exp + + +def _make_embedded_chain(chain_nodes, data_root): + return lattice.EmbeddedLattice( + logical_lattice=lattice.Chain( + dimensions=(len(chain_nodes),), + periodic=(False,), + data_root=data_root, + ), + chain_nodes=chain_nodes, + dimensions=(sum(len(chain) for chain in chain_nodes.values()),), + periodic=(False,), + ) diff --git a/tests/test_lattice_utils/test_experiment.py b/tests/test_lattice_utils/test_experiment.py new file mode 100644 index 0000000..69a0427 --- /dev/null +++ b/tests/test_lattice_utils/test_experiment.py @@ -0,0 +1,466 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import lzma +import pickle +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +import dimod +import numpy as np + +from dwave.experimental.lattice_utils import experiment, lattice +from tests.test_lattice_utils._helpers import _make_mock_sampler + + +class TestSamplerCall(unittest.TestCase): + def test_defaults(self): + sc = experiment.SamplerCall(run_index=0) + self.assertEqual(sc.run_index, 0) + self.assertIsNone(sc.bqm) + self.assertEqual(sc.shim_data, {}) + self.assertEqual(sc.logical_bqms, []) + self.assertEqual(sc.sampler_params, {}) + + def test_with_values(self): + bqm = dimod.BQM(vartype="SPIN") + sc = experiment.SamplerCall( + run_index=5, + bqm=bqm, + shim_data={"total_iterations": 1}, + logical_bqms=[bqm], + sampler_params={"num_reads": 100}, + ) + self.assertEqual(sc.run_index, 5) + self.assertIs(sc.bqm, bqm) + self.assertEqual(sc.shim_data["total_iterations"], 1) + + +class TestExperimentInit(unittest.TestCase): + def test_default_params(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + self.assertEqual(exp.param["signed_energy_scale"], 1.0) + self.assertEqual(exp.param["num_reads"], 100) + self.assertIs(exp.inst, chain) + + +class TestApplyParam(unittest.TestCase): + def test_data_path_with_schedule(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_schedule": [(0, 1), (5, 0.5)]}) + self.assertIn("asched", str(exp.data_path)) + + def test_apply_param_unknown_sampler_raises(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler(type_name="UnknownSampler") + exp = experiment.Experiment(inst=chain, sampler=sampler) + with self.assertRaises(TypeError): + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + + def test_apply_param_sets_run_index_zero(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + self.assertEqual(exp.run_index, 0) + + def test_apply_param_resumes_from_existing_iterations(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + for i in range(3): + fn = exp.data_path / f"iter{i:05d}.pkl.lzma" + fn.parent.mkdir(parents=True, exist_ok=True) + with lzma.open(fn, "wb") as f: + pickle.dump({}, f) + + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + self.assertEqual(exp.run_index, 3) + + +class TestShimdata(unittest.TestCase): + def test_initial_shim_no_embeddings(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.already_initialized = False + shim_data = exp._make_initial_shim() + self.assertEqual(shim_data["total_iterations"], 0) + self.assertNotIn("flux_biases", shim_data) + + def test_initial_shim_with_embeddings(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler(num_qubits=128) + exp = experiment.Experiment(inst=chain, sampler=sampler) + shim_data = exp._make_initial_shim() + self.assertIn("flux_biases", shim_data) + self.assertEqual(len(shim_data["flux_biases"]), 128) + + def test_initial_shim_with_preset_flux_biases(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler(num_qubits=128) + fb = np.ones(128) * 0.01 + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.param["flux_biases"] = fb + shim_data = exp._make_initial_shim() + np.testing.assert_array_almost_equal(shim_data["flux_biases"], fb) + + def test_load_shim_from_file(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.run_index = 1 + exp.data_path = Path(tmpdir) + + shim_data = {"total_iterations": 5, "flux_biases": np.zeros(10)} + data = {"shim_data": shim_data} + fn = Path(tmpdir) / "iter00000.pkl.lzma" + with lzma.open(fn, "wb") as f: + pickle.dump(data, f) + + loaded = exp._load_shim() + self.assertEqual(loaded["total_iterations"], 5) + + def test_load_shim_empty_file(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.run_index = 1 + exp.data_path = Path(tmpdir) + + fn = Path(tmpdir) / "iter00000.pkl.lzma" + fn.touch() + + with self.assertRaises(FileNotFoundError): + exp._load_shim() + + def test_load_shim_corrupted_file(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.run_index = 1 + exp.data_path = Path(tmpdir) + fn = Path(tmpdir) / "iter00000.pkl.lzma" + fn.write_bytes(b"not a valid lzma file") + + with self.assertRaises(OSError): + exp._load_shim() + + def test_load_shim_missing_file(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.data_path = Path(tmpdir) + exp.run_index = 1 + + with self.assertRaises(FileNotFoundError): + exp._load_shim() + + def test_get_shim_data_not_initialized(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.already_initialized = False + shim_data = exp._get_shim_data() + self.assertEqual(shim_data["total_iterations"], 0) + + +class TestCouplerShim(unittest.TestCase): + def test_coupler_shim_basic_update(self): + """rcs += step_size * (frust - mean(frust)) within each orbit bin. + + signed_energy_scale=0.5 keeps the effective coupler |J*rcs*scale| below the + truncation thresholds (>1 / <-2) so we can assert the raw update math. + """ + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + orbit_type="global", # all edges in one bin -> update is non-trivial + ) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler() + exp = experiment.Experiment( + inst=chain, + sampler=sampler, + config=experiment.ExperimentConfig(coupler_shim_step=0.01, signed_energy_scale=0.5), + ) + + sc = experiment.SamplerCall(run_index=0) + sc.logical_bqms = [chain.make_bqm()] + sc.shim_data = { + "total_iterations": 0, + "relative_coupler_strength": np.ones((1, chain.num_edges)), + } + # mean(frust) = 0.5, so delta = 0.01 * [-0.2, 0.0, 0.2] = [-0.002, 0.0, 0.002] + # Post-update mean(|rcs|) = 1.0 exactly -> renormalization is a no-op. + # Q = rcs * J(=1) * scale(=0.5) stays in [0.499, 0.501] -> no truncation. + results = {"CouplerFrustration": np.array([[0.3, 0.5, 0.7]])} + + exp._update_coupler_shim(sc, results) + + np.testing.assert_array_almost_equal( + sc.shim_data["relative_coupler_strength"], + np.array([[0.998, 1.0, 1.002]]), + ) + + def test_coupler_shim_singleton_orbits_is_noop(self): + """With singleton orbits, mean equals the single value, so rcs is unchanged.""" + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + orbit_type="singleton", + ) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler() + exp = experiment.Experiment( + inst=chain, + sampler=sampler, + config=experiment.ExperimentConfig(coupler_shim_step=0.01), + ) + + sc = experiment.SamplerCall(run_index=0) + sc.logical_bqms = [chain.make_bqm()] + rcs_before = np.ones((1, chain.num_edges)) + sc.shim_data = { + "total_iterations": 0, + "relative_coupler_strength": rcs_before.copy(), + } + results = {"CouplerFrustration": np.array([[0.3, 0.5, 0.7]])} + + exp._update_coupler_shim(sc, results) + + np.testing.assert_array_equal(sc.shim_data["relative_coupler_strength"], rcs_before) + + +class TestSaveLoadResults(unittest.TestCase): + def test_save_and_reload(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.data_path = Path(tmpdir) + exp.run_index = 0 + data = {"QubitMagnetization": np.zeros(4)} + exp._save_results(data) + fn = Path(tmpdir) / "iter00000.pkl.lzma" + self.assertTrue(fn.exists()) + + with lzma.open(fn, "rb") as f: + loaded = pickle.load(f) + np.testing.assert_array_equal(loaded["QubitMagnetization"], np.zeros(4)) + + def test_save_with_filename_and_run_index_raises(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.data_path = Path(tmpdir) + with self.assertRaises(ValueError): + exp._save_results({}, run_index=0, filename="test.pkl.lzma") + + def test_save_custom_filename(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.data_path = Path(tmpdir) + data = {"x": 1} + exp._save_results(data, filename="custom.pkl.lzma") + self.assertTrue((Path(tmpdir) / "custom.pkl.lzma").exists()) + + def test_load_results_ignore_shim(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + fn = exp.data_path / "iter00000.pkl.lzma" + fn.parent.mkdir(parents=True, exist_ok=True) + with lzma.open(fn, "wb") as f: + pickle.dump({"value": 0, "shim_data": {}}, f) + + results = exp.load_results(num_iterations=1, ignore_shim=True) + self.assertNotIn("shim_data", results[0]) + + def test_load_results_starting_iteration(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + for i in range(10): + fn = exp.data_path / f"iter{i:05d}.pkl.lzma" + fn.parent.mkdir(parents=True, exist_ok=True) + with lzma.open(fn, "wb") as f: + pickle.dump({"value": i, "shim_data": {}}, f) + + results = exp.load_results(num_iterations=3, start_iteration=2) + self.assertEqual(len(results), 3) + + def test_load_results_corrupted_lzma(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + fn = exp.data_path / "iter00000.pkl.lzma" + fn.parent.mkdir(parents=True, exist_ok=True) + fn.write_bytes(b"corrupted data") + + with self.assertRaises(lzma.LZMAError): + exp.load_results(num_iterations=1) + + def test_generate_data_type_conversions(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + sc = experiment.SamplerCall(run_index=0) + sc.shim_data = {"total_iterations": 1, "flux_biases": np.zeros(4)} + + results = { + "QubitMagnetization": np.array([0.1, 0.2, 0.3, 0.4]), + "Complex": np.array([1 + 2j, 3 + 4j]), + "ListData": [1, 2, 3], + } + savedata = exp._generate_data_to_save(sc, results) + self.assertEqual(savedata["QubitMagnetization"].dtype, np.float32) + self.assertEqual(savedata["Complex"].dtype, np.complex64) + self.assertEqual(savedata["shim_data"]["total_iterations"], 1) + + +class TestMakeBqm(unittest.TestCase): + def test_make_bqm_no_embeddings(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment( + inst=chain, + sampler=sampler, + config=experiment.ExperimentConfig(signed_energy_scale=0.5), + ) + sc = experiment.SamplerCall(run_index=0) + sc.logical_bqms = [chain.make_bqm()] + sc.shim_data = {"total_iterations": 0} + bqm = exp._make_bqm(sc) + for u, v in chain.edge_list: + self.assertAlmostEqual(bqm.quadratic[(u, v)], 0.5) + + def test_make_bqm_with_embeddings(self): + """Physical biases = logical_bias * relative_coupler_strength * signed_energy_scale.""" + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + chain.embedding_list = np.array([[10, 11, 12, 13]]) # offset to detect mapping + sampler = _make_mock_sampler() + exp = experiment.Experiment( + inst=chain, + sampler=sampler, + config=experiment.ExperimentConfig(signed_energy_scale=0.5), + ) + sc = experiment.SamplerCall(run_index=0) + sc.logical_bqms = [chain.make_bqm()] + sc.shim_data = { + "total_iterations": 0, + "relative_coupler_strength": np.full((1, chain.num_edges), 2.0), + } + + bqm = exp._make_bqm(sc) + + # Physical variables come from the embedding, not the logical indices. + self.assertEqual(set(bqm.variables), {10, 11, 12, 13}) + for u_log, v_log in chain.edge_list: + self.assertAlmostEqual(bqm.quadratic[(10 + u_log, 10 + v_log)], 1.0) + + self.assertEqual(len(bqm.quadratic), chain.num_edges) + + +class TestRunIteration(unittest.TestCase): + def test_run_iteration_basic(self): + """run_iteration() exercises the full pipeline: build call, sample, parse, shim, save.""" + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + exp = experiment.Experiment( + inst=chain, sampler=_make_mock_sampler(sync_response=True), max_iterations=1 + ) + chain._load_embeddings = mock.MagicMock() + finished = exp.run_iteration([{"signed_energy_scale": 1.0, "anneal_time": 1.0}]) + + self.assertFalse(finished) + result_files = list(exp.data_path.glob("iter*.pkl.lzma")) + self.assertEqual(len(result_files), 1) + + with lzma.open(result_files[0], "rb") as f: + data = pickle.load(f) + + self.assertIn("QubitMagnetization", data) + self.assertIn("CouplerCorrelation", data) + self.assertIn("shim_data", data) + self.assertEqual(data["shim_data"]["total_iterations"], 1) + + def test_run_iteration_returns_true_when_finished(self): + """run_iteration() returns True when max_iterations already reached.""" + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + exp = experiment.Experiment( + inst=chain, + sampler=_make_mock_sampler(sync_response=True), + config=experiment.ExperimentConfig(), + max_iterations=0, + ) + chain._load_embeddings = mock.MagicMock() + finished = exp.run_iteration([{"signed_energy_scale": 1.0, "anneal_time": 1.0}]) + + self.assertTrue(finished) + self.assertEqual(list(exp.data_path.glob("iter*.pkl.lzma")), []) + + +class TestFastAnnealExperiment(unittest.TestCase): + def test_default_params(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + config = experiment.FastAnnealExperimentConfig() + exp = experiment.Experiment(inst=chain, sampler=sampler, config=config) + self.assertTrue(exp.param.get("fast_anneal")) + self.assertEqual(exp.param["num_reads"], 100) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_lattice_utils/test_lattice.py b/tests/test_lattice_utils/test_lattice.py new file mode 100644 index 0000000..3afdeb6 --- /dev/null +++ b/tests/test_lattice_utils/test_lattice.py @@ -0,0 +1,459 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import tempfile +import unittest +from unittest import mock + +import dimod +import numpy as np +from dwave.samplers import SteepestDescentSolver + +from dwave.experimental.lattice_utils import lattice +from tests.test_lattice_utils._helpers import _make_embedded_chain, _make_triangular + + +class TestChain(unittest.TestCase): + def test_periodic(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + self.assertEqual(chain.num_spins, 6) + self.assertEqual(chain.num_edges, 6) + self.assertIn((5, 0), chain.edge_list) + + def test_non_periodic(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(6,), periodic=(False,), data_root=tmpdir) + self.assertEqual(chain.num_spins, 6) + self.assertEqual(chain.num_edges, 5) + self.assertNotIn((5, 0), chain.edge_list) + + def test_single_node_periodic(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(1,), periodic=(True,), data_root=tmpdir) + self.assertEqual(chain.num_spins, 1) + self.assertEqual(chain.num_edges, 0) + + def test_geometry_name(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + self.assertEqual(chain.geometry_name, "Chain") + + def test_default_periodic(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), data_root=tmpdir) + self.assertTrue(chain.periodic[0]) + + def test_edge_list_sorted(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(5,), periodic=(False,), data_root=tmpdir) + for u, v in chain.edge_list: + self.assertLess(u, v) + + def test_bqm_structure(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + self.assertEqual(len(bqm.variables), 4) + self.assertEqual(len(bqm.quadratic), 3) + for u, v in chain.edge_list: + self.assertAlmostEqual(bqm.quadratic[(u, v)], 1.0) + + def test_bqm_vartype(self): + with tempfile.TemporaryDirectory() as tmpdir: + bqm = lattice.Chain(dimensions=(3,), data_root=tmpdir).make_bqm() + self.assertEqual(bqm.vartype, dimod.SPIN) + + +class TestLattice(unittest.TestCase): + def test_orbit_singleton(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), orbit_type="singleton", data_root=tmpdir) + np.testing.assert_array_equal(chain.qubit_orbits, np.arange(4)) + np.testing.assert_array_equal(chain.coupler_orbits, np.arange(chain.num_edges)) + + def test_orbit_global(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), orbit_type="global", data_root=tmpdir) + np.testing.assert_array_equal(chain.qubit_orbits, np.zeros(4, dtype=int)) + np.testing.assert_array_equal( + chain.coupler_orbits, np.zeros(chain.num_edges, dtype=int) + ) + + def test_orbit_explicit(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(True,), + orbit_type="explicit", + qubit_orbits=np.array([0, 0, 1, 1]), + coupler_orbits=np.array([0, 0, 1, 1]), + data_root=tmpdir, + ) + np.testing.assert_array_equal(chain.qubit_orbits, [0, 0, 1, 1]) + + def test_unknown_orbit_type(self): + with tempfile.TemporaryDirectory() as tmpdir: + with self.assertRaises(ValueError): + lattice.Chain( + dimensions=(4,), periodic=(True,), orbit_type="bogus", data_root=tmpdir + ) + + def test_get_path_invalid_kind(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + with self.assertRaises(ValueError): + chain._get_path(None, "invalid") + + def test_standard_orbit_save_and_load(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(True,), + orbit_type="standard", + data_root=tmpdir, + ) + self.assertIsNotNone(chain.qubit_orbits) + self.assertIsNotNone(chain.coupler_orbits) + # Second instantiation should load from disk + chain2 = lattice.Chain( + dimensions=(4,), + periodic=(True,), + orbit_type="standard", + data_root=tmpdir, + ) + np.testing.assert_array_equal(chain.qubit_orbits, chain2.qubit_orbits) + + def test_embed_no_embeddings_found(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = mock.MagicMock() + type(sampler).__name__ = "MockDWaveSampler" + sampler.to_networkx_graph.return_value = chain.make_networkx_graph() + + with mock.patch( + "dwave.experimental.lattice_utils.lattice.lattice.find_multiple_embeddings", + return_value=[], + ): + with self.assertRaises(ValueError): + chain.embed_lattice(sampler, try_to_load=False, timeout=1) + + def test_embed_load_existing(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = mock.MagicMock() + type(sampler).__name__ = "MockDWaveSampler" + sampler.to_networkx_graph.return_value = chain.make_networkx_graph() + + embeddings = np.array([[0, 1, 2, 3]]) + chain._save_embeddings(sampler, embeddings) + + chain.embed_lattice(sampler, try_to_load=True, data_root=tmpdir) + np.testing.assert_array_equal(chain.embedding_list, embeddings) + + def test_embed_find_and_save(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = mock.MagicMock() + type(sampler).__name__ = "MockDWaveSampler" + sampler.to_networkx_graph.return_value = chain.make_networkx_graph() + + emb_dict = {i: i for i in range(4)} + with mock.patch( + "dwave.experimental.lattice_utils.lattice.lattice.find_multiple_embeddings", + return_value=[emb_dict], + ): + chain.embed_lattice(sampler, try_to_load=False, timeout=1, data_root=tmpdir) + # Verify embedding was found and saved + emb_path = chain._get_path("embedding", sampler_name="MockDWaveSampler") + self.assertTrue(emb_path.exists()) + + +class TestLatticeOptimize(unittest.TestCase): + """Tests for the ``Lattice.optimize`` instance method (vs. the free function + ``lattice.optimize`` covered by ``TestOptimizeFunction``).""" + + def test_plain_lattice_private_optimize_default_sampler(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + reference_energy_sampler_kwargs={"num_sweeps": 256, "num_reads": 16}, + ) + bqm = chain.make_bqm() + energy, sample, method = chain.optimize(bqm) + + self.assertEqual(energy, -3) + self.assertEqual(bqm.energy(sample), energy) + self.assertEqual(method, "ExponentialBackoffSimulatedAnnealingSampler") + + def test_plain_lattice_private_optimize_custom_sampler(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + reference_energy_sampler=dimod.ExactSolver(), + ) + bqm = chain.make_bqm() + energy, sample, method = chain.optimize(bqm) + + self.assertEqual(energy, -3) + self.assertEqual(bqm.energy(sample), energy) + self.assertEqual(method, "ExactSolver") + + def test_optimize_with_custom_sampler_steepest_descent(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + reference_energy_sampler=SteepestDescentSolver(), + reference_energy_sampler_kwargs={"initial_states": np.array([[1, -1, 1, -1]])}, + ) + bqm = chain.make_bqm() + energy, sample, method = chain.optimize(bqm) + + self.assertEqual(energy, -3) + self.assertAlmostEqual(bqm.energy(sample), energy) + self.assertEqual(method, "SteepestDescentSolver") + + def test_optimize_with_custom_exponential_backoff_params(self): + with tempfile.TemporaryDirectory() as tmpdir: + sampler = lattice.ExponentialBackoffSimulatedAnnealingSampler( + max_num_sweeps=512, min_num_sweeps=64 + ) + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + reference_energy_sampler=sampler, + reference_energy_sampler_kwargs={"num_reads": 16}, + ) + bqm = chain.make_bqm() + energy, sample, method = chain.optimize(bqm) + + self.assertEqual(energy, -3) + self.assertAlmostEqual(bqm.energy(sample), energy) + self.assertEqual(method, "ExponentialBackoffSimulatedAnnealingSampler") + self.assertEqual(sampler.max_num_sweeps, 512) + self.assertEqual(sampler.min_num_sweeps, 64) + + +class TestTriangular(unittest.TestCase): + def test_basic_construction(self): + with tempfile.TemporaryDirectory() as tmpdir: + tri = lattice.Triangular( + dimensions=(3, 3), + periodic=(True, False), + data_root=tmpdir, + orbit_type="singleton", + halve_boundary_couplers=False, + ) + self.assertEqual(tri.num_spins, 9) + self.assertGreater(tri.num_edges, 0) + self.assertEqual(tri.geometry_name, "Triangular") + + def test_coordinates(self): + with tempfile.TemporaryDirectory() as tmpdir: + tri = _make_triangular(tmpdir, 3, 3) + y, x = tri.coordinates(0) + self.assertEqual(y, 0) + self.assertEqual(x, 0) + y, x = tri.coordinates(4) + self.assertEqual(y, 1) + self.assertEqual(x, 1) + + def test_halve_boundary_couplers(self): + with tempfile.TemporaryDirectory() as tmpdir: + tri = _make_triangular( + tmpdir, 3, 3, periodic=(False, False), halve_boundary_couplers=True + ) + bqm = tri.make_bqm() + graph = tri.make_networkx_graph() + for u, v in tri.edge_list: + expected = 1.0 if (graph.degree[u] == 6 or graph.degree[v] == 6) else 0.5 + self.assertAlmostEqual(bqm.quadratic[(u, v)], expected) + + def test_periodicity(self): + with tempfile.TemporaryDirectory() as tmpdir: + tri = _make_triangular(tmpdir, 3, 3, periodic=(False, True)) + self.assertFalse(tri.periodic[0]) + self.assertTrue(tri.periodic[1]) + + +class TestDimerizedTriangular(unittest.TestCase): + def test_basic_construction(self): + with tempfile.TemporaryDirectory() as tmpdir: + dt = lattice.DimerizedTriangular( + dimensions=(3, 3), periodic=(True, False), orbit_type="singleton", data_root=tmpdir + ) + self.assertEqual(dt.geometry_name, "DimerizedTriangular") + self.assertIsNotNone(dt.logical_lattice) + self.assertEqual(dt.num_spins, 18) + + def test_chain_connectivity_self(self): + with tempfile.TemporaryDirectory() as tmpdir: + dt = lattice.DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + cc = dt.get_chain_connectivity(0) + self.assertEqual(cc, ((0, 1),)) + + def test_chain_connectivity_cases(self): + with tempfile.TemporaryDirectory() as tmpdir: + dt = lattice.DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + cases = [ + ((0,), ((0, 1),)), + ((0, 1), ((1, 0),)), + ((0, 3), ((1, 0),)), + ] + for args, expected in cases: + with self.subTest(args=args): + self.assertEqual(dt.get_chain_connectivity(*args), expected) + + +class TestEmbeddedLattice(unittest.TestCase): + def test_embed_sample(self): + with tempfile.TemporaryDirectory() as tmpdir: + dt = lattice.DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + logical_sample = np.array([1, -1, 1, -1, 1, -1, 1, -1, 1]) + embedded = dt.embed_sample(logical_sample) + self.assertEqual(len(embedded), dt.num_spins) + # Each chain should have the same value + for spin, chain in dt.chain_nodes.items(): + for node in chain: + self.assertEqual(embedded[node], logical_sample[spin]) + + def test_unembed_sample(self): + chain_nodes = {0: (0, 1, 2), 1: (3, 4, 5)} + with tempfile.TemporaryDirectory() as tmpdir: + embedded = lattice.EmbeddedLattice( + logical_lattice=lattice.Chain( + dimensions=(len(chain_nodes),), + periodic=(False,), + data_root=tmpdir, + ), + chain_nodes=chain_nodes, + dimensions=(sum(len(chain) for chain in chain_nodes.values()),), + periodic=(False,), + ) + physical_sample = np.array([1, 1, -1, -1, -1, 1]) + logical = embedded.unembed_sample(physical_sample) + np.testing.assert_array_equal(logical, np.array([1, -1])) + + def test_unembed_sample_breaks_ties_randomly(self): + with tempfile.TemporaryDirectory() as tmpdir: + embedded = _make_embedded_chain({0: (0, 1), 1: (2, 3)}, tmpdir) + physical_sample = np.array([1, -1, 1, -1]) + with mock.patch( + "dwave.experimental.lattice_utils.lattice.embedded_lattice.np.random.rand", + side_effect=[0.9, 0.1], + ): + logical = embedded.unembed_sample(physical_sample) + np.testing.assert_array_equal(logical, np.array([1, -1])) + + def test_unembed_sampleset(self): + with tempfile.TemporaryDirectory() as tmpdir: + embedded = _make_embedded_chain({0: (0, 1), 1: (2, 3)}, tmpdir) + samples = np.array( + [ + [1, 1, -1, -1], + [1, -1, 1, -1], + ] + ) + ss = dimod.SampleSet.from_samples(samples, vartype=dimod.SPIN, energy=0) + with mock.patch( + "dwave.experimental.lattice_utils.lattice.embedded_lattice.np.random.rand", + return_value=np.array([[0.2, 0.9], [0.2, 0.1]]), + ): + result = embedded.unembed_sampleset(ss) + np.testing.assert_array_equal(dimod.as_samples(result)[0], np.array([[1, -1], [1, -1]])) + + def test_connectivity_generic_self(self): + chain_nodes = {0: (10, 11, 12), 1: (20, 21, 22)} + with tempfile.TemporaryDirectory() as tmpdir: + el = _make_embedded_chain(chain_nodes, tmpdir) + cc = lattice.EmbeddedLattice.get_chain_connectivity(el, 0) + self.assertEqual(cc, ((0, 1), (0, 2), (1, 2))) + self.assertEqual( + {tuple(chain_nodes[0][index] for index in edge) for edge in cc}, + {(10, 11), (10, 12), (11, 12)}, + ) + + def test_nested_embedded_raises(self): + with tempfile.TemporaryDirectory() as tmpdir: + base_chain = lattice.Chain( + dimensions=(2,), + periodic=(False,), + data_root=tmpdir, + ) + embedded_once = lattice.EmbeddedLattice( + logical_lattice=base_chain, + chain_nodes={0: (0, 1), 1: (2, 3)}, + dimensions=(4,), + periodic=(False,), + ) + with self.assertRaises(NotImplementedError): + lattice.EmbeddedLattice( + logical_lattice=embedded_once, + chain_nodes={0: (0, 1), 1: (2, 3), 2: (4, 5), 3: (6, 7)}, + dimensions=(8,), + periodic=(False,), + ) + + +class TestOrbits(unittest.TestCase): + def test_reindex_basic(self): + mapping = {"a": 5, "b": 5, "c": 10} + result = lattice.reindex(mapping) + self.assertEqual(result, {'a': 0, 'b': 0, 'c': 1}) + + def test_signed_bqm_symmetry(self): + bqm = dimod.BQM(vartype="SPIN") + bqm.add_variable(0, 0.5) + bqm.add_variable(1, -0.3) + bqm.add_quadratic(0, 1, 1.0) + signed = lattice.make_signed_bqm(bqm) + self.assertAlmostEqual(signed.linear["p0"], 0.5) + self.assertAlmostEqual(signed.linear["m0"], -0.5) + + +class TestOptimizeFunction(unittest.TestCase): + """Tests for the free function ``lattice.optimize`` (vs. the ``Lattice.optimize`` + instance method covered by ``TestLatticeOptimize``).""" + + def test_plain_lattice(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + energy, sample, _ = lattice.optimize( + chain, bqm, sampler_kwargs={"num_sweeps": 256, "num_reads": 16} + ) + self.assertEqual(energy, -3) + self.assertEqual(bqm.energy(sample), energy) + + def test_embedded_lattice(self): + with tempfile.TemporaryDirectory() as tmpdir: + dt = lattice.DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + bqm = dt.make_bqm() + energy, sample, _ = lattice.optimize( + dt, bqm, sampler_kwargs={"num_sweeps": 256, "num_reads": 16} + ) + self.assertEqual(len(sample), dt.num_spins) + self.assertEqual(bqm.energy(sample), energy) + self.assertTrue(set(sample).issubset({-1, 1})) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_lattice_utils/test_observable.py b/tests/test_lattice_utils/test_observable.py new file mode 100644 index 0000000..2aeea18 --- /dev/null +++ b/tests/test_lattice_utils/test_observable.py @@ -0,0 +1,188 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import tempfile +import unittest +from pathlib import Path + +import dimod +import numpy as np + +from dwave.experimental.lattice_utils import lattice, observable +from tests.test_lattice_utils._helpers import _make_mock_experiment, _make_triangular + + +class TestQubitMagnetization(unittest.TestCase): + def test_qubit_magnetization(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + samples = np.array([[1, 1, -1, -1], [-1, -1, 1, 1]]) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(chain) + result = observable.QubitMagnetization().evaluate(exp, bqm, ss) + np.testing.assert_array_equal(result, [0.0, 0.0, 0.0, 0.0]) + + +class TestCouplerCorrelation(unittest.TestCase): + def test_coupler_correlation(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + exp = _make_mock_experiment(chain) + alt = np.tile([1, -1, 1, -1], (4, 1)) + ss_alt = dimod.SampleSet.from_samples_bqm(alt, bqm) + np.testing.assert_array_equal( + observable.CouplerCorrelation().evaluate(exp, bqm, ss_alt), + -np.ones(chain.num_edges), + ) + + +class TestCouplerFrustration(unittest.TestCase): + def test_coupler_frustration(self): + """All aligned (corr=1) -> frustration = 1.0""" + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + samples = np.ones((4, 4)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(chain) + np.testing.assert_array_almost_equal( + observable.CouplerFrustration().evaluate(exp, bqm, ss), np.ones(chain.num_edges) + ) + + +class TestSampleEnergy(unittest.TestCase): + def test_sample_energy(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + # All-ones: energy = sum of J for 3 edges = 3 + samples = np.ones((1, 4)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp_pos = _make_mock_experiment(chain, signed_energy_scale=1.0) + np.testing.assert_array_almost_equal( + observable.SampleEnergy().evaluate(exp_pos, bqm, ss), [3] + ) + + +class TestBitpackedSpins(unittest.TestCase): + def test_bitpacked_spins(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + samples = np.array([[1, -1, 1, -1], [-1, 1, -1, 1]]) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(chain) + packed, shape = observable.BitpackedSpins().evaluate(exp, bqm, ss) + self.assertEqual(shape, (2, 4)) + # Unpack and verify round-trip + unpacked = np.unpackbits(packed)[: shape[0] * shape[1]].reshape(shape) + np.testing.assert_array_equal(unpacked, np.equal(samples, 1)) + + +class TestReferenceEnergy(unittest.TestCase): + def test_reference_energy_save_load_roundtrip(self): + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "ref.txt" + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + sample = np.array([1, -1, 1, -1]) + obs = observable.ReferenceEnergy() + obs.save(path, -3, sample, "SA") + + exp = _make_mock_experiment(chain) + energy, loaded_sample, method = obs.load(exp, bqm, path) + self.assertEqual(energy, -3) + self.assertEqual(method, "SA") + np.testing.assert_array_equal(loaded_sample, sample) + + def test_reference_energy_evaluate_generates_and_caches(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + obs = observable.ReferenceEnergy() + + path1 = Path(tmpdir) / "ref_inst.txt" + energy1 = obs.evaluate(None, bqm, None, path=path1, inst=chain) + self.assertTrue(path1.exists()) + # Second call loads from cache — same value + energy1b = obs.evaluate(None, bqm, None, path=path1) + self.assertAlmostEqual(energy1, energy1b) + + exp = _make_mock_experiment(chain, run_index=0, num_random_instances=1) + path2 = Path(tmpdir) / "ref_exp.txt" + energy2 = obs.evaluate(exp, bqm, None, path=path2) + self.assertTrue(path2.exists()) + self.assertAlmostEqual(energy1, energy2) + + def test_reference_energy_update(self): + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "ref.txt" + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + obs = observable.ReferenceEnergy() + exp = _make_mock_experiment(chain) + + bad_sample = np.ones(4) + obs.save(path, bqm.energy(bad_sample), bad_sample, "SA") + + better = np.array([1, -1, 1, -1]) + obs.update(exp, bqm, better, path=path) + energy, _, _ = obs.load(exp, bqm, path) + self.assertAlmostEqual(energy, bqm.energy(better)) + + # Attempting to update with a worse sample raises ValueError + with self.assertRaises(ValueError): + obs.update(exp, bqm, bad_sample, path=path) + + +class TestKinks(unittest.TestCase): + def test_all_aligned(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + bqm = chain.make_bqm() + samples = np.ones((10, 6)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(chain) + result = observable.KinkKinkCorrelator().evaluate(exp, bqm, ss) + np.testing.assert_array_equal(result, np.zeros(6)) + + def test_mixed_pattern(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + bqm = chain.make_bqm() + # [1,1,-1,-1,1,1]: kink at sites 2,4 (domain walls) + samples = np.tile([1, 1, -1, -1, 1, 1], (20, 1)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(chain) + result = observable.KinkKinkCorrelator().evaluate(exp, bqm, ss) + expected = np.array([0.0, -0.25, 0.125, -0.25, 0.125, -0.25]) + np.testing.assert_array_almost_equal(result, expected) + + +class TestTriangularOP(unittest.TestCase): + def test_uniform_state_vanishes(self): + with tempfile.TemporaryDirectory() as tmpdir: + tri = _make_triangular(tmpdir, 3, 3, periodic=(True, False)) + bqm = tri.make_bqm() + samples = np.ones((5, 9)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(tri) + result = observable.TriangularOP().evaluate(exp, bqm, ss) + np.testing.assert_array_almost_equal(np.abs(result), np.zeros(5), decimal=10) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_lattice_utils/test_utils.py b/tests/test_lattice_utils/test_utils.py new file mode 100644 index 0000000..1b821ad --- /dev/null +++ b/tests/test_lattice_utils/test_utils.py @@ -0,0 +1,61 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np + +from dwave.experimental.lattice_utils.utils import ( + bootstrap, + confidence_interval, + generate_bootstrap_indices, +) + + +class TestUtils(unittest.TestCase): + def test_bootstrap_all_nan_skipnan(self): + rng = np.random.default_rng(seed=0) + result = bootstrap(np.array([np.nan, np.nan]), rng, repetitions=5, skipnan=True) + self.assertEqual(len(result), 5) + for val in result: + self.assertTrue(np.isnan(val)) + + def test_bootstrap_skipnan_false(self): + rng = np.random.default_rng(seed=0) + result = bootstrap(np.array([1.0, 2.0, np.nan]), rng, repetitions=5, skipnan=False) + self.assertEqual(len(result), 5) + + def test_bootstrap_custom_function(self): + rng = np.random.default_rng(seed=0) + result = bootstrap(np.arange(20), rng, repetitions=10, bootstrap_function=np.mean) + self.assertEqual(len(result), 10) + + def test_generate_bootstrap_indices_correct_count(self): + rng = np.random.default_rng(seed=0) + indices = list(generate_bootstrap_indices(10, 5, rng)) + self.assertEqual(len(indices), 5) + for idx in indices: + self.assertEqual(len(idx), 10) + self.assertTrue(np.all(idx >= 0)) + self.assertTrue(np.all(idx < 10)) + + def test_confidence_interval_width(self): + arr = np.arange(1000) + _, low1, high1 = confidence_interval(arr, width=0.5) + _, low2, high2 = confidence_interval(arr, width=0.99) + self.assertGreater(low2 + high2, low1 + high1) + + +if __name__ == "__main__": + unittest.main()