diff --git a/dwave/experimental/lattice_utils/__init__.py b/dwave/experimental/lattice_utils/__init__.py new file mode 100644 index 0000000..d8cb8ad --- /dev/null +++ b/dwave/experimental/lattice_utils/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2025 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 . 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..5f75b63 --- /dev/null +++ b/dwave/experimental/lattice_utils/experiment/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2025 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.samplercall import * +from dwave.experimental.lattice_utils.experiment.fast_anneal_experiment 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..331e81e --- /dev/null +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -0,0 +1,711 @@ +# Copyright 2025 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 lzma +import os +import pickle +import time +from pathlib import Path +from datetime import datetime +from typing import Any +from dataclasses import dataclass + +import dimod +import numpy as np + +from dwave.experimental.lattice_utils.lattice import Lattice +from dwave.experimental.lattice_utils.observable import ( + QubitMagnetization, + CouplerCorrelation, + CouplerFrustration, + SampleEnergy, + BitpackedSpins, + ReferenceEnergy, +) +from dwave.experimental.lattice_utils.experiment.samplercall import SamplerCall + +__all__ = ['Experiment', 'ExperimentConfig'] + +@dataclass +class ExperimentConfig: + """Container for the parameters that define an experiment.""" + 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 + +@dataclass +class ExperimentConfig: + """Container for the parameters that define an experiment.""" + + 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, + max_iterations: int | None = None, + config: ExperimentConfig, + ): + self.inst = inst + self.sampler = sampler + self.param = dict(vars(config)) + 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, + quiet: bool = True, + 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. + quiet: If false, prints a message when each result file is loaded. + ignore_shim: If true, the ``shimdata`` 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 not quiet: + print(f"Loaded {filename} at {datetime.now()}") + if result_fields is None: + result_fields = list(data.keys()) + if ignore_shim: + result_fields.remove("shimdata") + + 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] + for param_name, param_val in param.items(): + self.param[param_name] = param_val + + 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, **kwargs) -> 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. + + Returns: + A boolean value corresponding to whether or not the experiment is + finished. + """ + try: + self.inst._load_embeddings(self.sampler, **kwargs) + except FileNotFoundError as e: + raise FileNotFoundError("No Embedding Found: ", e) from e + + print( + f'\n{type(self.inst).__name__}={self.inst.dimensions}, J={self.param["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 = {} + + 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 * kwargs.get("scaling_factor", 1.0), + **call_dict[index].sampler_params, + ) + + if len(call_dict) == 0: + print(f"***\n***\nFINISHED for all {len(parameter_list)} parameterizations.\n***\n***") + return True + + # 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, quiet=True) + del response_dict[index] + made_progress = True + break + + if not made_progress: + time.sleep(0.1) # Waiting for results to come in + + 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 nominal 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.nominal_bqms[iemb] + ) + + results = {} + for observable in set(self.observables_to_collect): + results[observable.name] = [] + for iemb, sample_array in enumerate(sample_arrays): + bqm = call.nominal_bqms[iemb] + obs_result = observable.evaluate(self, bqm, sample_set[iemb]) + results[observable.name].append(obs_result) + + if type(results[observable.name][0]) == np.ndarray: + results[observable.name] = np.asarray(results[observable.name]) + + return results + + def _save_results( + self, + data_dict: dict[str, Any], + run_index: int | None = None, + quiet: bool = False, + 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 + + # 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) + + if not quiet: + print(f"Saved {filename} at {datetime.now()}") + + def _get_sorted_results_file_list(self) -> list[str]: + """Return result filenames sorted lexicographically.""" + fnlist = list(self.data_path.glob("iter*.pkl.lzma")) + fnlist.sort() + return [str(fn) for fn in fnlist] + + 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). + """ + energy_scale = self.param["energy_scale"] + + if "anneal_time" in self.param: + pathstring = f'energyscale{energy_scale:0.3}/atime{self.param["anneal_time"]:010.6f}us' + elif "anneal_schedule" in self.param: + pathstring = f'energyscale{energy_scale:0.3}/asched{self.param["anneal_schedule"]}' + else: + raise ValueError + + # 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_nominal_bqms(self) -> list[dimod.BQM]: + """Make nominal BQMs (one per embedding) for the experiment.""" + nominal_bqm = self.inst.make_nominal_bqm() + + if not hasattr(self.inst, "embedding_list"): + return [nominal_bqm] + + return [nominal_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.nominal_bqms = self._make_nominal_bqms() + sampler_call.shimdata = self._get_shimdata() + + # Here we can find out that we're finished. + if ( + self.max_iterations is not None + and sampler_call.shimdata["total_iterations"] >= self.max_iterations + ): + return None + + sampler_call.bqm = self._make_bqm(sampler_call) + sampler_call.sampler_params = self._make_sampler_params(shimdata=sampler_call.shimdata) + + 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 = parameter_list.copy() + 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"]] + + # We want the elements to be unique, of course. + ret_unique = [] + for entry in ret: + 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 type(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["shimdata"] = {} + for key in sampler_call.shimdata: + if type(sampler_call.shimdata[key]) == np.ndarray: + savedata["shimdata"][key] = sampler_call.shimdata[key].astype(np.float32) + elif type(sampler_call.shimdata[key]) == int: + savedata["shimdata"][key] = sampler_call.shimdata[key] + else: + savedata["shimdata"][key] = sampler_call.shimdata[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 "shimdata" in kwargs: + if "flux_biases" in kwargs["shimdata"]: + ret["flux_biases"] = list(kwargs["shimdata"]["flux_biases"]) + if "anneal_offsets" in kwargs["shimdata"]: + ret["anneal_offsets"] = list(kwargs["shimdata"]["anneal_offsets"]) + + if self.param.get("fast_anneal", False): + ret["fast_anneal"] = True + + ret["annealing_time"] = self.param["anneal_time"] + + return ret + + def _get_shimdata(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.""" + shimdata = {"total_iterations": 0} + if hasattr(self.inst, "embedding_list"): + num_embeddings = len(self.inst.embedding_list) + shimdata["flux_biases"] = np.zeros(self.sampler.properties["num_qubits"]) + shimdata["anneal_offsets"] = np.zeros(self.sampler.properties["num_qubits"]) + shimdata["relative_coupler_strength"] = np.ones((num_embeddings, self.inst.num_edges)) + + if self.param.get("flux_biases", None) is not None: + shimdata["flux_biases"] = self.param.get("flux_biases") + + return shimdata + + 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): + """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) + shimdata = data["shimdata"] + return shimdata + 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]): + """Update shim parameters according to shim data and parameters.""" + if "flux_biases" in sampler_call.shimdata 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.shimdata + and self.param.get("coupler_shim_step", 0) != 0 + ): + self._update_coupler_shim(sampler_call, results) + + sampler_call.shimdata["total_iterations"] += 1 + + def _update_flux_bias_shim(self, sampler_call: SamplerCall, results: dict[str, Any]): + """Update flux-bias shim values based on qubit magnetization.""" + target_magnetization = self.param["target_magnetization"] + qubit_magnetization = results["QubitMagnetization"] + flux_biases = sampler_call.shimdata["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 + energy_scale = self.param["energy_scale"] + relative_coupler_strength = sampler_call.shimdata["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.nominal_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] + nominal_values = np.array([bqm.quadratic[edge] for edge in self.inst.edge_list]) + coupler_signs = np.sign(nominal_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 nominal values. + # Let's try being more explicit. + for iemb in range(len(relative_coupler_strength)): + violators = ( + relative_coupler_strength[iemb, bin_edges] + * nominal_values[bin_edges] + * energy_scale + > 1 + ) + relative_coupler_strength[iemb, bin_edges[violators]] = ( + 0.99999 / nominal_values[bin_edges[violators]] / energy_scale + ) + + violators = ( + relative_coupler_strength[iemb, bin_edges] + * nominal_values[bin_edges] + * energy_scale + < -2 + ) + relative_coupler_strength[iemb, bin_edges[violators]] = ( + -1.99999 / nominal_values[bin_edges[violators]] / 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 nominal values. + # Let's try being more explicit. + for iemb in range(len(relative_coupler_strength)): + violators = ( + relative_coupler_strength[iemb, bin_edges] + * nominal_values[bin_edges] + * energy_scale + > 1 + ) + relative_coupler_strength[iemb, bin_edges[violators]] = ( + 0.99999 / nominal_values[bin_edges[violators]] / energy_scale + ) + + violators = ( + relative_coupler_strength[iemb, bin_edges] + * nominal_values[bin_edges] + * energy_scale + < -2 + ) + relative_coupler_strength[iemb, bin_edges[violators]] = ( + -1.99999 / nominal_values[bin_edges[violators]] / energy_scale + ) + + Q = nominal_values * relative_coupler_strength * 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.""" + energy_scale = self.param["energy_scale"] + bqm = dimod.BQM(vartype="SPIN") + if not hasattr(self.inst, "embedding_list"): + nominal_bqm = sampler_call.nominal_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 nominal_bqm.variables: + bqm.add_linear(v, nominal_bqm.linear[v]) + + for iedge, edge in enumerate(self.inst.edge_list): + bqm.add_quadratic(edge[0], edge[1], nominal_bqm.quadratic[*edge] * energy_scale) + + return bqm + + relative_coupler_strength = sampler_call.shimdata["relative_coupler_strength"] + for iemb, emb in enumerate(self.inst.embedding_list): + nominal_bqm = sampler_call.nominal_bqms[iemb].copy() + + for v in range(self.inst.num_spins): + # Don't touch degree-zero spins. Relevant to partial yield. + if nominal_bqm.degree(v) > 0: + bqm.add_linear(emb[v], nominal_bqm.linear[v]) + + for iedge, edge in enumerate(self.inst.edge_list): + bias = ( + nominal_bqm.quadratic[*edge] + * relative_coupler_strength[iemb, iedge] + * 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..59ca0b2 --- /dev/null +++ b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py @@ -0,0 +1,31 @@ +# Copyright 2025 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..cbd859b --- /dev/null +++ b/dwave/experimental/lattice_utils/experiment/samplercall.py @@ -0,0 +1,38 @@ +# Copyright 2025 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 + +import dimod + +__all__ = ['SamplerCall'] + + +@dataclass +class SamplerCall: + """Data class for managing asynchronous sampler calls.""" + + def __init__( + self, + run_index: int, + shimdata: dict | None = None, + bqm: dimod.BQM | None = None, + nominal_bqms: list | None = None, + sampler_params: dict | None = None, + ): + self.run_index: int = run_index + self.bqm: dimod.BQM | None = bqm + self.shimdata: dict = {} if shimdata is None else shimdata + self.nominal_bqms: list = [] if nominal_bqms is None else nominal_bqms + self.sampler_params: dict = {} if sampler_params is None else sampler_params diff --git a/dwave/experimental/lattice_utils/lattice/__init__.py b/dwave/experimental/lattice_utils/lattice/__init__.py new file mode 100644 index 0000000..22a583d --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/__init__.py @@ -0,0 +1,20 @@ +# Copyright 2025 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.lattice import * +from dwave.experimental.lattice_utils.lattice.chain import * +from dwave.experimental.lattice_utils.lattice.optimize import * +from dwave.experimental.lattice_utils.lattice.triangular import * +from dwave.experimental.lattice_utils.lattice.orbits import * +from dwave.experimental.lattice_utils.lattice.embedded_lattice 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..246cad0 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/chain.py @@ -0,0 +1,80 @@ +# Copyright 2025 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 Iterator +from pathlib import Path + +from numpy.typing import NDArray + +from dwave.experimental.lattice_utils.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, + ): + 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, + ) + + def generate_edges(self) -> Iterator[tuple[int, int]]: + """Yield edges for a 1D chain lattice. + + Returns: + An iterator 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 > 1: + 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..a33b0ba --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -0,0 +1,249 @@ +# Copyright 2025 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 itertools import combinations, product +from numbers import Integral +from collections.abc import Iterator, Hashable + +import dimod +import numpy as np +from numpy.typing import NDArray + +from dwave.experimental.lattice_utils.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]], + **kwargs, + ): + 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.") + self.chain_nodes: dict[tuple[int, Integral]] = chain_nodes + self.chain_coupling: float = -kwargs.pop("chain_strength", 2) + if not hasattr(self, "num_spins"): + self.num_spins = sum(len(c) for c in chain_nodes.values()) + kwargs.setdefault("periodic", self.logical_lattice.periodic) + super().__init__(**kwargs) + + 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 chainscorresponding 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) -> Iterator[tuple[Hashable, Hashable]]: + """Yield physical edges for the embedded lattice. + + Returns: + An iterator of tuples, where each tuple represents an edge between + two spins in the physical lattice. + """ + logical_bqm = self.logical_lattice.make_nominal_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_nominal_bqm(self, **kwargs) -> dimod.BQM: + """Construct and embed the nominal BQM. + + Args: + kwargs: Keyword arguments to pass to the logical lattice's + `make_nominal_bqm` method. + + Returns: + A dimod.BQM representing the embedded nominal 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_nominal_bqm(**kwargs)) + + def embed_bqm(self, logical_bqm: dimod.BQM) -> dimod.BQM: + """Embed a logical BQM onto the physical lattice. + + 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..a724923 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -0,0 +1,298 @@ +# Copyright 2025 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 +from pathlib import Path +from collections.abc import Iterator, Hashable +from abc import ABC, abstractmethod +import warnings + +import dimod +from minorminer.utils.parallel_embeddings import find_multiple_embeddings +import networkx as nx +import numpy as np +from numpy.typing import NDArray + +from dwave.experimental.lattice_utils.lattice.orbits import get_orbits +from dwave.experimental.lattice_utils.lattice.optimize import optimize + +__all__ = ['Lattice'] + + +class Lattice(ABC): + """An abstract base class for representing lattice geometries used in lattice-utils experiments. + + Subclasses are resonsible 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, + ): + self.dimensions = dimensions + self.data_root = Path(data_root) + + self.periodic = periodic if periodic is not None else tuple(False for _ in dimensions) + self.edge_list: list[tuple[Hashable, Hashable]] = list(self.generate_edges()) + + if not hasattr(self, "num_spins"): + raise AttributeError(f"{type(self).__name__} subclass must initialize self.num_spins") + + 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) -> Iterator[tuple[Hashable, Hashable]]: + """Yield the edges for this lattice.""" + raise NotImplementedError + + 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 = [], + **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. + """ + graph_bqm = dimod.to_networkx_graph(self.make_nominal_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 embedding file found.") + + 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("No embeddings found") + + 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_nominal_bqm(self) -> dimod.BQM: + """Construct a default nominal BQM coupling strength values set to +1. + + Args: + **kwargs: additional keyword arguments forwarded to subclass implementations. + Subclasses may use these to modify the construction of the nominal BQM. + + 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_nominal_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 not None and coupler_orbits is not None: + 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 _get_path( + self, + kind: str, + sampler_name: str | None = None, + extra_subdir: str | Path | None = None, + ) -> Path: + """Construct a standarized 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: str) -> 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)) + + def _make_networkx_graph(self) -> nx.Graph: + """Construct a NetworkX graph reprensetation of the lattice.""" + 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, **kwargs) -> tuple[float, NDArray, str]: + return optimize(lattice=self, bqm=bqm, **kwargs) diff --git a/dwave/experimental/lattice_utils/lattice/optimize.py b/dwave/experimental/lattice_utils/lattice/optimize.py new file mode 100644 index 0000000..1a91f38 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/optimize.py @@ -0,0 +1,115 @@ +# Copyright 2025 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 + +from dwave.samplers import SimulatedAnnealingSampler +import numpy as np +import dimod +from numpy.typing import NDArray + +__all__ = ['optimize', 'optimize_increasing_sa_sweeps'] + + +def optimize( + lattice: Lattice, + bqm: dimod.BQM, + sa_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. + sa_kwargs: Optional keyword arguments to pass to the simulated annealing + sampler, such as ``num_reads`` and ``num_sweeps``. + + Returns: + A tuple containing the best energy found, the corresponding sample as a + NumPy array, and a string indicating the optimization method used. + """ + if sa_kwargs is None: + sa_kwargs = {} + + # 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), sa_kwargs=sa_kwargs + ) + embedded_sample = lattice.embed_sample(logical_sample) + embedded_energy = bqm.energy(embedded_sample) + + return optimize_increasing_sa_sweeps(bqm, embedded_energy, embedded_sample) + + # If no special case, just use SA. + return optimize_increasing_sa_sweeps(bqm, sa_kwargs=sa_kwargs) + + +def optimize_increasing_sa_sweeps( + bqm: dimod.BQM, + reference_energy: float = np.inf, + reference_sample: NDArray | None = None, + sa_kwargs: dict[str, Any] | None = None, +) -> tuple[float, NDArray, str]: + """Optimize a BQM with simulated annealing and increasing sweep counts. + + Args: + bqm: The binary quadratic model to optimize. + reference_energy: An initial energy to compare against. If the best energy + found by SA is not better than this, the function will return without + increasing the number of sweeps. + reference_sample: An initial sample corresponding to the reference energy. + sa_kwargs: Optional keyword arguments to pass to the simulated annealing + sampler, such as ``num_reads`` and ``num_sweeps``. The ``num_sweeps`` + value will be overridden by this function as it increases exponentially. + + Returns: + A tuple containing the best energy found, the corresponding sample as a + NumPy array, and a string indicating the optimization method used. + """ + sa = SimulatedAnnealingSampler() + + if sa_kwargs is None: + sa_kwargs = {} + num_sweeps = sa_kwargs.get("num_sweeps", 256) + num_reads = sa_kwargs.get("num_reads", 256) + + while True: + sample_set = sa.sample(bqm, num_reads=num_reads, num_sweeps=num_sweeps) + energies = sample_set.data_vectors["energy"] + best = np.argmin(energies) + best_energy = energies[best] + + if best_energy < reference_energy: + reference_energy = best_energy + reference_sample = sample_set.record[best][0] + num_sweeps *= 2 + if num_sweeps > 1e3: + break + else: + break + + return reference_energy, reference_sample, "sa_exponential" diff --git a/dwave/experimental/lattice_utils/lattice/orbits.py b/dwave/experimental/lattice_utils/lattice/orbits.py new file mode 100644 index 0000000..40852c3 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/orbits.py @@ -0,0 +1,240 @@ +# Copyright 2025 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 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 = {h: [] for h in set(bqm.linear.values())} + mapping_mp = {h: [] for h in set(bqm.linear.values())} + mapping_J = {J: [] for J in set(bqm.quadratic.values())} + + 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 + + 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..6ec42e9 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -0,0 +1,289 @@ +# Copyright 2025 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 Iterator, Hashable +from pathlib import Path + +import networkx as nx +import numpy as np +from numpy.typing import NDArray +import dimod + +from dwave.experimental.lattice_utils.lattice.lattice import Lattice +from dwave.experimental.lattice_utils.lattice.embedded_lattice import EmbeddedLattice + +__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: Path | None = None, + 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_nominal_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) -> Iterator[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: + An iterator 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: Path | None = None, + 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 chainscorresponding 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..cb3ef4f --- /dev/null +++ b/dwave/experimental/lattice_utils/observable/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2025 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.observable import * +from dwave.experimental.lattice_utils.observable.kinks 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..9f879b8 --- /dev/null +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -0,0 +1,60 @@ +# Copyright 2025 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 dimod +from dimod import BQM, SampleSet +import numpy as np +from numpy.typing import NDArray + +from dwave.experimental.lattice_utils.observable.observable import Observable + +__all__ = ['KinkKinkCorrelator'] + + +class KinkKinkCorrelator(Observable): + """For 1D chains.""" + + def __init__(self): + super().__init__() + + 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["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..9e8c4b9 --- /dev/null +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -0,0 +1,285 @@ +# Copyright 2025 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 + +import numpy as np +from numpy.typing import NDArray +import dimod + +from dwave.experimental.lattice_utils.lattice import Lattice + +__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: + raise NotImplementedError + + +class QubitMagnetization(Observable): + """Compute the mean magnetization of each qubit.""" + + def evaluate( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample_set: dimod.SampleSet, + ) -> NDArray: + 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: + sample_array = dimod.as_samples(sample_set)[0].astype(float) + if len(experiment.inst.edge_list) == 0: + 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: + sample_array = dimod.as_samples(sample_set)[0].astype(float) + if len(experiment.inst.edge_list) == 0: + 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["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 ``energy_scale`` but include its sign. + """ + + def evaluate( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample_set: dimod.SampleSet, + ) -> NDArray: + return sample_set.data_vectors["energy"] * np.sign(experiment.param["energy_scale"]) + + +class BitpackedSpins(Observable): + """Return bitpacked spins and a tuple of the array size.""" + + def evaluate( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample_set: dimod.SampleSet, + ) -> tuple[NDArray, tuple[int, int]]: + 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.""" + if path is not None: + path = Path(path) + else: + path = get_reference_energy_path(experiment, bqm=bqm) + + 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.""" + if path is not None: + path = Path(path) + else: + path = get_reference_energy_path(experiment, bqm=bqm) + + 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.""" + 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, + 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. + """ + 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: + if path is None: + path = get_reference_energy_path(experiment, bqm=bqm) + self.save(path, new_energy, sample, reference_method_string) + print(f"Updated energy from {reference_energy} to {new_energy}.") + else: + raise ValueError + + +def get_reference_energy_path( + experiment: Experiment | None = None, + root: str | Path | None = None, + bqm: dimod.BQM | 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: + experiment: The experiment for which to get the reference energy path. + root: Optional root directory to use instead of the experiment's data root. + bqm: The BQM for which to get the reference energy path. + + Returns: + The path to the reference energy file. + """ + if bqm is None: + raise NotImplementedError("Must provide a BQM to get the reference energy path.") + + # Allow for generation of dummy experiment data without all the overhead, + # for running without an actual experiment. + if experiment is None: + 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..b6faac9 --- /dev/null +++ b/dwave/experimental/lattice_utils/observable/triangular.py @@ -0,0 +1,77 @@ +# Copyright 2025 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 numpy as np +from numpy.typing import NDArray +import dimod +from dimod import BQM + +from dwave.experimental.lattice_utils.observable.observable import Observable + +__all__ = ['TriangularOP'] + + +class TriangularOP(Observable): + """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..94341ab --- /dev/null +++ b/dwave/experimental/lattice_utils/utils.py @@ -0,0 +1,113 @@ +# Copyright 2025 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/example_1D_Ising_chain.py b/examples/example_1D_Ising_chain.py new file mode 100644 index 0000000..f9ea39f --- /dev/null +++ b/examples/example_1D_Ising_chain.py @@ -0,0 +1,343 @@ +# Copyright 2025 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. + +"""Example for 1D Ising chain.""" + +from pathlib import Path +import os + +from dwave.system import DWaveSampler +import matplotlib.pyplot as plt +from matplotlib.colors import to_rgb +import numpy as np + +from dwave.experimental.lattice_utils import lattice, experiment, observable +from dwave.experimental.lattice_utils.utils import bootstrap, confidence_interval + +rng = np.random.default_rng(seed=0) +# Set up a dict for collating kink densities +kd_dict = {} +kkc_dict = {} + +# Set up the parameters + +# Two samplers: an Advantage2 prototype and an Advantage system. +samplers = [ + DWaveSampler(solver="Advantage2_system1"), + DWaveSampler(solver="Advantage_system4.1"), +] + +NUM_SPINS = 256 + +# Two energy scales: one strong coupling and one weak coupling. +ENERGY_SCALES = (-1.8, 0.1) + +# Minimum anneal time is 5ns. 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.005, 50, 21), 6) + +errorbar_style = {"marker": '', "linestyle": '', "capsize": 2} +point_style = {"marker": 'o', "linestyle": ''} +# Create a folder to save figures in if it doesn't already exist +Path("figures").mkdir(exist_ok=True) + +data_root = Path(__file__).resolve().parents[1] + +for sampler in samplers: + + # 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. We will also set the orbit_type to 'standard', + # which will allow 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. + + # Here we will do some shimming: flux bias shim and coupler shim. We will + # run two energy scales: a very strong one (negative, ferromagnetic) and a + # very weak one (positive, antiferromagnetic). Positive and negative energy + # scales are equivalent by gauge transformation, but we run the strong coupling + # on the FM side because the maximum FM magnitude (-2) is larger than the + # maximum AFM magnitude (+1). + for energy_scale in ENERGY_SCALES: + config = experiment.FastAnnealExperimentConfig( + energy_scale=energy_scale, + coupler_shim_step=0.05, + flux_bias_shim_step=1e-6, + ) + exp = experiment.Experiment(inst=inst, sampler=sampler, max_iterations=5, config=config) + + # Every experiment has an attribute (a set) of observables to compute and + # save while the experiment runs.Here we can add non-default observables. + # In this case we will add the kink-kink correlator. The observable + # object is designed to provide a standard interface for adding whatever + # experiment-specific observables you might require. + exp.observables_to_collect.add(observable.KinkKinkCorrelator()) + + # Make parameter list. We will only vary anneal time. + parameter_list = [{"anneal_time": time} for time in ANNEAL_TIMES] + + for _ in range(20): #TODO clean this up? + done = exp.run_iteration(parameter_list) + if done: + break + + # Now we will run some analysis. Let's first just plot kink density as + # a function of annealing time. Kink density is the same as the average + # "FrustrationProbability" observable for a given coupler, which is already + # gathered by default since it is required for the coupler shim. + + # We will make some lists for the data we want to analyze, and for each + # iteration of the experiment we will load the results and append the + # observable to the list. + + frust = [] # average coupler frustration (kink density) + cshim = [] # coupler shim + fbshim = [] # flux bias shim + kkc = [] # kink-kink correlator + for param in parameter_list: + exp.apply_param(param) #TODO clean this up? + res = exp.load_results(num_iterations=1000) + + frust.append(np.array([np.mean(it["CouplerFrustration"]) for it in res])) + cshim.append( + np.asarray([it["shimdata"]["relative_coupler_strength"].ravel() for it in res]) + ) + fbshim.append(np.asarray([it["shimdata"]["flux_biases"] for it in res])) + kkc.append( + np.reshape(np.asarray([it["KinkKinkCorrelator"] for it in res]), (-1, NUM_SPINS)) + ) + + title = f"1D chain, {'x'.join([str(dim) for dim in inst.dimensions])}, J={exp.param["energy_scale"]}, {sampler.solver.name}" + fig, axes = plt.subplots(3, 3, figsize=(16, 10)) + fig.suptitle(title, fontsize=16) + rng = np.random.default_rng(0) + x = np.linspace(0, 2 * np.pi, 400) + plt.tight_layout() + plt.subplots_adjust(hspace=0.35, wspace=0.3, top=0.9) + + ax = axes[0, 0] + ax.loglog() + x_theory = ANNEAL_TIMES[0] + y_theory = np.mean(frust[0]) + theoryfit = np.polyfit( + np.log([x_theory, x_theory * 2]), np.log([y_theory, y_theory * (2**-0.5)]), 1 + ) + ax.plot( + ANNEAL_TIMES, + np.exp(np.polyval(theoryfit, np.log(ANNEAL_TIMES))), + linestyle="-", + color=[0.8, 0.8, 0.8], + label="theory", + ) + + M = np.asarray(frust) + + bs = np.asarray([bootstrap(_, rng, bootstrap_function=np.nanmedian) for _ in M]) + ci = np.asarray([confidence_interval(_) for _ in bs]) + + errorbar_handle = ax.errorbar( + ANNEAL_TIMES, + ci[:, 0], + yerr=[ci[:, 1], ci[:, 2]], + **errorbar_style, + ) + ax.plot( + ANNEAL_TIMES, + ci[:, 0], + color=errorbar_handle[0]._color, + markerfacecolor=np.array(to_rgb(errorbar_handle[0]._color)) / 2 + 0.5, + **point_style, + ) + + ax.set_title("Kink density (with ~t_a^{-1/2} guideline)") + ax.set_ylabel("kink density") + ax.set_xlabel("$t_a$ (μs)") + ax.set_ylim([5e-4, 5e-1]) + ax.set_xlim([0.002, 9e1]) + ax.grid(which="both", alpha=0.3) + + ax = axes[0, 1] + ax.loglog() + y = np.sqrt(np.asarray([np.mean(_**2) for _ in fbshim])) + ax.plot(ANNEAL_TIMES, y, marker="o", linestyle="-") + ax.set_title("RMS flux bias shim") + ax.set_xlabel("$t_a$ (μs)") + ax.set_ylabel("RMS flux bias") + ax.grid(which="both", alpha=0.3) + + ax = axes[0, 2] + ax.loglog() + y = np.sqrt(np.asarray([np.mean((_ - 1) ** 2) for _ in cshim])) + ax.plot(ANNEAL_TIMES, y, marker="o", linestyle="-") + ax.set_title("RMS coupler shim") + ax.set_xlabel("$t_a$ (μs)") + ax.set_ylabel("RMS coupler shim") + ax.grid(which="both", alpha=0.3) + + ax = axes[1, 0] + ax.plot(fbshim[0]) + ax.set_title(f"Flux bias shim, t_a={ANNEAL_TIMES[0]:.3f}μs") + ax.set_xlabel("Iteration") + ax.grid(which="both", alpha=0.3) + + ax = axes[1, 1] + ax.plot(fbshim[1]) + ax.set_title(f"Flux bias shim, t_a={ANNEAL_TIMES[1]:.3f}μs") + ax.set_xlabel("Iteration") + ax.grid(which="both", alpha=0.3) + + ax = axes[1, 2] + ax.plot(fbshim[6]) + ax.set_title(f"Flux bias shim, t_a={ANNEAL_TIMES[-1]:.3f}μs") + ax.set_xlabel("Iteration") + ax.grid(which="both", alpha=0.3) + + ax = axes[2, 0] + ax.plot(cshim[0]) + ax.set_title(f"Coupler shim, t_a={ANNEAL_TIMES[0]:.3f}μs") + ax.set_xlabel("Iteration") + ax.grid(which="both", alpha=0.3) + + ax = axes[2, 1] + ax.plot(cshim[1]) + ax.set_title(f"Coupler shim, t_a={ANNEAL_TIMES[1]:.3f}μs") + ax.set_xlabel("Iteration") + ax.grid(which="both", alpha=0.3) + + ax = axes[2, 2] + ax.plot(cshim[6]) + ax.set_title(f"Coupler shim, t_a={ANNEAL_TIMES[-1]:.3f}μs") + ax.set_xlabel("Iteration") + ax.grid(which="both", alpha=0.3) + + filename = title + for bad_symbol in "/: ;,": + filename = filename.replace(bad_symbol, "_") + fig.savefig(Path(os.getcwd()) / 'figures' / f"{filename}.png") + plt.show() + + # Put kink density in a dict so we can plot them all together. + kd_dict[sampler.solver.name, energy_scale] = np.asarray(frust) + kkc_dict[sampler.solver.name, energy_scale] = np.asarray(kkc) + +# Now plot the kink densities together, for a nice comparison. +fig2, ax2 = plt.subplots(1, 2, figsize=(10, 8)) +title = f"1D chain kink density, {'x'.join([str(dim) for dim in inst.dimensions])}" +fig2.suptitle(title, fontsize=16) + +for isampler, sampler in enumerate(samplers): + + for energy_scale in [-1.8, 0.1]: + M = kd_dict[sampler.solver.name, energy_scale] + + # Kink density plot + theoryx = ANNEAL_TIMES[0] + theoryy = np.mean(M[0]) + theoryfit = np.polyfit( + np.log([theoryx, theoryx * 2]), np.log([theoryy, theoryy * (2**-0.5)]), 1 + ) + ax2[isampler].plot( + ANNEAL_TIMES, + np.exp(np.polyval(theoryfit, np.log(ANNEAL_TIMES))), + linestyle="-", + color=[0.8, 0.8, 0.8], + label="theory", + ) + ax2[isampler].set_title(f"Kink density: {sampler.solver.name}") + ax2[isampler].loglog() + ax2[isampler].grid(which="both", alpha=0.3) + ax2[isampler].set_ylabel("Kink density") + ax2[isampler].set_xlabel("$t_a$ (μs)") + ax2[isampler].set_ylim([5e-4, 5e-1]) + ax2[isampler].set_xlim([0.002, 9e1]) + + x = ANNEAL_TIMES + M = kd_dict[sampler.solver.name, energy_scale] + bs = np.asarray([bootstrap(m, rng, bootstrap_function=np.nanmedian) for m in M]) + ci = np.asarray([confidence_interval(b) for b in bs]) + + errorbar_handle = ax2[isampler].errorbar( + ANNEAL_TIMES, ci[:, 0], yerr=[ci[:, 1], ci[:, 2]], **errorbar_style + ) + ax2[isampler].plot( + x, + ci[:, 0], + color=errorbar_handle[0]._color, + markerfacecolor=np.array(to_rgb(errorbar_handle[0]._color)) / 2 + 0.5, + **point_style, + ) + +filename = title +for bad_symbol in "/: ;,": + filename = filename.replace(bad_symbol, "_") +fig2.savefig(Path(os.getcwd()) / 'figures' / f"{filename}.png") +plt.show() + +# Now we will analyze the kink-kink correlator for the fastest anneals (5ns) +fig3, ax3 = plt.subplots(1, 2, figsize=(10, 8)) +dims = 'x'.join(map(str, inst.dimensions)) +time_ns = ANNEAL_TIMES[0] * 1000 +title = f"1D chain kink-kink correlator, {dims}, {time_ns:.1f} ns" +fig3.suptitle(title, fontsize=16) + +for isampler, sampler in enumerate(samplers): + + for energy_scale in [-1.8, 0.1]: + magnetization = kkc_dict[sampler.solver.name, energy_scale][0].T + kd = np.mean(kd_dict[sampler.solver.name, energy_scale][0]) + + # Kink density plot + ax3[isampler].grid(which="both", alpha=0.3) + ax3[isampler].set_title(f"Kink-kink correlator: {sampler.solver.name}") + ax3[isampler].set_ylabel("Kink-kink correlator") + ax3[isampler].set_xlabel("Normalized distance") + ax3[isampler].set_ylim([-0.15, 0.15]) + ax3[isampler].set_xlim([0.01, 1.5]) + + x = np.arange(NUM_SPINS) * kd + M = magnetization + + bs = np.asarray([bootstrap(m, rng, bootstrap_function=np.nanmedian) for m in M]) + ci = np.asarray([confidence_interval(i) for i in bs]) + + errorbar_handle = ax3[isampler].errorbar( + x, ci[:, 0], yerr=[ci[:, 1], ci[:, 2]], **errorbar_style + ) + ax3[isampler].plot( + x, + ci[:, 0], + marker='o', + linestyle='', + color=errorbar_handle[0]._color, + markerfacecolor=np.array(to_rgb(errorbar_handle[0]._color)) / 2 + 0.5, + ) + +filename = title +for bad_symbol in "/: ;,": + filename = filename.replace(bad_symbol, "_") +fig3.savefig(Path(os.getcwd()) / 'figures' / f"{filename}.png") +plt.show() diff --git a/examples/example_2D_geometric_frustration.py b/examples/example_2D_geometric_frustration.py new file mode 100644 index 0000000..2a68546 --- /dev/null +++ b/examples/example_2D_geometric_frustration.py @@ -0,0 +1,293 @@ +# Copyright 2025 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. + +"""Now we're doing to do a similar example but on triangular lattices, which are +embedded using two qubits per chain.""" + +from pathlib import Path +import os + +from dwave.system import DWaveSampler +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.colors import to_rgb + +from dwave.experimental.lattice_utils import lattice, experiment, observable +from dwave.experimental.lattice_utils.utils import bootstrap, confidence_interval + +rng = np.random.default_rng(seed=0) + +# Set up a dict for collating statistics +m_dict = {} +psi_dict = {} + +sampler = DWaveSampler(solver="Advantage2_system1") + +ANNEAL_TIMES = np.round(0.005 * np.logspace(0, 2, 17), 6) + +errorbar_style = {"marker": '', "linestyle": '', "capsize": 2} +point_style = {"marker": 'o', "linestyle": ''} +# Create a folder to save figures in if it doesn't already exist +Path("figures").mkdir(exist_ok=True) + +data_root = Path(__file__).resolve().parents[1] + +inst = lattice.DimerizedTriangular( + dimensions=(9, 12), + data_root=data_root, + periodic=(True, False), + orbit_type="explicit", + halve_boundary_couplers=True, + chain_strength=2, +) +inst.embed_lattice( + sampler, + max_number_of_embeddings=1, + timeout=10000, + remove_external_edges=True, + remove_odd_edges=True, + draw_reduced_graph=True, +) +# Now must make the orbits: chain and no-chain. +coupler_orbit = np.array( + [inst.make_nominal_bqm().quadratic[edge] == -2 for edge in inst.edge_list], + dtype=int, +) +qubit_orbit = np.ones(inst.num_spins, dtype=int) +inst.initialize_orbits(qubit_orbits=qubit_orbit, coupler_orbits=coupler_orbit) + +exp = experiment.FastAnnealExperiment( + inst=inst, + sampler=sampler, + num_reads=100, + readout_thermalization=100, + max_iterations=210, + results_root=Path("./results"), + automorph_embeddings=False, + energy_scale=0.8, + coupler_shim_step=0.1, + flux_bias_shim_step=5e-6, +) + +exp.observables_to_collect = [ + observable.QubitMagnetization(), + observable.CouplerCorrelation(), + observable.CouplerFrustration(), + observable.SampleEnergy(), + observable.TriangularOP(), + observable.ReferenceEnergy(), + observable.BitpackedSpins(), +] + +# Make parameter list +parameter_list = [{"anneal_time": rate} for rate in ANNEAL_TIMES] + +for _ in range(1000): + done = exp.run_iteration(parameter_list) + if done: + break + +# We will make some lists for the data we want to analyze, and for each iteration +# of the experiment we will load theresults and append the observable to the list. +frust = [] # average coupler frustration (kink density) +cshim = [] # coupler shim +fbshim = [] # flux bias shim +opmag = [] +psi = [] +ene = [] + +for param in parameter_list: + exp.apply_param(param) + res = exp.load_results(num_iterations=1000) + frust.append(np.array([np.mean(i["CouplerFrustration"]) for i in res])) + cshim.append(np.asarray([i["shimdata"]["relative_coupler_strength"].ravel() for i in res])) + fbshim.append(np.asarray([i["shimdata"]["flux_biases"].ravel() for i in res])) + opmag.append(np.array([np.mean(np.abs(i["TriangularOP"])) for i in res])) + ene.append(np.array([np.mean(i["SampleEnergy"]) for i in res])) + psi.append(np.asarray([i["TriangularOP"] for i in res])) + +title = ( + f"DimerizedTriangular, {'x'.join([str(dim) for dim in inst.dimensions])}, " + f"J={exp.param["energy_scale"]}, {sampler.solver.name}" +) +fig, axes = plt.subplots(3, 3, figsize=(16, 10)) +fig.suptitle(title, fontsize=16) +rng = np.random.default_rng(0) +x = np.linspace(0, 2 * np.pi, 400) +plt.tight_layout() +plt.subplots_adjust(hspace=0.35, wspace=0.3, top=0.9, left=0.07, bottom=0.07) + +ax = axes[0, 0] +ax.loglog() + +M = np.asarray(opmag) +bs = np.asarray([bootstrap(m, rng, bootstrap_function=np.nanmedian) for m in M]) +ci = np.asarray([confidence_interval(i) for i in bs]) + +errorbar_handle = ax.errorbar(ANNEAL_TIMES, ci[:, 0], yerr=[ci[:, 1], ci[:, 2]], **errorbar_style) +ax.plot( + ANNEAL_TIMES, + ci[:, 0], + color=errorbar_handle[0]._color, + markerfacecolor=np.array(to_rgb(errorbar_handle[0]._color)) / 2 + 0.5, + **point_style, +) + +ax.set_title("") +ax.set_ylabel("") +ax.set_xlabel("$t_a$ (μs)") +ax.set_xlim([0.002, 9e-1]) +ax.grid(which="both", alpha=0.3) + +ax = axes[0, 1] +ax.loglog() +y = np.sqrt(np.asarray([np.mean(_**2) for _ in fbshim])) +ax.plot(ANNEAL_TIMES, y, marker="o", linestyle="-") +ax.set_title("RMS flux bias shim") +ax.set_xlabel("$t_a$ (μs)") +ax.set_ylabel("RMS flux bias") +ax.grid(which="both", alpha=0.3) + +ax = axes[0, 2] +ax.loglog() +y = np.sqrt(np.asarray([np.mean((_ - 1) ** 2) for _ in cshim])) +ax.plot(ANNEAL_TIMES, y, marker="o", linestyle="-") +ax.set_title("RMS coupler shim") +ax.set_xlabel("$t_a$ (μs)") +ax.set_ylabel("RMS coupler shim") +ax.grid(which="both", alpha=0.3) + +ax = axes[1, 0] +ax.plot(fbshim[0]) +ax.set_title(f"Flux bias shim, t_a={ANNEAL_TIMES[0]:.3f}μs") +ax.set_xlabel("Iteration") +ax.grid(which="both", alpha=0.3) + +ax = axes[1, 1] +ax.plot(fbshim[1]) +ax.set_title(f"Flux bias shim, t_a={ANNEAL_TIMES[1]:.3f}μs") +ax.set_xlabel("Iteration") +ax.grid(which="both", alpha=0.3) + +ax = axes[1, 2] +ax.plot(fbshim[6]) +ax.set_title(f"Flux bias shim, t_a={ANNEAL_TIMES[-1]:.3f}μs") +ax.set_xlabel("Iteration") +ax.grid(which="both", alpha=0.3) + +ax = axes[2, 0] +ax.plot(cshim[0]) +ax.set_title(f"Coupler shim, t_a={ANNEAL_TIMES[0]:.3f}μs") +ax.set_xlabel("Iteration") +ax.grid(which="both", alpha=0.3) + +ax = axes[2, 1] +ax.plot(cshim[1]) +ax.set_title(f"Coupler shim, t_a={ANNEAL_TIMES[1]:.3f}μs") +ax.set_xlabel("Iteration") +ax.grid(which="both", alpha=0.3) + +ax = axes[2, 2] +ax.plot(cshim[6]) +ax.set_title(f"Coupler shim, t_a={ANNEAL_TIMES[-1]:.3f}μs") +ax.set_xlabel("Iteration") +ax.grid(which="both", alpha=0.3) + +filename = title +for bad_symbol in "/: ;,": + filename = filename.replace(bad_symbol, "_") +fig.savefig(Path(os.getcwd()) / 'figures' / f"{filename}.png") +plt.show() + +# Put kink density in a dict so we can plot them all together. +m_dict[sampler.solver.name] = np.asarray(opmag) +psi_dict[sampler.solver.name] = np.asarray(psi) + +# Now plot the order parameters together, for a nice comparison. +fig2, ax2 = plt.subplots(2, 1, figsize=(8, 12)) +title = f'Triangular, global orbit, J={exp.param["energy_scale"]}' +fig2.suptitle(title, fontsize=16) + +M = m_dict[sampler.solver.name] + +bs = np.asarray([bootstrap(m, rng, bootstrap_function=np.nanmedian) for m in M[:, :5]]) +ci = np.asarray([confidence_interval(i) for i in bs]) +errorbar_handle = ax2[0].errorbar( + ANNEAL_TIMES, + ci[:, 0], + yerr=[ci[:, 1], ci[:, 2]], + **errorbar_style, +) + +facecolor = np.array(to_rgb(errorbar_handle[0]._color)) / 2 + 0.5 +ax2[0].plot( + ANNEAL_TIMES, + ci[:, 0], + color=errorbar_handle[0]._color, + markerfacecolor=facecolor, + label="first 5 iterations of shim", + **point_style, +) +bs = np.asarray([bootstrap(m, rng, bootstrap_function=np.nanmedian) for m in M[:, -5:]]) +ci = np.asarray([confidence_interval(i) for i in bs]) +errorbar_handle = ax2[0].errorbar( + ANNEAL_TIMES, + ci[:, 0], + yerr=[ci[:, 1], ci[:, 2]], + **errorbar_style, +) +facecolor = np.array(to_rgb(errorbar_handle[0]._color)) / 2 + 0.5 +ax2[0].plot( + ANNEAL_TIMES, + ci[:, 0], + color=errorbar_handle[0]._color, + markerfacecolor=facecolor, + label="last 5 iterations of shim", + **point_style, +) + +ax2[0].loglog() +ax2[0].grid(which="both", alpha=0.3) +ax2[0].set_title(f": {sampler.solver.name}") +ax2[0].set_ylabel("") +ax2[0].set_xlabel("$t_a$ (μs)") +ax2[0].set_ylim([0.15, 1.2]) +ax2[0].legend() + +# And heatmaps of psi. +index = len(M) - 1 +M = psi_dict[sampler.solver.name][index][-10:].ravel() + +x = np.real(M) +y = np.imag(M) + +NUM_BINS = 41 +extent = (-2, 2, -1.95, 1.95) + +hb = ax2[1].hexbin(x, y, gridsize=NUM_BINS, cmap="inferno", extent=extent) +ax2[1].set_title(f"ψ, t_a={ANNEAL_TIMES[index]:.3f}μs") +cb = fig2.colorbar(hb, ax=ax2[1]) +cb.set_label("count") +ax2[1].plot([-1 / np.sqrt(3), 1 / np.sqrt(3)], [-1, 1], color=(0, 0, 0, 0.1), linestyle="-") +ax2[1].plot([-1 / np.sqrt(3), 1 / np.sqrt(3)], [1, -1], color=(0, 0, 0, 0.1), linestyle="-") +ax2[1].plot([-2 / np.sqrt(3), 2 / np.sqrt(3)], [0, 0], color=(0, 0, 0, 0.1), linestyle="-") +ax2[1].axis([-1.4, 1.4, -1.4, 1.4]) +ax2[1].set_aspect("equal", "box") + +filename = title +for bad_symbol in "/: ;,": + filename = filename.replace(bad_symbol, "_") +fig2.savefig(Path(os.getcwd()) / 'figures' / f"{filename}.png") + +plt.show() diff --git a/tests/test_lattice_utils.py b/tests/test_lattice_utils.py new file mode 100644 index 0000000..96198e1 --- /dev/null +++ b/tests/test_lattice_utils.py @@ -0,0 +1,1007 @@ +# Copyright 2025 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.utils import ( + bootstrap, + confidence_interval, + generate_bootstrap_indices, +) +from dwave.experimental.lattice_utils.lattice.chain import Chain +from dwave.experimental.lattice_utils.lattice.triangular import DimerizedTriangular, Triangular +from dwave.experimental.lattice_utils.lattice.embedded_lattice import EmbeddedLattice +from dwave.experimental.lattice_utils.lattice.orbits import make_signed_bqm, reindex +from dwave.experimental.lattice_utils.lattice.optimize import optimize +from dwave.experimental.lattice_utils.observable.observable import ( + BitpackedSpins, + CouplerCorrelation, + CouplerFrustration, + QubitMagnetization, + ReferenceEnergy, + SampleEnergy, + get_reference_energy_path, +) +from dwave.experimental.lattice_utils.observable.kinks import KinkKinkCorrelator +from dwave.experimental.lattice_utils.observable.triangular import TriangularOP +from dwave.experimental.lattice_utils.experiment.samplercall import SamplerCall +from dwave.experimental.lattice_utils.experiment.experiment import Experiment +from dwave.experimental.lattice_utils.experiment.fast_anneal_experiment import FastAnnealExperiment + + +def _make_triangular( + ly=3, + lx=3, + periodic=(True, False), + orbit_type="singleton", + halve_boundary_couplers=False, +): + return Triangular( + dimensions=(ly, lx), + periodic=periodic, + orbit_type=orbit_type, + halve_boundary_couplers=halve_boundary_couplers, + ) + + +def _make_mock_sampler(num_qubits=128, nodelist=None, solver_name="TestSolver"): + """Create a minimal mock sampler resembling DWaveSampler.""" + sampler = mock.MagicMock(spec=dimod.Sampler) + type(sampler).__name__ = "DWaveSampler" + 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 + return sampler + + +def _make_sync_sampler(n_cols=128, solver_name="TestSolver"): + """Sampler whose sample() immediately returns all-ones raw data (done=True). + + Mimics the async response interface used by DWaveSampler: .done() and + .samples() -> 2-D ndarray of shape (num_reads, n_cols). + """ + + class _Response: + def done(self): + return True + + def samples(self, sorted_by=None): + return np.ones((10, n_cols), dtype=float) + + sampler = mock.MagicMock() + type(sampler).__name__ = "DWaveSampler" + sampler.solver.name = solver_name + sampler.nodelist = list(range(n_cols)) + sampler.properties = {"num_qubits": n_cols} + sampler.sample.return_value = _Response() + return sampler + + +def _make_mock_experiment( + inst, energy_scale=1.0, run_index=0, num_random_instances=1, extra_params=None +): + """Return a lightweight mock Experiment with .inst and .param.""" + exp = mock.MagicMock() + exp.inst = inst + exp.param = {"energy_scale": energy_scale, "num_random_instances": num_random_instances} + exp.run_index = run_index + if extra_params: + exp.param.update(extra_params) + return exp + + +def _make_embedded_chain(chain_nodes): + return EmbeddedLattice( + logical_lattice=Chain( + dimensions=(len(chain_nodes),), + periodic=(False,), + ignore_embedding=True, + ), + chain_nodes=chain_nodes, + dimensions=(sum(len(chain) for chain in chain_nodes.values()),), + periodic=(False,), + ) + + +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) + + +class TestChain(unittest.TestCase): + def test_periodic(self): + chain = Chain(dimensions=(6,), periodic=(True,)) + self.assertEqual(chain.num_spins, 6) + self.assertEqual(chain.num_edges, 6) + self.assertIn((5, 0), chain.edge_list) + + def test_non_periodic(self): + chain = Chain(dimensions=(6,), periodic=(False,)) + 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): + chain = Chain(dimensions=(1,), periodic=(True,)) + self.assertEqual(chain.num_spins, 1) + self.assertEqual(chain.num_edges, 0) + + def test_two_node_periodic(self): + chain = Chain(dimensions=(2,), periodic=(True,)) + self.assertEqual(chain.num_edges, 2) + + def test_geometry_name(self): + chain = Chain(dimensions=(6,), periodic=(True,)) + self.assertEqual(chain.geometry_name, "Chain") + + +class TestLattice(unittest.TestCase): + def test_default_periodic(self): + chain = Chain(dimensions=(4,), periodic=(False,)) + self.assertFalse(chain.periodic[0]) + + def test_edge_list_sorted(self): + chain = Chain(dimensions=(5,), periodic=(False,)) + for u, v in chain.edge_list: + self.assertLess(u, v) + + def test_bqm_structure(self): + chain = Chain(dimensions=(4,), periodic=(False,)) + bqm = chain.make_nominal_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): + bqm = Chain(dimensions=(3,), periodic=(True,)).make_nominal_bqm() + self.assertEqual(bqm.vartype, dimod.SPIN) + + def test_orbit_singleton(self): + chain = Chain(dimensions=(4,), periodic=(True,), orbit_type="singleton") + 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): + chain = Chain(dimensions=(4,), periodic=(True,), orbit_type="global") + 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): + chain = Chain( + dimensions=(4,), + periodic=(True,), + orbit_type="explicit", + qubit_orbits=np.array([0, 0, 1, 1]), + coupler_orbits=np.array([0, 0, 1, 1]), + ) + np.testing.assert_array_equal(chain.qubit_orbits, [0, 0, 1, 1]) + + def test_unknown_orbit_type(self): + with self.assertRaises(ValueError): + Chain(dimensions=(4,), periodic=(True,), orbit_type="bogus") + + def test_get_path_invalid_kind(self): + chain = Chain(dimensions=(4,), periodic=(True,)) + with self.assertRaises(ValueError): + chain._get_path(None, "invalid") + + def test_standard_orbit_save_and_load(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain( + dimensions=(4,), + periodic=(True,), + orbit_type="standard", + lattice_data_root=Path(tmpdir), + ) + self.assertIsNotNone(chain.qubit_orbits) + self.assertIsNotNone(chain.coupler_orbits) + # Second instantiation should load from disk + chain2 = Chain( + dimensions=(4,), + periodic=(True,), + orbit_type="standard", + lattice_data_root=Path(tmpdir), + ) + np.testing.assert_array_equal(chain.qubit_orbits, chain2.qubit_orbits) + + def test_nested_embedded_raises(self): + dt = DimerizedTriangular(dimensions=(3, 3), periodic=(True, False), orbit_type="singleton") + # Fake a nested embedded lattice + dt.logical_lattice.logical_lattice = mock.MagicMock() + dt.orbit_type = "global" + with self.assertRaises(NotImplementedError): + dt.initialize_orbits() + + def test_embed_no_embeddings_found(self): + chain = Chain(dimensions=(4,), periodic=(True,)) + 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 = Chain(dimensions=(4,), periodic=(True,), lattice_data_root=Path(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, data_root=Path(tmpdir)) + + chain.embed_lattice(sampler, try_to_load=True, data_root=Path(tmpdir)) + np.testing.assert_array_equal(chain.embedding_list, embeddings) + + def test_embed_find_and_save(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,), lattice_data_root=Path(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=Path(tmpdir)) + # Verify embedding was found and saved + emb_path = chain._get_path(Path(tmpdir), "embedding", sampler_name="MockDWaveSampler") + self.assertTrue(emb_path.exists()) + + +class TestTriangular(unittest.TestCase): + def test_basic_construction(self): + tri = _make_triangular(3, 3) + self.assertEqual(tri.num_spins, 9) + self.assertGreater(tri.num_edges, 0) + self.assertEqual(tri.geometry_name, "Triangular") + + def test_coordinates(self): + tri = _make_triangular(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): + tri = _make_triangular(3, 3, periodic=(False, False), halve_boundary_couplers=True) + bqm = tri.make_nominal_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): + tri = _make_triangular(3, 3, periodic=(False, True)) + self.assertFalse(tri.periodic[0]) + self.assertTrue(tri.periodic[1]) + + +class TestDimerizedTriangular(unittest.TestCase): + def test_basic_construction(self): + dt = DimerizedTriangular(dimensions=(3, 3), periodic=(True, False), orbit_type="singleton") + self.assertEqual(dt.geometry_name, "DimerizedTriangular") + self.assertIsNotNone(dt.logical_lattice) + self.assertEqual(dt.num_spins, 18) + + def test_chain_connectivity_self(self): + dt = DimerizedTriangular(dimensions=(3, 3), periodic=(True, False), orbit_type="singleton") + cc = dt.get_chain_connectivity(0) + self.assertEqual(cc, ((0, 1),)) + + def test_chain_connectivity_cases(self): + dt = DimerizedTriangular(dimensions=(3, 3), periodic=(True, False), orbit_type="singleton") + 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): + dt = DimerizedTriangular(dimensions=(3, 3), periodic=(True, False), orbit_type="singleton") + 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): + embedded = _make_embedded_chain({0: (0, 1, 2), 1: (3, 4, 5)}) + 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): + embedded = _make_embedded_chain({0: (0, 1), 1: (2, 3)}) + 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): + embedded = _make_embedded_chain({0: (0, 1), 1: (2, 3)}) + 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): + # Use a simple embedded lattice with chain_nodes of length 3 + chain_nodes = {0: (10, 11, 12), 1: (20, 21, 22)} + el = _make_embedded_chain(chain_nodes) + # Generic self-connectivity: all combinations within the chain + cc = 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)}, + ) + + +class TestOrbits(unittest.TestCase): + def test_reindex_basic(self): + mapping = {"a": 5, "b": 5, "c": 10} + result = 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 = make_signed_bqm(bqm) + self.assertAlmostEqual(signed.linear["p0"], 0.5) + self.assertAlmostEqual(signed.linear["m0"], -0.5) + + +class TestOptimize(unittest.TestCase): + def test_plain_lattice(self): + chain = Chain(dimensions=(4,), periodic=(False,)) + bqm = chain.make_nominal_bqm() + energy, sample, _ = optimize(chain, bqm, sa_kwargs={"num_sweeps": 256, "num_reads": 16}) + self.assertEqual(energy, -3.0) + self.assertEqual(bqm.energy(sample), energy) + + def test_embedded_lattice(self): + dt = DimerizedTriangular(dimensions=(3, 3), periodic=(True, False), orbit_type="singleton") + bqm = dt.make_nominal_bqm() + energy, sample, _ = optimize(dt, bqm, sa_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})) + + +class TestObservables(unittest.TestCase): + def test_qubit_magnetization(self): + chain = Chain(dimensions=(4,), periodic=(False,)) + bqm = chain.make_nominal_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 = QubitMagnetization().evaluate(exp, bqm, ss) + np.testing.assert_array_equal(result, [0.0, 0.0, 0.0, 0.0]) + + def test_coupler_correlation(self): + chain = Chain(dimensions=(4,), periodic=(False,)) + bqm = chain.make_nominal_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( + CouplerCorrelation().evaluate(exp, bqm, ss_alt), -np.ones(chain.num_edges) + ) + + def test_coupler_frustration(self): + # All aligned (corr=1) -> frustration = 1.0 + chain = Chain(dimensions=(4,), periodic=(False,)) + bqm = chain.make_nominal_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( + CouplerFrustration().evaluate(exp, bqm, ss), np.ones(chain.num_edges) + ) + + def test_sample_energy(self): + chain = Chain(dimensions=(4,), periodic=(False,)) + bqm = chain.make_nominal_bqm() + # All-ones: energy = sum of J for 3 edges = 3.0 + samples = np.ones((1, 4)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp_pos = _make_mock_experiment(chain, energy_scale=1.0) + np.testing.assert_array_almost_equal(SampleEnergy().evaluate(exp_pos, bqm, ss), [3.0]) + + def test_bitpacked_spins(self): + chain = Chain(dimensions=(4,), periodic=(False,)) + bqm = chain.make_nominal_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 = 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)) + + def test_reference_energy_save_load_roundtrip(self): + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "ref.txt" + chain = Chain(dimensions=(4,), periodic=(False,)) + bqm = chain.make_nominal_bqm() + sample = np.array([1, -1, 1, -1]) + obs = ReferenceEnergy() + obs.save(path, -3.0, sample, "SA") + + exp = _make_mock_experiment(chain) + energy, loaded_sample, method = obs.load(exp, bqm, path) + self.assertAlmostEqual(energy, -3.0) + 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 = Chain(dimensions=(4,), periodic=(False,), lattice_data_root=Path(tmpdir)) + bqm = chain.make_nominal_bqm() + obs = 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 = Chain(dimensions=(4,), periodic=(False,)) + bqm = chain.make_nominal_bqm() + obs = 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) + + def test_reference_energy_path(self): + chain = Chain(dimensions=(4,), periodic=(True,)) + bqm = chain.make_nominal_bqm() + exp = _make_mock_experiment(chain) + + with self.assertRaises(NotImplementedError): + get_reference_energy_path(experiment=None, bqm=None) + + path = get_reference_energy_path(experiment=exp, bqm=bqm) + self.assertTrue(str(path).endswith(".txt")) + + # Via dummy data dict (experiment=None) + dummy = {"run_index": 0, "num_random_instances": 1, "inst": chain} + path2 = get_reference_energy_path(bqm=bqm, dummy_experiment_data_dict=dummy) + self.assertTrue(str(path2).endswith(".txt")) + + with tempfile.TemporaryDirectory() as tmpdir: + path3 = get_reference_energy_path(experiment=exp, bqm=bqm, root=tmpdir) + self.assertIn(tmpdir, str(path3)) + + +class TestKinks(unittest.TestCase): + def test_all_aligned(self): + chain = Chain(dimensions=(6,), periodic=(True,)) + bqm = chain.make_nominal_bqm() + samples = np.ones((10, 6)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(chain) + result = KinkKinkCorrelator().evaluate(exp, bqm, ss) + # All neighbors aligned -> every site is a "kink" (K=1 everywhere) + np.testing.assert_array_equal(result, np.zeros(6)) + + def test_mixed_pattern(self): + chain = Chain(dimensions=(6,), periodic=(True,)) + bqm = chain.make_nominal_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 = 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): + tri = _make_triangular(3, 3, periodic=(True, False)) + bqm = tri.make_nominal_bqm() + samples = np.ones((5, 9)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(tri) + result = TriangularOP().evaluate(exp, bqm, ss) + # Uniform spins: equal sublattice mags cancel + np.testing.assert_array_almost_equal(np.abs(result), np.zeros(5), decimal=10) + + def test_evaluate_embedded(self): + dt = DimerizedTriangular(dimensions=(3, 3), periodic=(True, False), orbit_type="singleton") + bqm = dt.make_nominal_bqm() + # Uniform embedded spins + samples = np.ones((5, dt.num_spins)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(dt) + result = TriangularOP().evaluate(exp, bqm, ss) + np.testing.assert_array_almost_equal(np.abs(result), np.zeros(5), decimal=10) + + +class TestSamplerCall(unittest.TestCase): + def test_defaults(self): + sc = SamplerCall(run_index=0) + self.assertEqual(sc.run_index, 0) + self.assertIsNone(sc.bqm) + self.assertEqual(sc.shimdata, {}) + self.assertEqual(sc.nominal_bqms, []) + self.assertEqual(sc.sampler_params, {}) + + def test_with_values(self): + bqm = dimod.BQM(vartype="SPIN") + sc = SamplerCall( + run_index=5, + bqm=bqm, + shimdata={"total_iterations": 1}, + nominal_bqms=[bqm], + sampler_params={"num_reads": 100}, + ) + self.assertEqual(sc.run_index, 5) + self.assertIs(sc.bqm, bqm) + self.assertEqual(sc.shimdata["total_iterations"], 1) + + +class TestExperiment(unittest.TestCase): + def test_default_params(self): + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler) + self.assertEqual(exp.param["energy_scale"], 1.0) + self.assertEqual(exp.param["num_reads"], 100) + self.assertIs(exp.inst, chain) + + def test_results_root(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir) + self.assertEqual(exp.experiment_results_root, Path(tmpdir).resolve()) + + def test_data_path_with_schedule(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir, anneal_time=5.0) + exp.param["anneal_schedule"] = [(0, 1), (5, 0.5)] + del exp.param["anneal_time"] + exp.apply_param({"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 = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + type(sampler).__name__ = "UnknownSampler" + exp = Experiment(chain, sampler, results_root=tmpdir) + with self.assertRaises(TypeError): + exp.apply_param({"energy_scale": 1.0, "anneal_time": 1.0}) + + def test_apply_param_no_anneal_or_schedule_raises(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir) + del exp.param["anneal_time"] + with self.assertRaises(ValueError): + exp.apply_param({"energy_scale": 1.0}) + + def test_initial_shim_no_embeddings(self): + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler) + exp.already_initialized = False + shimdata = exp._make_initial_shim() + self.assertEqual(shimdata["total_iterations"], 0) + self.assertNotIn("flux_biases", shimdata) + + def test_initial_shim_with_embeddings(self): + chain = Chain(dimensions=(4,), periodic=(True,)) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler(num_qubits=128) + exp = Experiment(chain, sampler) + shimdata = exp._make_initial_shim() + self.assertIn("flux_biases", shimdata) + self.assertEqual(len(shimdata["flux_biases"]), 128) + + def test_initial_shim_with_preset_flux_biases(self): + chain = Chain(dimensions=(4,), periodic=(True,)) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler(num_qubits=128) + fb = np.ones(128) * 0.01 + exp = Experiment(chain, sampler) + exp.param["flux_biases"] = fb + shimdata = exp._make_initial_shim() + np.testing.assert_array_almost_equal(shimdata["flux_biases"], fb) + + def test_load_shim_from_file(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir) + exp.data_path = Path(tmpdir) + exp.run_index = 1 + + shimdata = {"total_iterations": 5, "flux_biases": np.zeros(10)} + data = {"shimdata": shimdata} + 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 = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir) + exp.data_path = Path(tmpdir) + exp.run_index = 1 + 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 = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir) + exp.data_path = Path(tmpdir) + exp.run_index = 1 + 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 = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir) + exp.data_path = Path(tmpdir) + exp.run_index = 1 + + # No file exists at all - patch getsize to not fail early + with mock.patch("os.path.getsize", return_value=100): + with self.assertRaises(FileNotFoundError): + exp._load_shim() + + def test_save_and_reload(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir) + 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 = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir) + 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 = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir) + 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_apply_param_sets_run_index_zero(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir, anneal_time=1.0) + exp.apply_param({"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 = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir, anneal_time=1.0) + exp.apply_param({"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({"energy_scale": 1.0, "anneal_time": 1.0}) + self.assertEqual(exp.run_index, 3) + + def test_load_results_ignore_shim(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir, anneal_time=1.0) + exp.apply_param({"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, "shimdata": {}}, f) + + results = exp.load_results(num_iterations=1, ignore_shim=True) + self.assertNotIn("shimdata", results[0]) + + def test_load_results_starting_iteration(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir, anneal_time=1.0) + exp.apply_param({"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, "shimdata": {}}, f) + + results = exp.load_results(num_iterations=3, starting_iteration=2) + self.assertEqual(len(results), 3) + + def test_load_results_corrupted_lzma(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, results_root=tmpdir, anneal_time=1.0) + exp.apply_param({"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): + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler) + sc = SamplerCall(run_index=0) + sc.shimdata = {"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["shimdata"]["total_iterations"], 1) + + def test_make_bqm_no_embeddings(self): + chain = Chain(dimensions=(4,), periodic=(False,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, energy_scale=0.5) + sc = SamplerCall(run_index=0) + sc.nominal_bqms = [chain.make_nominal_bqm()] + sc.shimdata = {"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): + chain = Chain(dimensions=(4,), periodic=(False,)) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, energy_scale=1.0) + sc = SamplerCall(run_index=0) + sc.nominal_bqms = [chain.make_nominal_bqm()] + sc.shimdata = { + "total_iterations": 0, + "relative_coupler_strength": np.ones((1, chain.num_edges)), + } + + bqm = exp._make_bqm(sc) + self.assertGreater(len(bqm.quadratic), 0) + + def test_run_iteration_basic(self): + """run_iteration() exercises the full pipeline: build call, sample, parse, shim, save.""" + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), lattice_data_root=Path(tmpdir)) + exp = Experiment( + chain, _make_sync_sampler(), results_root=tmpdir, anneal_time=1.0, max_iterations=1 + ) + chain._load_embeddings = mock.MagicMock() + finished = exp.run_iteration([{"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("shimdata", data) + self.assertEqual(data["shimdata"]["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 = Chain(dimensions=(4,), periodic=(False,), lattice_data_root=Path(tmpdir)) + exp = Experiment( + chain, _make_sync_sampler(), results_root=tmpdir, anneal_time=1.0, max_iterations=0 + ) + chain._load_embeddings = mock.MagicMock() + finished = exp.run_iteration([{"energy_scale": 1.0, "anneal_time": 1.0}]) + + self.assertTrue(finished) + self.assertEqual(list(exp.data_path.glob("iter*.pkl.lzma")), []) + + def test_flux_bias_shim_basic_update(self): + chain = Chain(dimensions=(4,), periodic=(False,)) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, flux_bias_shim_step=0.001) + + sc = SamplerCall(run_index=0) + sc.shimdata = {"flux_biases": np.zeros(128), "total_iterations": 0} + results = {"QubitMagnetization": np.array([0.1, -0.1, 0.2, -0.2])} + exp._update_flux_bias_shim(sc, results) + self.assertFalse(np.all(sc.shimdata["flux_biases"] == 0)) + + def test_coupler_shim_basic_update(self): + chain = Chain(dimensions=(4,), periodic=(False,)) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler, coupler_shim_step=0.01, energy_scale=1.0) + + sc = SamplerCall(run_index=0) + bqm = chain.make_nominal_bqm() + sc.nominal_bqms = [bqm] + sc.shimdata = { + "total_iterations": 0, + "relative_coupler_strength": np.ones((1, chain.num_edges)), + } + results = {"CouplerFrustration": np.random.rand(1, chain.num_edges)} + exp._update_coupler_shim(sc, results) + self.assertEqual(sc.shimdata["relative_coupler_strength"].shape, (1, chain.num_edges)) + + def test_get_shimdata_not_initialized(self): + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler) + exp.already_initialized = False + shimdata = exp._get_shimdata() + self.assertEqual(shimdata["total_iterations"], 0) + + +class TestFastAnnealExperiment(unittest.TestCase): + def test_default_params(self): + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = FastAnnealExperiment(chain, sampler) + self.assertTrue(exp.param.get("fast_anneal")) + self.assertEqual(exp.param["num_reads"], 100) + + def test_observables(self): + obs_names = {type(o).__name__ for o in FastAnnealExperiment.observables_to_collect} + self.assertIn("QubitMagnetization", obs_names) + self.assertIn("SampleEnergy", obs_names) + self.assertIn("ReferenceEnergy", obs_names) + + +if __name__ == "__main__": + unittest.main()