From 7dac34b1536f15c14e4619dcdc7325b207808558 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Fri, 13 Mar 2026 13:41:22 -0700 Subject: [PATCH 01/13] Add lattice-utils to dwave-experimental --- dwave/experimental/lattice_utils/__init__.py | 19 + .../lattice_utils/experiment/__init__.py | 17 + .../lattice_utils/experiment/experiment.py | 787 ++++++++++++ .../experiment/fast_anneal_experiment.py | 55 + .../lattice_utils/experiment/samplercall.py | 38 + .../lattice_utils/lattice/__init__.py | 18 + .../lattice_utils/lattice/automorphism.py | 1110 +++++++++++++++++ .../lattice_utils/lattice/chain.py | 33 + .../lattice_utils/lattice/embedded_lattice.py | 172 +++ .../lattice_utils/lattice/lattice.py | 289 +++++ .../lattice_utils/lattice/optimize.py | 84 ++ .../lattice_utils/lattice/orbits.py | 239 ++++ .../lattice_utils/lattice/triangular.py | 159 +++ .../lattice_utils/observable/__init__.py | 17 + .../lattice_utils/observable/kinks.py | 49 + .../lattice_utils/observable/observable.py | 255 ++++ .../lattice_utils/observable/triangular.py | 61 + dwave/experimental/lattice_utils/utils.py | 65 + examples/example_1D_Ising_chain.py | 348 ++++++ examples/example_2D_geometric_frustration.py | 288 +++++ 20 files changed, 4103 insertions(+) create mode 100644 dwave/experimental/lattice_utils/__init__.py create mode 100644 dwave/experimental/lattice_utils/experiment/__init__.py create mode 100644 dwave/experimental/lattice_utils/experiment/experiment.py create mode 100644 dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py create mode 100644 dwave/experimental/lattice_utils/experiment/samplercall.py create mode 100644 dwave/experimental/lattice_utils/lattice/__init__.py create mode 100644 dwave/experimental/lattice_utils/lattice/automorphism.py create mode 100644 dwave/experimental/lattice_utils/lattice/chain.py create mode 100644 dwave/experimental/lattice_utils/lattice/embedded_lattice.py create mode 100644 dwave/experimental/lattice_utils/lattice/lattice.py create mode 100644 dwave/experimental/lattice_utils/lattice/optimize.py create mode 100644 dwave/experimental/lattice_utils/lattice/orbits.py create mode 100644 dwave/experimental/lattice_utils/lattice/triangular.py create mode 100644 dwave/experimental/lattice_utils/observable/__init__.py create mode 100644 dwave/experimental/lattice_utils/observable/kinks.py create mode 100644 dwave/experimental/lattice_utils/observable/observable.py create mode 100644 dwave/experimental/lattice_utils/observable/triangular.py create mode 100644 dwave/experimental/lattice_utils/utils.py create mode 100644 examples/example_1D_Ising_chain.py create mode 100644 examples/example_2D_geometric_frustration.py diff --git a/dwave/experimental/lattice_utils/__init__.py b/dwave/experimental/lattice_utils/__init__.py new file mode 100644 index 0000000..f4c17b8 --- /dev/null +++ b/dwave/experimental/lattice_utils/__init__.py @@ -0,0 +1,19 @@ +# 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 latqa.experiment import * +# from latqa.lattice import * +# from latqa.observable import * +# from latqa.analysis import * +# from latqa._paths import * 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..260b545 --- /dev/null +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -0,0 +1,787 @@ +# 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 + +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'] + +class Experiment(): + """Base class for experiment in LatQA.""" + + default_parameters = { + "energy_scale": 1.0, + "num_reads": 100, + "anneal_time": 1.0, + "num_random_instances": 1, + "readout_thermalization": 100, + "flux_bias_shim_step": 0.0, + "coupler_shim_step": 0.0, + "anneal_offset_shim_step": 0.0, + "target_magnetization": 0.0, + } + + observables_to_collect = { + QubitMagnetization(), + CouplerCorrelation(), + CouplerFrustration(), + SampleEnergy(), + BitpackedSpins(), + ReferenceEnergy(), + } + def __init__(self, inst: Lattice, sampler: dimod.Sampler, **kwargs): + self.inst: Lattice = inst + self.sampler: dimod.Sampler = sampler + self.param: dict = {} + self.already_initialized: bool = False # until evidenced + + # Any forced types that might be required. + self.typedict: dict[str, float] = {"energy_scale": float} + + self.experiment_results_root: Path = Path( + kwargs.get("results_root", Path.cwd() / "results") + ).resolve() + + self.loop_data_files: int = kwargs.get("loop_data_files", 1000000000) + self.max_iterations: int | None = kwargs.get("max_iterations", None) + + # Apply parameters + for field, default in self.default_parameters.items(): + value = kwargs.get(field, default) + + if field in kwargs and field in self.typedict: + value = self.typedict[field](value) + + self.param[field] = value + + + def load_results( + self, + num_iterations: int = 100, + starting_iteration: int | None = None, + result_fields: dict[str, Any] | None = None, + quiet: bool = True, + ignore_shim: bool = False, + mod: int | None = None, + ) -> dict[str, Any]: + """Load results from the highest-numbered iterations of the experiment. + + Args: + num_iterations: Maximum number of iterations to load. + starting_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. + mod: If provided, the returned list is truncated so its length is + divisable by ``mod``. + + Returns: + A list of dictionaries containing the results for each iteration. + """ + self._set_run_index() # this is the NEXT run index. + + fnlist = self._get_sorted_results_file_list() + if starting_iteration is not None: + fnlist = fnlist[max(starting_iteration, 0):max(starting_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}) + + if mod is not None: + return results[:mod * (len(results) // mod)] + + return results + + def apply_param(self, param: dict[str, float]) -> bool: + """Apply a parameter configuration to the experiment. + + Args: + param: Dictionary of parameter values to apply to the experiment. + Keys correspond to parameter names and values are converted to + the appropriate type if a converted is defined in ``self.typedict``. + + Returns: + A boolean value corresponding to whether or not the experiment has + already completed all required iterations for this parameter + configuration. + """ + # Reformat for filename consistency; can be important for loading data. + param = self._format_parameter_list([param])[0] + for param_name, param_val in param.items(): + if param_name in self.typedict: + self.param[param_name] = self.typedict[param_name](param_val) + else: + self.param[param_name] = param_val + + self.data_path = self.experiment_results_root / self._get_relative_data_path() + self.already_initialized = self._set_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: + for index, val in response_dict.items(): + + if val.done(): + self.apply_param(parameter_list[index]) + + # Get the results and update the shim + results = self.parse_results(call_dict[index], response_dict[index]) + 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] + break + 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, cinluding the nominal BQMs and any + applied spin-reversal transform. + 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] + + if call.spin_reversal_transform is not None: + for iemb, emb in enumerate(embedding_list): + for iv, v in enumerate(emb): + if call.spin_reversal_transform[v]: + sample_arrays[iemb][:, iv] *= -1 + else: + sample_arrays = [response.samples(sorted_by=None)[:, np.arange(self.inst.num_spins)]] + + if call.spin_reversal_transform is not None: + raise NotImplementedError + + 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*")) + fnlist.sort() + return [str(fn) for fn in fnlist] + + def _set_run_index(self) -> bool: + """Set the run index (data file index, iteration number) for the parameterization. + + Uses the name of the most recently modified file in + the folder. Also creates the data path if it doesn't exist. Return value + is boolean, whether the iterations have already been started or not. + """ + # If the path doesn't exist, create it and set the run index to zero. + if self.data_path.exists() is False: + os.makedirs(self.data_path) + self.run_index = 0 + return False + + fnlist = list(self.data_path.glob("iter*")) + if not fnlist: + self.run_index = 0 + return False + + # Folder exists and is nonempty. + # Only complication is looping; we set the flag to True if we overflow mod + fnlist.sort(key=lambda x: os.path.getmtime(x)) + latest_file_iter = int(fnlist[-1].stem.split(".")[0][4:]) + self.run_index = np.max([0, np.mod(latest_file_iter + 1, self.loop_data_files)]) + if self.run_index < latest_file_iter: + self.has_looped = True + + return True + + 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 type(self.sampler).__name__ in [ + "SimulatedAnnealingSampler", + "SimulatedQuantumAnnealingSampler", + "RotorSampler", + "ParallelTemperingSampler", + "DiscreteSimulatedBifurcationSampler", + ]: + # Applies to Monte Carlo dynamics, for example + pathstring = f'energyscale{energy_scale:0.3}/nsweeps{self.param["num_sweeps"]:010d}mcs' + elif "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(seed=self.run_index) + + 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() + sampler_call.spin_reversal_transform = self._get_spin_reversal_transform() + + # Here we can find out that we're finished. + if 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 "target_s" in entry: + entry["target_s"] = np.round(entry["target_s"], 4) + 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 "num_sweeps" in self.param: + ret["num_sweeps"] = self.param["num_sweeps"] + elif "target_s" in self.param and "dwell_time" in self.param: # fast reverse anneal + ret["x_target_s"] = self.param["target_s"] + ret["x_dwell_time"] = self.param["dwell_time"] + ret["anneal_schedule"] = self.param["anneal_schedule"] + elif "anneal_schedule" in self.param: + ret["anneal_schedule"] = self.param["anneal_schedule"] + elif self.param.get("fast_anneal", False): + ret["fast_anneal"] = True + ret["annealing_time"] = self.param["anneal_time"] + else: + ret["annealing_time"] = self.param["anneal_time"] + + if "reinitialize_state" in self.param and ret["anneal_schedule"][0][1] == 1: + ret["reinitialize_state"] = self.param.get("reinitialize_state", False) + if ret["reinitialize_state"] is None: + ret["reinitialize_state"] = False + + if self.param.get("initial_state", None) is not None and ret["anneal_schedule"][0][1] == 1: + ret["initial_state"] = self.param["initial_state"] + elif "reinitialize_state" in self.param and ret["anneal_schedule"][0][1] == 1: + # Set to None, meaning that it will be randomized. + ret["initial_state"] = { + qubit: np.random.randint(2) * 2 - 1 + for qubit in self.inst.embedding_list.ravel() + } + + 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.""" + file_name = f"iter{np.mod(self.run_index - 1, self.loop_data_files):05d}.pkl.lzma" + return self.data_path / file_name + + 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.""" + # Ok, let's use orbits and stuff. + 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. + assert ( + len(sampler_call.nominal_bqms) == 1 + or sampler_call.nominal_bqms[0] == sampler_call.nominal_bqms[1] + ), "Case for distinct embedded BQMs not implemented yet." + + bqm = sampler_call.nominal_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() + + if sampler_call.spin_reversal_transform is not None: + for iv, v in enumerate(emb): + if sampler_call.spin_reversal_transform[v]: + nominal_bqm.flip_variable(iv) + + 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 + + def _get_spin_reversal_transform(self) -> dict[int, bool] | None: + """Generate a spin-reversal transform for the current sampler.""" + if self.param.get("spin_reversal_transform", False): + seed = self.param.get("spin_reversal_transform_seed", None) + rng = np.random.default_rng(seed) + flips = rng.random(len(self.sampler.nodelist)) > 0.5 + + spin_reversal_transform = {v: False for v in self.sampler.nodelist} + for iv, v in enumerate(self.sampler.nodelist): + if flips[iv]: + spin_reversal_transform[v] = True + + return spin_reversal_transform + + return None 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..7e306f5 --- /dev/null +++ b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py @@ -0,0 +1,55 @@ +# 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 import Experiment +from dwave.experimental.lattice_utils.observable import ( + QubitMagnetization, + CouplerCorrelation, + CouplerFrustration, + SampleEnergy, + BitpackedSpins, + ReferenceEnergy +) + +__all__ = ['FastAnnealExperiment'] + +class FastAnnealExperiment(Experiment): + # Set default parameters + default_parameters = { + "energy_scale": 1.0, + "automorph_embeddings": False, + "spin_reversal_transform": False, + "spin_reversal_transform_seed": None, + "num_reads": 100, + "num_random_instances": None, + "readout_thermalization": 100, + "fast_anneal": True, + "anneal_time": 1.0, + "flux_bias_shim_step": 0.0, + "coupler_shim_step": 0.0, + "coupler_damp": 0.0, + "anneal_offset_shim_step": 0.0, + "anneal_offset_damp": 0.0, + "individual_qubit_anneal_offsets": None, + "target_magnetization": 0.0, + "logical_software": False, + } + observables_to_collect = { + QubitMagnetization(), + CouplerCorrelation(), + CouplerFrustration(), + SampleEnergy(), + BitpackedSpins(), + ReferenceEnergy(), + } diff --git a/dwave/experimental/lattice_utils/experiment/samplercall.py b/dwave/experimental/lattice_utils/experiment/samplercall.py new file mode 100644 index 0000000..5fb430d --- /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 + self.spin_reversal_transform: None = None diff --git a/dwave/experimental/lattice_utils/lattice/__init__.py b/dwave/experimental/lattice_utils/lattice/__init__.py new file mode 100644 index 0000000..18759d2 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/__init__.py @@ -0,0 +1,18 @@ +# 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 * diff --git a/dwave/experimental/lattice_utils/lattice/automorphism.py b/dwave/experimental/lattice_utils/lattice/automorphism.py new file mode 100644 index 0000000..e22996a --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/automorphism.py @@ -0,0 +1,1110 @@ +# 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 import deque +from collections.abc import Hashable +from dataclasses import dataclass +from enum import Enum, auto +import hashlib +from itertools import chain +import random +from typing import Mapping + +import networkx as nx +import numpy as np +from numpy.typing import NDArray + + +@dataclass +class ComponentInfo: + """Container for per-component data used during automorphism discovery on disjoint graphs.""" + u_vector: list + nodes: NDArray + best_perm: NDArray + +class EnterMode(Enum): + """Controls when the ``_enter()`` function attempts to compose new automorphisms.""" + RECURSE = auto() + RECURSE_ONCE = auto() + NO_RECURSE = auto() + +class SchreierContext: + """This object holds mutable states used throughout the automorphism calculation. + + Args: + graph: A NetworkX Graph object representing the input graph. + num_samples: Number of samples to use for generating new coset representatives + from the existing set. If not provided, all coset representatives are used. + seed: Seed used for reproducibility. Defaults to 42. + """ + def __init__(self, graph: nx.Graph, num_samples: int | None = None, seed: int = 42) -> None: + original_nodes_sorted = sorted(graph.nodes()) + self._index_to_node: dict[int, Hashable] = { + new: old for new, old in enumerate(original_nodes_sorted) + } + self._node_to_index: dict[Hashable, int] = { + old: new for new, old in enumerate(original_nodes_sorted) + } + graph = nx.relabel_nodes(graph, self._node_to_index) # relabel nodes contiguously (0...n-1) + + self._nodes: list[int] = list(graph.nodes()) + self._num_nodes: int = graph.number_of_nodes() + self._graph_edges: list[tuple[int, int]] = list(graph.edges()) + self._neighbours: list[set[int]] = [set(graph.neighbors(i)) for i in range(self._num_nodes)] + self._graph: nx.Graph = graph + + self._num_samples: int | None = num_samples + self._rng: random.Random = random.Random(seed) + + self._leaf_nodes: int = 0 + self._nodes_reached: int = 0 + self._depth: int = 0 + + self._u_map: dict[np.intp, int] = {} + self._u_len: int = 0 + self._u_vector: list = [] + self._u_vector_inv: list[list[NDArray[np.intp]]] = [] + + self._identity: NDArray[np.intp] = np.arange(self._num_nodes, dtype=np.intp) + + self._best_perm: NDArray = np.arange(self._num_nodes) + self._best_perm_exist: bool = False + self._compare_adj: bool = False + self._trace_history: list = [] + + self._in_colors_adj: bytearray = bytearray(self._num_nodes) + self._in_refine_stack: bytearray = bytearray(self._num_nodes) + + self._color_degree: list[int] = [0] * self._num_nodes + self._min_color_degree: list[int] = [0] * self._num_nodes + self._max_color_degree: list[int] = [0] * self._num_nodes + self._active_vertices: list[list[int]] = [[] for _ in range(self._num_nodes)] + + if self._num_nodes <= 65535: + self._color_dtype: np.dtype = np.uint16 + else: + self._color_dtype: np.dtype = np.uint32 + + @property + def leaf_nodes(self) -> int: + """Number of leaf nodes encountered in the search tree.""" + return self._leaf_nodes + + @property + def nodes_reached(self) -> int: + """Total number of nodes reached during traversal of the search tree.""" + return self._nodes_reached + + @property + def index_to_node(self) -> dict[int, Hashable]: + """The mapping from the basis of relabelled nodes (0...n-1) to the original + node labels.""" + return self._index_to_node + + @property + def node_to_index(self) -> dict[Hashable, int]: + """The mapping from the original node labels to the basis of relabelled + nodes (0...n-1).""" + return self._node_to_index + + @property + def u_map(self) -> dict[np.intp, int]: + """Map from coset representative group index to stabilizer index.""" + return self._u_map + + @property + def u_vector(self) -> list[list[NDArray[np.intp]]]: + """Coset representatives grouped by stabilizer index.""" + return self._u_vector + + @property + def num_automorphisms(self) -> int: + """Number of automorphisms implied by u_vector.""" + if self._u_vector: + return int(np.prod([len(u_i) + 1 for u_i in self._u_vector], dtype=object)) + else: + return 1 + + @property + def vertex_orbits(self) -> list[list[int]]: + """Vertex orbits induced by the coset representatives in u_vector and returned + in the basis of relabelled nodes (0...n-1).""" + return vertex_orbits(self._u_vector, self._nodes) + + @property + def vertex_orbits_original_labels(self) -> list[list[Hashable]]: + """Vertex orbits induced by the coset representatives in u_vector and returned + with the original node labels.""" + return vertex_orbits(self._u_vector, self._nodes, index_to_node=self._index_to_node) + + @property + def edge_orbits(self) -> list[list[int]]: + """Edge orbits induced by the coset representatives in u_vector and returned + in the basis of relabelled nodes (0...n-1).""" + return edge_orbits(self._u_vector, self._graph_edges) + + @property + def edge_orbits_original_labels(self) -> list[list[Hashable]]: + """Edge orbits induced by the coset representatives in u_vector and returned + with the original node labels.""" + return edge_orbits(self._u_vector, self._graph_edges, index_to_node=self._index_to_node) + + def _test_composability(self, g: NDArray[np.intp]) -> tuple[int, NDArray[np.intp]]: + """Test if an automorphism is composable from coset representatives. + + Based on Algorithm 6.10 from Kreher, D. L., & Stinson, D. R. (1999). + Combinatorial algorithms: Generation, enumeration, and search. + + Modified to use a mask to skip sifting by identity permutations, which + have no effect. + + Args: + g: A permutation represented as a list of integers in one-line notation. + + Returns: + A tuple (i, g_reduced) where i is the index of the first base position + that could not be sifted. If ``g`` is completely sifted the returned index + equals ``self._num_nodes``. ``g_reduced`` is the permutation obtained after + sifting through all positions up to (but not including) the returned + index. + """ + mask = (g != self._identity) + index = mask.argmax() + next_diff = 0 + + while mask[index]: + next_diff += index + if next_diff not in self._u_map: + return next_diff, g + + for i, h in enumerate(self._u_vector[self._u_map[next_diff]]): + if h[next_diff] == g[next_diff]: + break + else: + return next_diff, g + + g = self._u_vector_inv[self._u_map[next_diff]][i][g] + mask = (g[next_diff:] != self._identity[next_diff:]) + index = mask.argmax() + + return self._num_nodes, g + + def _enter(self, g: NDArray[np.intp], mode: EnterMode = EnterMode.RECURSE) -> None: + """Add automorphism if it can't be composed from coset representatives. + + Based on Algorithm 6.11 from Kreher, D. L., & Stinson, D. R. (1999). + Combinatorial algorithms: Generation, enumeration, and search. + + If an automorphism can't be composed from existing coset representatives + it is added as a new coset representative to u_vector. Depending on the + setting of ``mode``, ``_enter()`` is called recursively to attempt to + compose additional coset representatives from the composition between + the newly-discovered coset representative and existing coset representatives. + + The automorphisms discovered will result in pruning comparable to nauty, + as measured by comparing the total number of search tree nodes visited + for zephyr graphs of various sizes. + + Args: + g: A permutation represented as a list of integers in one-line notation. + mode: Specifies if recursive calls to ``enter()`` are performed to attempt + to compose new automorphisms. The setting ``EnterMode.RECURSE_ONCE`` + results in a single call to ``enter()`` per coset representative where + no further attempts to compose automorphisms occur. + """ + i, g = self._test_composability(g) + if i == self._num_nodes: + return + + if i not in self._u_map: + self._u_map[i] = self._u_len + self._u_len += 1 + self._u_vector.append([]) + self._u_vector_inv.append([]) + + self._u_vector[self._u_map[i]].append(g) + self._u_vector_inv[self._u_map[i]].append(inv(self._num_nodes, g)) + + if mode is EnterMode.NO_RECURSE: + return + + for u_i in self._u_vector: + for h in u_i: + f = mult(g, h) + if mode is EnterMode.RECURSE_ONCE: + self._enter(f, mode=EnterMode.NO_RECURSE) + else: + self._enter(f) + + def _refine( + self, + partition: list[set[int]], + trace: NDArray[np.integer], + color: NDArray[np.integer], + num_colors: int, + individualized_vertex: int | None = None, + ) -> None: + """Perform color refinement on the current partition until an equitable + coloring is reached. + + This procedure implements the 1-dimensional Weisfeiler-Leman (WL) refinement, + following Algorithms 2 and 3 of Berkholz (2016), *Tight lower and upper bounds + for the complexity of canonical color refinement*. + + A refinement stack is initialized with either: + • all color classes (if no vertex has been individualized), or + • the color class of the individualized vertex. + + For each color class popped from the stack, the algorithm computes the + color-degree of every vertex: the number of neighbours it has in the refining + color class. These color-degrees determine how each color class should be + split. If a color class contains vertices with differing color-degrees, it is + partitioned into new color classes, and the smaller subcells are pushed onto + the refinement stack. + + The process continues until no color class can be further refined, yielding an + equitable coloring. + + If a vertex was individualized prior to this refinement step, only the + color class containing that vertex needs to be placed on the refinement + stack initially, since only colors adjacent to that color can be affected. + + For performance reasons, ``num_colors`` is passed as a single-element list + so that updates to the number of colors persist across calls without having + to return anything. + + Args: + partition: The current partition structure, represented as a list of sets of vertices + ordered by color. + trace: A list of the sizes of each partition cell (color class), ordered by color. + color: An array mapping each vertex to its current color. + num_colors: The current number of colors in the partition. + individualized_vertex: The vertex individualized prior to this refinement step, if any. + + Returns: + The new number of colors, the updated trace array, and the updated color array. + """ + neighbours = self._neighbours + color_degree = self._color_degree + min_color_degree = self._min_color_degree + max_color_degree = self._max_color_degree + active_vertices = self._active_vertices + in_refine_stack = self._in_refine_stack + in_colors_adj = self._in_colors_adj + + colors_adj = [] + + if individualized_vertex is None: + refine_stack = list(range(num_colors)) + else: + refine_stack = [color[individualized_vertex]] + num_colors = [num_colors] # mutable container so ``_split_up_color()`` can increment it + + for v in refine_stack: + in_refine_stack[v] = 1 + + while refine_stack: + refinement_color = refine_stack.pop() + in_refine_stack[refinement_color] = 0 + + for v in partition[refinement_color]: + for w in neighbours[v]: + color_degree[w] += 1 + cw = color[w] + if color_degree[w] == 1: + active_vertices[cw].append(w) + if in_colors_adj[cw] == 0: + colors_adj.append(cw) + in_colors_adj[cw] = 1 + if color_degree[w] > max_color_degree[cw]: + max_color_degree[cw] = color_degree[w] + + for c in colors_adj: + if trace[c] != len(active_vertices[c]): + min_color_degree[c] = 0 + else: + min_color_degree[c] = max_color_degree[c] + for v in active_vertices[c]: + if color_degree[v] < min_color_degree[c]: + min_color_degree[c] = color_degree[v] + + colors_to_split = [] + for c in colors_adj: + if min_color_degree[c] < max_color_degree[c]: + colors_to_split.append(c) + + for color_to_split in sorted(colors_to_split): + self._split_up_color( + color_to_split=color_to_split, + partition=partition, + color=color, + trace=trace, + active_vertices=active_vertices, + color_degree=color_degree, + min_degree=min_color_degree[color_to_split], + max_degree=max_color_degree[color_to_split], + refine_stack=refine_stack, + in_refine_stack=in_refine_stack, + num_colors=num_colors, + ) + + ## reset attributes + for c in colors_adj: + for v in active_vertices[c]: + color_degree[v] = 0 + max_color_degree[c] = 0 + active_vertices[c] = [] + in_colors_adj[c] = 0 + colors_adj = [] + + return num_colors[0], trace, color + + def _split_up_color( + self, + *, + color_to_split: int, + partition: list[set[int]], + color: NDArray[np.integer], + trace: NDArray[np.integer], + active_vertices: list[list[int]], + color_degree: list[int], + min_degree: int, + max_degree: int, + refine_stack: list[int], + in_refine_stack: bytearray, + num_colors: list[int], + ) -> None: + """Splits a color class into subcells based on the color-degrees of its vertices. + + Based on algorithm 3 of Berkholz (2016), *Tight lower and upper bounds + for the complexity of canonical color refinement*. + + Given a color class ``color_to_split`` whose vertices exhibit differing + color-degrees with respect to the current refining color, this routine + partitions that class into new color classes. Vertices with the same + color-degree remain together, while vertices with different degrees are + assigned fresh color identifiers. + + The largest resulting subcell retains the original color label, while + all smaller subcells are assigned new colors and pushed onto the + refinement stack (Hopcroft's trick). The partition structure, trace array, + number of colors, and vertex-to-color mapping are updated in place. + + Args: + color_to_split: The color class to be split. + partition: The current partition structure, represented as a list of sets of vertices + ordered by color. + color: An array mapping each vertex to its current color. + trace: A list of the sizes of each partition cell (color class), ordered by color. + active_vertices: Lists of vertices adjacent to the color class being split, + ordered by color. + color_degree: The color-degree of each vertex. + min_degree: Minimum color-degree among vertices in the color class being split. + max_degree: Maximum color-degree among vertices in the color class being split. + refine_stack: Stack of colors scheduled for refinement. + in_refine_stack: Flags indicating which colors are already on the stack. + num_colors: The number of colors, used to determine the next color label to assign to + newly-refined cells. Stored as a single-element list so that updates persist across + calls. + """ + degree_to_new_color = [0] * (max_degree + 1) + num_color_degree = [0] * (max_degree + 1) + num_color_degree[0] = trace[color_to_split] - len(active_vertices[color_to_split]) + + for v in active_vertices[color_to_split]: + num_color_degree[color_degree[v]] += 1 + + largest_subcell_degree = 0 + for i in range(1, max_degree + 1): + if num_color_degree[i] > num_color_degree[largest_subcell_degree]: + largest_subcell_degree = i + + for i in range(max_degree + 1): + if num_color_degree[i] > 0: + if i == min_degree: + degree_to_new_color[i] = color_to_split + if not in_refine_stack[color_to_split] and i != largest_subcell_degree: + refine_stack.append(degree_to_new_color[i]) + in_refine_stack[degree_to_new_color[i]] = 1 + else: + degree_to_new_color[i] = num_colors[0] + partition[num_colors[0]] = set() + if in_refine_stack[color_to_split] or i != largest_subcell_degree: + refine_stack.append(degree_to_new_color[i]) + in_refine_stack[degree_to_new_color[i]] = 1 + num_colors[0] += 1 + + for v in active_vertices[color_to_split]: + new_color = degree_to_new_color[color_degree[v]] + if new_color != color_to_split: + partition[color_to_split] = partition[color_to_split] - {v} # must create new obj + partition[new_color].add(v) + trace[color_to_split] -= 1 + trace[new_color] += 1 + color[v] = new_color + + def _canon( + self, + partition: list[set[int]], + trace: NDArray[np.integer], + color: NDArray[np.integer], + num_colors: int, + individualized_vertex: int | None = None, + ) -> None: + """Generate search tree based on iterative color refinement and vertex + individualization. + + Loosely based on Algorithm 7.9 from Kreher, D. L., & Stinson, D. R. (1999). + Combinatorial algorithms: Generation, enumeration, and search. Additional + data structures are used to efficiently track the number of vertices + belonging to each color, vertex colors, and number of colors. Additionally, + the most recently individualized vertex is tracked and used to perform + color refinement more efficiently. + + Color refinement is performed iteratively on a graph until a discrete + coloring is achieved. If the coloring is not discrete after refinement, + vertices belonging to the same color are individualized, meaning that they + are assigned a new color, often breaking the symmetry of the graph and + allowing a subsequent color refinement step to produce further refinement. + + By default, graph comparisons using adjacency matrices are not performed, as + this becomes a bottleneck for even modestly sized graphs. Instead, the + ``trace`` for each graph is compared, which corresponds to the number of + vertices belonging to each color, ordered by color. This check is orders of + magnitude faster and has been found to have identical pruning capability + for graphs of interest, such as chimera, pegasus, and zephyr graphs, as + well as the disjoint compositions of smaller and simpler graphs as may + be encountered when doing parallel embeddings. + + If a graph has more than one component, comparisons using adjacency matrices are + used. This enables isomorphism detection between components, and in turn + a more efficient approach to generating the full automorphism group, which + may contain many automorphisms between isomorphic components. + + Kreher and Stinson perform comprehensive pruning by changing the base of + the left transversals to coincide with the current permutation order up to the + first non-discrete partition cell, or first split. At the cost of performing + this base change, it allows pruning to be performed by only considering + the left transversal with a stabilizer index equal to the index of the first + split. In practice, changing the base at each node of the search tree + becomes prohibitively expensive even more mostly sized graphs, and instead + the approach taken here is to avoid base changes, but instead to more carefully + evaluate which coset representatives to use for pruning. This is done by + ignoring the automorphisms that do not respect the current partition structure. + + Args: + partition: The current partition structure, represented as a list of + sets of vertices ordered by color. + trace: The number of vertices belonging to each color, ordered by color. + color: A map from each vertex to its color. + num_colors: The number of unique colors, equivalent to the number + of cells in the partition. + individualized_vertex: The most recently individualized vertex. + """ + self._nodes_reached += 1 + self._depth += 1 + + num_colors, trace, color = self._refine( + partition, + trace, + color, + num_colors, + individualized_vertex=individualized_vertex + ) + + if not self._best_perm_exist: + self._trace_history.append(trace.tobytes()) + + # first non-singleton block index + first_split = self._num_nodes - 1 + for i, block in enumerate(partition): + if len(block) > 1: + first_split = i + break + + compare_result = 2 + if self._best_perm_exist: # if a leaf node has been reached previously + + if self._compare_adj: + perm_candidate = list(chain.from_iterable(p for p in partition if p is not None)) + compare_result = self._compare(perm_candidate, first_split) + else: + compare_result = trace.tobytes() == self._trace_history[self._depth - 1] + + if compare_result == 0: + return + + if first_split == self._num_nodes - 1: # leaf node reached + self._leaf_nodes += 1 + + if not self._best_perm_exist: + self._best_perm_exist = True + self._best_perm[:] = list(chain.from_iterable(partition)) + + elif compare_result == 2: + perm_candidate = list(chain.from_iterable(partition)) + self._best_perm[:] = perm_candidate + + elif compare_result == 1: + perm_transformed = np.empty(self._num_nodes, dtype=np.intp) + perm_candidate = list(chain.from_iterable(partition)) + perm_transformed[perm_candidate] = self._best_perm + self._enter(perm_transformed) + + return + + candidates = sorted(partition[first_split]) + remaining_in_block = partition[first_split] + updated_partition = partition + trace[first_split] -= 1 + trace[num_colors] = 1 + + while candidates: + vertex = next(iter(candidates)) + updated_partition[first_split] = remaining_in_block - {vertex} + updated_partition[num_colors] = {vertex} + individualized_partition = list(updated_partition) # copy outer list + color[vertex] = num_colors # updated individualized cell + trace_copy = np.array(trace) + color_copy = np.array(color) + + self._canon( + individualized_partition, + trace_copy, + color_copy, + num_colors + 1, + individualized_vertex=vertex + ) + + color[vertex] = first_split + candidates.remove(vertex) + + # prune the search tree using automorphisms + for stab_index, u_index in self._u_map.items(): + if stab_index > vertex: # these automorphisms map vertex to itself + continue + + for g in self._u_vector[u_index]: + if g[vertex] not in candidates: + continue + + # automorphism must respect current partition structure + for w in candidates: + if color[w] != color[g[w]]: + break + else: + candidates.remove(g[vertex]) + + self._depth -= 1 + + def _compare(self, perm: NDArray[np.intp], first_split: int) -> int: + """Compare canonical adjacency matrix against itself under a partial permutation. + + At the first differing entry, returns whether the partial permutatation has + a greater or lesser value, otherwise it returns that they are equal. + + Based on Algorithm 7.6 from Kreher, D. L., & Stinson, D. R. (1999). + Combinatorial algorithms: Generation, enumeration, and search. + + Args: + perm: The permutation of the adjacency matrix to compare the canonical + adjacency matrix against. + first_split: The index of the first block of the partition containing + more than one vertex, defining the size of the partial permutation + of perm to use. + + Returns: + An integer 0, 1, or 2 depending on whether the partial permutation + perm results in an adjacency matrix which is less than, equal to, or + greater than the canonical adjacency matrix, respectively. + """ + neighbours = self._neighbours + best_perm = self._best_perm + for j in range(1, first_split): + neighbours_best_j = neighbours[best_perm[j]] + neighbours_pi_j = neighbours[perm[j]] + for i in range(j): + bit_best = 1 if best_perm[i] in neighbours_best_j else 0 + bit_pi = 1 if perm[i] in neighbours_pi_j else 0 + if bit_best < bit_pi: + return 0 + if bit_best > bit_pi: + return 2 + return 1 + + def _certificate(self) -> bytes: + """Generate a canonical certificate for a graph. + + Based on the permutation ``self.best_perm`` that minimizes the binary value + of the upper triangular portion of the adjacency matrix of the graph, + as found by comparing leaf nodes of the search tree during the search for + automorphisms. + + Returns: + cert_hash: a hash object of the canonical adjacency bitstring. + """ + cert_hash = hashlib.sha256() + neighbours = self._neighbours + best_perm = self._best_perm + + for j in range(1, self._num_nodes): + neighbours_best_j = neighbours[best_perm[j]] + + for i in range(j): + bit = 1 if best_perm[i] in neighbours_best_j else 0 + cert_hash.update(bytes([bit])) + + return cert_hash.digest() + + def _initial_partition( + self, + graph_coloring=None + ) -> tuple[list[set[int] | None], np.ndarray, np.ndarray, int]: + """Initialize the initial partition for a graph. + + Currently this only supports graphs whose vertices are initially the same + color, but could be expanded in the future to accommodate graphs with a + non-trivial initial vertex coloring. + + Returns: + partition: The initial partition structure, represented as a list of sets of vertices + ordered by color. + trace: A list of the sizes of each partition cell (color class), ordered by color. + color: An array mapping each vertex to its current color. + num_colors: The number of colors in the initial partition. + """ + trace = np.zeros(self._num_nodes, dtype=self._color_dtype) + color = np.zeros(self._num_nodes, dtype=self._color_dtype) + + if graph_coloring is None: + partition = [set(self._nodes)] + [None] * (self._num_nodes - 1) + trace[0] = self._num_nodes + num_colors = 1 + else: + num_colors = max(graph_coloring.values()) + 1 + partition = [set() for _ in range(num_colors)] + [None] * (self._num_nodes - num_colors) + for node_label, node_color in graph_coloring.items(): + node_index = self.node_to_index[node_label] # TODO ensure this works with colors not starting at 0..n-1 + partition[node_color].add(node_index) + + for i, nodes in enumerate(partition[:num_colors]): + color[list(nodes)] = i + trace[i] = len(nodes) + + #num_colors = i + 1 + #partition = partition + [None] * (self._num_nodes - num_colors) + + return partition, trace, color, num_colors + + +def vertex_orbits( + u_vector: list[list[NDArray[np.intp]]], + nodes: list[int], + index_to_node: Mapping[int, int] | None = None, +) -> list[list[int]]: + """Calculate vertex orbits using breadth-first search. + + If ``u_vector`` contains no coset representatives, trivial orbits are returned. + + Args: + u_vector: Coset representatives grouped by stabilizer index. + nodes: List of vertex indices used to return trivial orbits when ``u_vector`` is empty. + index_to_node: An optional dictionary for returning orbits with their original node labels. + + Returns: + A list of orbits, each orbit is a list of vertex indices. + + Example: + >>> import numpy as np + >>> from dwave.experimental.automorphism import vertex_orbits + ... + >>> u_vector = [ + ... [np.array([0, 1, 4, 3, 2, 6, 5, 7])], + ... [np.array([2, 1, 4, 3, 0, 7, 5, 6]), np.array([4, 1, 0, 3, 2, 6, 7, 5])], + ... [np.array([0, 3, 2, 1, 4, 5, 6, 7])], + ... ] + >>> nodes = list(range(8)) + >>> vertex_orbits(u_vector, nodes) + [[0, 2, 4], [1, 3], [5, 6, 7]] + """ + if not u_vector: + return [[x] for x in nodes] + + if not all(isinstance(sublist, list) for sublist in u_vector): + raise ValueError("u_vector must be a list of lists.") + + if isinstance(nodes, np.ndarray): + nodes = nodes.tolist() + + if not isinstance(nodes, list) or not all(isinstance(n, int) for n in nodes): + raise ValueError("nodes must be a list of integers.") + + visited = set() + orbits = [] + num_nodes = len(nodes) + generators = [g for u_vector_i in u_vector for g in u_vector_i] + generators.append(np.arange(num_nodes)) + label = (lambda x: index_to_node[x]) if index_to_node is not None else int + + for v_start in nodes: + if v_start in visited: + continue + + visited.add(v_start) + orb = [label(v_start)] + + q = deque([v_start]) + while q: + v_current = q.popleft() + + for g in generators: + v_current = g[v_current] + if v_current not in visited: + visited.add(v_current) + q.append(v_current) + orb.append(label(v_current)) + orb.sort() + orbits.append(orb) + + orbits.sort() + return orbits + + +def edge_orbits( + u_vector: list[list[NDArray[np.intp]]], + edges: list[tuple[int, int]], + index_to_node: Mapping[int, int] | None = None, +) -> list[list[int]]: + """Calculate edge orbits using breadth-first search. + + Args: + u_vector: Coset representatives grouped by stabilizer index. + edges: List of graph edges as tuples of vertex index pairs. + + Returns: + A list of orbits, each orbit is a list of edges (tuples of vertex index pairs). + + Example: + >>> import numpy as np + >>> from dwave.experimental.automorphism import edge_orbits + ... + >>> u_vector = [ + ... [np.array([0, 1, 4, 3, 2, 6, 5, 7])], + ... [np.array([2, 1, 4, 3, 0, 7, 5, 6]), np.array([4, 1, 0, 3, 2, 6, 7, 5])], + ... [np.array([0, 3, 2, 1, 4, 5, 6, 7])], + ... ] + >>> edges = [ + ... (0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6), + ... (6, 7), (7, 0), (0, 3), (1, 4), (2, 6), (5, 7) + ... ] + >>> orbits = edge_orbits(u_vector, edges) + >>> orbits[0] + [(0, 1), (0, 3), (1, 2), (1, 4), (2, 3), (3, 4)] + >>> orbits[1:] + [[(0, 7), (2, 6), (4, 5)], [(5, 6), (5, 7), (6, 7)]] + """ + if not u_vector: + return [[x] for x in edges] + + if not all(isinstance(sublist, list) for sublist in u_vector): + raise ValueError("u_vector must be a list of lists.") + + if not isinstance(edges, list) or not all(isinstance(e, tuple) for e in edges): + raise TypeError("edges must be a list of tuples") + + visited = set() + orbits = [] + generators = [g for u_vector_i in u_vector for g in u_vector_i] + label = (lambda x: index_to_node[x]) if index_to_node is not None else int + + for u_start, v_start in edges: + e_start = (u_start, v_start) if u_start < v_start else (v_start, u_start) + + if e_start in visited: + continue + + visited.add(e_start) + orb = [tuple(label(x) for x in e_start)] + + q = deque([e_start]) + while q: + u, v = q.popleft() + for g in generators: + e_current = (g[u], g[v]) if g[u] < g[v] else (g[v], g[u]) + + if e_current not in visited: + visited.add(e_current) + q.append(e_current) + orb.append(tuple(label(x) for x in e_current)) + + orb.sort() + orbits.append(orb) + + orbits.sort() + return orbits + + +def sample_automorphisms( + u_vector: list[list[NDArray[np.intp]]], + num_samples: int = 1, + seed: int | None = None, +) -> list[NDArray[np.intp]]: + """Uniformly sample automorphisms from the Schreier-Sims representation. + + Randomly samples one coset representative from each non-trivial left + transversal and takes the product, guaranteeing uniform sampling. The + automorphisms can be composed uniformly regardless of the ordering of + the left transversals in 'u_vector'. All products involving identity + automorphisms are ignored. + + Args: + u_vector: Coset representatives grouped by stabilizer index. + num_samples: The number of automorphisms to return. + seed: Random seed for reproducibility. + + Returns: + A list of uniformly sampled automorphisms in one-line notation. + + Example: + >>> import networkx as nx + >>> from dwave.experimental.automorphism import schreier_rep, sample_automorphisms + ... + >>> graph = nx.cycle_graph(8) + >>> result = schreier_rep(graph) + >>> sample_automorphisms(result.u_vector, seed=42) + [array([3, 4, 5, 6, 7, 0, 1, 2])] + >>> sample_automorphisms(result.u_vector, num_samples=2, seed=42) + [array([3, 4, 5, 6, 7, 0, 1, 2]), array([6, 5, 4, 3, 2, 1, 0, 7])] + """ + rng = np.random.default_rng(seed) + num_nodes = len(u_vector[0][0]) + u_counts = [len(u_i) for u_i in u_vector] + sampled_automorphisms = [] + + for _ in range(num_samples): + sample_indices = rng.integers(low=-1, high=u_counts) + g_product = np.arange(num_nodes) + + for i, u_i in enumerate(u_vector): + if sample_indices[i] >= 0: + g = u_i[sample_indices[i]] + g_product = mult(g, g_product) + + sampled_automorphisms.append(g_product) + + return sampled_automorphisms + + +def mult(alpha: NDArray[np.intp], beta: NDArray[np.intp]) -> NDArray[np.intp]: + """Compose two permutations in one-line notation, alpha after beta. + + Args: + alpha: A permutation represented as a list of integers in one-line notation. + beta: Another permutation of the same length. + + Returns: + The composition alpha ∘ beta in one-line notation. + + Example: + >>> import numpy as np + >>> from dwave.experimental.automorphism import mult + ... + >>> alpha = np.array([2,0,1], dtype=np.intp) # (0,2,1): 0->2, 1->0, 2->1 + >>> beta = np.array([1,2,0], dtype=np.intp) # (0,1,2): 0->1, 1->2, 2->0 + >>> mult(alpha, beta) + array([0, 1, 2]) + """ + return alpha[beta] + + +def inv(n: int, alpha: NDArray[np.intp]) -> NDArray[np.intp]: + """Calculate the inverse of a permutation in one-line notation. + + Args: + n: Length of permutation alpha. + alpha: A permutation represented as a list of integers in one-line notation. + + Returns: + The inverse of alpha in one-line notation. + + Example: + >>> import numpy as np + >>> from dwave.experimental.automorphism import inv + ... + >>> alpha = np.array([2,0,1], dtype=np.intp) # (0,2,1): 0->2, 1->0, 2->1 + >>> inv(3, alpha) + array([1, 2, 0]) + """ + alpha_inv = np.empty(n, dtype=np.intp) + alpha_inv[alpha] = np.arange(n, dtype=alpha_inv.dtype) + return alpha_inv + + +def schreier_rep( + graph: nx.Graph, + num_samples: int | None = None, + seed: int = 42, + graph_coloring=None +) -> SchreierContext: + """Compute Schreier representatives and orbits for a graph. + + Builds a depth-first search tree, iteratively performing color refinement + and vertex individualization until leaf nodes are reached where all graph + vertices are uniquely colored. Leaf nodes with identical adjacency matrices + represent graph automorphisms. Discovered automorphisms are used to prune + the search tree. + + If graphs have more than one component, automorphisms are found for each + individual component, and automorphisms between components are determined + by considering which components are isomorphic. Since the number of automorphisms + between isomorphic components scales factorially with the number of components, + this is significantly faster than naively performing refinement-individualization + over the whole graph. It would be possible to update ``u_vector`` directly + without using ``enter()``, which in principle should be even faster, and should + be the first place to look if further performance improvements are required. + + Args: + graph: A NetworkX Graph object representing the input graph containing + the following methods: + - ``nodes()``: iterable of all nodes + - ``number_of_nodes()``: total number of nodes + - ``edges()``: iterable of all edges + - ``neighbors()``: iterable of all neighbours for a given node + num_samples: Number of samples to use for generating new coset representatives + from the existing set. If not provided, all coset representatives are used. + seed: Random seed for reproducibility. Defaults to 42. + """ + if nx.number_connected_components(graph) == 1: + ctx = SchreierContext(graph, num_samples=num_samples, seed=seed) + initial_partition, trace, color, num_colors = ctx._initial_partition(graph_coloring) + + ctx._canon(initial_partition, trace, color, num_colors) + return ctx + + # relabel vertices so components have contiguous labels + index_to_node = {} + node_to_index = {} + next_label = 0 + + component_vertices = list(nx.connected_components(graph)) + for vertices in component_vertices: + for vertex in sorted(vertices): + node_to_index[vertex] = next_label + index_to_node[next_label] = vertex + next_label += 1 + + graph = nx.relabel_nodes(graph, node_to_index, copy=True) + + # enter component automorphisms into global graph + ctx = SchreierContext(graph, num_samples=num_samples, seed=seed) + ctx._index_to_node = index_to_node + ctx._node_to_index = node_to_index + + # group isomorphic components together + components = [ctx._graph.subgraph(c).copy() for c in nx.connected_components(ctx._graph)] + + unique_components = {} + for comp in components: + ctx_comp = SchreierContext(comp, num_samples=num_samples, seed=seed) + ctx_comp._compare_adj = True + + initial_partition, trace, color, num_colors = ctx_comp._initial_partition() + ctx_comp._canon(initial_partition, trace, color, num_colors) + + ctx._nodes_reached += ctx_comp.nodes_reached # update the global search tree statistics + ctx._leaf_nodes += ctx_comp.leaf_nodes + + unique_components.setdefault(ctx_comp._certificate(), []).append( + ComponentInfo(ctx_comp._u_vector, np.array(sorted(comp.nodes())), ctx_comp._best_perm) + ) + + # enter the local automorphisms + graph_nnodes = ctx._graph.number_of_nodes() + for identical_components in unique_components.values(): + for comp in identical_components: + for u in chain.from_iterable(comp.u_vector): + u_global = np.arange(graph_nnodes) + u_global[comp.nodes] = u_global[comp.nodes][u] + ctx._enter(u_global, mode=EnterMode.NO_RECURSE) + + # enter swap automorphisms + for comps in unique_components.values(): + for i in range(len(comps) - 1): + i_nodes = comps[i].nodes + j_nodes = comps[i + 1].nodes + + # swap automorphisms must be entered in the canonical basis + i_canon_perm = comps[i].best_perm + j_canon_perm = comps[i + 1].best_perm + i_canon = i_nodes[i_canon_perm] + j_canon = j_nodes[j_canon_perm] + + u_global = np.arange(graph_nnodes) + u_global[i_canon], u_global[j_canon] = u_global[j_canon], u_global[i_canon] + ctx._enter(u_global, mode=EnterMode.RECURSE_ONCE) + + return ctx + +def array_to_cycle( + array: NDArray[np.intp], + index_to_node: Mapping[int, Hashable] | None = None +) -> str: + """Convert an array in one-line notation to a string in cycle notation. + + Based on Algorithm 6.4 from Kreher, D. L., & Stinson, D. R. (1999). + Combinatorial algorithms: Generation, enumeration, and search. + + Args: + array: The permutation in one-line notation. + index_to_node: An optional relabelling dictionary. By default, array indices + are used. + + Returns: + The permutation as a string in cycle notation. + + Example: + >>> import numpy as np + >>> from dwave.experimental.automorphism import array_to_cycle + ... + >>> alpha = np.array([2,0,1], dtype=np.intp) # (0,2,1): 0->2, 1->0, 2->1 + >>> array_to_cycle(alpha) + '(0,2,1)' + >>> array_to_cycle(np.array([2,0,1]), index_to_node={0: 5, 1: 7, 2: 9}) + '(5,9,7)' + """ + if index_to_node is not None: + expected = set(range(len(array))) + if index_to_node.keys() != expected: + missing = expected - index_to_node.keys() + raise ValueError(f"index_to_node missing keys: {missing}") + + label = (lambda x: str(index_to_node[x])) if index_to_node is not None else str + unvisited = [True] * len(array) + cycle_parts = [] + + for i in range(len(array)): + if unvisited[i]: + cycle_parts.append('(') + cycle_parts.append(label(i)) + unvisited[i] = False + j = i + + while unvisited[array[j]]: + cycle_parts.append(',') + j = array[j] + cycle_parts.append(label(j)) + unvisited[j] = False + + cycle_parts.append(')') + return ''.join(cycle_parts) diff --git a/dwave/experimental/lattice_utils/lattice/chain.py b/dwave/experimental/lattice_utils/lattice/chain.py new file mode 100644 index 0000000..d2160ab --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/chain.py @@ -0,0 +1,33 @@ +# 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 dwave.experimental.lattice_utils.lattice import Lattice + +__all__ = ['Chain'] + +class Chain(Lattice): + + def __init__(self, **kwargs): + periodic: tuple[bool, ...] = kwargs.pop("periodic", (True,)) + self.geometry_name: str = "Chain" + super().__init__(periodic=periodic, **kwargs) + + def generate_edges(self) -> Iterator[tuple[int, int]]: + """Yield edges for a 1D chain lattice.""" + for i in range(self.dimensions[0] - 1): + yield (i, i + 1) + if self.periodic[0]: + yield (i + 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..3a47ce2 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -0,0 +1,172 @@ +# 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. + +"""What do we want here? + +- A class that handles embedded models. +- Perhaps called FixedEmbeddingModel +- Should have a function called get_chain_connections, which takes an edge and +returns a list or tuple of adjacent indices. For example, in the 3D dimer class, +x,y, and z-couplings would return +((1,1)) +((0,0)) +((0,1),(1,0)) +respectively. +A chain coupling (edge=(u,v) where u==v) can, in this case, return ((0,1)). +""" +from itertools import combinations, product +from numbers import Integral +from typing import Any +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): + """Specifics should depend on the embedding. Should this have a logical + model as an attribute? Let's try.""" + def __init__( + self, + logical_lattice_class: Lattice, + logical_lattice_kwargs: dict[str, Any], + chain_nodes: dict[int, tuple[int, Integral]], + **kwargs, + ): + self.logical_lattice: Lattice = logical_lattice_class(**logical_lattice_kwargs) + self.chain_nodes: dict[tuple[int, Integral]] = chain_nodes + self.chain_coupling: float = -kwargs.pop("chain_strength", 2) + kwargs.setdefault("periodic", self.logical_lattice.periodic) + super().__init__(**kwargs) + + def get_chain_connectivity(self, u, v=None): + """Should also work for chains! These can be thought of as self-loops.""" + # 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.""" + 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.""" + 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.""" + # 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.""" + 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.""" + 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)).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.""" + 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.""" + 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..3bf53ee --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/lattice.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. + +import os +from pathlib import Path +from collections.abc import Hashable +from numbers import Integral + +import dimod +#import matplotlib.pyplot as plt +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(): + """Base class for instances in LatQA.""" + + def __init__(self, **kwargs): + + self.dimensions: tuple[int, ...] = kwargs["dimensions"] + self.lattice_data_root: Path = kwargs.get("lattice_data_root", Path.cwd() / "lattice_data") + + self.periodic: tuple[bool, ...] = kwargs.get("periodic", [False] * len(self.dimensions)) + self.edge_list: list[tuple[Hashable, Hashable]] = list(self.generate_edges()) + if len(self.edge_list) > 0: + self.num_spins: Integral = np.max(np.asarray(self.edge_list)) + 1 + + self.num_edges: int = len(self.edge_list) + self.orbit_type: str = kwargs.get("orbit_type", "singleton") + self.initialize_orbits( + kwargs.get("qubit_orbits"), + kwargs.get("coupler_orbits"), + ) + + def embed_lattice( + self, + sampler: dimod.Sampler, + try_to_load: bool = True, + timeout: int = 10, + data_root: str | Path | None = None, + 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. + data_root: Root directory for loading and saving embedding data. + 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, data_root) + filename = self._make_filename( + "embedding", + data_root=data_root, + sampler=sampler, + ) + print(f"Loaded embedding from file {filename}") + return + except FileNotFoundError: + pass + + 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, data_root=data_root) + + return + + def make_nominal_bqm(self, **kwargs) -> 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. + -``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"): + if hasattr(self.logical_lattice, "logical_lattice"): + raise NotImplementedError # Nested embedded lattices not supported. + 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: + print('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: + assert len(qubit_orbits) == self.num_spins + assert len(coupler_orbits) == self.num_edges + 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, + root: Path | None, + 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 + if root is None: + root = Path(__file__).parent.parent / "data" + + if sampler_name is None: + path = Path(root) / kind / class_subdir / self._get_size_pathstring() + else: + path = Path(root) / kind / class_subdir / sampler_name / self._get_size_pathstring() + + return path.with_suffix(".txt") + + def _make_filename( + self, + kind: str, + sampler: dimod.Sampler | None = None, + data_root: str | Path | None = None, + ) -> Path: + """Construct a data filename for the specified sampler and data type.""" + if data_root is None: + data_root = self.lattice_data_root + if sampler is None: + return self._get_path(data_root, kind) + + if type(sampler).__name__ == "MockDWaveSampler": + return self._get_path(data_root, kind, sampler_name="MockDWaveSampler") + return self._get_path(data_root, kind, sampler_name=sampler.solver.name) + + def _save_embeddings( + self, + sampler: dimod.Sampler, + embeddings: NDArray, + data_root: str | Path | None = None, + ) -> None: + """Save embedding data to disk.""" + cache_filename = self._make_filename("embedding", sampler=sampler, data_root=data_root) + os.makedirs(cache_filename.parent, exist_ok=True) + np.savetxt(cache_filename, embeddings, fmt="%d") + print(f"Saved {len(embeddings)} embeddings to file {cache_filename}") + + def _load_embeddings(self, sampler: str, data_root: str | Path | None = None, **kwargs) -> None: + """Load embedding data.""" + filename = self._make_filename("embedding", sampler=sampler, data_root=data_root) + self.embedding_list = np.atleast_2d(np.loadtxt(filename, dtype=int)) + + def _save_orbits(self, data_root: str | Path | None = None) -> None: + """Save qubit and coupler orbits to disk.""" + cache_filename = self._make_filename("orbits", data_root=data_root) + 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") + print(f"Saved orbits to folder {cache_dir}") + + def _load_orbits(self, data_root: str | Path | None = None, **kwargs) -> None: + """Load qubit and coupler orbits.""" + cache_filename = self._make_filename("orbits", data_root=data_root) + 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) + print(f'Loaded orbits from {cache_dir}') + + 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..fe4bbbe --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/optimize.py @@ -0,0 +1,84 @@ +# 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 time +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]: + """Should take the lattice, from which we can infer the appropriate action.""" + 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]: + """Run SA with exponentially increasing sweep counts until no improvement is achieved.""" + start = time.time() + 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: + print(f"Running SA with {num_sweeps} sweeps;", end=" ") + 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] + print(f"best energy is {best_energy}. ") + + 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 + + end = time.time() + print(f"took {end - start:2f}s") + 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..e3ad017 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/orbits.py @@ -0,0 +1,239 @@ +# 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.lattice_utils.lattice.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. + """ + 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. + """ + # Nodes and edges added in a seemingly ugly way in order to get the order right. + ret = dimod.BinaryQuadraticModel(vartype="SPIN") + for var in bqm.variables: + ret.add_variable(f"p{var}", bqm.linear[var]) + for var in bqm.variables: + 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)]) + for u, v in bqm.quadratic: + ret.add_quadratic(f"m{u}", f"m{v}", bqm.quadratic[(u, v)]) + for u, v in bqm.quadratic: + ret.add_quadratic(f"p{u}", f"m{v}", -bqm.quadratic[(u, v)]) + for u, v in bqm.quadratic: + 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. We also need to keep track of self-symmetric pairs + of spins. + + 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. + """ + # Combine coupler orbits so that O(p1p2)=O(m1m2) and O(p1m2)=O(m1p2) + 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"m{v}", f"p{u}")] = min( + signed_coupler_orbits[(f"m{v}", f"p{u}")], + signed_coupler_orbits[(f"m{u}", f"p{v}")], + ) + signed_coupler_orbits[(f"m{u}", f"p{v}")] = signed_coupler_orbits[(f"m{v}", f"p{u}")] + + qubit_orbits = {} + for v in bqm.linear: + qubit_orbits[v] = signed_qubit_orbits[(f"p{v}")] + + coupler_orbits = {} + for u, v in bqm.quadratic: + coupler_orbits[(u, v)] = signed_coupler_orbits[(f"p{u}", 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)`` mappig 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..94c6647 --- /dev/null +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -0,0 +1,159 @@ +# 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 + +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'] + +# For triangular, really for triangular AFM. Will include explicit dimerized +# version in square lattice. +class Triangular(Lattice): + """FM or AFM, like Chain.""" + + def __init__(self, **kwargs): + periodic = kwargs.pop("periodic", (True, False)) + self.geometry_name: str = "Triangular" + self.halve_boundary_couplers: bool = kwargs.pop("halve_boundary_couplers", False) + 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__(periodic=periodic, **kwargs) + assert self.periodic[0] is False or self.dimensions[0] % 3 == 0 + assert self.periodic[1] is False or self.dimensions[1] % 3 == 0 + + def coordinates(self, node: int) -> tuple[int, int]: + """Get y,x coordinates, traversing y first.""" + Ly = self.dimensions[0] + return node % Ly, node // Ly + + def make_nominal_bqm(self, **kwargs) -> dimod.BQM: + """Accommodate the possibility of halving boundary couplers.""" + 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 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. + """ + 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): + + def __init__(self, logical_lattice_class: Lattice = Triangular, **kwargs): + chain_nodes = { + v: (v, v + np.prod(kwargs["dimensions"])) + for v in range(np.prod(kwargs["dimensions"])) + } + self.geometry_name: str = "DimerizedTriangular" + logical_lattice_kwargs = kwargs.copy() + logical_lattice_kwargs.update({"ignore_embedding": True}) + super().__init__( + logical_lattice_class=logical_lattice_class, + logical_lattice_kwargs=logical_lattice_kwargs, + chain_nodes=chain_nodes, + **kwargs + ) + self.halve_boundary_couplers: bool = self.logical_lattice.halve_boundary_couplers + + def get_chain_connectivity( + self, + u: Hashable, + v: Hashable | None = None, + ) -> Iterator[tuple[int, int]]: + """Should also work for chains! These can be thought of as self-loops.""" + 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..c2a952b --- /dev/null +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -0,0 +1,49 @@ +# 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.""" + sample_array = dimod.as_samples(sample_set)[0] + + S = np.roll(sample_array, 1, axis=1) + K = np.multiply(S, sample_array) == np.sign(experiment.param["energy_scale"]) + L = K.shape[-1] + K = np.reshape(K, (-1, L)) + kink_density = np.mean(K) + + CKK = np.zeros((K.shape[-1],)) + + for R in range(1, L): + KR = np.roll(K, R, axis=1) + CKK[R] = np.mean(np.multiply(K, KR)) - np.power(np.mean(K), 2) + + CKK /= kink_density**2 + + return CKK diff --git a/dwave/experimental/lattice_utils/observable/observable.py b/dwave/experimental/lattice_utils/observable/observable.py new file mode 100644 index 0000000..6d5eaed --- /dev/null +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -0,0 +1,255 @@ +# 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 + +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', +] + +class Observable(ABC): + """The observable class does not take any parameters. Its primary + functionality is through the required 'evaluate' method, which requires + parameters 'experiment' and 'bqm' defining the context, and 'sample_set' + which provides the samples on which we compute the observable. Output is a + numpy array of arbitrary type (usually float). + """ + def __init__(self): + self.name: str = type(self).__name__ + + @abstractmethod + def evaluate( + self, + experiment: Experiment, + bqm: dimod.BQM, + sample_set: dimod.SampleSet, + ) -> NDArray: + 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 [] + row, col = np.asarray(experiment.inst.edge_list).T + + # Surprisingly, it's faster to multiply the whole matrix. + 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 [] + row, col = np.asarray(experiment.inst.edge_list).T + + # Surprisingly, it's faster to multiply the whole matrix. + 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: + + if path is None: + 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) + else: + energy, sample, method_string = inst._optimize(bqm) + + 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 None: + 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.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. + """ + 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 needs to be fixed if you have something not in the instance + pathstring that needs to be taken into account, for example if the ground-state + energies depend on the chip. + """ + if bqm is None: + raise NotImplementedError # defunct. + + # 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"].lattice_data_root + else: + root = Path(root) + + path = root / "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..4d12bb6 --- /dev/null +++ b/dwave/experimental/lattice_utils/observable/triangular.py @@ -0,0 +1,61 @@ +# 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. + +"""For triangular order parameters""" + +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 triangular lattices. Unembeds if possible.""" + def evaluate( + self, + experiment: Experiment, + bqm: BQM, + sample_set: dimod.SampleSet, + ) -> NDArray: + + if hasattr(experiment.inst, "logical_lattice"): + # If the lattice is an embedded lattice + lbqm = experiment.inst.unembed_bqm(bqm) + + # unembed the sample set. + 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 edge in lbqm.quadratic: + assert triangular_sublattice[edge[0]] != triangular_sublattice[edge[1]] + + 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 + op = np.matmul(sublattice_mags, angles).ravel() / np.sqrt(3) + + return op diff --git a/dwave/experimental/lattice_utils/utils.py b/dwave/experimental/lattice_utils/utils.py new file mode 100644 index 0000000..754d899 --- /dev/null +++ b/dwave/experimental/lattice_utils/utils.py @@ -0,0 +1,65 @@ +# 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, + repetitions: int = 200, + bootstrap_function: Callable[[NDArray], float] = np.nanmedian, + seed: int | None = None, + skipnan: bool = True, +) -> list[float]: + """Compute bootstrap estimates of a 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, seed=seed): + output.append(bootstrap_function(array[inds])) + + return output + + +def generate_bootstrap_indices( + size: int, + repetitions: int, + seed: int | None = None, +) -> Iterator[NDArray]: + """Generate resampled indices.""" + np.random.seed(seed) + for _ in range(repetitions): + inds = np.random.choice(range(size), replace=True, size=size) + yield inds + + +def confidence_interval(array: NDArray, width: float = 0.95) -> tuple[float, float, float]: + """Ravel and take the quantiles; return median and error bar lengths.""" + 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..b42cd8f --- /dev/null +++ b/examples/example_1D_Ising_chain.py @@ -0,0 +1,348 @@ +# 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 + +# 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_system3.1"), + 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) + +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,), + periodic=(True,), + sampler=sampler, + 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: + exp = experiment.FastAnnealExperiment( + inst=inst, + sampler=sampler, + loop_data_files=30, + max_iterations=5, + energy_scale=energy_scale, + coupler_shim_step=0.05, + flux_bias_shim_step=1e-6, + ) + # 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 (CITE). 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): + 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) + 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", + ) + + #x = ANNEAL_TIMES + M = np.asarray(frust) + bs = np.asarray([bootstrap(_, bootstrap_function=np.nanmedian, seed=None) 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, bootstrap_function=np.nanmedian, seed=None) 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, bootstrap_function=np.nanmedian, seed=None) 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..6dac014 --- /dev/null +++ b/examples/example_2D_geometric_frustration.py @@ -0,0 +1,288 @@ +# 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 + +# Set up a dict for collating statistics +m_dict = {} +psi_dict = {} + +# Just an Advantage2 prototype. +sampler = DWaveSampler(solver="Advantage2_system3.1") + +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) + +inst = lattice.DimerizedTriangular( + dimensions=(9, 12), + periodic=(True, False), + sampler=sampler, + 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, bootstrap_function=np.nanmedian, seed=None) 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, bootstrap_function=np.nanmedian, seed=None) 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, bootstrap_function=np.nanmedian, seed=None) 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() From 23938bfe8929883412e075015f105ddf38dfad2a Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 15 Apr 2026 09:57:34 -0700 Subject: [PATCH 02/13] Add unit tests -move initialization of num_spins to subclass --- .../lattice_utils/lattice/chain.py | 9 +- .../lattice_utils/lattice/embedded_lattice.py | 4 +- .../lattice_utils/lattice/lattice.py | 11 +- .../lattice_utils/lattice/triangular.py | 2 + tests/test_lattice_utils.py | 1042 +++++++++++++++++ 5 files changed, 1057 insertions(+), 11 deletions(-) create mode 100644 tests/test_lattice_utils.py diff --git a/dwave/experimental/lattice_utils/lattice/chain.py b/dwave/experimental/lattice_utils/lattice/chain.py index d2160ab..f7a5cb0 100644 --- a/dwave/experimental/lattice_utils/lattice/chain.py +++ b/dwave/experimental/lattice_utils/lattice/chain.py @@ -23,11 +23,14 @@ class Chain(Lattice): def __init__(self, **kwargs): periodic: tuple[bool, ...] = kwargs.pop("periodic", (True,)) self.geometry_name: str = "Chain" + self.num_spins = kwargs["dimensions"][0] super().__init__(periodic=periodic, **kwargs) def generate_edges(self) -> Iterator[tuple[int, int]]: """Yield edges for a 1D chain lattice.""" - for i in range(self.dimensions[0] - 1): + n = self.dimensions[0] + for i in range(n - 1): yield (i, i + 1) - if self.periodic[0]: - yield (i + 1, 0) + + 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 index 3a47ce2..13d8e21 100644 --- a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -51,6 +51,8 @@ def __init__( self.logical_lattice: Lattice = logical_lattice_class(**logical_lattice_kwargs) 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) @@ -151,7 +153,7 @@ def unembed_sampleset(self, sampleset: dimod.SampleSet) -> dimod.SampleSet: for v in range(len(self.chain_nodes)) ] ) - voted_samples = np.sign(voted_samples + np.random.rand(*voted_samples.shape)).T + 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) diff --git a/dwave/experimental/lattice_utils/lattice/lattice.py b/dwave/experimental/lattice_utils/lattice/lattice.py index 3bf53ee..a3f5f89 100644 --- a/dwave/experimental/lattice_utils/lattice/lattice.py +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -18,7 +18,6 @@ from numbers import Integral import dimod -#import matplotlib.pyplot as plt from minorminer.utils.parallel_embeddings import find_multiple_embeddings import networkx as nx import numpy as np @@ -39,15 +38,13 @@ def __init__(self, **kwargs): self.periodic: tuple[bool, ...] = kwargs.get("periodic", [False] * len(self.dimensions)) self.edge_list: list[tuple[Hashable, Hashable]] = list(self.generate_edges()) - if len(self.edge_list) > 0: - self.num_spins: Integral = np.max(np.asarray(self.edge_list)) + 1 + + 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 = kwargs.get("orbit_type", "singleton") - self.initialize_orbits( - kwargs.get("qubit_orbits"), - kwargs.get("coupler_orbits"), - ) + self.initialize_orbits(kwargs.get("qubit_orbits"), kwargs.get("coupler_orbits")) def embed_lattice( self, diff --git a/dwave/experimental/lattice_utils/lattice/triangular.py b/dwave/experimental/lattice_utils/lattice/triangular.py index 94c6647..a462547 100644 --- a/dwave/experimental/lattice_utils/lattice/triangular.py +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -33,6 +33,7 @@ def __init__(self, **kwargs): periodic = kwargs.pop("periodic", (True, False)) self.geometry_name: str = "Triangular" self.halve_boundary_couplers: bool = kwargs.pop("halve_boundary_couplers", False) + self.num_spins = kwargs["dimensions"][0] * kwargs["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 @@ -117,6 +118,7 @@ def __init__(self, logical_lattice_class: Lattice = Triangular, **kwargs): for v in range(np.prod(kwargs["dimensions"])) } self.geometry_name: str = "DimerizedTriangular" + self.num_spins = 2 * int(np.prod(kwargs["dimensions"])) logical_lattice_kwargs = kwargs.copy() logical_lattice_kwargs.update({"ignore_embedding": True}) super().__init__( diff --git a/tests/test_lattice_utils.py b/tests/test_lattice_utils.py new file mode 100644 index 0000000..b43504d --- /dev/null +++ b/tests/test_lattice_utils.py @@ -0,0 +1,1042 @@ +# 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) + + s = mock.MagicMock() + type(s).__name__ = "DWaveSampler" + s.solver.name = solver_name + s.nodelist = list(range(n_cols)) + s.properties = {"num_qubits": n_cols} + s.sample.return_value = _Response() + return s + + +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_class=Chain, + logical_lattice_kwargs={ + "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): + result = bootstrap(np.array([np.nan, np.nan]), repetitions=5, skipnan=True) + self.assertEqual(len(result), 5) + for val in result: + self.assertTrue(np.isnan(val)) + + def test_bootstrap_skipnan_false(self): + result = bootstrap(np.array([1.0, 2.0, np.nan]), repetitions=5, skipnan=False) + self.assertEqual(len(result), 5) + + def test_bootstrap_custom_function(self): + result = bootstrap(np.arange(20), repetitions=10, bootstrap_function=np.mean, seed=0) + self.assertEqual(len(result), 10) + + def test_bootstrap_seed_reproducibility(self): + r1 = bootstrap(np.arange(10), repetitions=20, seed=123) + r2 = bootstrap(np.arange(10), repetitions=20, seed=123) + np.testing.assert_array_equal(r1, r2) + + def test_generate_bootstrap_indices_correct_count(self): + indices = list(generate_bootstrap_indices(10, 5, seed=0)) + 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_spin_reversal_disabled(self): + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler) + self.assertIsNone(exp._get_spin_reversal_transform()) + + def test_spin_reversal_enabled_with_seed(self): + chain = Chain(dimensions=(4,), periodic=(True,)) + sampler = _make_mock_sampler(num_qubits=8) + exp = Experiment(chain, sampler) + exp.param["spin_reversal_transform"] = True + exp.param["spin_reversal_transform_seed"] = 42 + srt1 = exp._get_spin_reversal_transform() + srt2 = exp._get_spin_reversal_transform() + self.assertEqual(srt1, srt2) + + 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)), + } + sc.spin_reversal_transform = None + + 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_parse_results_with_spin_reversal(self): + chain = Chain(dimensions=(4,), periodic=(False,)) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler() + exp = Experiment(chain, sampler) + exp.run_index = 0 + bqm = chain.make_nominal_bqm() + samples = np.ones((10, 4)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + + sc = SamplerCall(run_index=0) + sc.nominal_bqms = [bqm] + sc.spin_reversal_transform = {0: True, 1: False, 2: True, 3: False} + results = exp.parse_results(sc, ss) + self.assertIn("QubitMagnetization", results) + + 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() From 759eca568278d44b199869c262c06d1f4f9824c9 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 7 May 2026 11:08:56 -0700 Subject: [PATCH 03/13] Improve docstrings and formatting -added more comprehensive docstrings -fixed formatting issues --- dwave/experimental/lattice_utils/__init__.py | 8 +- .../lattice_utils/experiment/experiment.py | 59 +++--- .../experiment/fast_anneal_experiment.py | 3 +- .../lattice_utils/experiment/samplercall.py | 1 + .../lattice_utils/lattice/__init__.py | 2 + .../lattice_utils/lattice/chain.py | 56 ++++- .../lattice_utils/lattice/embedded_lattice.py | 133 +++++++++--- .../lattice_utils/lattice/lattice.py | 151 +++++++------ .../lattice_utils/lattice/optimize.py | 53 ++++- .../lattice_utils/lattice/orbits.py | 57 ++--- .../lattice_utils/lattice/triangular.py | 198 ++++++++++++++---- .../lattice_utils/observable/kinks.py | 1 + .../lattice_utils/observable/observable.py | 29 ++- .../lattice_utils/observable/triangular.py | 2 + dwave/experimental/lattice_utils/utils.py | 64 +++++- 15 files changed, 584 insertions(+), 233 deletions(-) diff --git a/dwave/experimental/lattice_utils/__init__.py b/dwave/experimental/lattice_utils/__init__.py index f4c17b8..d8cb8ad 100644 --- a/dwave/experimental/lattice_utils/__init__.py +++ b/dwave/experimental/lattice_utils/__init__.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -# from latqa.experiment import * -# from latqa.lattice import * -# from latqa.observable import * -# from latqa.analysis import * -# from latqa._paths import * +from . import experiment, lattice, observable + +__all__ = ["experiment", "lattice", "observable"] diff --git a/dwave/experimental/lattice_utils/experiment/experiment.py b/dwave/experimental/lattice_utils/experiment/experiment.py index 260b545..3e1a81e 100644 --- a/dwave/experimental/lattice_utils/experiment/experiment.py +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -31,13 +31,14 @@ CouplerFrustration, SampleEnergy, BitpackedSpins, - ReferenceEnergy + ReferenceEnergy, ) from dwave.experimental.lattice_utils.experiment.samplercall import SamplerCall __all__ = ['Experiment'] -class Experiment(): + +class Experiment: """Base class for experiment in LatQA.""" default_parameters = { @@ -60,6 +61,7 @@ class Experiment(): BitpackedSpins(), ReferenceEnergy(), } + def __init__(self, inst: Lattice, sampler: dimod.Sampler, **kwargs): self.inst: Lattice = inst self.sampler: dimod.Sampler = sampler @@ -85,7 +87,6 @@ def __init__(self, inst: Lattice, sampler: dimod.Sampler, **kwargs): self.param[field] = value - def load_results( self, num_iterations: int = 100, @@ -103,7 +104,7 @@ def load_results( 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. + ``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. @@ -117,7 +118,9 @@ def load_results( fnlist = self._get_sorted_results_file_list() if starting_iteration is not None: - fnlist = fnlist[max(starting_iteration, 0):max(starting_iteration + num_iterations, 0)] + fnlist = fnlist[ + max(starting_iteration, 0) : max(starting_iteration + num_iterations, 0) + ] else: fnlist = fnlist[-num_iterations:] @@ -140,7 +143,7 @@ def load_results( results.append({k: data[k] for k in result_fields}) if mod is not None: - return results[:mod * (len(results) // mod)] + return results[: mod * (len(results) // mod)] return results @@ -261,8 +264,7 @@ def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[st sample_set = {} for iemb, sample_array in enumerate(sample_arrays): sample_set[iemb] = dimod.SampleSet.from_samples_bqm( - sample_array, - call.nominal_bqms[iemb] + sample_array, call.nominal_bqms[iemb] ) results = {} @@ -391,15 +393,17 @@ def _get_parameter_pathstring(self) -> str: 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() - ]) + 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(seed=self.run_index) + nominal_bqm = self.inst.make_nominal_bqm() if not hasattr(self.inst, "embedding_list"): return [nominal_bqm] @@ -521,8 +525,7 @@ def _make_sampler_params(self, **kwargs) -> dict[str, Any]: elif "reinitialize_state" in self.param and ret["anneal_schedule"][0][1] == 1: # Set to None, meaning that it will be randomized. ret["initial_state"] = { - qubit: np.random.randint(2) * 2 - 1 - for qubit in self.inst.embedding_list.ravel() + qubit: np.random.randint(2) * 2 - 1 for qubit in self.inst.embedding_list.ravel() } return ret @@ -644,13 +647,11 @@ def _update_coupler_shim( # Damp the couplers (push toward default value) if "coupler_damp" in self.param and self.param["coupler_damp"] > 0: - excess = ( + excess = relative_coupler_strength[:, bin_edges] - np.mean( 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"] + np.multiply(coupler_signs[bin_edges], excess) * self.param["coupler_damp"] ) # New truncation method... previous is buggy when we mix signs of nominal values. @@ -663,9 +664,7 @@ def _update_coupler_shim( > 1 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - 0.99999 - / nominal_values[bin_edges[violators]] - / energy_scale + 0.99999 / nominal_values[bin_edges[violators]] / energy_scale ) violators = ( @@ -675,9 +674,7 @@ def _update_coupler_shim( < -2 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - -1.99999 - / nominal_values[bin_edges[violators]] - / energy_scale + -1.99999 / nominal_values[bin_edges[violators]] / energy_scale ) # Renormalize each orbit after truncation @@ -702,9 +699,7 @@ def _update_coupler_shim( > 1 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - 0.99999 - / nominal_values[bin_edges[violators]] - / energy_scale + 0.99999 / nominal_values[bin_edges[violators]] / energy_scale ) violators = ( @@ -714,9 +709,7 @@ def _update_coupler_shim( < -2 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - -1.99999 - / nominal_values[bin_edges[violators]] - / energy_scale + -1.99999 / nominal_values[bin_edges[violators]] / energy_scale ) Q = nominal_values * relative_coupler_strength * energy_scale @@ -726,7 +719,7 @@ def _update_coupler_shim( 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.""" diff --git a/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py index 7e306f5..6e67cf6 100644 --- a/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py +++ b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py @@ -19,11 +19,12 @@ CouplerFrustration, SampleEnergy, BitpackedSpins, - ReferenceEnergy + ReferenceEnergy, ) __all__ = ['FastAnnealExperiment'] + class FastAnnealExperiment(Experiment): # Set default parameters default_parameters = { diff --git a/dwave/experimental/lattice_utils/experiment/samplercall.py b/dwave/experimental/lattice_utils/experiment/samplercall.py index 5fb430d..458ba91 100644 --- a/dwave/experimental/lattice_utils/experiment/samplercall.py +++ b/dwave/experimental/lattice_utils/experiment/samplercall.py @@ -18,6 +18,7 @@ __all__ = ['SamplerCall'] + @dataclass class SamplerCall: """Data class for managing asynchronous sampler calls.""" diff --git a/dwave/experimental/lattice_utils/lattice/__init__.py b/dwave/experimental/lattice_utils/lattice/__init__.py index 18759d2..22a583d 100644 --- a/dwave/experimental/lattice_utils/lattice/__init__.py +++ b/dwave/experimental/lattice_utils/lattice/__init__.py @@ -16,3 +16,5 @@ 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 index f7a5cb0..246cad0 100644 --- a/dwave/experimental/lattice_utils/lattice/chain.py +++ b/dwave/experimental/lattice_utils/lattice/chain.py @@ -13,21 +13,65 @@ # 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. - def __init__(self, **kwargs): - periodic: tuple[bool, ...] = kwargs.pop("periodic", (True,)) - self.geometry_name: str = "Chain" - self.num_spins = kwargs["dimensions"][0] - super().__init__(periodic=periodic, **kwargs) + 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.""" + """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) diff --git a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py index 13d8e21..1e4cffa 100644 --- a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -12,22 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""What do we want here? - -- A class that handles embedded models. -- Perhaps called FixedEmbeddingModel -- Should have a function called get_chain_connections, which takes an edge and -returns a list or tuple of adjacent indices. For example, in the 3D dimer class, -x,y, and z-couplings would return -((1,1)) -((0,0)) -((0,1),(1,0)) -respectively. -A chain coupling (edge=(u,v) where u==v) can, in this case, return ((0,1)). -""" + from itertools import combinations, product from numbers import Integral -from typing import Any from collections.abc import Iterator, Hashable import dimod @@ -38,17 +25,37 @@ __all__ = ['EmbeddedLattice'] + class EmbeddedLattice(Lattice): - """Specifics should depend on the embedding. Should this have a logical - model as an attribute? Let's try.""" + """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_class: Lattice, - logical_lattice_kwargs: dict[str, Any], + logical_lattice: Lattice, chain_nodes: dict[int, tuple[int, Integral]], **kwargs, ): - self.logical_lattice: Lattice = logical_lattice_class(**logical_lattice_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"): @@ -56,8 +63,24 @@ def __init__( kwargs.setdefault("periodic", self.logical_lattice.periodic) super().__init__(**kwargs) - def get_chain_connectivity(self, u, v=None): - """Should also work for chains! These can be thought of as self-loops.""" + 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)) @@ -66,7 +89,12 @@ def get_chain_connectivity(self, u, v=None): 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.""" + """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. @@ -82,7 +110,15 @@ def generate_edges(self) -> Iterator[tuple[Hashable, Hashable]]: yield u_chain[edge[0]], v_chain[edge[1]] def make_nominal_bqm(self, **kwargs) -> dimod.BQM: - """Construct and embed the nominal 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) @@ -90,7 +126,16 @@ def make_nominal_bqm(self, **kwargs) -> dimod.BQM: 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.""" + """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())) @@ -124,7 +169,15 @@ def embed_bqm(self, logical_bqm: dimod.BQM) -> dimod.BQM: return embedded_bqm def unembed_bqm(self, embedded_bqm: dimod.BQM) -> dimod.BQM: - """Unembed an embedded BQM back onto the logical variable space.""" + """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) @@ -138,13 +191,21 @@ def unembed_bqm(self, embedded_bqm: dimod.BQM) -> dimod.BQM: for u, v in embedded_bqm.quadratic: if which_spin[u] != which_spin[v]: - bias_uv = embedded_bqm.quadratic[u, 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.""" + """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( @@ -158,7 +219,14 @@ def unembed_sampleset(self, sampleset: dimod.SampleSet) -> dimod.SampleSet: 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.""" + """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] @@ -166,7 +234,14 @@ def embed_sample(self, sample: NDArray) -> NDArray: return ret def unembed_sample(self, sample: NDArray) -> NDArray: - """Unembed a physical sample using majority vote with random tie-breaking.""" + """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) diff --git a/dwave/experimental/lattice_utils/lattice/lattice.py b/dwave/experimental/lattice_utils/lattice/lattice.py index a3f5f89..af59079 100644 --- a/dwave/experimental/lattice_utils/lattice/lattice.py +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -14,8 +14,9 @@ import os from pathlib import Path -from collections.abc import Hashable -from numbers import Integral +from collections.abc import Iterator, Hashable +from abc import ABC, abstractmethod +import warnings import dimod from minorminer.utils.parallel_embeddings import find_multiple_embeddings @@ -28,30 +29,62 @@ __all__ = ['Lattice'] -class Lattice(): - """Base class for instances in LatQA.""" - def __init__(self, **kwargs): +class Lattice(ABC): + """An abstract base class for representing lattice geometries used in lattice-utils experiments. - self.dimensions: tuple[int, ...] = kwargs["dimensions"] - self.lattice_data_root: Path = kwargs.get("lattice_data_root", Path.cwd() / "lattice_data") + Subclasses are resonsible for defining the lattice geometry itself. In particular, + a subclass must: - self.periodic: tuple[bool, ...] = kwargs.get("periodic", [False] * len(self.dimensions)) + - 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 = kwargs.get("orbit_type", "singleton") - self.initialize_orbits(kwargs.get("qubit_orbits"), kwargs.get("coupler_orbits")) + 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, - data_root: str | Path | None = None, max_number_of_embeddings: int | None = None, min_number_of_embeddings: int = 1, exclude_qubits: list = [], @@ -64,7 +97,6 @@ def embed_lattice( 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. - data_root: Root directory for loading and saving embedding data. 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 @@ -76,33 +108,26 @@ def embed_lattice( if try_to_load: try: - self._load_embeddings(sampler, data_root) - filename = self._make_filename( - "embedding", - data_root=data_root, - sampler=sampler, - ) - print(f"Loaded embedding from file {filename}") + self._load_embeddings(sampler) return except FileNotFoundError: - pass + 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} + embedder_kwargs={'timeout': timeout}, ) if not embedding_dicts: - raise ValueError("No Embeddings Found") + 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, data_root=data_root) - - return + self._save_embeddings(sampler, embeddings) - def make_nominal_bqm(self, **kwargs) -> dimod.BQM: + #def make_nominal_bqm(self, **kwargs) -> dimod.BQM: + def make_nominal_bqm(self) -> dimod.BQM: """Construct a default nominal BQM coupling strength values set to +1. Args: @@ -135,21 +160,20 @@ def initialize_orbits( orbit and all chain couplers in another. -``standard``: Load previously computed automorphism-based orbits, or compute them and save them if unavailable. - -``explicit``: use the orbit assignments provided via ``qubit_orbits`` + -``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 + 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"): - if hasattr(self.logical_lattice, "logical_lattice"): - raise NotImplementedError # Nested embedded lattices not supported. 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) @@ -163,7 +187,7 @@ def initialize_orbits( try: self._load_orbits() except FileNotFoundError: - print('Calculating orbits...') + # calculating orbits bqm = self.make_nominal_bqm() self.qubit_orbits, self.coupler_orbits = get_orbits(bqm, self.edge_list) self._save_orbits() @@ -174,19 +198,25 @@ def initialize_orbits( elif self.orbit_type == "explicit": if qubit_orbits is not None and coupler_orbits is not None: - assert len(qubit_orbits) == self.num_spins - assert len(coupler_orbits) == self.num_edges + 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}.' \ + f'Unknown orbit type {self.orbit_type}. ' 'Must be "global", "standard", "singleton", or "explicit".' ) def _get_path( self, - root: Path | None, kind: str, sampler_name: str | None = None, extra_subdir: str | Path | None = None, @@ -198,66 +228,49 @@ def _get_path( class_subdir = Path(self.geometry_name) if extra_subdir is not None: class_subdir = class_subdir / extra_subdir - if root is None: - root = Path(__file__).parent.parent / "data" - if sampler_name is None: - path = Path(root) / kind / class_subdir / self._get_size_pathstring() - else: - path = Path(root) / kind / class_subdir / sampler_name / self._get_size_pathstring() + base_dir = self.data_root / "lattice_data" / kind / class_subdir + if sampler_name is not None: + base_dir = base_dir / sampler_name - return path.with_suffix(".txt") + filename = f"{self._get_size_pathstring()}.txt" + return base_dir / filename - def _make_filename( - self, - kind: str, - sampler: dimod.Sampler | None = None, - data_root: str | Path | None = None, - ) -> Path: + def _make_filename(self, kind: str, sampler: dimod.Sampler | None = None) -> Path: """Construct a data filename for the specified sampler and data type.""" - if data_root is None: - data_root = self.lattice_data_root if sampler is None: - return self._get_path(data_root, kind) + return self._get_path(kind) if type(sampler).__name__ == "MockDWaveSampler": - return self._get_path(data_root, kind, sampler_name="MockDWaveSampler") - return self._get_path(data_root, kind, sampler_name=sampler.solver.name) + 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, - data_root: str | Path | None = None, - ) -> None: + def _save_embeddings(self, sampler: dimod.Sampler, embeddings: NDArray) -> None: """Save embedding data to disk.""" - cache_filename = self._make_filename("embedding", sampler=sampler, data_root=data_root) + cache_filename = self._make_filename("embedding", sampler=sampler) os.makedirs(cache_filename.parent, exist_ok=True) np.savetxt(cache_filename, embeddings, fmt="%d") - print(f"Saved {len(embeddings)} embeddings to file {cache_filename}") - def _load_embeddings(self, sampler: str, data_root: str | Path | None = None, **kwargs) -> None: + def _load_embeddings(self, sampler: str) -> None: """Load embedding data.""" - filename = self._make_filename("embedding", sampler=sampler, data_root=data_root) + filename = self._make_filename("embedding", sampler=sampler) self.embedding_list = np.atleast_2d(np.loadtxt(filename, dtype=int)) - def _save_orbits(self, data_root: str | Path | None = None) -> None: + def _save_orbits(self) -> None: """Save qubit and coupler orbits to disk.""" - cache_filename = self._make_filename("orbits", data_root=data_root) + 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") - print(f"Saved orbits to folder {cache_dir}") - def _load_orbits(self, data_root: str | Path | None = None, **kwargs) -> None: + def _load_orbits(self) -> None: """Load qubit and coupler orbits.""" - cache_filename = self._make_filename("orbits", data_root=data_root) + 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) - print(f'Loaded orbits from {cache_dir}') def _get_instance_pathstring(self) -> str: """Construct an instance-specific pathstring. diff --git a/dwave/experimental/lattice_utils/lattice/optimize.py b/dwave/experimental/lattice_utils/lattice/optimize.py index fe4bbbe..1a91f38 100644 --- a/dwave/experimental/lattice_utils/lattice/optimize.py +++ b/dwave/experimental/lattice_utils/lattice/optimize.py @@ -13,7 +13,6 @@ # limitations under the License. from __future__ import annotations -import time from typing import Any from dwave.samplers import SimulatedAnnealingSampler @@ -23,21 +22,42 @@ __all__ = ['optimize', 'optimize_increasing_sa_sweeps'] + def optimize( lattice: Lattice, bqm: dimod.BQM, sa_kwargs: dict[str, Any] | None = None, ) -> tuple[float, NDArray, str]: - """Should take the lattice, from which we can infer the appropriate action.""" + """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 + lattice.logical_lattice, lattice.unembed_bqm(bqm), sa_kwargs=sa_kwargs ) embedded_sample = lattice.embed_sample(logical_sample) embedded_energy = bqm.energy(embedded_sample) @@ -47,14 +67,29 @@ def optimize( # 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]: - """Run SA with exponentially increasing sweep counts until no improvement is achieved.""" - start = time.time() + """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: @@ -63,12 +98,10 @@ def optimize_increasing_sa_sweeps( num_reads = sa_kwargs.get("num_reads", 256) while True: - print(f"Running SA with {num_sweeps} sweeps;", end=" ") 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] - print(f"best energy is {best_energy}. ") if best_energy < reference_energy: reference_energy = best_energy @@ -79,6 +112,4 @@ def optimize_increasing_sa_sweeps( else: break - end = time.time() - print(f"took {end - start:2f}s") 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 index e3ad017..30f6b80 100644 --- a/dwave/experimental/lattice_utils/lattice/orbits.py +++ b/dwave/experimental/lattice_utils/lattice/orbits.py @@ -29,11 +29,16 @@ '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()} @@ -41,34 +46,28 @@ def reindex(mapping: dict[Hashable, int]) -> dict[Hashable, int]: 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) + 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. - """ - # Nodes and edges added in a seemingly ugly way in order to get the order right. + """ ret = dimod.BinaryQuadraticModel(vartype="SPIN") for var in bqm.variables: ret.add_variable(f"p{var}", bqm.linear[var]) - for var in bqm.variables: 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)]) - for u, v in bqm.quadratic: ret.add_quadratic(f"m{u}", f"m{v}", bqm.quadratic[(u, v)]) - for u, v in bqm.quadratic: ret.add_quadratic(f"p{u}", f"m{v}", -bqm.quadratic[(u, v)]) - for u, v in bqm.quadratic: ret.add_quadratic(f"m{u}", f"p{v}", -bqm.quadratic[(u, v)]) return ret @@ -94,7 +93,7 @@ def get_bqm_orbits( 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. + edge to an integer orbit label. """ # The function first adds auxiliary elements to a BQM graph = nx.Graph() @@ -170,20 +169,24 @@ def get_unsigned_bqm_orbits( """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. We also need to keep track of self-symmetric pairs - of spins. + 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: + 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. + each coupling to its orbit index. """ - # Combine coupler orbits so that O(p1p2)=O(m1m2) and O(p1m2)=O(m1p2) + 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}")], @@ -191,20 +194,18 @@ def get_unsigned_bqm_orbits( ) signed_coupler_orbits[(f"m{u}", f"m{v}")] = signed_coupler_orbits[(f"p{u}", f"p{v}")] - signed_coupler_orbits[(f"m{v}", f"p{u}")] = min( - signed_coupler_orbits[(f"m{v}", f"p{u}")], + 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"m{v}", f"p{u}")] + 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}")] - coupler_orbits = {} - for u, v in bqm.quadratic: - coupler_orbits[(u, v)] = signed_coupler_orbits[(f"p{u}", f"p{v}")] - return reindex(qubit_orbits), reindex(coupler_orbits) @@ -214,12 +215,12 @@ def get_orbits(bqm: dimod.BQM, edge_list: list[int, int]) -> tuple[NDArray, NDAr 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)`` mappig each entry of ``edge_list`` + 1-D array of length ``len(edge_list)`` mapping each entry of ``edge_list`` to an orbit index. """ signed_bqm = make_signed_bqm(bqm) diff --git a/dwave/experimental/lattice_utils/lattice/triangular.py b/dwave/experimental/lattice_utils/lattice/triangular.py index a462547..fe18ad2 100644 --- a/dwave/experimental/lattice_utils/lattice/triangular.py +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -13,6 +13,7 @@ # limitations under the License. from collections.abc import Iterator, Hashable +from pathlib import Path import networkx as nx import numpy as np @@ -24,31 +25,89 @@ __all__ = ['Triangular', 'DimerizedTriangular'] -# For triangular, really for triangular AFM. Will include explicit dimerized -# version in square lattice. + class Triangular(Lattice): - """FM or AFM, like Chain.""" + """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, **kwargs): - periodic = kwargs.pop("periodic", (True, False)) + 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 = kwargs.pop("halve_boundary_couplers", False) - self.num_spins = kwargs["dimensions"][0] * kwargs["dimensions"][1] + 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__(periodic=periodic, **kwargs) - assert self.periodic[0] is False or self.dimensions[0] % 3 == 0 - assert self.periodic[1] is False or self.dimensions[1] % 3 == 0 + 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]: - """Get y,x coordinates, traversing y first.""" - Ly = self.dimensions[0] - return node % Ly, node // Ly + """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. - def make_nominal_bqm(self, **kwargs) -> dimod.BQM: - """Accommodate the possibility of halving boundary couplers.""" + Returns: + A dimod.BQM representing the nominal triangular lattice. + """ graph = self._make_networkx_graph() bqm = dimod.BQM(vartype="SPIN") @@ -63,13 +122,17 @@ def make_nominal_bqm(self, **kwargs) -> dimod.BQM: return bqm def generate_edges(self) -> Iterator[tuple[int, int]]: - """Yield edges for the 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 + """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 + length_y, length_x = self.dimensions graph = nx.Graph() for x in range(length_x): @@ -110,22 +173,75 @@ def generate_edges(self) -> Iterator[tuple[int, int]]: 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, + ) - def __init__(self, logical_lattice_class: Lattice = Triangular, **kwargs): - chain_nodes = { - v: (v, v + np.prod(kwargs["dimensions"])) - for v in range(np.prod(kwargs["dimensions"])) - } - self.geometry_name: str = "DimerizedTriangular" - self.num_spins = 2 * int(np.prod(kwargs["dimensions"])) - logical_lattice_kwargs = kwargs.copy() - logical_lattice_kwargs.update({"ignore_embedding": True}) super().__init__( - logical_lattice_class=logical_lattice_class, - logical_lattice_kwargs=logical_lattice_kwargs, + logical_lattice=logical_lattice, chain_nodes=chain_nodes, - **kwargs + 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 @@ -133,8 +249,20 @@ def get_chain_connectivity( self, u: Hashable, v: Hashable | None = None, - ) -> Iterator[tuple[int, int]]: - """Should also work for chains! These can be thought of as self-loops.""" + ) -> 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. @@ -151,7 +279,7 @@ def get_chain_connectivity( 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 ux == vx - 1 or vx == 0: # (ux == self.dimensions[1] - 1 and self.periodic[1]): if uy == vy: return ((1, 0),) return ((0, 1),) diff --git a/dwave/experimental/lattice_utils/observable/kinks.py b/dwave/experimental/lattice_utils/observable/kinks.py index c2a952b..a9583e2 100644 --- a/dwave/experimental/lattice_utils/observable/kinks.py +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -22,6 +22,7 @@ __all__ = ['KinkKinkCorrelator'] + class KinkKinkCorrelator(Observable): """For 1D chains.""" diff --git a/dwave/experimental/lattice_utils/observable/observable.py b/dwave/experimental/lattice_utils/observable/observable.py index 6d5eaed..e7380d8 100644 --- a/dwave/experimental/lattice_utils/observable/observable.py +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -33,6 +33,7 @@ 'ReferenceEnergy', ] + class Observable(ABC): """The observable class does not take any parameters. Its primary functionality is through the required 'evaluate' method, which requires @@ -40,6 +41,7 @@ class Observable(ABC): which provides the samples on which we compute the observable. Output is a numpy array of arbitrary type (usually float). """ + def __init__(self): self.name: str = type(self).__name__ @@ -55,6 +57,7 @@ def evaluate( class QubitMagnetization(Observable): """Compute the mean magnetization of each qubit.""" + def evaluate( self, experiment: Experiment, @@ -67,6 +70,7 @@ def evaluate( class CouplerCorrelation(Observable): """Compute pairwise spin correlations for each coupler.""" + def evaluate( self, experiment: Experiment, @@ -85,6 +89,7 @@ def evaluate( class CouplerFrustration(Observable): """Compute the mean coupler frustration for each edge.""" + def evaluate( self, experiment: Experiment, @@ -98,10 +103,9 @@ def evaluate( # Surprisingly, it's faster to multiply the whole matrix. 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"]) - ) + 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 @@ -111,6 +115,7 @@ class SampleEnergy(Observable): Energies exclude the magnitude of ``energy_scale`` but include its sign. """ + def evaluate( self, experiment: Experiment, @@ -122,6 +127,7 @@ def evaluate( class BitpackedSpins(Observable): """Return bitpacked spins and a tuple of the array size.""" + def evaluate( self, experiment: Experiment, @@ -140,7 +146,9 @@ def evaluate( class ReferenceEnergy(Observable): """Return a cached reference energy, computing it and saving it if needed.""" - def evaluate(self, + + def evaluate( + self, experiment: Experiment, bqm: dimod.BQM, sample_set: dimod.SampleSet, @@ -219,7 +227,7 @@ def get_reference_energy_path( 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 needs to be fixed if you have something not in the instance pathstring that needs to be taken into account, for example if the ground-state energies depend on the chip. @@ -239,11 +247,16 @@ def get_reference_energy_path( } if root is None: - root = experiment_data_dict["inst"].lattice_data_root + root = experiment_data_dict["inst"].data_root else: root = Path(root) - path = root / "reference_energies" / experiment_data_dict["inst"]._get_instance_pathstring() + 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( diff --git a/dwave/experimental/lattice_utils/observable/triangular.py b/dwave/experimental/lattice_utils/observable/triangular.py index 4d12bb6..89ef7b2 100644 --- a/dwave/experimental/lattice_utils/observable/triangular.py +++ b/dwave/experimental/lattice_utils/observable/triangular.py @@ -24,8 +24,10 @@ __all__ = ['TriangularOP'] + class TriangularOP(Observable): """For triangular lattices. Unembeds if possible.""" + def evaluate( self, experiment: Experiment, diff --git a/dwave/experimental/lattice_utils/utils.py b/dwave/experimental/lattice_utils/utils.py index 754d899..94341ab 100644 --- a/dwave/experimental/lattice_utils/utils.py +++ b/dwave/experimental/lattice_utils/utils.py @@ -17,14 +17,33 @@ 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, - seed: int | None = None, skipnan: bool = True, ) -> list[float]: - """Compute bootstrap estimates of a statistic.""" + """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)] @@ -33,7 +52,7 @@ def bootstrap( output = [] if len(array) > 0: - for inds in generate_bootstrap_indices(array.size, repetitions, seed=seed): + for inds in generate_bootstrap_indices(array.size, repetitions, rng): output.append(bootstrap_function(array[inds])) return output @@ -42,17 +61,46 @@ def bootstrap( def generate_bootstrap_indices( size: int, repetitions: int, - seed: int | None = None, + rng: np.random.Generator, ) -> Iterator[NDArray]: - """Generate resampled indices.""" - np.random.seed(seed) + """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 = np.random.choice(range(size), replace=True, size=size) + inds = rng.choice(range(size), replace=True, size=size) yield inds def confidence_interval(array: NDArray, width: float = 0.95) -> tuple[float, float, float]: - """Ravel and take the quantiles; return median and error bar lengths.""" + """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 From b6bf8f6c28c1b0252c1c82b54db46d8ff8fbc108 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 7 May 2026 11:08:56 -0700 Subject: [PATCH 04/13] Improve docstrings and formatting -added more comprehensive docstrings -fixed formatting issues --- dwave/experimental/lattice_utils/__init__.py | 8 +- .../lattice_utils/experiment/experiment.py | 59 +++--- .../experiment/fast_anneal_experiment.py | 3 +- .../lattice_utils/experiment/samplercall.py | 1 + .../lattice_utils/lattice/__init__.py | 2 + .../lattice_utils/lattice/chain.py | 56 ++++- .../lattice_utils/lattice/embedded_lattice.py | 133 +++++++++--- .../lattice_utils/lattice/lattice.py | 151 +++++++------ .../lattice_utils/lattice/optimize.py | 53 ++++- .../lattice_utils/lattice/orbits.py | 57 ++--- .../lattice_utils/lattice/triangular.py | 198 ++++++++++++++---- .../lattice_utils/observable/kinks.py | 1 + .../lattice_utils/observable/observable.py | 29 ++- .../lattice_utils/observable/triangular.py | 2 + dwave/experimental/lattice_utils/utils.py | 64 +++++- examples/example_1D_Ising_chain.py | 30 ++- examples/example_2D_geometric_frustration.py | 23 +- tests/test_lattice_utils.py | 80 ++++--- 18 files changed, 651 insertions(+), 299 deletions(-) diff --git a/dwave/experimental/lattice_utils/__init__.py b/dwave/experimental/lattice_utils/__init__.py index f4c17b8..d8cb8ad 100644 --- a/dwave/experimental/lattice_utils/__init__.py +++ b/dwave/experimental/lattice_utils/__init__.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -# from latqa.experiment import * -# from latqa.lattice import * -# from latqa.observable import * -# from latqa.analysis import * -# from latqa._paths import * +from . import experiment, lattice, observable + +__all__ = ["experiment", "lattice", "observable"] diff --git a/dwave/experimental/lattice_utils/experiment/experiment.py b/dwave/experimental/lattice_utils/experiment/experiment.py index 260b545..3e1a81e 100644 --- a/dwave/experimental/lattice_utils/experiment/experiment.py +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -31,13 +31,14 @@ CouplerFrustration, SampleEnergy, BitpackedSpins, - ReferenceEnergy + ReferenceEnergy, ) from dwave.experimental.lattice_utils.experiment.samplercall import SamplerCall __all__ = ['Experiment'] -class Experiment(): + +class Experiment: """Base class for experiment in LatQA.""" default_parameters = { @@ -60,6 +61,7 @@ class Experiment(): BitpackedSpins(), ReferenceEnergy(), } + def __init__(self, inst: Lattice, sampler: dimod.Sampler, **kwargs): self.inst: Lattice = inst self.sampler: dimod.Sampler = sampler @@ -85,7 +87,6 @@ def __init__(self, inst: Lattice, sampler: dimod.Sampler, **kwargs): self.param[field] = value - def load_results( self, num_iterations: int = 100, @@ -103,7 +104,7 @@ def load_results( 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. + ``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. @@ -117,7 +118,9 @@ def load_results( fnlist = self._get_sorted_results_file_list() if starting_iteration is not None: - fnlist = fnlist[max(starting_iteration, 0):max(starting_iteration + num_iterations, 0)] + fnlist = fnlist[ + max(starting_iteration, 0) : max(starting_iteration + num_iterations, 0) + ] else: fnlist = fnlist[-num_iterations:] @@ -140,7 +143,7 @@ def load_results( results.append({k: data[k] for k in result_fields}) if mod is not None: - return results[:mod * (len(results) // mod)] + return results[: mod * (len(results) // mod)] return results @@ -261,8 +264,7 @@ def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[st sample_set = {} for iemb, sample_array in enumerate(sample_arrays): sample_set[iemb] = dimod.SampleSet.from_samples_bqm( - sample_array, - call.nominal_bqms[iemb] + sample_array, call.nominal_bqms[iemb] ) results = {} @@ -391,15 +393,17 @@ def _get_parameter_pathstring(self) -> str: 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() - ]) + 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(seed=self.run_index) + nominal_bqm = self.inst.make_nominal_bqm() if not hasattr(self.inst, "embedding_list"): return [nominal_bqm] @@ -521,8 +525,7 @@ def _make_sampler_params(self, **kwargs) -> dict[str, Any]: elif "reinitialize_state" in self.param and ret["anneal_schedule"][0][1] == 1: # Set to None, meaning that it will be randomized. ret["initial_state"] = { - qubit: np.random.randint(2) * 2 - 1 - for qubit in self.inst.embedding_list.ravel() + qubit: np.random.randint(2) * 2 - 1 for qubit in self.inst.embedding_list.ravel() } return ret @@ -644,13 +647,11 @@ def _update_coupler_shim( # Damp the couplers (push toward default value) if "coupler_damp" in self.param and self.param["coupler_damp"] > 0: - excess = ( + excess = relative_coupler_strength[:, bin_edges] - np.mean( 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"] + np.multiply(coupler_signs[bin_edges], excess) * self.param["coupler_damp"] ) # New truncation method... previous is buggy when we mix signs of nominal values. @@ -663,9 +664,7 @@ def _update_coupler_shim( > 1 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - 0.99999 - / nominal_values[bin_edges[violators]] - / energy_scale + 0.99999 / nominal_values[bin_edges[violators]] / energy_scale ) violators = ( @@ -675,9 +674,7 @@ def _update_coupler_shim( < -2 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - -1.99999 - / nominal_values[bin_edges[violators]] - / energy_scale + -1.99999 / nominal_values[bin_edges[violators]] / energy_scale ) # Renormalize each orbit after truncation @@ -702,9 +699,7 @@ def _update_coupler_shim( > 1 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - 0.99999 - / nominal_values[bin_edges[violators]] - / energy_scale + 0.99999 / nominal_values[bin_edges[violators]] / energy_scale ) violators = ( @@ -714,9 +709,7 @@ def _update_coupler_shim( < -2 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - -1.99999 - / nominal_values[bin_edges[violators]] - / energy_scale + -1.99999 / nominal_values[bin_edges[violators]] / energy_scale ) Q = nominal_values * relative_coupler_strength * energy_scale @@ -726,7 +719,7 @@ def _update_coupler_shim( 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.""" diff --git a/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py index 7e306f5..6e67cf6 100644 --- a/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py +++ b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py @@ -19,11 +19,12 @@ CouplerFrustration, SampleEnergy, BitpackedSpins, - ReferenceEnergy + ReferenceEnergy, ) __all__ = ['FastAnnealExperiment'] + class FastAnnealExperiment(Experiment): # Set default parameters default_parameters = { diff --git a/dwave/experimental/lattice_utils/experiment/samplercall.py b/dwave/experimental/lattice_utils/experiment/samplercall.py index 5fb430d..458ba91 100644 --- a/dwave/experimental/lattice_utils/experiment/samplercall.py +++ b/dwave/experimental/lattice_utils/experiment/samplercall.py @@ -18,6 +18,7 @@ __all__ = ['SamplerCall'] + @dataclass class SamplerCall: """Data class for managing asynchronous sampler calls.""" diff --git a/dwave/experimental/lattice_utils/lattice/__init__.py b/dwave/experimental/lattice_utils/lattice/__init__.py index 18759d2..22a583d 100644 --- a/dwave/experimental/lattice_utils/lattice/__init__.py +++ b/dwave/experimental/lattice_utils/lattice/__init__.py @@ -16,3 +16,5 @@ 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 index f7a5cb0..246cad0 100644 --- a/dwave/experimental/lattice_utils/lattice/chain.py +++ b/dwave/experimental/lattice_utils/lattice/chain.py @@ -13,21 +13,65 @@ # 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. - def __init__(self, **kwargs): - periodic: tuple[bool, ...] = kwargs.pop("periodic", (True,)) - self.geometry_name: str = "Chain" - self.num_spins = kwargs["dimensions"][0] - super().__init__(periodic=periodic, **kwargs) + 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.""" + """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) diff --git a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py index 13d8e21..1e4cffa 100644 --- a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -12,22 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""What do we want here? - -- A class that handles embedded models. -- Perhaps called FixedEmbeddingModel -- Should have a function called get_chain_connections, which takes an edge and -returns a list or tuple of adjacent indices. For example, in the 3D dimer class, -x,y, and z-couplings would return -((1,1)) -((0,0)) -((0,1),(1,0)) -respectively. -A chain coupling (edge=(u,v) where u==v) can, in this case, return ((0,1)). -""" + from itertools import combinations, product from numbers import Integral -from typing import Any from collections.abc import Iterator, Hashable import dimod @@ -38,17 +25,37 @@ __all__ = ['EmbeddedLattice'] + class EmbeddedLattice(Lattice): - """Specifics should depend on the embedding. Should this have a logical - model as an attribute? Let's try.""" + """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_class: Lattice, - logical_lattice_kwargs: dict[str, Any], + logical_lattice: Lattice, chain_nodes: dict[int, tuple[int, Integral]], **kwargs, ): - self.logical_lattice: Lattice = logical_lattice_class(**logical_lattice_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"): @@ -56,8 +63,24 @@ def __init__( kwargs.setdefault("periodic", self.logical_lattice.periodic) super().__init__(**kwargs) - def get_chain_connectivity(self, u, v=None): - """Should also work for chains! These can be thought of as self-loops.""" + 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)) @@ -66,7 +89,12 @@ def get_chain_connectivity(self, u, v=None): 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.""" + """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. @@ -82,7 +110,15 @@ def generate_edges(self) -> Iterator[tuple[Hashable, Hashable]]: yield u_chain[edge[0]], v_chain[edge[1]] def make_nominal_bqm(self, **kwargs) -> dimod.BQM: - """Construct and embed the nominal 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) @@ -90,7 +126,16 @@ def make_nominal_bqm(self, **kwargs) -> dimod.BQM: 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.""" + """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())) @@ -124,7 +169,15 @@ def embed_bqm(self, logical_bqm: dimod.BQM) -> dimod.BQM: return embedded_bqm def unembed_bqm(self, embedded_bqm: dimod.BQM) -> dimod.BQM: - """Unembed an embedded BQM back onto the logical variable space.""" + """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) @@ -138,13 +191,21 @@ def unembed_bqm(self, embedded_bqm: dimod.BQM) -> dimod.BQM: for u, v in embedded_bqm.quadratic: if which_spin[u] != which_spin[v]: - bias_uv = embedded_bqm.quadratic[u, 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.""" + """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( @@ -158,7 +219,14 @@ def unembed_sampleset(self, sampleset: dimod.SampleSet) -> dimod.SampleSet: 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.""" + """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] @@ -166,7 +234,14 @@ def embed_sample(self, sample: NDArray) -> NDArray: return ret def unembed_sample(self, sample: NDArray) -> NDArray: - """Unembed a physical sample using majority vote with random tie-breaking.""" + """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) diff --git a/dwave/experimental/lattice_utils/lattice/lattice.py b/dwave/experimental/lattice_utils/lattice/lattice.py index a3f5f89..af59079 100644 --- a/dwave/experimental/lattice_utils/lattice/lattice.py +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -14,8 +14,9 @@ import os from pathlib import Path -from collections.abc import Hashable -from numbers import Integral +from collections.abc import Iterator, Hashable +from abc import ABC, abstractmethod +import warnings import dimod from minorminer.utils.parallel_embeddings import find_multiple_embeddings @@ -28,30 +29,62 @@ __all__ = ['Lattice'] -class Lattice(): - """Base class for instances in LatQA.""" - def __init__(self, **kwargs): +class Lattice(ABC): + """An abstract base class for representing lattice geometries used in lattice-utils experiments. - self.dimensions: tuple[int, ...] = kwargs["dimensions"] - self.lattice_data_root: Path = kwargs.get("lattice_data_root", Path.cwd() / "lattice_data") + Subclasses are resonsible for defining the lattice geometry itself. In particular, + a subclass must: - self.periodic: tuple[bool, ...] = kwargs.get("periodic", [False] * len(self.dimensions)) + - 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 = kwargs.get("orbit_type", "singleton") - self.initialize_orbits(kwargs.get("qubit_orbits"), kwargs.get("coupler_orbits")) + 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, - data_root: str | Path | None = None, max_number_of_embeddings: int | None = None, min_number_of_embeddings: int = 1, exclude_qubits: list = [], @@ -64,7 +97,6 @@ def embed_lattice( 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. - data_root: Root directory for loading and saving embedding data. 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 @@ -76,33 +108,26 @@ def embed_lattice( if try_to_load: try: - self._load_embeddings(sampler, data_root) - filename = self._make_filename( - "embedding", - data_root=data_root, - sampler=sampler, - ) - print(f"Loaded embedding from file {filename}") + self._load_embeddings(sampler) return except FileNotFoundError: - pass + 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} + embedder_kwargs={'timeout': timeout}, ) if not embedding_dicts: - raise ValueError("No Embeddings Found") + 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, data_root=data_root) - - return + self._save_embeddings(sampler, embeddings) - def make_nominal_bqm(self, **kwargs) -> dimod.BQM: + #def make_nominal_bqm(self, **kwargs) -> dimod.BQM: + def make_nominal_bqm(self) -> dimod.BQM: """Construct a default nominal BQM coupling strength values set to +1. Args: @@ -135,21 +160,20 @@ def initialize_orbits( orbit and all chain couplers in another. -``standard``: Load previously computed automorphism-based orbits, or compute them and save them if unavailable. - -``explicit``: use the orbit assignments provided via ``qubit_orbits`` + -``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 + 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"): - if hasattr(self.logical_lattice, "logical_lattice"): - raise NotImplementedError # Nested embedded lattices not supported. 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) @@ -163,7 +187,7 @@ def initialize_orbits( try: self._load_orbits() except FileNotFoundError: - print('Calculating orbits...') + # calculating orbits bqm = self.make_nominal_bqm() self.qubit_orbits, self.coupler_orbits = get_orbits(bqm, self.edge_list) self._save_orbits() @@ -174,19 +198,25 @@ def initialize_orbits( elif self.orbit_type == "explicit": if qubit_orbits is not None and coupler_orbits is not None: - assert len(qubit_orbits) == self.num_spins - assert len(coupler_orbits) == self.num_edges + 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}.' \ + f'Unknown orbit type {self.orbit_type}. ' 'Must be "global", "standard", "singleton", or "explicit".' ) def _get_path( self, - root: Path | None, kind: str, sampler_name: str | None = None, extra_subdir: str | Path | None = None, @@ -198,66 +228,49 @@ def _get_path( class_subdir = Path(self.geometry_name) if extra_subdir is not None: class_subdir = class_subdir / extra_subdir - if root is None: - root = Path(__file__).parent.parent / "data" - if sampler_name is None: - path = Path(root) / kind / class_subdir / self._get_size_pathstring() - else: - path = Path(root) / kind / class_subdir / sampler_name / self._get_size_pathstring() + base_dir = self.data_root / "lattice_data" / kind / class_subdir + if sampler_name is not None: + base_dir = base_dir / sampler_name - return path.with_suffix(".txt") + filename = f"{self._get_size_pathstring()}.txt" + return base_dir / filename - def _make_filename( - self, - kind: str, - sampler: dimod.Sampler | None = None, - data_root: str | Path | None = None, - ) -> Path: + def _make_filename(self, kind: str, sampler: dimod.Sampler | None = None) -> Path: """Construct a data filename for the specified sampler and data type.""" - if data_root is None: - data_root = self.lattice_data_root if sampler is None: - return self._get_path(data_root, kind) + return self._get_path(kind) if type(sampler).__name__ == "MockDWaveSampler": - return self._get_path(data_root, kind, sampler_name="MockDWaveSampler") - return self._get_path(data_root, kind, sampler_name=sampler.solver.name) + 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, - data_root: str | Path | None = None, - ) -> None: + def _save_embeddings(self, sampler: dimod.Sampler, embeddings: NDArray) -> None: """Save embedding data to disk.""" - cache_filename = self._make_filename("embedding", sampler=sampler, data_root=data_root) + cache_filename = self._make_filename("embedding", sampler=sampler) os.makedirs(cache_filename.parent, exist_ok=True) np.savetxt(cache_filename, embeddings, fmt="%d") - print(f"Saved {len(embeddings)} embeddings to file {cache_filename}") - def _load_embeddings(self, sampler: str, data_root: str | Path | None = None, **kwargs) -> None: + def _load_embeddings(self, sampler: str) -> None: """Load embedding data.""" - filename = self._make_filename("embedding", sampler=sampler, data_root=data_root) + filename = self._make_filename("embedding", sampler=sampler) self.embedding_list = np.atleast_2d(np.loadtxt(filename, dtype=int)) - def _save_orbits(self, data_root: str | Path | None = None) -> None: + def _save_orbits(self) -> None: """Save qubit and coupler orbits to disk.""" - cache_filename = self._make_filename("orbits", data_root=data_root) + 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") - print(f"Saved orbits to folder {cache_dir}") - def _load_orbits(self, data_root: str | Path | None = None, **kwargs) -> None: + def _load_orbits(self) -> None: """Load qubit and coupler orbits.""" - cache_filename = self._make_filename("orbits", data_root=data_root) + 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) - print(f'Loaded orbits from {cache_dir}') def _get_instance_pathstring(self) -> str: """Construct an instance-specific pathstring. diff --git a/dwave/experimental/lattice_utils/lattice/optimize.py b/dwave/experimental/lattice_utils/lattice/optimize.py index fe4bbbe..1a91f38 100644 --- a/dwave/experimental/lattice_utils/lattice/optimize.py +++ b/dwave/experimental/lattice_utils/lattice/optimize.py @@ -13,7 +13,6 @@ # limitations under the License. from __future__ import annotations -import time from typing import Any from dwave.samplers import SimulatedAnnealingSampler @@ -23,21 +22,42 @@ __all__ = ['optimize', 'optimize_increasing_sa_sweeps'] + def optimize( lattice: Lattice, bqm: dimod.BQM, sa_kwargs: dict[str, Any] | None = None, ) -> tuple[float, NDArray, str]: - """Should take the lattice, from which we can infer the appropriate action.""" + """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 + lattice.logical_lattice, lattice.unembed_bqm(bqm), sa_kwargs=sa_kwargs ) embedded_sample = lattice.embed_sample(logical_sample) embedded_energy = bqm.energy(embedded_sample) @@ -47,14 +67,29 @@ def optimize( # 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]: - """Run SA with exponentially increasing sweep counts until no improvement is achieved.""" - start = time.time() + """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: @@ -63,12 +98,10 @@ def optimize_increasing_sa_sweeps( num_reads = sa_kwargs.get("num_reads", 256) while True: - print(f"Running SA with {num_sweeps} sweeps;", end=" ") 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] - print(f"best energy is {best_energy}. ") if best_energy < reference_energy: reference_energy = best_energy @@ -79,6 +112,4 @@ def optimize_increasing_sa_sweeps( else: break - end = time.time() - print(f"took {end - start:2f}s") 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 index e3ad017..30f6b80 100644 --- a/dwave/experimental/lattice_utils/lattice/orbits.py +++ b/dwave/experimental/lattice_utils/lattice/orbits.py @@ -29,11 +29,16 @@ '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()} @@ -41,34 +46,28 @@ def reindex(mapping: dict[Hashable, int]) -> dict[Hashable, int]: 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) + 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. - """ - # Nodes and edges added in a seemingly ugly way in order to get the order right. + """ ret = dimod.BinaryQuadraticModel(vartype="SPIN") for var in bqm.variables: ret.add_variable(f"p{var}", bqm.linear[var]) - for var in bqm.variables: 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)]) - for u, v in bqm.quadratic: ret.add_quadratic(f"m{u}", f"m{v}", bqm.quadratic[(u, v)]) - for u, v in bqm.quadratic: ret.add_quadratic(f"p{u}", f"m{v}", -bqm.quadratic[(u, v)]) - for u, v in bqm.quadratic: ret.add_quadratic(f"m{u}", f"p{v}", -bqm.quadratic[(u, v)]) return ret @@ -94,7 +93,7 @@ def get_bqm_orbits( 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. + edge to an integer orbit label. """ # The function first adds auxiliary elements to a BQM graph = nx.Graph() @@ -170,20 +169,24 @@ def get_unsigned_bqm_orbits( """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. We also need to keep track of self-symmetric pairs - of spins. + 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: + 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. + each coupling to its orbit index. """ - # Combine coupler orbits so that O(p1p2)=O(m1m2) and O(p1m2)=O(m1p2) + 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}")], @@ -191,20 +194,18 @@ def get_unsigned_bqm_orbits( ) signed_coupler_orbits[(f"m{u}", f"m{v}")] = signed_coupler_orbits[(f"p{u}", f"p{v}")] - signed_coupler_orbits[(f"m{v}", f"p{u}")] = min( - signed_coupler_orbits[(f"m{v}", f"p{u}")], + 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"m{v}", f"p{u}")] + 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}")] - coupler_orbits = {} - for u, v in bqm.quadratic: - coupler_orbits[(u, v)] = signed_coupler_orbits[(f"p{u}", f"p{v}")] - return reindex(qubit_orbits), reindex(coupler_orbits) @@ -214,12 +215,12 @@ def get_orbits(bqm: dimod.BQM, edge_list: list[int, int]) -> tuple[NDArray, NDAr 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)`` mappig each entry of ``edge_list`` + 1-D array of length ``len(edge_list)`` mapping each entry of ``edge_list`` to an orbit index. """ signed_bqm = make_signed_bqm(bqm) diff --git a/dwave/experimental/lattice_utils/lattice/triangular.py b/dwave/experimental/lattice_utils/lattice/triangular.py index a462547..fe18ad2 100644 --- a/dwave/experimental/lattice_utils/lattice/triangular.py +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -13,6 +13,7 @@ # limitations under the License. from collections.abc import Iterator, Hashable +from pathlib import Path import networkx as nx import numpy as np @@ -24,31 +25,89 @@ __all__ = ['Triangular', 'DimerizedTriangular'] -# For triangular, really for triangular AFM. Will include explicit dimerized -# version in square lattice. + class Triangular(Lattice): - """FM or AFM, like Chain.""" + """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, **kwargs): - periodic = kwargs.pop("periodic", (True, False)) + 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 = kwargs.pop("halve_boundary_couplers", False) - self.num_spins = kwargs["dimensions"][0] * kwargs["dimensions"][1] + 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__(periodic=periodic, **kwargs) - assert self.periodic[0] is False or self.dimensions[0] % 3 == 0 - assert self.periodic[1] is False or self.dimensions[1] % 3 == 0 + 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]: - """Get y,x coordinates, traversing y first.""" - Ly = self.dimensions[0] - return node % Ly, node // Ly + """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. - def make_nominal_bqm(self, **kwargs) -> dimod.BQM: - """Accommodate the possibility of halving boundary couplers.""" + Returns: + A dimod.BQM representing the nominal triangular lattice. + """ graph = self._make_networkx_graph() bqm = dimod.BQM(vartype="SPIN") @@ -63,13 +122,17 @@ def make_nominal_bqm(self, **kwargs) -> dimod.BQM: return bqm def generate_edges(self) -> Iterator[tuple[int, int]]: - """Yield edges for the 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 + """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 + length_y, length_x = self.dimensions graph = nx.Graph() for x in range(length_x): @@ -110,22 +173,75 @@ def generate_edges(self) -> Iterator[tuple[int, int]]: 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, + ) - def __init__(self, logical_lattice_class: Lattice = Triangular, **kwargs): - chain_nodes = { - v: (v, v + np.prod(kwargs["dimensions"])) - for v in range(np.prod(kwargs["dimensions"])) - } - self.geometry_name: str = "DimerizedTriangular" - self.num_spins = 2 * int(np.prod(kwargs["dimensions"])) - logical_lattice_kwargs = kwargs.copy() - logical_lattice_kwargs.update({"ignore_embedding": True}) super().__init__( - logical_lattice_class=logical_lattice_class, - logical_lattice_kwargs=logical_lattice_kwargs, + logical_lattice=logical_lattice, chain_nodes=chain_nodes, - **kwargs + 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 @@ -133,8 +249,20 @@ def get_chain_connectivity( self, u: Hashable, v: Hashable | None = None, - ) -> Iterator[tuple[int, int]]: - """Should also work for chains! These can be thought of as self-loops.""" + ) -> 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. @@ -151,7 +279,7 @@ def get_chain_connectivity( 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 ux == vx - 1 or vx == 0: # (ux == self.dimensions[1] - 1 and self.periodic[1]): if uy == vy: return ((1, 0),) return ((0, 1),) diff --git a/dwave/experimental/lattice_utils/observable/kinks.py b/dwave/experimental/lattice_utils/observable/kinks.py index c2a952b..a9583e2 100644 --- a/dwave/experimental/lattice_utils/observable/kinks.py +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -22,6 +22,7 @@ __all__ = ['KinkKinkCorrelator'] + class KinkKinkCorrelator(Observable): """For 1D chains.""" diff --git a/dwave/experimental/lattice_utils/observable/observable.py b/dwave/experimental/lattice_utils/observable/observable.py index 6d5eaed..e7380d8 100644 --- a/dwave/experimental/lattice_utils/observable/observable.py +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -33,6 +33,7 @@ 'ReferenceEnergy', ] + class Observable(ABC): """The observable class does not take any parameters. Its primary functionality is through the required 'evaluate' method, which requires @@ -40,6 +41,7 @@ class Observable(ABC): which provides the samples on which we compute the observable. Output is a numpy array of arbitrary type (usually float). """ + def __init__(self): self.name: str = type(self).__name__ @@ -55,6 +57,7 @@ def evaluate( class QubitMagnetization(Observable): """Compute the mean magnetization of each qubit.""" + def evaluate( self, experiment: Experiment, @@ -67,6 +70,7 @@ def evaluate( class CouplerCorrelation(Observable): """Compute pairwise spin correlations for each coupler.""" + def evaluate( self, experiment: Experiment, @@ -85,6 +89,7 @@ def evaluate( class CouplerFrustration(Observable): """Compute the mean coupler frustration for each edge.""" + def evaluate( self, experiment: Experiment, @@ -98,10 +103,9 @@ def evaluate( # Surprisingly, it's faster to multiply the whole matrix. 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"]) - ) + 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 @@ -111,6 +115,7 @@ class SampleEnergy(Observable): Energies exclude the magnitude of ``energy_scale`` but include its sign. """ + def evaluate( self, experiment: Experiment, @@ -122,6 +127,7 @@ def evaluate( class BitpackedSpins(Observable): """Return bitpacked spins and a tuple of the array size.""" + def evaluate( self, experiment: Experiment, @@ -140,7 +146,9 @@ def evaluate( class ReferenceEnergy(Observable): """Return a cached reference energy, computing it and saving it if needed.""" - def evaluate(self, + + def evaluate( + self, experiment: Experiment, bqm: dimod.BQM, sample_set: dimod.SampleSet, @@ -219,7 +227,7 @@ def get_reference_energy_path( 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 needs to be fixed if you have something not in the instance pathstring that needs to be taken into account, for example if the ground-state energies depend on the chip. @@ -239,11 +247,16 @@ def get_reference_energy_path( } if root is None: - root = experiment_data_dict["inst"].lattice_data_root + root = experiment_data_dict["inst"].data_root else: root = Path(root) - path = root / "reference_energies" / experiment_data_dict["inst"]._get_instance_pathstring() + 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( diff --git a/dwave/experimental/lattice_utils/observable/triangular.py b/dwave/experimental/lattice_utils/observable/triangular.py index 4d12bb6..89ef7b2 100644 --- a/dwave/experimental/lattice_utils/observable/triangular.py +++ b/dwave/experimental/lattice_utils/observable/triangular.py @@ -24,8 +24,10 @@ __all__ = ['TriangularOP'] + class TriangularOP(Observable): """For triangular lattices. Unembeds if possible.""" + def evaluate( self, experiment: Experiment, diff --git a/dwave/experimental/lattice_utils/utils.py b/dwave/experimental/lattice_utils/utils.py index 754d899..94341ab 100644 --- a/dwave/experimental/lattice_utils/utils.py +++ b/dwave/experimental/lattice_utils/utils.py @@ -17,14 +17,33 @@ 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, - seed: int | None = None, skipnan: bool = True, ) -> list[float]: - """Compute bootstrap estimates of a statistic.""" + """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)] @@ -33,7 +52,7 @@ def bootstrap( output = [] if len(array) > 0: - for inds in generate_bootstrap_indices(array.size, repetitions, seed=seed): + for inds in generate_bootstrap_indices(array.size, repetitions, rng): output.append(bootstrap_function(array[inds])) return output @@ -42,17 +61,46 @@ def bootstrap( def generate_bootstrap_indices( size: int, repetitions: int, - seed: int | None = None, + rng: np.random.Generator, ) -> Iterator[NDArray]: - """Generate resampled indices.""" - np.random.seed(seed) + """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 = np.random.choice(range(size), replace=True, size=size) + inds = rng.choice(range(size), replace=True, size=size) yield inds def confidence_interval(array: NDArray, width: float = 0.95) -> tuple[float, float, float]: - """Ravel and take the quantiles; return median and error bar lengths.""" + """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 diff --git a/examples/example_1D_Ising_chain.py b/examples/example_1D_Ising_chain.py index b42cd8f..e433ced 100644 --- a/examples/example_1D_Ising_chain.py +++ b/examples/example_1D_Ising_chain.py @@ -25,6 +25,7 @@ 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 = {} @@ -38,6 +39,7 @@ ] NUM_SPINS = 256 +#NUM_SPINS = 8 # Two energy scales: one strong coupling and one weak coupling. ENERGY_SCALES = (-1.8, 0.1) @@ -51,13 +53,15 @@ # 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,), - sampler=sampler, orbit_type="standard", ) @@ -134,7 +138,7 @@ 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) + x = np.linspace(0, 2 * np.pi, 400) plt.tight_layout() plt.subplots_adjust(hspace=0.35, wspace=0.3, top=0.9) @@ -153,9 +157,9 @@ label="theory", ) - #x = ANNEAL_TIMES M = np.asarray(frust) - bs = np.asarray([bootstrap(_, bootstrap_function=np.nanmedian, seed=None) for _ in M]) + + bs = np.asarray([bootstrap(_, rng, bootstrap_function=np.nanmedian) for _ in M]) ci = np.asarray([confidence_interval(_) for _ in bs]) errorbar_handle = ax.errorbar( @@ -276,14 +280,11 @@ x = ANNEAL_TIMES M = kd_dict[sampler.solver.name, energy_scale] - bs = np.asarray([bootstrap(m, bootstrap_function=np.nanmedian, seed=None) for m in M]) + 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 + ANNEAL_TIMES, ci[:, 0], yerr=[ci[:, 1], ci[:, 2]], **errorbar_style ) ax2[isampler].plot( x, @@ -303,7 +304,7 @@ 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" +title = f"1D chain kink-kink correlator, {dims}, {time_ns:.1f} ns" fig3.suptitle(title, fontsize=16) for isampler, sampler in enumerate(samplers): @@ -323,14 +324,11 @@ x = np.arange(NUM_SPINS) * kd M = magnetization - bs = np.asarray([bootstrap(m, bootstrap_function=np.nanmedian, seed=None) for m in M]) + 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 + x, ci[:, 0], yerr=[ci[:, 1], ci[:, 2]], **errorbar_style ) ax3[isampler].plot( x, @@ -338,7 +336,7 @@ marker='o', linestyle='', color=errorbar_handle[0]._color, - markerfacecolor=np.array(to_rgb(errorbar_handle[0]._color)) / 2 + 0.5 + markerfacecolor=np.array(to_rgb(errorbar_handle[0]._color)) / 2 + 0.5, ) filename = title diff --git a/examples/example_2D_geometric_frustration.py b/examples/example_2D_geometric_frustration.py index 6dac014..2a68546 100644 --- a/examples/example_2D_geometric_frustration.py +++ b/examples/example_2D_geometric_frustration.py @@ -26,12 +26,13 @@ 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 = {} -# Just an Advantage2 prototype. -sampler = DWaveSampler(solver="Advantage2_system3.1") +sampler = DWaveSampler(solver="Advantage2_system1") ANNEAL_TIMES = np.round(0.005 * np.logspace(0, 2, 17), 6) @@ -40,10 +41,12 @@ # 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), - sampler=sampler, orbit_type="explicit", halve_boundary_couplers=True, chain_strength=2, @@ -114,12 +117,14 @@ 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])}, " \ +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) +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) @@ -127,7 +132,7 @@ ax.loglog() M = np.asarray(opmag) -bs = np.asarray([bootstrap(m, bootstrap_function=np.nanmedian, seed=None) for m in M]) +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) @@ -211,12 +216,12 @@ # 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"]}' +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, bootstrap_function=np.nanmedian, seed=None) for m in M[:, :5]]) +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, @@ -234,7 +239,7 @@ label="first 5 iterations of shim", **point_style, ) -bs = np.asarray([bootstrap(m, bootstrap_function=np.nanmedian, seed=None) for m in M[:, -5:]]) +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, diff --git a/tests/test_lattice_utils.py b/tests/test_lattice_utils.py index b43504d..7a7df43 100644 --- a/tests/test_lattice_utils.py +++ b/tests/test_lattice_utils.py @@ -82,6 +82,7 @@ def _make_sync_sampler(n_cols=128, solver_name="TestSolver"): 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 @@ -89,21 +90,17 @@ def done(self): def samples(self, sorted_by=None): return np.ones((10, n_cols), dtype=float) - s = mock.MagicMock() - type(s).__name__ = "DWaveSampler" - s.solver.name = solver_name - s.nodelist = list(range(n_cols)) - s.properties = {"num_qubits": n_cols} - s.sample.return_value = _Response() - return s + 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 + 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() @@ -117,12 +114,11 @@ def _make_mock_experiment( def _make_embedded_chain(chain_nodes): return EmbeddedLattice( - logical_lattice_class=Chain, - logical_lattice_kwargs={ - "dimensions": (len(chain_nodes),), - "periodic": (False,), - "ignore_embedding": True, - }, + 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,), @@ -131,26 +127,25 @@ def _make_embedded_chain(chain_nodes): class TestUtils(unittest.TestCase): def test_bootstrap_all_nan_skipnan(self): - result = bootstrap(np.array([np.nan, np.nan]), repetitions=5, skipnan=True) + 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): - result = bootstrap(np.array([1.0, 2.0, np.nan]), repetitions=5, skipnan=False) + 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): - result = bootstrap(np.arange(20), repetitions=10, bootstrap_function=np.mean, seed=0) + 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_bootstrap_seed_reproducibility(self): - r1 = bootstrap(np.arange(10), repetitions=20, seed=123) - r2 = bootstrap(np.arange(10), repetitions=20, seed=123) - np.testing.assert_array_equal(r1, r2) - def test_generate_bootstrap_indices_correct_count(self): - indices = list(generate_bootstrap_indices(10, 5, seed=0)) + 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) @@ -307,8 +302,7 @@ def test_embed_find_and_save(self): "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)) + 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()) @@ -343,6 +337,7 @@ def test_periodicity(self): 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") @@ -396,10 +391,12 @@ def test_unembed_sample_breaks_ties_randomly(self): 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], - ]) + 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", @@ -492,9 +489,7 @@ def test_sample_energy(self): 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] - ) + 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,)) @@ -505,7 +500,7 @@ def test_bitpacked_spins(self): 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) + 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): @@ -943,8 +938,9 @@ 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) + 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}]) @@ -962,8 +958,9 @@ 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) + 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}]) @@ -1023,6 +1020,7 @@ def test_get_shimdata_not_initialized(self): 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,)) From 8122131af860ab201a435c42d6ce28f292cb4313 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 12 May 2026 10:27:17 -0700 Subject: [PATCH 05/13] Remove temporary automorphism file and additional cleanup -removed spin reversal transform functionality -uses automorphism module already in dwave-experimental -uses data classes for experiment configs --- .../lattice_utils/experiment/experiment.py | 283 ++--- .../experiment/fast_anneal_experiment.py | 53 +- .../lattice_utils/experiment/samplercall.py | 1 - .../lattice_utils/lattice/automorphism.py | 1110 ----------------- .../lattice_utils/lattice/lattice.py | 1 - .../lattice_utils/lattice/orbits.py | 2 +- .../lattice_utils/observable/kinks.py | 36 +- .../lattice_utils/observable/observable.py | 69 +- .../lattice_utils/observable/triangular.py | 32 +- examples/example_1D_Ising_chain.py | 19 +- tests/test_lattice_utils.py | 33 - 11 files changed, 214 insertions(+), 1425 deletions(-) delete mode 100644 dwave/experimental/lattice_utils/lattice/automorphism.py diff --git a/dwave/experimental/lattice_utils/experiment/experiment.py b/dwave/experimental/lattice_utils/experiment/experiment.py index 3e1a81e..8cc62f1 100644 --- a/dwave/experimental/lattice_utils/experiment/experiment.py +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -20,6 +20,7 @@ from pathlib import Path from datetime import datetime from typing import Any +from dataclasses import dataclass import dimod import numpy as np @@ -35,72 +36,74 @@ ) from dwave.experimental.lattice_utils.experiment.samplercall import SamplerCall -__all__ = ['Experiment'] +__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 class Experiment: - """Base class for experiment in LatQA.""" - - default_parameters = { - "energy_scale": 1.0, - "num_reads": 100, - "anneal_time": 1.0, - "num_random_instances": 1, - "readout_thermalization": 100, - "flux_bias_shim_step": 0.0, - "coupler_shim_step": 0.0, - "anneal_offset_shim_step": 0.0, - "target_magnetization": 0.0, - } - - observables_to_collect = { - QubitMagnetization(), - CouplerCorrelation(), - CouplerFrustration(), - SampleEnergy(), - BitpackedSpins(), - ReferenceEnergy(), - } - - def __init__(self, inst: Lattice, sampler: dimod.Sampler, **kwargs): - self.inst: Lattice = inst - self.sampler: dimod.Sampler = sampler - self.param: dict = {} - self.already_initialized: bool = False # until evidenced - - # Any forced types that might be required. - self.typedict: dict[str, float] = {"energy_scale": float} - - self.experiment_results_root: Path = Path( - kwargs.get("results_root", Path.cwd() / "results") - ).resolve() - - self.loop_data_files: int = kwargs.get("loop_data_files", 1000000000) - self.max_iterations: int | None = kwargs.get("max_iterations", None) - - # Apply parameters - for field, default in self.default_parameters.items(): - value = kwargs.get(field, default) - - if field in kwargs and field in self.typedict: - value = self.typedict[field](value) - - self.param[field] = value + """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, - starting_iteration: int | None = None, - result_fields: dict[str, Any] | None = None, + start_iteration: int | None = None, + result_fields: list[str] | None = None, quiet: bool = True, ignore_shim: bool = False, - mod: int | None = None, - ) -> dict[str, Any]: + ) -> list[dict[str, Any]]: """Load results from the highest-numbered iterations of the experiment. Args: num_iterations: Maximum number of iterations to load. - starting_iteration: If provided, load results starting from this + 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 @@ -108,19 +111,13 @@ def load_results( 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. - mod: If provided, the returned list is truncated so its length is - divisable by ``mod``. Returns: A list of dictionaries containing the results for each iteration. """ - self._set_run_index() # this is the NEXT run index. - fnlist = self._get_sorted_results_file_list() - if starting_iteration is not None: - fnlist = fnlist[ - max(starting_iteration, 0) : max(starting_iteration + num_iterations, 0) - ] + if start_iteration is not None: + fnlist = fnlist[max(start_iteration, 0) : max(start_iteration + num_iterations, 0)] else: fnlist = fnlist[-num_iterations:] @@ -142,34 +139,23 @@ def load_results( results.append({k: data[k] for k in result_fields}) - if mod is not None: - return results[: mod * (len(results) // mod)] - return results - def apply_param(self, param: dict[str, float]) -> bool: + 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. - Keys correspond to parameter names and values are converted to - the appropriate type if a converted is defined in ``self.typedict``. - - Returns: - A boolean value corresponding to whether or not the experiment has - already completed all required iterations for this parameter - configuration. """ - # Reformat for filename consistency; can be important for loading data. param = self._format_parameter_list([param])[0] for param_name, param_val in param.items(): - if param_name in self.typedict: - self.param[param_name] = self.typedict[param_name](param_val) - else: - self.param[param_name] = param_val + self.param[param_name] = param_val self.data_path = self.experiment_results_root / self._get_relative_data_path() - self.already_initialized = self._set_run_index() + 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``. @@ -217,18 +203,21 @@ def run_iteration(self, parameter_list: list, **kwargs) -> bool: # 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]) - - # Get the results and update the shim - results = self.parse_results(call_dict[index], response_dict[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 @@ -237,8 +226,7 @@ def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[st """Parse a sampler response into per-embedding observable results. Args: - call: Sampler call metadata, cinluding the nominal BQMs and any - applied spin-reversal transform. + call: Sampler call metadata, inluding the nominal BQMs response: Raw sample set returned by the sampler. Returns: @@ -249,18 +237,9 @@ def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[st embedding_list = self.inst.embedding_list myarr = response.samples(sorted_by=None) sample_arrays = [myarr[:, emb].copy() for emb in embedding_list] - - if call.spin_reversal_transform is not None: - for iemb, emb in enumerate(embedding_list): - for iv, v in enumerate(emb): - if call.spin_reversal_transform[v]: - sample_arrays[iemb][:, iv] *= -1 else: sample_arrays = [response.samples(sorted_by=None)[:, np.arange(self.inst.num_spins)]] - if call.spin_reversal_transform is not None: - raise NotImplementedError - sample_set = {} for iemb, sample_array in enumerate(sample_arrays): sample_set[iemb] = dimod.SampleSet.from_samples_bqm( @@ -308,37 +287,30 @@ def _save_results( def _get_sorted_results_file_list(self) -> list[str]: """Return result filenames sorted lexicographically.""" - fnlist = list(self.data_path.glob("iter*")) + fnlist = list(self.data_path.glob("iter*.pkl.lzma")) fnlist.sort() return [str(fn) for fn in fnlist] - def _set_run_index(self) -> bool: - """Set the run index (data file index, iteration number) for the parameterization. - - Uses the name of the most recently modified file in - the folder. Also creates the data path if it doesn't exist. Return value - is boolean, whether the iterations have already been started or not. - """ - # If the path doesn't exist, create it and set the run index to zero. - if self.data_path.exists() is False: - os.makedirs(self.data_path) - self.run_index = 0 - return False + 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*")) + fnlist = list(self.data_path.glob("iter*.pkl.lzma")) if not fnlist: - self.run_index = 0 - return False + return 0, False + + latest_file_iter = max(int(fn.stem.split(".")[0][4:]) for fn in fnlist) + return latest_file_iter + 1, True - # Folder exists and is nonempty. - # Only complication is looping; we set the flag to True if we overflow mod - fnlist.sort(key=lambda x: os.path.getmtime(x)) - latest_file_iter = int(fnlist[-1].stem.split(".")[0][4:]) - self.run_index = np.max([0, np.mod(latest_file_iter + 1, self.loop_data_files)]) - if self.run_index < latest_file_iter: - self.has_looped = 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.") - return True + 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. @@ -368,16 +340,7 @@ def _get_parameter_pathstring(self) -> str: """ energy_scale = self.param["energy_scale"] - if type(self.sampler).__name__ in [ - "SimulatedAnnealingSampler", - "SimulatedQuantumAnnealingSampler", - "RotorSampler", - "ParallelTemperingSampler", - "DiscreteSimulatedBifurcationSampler", - ]: - # Applies to Monte Carlo dynamics, for example - pathstring = f'energyscale{energy_scale:0.3}/nsweeps{self.param["num_sweeps"]:010d}mcs' - elif "anneal_time" in self.param: + 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"]}' @@ -418,10 +381,12 @@ def _build_sampler_call(self) -> None | SamplerCall: sampler_call = SamplerCall(run_index=self.run_index) sampler_call.nominal_bqms = self._make_nominal_bqms() sampler_call.shimdata = self._get_shimdata() - sampler_call.spin_reversal_transform = self._get_spin_reversal_transform() # Here we can find out that we're finished. - if sampler_call.shimdata["total_iterations"] >= self.max_iterations: + 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) @@ -440,8 +405,6 @@ def _format_parameter_list( """ ret = parameter_list.copy() for entry in ret: - if "target_s" in entry: - entry["target_s"] = np.round(entry["target_s"], 4) if "anneal_time" in entry: entry["anneal_time"] = np.round(entry["anneal_time"], 6) if "anneal_schedule" in entry: @@ -501,32 +464,10 @@ def _make_sampler_params(self, **kwargs) -> dict[str, Any]: if "anneal_offsets" in kwargs["shimdata"]: ret["anneal_offsets"] = list(kwargs["shimdata"]["anneal_offsets"]) - if "num_sweeps" in self.param: - ret["num_sweeps"] = self.param["num_sweeps"] - elif "target_s" in self.param and "dwell_time" in self.param: # fast reverse anneal - ret["x_target_s"] = self.param["target_s"] - ret["x_dwell_time"] = self.param["dwell_time"] - ret["anneal_schedule"] = self.param["anneal_schedule"] - elif "anneal_schedule" in self.param: - ret["anneal_schedule"] = self.param["anneal_schedule"] - elif self.param.get("fast_anneal", False): + if self.param.get("fast_anneal", False): ret["fast_anneal"] = True - ret["annealing_time"] = self.param["anneal_time"] - else: - ret["annealing_time"] = self.param["anneal_time"] - - if "reinitialize_state" in self.param and ret["anneal_schedule"][0][1] == 1: - ret["reinitialize_state"] = self.param.get("reinitialize_state", False) - if ret["reinitialize_state"] is None: - ret["reinitialize_state"] = False - if self.param.get("initial_state", None) is not None and ret["anneal_schedule"][0][1] == 1: - ret["initial_state"] = self.param["initial_state"] - elif "reinitialize_state" in self.param and ret["anneal_schedule"][0][1] == 1: - # Set to None, meaning that it will be randomized. - ret["initial_state"] = { - qubit: np.random.randint(2) * 2 - 1 for qubit in self.inst.embedding_list.ravel() - } + ret["annealing_time"] = self.param["anneal_time"] return ret @@ -552,8 +493,7 @@ def _make_initial_shim(self) -> dict[str, Any]: def _get_latest_iteration_filename(self) -> Path: """Return the filename of the most recently completed iteration.""" - file_name = f"iter{np.mod(self.run_index - 1, self.loop_data_files):05d}.pkl.lzma" - return self.data_path / file_name + 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.""" @@ -615,7 +555,6 @@ def _update_coupler_shim( step_size: float | None = None, ) -> None: """Update relative coupler strength based on measured frustration.""" - # Ok, let's use orbits and stuff. orbits = self.inst.coupler_orbits energy_scale = self.param["energy_scale"] relative_coupler_strength = sampler_call.shimdata["relative_coupler_strength"] @@ -630,12 +569,11 @@ def _update_coupler_shim( normalization_basis = np.ones_like(orbits, dtype=bool) # Assume we have multiple embeddings of the same BQM. - assert ( - len(sampler_call.nominal_bqms) == 1 - or sampler_call.nominal_bqms[0] == sampler_call.nominal_bqms[1] - ), "Case for distinct embedded BQMs not implemented yet." + 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 = sampler_call.nominal_bqms[0] + 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): @@ -743,11 +681,6 @@ def _make_bqm(self, sampler_call: SamplerCall) -> dimod.BQM: for iemb, emb in enumerate(self.inst.embedding_list): nominal_bqm = sampler_call.nominal_bqms[iemb].copy() - if sampler_call.spin_reversal_transform is not None: - for iv, v in enumerate(emb): - if sampler_call.spin_reversal_transform[v]: - nominal_bqm.flip_variable(iv) - for v in range(self.inst.num_spins): # Don't touch degree-zero spins. Relevant to partial yield. if nominal_bqm.degree(v) > 0: @@ -762,19 +695,3 @@ def _make_bqm(self, sampler_call: SamplerCall) -> dimod.BQM: bqm.add_quadratic(emb[edge[0]], emb[edge[1]], bias) return bqm - - def _get_spin_reversal_transform(self) -> dict[int, bool] | None: - """Generate a spin-reversal transform for the current sampler.""" - if self.param.get("spin_reversal_transform", False): - seed = self.param.get("spin_reversal_transform_seed", None) - rng = np.random.default_rng(seed) - flips = rng.random(len(self.sampler.nodelist)) > 0.5 - - spin_reversal_transform = {v: False for v in self.sampler.nodelist} - for iv, v in enumerate(self.sampler.nodelist): - if flips[iv]: - spin_reversal_transform[v] = True - - return spin_reversal_transform - - return None diff --git a/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py index 6e67cf6..59ca0b2 100644 --- a/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py +++ b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py @@ -12,45 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dwave.experimental.lattice_utils.experiment import Experiment -from dwave.experimental.lattice_utils.observable import ( - QubitMagnetization, - CouplerCorrelation, - CouplerFrustration, - SampleEnergy, - BitpackedSpins, - ReferenceEnergy, -) +from dataclasses import dataclass -__all__ = ['FastAnnealExperiment'] +from dwave.experimental.lattice_utils.experiment import ExperimentConfig +__all__ = ['FastAnnealExperimentConfig'] -class FastAnnealExperiment(Experiment): - # Set default parameters - default_parameters = { - "energy_scale": 1.0, - "automorph_embeddings": False, - "spin_reversal_transform": False, - "spin_reversal_transform_seed": None, - "num_reads": 100, - "num_random_instances": None, - "readout_thermalization": 100, - "fast_anneal": True, - "anneal_time": 1.0, - "flux_bias_shim_step": 0.0, - "coupler_shim_step": 0.0, - "coupler_damp": 0.0, - "anneal_offset_shim_step": 0.0, - "anneal_offset_damp": 0.0, - "individual_qubit_anneal_offsets": None, - "target_magnetization": 0.0, - "logical_software": False, - } - observables_to_collect = { - QubitMagnetization(), - CouplerCorrelation(), - CouplerFrustration(), - SampleEnergy(), - BitpackedSpins(), - ReferenceEnergy(), - } + +@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 index 458ba91..cbd859b 100644 --- a/dwave/experimental/lattice_utils/experiment/samplercall.py +++ b/dwave/experimental/lattice_utils/experiment/samplercall.py @@ -36,4 +36,3 @@ def __init__( 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 - self.spin_reversal_transform: None = None diff --git a/dwave/experimental/lattice_utils/lattice/automorphism.py b/dwave/experimental/lattice_utils/lattice/automorphism.py deleted file mode 100644 index e22996a..0000000 --- a/dwave/experimental/lattice_utils/lattice/automorphism.py +++ /dev/null @@ -1,1110 +0,0 @@ -# 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 import deque -from collections.abc import Hashable -from dataclasses import dataclass -from enum import Enum, auto -import hashlib -from itertools import chain -import random -from typing import Mapping - -import networkx as nx -import numpy as np -from numpy.typing import NDArray - - -@dataclass -class ComponentInfo: - """Container for per-component data used during automorphism discovery on disjoint graphs.""" - u_vector: list - nodes: NDArray - best_perm: NDArray - -class EnterMode(Enum): - """Controls when the ``_enter()`` function attempts to compose new automorphisms.""" - RECURSE = auto() - RECURSE_ONCE = auto() - NO_RECURSE = auto() - -class SchreierContext: - """This object holds mutable states used throughout the automorphism calculation. - - Args: - graph: A NetworkX Graph object representing the input graph. - num_samples: Number of samples to use for generating new coset representatives - from the existing set. If not provided, all coset representatives are used. - seed: Seed used for reproducibility. Defaults to 42. - """ - def __init__(self, graph: nx.Graph, num_samples: int | None = None, seed: int = 42) -> None: - original_nodes_sorted = sorted(graph.nodes()) - self._index_to_node: dict[int, Hashable] = { - new: old for new, old in enumerate(original_nodes_sorted) - } - self._node_to_index: dict[Hashable, int] = { - old: new for new, old in enumerate(original_nodes_sorted) - } - graph = nx.relabel_nodes(graph, self._node_to_index) # relabel nodes contiguously (0...n-1) - - self._nodes: list[int] = list(graph.nodes()) - self._num_nodes: int = graph.number_of_nodes() - self._graph_edges: list[tuple[int, int]] = list(graph.edges()) - self._neighbours: list[set[int]] = [set(graph.neighbors(i)) for i in range(self._num_nodes)] - self._graph: nx.Graph = graph - - self._num_samples: int | None = num_samples - self._rng: random.Random = random.Random(seed) - - self._leaf_nodes: int = 0 - self._nodes_reached: int = 0 - self._depth: int = 0 - - self._u_map: dict[np.intp, int] = {} - self._u_len: int = 0 - self._u_vector: list = [] - self._u_vector_inv: list[list[NDArray[np.intp]]] = [] - - self._identity: NDArray[np.intp] = np.arange(self._num_nodes, dtype=np.intp) - - self._best_perm: NDArray = np.arange(self._num_nodes) - self._best_perm_exist: bool = False - self._compare_adj: bool = False - self._trace_history: list = [] - - self._in_colors_adj: bytearray = bytearray(self._num_nodes) - self._in_refine_stack: bytearray = bytearray(self._num_nodes) - - self._color_degree: list[int] = [0] * self._num_nodes - self._min_color_degree: list[int] = [0] * self._num_nodes - self._max_color_degree: list[int] = [0] * self._num_nodes - self._active_vertices: list[list[int]] = [[] for _ in range(self._num_nodes)] - - if self._num_nodes <= 65535: - self._color_dtype: np.dtype = np.uint16 - else: - self._color_dtype: np.dtype = np.uint32 - - @property - def leaf_nodes(self) -> int: - """Number of leaf nodes encountered in the search tree.""" - return self._leaf_nodes - - @property - def nodes_reached(self) -> int: - """Total number of nodes reached during traversal of the search tree.""" - return self._nodes_reached - - @property - def index_to_node(self) -> dict[int, Hashable]: - """The mapping from the basis of relabelled nodes (0...n-1) to the original - node labels.""" - return self._index_to_node - - @property - def node_to_index(self) -> dict[Hashable, int]: - """The mapping from the original node labels to the basis of relabelled - nodes (0...n-1).""" - return self._node_to_index - - @property - def u_map(self) -> dict[np.intp, int]: - """Map from coset representative group index to stabilizer index.""" - return self._u_map - - @property - def u_vector(self) -> list[list[NDArray[np.intp]]]: - """Coset representatives grouped by stabilizer index.""" - return self._u_vector - - @property - def num_automorphisms(self) -> int: - """Number of automorphisms implied by u_vector.""" - if self._u_vector: - return int(np.prod([len(u_i) + 1 for u_i in self._u_vector], dtype=object)) - else: - return 1 - - @property - def vertex_orbits(self) -> list[list[int]]: - """Vertex orbits induced by the coset representatives in u_vector and returned - in the basis of relabelled nodes (0...n-1).""" - return vertex_orbits(self._u_vector, self._nodes) - - @property - def vertex_orbits_original_labels(self) -> list[list[Hashable]]: - """Vertex orbits induced by the coset representatives in u_vector and returned - with the original node labels.""" - return vertex_orbits(self._u_vector, self._nodes, index_to_node=self._index_to_node) - - @property - def edge_orbits(self) -> list[list[int]]: - """Edge orbits induced by the coset representatives in u_vector and returned - in the basis of relabelled nodes (0...n-1).""" - return edge_orbits(self._u_vector, self._graph_edges) - - @property - def edge_orbits_original_labels(self) -> list[list[Hashable]]: - """Edge orbits induced by the coset representatives in u_vector and returned - with the original node labels.""" - return edge_orbits(self._u_vector, self._graph_edges, index_to_node=self._index_to_node) - - def _test_composability(self, g: NDArray[np.intp]) -> tuple[int, NDArray[np.intp]]: - """Test if an automorphism is composable from coset representatives. - - Based on Algorithm 6.10 from Kreher, D. L., & Stinson, D. R. (1999). - Combinatorial algorithms: Generation, enumeration, and search. - - Modified to use a mask to skip sifting by identity permutations, which - have no effect. - - Args: - g: A permutation represented as a list of integers in one-line notation. - - Returns: - A tuple (i, g_reduced) where i is the index of the first base position - that could not be sifted. If ``g`` is completely sifted the returned index - equals ``self._num_nodes``. ``g_reduced`` is the permutation obtained after - sifting through all positions up to (but not including) the returned - index. - """ - mask = (g != self._identity) - index = mask.argmax() - next_diff = 0 - - while mask[index]: - next_diff += index - if next_diff not in self._u_map: - return next_diff, g - - for i, h in enumerate(self._u_vector[self._u_map[next_diff]]): - if h[next_diff] == g[next_diff]: - break - else: - return next_diff, g - - g = self._u_vector_inv[self._u_map[next_diff]][i][g] - mask = (g[next_diff:] != self._identity[next_diff:]) - index = mask.argmax() - - return self._num_nodes, g - - def _enter(self, g: NDArray[np.intp], mode: EnterMode = EnterMode.RECURSE) -> None: - """Add automorphism if it can't be composed from coset representatives. - - Based on Algorithm 6.11 from Kreher, D. L., & Stinson, D. R. (1999). - Combinatorial algorithms: Generation, enumeration, and search. - - If an automorphism can't be composed from existing coset representatives - it is added as a new coset representative to u_vector. Depending on the - setting of ``mode``, ``_enter()`` is called recursively to attempt to - compose additional coset representatives from the composition between - the newly-discovered coset representative and existing coset representatives. - - The automorphisms discovered will result in pruning comparable to nauty, - as measured by comparing the total number of search tree nodes visited - for zephyr graphs of various sizes. - - Args: - g: A permutation represented as a list of integers in one-line notation. - mode: Specifies if recursive calls to ``enter()`` are performed to attempt - to compose new automorphisms. The setting ``EnterMode.RECURSE_ONCE`` - results in a single call to ``enter()`` per coset representative where - no further attempts to compose automorphisms occur. - """ - i, g = self._test_composability(g) - if i == self._num_nodes: - return - - if i not in self._u_map: - self._u_map[i] = self._u_len - self._u_len += 1 - self._u_vector.append([]) - self._u_vector_inv.append([]) - - self._u_vector[self._u_map[i]].append(g) - self._u_vector_inv[self._u_map[i]].append(inv(self._num_nodes, g)) - - if mode is EnterMode.NO_RECURSE: - return - - for u_i in self._u_vector: - for h in u_i: - f = mult(g, h) - if mode is EnterMode.RECURSE_ONCE: - self._enter(f, mode=EnterMode.NO_RECURSE) - else: - self._enter(f) - - def _refine( - self, - partition: list[set[int]], - trace: NDArray[np.integer], - color: NDArray[np.integer], - num_colors: int, - individualized_vertex: int | None = None, - ) -> None: - """Perform color refinement on the current partition until an equitable - coloring is reached. - - This procedure implements the 1-dimensional Weisfeiler-Leman (WL) refinement, - following Algorithms 2 and 3 of Berkholz (2016), *Tight lower and upper bounds - for the complexity of canonical color refinement*. - - A refinement stack is initialized with either: - • all color classes (if no vertex has been individualized), or - • the color class of the individualized vertex. - - For each color class popped from the stack, the algorithm computes the - color-degree of every vertex: the number of neighbours it has in the refining - color class. These color-degrees determine how each color class should be - split. If a color class contains vertices with differing color-degrees, it is - partitioned into new color classes, and the smaller subcells are pushed onto - the refinement stack. - - The process continues until no color class can be further refined, yielding an - equitable coloring. - - If a vertex was individualized prior to this refinement step, only the - color class containing that vertex needs to be placed on the refinement - stack initially, since only colors adjacent to that color can be affected. - - For performance reasons, ``num_colors`` is passed as a single-element list - so that updates to the number of colors persist across calls without having - to return anything. - - Args: - partition: The current partition structure, represented as a list of sets of vertices - ordered by color. - trace: A list of the sizes of each partition cell (color class), ordered by color. - color: An array mapping each vertex to its current color. - num_colors: The current number of colors in the partition. - individualized_vertex: The vertex individualized prior to this refinement step, if any. - - Returns: - The new number of colors, the updated trace array, and the updated color array. - """ - neighbours = self._neighbours - color_degree = self._color_degree - min_color_degree = self._min_color_degree - max_color_degree = self._max_color_degree - active_vertices = self._active_vertices - in_refine_stack = self._in_refine_stack - in_colors_adj = self._in_colors_adj - - colors_adj = [] - - if individualized_vertex is None: - refine_stack = list(range(num_colors)) - else: - refine_stack = [color[individualized_vertex]] - num_colors = [num_colors] # mutable container so ``_split_up_color()`` can increment it - - for v in refine_stack: - in_refine_stack[v] = 1 - - while refine_stack: - refinement_color = refine_stack.pop() - in_refine_stack[refinement_color] = 0 - - for v in partition[refinement_color]: - for w in neighbours[v]: - color_degree[w] += 1 - cw = color[w] - if color_degree[w] == 1: - active_vertices[cw].append(w) - if in_colors_adj[cw] == 0: - colors_adj.append(cw) - in_colors_adj[cw] = 1 - if color_degree[w] > max_color_degree[cw]: - max_color_degree[cw] = color_degree[w] - - for c in colors_adj: - if trace[c] != len(active_vertices[c]): - min_color_degree[c] = 0 - else: - min_color_degree[c] = max_color_degree[c] - for v in active_vertices[c]: - if color_degree[v] < min_color_degree[c]: - min_color_degree[c] = color_degree[v] - - colors_to_split = [] - for c in colors_adj: - if min_color_degree[c] < max_color_degree[c]: - colors_to_split.append(c) - - for color_to_split in sorted(colors_to_split): - self._split_up_color( - color_to_split=color_to_split, - partition=partition, - color=color, - trace=trace, - active_vertices=active_vertices, - color_degree=color_degree, - min_degree=min_color_degree[color_to_split], - max_degree=max_color_degree[color_to_split], - refine_stack=refine_stack, - in_refine_stack=in_refine_stack, - num_colors=num_colors, - ) - - ## reset attributes - for c in colors_adj: - for v in active_vertices[c]: - color_degree[v] = 0 - max_color_degree[c] = 0 - active_vertices[c] = [] - in_colors_adj[c] = 0 - colors_adj = [] - - return num_colors[0], trace, color - - def _split_up_color( - self, - *, - color_to_split: int, - partition: list[set[int]], - color: NDArray[np.integer], - trace: NDArray[np.integer], - active_vertices: list[list[int]], - color_degree: list[int], - min_degree: int, - max_degree: int, - refine_stack: list[int], - in_refine_stack: bytearray, - num_colors: list[int], - ) -> None: - """Splits a color class into subcells based on the color-degrees of its vertices. - - Based on algorithm 3 of Berkholz (2016), *Tight lower and upper bounds - for the complexity of canonical color refinement*. - - Given a color class ``color_to_split`` whose vertices exhibit differing - color-degrees with respect to the current refining color, this routine - partitions that class into new color classes. Vertices with the same - color-degree remain together, while vertices with different degrees are - assigned fresh color identifiers. - - The largest resulting subcell retains the original color label, while - all smaller subcells are assigned new colors and pushed onto the - refinement stack (Hopcroft's trick). The partition structure, trace array, - number of colors, and vertex-to-color mapping are updated in place. - - Args: - color_to_split: The color class to be split. - partition: The current partition structure, represented as a list of sets of vertices - ordered by color. - color: An array mapping each vertex to its current color. - trace: A list of the sizes of each partition cell (color class), ordered by color. - active_vertices: Lists of vertices adjacent to the color class being split, - ordered by color. - color_degree: The color-degree of each vertex. - min_degree: Minimum color-degree among vertices in the color class being split. - max_degree: Maximum color-degree among vertices in the color class being split. - refine_stack: Stack of colors scheduled for refinement. - in_refine_stack: Flags indicating which colors are already on the stack. - num_colors: The number of colors, used to determine the next color label to assign to - newly-refined cells. Stored as a single-element list so that updates persist across - calls. - """ - degree_to_new_color = [0] * (max_degree + 1) - num_color_degree = [0] * (max_degree + 1) - num_color_degree[0] = trace[color_to_split] - len(active_vertices[color_to_split]) - - for v in active_vertices[color_to_split]: - num_color_degree[color_degree[v]] += 1 - - largest_subcell_degree = 0 - for i in range(1, max_degree + 1): - if num_color_degree[i] > num_color_degree[largest_subcell_degree]: - largest_subcell_degree = i - - for i in range(max_degree + 1): - if num_color_degree[i] > 0: - if i == min_degree: - degree_to_new_color[i] = color_to_split - if not in_refine_stack[color_to_split] and i != largest_subcell_degree: - refine_stack.append(degree_to_new_color[i]) - in_refine_stack[degree_to_new_color[i]] = 1 - else: - degree_to_new_color[i] = num_colors[0] - partition[num_colors[0]] = set() - if in_refine_stack[color_to_split] or i != largest_subcell_degree: - refine_stack.append(degree_to_new_color[i]) - in_refine_stack[degree_to_new_color[i]] = 1 - num_colors[0] += 1 - - for v in active_vertices[color_to_split]: - new_color = degree_to_new_color[color_degree[v]] - if new_color != color_to_split: - partition[color_to_split] = partition[color_to_split] - {v} # must create new obj - partition[new_color].add(v) - trace[color_to_split] -= 1 - trace[new_color] += 1 - color[v] = new_color - - def _canon( - self, - partition: list[set[int]], - trace: NDArray[np.integer], - color: NDArray[np.integer], - num_colors: int, - individualized_vertex: int | None = None, - ) -> None: - """Generate search tree based on iterative color refinement and vertex - individualization. - - Loosely based on Algorithm 7.9 from Kreher, D. L., & Stinson, D. R. (1999). - Combinatorial algorithms: Generation, enumeration, and search. Additional - data structures are used to efficiently track the number of vertices - belonging to each color, vertex colors, and number of colors. Additionally, - the most recently individualized vertex is tracked and used to perform - color refinement more efficiently. - - Color refinement is performed iteratively on a graph until a discrete - coloring is achieved. If the coloring is not discrete after refinement, - vertices belonging to the same color are individualized, meaning that they - are assigned a new color, often breaking the symmetry of the graph and - allowing a subsequent color refinement step to produce further refinement. - - By default, graph comparisons using adjacency matrices are not performed, as - this becomes a bottleneck for even modestly sized graphs. Instead, the - ``trace`` for each graph is compared, which corresponds to the number of - vertices belonging to each color, ordered by color. This check is orders of - magnitude faster and has been found to have identical pruning capability - for graphs of interest, such as chimera, pegasus, and zephyr graphs, as - well as the disjoint compositions of smaller and simpler graphs as may - be encountered when doing parallel embeddings. - - If a graph has more than one component, comparisons using adjacency matrices are - used. This enables isomorphism detection between components, and in turn - a more efficient approach to generating the full automorphism group, which - may contain many automorphisms between isomorphic components. - - Kreher and Stinson perform comprehensive pruning by changing the base of - the left transversals to coincide with the current permutation order up to the - first non-discrete partition cell, or first split. At the cost of performing - this base change, it allows pruning to be performed by only considering - the left transversal with a stabilizer index equal to the index of the first - split. In practice, changing the base at each node of the search tree - becomes prohibitively expensive even more mostly sized graphs, and instead - the approach taken here is to avoid base changes, but instead to more carefully - evaluate which coset representatives to use for pruning. This is done by - ignoring the automorphisms that do not respect the current partition structure. - - Args: - partition: The current partition structure, represented as a list of - sets of vertices ordered by color. - trace: The number of vertices belonging to each color, ordered by color. - color: A map from each vertex to its color. - num_colors: The number of unique colors, equivalent to the number - of cells in the partition. - individualized_vertex: The most recently individualized vertex. - """ - self._nodes_reached += 1 - self._depth += 1 - - num_colors, trace, color = self._refine( - partition, - trace, - color, - num_colors, - individualized_vertex=individualized_vertex - ) - - if not self._best_perm_exist: - self._trace_history.append(trace.tobytes()) - - # first non-singleton block index - first_split = self._num_nodes - 1 - for i, block in enumerate(partition): - if len(block) > 1: - first_split = i - break - - compare_result = 2 - if self._best_perm_exist: # if a leaf node has been reached previously - - if self._compare_adj: - perm_candidate = list(chain.from_iterable(p for p in partition if p is not None)) - compare_result = self._compare(perm_candidate, first_split) - else: - compare_result = trace.tobytes() == self._trace_history[self._depth - 1] - - if compare_result == 0: - return - - if first_split == self._num_nodes - 1: # leaf node reached - self._leaf_nodes += 1 - - if not self._best_perm_exist: - self._best_perm_exist = True - self._best_perm[:] = list(chain.from_iterable(partition)) - - elif compare_result == 2: - perm_candidate = list(chain.from_iterable(partition)) - self._best_perm[:] = perm_candidate - - elif compare_result == 1: - perm_transformed = np.empty(self._num_nodes, dtype=np.intp) - perm_candidate = list(chain.from_iterable(partition)) - perm_transformed[perm_candidate] = self._best_perm - self._enter(perm_transformed) - - return - - candidates = sorted(partition[first_split]) - remaining_in_block = partition[first_split] - updated_partition = partition - trace[first_split] -= 1 - trace[num_colors] = 1 - - while candidates: - vertex = next(iter(candidates)) - updated_partition[first_split] = remaining_in_block - {vertex} - updated_partition[num_colors] = {vertex} - individualized_partition = list(updated_partition) # copy outer list - color[vertex] = num_colors # updated individualized cell - trace_copy = np.array(trace) - color_copy = np.array(color) - - self._canon( - individualized_partition, - trace_copy, - color_copy, - num_colors + 1, - individualized_vertex=vertex - ) - - color[vertex] = first_split - candidates.remove(vertex) - - # prune the search tree using automorphisms - for stab_index, u_index in self._u_map.items(): - if stab_index > vertex: # these automorphisms map vertex to itself - continue - - for g in self._u_vector[u_index]: - if g[vertex] not in candidates: - continue - - # automorphism must respect current partition structure - for w in candidates: - if color[w] != color[g[w]]: - break - else: - candidates.remove(g[vertex]) - - self._depth -= 1 - - def _compare(self, perm: NDArray[np.intp], first_split: int) -> int: - """Compare canonical adjacency matrix against itself under a partial permutation. - - At the first differing entry, returns whether the partial permutatation has - a greater or lesser value, otherwise it returns that they are equal. - - Based on Algorithm 7.6 from Kreher, D. L., & Stinson, D. R. (1999). - Combinatorial algorithms: Generation, enumeration, and search. - - Args: - perm: The permutation of the adjacency matrix to compare the canonical - adjacency matrix against. - first_split: The index of the first block of the partition containing - more than one vertex, defining the size of the partial permutation - of perm to use. - - Returns: - An integer 0, 1, or 2 depending on whether the partial permutation - perm results in an adjacency matrix which is less than, equal to, or - greater than the canonical adjacency matrix, respectively. - """ - neighbours = self._neighbours - best_perm = self._best_perm - for j in range(1, first_split): - neighbours_best_j = neighbours[best_perm[j]] - neighbours_pi_j = neighbours[perm[j]] - for i in range(j): - bit_best = 1 if best_perm[i] in neighbours_best_j else 0 - bit_pi = 1 if perm[i] in neighbours_pi_j else 0 - if bit_best < bit_pi: - return 0 - if bit_best > bit_pi: - return 2 - return 1 - - def _certificate(self) -> bytes: - """Generate a canonical certificate for a graph. - - Based on the permutation ``self.best_perm`` that minimizes the binary value - of the upper triangular portion of the adjacency matrix of the graph, - as found by comparing leaf nodes of the search tree during the search for - automorphisms. - - Returns: - cert_hash: a hash object of the canonical adjacency bitstring. - """ - cert_hash = hashlib.sha256() - neighbours = self._neighbours - best_perm = self._best_perm - - for j in range(1, self._num_nodes): - neighbours_best_j = neighbours[best_perm[j]] - - for i in range(j): - bit = 1 if best_perm[i] in neighbours_best_j else 0 - cert_hash.update(bytes([bit])) - - return cert_hash.digest() - - def _initial_partition( - self, - graph_coloring=None - ) -> tuple[list[set[int] | None], np.ndarray, np.ndarray, int]: - """Initialize the initial partition for a graph. - - Currently this only supports graphs whose vertices are initially the same - color, but could be expanded in the future to accommodate graphs with a - non-trivial initial vertex coloring. - - Returns: - partition: The initial partition structure, represented as a list of sets of vertices - ordered by color. - trace: A list of the sizes of each partition cell (color class), ordered by color. - color: An array mapping each vertex to its current color. - num_colors: The number of colors in the initial partition. - """ - trace = np.zeros(self._num_nodes, dtype=self._color_dtype) - color = np.zeros(self._num_nodes, dtype=self._color_dtype) - - if graph_coloring is None: - partition = [set(self._nodes)] + [None] * (self._num_nodes - 1) - trace[0] = self._num_nodes - num_colors = 1 - else: - num_colors = max(graph_coloring.values()) + 1 - partition = [set() for _ in range(num_colors)] + [None] * (self._num_nodes - num_colors) - for node_label, node_color in graph_coloring.items(): - node_index = self.node_to_index[node_label] # TODO ensure this works with colors not starting at 0..n-1 - partition[node_color].add(node_index) - - for i, nodes in enumerate(partition[:num_colors]): - color[list(nodes)] = i - trace[i] = len(nodes) - - #num_colors = i + 1 - #partition = partition + [None] * (self._num_nodes - num_colors) - - return partition, trace, color, num_colors - - -def vertex_orbits( - u_vector: list[list[NDArray[np.intp]]], - nodes: list[int], - index_to_node: Mapping[int, int] | None = None, -) -> list[list[int]]: - """Calculate vertex orbits using breadth-first search. - - If ``u_vector`` contains no coset representatives, trivial orbits are returned. - - Args: - u_vector: Coset representatives grouped by stabilizer index. - nodes: List of vertex indices used to return trivial orbits when ``u_vector`` is empty. - index_to_node: An optional dictionary for returning orbits with their original node labels. - - Returns: - A list of orbits, each orbit is a list of vertex indices. - - Example: - >>> import numpy as np - >>> from dwave.experimental.automorphism import vertex_orbits - ... - >>> u_vector = [ - ... [np.array([0, 1, 4, 3, 2, 6, 5, 7])], - ... [np.array([2, 1, 4, 3, 0, 7, 5, 6]), np.array([4, 1, 0, 3, 2, 6, 7, 5])], - ... [np.array([0, 3, 2, 1, 4, 5, 6, 7])], - ... ] - >>> nodes = list(range(8)) - >>> vertex_orbits(u_vector, nodes) - [[0, 2, 4], [1, 3], [5, 6, 7]] - """ - if not u_vector: - return [[x] for x in nodes] - - if not all(isinstance(sublist, list) for sublist in u_vector): - raise ValueError("u_vector must be a list of lists.") - - if isinstance(nodes, np.ndarray): - nodes = nodes.tolist() - - if not isinstance(nodes, list) or not all(isinstance(n, int) for n in nodes): - raise ValueError("nodes must be a list of integers.") - - visited = set() - orbits = [] - num_nodes = len(nodes) - generators = [g for u_vector_i in u_vector for g in u_vector_i] - generators.append(np.arange(num_nodes)) - label = (lambda x: index_to_node[x]) if index_to_node is not None else int - - for v_start in nodes: - if v_start in visited: - continue - - visited.add(v_start) - orb = [label(v_start)] - - q = deque([v_start]) - while q: - v_current = q.popleft() - - for g in generators: - v_current = g[v_current] - if v_current not in visited: - visited.add(v_current) - q.append(v_current) - orb.append(label(v_current)) - orb.sort() - orbits.append(orb) - - orbits.sort() - return orbits - - -def edge_orbits( - u_vector: list[list[NDArray[np.intp]]], - edges: list[tuple[int, int]], - index_to_node: Mapping[int, int] | None = None, -) -> list[list[int]]: - """Calculate edge orbits using breadth-first search. - - Args: - u_vector: Coset representatives grouped by stabilizer index. - edges: List of graph edges as tuples of vertex index pairs. - - Returns: - A list of orbits, each orbit is a list of edges (tuples of vertex index pairs). - - Example: - >>> import numpy as np - >>> from dwave.experimental.automorphism import edge_orbits - ... - >>> u_vector = [ - ... [np.array([0, 1, 4, 3, 2, 6, 5, 7])], - ... [np.array([2, 1, 4, 3, 0, 7, 5, 6]), np.array([4, 1, 0, 3, 2, 6, 7, 5])], - ... [np.array([0, 3, 2, 1, 4, 5, 6, 7])], - ... ] - >>> edges = [ - ... (0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6), - ... (6, 7), (7, 0), (0, 3), (1, 4), (2, 6), (5, 7) - ... ] - >>> orbits = edge_orbits(u_vector, edges) - >>> orbits[0] - [(0, 1), (0, 3), (1, 2), (1, 4), (2, 3), (3, 4)] - >>> orbits[1:] - [[(0, 7), (2, 6), (4, 5)], [(5, 6), (5, 7), (6, 7)]] - """ - if not u_vector: - return [[x] for x in edges] - - if not all(isinstance(sublist, list) for sublist in u_vector): - raise ValueError("u_vector must be a list of lists.") - - if not isinstance(edges, list) or not all(isinstance(e, tuple) for e in edges): - raise TypeError("edges must be a list of tuples") - - visited = set() - orbits = [] - generators = [g for u_vector_i in u_vector for g in u_vector_i] - label = (lambda x: index_to_node[x]) if index_to_node is not None else int - - for u_start, v_start in edges: - e_start = (u_start, v_start) if u_start < v_start else (v_start, u_start) - - if e_start in visited: - continue - - visited.add(e_start) - orb = [tuple(label(x) for x in e_start)] - - q = deque([e_start]) - while q: - u, v = q.popleft() - for g in generators: - e_current = (g[u], g[v]) if g[u] < g[v] else (g[v], g[u]) - - if e_current not in visited: - visited.add(e_current) - q.append(e_current) - orb.append(tuple(label(x) for x in e_current)) - - orb.sort() - orbits.append(orb) - - orbits.sort() - return orbits - - -def sample_automorphisms( - u_vector: list[list[NDArray[np.intp]]], - num_samples: int = 1, - seed: int | None = None, -) -> list[NDArray[np.intp]]: - """Uniformly sample automorphisms from the Schreier-Sims representation. - - Randomly samples one coset representative from each non-trivial left - transversal and takes the product, guaranteeing uniform sampling. The - automorphisms can be composed uniformly regardless of the ordering of - the left transversals in 'u_vector'. All products involving identity - automorphisms are ignored. - - Args: - u_vector: Coset representatives grouped by stabilizer index. - num_samples: The number of automorphisms to return. - seed: Random seed for reproducibility. - - Returns: - A list of uniformly sampled automorphisms in one-line notation. - - Example: - >>> import networkx as nx - >>> from dwave.experimental.automorphism import schreier_rep, sample_automorphisms - ... - >>> graph = nx.cycle_graph(8) - >>> result = schreier_rep(graph) - >>> sample_automorphisms(result.u_vector, seed=42) - [array([3, 4, 5, 6, 7, 0, 1, 2])] - >>> sample_automorphisms(result.u_vector, num_samples=2, seed=42) - [array([3, 4, 5, 6, 7, 0, 1, 2]), array([6, 5, 4, 3, 2, 1, 0, 7])] - """ - rng = np.random.default_rng(seed) - num_nodes = len(u_vector[0][0]) - u_counts = [len(u_i) for u_i in u_vector] - sampled_automorphisms = [] - - for _ in range(num_samples): - sample_indices = rng.integers(low=-1, high=u_counts) - g_product = np.arange(num_nodes) - - for i, u_i in enumerate(u_vector): - if sample_indices[i] >= 0: - g = u_i[sample_indices[i]] - g_product = mult(g, g_product) - - sampled_automorphisms.append(g_product) - - return sampled_automorphisms - - -def mult(alpha: NDArray[np.intp], beta: NDArray[np.intp]) -> NDArray[np.intp]: - """Compose two permutations in one-line notation, alpha after beta. - - Args: - alpha: A permutation represented as a list of integers in one-line notation. - beta: Another permutation of the same length. - - Returns: - The composition alpha ∘ beta in one-line notation. - - Example: - >>> import numpy as np - >>> from dwave.experimental.automorphism import mult - ... - >>> alpha = np.array([2,0,1], dtype=np.intp) # (0,2,1): 0->2, 1->0, 2->1 - >>> beta = np.array([1,2,0], dtype=np.intp) # (0,1,2): 0->1, 1->2, 2->0 - >>> mult(alpha, beta) - array([0, 1, 2]) - """ - return alpha[beta] - - -def inv(n: int, alpha: NDArray[np.intp]) -> NDArray[np.intp]: - """Calculate the inverse of a permutation in one-line notation. - - Args: - n: Length of permutation alpha. - alpha: A permutation represented as a list of integers in one-line notation. - - Returns: - The inverse of alpha in one-line notation. - - Example: - >>> import numpy as np - >>> from dwave.experimental.automorphism import inv - ... - >>> alpha = np.array([2,0,1], dtype=np.intp) # (0,2,1): 0->2, 1->0, 2->1 - >>> inv(3, alpha) - array([1, 2, 0]) - """ - alpha_inv = np.empty(n, dtype=np.intp) - alpha_inv[alpha] = np.arange(n, dtype=alpha_inv.dtype) - return alpha_inv - - -def schreier_rep( - graph: nx.Graph, - num_samples: int | None = None, - seed: int = 42, - graph_coloring=None -) -> SchreierContext: - """Compute Schreier representatives and orbits for a graph. - - Builds a depth-first search tree, iteratively performing color refinement - and vertex individualization until leaf nodes are reached where all graph - vertices are uniquely colored. Leaf nodes with identical adjacency matrices - represent graph automorphisms. Discovered automorphisms are used to prune - the search tree. - - If graphs have more than one component, automorphisms are found for each - individual component, and automorphisms between components are determined - by considering which components are isomorphic. Since the number of automorphisms - between isomorphic components scales factorially with the number of components, - this is significantly faster than naively performing refinement-individualization - over the whole graph. It would be possible to update ``u_vector`` directly - without using ``enter()``, which in principle should be even faster, and should - be the first place to look if further performance improvements are required. - - Args: - graph: A NetworkX Graph object representing the input graph containing - the following methods: - - ``nodes()``: iterable of all nodes - - ``number_of_nodes()``: total number of nodes - - ``edges()``: iterable of all edges - - ``neighbors()``: iterable of all neighbours for a given node - num_samples: Number of samples to use for generating new coset representatives - from the existing set. If not provided, all coset representatives are used. - seed: Random seed for reproducibility. Defaults to 42. - """ - if nx.number_connected_components(graph) == 1: - ctx = SchreierContext(graph, num_samples=num_samples, seed=seed) - initial_partition, trace, color, num_colors = ctx._initial_partition(graph_coloring) - - ctx._canon(initial_partition, trace, color, num_colors) - return ctx - - # relabel vertices so components have contiguous labels - index_to_node = {} - node_to_index = {} - next_label = 0 - - component_vertices = list(nx.connected_components(graph)) - for vertices in component_vertices: - for vertex in sorted(vertices): - node_to_index[vertex] = next_label - index_to_node[next_label] = vertex - next_label += 1 - - graph = nx.relabel_nodes(graph, node_to_index, copy=True) - - # enter component automorphisms into global graph - ctx = SchreierContext(graph, num_samples=num_samples, seed=seed) - ctx._index_to_node = index_to_node - ctx._node_to_index = node_to_index - - # group isomorphic components together - components = [ctx._graph.subgraph(c).copy() for c in nx.connected_components(ctx._graph)] - - unique_components = {} - for comp in components: - ctx_comp = SchreierContext(comp, num_samples=num_samples, seed=seed) - ctx_comp._compare_adj = True - - initial_partition, trace, color, num_colors = ctx_comp._initial_partition() - ctx_comp._canon(initial_partition, trace, color, num_colors) - - ctx._nodes_reached += ctx_comp.nodes_reached # update the global search tree statistics - ctx._leaf_nodes += ctx_comp.leaf_nodes - - unique_components.setdefault(ctx_comp._certificate(), []).append( - ComponentInfo(ctx_comp._u_vector, np.array(sorted(comp.nodes())), ctx_comp._best_perm) - ) - - # enter the local automorphisms - graph_nnodes = ctx._graph.number_of_nodes() - for identical_components in unique_components.values(): - for comp in identical_components: - for u in chain.from_iterable(comp.u_vector): - u_global = np.arange(graph_nnodes) - u_global[comp.nodes] = u_global[comp.nodes][u] - ctx._enter(u_global, mode=EnterMode.NO_RECURSE) - - # enter swap automorphisms - for comps in unique_components.values(): - for i in range(len(comps) - 1): - i_nodes = comps[i].nodes - j_nodes = comps[i + 1].nodes - - # swap automorphisms must be entered in the canonical basis - i_canon_perm = comps[i].best_perm - j_canon_perm = comps[i + 1].best_perm - i_canon = i_nodes[i_canon_perm] - j_canon = j_nodes[j_canon_perm] - - u_global = np.arange(graph_nnodes) - u_global[i_canon], u_global[j_canon] = u_global[j_canon], u_global[i_canon] - ctx._enter(u_global, mode=EnterMode.RECURSE_ONCE) - - return ctx - -def array_to_cycle( - array: NDArray[np.intp], - index_to_node: Mapping[int, Hashable] | None = None -) -> str: - """Convert an array in one-line notation to a string in cycle notation. - - Based on Algorithm 6.4 from Kreher, D. L., & Stinson, D. R. (1999). - Combinatorial algorithms: Generation, enumeration, and search. - - Args: - array: The permutation in one-line notation. - index_to_node: An optional relabelling dictionary. By default, array indices - are used. - - Returns: - The permutation as a string in cycle notation. - - Example: - >>> import numpy as np - >>> from dwave.experimental.automorphism import array_to_cycle - ... - >>> alpha = np.array([2,0,1], dtype=np.intp) # (0,2,1): 0->2, 1->0, 2->1 - >>> array_to_cycle(alpha) - '(0,2,1)' - >>> array_to_cycle(np.array([2,0,1]), index_to_node={0: 5, 1: 7, 2: 9}) - '(5,9,7)' - """ - if index_to_node is not None: - expected = set(range(len(array))) - if index_to_node.keys() != expected: - missing = expected - index_to_node.keys() - raise ValueError(f"index_to_node missing keys: {missing}") - - label = (lambda x: str(index_to_node[x])) if index_to_node is not None else str - unvisited = [True] * len(array) - cycle_parts = [] - - for i in range(len(array)): - if unvisited[i]: - cycle_parts.append('(') - cycle_parts.append(label(i)) - unvisited[i] = False - j = i - - while unvisited[array[j]]: - cycle_parts.append(',') - j = array[j] - cycle_parts.append(label(j)) - unvisited[j] = False - - cycle_parts.append(')') - return ''.join(cycle_parts) diff --git a/dwave/experimental/lattice_utils/lattice/lattice.py b/dwave/experimental/lattice_utils/lattice/lattice.py index af59079..a724923 100644 --- a/dwave/experimental/lattice_utils/lattice/lattice.py +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -126,7 +126,6 @@ def embed_lattice( if len(embeddings) >= min_number_of_embeddings and np.prod(embeddings.shape): self._save_embeddings(sampler, embeddings) - #def make_nominal_bqm(self, **kwargs) -> dimod.BQM: def make_nominal_bqm(self) -> dimod.BQM: """Construct a default nominal BQM coupling strength values set to +1. diff --git a/dwave/experimental/lattice_utils/lattice/orbits.py b/dwave/experimental/lattice_utils/lattice/orbits.py index 30f6b80..97d4664 100644 --- a/dwave/experimental/lattice_utils/lattice/orbits.py +++ b/dwave/experimental/lattice_utils/lattice/orbits.py @@ -19,7 +19,7 @@ import numpy as np from numpy.typing import NDArray -from dwave.experimental.lattice_utils.lattice.automorphism import schreier_rep +from dwave.experimental.automorphism import schreier_rep __all__ = [ 'reindex', diff --git a/dwave/experimental/lattice_utils/observable/kinks.py b/dwave/experimental/lattice_utils/observable/kinks.py index a9583e2..43e5bcf 100644 --- a/dwave/experimental/lattice_utils/observable/kinks.py +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -30,21 +30,31 @@ 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.""" - sample_array = dimod.as_samples(sample_set)[0] + """Compute the kink-kink correlator for 1D spin chains. - S = np.roll(sample_array, 1, axis=1) - K = np.multiply(S, sample_array) == np.sign(experiment.param["energy_scale"]) - L = K.shape[-1] - K = np.reshape(K, (-1, L)) - kink_density = np.mean(K) + 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. - CKK = np.zeros((K.shape[-1],)) + Returns: + A numpy array containing the kink-kink correlator values for each sample. + """ + samples = dimod.as_samples(sample_set)[0] - for R in range(1, L): - KR = np.roll(K, R, axis=1) - CKK[R] = np.mean(np.multiply(K, KR)) - np.power(np.mean(K), 2) + 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) - CKK /= kink_density**2 + kink_kink_correlator = np.zeros((kink_mask.shape[-1],)) - return CKK + 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 index e7380d8..8365a67 100644 --- a/dwave/experimental/lattice_utils/observable/observable.py +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -15,7 +15,7 @@ from __future__ import annotations from abc import ABC, abstractmethod from pathlib import Path -from typing import Any +from typing import Any, TypeAlias import numpy as np from numpy.typing import NDArray @@ -33,15 +33,19 @@ 'ReferenceEnergy', ] +ObservableResult: TypeAlias = ( + NDArray + | float + | int + | tuple[NDArray, tuple[int, int]] +) class Observable(ABC): - """The observable class does not take any parameters. Its primary - functionality is through the required 'evaluate' method, which requires - parameters 'experiment' and 'bqm' defining the context, and 'sample_set' - which provides the samples on which we compute the observable. Output is a - numpy array of arbitrary type (usually float). + """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__ @@ -51,7 +55,7 @@ def evaluate( experiment: Experiment, bqm: dimod.BQM, sample_set: dimod.SampleSet, - ) -> NDArray: + ) -> ObservableResult: raise NotImplementedError @@ -79,10 +83,9 @@ def evaluate( ) -> NDArray: sample_array = dimod.as_samples(sample_set)[0].astype(float) if len(experiment.inst.edge_list) == 0: - return [] - row, col = np.asarray(experiment.inst.edge_list).T + return np.empty(0, dtype=float) - # Surprisingly, it's faster to multiply the whole matrix. + 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 @@ -98,10 +101,9 @@ def evaluate( ) -> NDArray: sample_array = dimod.as_samples(sample_set)[0].astype(float) if len(experiment.inst.edge_list) == 0: - return [] - row, col = np.asarray(experiment.inst.edge_list).T + return np.empty(0, dtype=float) - # Surprisingly, it's faster to multiply the whole matrix. + 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] @@ -155,8 +157,10 @@ def evaluate( path: str | Path | None = None, inst: Lattice | None = None, ) -> float: - - if path is None: + """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(): @@ -166,8 +170,12 @@ def evaluate( # And if we can't load, we generate a reference sample. if experiment is not None: energy, sample, method_string = experiment.inst._optimize(bqm) - else: + 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) @@ -180,8 +188,11 @@ def load( path: str | Path | None = None, ) -> tuple[float, NDArray, str]: """Load and get the full data tuple, not just the energy.""" - if path is None: + 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()) @@ -192,6 +203,7 @@ def load( 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="") @@ -207,8 +219,10 @@ def update( 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. """ - reference_energy, _, reference_method_string = self.load(experiment, bqm, path) + 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: @@ -228,12 +242,19 @@ def get_reference_energy_path( ) -> Path: """Return the path to the reference energy file for the given experiment and BQM. - This needs to be fixed if you have something not in the instance - pathstring that needs to be taken into account, for example if the ground-state - energies depend on the chip. + 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 # defunct. + 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. @@ -258,7 +279,7 @@ def get_reference_energy_path( / experiment_data_dict["inst"]._get_instance_pathstring() ) - # Use hash. BQM is not hashable so use the experiment.inst data to generate a tuple. + # 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 ) diff --git a/dwave/experimental/lattice_utils/observable/triangular.py b/dwave/experimental/lattice_utils/observable/triangular.py index 89ef7b2..b6faac9 100644 --- a/dwave/experimental/lattice_utils/observable/triangular.py +++ b/dwave/experimental/lattice_utils/observable/triangular.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""For triangular order parameters""" - from __future__ import annotations import numpy as np from numpy.typing import NDArray @@ -26,7 +24,7 @@ class TriangularOP(Observable): - """For triangular lattices. Unembeds if possible.""" + """For calculating the order parameter of triangular lattices.""" def evaluate( self, @@ -34,12 +32,24 @@ def evaluate( 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"): - # If the lattice is an embedded lattice lbqm = experiment.inst.unembed_bqm(bqm) - # unembed the sample set. lss = experiment.inst.unembed_sampleset(sample_set) triangular_sublattice = experiment.inst.logical_lattice.sublattice else: @@ -48,8 +58,12 @@ def evaluate( sample_array = dimod.as_samples(lss)[0] - for edge in lbqm.quadratic: - assert triangular_sublattice[edge[0]] != triangular_sublattice[edge[1]] + 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): @@ -58,6 +72,6 @@ def evaluate( ) angles = np.array(np.exp([0.0, 1.0j * 4 * np.pi / 3, 1.0j * 2 * np.pi / 3])).T - op = np.matmul(sublattice_mags, angles).ravel() / np.sqrt(3) + order_parameter = np.matmul(sublattice_mags, angles).ravel() / np.sqrt(3) - return op + return order_parameter diff --git a/examples/example_1D_Ising_chain.py b/examples/example_1D_Ising_chain.py index e433ced..f9ea39f 100644 --- a/examples/example_1D_Ising_chain.py +++ b/examples/example_1D_Ising_chain.py @@ -34,12 +34,11 @@ # Two samplers: an Advantage2 prototype and an Advantage system. samplers = [ - DWaveSampler(solver="Advantage2_system3.1"), + DWaveSampler(solver="Advantage2_system1"), DWaveSampler(solver="Advantage_system4.1"), ] NUM_SPINS = 256 -#NUM_SPINS = 8 # Two energy scales: one strong coupling and one weak coupling. ENERGY_SCALES = (-1.8, 0.1) @@ -79,23 +78,21 @@ # 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 + # 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: - exp = experiment.FastAnnealExperiment( - inst=inst, - sampler=sampler, - loop_data_files=30, - max_iterations=5, + 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 (CITE). The observable + # 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()) @@ -103,7 +100,7 @@ # Make parameter list. We will only vary anneal time. parameter_list = [{"anneal_time": time} for time in ANNEAL_TIMES] - for _ in range(20): + for _ in range(20): #TODO clean this up? done = exp.run_iteration(parameter_list) if done: break @@ -122,7 +119,7 @@ fbshim = [] # flux bias shim kkc = [] # kink-kink correlator for param in parameter_list: - exp.apply_param(param) + 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])) diff --git a/tests/test_lattice_utils.py b/tests/test_lattice_utils.py index 7a7df43..96198e1 100644 --- a/tests/test_lattice_utils.py +++ b/tests/test_lattice_utils.py @@ -690,22 +690,6 @@ def test_apply_param_no_anneal_or_schedule_raises(self): with self.assertRaises(ValueError): exp.apply_param({"energy_scale": 1.0}) - def test_spin_reversal_disabled(self): - chain = Chain(dimensions=(4,), periodic=(True,)) - sampler = _make_mock_sampler() - exp = Experiment(chain, sampler) - self.assertIsNone(exp._get_spin_reversal_transform()) - - def test_spin_reversal_enabled_with_seed(self): - chain = Chain(dimensions=(4,), periodic=(True,)) - sampler = _make_mock_sampler(num_qubits=8) - exp = Experiment(chain, sampler) - exp.param["spin_reversal_transform"] = True - exp.param["spin_reversal_transform_seed"] = 42 - srt1 = exp._get_spin_reversal_transform() - srt2 = exp._get_spin_reversal_transform() - self.assertEqual(srt1, srt2) - def test_initial_shim_no_embeddings(self): chain = Chain(dimensions=(4,), periodic=(True,)) sampler = _make_mock_sampler() @@ -929,7 +913,6 @@ def test_make_bqm_with_embeddings(self): "total_iterations": 0, "relative_coupler_strength": np.ones((1, chain.num_edges)), } - sc.spin_reversal_transform = None bqm = exp._make_bqm(sc) self.assertGreater(len(bqm.quadratic), 0) @@ -996,22 +979,6 @@ def test_coupler_shim_basic_update(self): exp._update_coupler_shim(sc, results) self.assertEqual(sc.shimdata["relative_coupler_strength"].shape, (1, chain.num_edges)) - def test_parse_results_with_spin_reversal(self): - chain = Chain(dimensions=(4,), periodic=(False,)) - chain.embedding_list = np.array([[0, 1, 2, 3]]) - sampler = _make_mock_sampler() - exp = Experiment(chain, sampler) - exp.run_index = 0 - bqm = chain.make_nominal_bqm() - samples = np.ones((10, 4)) - ss = dimod.SampleSet.from_samples_bqm(samples, bqm) - - sc = SamplerCall(run_index=0) - sc.nominal_bqms = [bqm] - sc.spin_reversal_transform = {0: True, 1: False, 2: True, 3: False} - results = exp.parse_results(sc, ss) - self.assertIn("QubitMagnetization", results) - def test_get_shimdata_not_initialized(self): chain = Chain(dimensions=(4,), periodic=(True,)) sampler = _make_mock_sampler() From 0ef6a528151d0874b0bbb37d619ce3b5772ae7a6 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 12 May 2026 10:27:17 -0700 Subject: [PATCH 06/13] Remove temporary automorphism file and additional cleanup -removed spin reversal transform functionality -uses automorphism module already in dwave-experimental -uses data classes for experiment configs -formatted using black --- .../lattice_utils/experiment/experiment.py | 285 ++--- .../experiment/fast_anneal_experiment.py | 53 +- .../lattice_utils/experiment/samplercall.py | 1 - .../lattice_utils/lattice/automorphism.py | 1110 ----------------- .../lattice_utils/lattice/embedded_lattice.py | 2 +- .../lattice_utils/lattice/lattice.py | 1 - .../lattice_utils/lattice/orbits.py | 4 +- .../lattice_utils/lattice/triangular.py | 6 +- .../lattice_utils/observable/kinks.py | 36 +- .../lattice_utils/observable/observable.py | 64 +- .../lattice_utils/observable/triangular.py | 32 +- examples/example_1D_Ising_chain.py | 19 +- tests/test_lattice_utils.py | 33 - 13 files changed, 217 insertions(+), 1429 deletions(-) delete mode 100644 dwave/experimental/lattice_utils/lattice/automorphism.py diff --git a/dwave/experimental/lattice_utils/experiment/experiment.py b/dwave/experimental/lattice_utils/experiment/experiment.py index 3e1a81e..5d53c7b 100644 --- a/dwave/experimental/lattice_utils/experiment/experiment.py +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -20,6 +20,7 @@ from pathlib import Path from datetime import datetime from typing import Any +from dataclasses import dataclass import dimod import numpy as np @@ -35,72 +36,76 @@ ) from dwave.experimental.lattice_utils.experiment.samplercall import SamplerCall -__all__ = ['Experiment'] +__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 class Experiment: - """Base class for experiment in LatQA.""" - - default_parameters = { - "energy_scale": 1.0, - "num_reads": 100, - "anneal_time": 1.0, - "num_random_instances": 1, - "readout_thermalization": 100, - "flux_bias_shim_step": 0.0, - "coupler_shim_step": 0.0, - "anneal_offset_shim_step": 0.0, - "target_magnetization": 0.0, - } - - observables_to_collect = { - QubitMagnetization(), - CouplerCorrelation(), - CouplerFrustration(), - SampleEnergy(), - BitpackedSpins(), - ReferenceEnergy(), - } - - def __init__(self, inst: Lattice, sampler: dimod.Sampler, **kwargs): - self.inst: Lattice = inst - self.sampler: dimod.Sampler = sampler - self.param: dict = {} - self.already_initialized: bool = False # until evidenced - - # Any forced types that might be required. - self.typedict: dict[str, float] = {"energy_scale": float} - - self.experiment_results_root: Path = Path( - kwargs.get("results_root", Path.cwd() / "results") - ).resolve() - - self.loop_data_files: int = kwargs.get("loop_data_files", 1000000000) - self.max_iterations: int | None = kwargs.get("max_iterations", None) - - # Apply parameters - for field, default in self.default_parameters.items(): - value = kwargs.get(field, default) - - if field in kwargs and field in self.typedict: - value = self.typedict[field](value) - - self.param[field] = value + """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, - starting_iteration: int | None = None, - result_fields: dict[str, Any] | None = None, + start_iteration: int | None = None, + result_fields: list[str] | None = None, quiet: bool = True, ignore_shim: bool = False, - mod: int | None = None, - ) -> dict[str, Any]: + ) -> list[dict[str, Any]]: """Load results from the highest-numbered iterations of the experiment. Args: num_iterations: Maximum number of iterations to load. - starting_iteration: If provided, load results starting from this + 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 @@ -108,19 +113,13 @@ def load_results( 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. - mod: If provided, the returned list is truncated so its length is - divisable by ``mod``. Returns: A list of dictionaries containing the results for each iteration. """ - self._set_run_index() # this is the NEXT run index. - fnlist = self._get_sorted_results_file_list() - if starting_iteration is not None: - fnlist = fnlist[ - max(starting_iteration, 0) : max(starting_iteration + num_iterations, 0) - ] + if start_iteration is not None: + fnlist = fnlist[max(start_iteration, 0) : max(start_iteration + num_iterations, 0)] else: fnlist = fnlist[-num_iterations:] @@ -142,34 +141,23 @@ def load_results( results.append({k: data[k] for k in result_fields}) - if mod is not None: - return results[: mod * (len(results) // mod)] - return results - def apply_param(self, param: dict[str, float]) -> bool: + 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. - Keys correspond to parameter names and values are converted to - the appropriate type if a converted is defined in ``self.typedict``. - - Returns: - A boolean value corresponding to whether or not the experiment has - already completed all required iterations for this parameter - configuration. """ - # Reformat for filename consistency; can be important for loading data. param = self._format_parameter_list([param])[0] for param_name, param_val in param.items(): - if param_name in self.typedict: - self.param[param_name] = self.typedict[param_name](param_val) - else: - self.param[param_name] = param_val + self.param[param_name] = param_val self.data_path = self.experiment_results_root / self._get_relative_data_path() - self.already_initialized = self._set_run_index() + 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``. @@ -217,18 +205,21 @@ def run_iteration(self, parameter_list: list, **kwargs) -> bool: # 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]) - - # Get the results and update the shim - results = self.parse_results(call_dict[index], response_dict[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 @@ -237,8 +228,7 @@ def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[st """Parse a sampler response into per-embedding observable results. Args: - call: Sampler call metadata, cinluding the nominal BQMs and any - applied spin-reversal transform. + call: Sampler call metadata, inluding the nominal BQMs response: Raw sample set returned by the sampler. Returns: @@ -249,18 +239,9 @@ def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[st embedding_list = self.inst.embedding_list myarr = response.samples(sorted_by=None) sample_arrays = [myarr[:, emb].copy() for emb in embedding_list] - - if call.spin_reversal_transform is not None: - for iemb, emb in enumerate(embedding_list): - for iv, v in enumerate(emb): - if call.spin_reversal_transform[v]: - sample_arrays[iemb][:, iv] *= -1 else: sample_arrays = [response.samples(sorted_by=None)[:, np.arange(self.inst.num_spins)]] - if call.spin_reversal_transform is not None: - raise NotImplementedError - sample_set = {} for iemb, sample_array in enumerate(sample_arrays): sample_set[iemb] = dimod.SampleSet.from_samples_bqm( @@ -308,37 +289,30 @@ def _save_results( def _get_sorted_results_file_list(self) -> list[str]: """Return result filenames sorted lexicographically.""" - fnlist = list(self.data_path.glob("iter*")) + fnlist = list(self.data_path.glob("iter*.pkl.lzma")) fnlist.sort() return [str(fn) for fn in fnlist] - def _set_run_index(self) -> bool: - """Set the run index (data file index, iteration number) for the parameterization. + 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 - Uses the name of the most recently modified file in - the folder. Also creates the data path if it doesn't exist. Return value - is boolean, whether the iterations have already been started or not. - """ - # If the path doesn't exist, create it and set the run index to zero. - if self.data_path.exists() is False: - os.makedirs(self.data_path) - self.run_index = 0 - return False - - fnlist = list(self.data_path.glob("iter*")) + fnlist = list(self.data_path.glob("iter*.pkl.lzma")) if not fnlist: - self.run_index = 0 - return False + return 0, False + + latest_file_iter = max(int(fn.stem.split(".")[0][4:]) for fn in fnlist) + return latest_file_iter + 1, True - # Folder exists and is nonempty. - # Only complication is looping; we set the flag to True if we overflow mod - fnlist.sort(key=lambda x: os.path.getmtime(x)) - latest_file_iter = int(fnlist[-1].stem.split(".")[0][4:]) - self.run_index = np.max([0, np.mod(latest_file_iter + 1, self.loop_data_files)]) - if self.run_index < latest_file_iter: - self.has_looped = 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.") - return True + 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. @@ -368,16 +342,7 @@ def _get_parameter_pathstring(self) -> str: """ energy_scale = self.param["energy_scale"] - if type(self.sampler).__name__ in [ - "SimulatedAnnealingSampler", - "SimulatedQuantumAnnealingSampler", - "RotorSampler", - "ParallelTemperingSampler", - "DiscreteSimulatedBifurcationSampler", - ]: - # Applies to Monte Carlo dynamics, for example - pathstring = f'energyscale{energy_scale:0.3}/nsweeps{self.param["num_sweeps"]:010d}mcs' - elif "anneal_time" in self.param: + 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"]}' @@ -418,10 +383,12 @@ def _build_sampler_call(self) -> None | SamplerCall: sampler_call = SamplerCall(run_index=self.run_index) sampler_call.nominal_bqms = self._make_nominal_bqms() sampler_call.shimdata = self._get_shimdata() - sampler_call.spin_reversal_transform = self._get_spin_reversal_transform() # Here we can find out that we're finished. - if sampler_call.shimdata["total_iterations"] >= self.max_iterations: + 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) @@ -440,8 +407,6 @@ def _format_parameter_list( """ ret = parameter_list.copy() for entry in ret: - if "target_s" in entry: - entry["target_s"] = np.round(entry["target_s"], 4) if "anneal_time" in entry: entry["anneal_time"] = np.round(entry["anneal_time"], 6) if "anneal_schedule" in entry: @@ -501,32 +466,10 @@ def _make_sampler_params(self, **kwargs) -> dict[str, Any]: if "anneal_offsets" in kwargs["shimdata"]: ret["anneal_offsets"] = list(kwargs["shimdata"]["anneal_offsets"]) - if "num_sweeps" in self.param: - ret["num_sweeps"] = self.param["num_sweeps"] - elif "target_s" in self.param and "dwell_time" in self.param: # fast reverse anneal - ret["x_target_s"] = self.param["target_s"] - ret["x_dwell_time"] = self.param["dwell_time"] - ret["anneal_schedule"] = self.param["anneal_schedule"] - elif "anneal_schedule" in self.param: - ret["anneal_schedule"] = self.param["anneal_schedule"] - elif self.param.get("fast_anneal", False): + if self.param.get("fast_anneal", False): ret["fast_anneal"] = True - ret["annealing_time"] = self.param["anneal_time"] - else: - ret["annealing_time"] = self.param["anneal_time"] - if "reinitialize_state" in self.param and ret["anneal_schedule"][0][1] == 1: - ret["reinitialize_state"] = self.param.get("reinitialize_state", False) - if ret["reinitialize_state"] is None: - ret["reinitialize_state"] = False - - if self.param.get("initial_state", None) is not None and ret["anneal_schedule"][0][1] == 1: - ret["initial_state"] = self.param["initial_state"] - elif "reinitialize_state" in self.param and ret["anneal_schedule"][0][1] == 1: - # Set to None, meaning that it will be randomized. - ret["initial_state"] = { - qubit: np.random.randint(2) * 2 - 1 for qubit in self.inst.embedding_list.ravel() - } + ret["annealing_time"] = self.param["anneal_time"] return ret @@ -552,8 +495,7 @@ def _make_initial_shim(self) -> dict[str, Any]: def _get_latest_iteration_filename(self) -> Path: """Return the filename of the most recently completed iteration.""" - file_name = f"iter{np.mod(self.run_index - 1, self.loop_data_files):05d}.pkl.lzma" - return self.data_path / file_name + 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.""" @@ -615,7 +557,6 @@ def _update_coupler_shim( step_size: float | None = None, ) -> None: """Update relative coupler strength based on measured frustration.""" - # Ok, let's use orbits and stuff. orbits = self.inst.coupler_orbits energy_scale = self.param["energy_scale"] relative_coupler_strength = sampler_call.shimdata["relative_coupler_strength"] @@ -630,12 +571,11 @@ def _update_coupler_shim( normalization_basis = np.ones_like(orbits, dtype=bool) # Assume we have multiple embeddings of the same BQM. - assert ( - len(sampler_call.nominal_bqms) == 1 - or sampler_call.nominal_bqms[0] == sampler_call.nominal_bqms[1] - ), "Case for distinct embedded BQMs not implemented yet." + 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 = sampler_call.nominal_bqms[0] + 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): @@ -743,11 +683,6 @@ def _make_bqm(self, sampler_call: SamplerCall) -> dimod.BQM: for iemb, emb in enumerate(self.inst.embedding_list): nominal_bqm = sampler_call.nominal_bqms[iemb].copy() - if sampler_call.spin_reversal_transform is not None: - for iv, v in enumerate(emb): - if sampler_call.spin_reversal_transform[v]: - nominal_bqm.flip_variable(iv) - for v in range(self.inst.num_spins): # Don't touch degree-zero spins. Relevant to partial yield. if nominal_bqm.degree(v) > 0: @@ -762,19 +697,3 @@ def _make_bqm(self, sampler_call: SamplerCall) -> dimod.BQM: bqm.add_quadratic(emb[edge[0]], emb[edge[1]], bias) return bqm - - def _get_spin_reversal_transform(self) -> dict[int, bool] | None: - """Generate a spin-reversal transform for the current sampler.""" - if self.param.get("spin_reversal_transform", False): - seed = self.param.get("spin_reversal_transform_seed", None) - rng = np.random.default_rng(seed) - flips = rng.random(len(self.sampler.nodelist)) > 0.5 - - spin_reversal_transform = {v: False for v in self.sampler.nodelist} - for iv, v in enumerate(self.sampler.nodelist): - if flips[iv]: - spin_reversal_transform[v] = True - - return spin_reversal_transform - - return None diff --git a/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py index 6e67cf6..59ca0b2 100644 --- a/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py +++ b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py @@ -12,45 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dwave.experimental.lattice_utils.experiment import Experiment -from dwave.experimental.lattice_utils.observable import ( - QubitMagnetization, - CouplerCorrelation, - CouplerFrustration, - SampleEnergy, - BitpackedSpins, - ReferenceEnergy, -) +from dataclasses import dataclass -__all__ = ['FastAnnealExperiment'] +from dwave.experimental.lattice_utils.experiment import ExperimentConfig +__all__ = ['FastAnnealExperimentConfig'] -class FastAnnealExperiment(Experiment): - # Set default parameters - default_parameters = { - "energy_scale": 1.0, - "automorph_embeddings": False, - "spin_reversal_transform": False, - "spin_reversal_transform_seed": None, - "num_reads": 100, - "num_random_instances": None, - "readout_thermalization": 100, - "fast_anneal": True, - "anneal_time": 1.0, - "flux_bias_shim_step": 0.0, - "coupler_shim_step": 0.0, - "coupler_damp": 0.0, - "anneal_offset_shim_step": 0.0, - "anneal_offset_damp": 0.0, - "individual_qubit_anneal_offsets": None, - "target_magnetization": 0.0, - "logical_software": False, - } - observables_to_collect = { - QubitMagnetization(), - CouplerCorrelation(), - CouplerFrustration(), - SampleEnergy(), - BitpackedSpins(), - ReferenceEnergy(), - } + +@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 index 458ba91..cbd859b 100644 --- a/dwave/experimental/lattice_utils/experiment/samplercall.py +++ b/dwave/experimental/lattice_utils/experiment/samplercall.py @@ -36,4 +36,3 @@ def __init__( 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 - self.spin_reversal_transform: None = None diff --git a/dwave/experimental/lattice_utils/lattice/automorphism.py b/dwave/experimental/lattice_utils/lattice/automorphism.py deleted file mode 100644 index e22996a..0000000 --- a/dwave/experimental/lattice_utils/lattice/automorphism.py +++ /dev/null @@ -1,1110 +0,0 @@ -# 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 import deque -from collections.abc import Hashable -from dataclasses import dataclass -from enum import Enum, auto -import hashlib -from itertools import chain -import random -from typing import Mapping - -import networkx as nx -import numpy as np -from numpy.typing import NDArray - - -@dataclass -class ComponentInfo: - """Container for per-component data used during automorphism discovery on disjoint graphs.""" - u_vector: list - nodes: NDArray - best_perm: NDArray - -class EnterMode(Enum): - """Controls when the ``_enter()`` function attempts to compose new automorphisms.""" - RECURSE = auto() - RECURSE_ONCE = auto() - NO_RECURSE = auto() - -class SchreierContext: - """This object holds mutable states used throughout the automorphism calculation. - - Args: - graph: A NetworkX Graph object representing the input graph. - num_samples: Number of samples to use for generating new coset representatives - from the existing set. If not provided, all coset representatives are used. - seed: Seed used for reproducibility. Defaults to 42. - """ - def __init__(self, graph: nx.Graph, num_samples: int | None = None, seed: int = 42) -> None: - original_nodes_sorted = sorted(graph.nodes()) - self._index_to_node: dict[int, Hashable] = { - new: old for new, old in enumerate(original_nodes_sorted) - } - self._node_to_index: dict[Hashable, int] = { - old: new for new, old in enumerate(original_nodes_sorted) - } - graph = nx.relabel_nodes(graph, self._node_to_index) # relabel nodes contiguously (0...n-1) - - self._nodes: list[int] = list(graph.nodes()) - self._num_nodes: int = graph.number_of_nodes() - self._graph_edges: list[tuple[int, int]] = list(graph.edges()) - self._neighbours: list[set[int]] = [set(graph.neighbors(i)) for i in range(self._num_nodes)] - self._graph: nx.Graph = graph - - self._num_samples: int | None = num_samples - self._rng: random.Random = random.Random(seed) - - self._leaf_nodes: int = 0 - self._nodes_reached: int = 0 - self._depth: int = 0 - - self._u_map: dict[np.intp, int] = {} - self._u_len: int = 0 - self._u_vector: list = [] - self._u_vector_inv: list[list[NDArray[np.intp]]] = [] - - self._identity: NDArray[np.intp] = np.arange(self._num_nodes, dtype=np.intp) - - self._best_perm: NDArray = np.arange(self._num_nodes) - self._best_perm_exist: bool = False - self._compare_adj: bool = False - self._trace_history: list = [] - - self._in_colors_adj: bytearray = bytearray(self._num_nodes) - self._in_refine_stack: bytearray = bytearray(self._num_nodes) - - self._color_degree: list[int] = [0] * self._num_nodes - self._min_color_degree: list[int] = [0] * self._num_nodes - self._max_color_degree: list[int] = [0] * self._num_nodes - self._active_vertices: list[list[int]] = [[] for _ in range(self._num_nodes)] - - if self._num_nodes <= 65535: - self._color_dtype: np.dtype = np.uint16 - else: - self._color_dtype: np.dtype = np.uint32 - - @property - def leaf_nodes(self) -> int: - """Number of leaf nodes encountered in the search tree.""" - return self._leaf_nodes - - @property - def nodes_reached(self) -> int: - """Total number of nodes reached during traversal of the search tree.""" - return self._nodes_reached - - @property - def index_to_node(self) -> dict[int, Hashable]: - """The mapping from the basis of relabelled nodes (0...n-1) to the original - node labels.""" - return self._index_to_node - - @property - def node_to_index(self) -> dict[Hashable, int]: - """The mapping from the original node labels to the basis of relabelled - nodes (0...n-1).""" - return self._node_to_index - - @property - def u_map(self) -> dict[np.intp, int]: - """Map from coset representative group index to stabilizer index.""" - return self._u_map - - @property - def u_vector(self) -> list[list[NDArray[np.intp]]]: - """Coset representatives grouped by stabilizer index.""" - return self._u_vector - - @property - def num_automorphisms(self) -> int: - """Number of automorphisms implied by u_vector.""" - if self._u_vector: - return int(np.prod([len(u_i) + 1 for u_i in self._u_vector], dtype=object)) - else: - return 1 - - @property - def vertex_orbits(self) -> list[list[int]]: - """Vertex orbits induced by the coset representatives in u_vector and returned - in the basis of relabelled nodes (0...n-1).""" - return vertex_orbits(self._u_vector, self._nodes) - - @property - def vertex_orbits_original_labels(self) -> list[list[Hashable]]: - """Vertex orbits induced by the coset representatives in u_vector and returned - with the original node labels.""" - return vertex_orbits(self._u_vector, self._nodes, index_to_node=self._index_to_node) - - @property - def edge_orbits(self) -> list[list[int]]: - """Edge orbits induced by the coset representatives in u_vector and returned - in the basis of relabelled nodes (0...n-1).""" - return edge_orbits(self._u_vector, self._graph_edges) - - @property - def edge_orbits_original_labels(self) -> list[list[Hashable]]: - """Edge orbits induced by the coset representatives in u_vector and returned - with the original node labels.""" - return edge_orbits(self._u_vector, self._graph_edges, index_to_node=self._index_to_node) - - def _test_composability(self, g: NDArray[np.intp]) -> tuple[int, NDArray[np.intp]]: - """Test if an automorphism is composable from coset representatives. - - Based on Algorithm 6.10 from Kreher, D. L., & Stinson, D. R. (1999). - Combinatorial algorithms: Generation, enumeration, and search. - - Modified to use a mask to skip sifting by identity permutations, which - have no effect. - - Args: - g: A permutation represented as a list of integers in one-line notation. - - Returns: - A tuple (i, g_reduced) where i is the index of the first base position - that could not be sifted. If ``g`` is completely sifted the returned index - equals ``self._num_nodes``. ``g_reduced`` is the permutation obtained after - sifting through all positions up to (but not including) the returned - index. - """ - mask = (g != self._identity) - index = mask.argmax() - next_diff = 0 - - while mask[index]: - next_diff += index - if next_diff not in self._u_map: - return next_diff, g - - for i, h in enumerate(self._u_vector[self._u_map[next_diff]]): - if h[next_diff] == g[next_diff]: - break - else: - return next_diff, g - - g = self._u_vector_inv[self._u_map[next_diff]][i][g] - mask = (g[next_diff:] != self._identity[next_diff:]) - index = mask.argmax() - - return self._num_nodes, g - - def _enter(self, g: NDArray[np.intp], mode: EnterMode = EnterMode.RECURSE) -> None: - """Add automorphism if it can't be composed from coset representatives. - - Based on Algorithm 6.11 from Kreher, D. L., & Stinson, D. R. (1999). - Combinatorial algorithms: Generation, enumeration, and search. - - If an automorphism can't be composed from existing coset representatives - it is added as a new coset representative to u_vector. Depending on the - setting of ``mode``, ``_enter()`` is called recursively to attempt to - compose additional coset representatives from the composition between - the newly-discovered coset representative and existing coset representatives. - - The automorphisms discovered will result in pruning comparable to nauty, - as measured by comparing the total number of search tree nodes visited - for zephyr graphs of various sizes. - - Args: - g: A permutation represented as a list of integers in one-line notation. - mode: Specifies if recursive calls to ``enter()`` are performed to attempt - to compose new automorphisms. The setting ``EnterMode.RECURSE_ONCE`` - results in a single call to ``enter()`` per coset representative where - no further attempts to compose automorphisms occur. - """ - i, g = self._test_composability(g) - if i == self._num_nodes: - return - - if i not in self._u_map: - self._u_map[i] = self._u_len - self._u_len += 1 - self._u_vector.append([]) - self._u_vector_inv.append([]) - - self._u_vector[self._u_map[i]].append(g) - self._u_vector_inv[self._u_map[i]].append(inv(self._num_nodes, g)) - - if mode is EnterMode.NO_RECURSE: - return - - for u_i in self._u_vector: - for h in u_i: - f = mult(g, h) - if mode is EnterMode.RECURSE_ONCE: - self._enter(f, mode=EnterMode.NO_RECURSE) - else: - self._enter(f) - - def _refine( - self, - partition: list[set[int]], - trace: NDArray[np.integer], - color: NDArray[np.integer], - num_colors: int, - individualized_vertex: int | None = None, - ) -> None: - """Perform color refinement on the current partition until an equitable - coloring is reached. - - This procedure implements the 1-dimensional Weisfeiler-Leman (WL) refinement, - following Algorithms 2 and 3 of Berkholz (2016), *Tight lower and upper bounds - for the complexity of canonical color refinement*. - - A refinement stack is initialized with either: - • all color classes (if no vertex has been individualized), or - • the color class of the individualized vertex. - - For each color class popped from the stack, the algorithm computes the - color-degree of every vertex: the number of neighbours it has in the refining - color class. These color-degrees determine how each color class should be - split. If a color class contains vertices with differing color-degrees, it is - partitioned into new color classes, and the smaller subcells are pushed onto - the refinement stack. - - The process continues until no color class can be further refined, yielding an - equitable coloring. - - If a vertex was individualized prior to this refinement step, only the - color class containing that vertex needs to be placed on the refinement - stack initially, since only colors adjacent to that color can be affected. - - For performance reasons, ``num_colors`` is passed as a single-element list - so that updates to the number of colors persist across calls without having - to return anything. - - Args: - partition: The current partition structure, represented as a list of sets of vertices - ordered by color. - trace: A list of the sizes of each partition cell (color class), ordered by color. - color: An array mapping each vertex to its current color. - num_colors: The current number of colors in the partition. - individualized_vertex: The vertex individualized prior to this refinement step, if any. - - Returns: - The new number of colors, the updated trace array, and the updated color array. - """ - neighbours = self._neighbours - color_degree = self._color_degree - min_color_degree = self._min_color_degree - max_color_degree = self._max_color_degree - active_vertices = self._active_vertices - in_refine_stack = self._in_refine_stack - in_colors_adj = self._in_colors_adj - - colors_adj = [] - - if individualized_vertex is None: - refine_stack = list(range(num_colors)) - else: - refine_stack = [color[individualized_vertex]] - num_colors = [num_colors] # mutable container so ``_split_up_color()`` can increment it - - for v in refine_stack: - in_refine_stack[v] = 1 - - while refine_stack: - refinement_color = refine_stack.pop() - in_refine_stack[refinement_color] = 0 - - for v in partition[refinement_color]: - for w in neighbours[v]: - color_degree[w] += 1 - cw = color[w] - if color_degree[w] == 1: - active_vertices[cw].append(w) - if in_colors_adj[cw] == 0: - colors_adj.append(cw) - in_colors_adj[cw] = 1 - if color_degree[w] > max_color_degree[cw]: - max_color_degree[cw] = color_degree[w] - - for c in colors_adj: - if trace[c] != len(active_vertices[c]): - min_color_degree[c] = 0 - else: - min_color_degree[c] = max_color_degree[c] - for v in active_vertices[c]: - if color_degree[v] < min_color_degree[c]: - min_color_degree[c] = color_degree[v] - - colors_to_split = [] - for c in colors_adj: - if min_color_degree[c] < max_color_degree[c]: - colors_to_split.append(c) - - for color_to_split in sorted(colors_to_split): - self._split_up_color( - color_to_split=color_to_split, - partition=partition, - color=color, - trace=trace, - active_vertices=active_vertices, - color_degree=color_degree, - min_degree=min_color_degree[color_to_split], - max_degree=max_color_degree[color_to_split], - refine_stack=refine_stack, - in_refine_stack=in_refine_stack, - num_colors=num_colors, - ) - - ## reset attributes - for c in colors_adj: - for v in active_vertices[c]: - color_degree[v] = 0 - max_color_degree[c] = 0 - active_vertices[c] = [] - in_colors_adj[c] = 0 - colors_adj = [] - - return num_colors[0], trace, color - - def _split_up_color( - self, - *, - color_to_split: int, - partition: list[set[int]], - color: NDArray[np.integer], - trace: NDArray[np.integer], - active_vertices: list[list[int]], - color_degree: list[int], - min_degree: int, - max_degree: int, - refine_stack: list[int], - in_refine_stack: bytearray, - num_colors: list[int], - ) -> None: - """Splits a color class into subcells based on the color-degrees of its vertices. - - Based on algorithm 3 of Berkholz (2016), *Tight lower and upper bounds - for the complexity of canonical color refinement*. - - Given a color class ``color_to_split`` whose vertices exhibit differing - color-degrees with respect to the current refining color, this routine - partitions that class into new color classes. Vertices with the same - color-degree remain together, while vertices with different degrees are - assigned fresh color identifiers. - - The largest resulting subcell retains the original color label, while - all smaller subcells are assigned new colors and pushed onto the - refinement stack (Hopcroft's trick). The partition structure, trace array, - number of colors, and vertex-to-color mapping are updated in place. - - Args: - color_to_split: The color class to be split. - partition: The current partition structure, represented as a list of sets of vertices - ordered by color. - color: An array mapping each vertex to its current color. - trace: A list of the sizes of each partition cell (color class), ordered by color. - active_vertices: Lists of vertices adjacent to the color class being split, - ordered by color. - color_degree: The color-degree of each vertex. - min_degree: Minimum color-degree among vertices in the color class being split. - max_degree: Maximum color-degree among vertices in the color class being split. - refine_stack: Stack of colors scheduled for refinement. - in_refine_stack: Flags indicating which colors are already on the stack. - num_colors: The number of colors, used to determine the next color label to assign to - newly-refined cells. Stored as a single-element list so that updates persist across - calls. - """ - degree_to_new_color = [0] * (max_degree + 1) - num_color_degree = [0] * (max_degree + 1) - num_color_degree[0] = trace[color_to_split] - len(active_vertices[color_to_split]) - - for v in active_vertices[color_to_split]: - num_color_degree[color_degree[v]] += 1 - - largest_subcell_degree = 0 - for i in range(1, max_degree + 1): - if num_color_degree[i] > num_color_degree[largest_subcell_degree]: - largest_subcell_degree = i - - for i in range(max_degree + 1): - if num_color_degree[i] > 0: - if i == min_degree: - degree_to_new_color[i] = color_to_split - if not in_refine_stack[color_to_split] and i != largest_subcell_degree: - refine_stack.append(degree_to_new_color[i]) - in_refine_stack[degree_to_new_color[i]] = 1 - else: - degree_to_new_color[i] = num_colors[0] - partition[num_colors[0]] = set() - if in_refine_stack[color_to_split] or i != largest_subcell_degree: - refine_stack.append(degree_to_new_color[i]) - in_refine_stack[degree_to_new_color[i]] = 1 - num_colors[0] += 1 - - for v in active_vertices[color_to_split]: - new_color = degree_to_new_color[color_degree[v]] - if new_color != color_to_split: - partition[color_to_split] = partition[color_to_split] - {v} # must create new obj - partition[new_color].add(v) - trace[color_to_split] -= 1 - trace[new_color] += 1 - color[v] = new_color - - def _canon( - self, - partition: list[set[int]], - trace: NDArray[np.integer], - color: NDArray[np.integer], - num_colors: int, - individualized_vertex: int | None = None, - ) -> None: - """Generate search tree based on iterative color refinement and vertex - individualization. - - Loosely based on Algorithm 7.9 from Kreher, D. L., & Stinson, D. R. (1999). - Combinatorial algorithms: Generation, enumeration, and search. Additional - data structures are used to efficiently track the number of vertices - belonging to each color, vertex colors, and number of colors. Additionally, - the most recently individualized vertex is tracked and used to perform - color refinement more efficiently. - - Color refinement is performed iteratively on a graph until a discrete - coloring is achieved. If the coloring is not discrete after refinement, - vertices belonging to the same color are individualized, meaning that they - are assigned a new color, often breaking the symmetry of the graph and - allowing a subsequent color refinement step to produce further refinement. - - By default, graph comparisons using adjacency matrices are not performed, as - this becomes a bottleneck for even modestly sized graphs. Instead, the - ``trace`` for each graph is compared, which corresponds to the number of - vertices belonging to each color, ordered by color. This check is orders of - magnitude faster and has been found to have identical pruning capability - for graphs of interest, such as chimera, pegasus, and zephyr graphs, as - well as the disjoint compositions of smaller and simpler graphs as may - be encountered when doing parallel embeddings. - - If a graph has more than one component, comparisons using adjacency matrices are - used. This enables isomorphism detection between components, and in turn - a more efficient approach to generating the full automorphism group, which - may contain many automorphisms between isomorphic components. - - Kreher and Stinson perform comprehensive pruning by changing the base of - the left transversals to coincide with the current permutation order up to the - first non-discrete partition cell, or first split. At the cost of performing - this base change, it allows pruning to be performed by only considering - the left transversal with a stabilizer index equal to the index of the first - split. In practice, changing the base at each node of the search tree - becomes prohibitively expensive even more mostly sized graphs, and instead - the approach taken here is to avoid base changes, but instead to more carefully - evaluate which coset representatives to use for pruning. This is done by - ignoring the automorphisms that do not respect the current partition structure. - - Args: - partition: The current partition structure, represented as a list of - sets of vertices ordered by color. - trace: The number of vertices belonging to each color, ordered by color. - color: A map from each vertex to its color. - num_colors: The number of unique colors, equivalent to the number - of cells in the partition. - individualized_vertex: The most recently individualized vertex. - """ - self._nodes_reached += 1 - self._depth += 1 - - num_colors, trace, color = self._refine( - partition, - trace, - color, - num_colors, - individualized_vertex=individualized_vertex - ) - - if not self._best_perm_exist: - self._trace_history.append(trace.tobytes()) - - # first non-singleton block index - first_split = self._num_nodes - 1 - for i, block in enumerate(partition): - if len(block) > 1: - first_split = i - break - - compare_result = 2 - if self._best_perm_exist: # if a leaf node has been reached previously - - if self._compare_adj: - perm_candidate = list(chain.from_iterable(p for p in partition if p is not None)) - compare_result = self._compare(perm_candidate, first_split) - else: - compare_result = trace.tobytes() == self._trace_history[self._depth - 1] - - if compare_result == 0: - return - - if first_split == self._num_nodes - 1: # leaf node reached - self._leaf_nodes += 1 - - if not self._best_perm_exist: - self._best_perm_exist = True - self._best_perm[:] = list(chain.from_iterable(partition)) - - elif compare_result == 2: - perm_candidate = list(chain.from_iterable(partition)) - self._best_perm[:] = perm_candidate - - elif compare_result == 1: - perm_transformed = np.empty(self._num_nodes, dtype=np.intp) - perm_candidate = list(chain.from_iterable(partition)) - perm_transformed[perm_candidate] = self._best_perm - self._enter(perm_transformed) - - return - - candidates = sorted(partition[first_split]) - remaining_in_block = partition[first_split] - updated_partition = partition - trace[first_split] -= 1 - trace[num_colors] = 1 - - while candidates: - vertex = next(iter(candidates)) - updated_partition[first_split] = remaining_in_block - {vertex} - updated_partition[num_colors] = {vertex} - individualized_partition = list(updated_partition) # copy outer list - color[vertex] = num_colors # updated individualized cell - trace_copy = np.array(trace) - color_copy = np.array(color) - - self._canon( - individualized_partition, - trace_copy, - color_copy, - num_colors + 1, - individualized_vertex=vertex - ) - - color[vertex] = first_split - candidates.remove(vertex) - - # prune the search tree using automorphisms - for stab_index, u_index in self._u_map.items(): - if stab_index > vertex: # these automorphisms map vertex to itself - continue - - for g in self._u_vector[u_index]: - if g[vertex] not in candidates: - continue - - # automorphism must respect current partition structure - for w in candidates: - if color[w] != color[g[w]]: - break - else: - candidates.remove(g[vertex]) - - self._depth -= 1 - - def _compare(self, perm: NDArray[np.intp], first_split: int) -> int: - """Compare canonical adjacency matrix against itself under a partial permutation. - - At the first differing entry, returns whether the partial permutatation has - a greater or lesser value, otherwise it returns that they are equal. - - Based on Algorithm 7.6 from Kreher, D. L., & Stinson, D. R. (1999). - Combinatorial algorithms: Generation, enumeration, and search. - - Args: - perm: The permutation of the adjacency matrix to compare the canonical - adjacency matrix against. - first_split: The index of the first block of the partition containing - more than one vertex, defining the size of the partial permutation - of perm to use. - - Returns: - An integer 0, 1, or 2 depending on whether the partial permutation - perm results in an adjacency matrix which is less than, equal to, or - greater than the canonical adjacency matrix, respectively. - """ - neighbours = self._neighbours - best_perm = self._best_perm - for j in range(1, first_split): - neighbours_best_j = neighbours[best_perm[j]] - neighbours_pi_j = neighbours[perm[j]] - for i in range(j): - bit_best = 1 if best_perm[i] in neighbours_best_j else 0 - bit_pi = 1 if perm[i] in neighbours_pi_j else 0 - if bit_best < bit_pi: - return 0 - if bit_best > bit_pi: - return 2 - return 1 - - def _certificate(self) -> bytes: - """Generate a canonical certificate for a graph. - - Based on the permutation ``self.best_perm`` that minimizes the binary value - of the upper triangular portion of the adjacency matrix of the graph, - as found by comparing leaf nodes of the search tree during the search for - automorphisms. - - Returns: - cert_hash: a hash object of the canonical adjacency bitstring. - """ - cert_hash = hashlib.sha256() - neighbours = self._neighbours - best_perm = self._best_perm - - for j in range(1, self._num_nodes): - neighbours_best_j = neighbours[best_perm[j]] - - for i in range(j): - bit = 1 if best_perm[i] in neighbours_best_j else 0 - cert_hash.update(bytes([bit])) - - return cert_hash.digest() - - def _initial_partition( - self, - graph_coloring=None - ) -> tuple[list[set[int] | None], np.ndarray, np.ndarray, int]: - """Initialize the initial partition for a graph. - - Currently this only supports graphs whose vertices are initially the same - color, but could be expanded in the future to accommodate graphs with a - non-trivial initial vertex coloring. - - Returns: - partition: The initial partition structure, represented as a list of sets of vertices - ordered by color. - trace: A list of the sizes of each partition cell (color class), ordered by color. - color: An array mapping each vertex to its current color. - num_colors: The number of colors in the initial partition. - """ - trace = np.zeros(self._num_nodes, dtype=self._color_dtype) - color = np.zeros(self._num_nodes, dtype=self._color_dtype) - - if graph_coloring is None: - partition = [set(self._nodes)] + [None] * (self._num_nodes - 1) - trace[0] = self._num_nodes - num_colors = 1 - else: - num_colors = max(graph_coloring.values()) + 1 - partition = [set() for _ in range(num_colors)] + [None] * (self._num_nodes - num_colors) - for node_label, node_color in graph_coloring.items(): - node_index = self.node_to_index[node_label] # TODO ensure this works with colors not starting at 0..n-1 - partition[node_color].add(node_index) - - for i, nodes in enumerate(partition[:num_colors]): - color[list(nodes)] = i - trace[i] = len(nodes) - - #num_colors = i + 1 - #partition = partition + [None] * (self._num_nodes - num_colors) - - return partition, trace, color, num_colors - - -def vertex_orbits( - u_vector: list[list[NDArray[np.intp]]], - nodes: list[int], - index_to_node: Mapping[int, int] | None = None, -) -> list[list[int]]: - """Calculate vertex orbits using breadth-first search. - - If ``u_vector`` contains no coset representatives, trivial orbits are returned. - - Args: - u_vector: Coset representatives grouped by stabilizer index. - nodes: List of vertex indices used to return trivial orbits when ``u_vector`` is empty. - index_to_node: An optional dictionary for returning orbits with their original node labels. - - Returns: - A list of orbits, each orbit is a list of vertex indices. - - Example: - >>> import numpy as np - >>> from dwave.experimental.automorphism import vertex_orbits - ... - >>> u_vector = [ - ... [np.array([0, 1, 4, 3, 2, 6, 5, 7])], - ... [np.array([2, 1, 4, 3, 0, 7, 5, 6]), np.array([4, 1, 0, 3, 2, 6, 7, 5])], - ... [np.array([0, 3, 2, 1, 4, 5, 6, 7])], - ... ] - >>> nodes = list(range(8)) - >>> vertex_orbits(u_vector, nodes) - [[0, 2, 4], [1, 3], [5, 6, 7]] - """ - if not u_vector: - return [[x] for x in nodes] - - if not all(isinstance(sublist, list) for sublist in u_vector): - raise ValueError("u_vector must be a list of lists.") - - if isinstance(nodes, np.ndarray): - nodes = nodes.tolist() - - if not isinstance(nodes, list) or not all(isinstance(n, int) for n in nodes): - raise ValueError("nodes must be a list of integers.") - - visited = set() - orbits = [] - num_nodes = len(nodes) - generators = [g for u_vector_i in u_vector for g in u_vector_i] - generators.append(np.arange(num_nodes)) - label = (lambda x: index_to_node[x]) if index_to_node is not None else int - - for v_start in nodes: - if v_start in visited: - continue - - visited.add(v_start) - orb = [label(v_start)] - - q = deque([v_start]) - while q: - v_current = q.popleft() - - for g in generators: - v_current = g[v_current] - if v_current not in visited: - visited.add(v_current) - q.append(v_current) - orb.append(label(v_current)) - orb.sort() - orbits.append(orb) - - orbits.sort() - return orbits - - -def edge_orbits( - u_vector: list[list[NDArray[np.intp]]], - edges: list[tuple[int, int]], - index_to_node: Mapping[int, int] | None = None, -) -> list[list[int]]: - """Calculate edge orbits using breadth-first search. - - Args: - u_vector: Coset representatives grouped by stabilizer index. - edges: List of graph edges as tuples of vertex index pairs. - - Returns: - A list of orbits, each orbit is a list of edges (tuples of vertex index pairs). - - Example: - >>> import numpy as np - >>> from dwave.experimental.automorphism import edge_orbits - ... - >>> u_vector = [ - ... [np.array([0, 1, 4, 3, 2, 6, 5, 7])], - ... [np.array([2, 1, 4, 3, 0, 7, 5, 6]), np.array([4, 1, 0, 3, 2, 6, 7, 5])], - ... [np.array([0, 3, 2, 1, 4, 5, 6, 7])], - ... ] - >>> edges = [ - ... (0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6), - ... (6, 7), (7, 0), (0, 3), (1, 4), (2, 6), (5, 7) - ... ] - >>> orbits = edge_orbits(u_vector, edges) - >>> orbits[0] - [(0, 1), (0, 3), (1, 2), (1, 4), (2, 3), (3, 4)] - >>> orbits[1:] - [[(0, 7), (2, 6), (4, 5)], [(5, 6), (5, 7), (6, 7)]] - """ - if not u_vector: - return [[x] for x in edges] - - if not all(isinstance(sublist, list) for sublist in u_vector): - raise ValueError("u_vector must be a list of lists.") - - if not isinstance(edges, list) or not all(isinstance(e, tuple) for e in edges): - raise TypeError("edges must be a list of tuples") - - visited = set() - orbits = [] - generators = [g for u_vector_i in u_vector for g in u_vector_i] - label = (lambda x: index_to_node[x]) if index_to_node is not None else int - - for u_start, v_start in edges: - e_start = (u_start, v_start) if u_start < v_start else (v_start, u_start) - - if e_start in visited: - continue - - visited.add(e_start) - orb = [tuple(label(x) for x in e_start)] - - q = deque([e_start]) - while q: - u, v = q.popleft() - for g in generators: - e_current = (g[u], g[v]) if g[u] < g[v] else (g[v], g[u]) - - if e_current not in visited: - visited.add(e_current) - q.append(e_current) - orb.append(tuple(label(x) for x in e_current)) - - orb.sort() - orbits.append(orb) - - orbits.sort() - return orbits - - -def sample_automorphisms( - u_vector: list[list[NDArray[np.intp]]], - num_samples: int = 1, - seed: int | None = None, -) -> list[NDArray[np.intp]]: - """Uniformly sample automorphisms from the Schreier-Sims representation. - - Randomly samples one coset representative from each non-trivial left - transversal and takes the product, guaranteeing uniform sampling. The - automorphisms can be composed uniformly regardless of the ordering of - the left transversals in 'u_vector'. All products involving identity - automorphisms are ignored. - - Args: - u_vector: Coset representatives grouped by stabilizer index. - num_samples: The number of automorphisms to return. - seed: Random seed for reproducibility. - - Returns: - A list of uniformly sampled automorphisms in one-line notation. - - Example: - >>> import networkx as nx - >>> from dwave.experimental.automorphism import schreier_rep, sample_automorphisms - ... - >>> graph = nx.cycle_graph(8) - >>> result = schreier_rep(graph) - >>> sample_automorphisms(result.u_vector, seed=42) - [array([3, 4, 5, 6, 7, 0, 1, 2])] - >>> sample_automorphisms(result.u_vector, num_samples=2, seed=42) - [array([3, 4, 5, 6, 7, 0, 1, 2]), array([6, 5, 4, 3, 2, 1, 0, 7])] - """ - rng = np.random.default_rng(seed) - num_nodes = len(u_vector[0][0]) - u_counts = [len(u_i) for u_i in u_vector] - sampled_automorphisms = [] - - for _ in range(num_samples): - sample_indices = rng.integers(low=-1, high=u_counts) - g_product = np.arange(num_nodes) - - for i, u_i in enumerate(u_vector): - if sample_indices[i] >= 0: - g = u_i[sample_indices[i]] - g_product = mult(g, g_product) - - sampled_automorphisms.append(g_product) - - return sampled_automorphisms - - -def mult(alpha: NDArray[np.intp], beta: NDArray[np.intp]) -> NDArray[np.intp]: - """Compose two permutations in one-line notation, alpha after beta. - - Args: - alpha: A permutation represented as a list of integers in one-line notation. - beta: Another permutation of the same length. - - Returns: - The composition alpha ∘ beta in one-line notation. - - Example: - >>> import numpy as np - >>> from dwave.experimental.automorphism import mult - ... - >>> alpha = np.array([2,0,1], dtype=np.intp) # (0,2,1): 0->2, 1->0, 2->1 - >>> beta = np.array([1,2,0], dtype=np.intp) # (0,1,2): 0->1, 1->2, 2->0 - >>> mult(alpha, beta) - array([0, 1, 2]) - """ - return alpha[beta] - - -def inv(n: int, alpha: NDArray[np.intp]) -> NDArray[np.intp]: - """Calculate the inverse of a permutation in one-line notation. - - Args: - n: Length of permutation alpha. - alpha: A permutation represented as a list of integers in one-line notation. - - Returns: - The inverse of alpha in one-line notation. - - Example: - >>> import numpy as np - >>> from dwave.experimental.automorphism import inv - ... - >>> alpha = np.array([2,0,1], dtype=np.intp) # (0,2,1): 0->2, 1->0, 2->1 - >>> inv(3, alpha) - array([1, 2, 0]) - """ - alpha_inv = np.empty(n, dtype=np.intp) - alpha_inv[alpha] = np.arange(n, dtype=alpha_inv.dtype) - return alpha_inv - - -def schreier_rep( - graph: nx.Graph, - num_samples: int | None = None, - seed: int = 42, - graph_coloring=None -) -> SchreierContext: - """Compute Schreier representatives and orbits for a graph. - - Builds a depth-first search tree, iteratively performing color refinement - and vertex individualization until leaf nodes are reached where all graph - vertices are uniquely colored. Leaf nodes with identical adjacency matrices - represent graph automorphisms. Discovered automorphisms are used to prune - the search tree. - - If graphs have more than one component, automorphisms are found for each - individual component, and automorphisms between components are determined - by considering which components are isomorphic. Since the number of automorphisms - between isomorphic components scales factorially with the number of components, - this is significantly faster than naively performing refinement-individualization - over the whole graph. It would be possible to update ``u_vector`` directly - without using ``enter()``, which in principle should be even faster, and should - be the first place to look if further performance improvements are required. - - Args: - graph: A NetworkX Graph object representing the input graph containing - the following methods: - - ``nodes()``: iterable of all nodes - - ``number_of_nodes()``: total number of nodes - - ``edges()``: iterable of all edges - - ``neighbors()``: iterable of all neighbours for a given node - num_samples: Number of samples to use for generating new coset representatives - from the existing set. If not provided, all coset representatives are used. - seed: Random seed for reproducibility. Defaults to 42. - """ - if nx.number_connected_components(graph) == 1: - ctx = SchreierContext(graph, num_samples=num_samples, seed=seed) - initial_partition, trace, color, num_colors = ctx._initial_partition(graph_coloring) - - ctx._canon(initial_partition, trace, color, num_colors) - return ctx - - # relabel vertices so components have contiguous labels - index_to_node = {} - node_to_index = {} - next_label = 0 - - component_vertices = list(nx.connected_components(graph)) - for vertices in component_vertices: - for vertex in sorted(vertices): - node_to_index[vertex] = next_label - index_to_node[next_label] = vertex - next_label += 1 - - graph = nx.relabel_nodes(graph, node_to_index, copy=True) - - # enter component automorphisms into global graph - ctx = SchreierContext(graph, num_samples=num_samples, seed=seed) - ctx._index_to_node = index_to_node - ctx._node_to_index = node_to_index - - # group isomorphic components together - components = [ctx._graph.subgraph(c).copy() for c in nx.connected_components(ctx._graph)] - - unique_components = {} - for comp in components: - ctx_comp = SchreierContext(comp, num_samples=num_samples, seed=seed) - ctx_comp._compare_adj = True - - initial_partition, trace, color, num_colors = ctx_comp._initial_partition() - ctx_comp._canon(initial_partition, trace, color, num_colors) - - ctx._nodes_reached += ctx_comp.nodes_reached # update the global search tree statistics - ctx._leaf_nodes += ctx_comp.leaf_nodes - - unique_components.setdefault(ctx_comp._certificate(), []).append( - ComponentInfo(ctx_comp._u_vector, np.array(sorted(comp.nodes())), ctx_comp._best_perm) - ) - - # enter the local automorphisms - graph_nnodes = ctx._graph.number_of_nodes() - for identical_components in unique_components.values(): - for comp in identical_components: - for u in chain.from_iterable(comp.u_vector): - u_global = np.arange(graph_nnodes) - u_global[comp.nodes] = u_global[comp.nodes][u] - ctx._enter(u_global, mode=EnterMode.NO_RECURSE) - - # enter swap automorphisms - for comps in unique_components.values(): - for i in range(len(comps) - 1): - i_nodes = comps[i].nodes - j_nodes = comps[i + 1].nodes - - # swap automorphisms must be entered in the canonical basis - i_canon_perm = comps[i].best_perm - j_canon_perm = comps[i + 1].best_perm - i_canon = i_nodes[i_canon_perm] - j_canon = j_nodes[j_canon_perm] - - u_global = np.arange(graph_nnodes) - u_global[i_canon], u_global[j_canon] = u_global[j_canon], u_global[i_canon] - ctx._enter(u_global, mode=EnterMode.RECURSE_ONCE) - - return ctx - -def array_to_cycle( - array: NDArray[np.intp], - index_to_node: Mapping[int, Hashable] | None = None -) -> str: - """Convert an array in one-line notation to a string in cycle notation. - - Based on Algorithm 6.4 from Kreher, D. L., & Stinson, D. R. (1999). - Combinatorial algorithms: Generation, enumeration, and search. - - Args: - array: The permutation in one-line notation. - index_to_node: An optional relabelling dictionary. By default, array indices - are used. - - Returns: - The permutation as a string in cycle notation. - - Example: - >>> import numpy as np - >>> from dwave.experimental.automorphism import array_to_cycle - ... - >>> alpha = np.array([2,0,1], dtype=np.intp) # (0,2,1): 0->2, 1->0, 2->1 - >>> array_to_cycle(alpha) - '(0,2,1)' - >>> array_to_cycle(np.array([2,0,1]), index_to_node={0: 5, 1: 7, 2: 9}) - '(5,9,7)' - """ - if index_to_node is not None: - expected = set(range(len(array))) - if index_to_node.keys() != expected: - missing = expected - index_to_node.keys() - raise ValueError(f"index_to_node missing keys: {missing}") - - label = (lambda x: str(index_to_node[x])) if index_to_node is not None else str - unvisited = [True] * len(array) - cycle_parts = [] - - for i in range(len(array)): - if unvisited[i]: - cycle_parts.append('(') - cycle_parts.append(label(i)) - unvisited[i] = False - j = i - - while unvisited[array[j]]: - cycle_parts.append(',') - j = array[j] - cycle_parts.append(label(j)) - unvisited[j] = False - - cycle_parts.append(')') - return ''.join(cycle_parts) diff --git a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py index 1e4cffa..a33b0ba 100644 --- a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -90,7 +90,7 @@ def get_chain_connectivity( 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. diff --git a/dwave/experimental/lattice_utils/lattice/lattice.py b/dwave/experimental/lattice_utils/lattice/lattice.py index af59079..a724923 100644 --- a/dwave/experimental/lattice_utils/lattice/lattice.py +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -126,7 +126,6 @@ def embed_lattice( if len(embeddings) >= min_number_of_embeddings and np.prod(embeddings.shape): self._save_embeddings(sampler, embeddings) - #def make_nominal_bqm(self, **kwargs) -> dimod.BQM: def make_nominal_bqm(self) -> dimod.BQM: """Construct a default nominal BQM coupling strength values set to +1. diff --git a/dwave/experimental/lattice_utils/lattice/orbits.py b/dwave/experimental/lattice_utils/lattice/orbits.py index 30f6b80..40852c3 100644 --- a/dwave/experimental/lattice_utils/lattice/orbits.py +++ b/dwave/experimental/lattice_utils/lattice/orbits.py @@ -19,7 +19,7 @@ import numpy as np from numpy.typing import NDArray -from dwave.experimental.lattice_utils.lattice.automorphism import schreier_rep +from dwave.experimental.automorphism import schreier_rep __all__ = [ 'reindex', @@ -186,7 +186,7 @@ def get_unsigned_bqm_orbits( each original variable to its orbit index and ``coupler_orbits`` maps each coupling to its orbit index. """ - coupler_orbits={} + 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}")], diff --git a/dwave/experimental/lattice_utils/lattice/triangular.py b/dwave/experimental/lattice_utils/lattice/triangular.py index fe18ad2..6ec42e9 100644 --- a/dwave/experimental/lattice_utils/lattice/triangular.py +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -28,10 +28,10 @@ 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. @@ -176,7 +176,7 @@ def generate_edges(self) -> Iterator[tuple[int, int]]: 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. diff --git a/dwave/experimental/lattice_utils/observable/kinks.py b/dwave/experimental/lattice_utils/observable/kinks.py index a9583e2..9f879b8 100644 --- a/dwave/experimental/lattice_utils/observable/kinks.py +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -30,21 +30,31 @@ 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.""" - sample_array = dimod.as_samples(sample_set)[0] + """Compute the kink-kink correlator for 1D spin chains. - S = np.roll(sample_array, 1, axis=1) - K = np.multiply(S, sample_array) == np.sign(experiment.param["energy_scale"]) - L = K.shape[-1] - K = np.reshape(K, (-1, L)) - kink_density = np.mean(K) + 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. - CKK = np.zeros((K.shape[-1],)) + Returns: + A numpy array containing the kink-kink correlator values for each sample. + """ + samples = dimod.as_samples(sample_set)[0] - for R in range(1, L): - KR = np.roll(K, R, axis=1) - CKK[R] = np.mean(np.multiply(K, KR)) - np.power(np.mean(K), 2) + 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) - CKK /= kink_density**2 + kink_kink_correlator = np.zeros((kink_mask.shape[-1],)) - return CKK + 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 index e7380d8..ac2f5ff 100644 --- a/dwave/experimental/lattice_utils/observable/observable.py +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -15,7 +15,7 @@ from __future__ import annotations from abc import ABC, abstractmethod from pathlib import Path -from typing import Any +from typing import Any, TypeAlias import numpy as np from numpy.typing import NDArray @@ -33,13 +33,14 @@ 'ReferenceEnergy', ] +ObservableResult: TypeAlias = NDArray | float | int | tuple[NDArray, tuple[int, int]] + class Observable(ABC): - """The observable class does not take any parameters. Its primary - functionality is through the required 'evaluate' method, which requires - parameters 'experiment' and 'bqm' defining the context, and 'sample_set' - which provides the samples on which we compute the observable. Output is a - numpy array of arbitrary type (usually float). + """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): @@ -51,7 +52,7 @@ def evaluate( experiment: Experiment, bqm: dimod.BQM, sample_set: dimod.SampleSet, - ) -> NDArray: + ) -> ObservableResult: raise NotImplementedError @@ -79,10 +80,9 @@ def evaluate( ) -> NDArray: sample_array = dimod.as_samples(sample_set)[0].astype(float) if len(experiment.inst.edge_list) == 0: - return [] - row, col = np.asarray(experiment.inst.edge_list).T + return np.empty(0, dtype=float) - # Surprisingly, it's faster to multiply the whole matrix. + 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 @@ -98,10 +98,9 @@ def evaluate( ) -> NDArray: sample_array = dimod.as_samples(sample_set)[0].astype(float) if len(experiment.inst.edge_list) == 0: - return [] - row, col = np.asarray(experiment.inst.edge_list).T + return np.empty(0, dtype=float) - # Surprisingly, it's faster to multiply the whole matrix. + 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] @@ -155,8 +154,10 @@ def evaluate( path: str | Path | None = None, inst: Lattice | None = None, ) -> float: - - if path is None: + """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(): @@ -166,8 +167,12 @@ def evaluate( # And if we can't load, we generate a reference sample. if experiment is not None: energy, sample, method_string = experiment.inst._optimize(bqm) - else: + 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) @@ -180,8 +185,11 @@ def load( path: str | Path | None = None, ) -> tuple[float, NDArray, str]: """Load and get the full data tuple, not just the energy.""" - if path is None: + 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()) @@ -192,6 +200,7 @@ def load( 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="") @@ -207,8 +216,10 @@ def update( 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. """ - reference_energy, _, reference_method_string = self.load(experiment, bqm, path) + 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: @@ -228,12 +239,19 @@ def get_reference_energy_path( ) -> Path: """Return the path to the reference energy file for the given experiment and BQM. - This needs to be fixed if you have something not in the instance - pathstring that needs to be taken into account, for example if the ground-state - energies depend on the chip. + 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 # defunct. + 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. @@ -258,7 +276,7 @@ def get_reference_energy_path( / experiment_data_dict["inst"]._get_instance_pathstring() ) - # Use hash. BQM is not hashable so use the experiment.inst data to generate a tuple. + # 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 ) diff --git a/dwave/experimental/lattice_utils/observable/triangular.py b/dwave/experimental/lattice_utils/observable/triangular.py index 89ef7b2..b6faac9 100644 --- a/dwave/experimental/lattice_utils/observable/triangular.py +++ b/dwave/experimental/lattice_utils/observable/triangular.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""For triangular order parameters""" - from __future__ import annotations import numpy as np from numpy.typing import NDArray @@ -26,7 +24,7 @@ class TriangularOP(Observable): - """For triangular lattices. Unembeds if possible.""" + """For calculating the order parameter of triangular lattices.""" def evaluate( self, @@ -34,12 +32,24 @@ def evaluate( 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"): - # If the lattice is an embedded lattice lbqm = experiment.inst.unembed_bqm(bqm) - # unembed the sample set. lss = experiment.inst.unembed_sampleset(sample_set) triangular_sublattice = experiment.inst.logical_lattice.sublattice else: @@ -48,8 +58,12 @@ def evaluate( sample_array = dimod.as_samples(lss)[0] - for edge in lbqm.quadratic: - assert triangular_sublattice[edge[0]] != triangular_sublattice[edge[1]] + 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): @@ -58,6 +72,6 @@ def evaluate( ) angles = np.array(np.exp([0.0, 1.0j * 4 * np.pi / 3, 1.0j * 2 * np.pi / 3])).T - op = np.matmul(sublattice_mags, angles).ravel() / np.sqrt(3) + order_parameter = np.matmul(sublattice_mags, angles).ravel() / np.sqrt(3) - return op + return order_parameter diff --git a/examples/example_1D_Ising_chain.py b/examples/example_1D_Ising_chain.py index e433ced..f9ea39f 100644 --- a/examples/example_1D_Ising_chain.py +++ b/examples/example_1D_Ising_chain.py @@ -34,12 +34,11 @@ # Two samplers: an Advantage2 prototype and an Advantage system. samplers = [ - DWaveSampler(solver="Advantage2_system3.1"), + DWaveSampler(solver="Advantage2_system1"), DWaveSampler(solver="Advantage_system4.1"), ] NUM_SPINS = 256 -#NUM_SPINS = 8 # Two energy scales: one strong coupling and one weak coupling. ENERGY_SCALES = (-1.8, 0.1) @@ -79,23 +78,21 @@ # 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 + # 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: - exp = experiment.FastAnnealExperiment( - inst=inst, - sampler=sampler, - loop_data_files=30, - max_iterations=5, + 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 (CITE). The observable + # 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()) @@ -103,7 +100,7 @@ # Make parameter list. We will only vary anneal time. parameter_list = [{"anneal_time": time} for time in ANNEAL_TIMES] - for _ in range(20): + for _ in range(20): #TODO clean this up? done = exp.run_iteration(parameter_list) if done: break @@ -122,7 +119,7 @@ fbshim = [] # flux bias shim kkc = [] # kink-kink correlator for param in parameter_list: - exp.apply_param(param) + 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])) diff --git a/tests/test_lattice_utils.py b/tests/test_lattice_utils.py index 7a7df43..96198e1 100644 --- a/tests/test_lattice_utils.py +++ b/tests/test_lattice_utils.py @@ -690,22 +690,6 @@ def test_apply_param_no_anneal_or_schedule_raises(self): with self.assertRaises(ValueError): exp.apply_param({"energy_scale": 1.0}) - def test_spin_reversal_disabled(self): - chain = Chain(dimensions=(4,), periodic=(True,)) - sampler = _make_mock_sampler() - exp = Experiment(chain, sampler) - self.assertIsNone(exp._get_spin_reversal_transform()) - - def test_spin_reversal_enabled_with_seed(self): - chain = Chain(dimensions=(4,), periodic=(True,)) - sampler = _make_mock_sampler(num_qubits=8) - exp = Experiment(chain, sampler) - exp.param["spin_reversal_transform"] = True - exp.param["spin_reversal_transform_seed"] = 42 - srt1 = exp._get_spin_reversal_transform() - srt2 = exp._get_spin_reversal_transform() - self.assertEqual(srt1, srt2) - def test_initial_shim_no_embeddings(self): chain = Chain(dimensions=(4,), periodic=(True,)) sampler = _make_mock_sampler() @@ -929,7 +913,6 @@ def test_make_bqm_with_embeddings(self): "total_iterations": 0, "relative_coupler_strength": np.ones((1, chain.num_edges)), } - sc.spin_reversal_transform = None bqm = exp._make_bqm(sc) self.assertGreater(len(bqm.quadratic), 0) @@ -996,22 +979,6 @@ def test_coupler_shim_basic_update(self): exp._update_coupler_shim(sc, results) self.assertEqual(sc.shimdata["relative_coupler_strength"].shape, (1, chain.num_edges)) - def test_parse_results_with_spin_reversal(self): - chain = Chain(dimensions=(4,), periodic=(False,)) - chain.embedding_list = np.array([[0, 1, 2, 3]]) - sampler = _make_mock_sampler() - exp = Experiment(chain, sampler) - exp.run_index = 0 - bqm = chain.make_nominal_bqm() - samples = np.ones((10, 4)) - ss = dimod.SampleSet.from_samples_bqm(samples, bqm) - - sc = SamplerCall(run_index=0) - sc.nominal_bqms = [bqm] - sc.spin_reversal_transform = {0: True, 1: False, 2: True, 3: False} - results = exp.parse_results(sc, ss) - self.assertIn("QubitMagnetization", results) - def test_get_shimdata_not_initialized(self): chain = Chain(dimensions=(4,), periodic=(True,)) sampler = _make_mock_sampler() From f364d2d8c8867345d66254696651502b78b7e5ac Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 8 Jun 2026 11:45:01 -0700 Subject: [PATCH 07/13] Add progress bars and missing docstrings -added progress bars using TQDM -added docstrings to the evaluation methods in observable.py -misc cleanup --- dwave/experimental/lattice_utils/__init__.py | 2 +- .../lattice_utils/experiment/experiment.py | 141 +++++++++++++++--- .../lattice_utils/lattice/chain.py | 8 +- .../lattice_utils/lattice/embedded_lattice.py | 9 +- .../lattice_utils/lattice/lattice.py | 16 +- .../lattice_utils/lattice/optimize.py | 2 +- .../lattice_utils/lattice/orbits.py | 13 +- .../lattice_utils/lattice/triangular.py | 40 ++--- .../lattice_utils/observable/kinks.py | 5 +- .../lattice_utils/observable/observable.py | 95 ++++++++++-- .../lattice_utils/observable/triangular.py | 2 +- examples/example_1D_Ising_chain.py | 39 +++-- examples/example_2D_geometric_frustration.py | 4 +- examples/mca_shim_AO_FB.py | 7 +- 14 files changed, 286 insertions(+), 97 deletions(-) diff --git a/dwave/experimental/lattice_utils/__init__.py b/dwave/experimental/lattice_utils/__init__.py index d8cb8ad..fb1e903 100644 --- a/dwave/experimental/lattice_utils/__init__.py +++ b/dwave/experimental/lattice_utils/__init__.py @@ -12,6 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from . import experiment, lattice, observable +from dwave.experimental.lattice_utils import experiment, lattice, observable __all__ = ["experiment", "lattice", "observable"] diff --git a/dwave/experimental/lattice_utils/experiment/experiment.py b/dwave/experimental/lattice_utils/experiment/experiment.py index 331e81e..c2dae1c 100644 --- a/dwave/experimental/lattice_utils/experiment/experiment.py +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -24,6 +24,7 @@ import dimod import numpy as np +from tqdm.auto import tqdm from dwave.experimental.lattice_utils.lattice import Lattice from dwave.experimental.lattice_utils.observable import ( @@ -38,18 +39,10 @@ __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 +DW_TEAL = "#17bebb" +DW_BLUE = "#2a7de1" +DW_ORANGE = "#f37820" + @dataclass class ExperimentConfig: @@ -171,7 +164,12 @@ def apply_param(self, param: dict[str, float]) -> None: 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: + def run_iteration( + self, + parameter_list: list, + progress: bool = False, + scaling_factor: float = 1.0, + ) -> bool: """Run one experiment iteration for each parameter set in ``parameter_list``. For each parametrization, this method applies the parameters, builds the @@ -180,26 +178,36 @@ def run_iteration(self, parameter_list: list, **kwargs) -> bool: Args: parameter_list: List of parameter dictionaries to run. + progress: If true, displays a progress bar for waiting on results. + scaling_factor: A multiplicative factor to apply to the BQM before sampling. Returns: A boolean value corresponding to whether or not the experiment is finished. """ try: - self.inst._load_embeddings(self.sampler, **kwargs) + self.inst._load_embeddings(self.sampler) 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()})" + tqdm.write( + f"\n{type(self.inst).__name__}={self.inst.dimensions}, " + f"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 = {} + create_bar = self._make_progress_bar( + total=len(parameter_list), + desc="Creating sampler calls", + colour=DW_BLUE, + enabled=progress, + ) + for index, param in enumerate(parameter_list): self.apply_param(param) call_dict[index] = self._build_sampler_call() @@ -207,14 +215,29 @@ def run_iteration(self, parameter_list: list, **kwargs) -> bool: call_dict.pop(index) else: response_dict[index] = self.sampler.sample( - call_dict[index].bqm * kwargs.get("scaling_factor", 1.0), + call_dict[index].bqm * scaling_factor, **call_dict[index].sampler_params, ) + if create_bar is not None: + create_bar.update() + + if create_bar is not None: + create_bar.close() if len(call_dict) == 0: - print(f"***\n***\nFINISHED for all {len(parameter_list)} parameterizations.\n***\n***") + if progress: + tqdm.write( + f"***\n***\nFINISHED for all {len(parameter_list)} parameterizations.\n***\n***" + ) return True + wait_bar = self._make_progress_bar( + total=len(call_dict), + desc=" Awaiting/parsing data", + colour=DW_TEAL, + enabled=progress, + ) + # Get and manage all the results while response_dict: made_progress = False @@ -227,6 +250,8 @@ def run_iteration(self, parameter_list: list, **kwargs) -> bool: self._update_shim(call_dict[index], results) savedata = self._generate_data_to_save(call_dict[index], results) self._save_results(savedata, quiet=True) + if wait_bar is not None: + wait_bar.update() del response_dict[index] made_progress = True break @@ -234,6 +259,10 @@ def run_iteration(self, parameter_list: list, **kwargs) -> bool: if not made_progress: time.sleep(0.1) # Waiting for results to come in + if wait_bar is not None: + wait_bar.close() + + self._print_iteration_status(call_dict, len(parameter_list), enabled=progress) return False def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[str, Any]: @@ -273,6 +302,78 @@ def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[st return results + def _make_progress_bar( + self, + *, + total: int, + desc: str, + colour: str, + enabled: bool, + bar_format: str | None = None, + initial: int | float = 0, + ) -> tqdm | None: + """Create a tqdm progress bar with consistent formatting.""" + if not enabled: + return None + + if bar_format is None: + bar_width = min(100, max(total, 20)) + bar_format = f"{{desc}}: |{{bar:{bar_width}}}{{r_bar}}{{bar:-{bar_width}b}}" + + return tqdm( + total=total, + initial=initial, + desc=desc, + bar_format=bar_format, + colour=colour, + ) + + def _print_iteration_status( + self, + call_dict: dict[int, SamplerCall], + num_params: int, + enabled: bool, + ) -> None: + """Print a summary of the iteration status, including progress and iteration ranges.""" + if not enabled: + return + iteration_range = ( + f"Iteration range " + f"{min(call.shimdata['total_iterations'] for call in call_dict.values())}-" + f"{max(call.shimdata['total_iterations'] for call in call_dict.values())} " + ) + if self.max_iterations is None: + tqdm.write(" Total progress: " + iteration_range) + return + + total = num_params * self.max_iterations + progress_value = ( + sum(call.shimdata["total_iterations"] for call in call_dict.values()) + + (num_params - len(call_dict)) * self.max_iterations + ) + + progress_string = ( + f"{progress_value / total * 100:.1f}% " + f"Iteration range " + f"{min(call.shimdata['total_iterations'] for call in call_dict.values())}-" + f"{max(call.shimdata['total_iterations'] for call in call_dict.values())} " + f"of {self.max_iterations} " + f"({num_params - len(call_dict)} of {num_params} parameters finished)" + ) + + bar_width = min(100, max(num_params, 20)) + bar_format = f"{{desc}}: |{{bar:{bar_width}}}| {progress_string}" + + total_bar = self._make_progress_bar( + total=total, + desc=" Total progress", + bar_format=bar_format, + colour=DW_ORANGE, + enabled=enabled, + initial=progress_value, + ) + total_bar.close() + def _save_results( self, data_dict: dict[str, Any], diff --git a/dwave/experimental/lattice_utils/lattice/chain.py b/dwave/experimental/lattice_utils/lattice/chain.py index 246cad0..39391ab 100644 --- a/dwave/experimental/lattice_utils/lattice/chain.py +++ b/dwave/experimental/lattice_utils/lattice/chain.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections.abc import Iterator +from collections.abc import Generator from pathlib import Path from numpy.typing import NDArray @@ -27,7 +27,7 @@ class 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. + non-periodic (open chain) based on the ``periodic`` parameter. Args: dimensions: One-element tuple giving the number of spins in the chain. @@ -65,11 +65,11 @@ def __init__( coupler_orbits=coupler_orbits, ) - def generate_edges(self) -> Iterator[tuple[int, int]]: + def generate_edges(self) -> Generator[tuple[int, int]]: """Yield edges for a 1D chain lattice. Returns: - An iterator of tuples, where each tuple represents an edge between + A generator of tuples, where each tuple represents an edge between two spins in the chain. """ n = self.dimensions[0] diff --git a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py index a33b0ba..842ad46 100644 --- a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -15,7 +15,7 @@ from itertools import combinations, product from numbers import Integral -from collections.abc import Iterator, Hashable +from collections.abc import Generator, Hashable import dimod import numpy as np @@ -56,6 +56,7 @@ def __init__( 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"): @@ -76,7 +77,7 @@ def get_chain_connectivity( 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 + in the chains corresponding to u and v that should be connected. For a chain edge (u == v or v is None), this will return pairs of indices within the same chain. For a logical edge (u != v), this will return pairs of indices between the two chains. @@ -88,11 +89,11 @@ def get_chain_connectivity( # 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]]: + def generate_edges(self) -> Generator[tuple[Hashable, Hashable]]: """Yield physical edges for the embedded lattice. Returns: - An iterator of tuples, where each tuple represents an edge between + A generator of tuples, where each tuple represents an edge between two spins in the physical lattice. """ logical_bqm = self.logical_lattice.make_nominal_bqm() diff --git a/dwave/experimental/lattice_utils/lattice/lattice.py b/dwave/experimental/lattice_utils/lattice/lattice.py index a724923..63db7a3 100644 --- a/dwave/experimental/lattice_utils/lattice/lattice.py +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -14,7 +14,7 @@ import os from pathlib import Path -from collections.abc import Iterator, Hashable +from collections.abc import Generator, Hashable from abc import ABC, abstractmethod import warnings @@ -66,6 +66,12 @@ def __init__( self.data_root = Path(data_root) self.periodic = periodic if periodic is not None else tuple(False for _ in dimensions) + if len(self.periodic) != len(self.dimensions): + raise ValueError( + f"periodic and dimensions must have the same length: " + f"got {len(self.periodic)} and {len(self.dimensions)}." + ) + self.edge_list: list[tuple[Hashable, Hashable]] = list(self.generate_edges()) if not hasattr(self, "num_spins"): @@ -76,9 +82,8 @@ def __init__( self.initialize_orbits(qubit_orbits, coupler_orbits) @abstractmethod - def generate_edges(self) -> Iterator[tuple[Hashable, Hashable]]: + def generate_edges(self) -> Generator[tuple[Hashable, Hashable]]: """Yield the edges for this lattice.""" - raise NotImplementedError def embed_lattice( self, @@ -87,7 +92,7 @@ def embed_lattice( timeout: int = 10, max_number_of_embeddings: int | None = None, min_number_of_embeddings: int = 1, - exclude_qubits: list = [], + exclude_qubits: list | None = None, **kwargs, ) -> None: """Find or load embeddings onto the sampler graph. @@ -102,6 +107,9 @@ def embed_lattice( exclude_qubits: Qubits to remove from the sampler graph before searching for embeddings. """ + if exclude_qubits is None: + exclude_qubits = [] + graph_bqm = dimod.to_networkx_graph(self.make_nominal_bqm()) graph_sampler = sampler.to_networkx_graph() graph_sampler.remove_nodes_from(exclude_qubits) diff --git a/dwave/experimental/lattice_utils/lattice/optimize.py b/dwave/experimental/lattice_utils/lattice/optimize.py index 1a91f38..11943d8 100644 --- a/dwave/experimental/lattice_utils/lattice/optimize.py +++ b/dwave/experimental/lattice_utils/lattice/optimize.py @@ -88,7 +88,7 @@ def optimize_increasing_sa_sweeps( Returns: A tuple containing the best energy found, the corresponding sample as a - NumPy array, and a string indicating the optimization method used. + NumPy array, and a string indicating the optimization method used. """ sa = SimulatedAnnealingSampler() diff --git a/dwave/experimental/lattice_utils/lattice/orbits.py b/dwave/experimental/lattice_utils/lattice/orbits.py index 40852c3..cdcc812 100644 --- a/dwave/experimental/lattice_utils/lattice/orbits.py +++ b/dwave/experimental/lattice_utils/lattice/orbits.py @@ -13,6 +13,7 @@ # limitations under the License. from collections.abc import Hashable +from collections import defaultdict import dimod import networkx as nx @@ -37,7 +38,7 @@ def reindex(mapping: dict[Hashable, int]) -> dict[Hashable, int]: mapping: Dictionary whose values represent indices or labels. Returns: - A new dictionary with the same keys as `mapping` but with values reindexed + 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()))} @@ -113,9 +114,9 @@ def get_bqm_orbits( 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())} + mapping_h = defaultdict(list) + mapping_mp = defaultdict(list) + mapping_J = defaultdict(list) for p, q in bqm.linear.items(): mapping_h[q].append(f"hnode_{p}") @@ -213,8 +214,8 @@ def get_orbits(bqm: dimod.BQM, edge_list: list[int, int]) -> tuple[NDArray, NDAr """Provide a bqm and receive a set of usable orbits derived from the signed BQM. Args: - bqm: Ising model to analyze - edge_list + bqm: Ising model to analyze. + edge_list: List of edges from the original unsigned BQM. Returns: A tuple ``(qubit_orbits_array, coupler_orbits_array)`` where diff --git a/dwave/experimental/lattice_utils/lattice/triangular.py b/dwave/experimental/lattice_utils/lattice/triangular.py index 6ec42e9..b83622e 100644 --- a/dwave/experimental/lattice_utils/lattice/triangular.py +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections.abc import Iterator, Hashable +from collections.abc import Generator, Hashable from pathlib import Path import networkx as nx @@ -37,13 +37,16 @@ class Triangular(Lattice): 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. + 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. """ @@ -61,6 +64,7 @@ def __init__( ): 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] @@ -68,14 +72,6 @@ def __init__( 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." @@ -84,6 +80,14 @@ def __init__( raise ValueError( "For Triangular with periodic[1]=True, dimensions[1] must be divisible by 3." ) + super().__init__( + dimensions=dimensions, + periodic=periodic, + data_root=data_root, + orbit_type=orbit_type, + qubit_orbits=qubit_orbits, + coupler_orbits=coupler_orbits, + ) def coordinates(self, node: int) -> tuple[int, int]: """Return the coordinates of a node in the lattice given its index. @@ -121,7 +125,7 @@ def make_nominal_bqm(self) -> dimod.BQM: return bqm - def generate_edges(self) -> Iterator[tuple[int, int]]: + def generate_edges(self) -> Generator[tuple[int, int]]: """Yield edges for the triangular lattice and initialize coordinate attributes. y is the first dimension, x is the second. Edges are straight along @@ -129,7 +133,7 @@ def generate_edges(self) -> Iterator[tuple[int, int]]: not periodic. Returns: - An iterator of tuples, where each tuple represents an edge between + A generator of tuples, where each tuple represents an edge between two spins in the lattice. """ length_y, length_x = self.dimensions diff --git a/dwave/experimental/lattice_utils/observable/kinks.py b/dwave/experimental/lattice_utils/observable/kinks.py index 9f879b8..65495cc 100644 --- a/dwave/experimental/lattice_utils/observable/kinks.py +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -24,10 +24,7 @@ class KinkKinkCorrelator(Observable): - """For 1D chains.""" - - def __init__(self): - super().__init__() + """A class for computing the kink-kink correlator for 1D spin chains.""" def evaluate(self, experiment: Experiment, bqm: BQM, sample_set: SampleSet) -> NDArray: """Compute the kink-kink correlator for 1D spin chains. diff --git a/dwave/experimental/lattice_utils/observable/observable.py b/dwave/experimental/lattice_utils/observable/observable.py index 9e8c4b9..4cc9ebf 100644 --- a/dwave/experimental/lattice_utils/observable/observable.py +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -42,6 +42,7 @@ class Observable(ABC): 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__ @@ -52,7 +53,7 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> ObservableResult: - raise NotImplementedError + pass class QubitMagnetization(Observable): @@ -64,6 +65,16 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> NDArray: + """Return per-qubit mean spin values over the provided samples. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples used to compute the magnetization. + + Returns: + A numpy array containing the mean magnetization for each qubit. + """ sample_array = dimod.as_samples(sample_set)[0].astype(float) return np.mean(sample_array, axis=0) @@ -77,6 +88,16 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> NDArray: + """Return per-coupler pairwise spin correlations over the provided samples. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples used to compute the coupler correlations. + + Returns: + A numpy array containing the pairwise spin correlations for each coupler. + """ sample_array = dimod.as_samples(sample_set)[0].astype(float) if len(experiment.inst.edge_list) == 0: return np.empty(0, dtype=float) @@ -95,6 +116,16 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> NDArray: + """Return the mean coupler frustration over the provided samples. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples used to compute the mean coupler frustration. + + Returns: + A numpy array containing the mean coupler frustration for each edge. + """ sample_array = dimod.as_samples(sample_set)[0].astype(float) if len(experiment.inst.edge_list) == 0: return np.empty(0, dtype=float) @@ -120,11 +151,22 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> NDArray: + """Return signed sample energies from the sample set. + + Args: + experiment: Experiment context providing the ``energy_scale`` sign. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples containing energy data. + + Returns: + A numpy array containing the sample energies multiplied by the sign + of ``energy_scale``. + """ 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.""" + """Compute bitpacked spins.""" def evaluate( self, @@ -132,6 +174,17 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> tuple[NDArray, tuple[int, int]]: + """Return bitpacked spin samples and their original array shape. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples containing the spin values to unpack. + + Returns: + A tuple containing the bitpacked spin array and the original array + shape. + """ sample_array = dimod.as_samples(sample_set)[0] # Bitpack solutions @@ -153,11 +206,25 @@ def evaluate( path: str | Path | None = None, inst: Lattice | None = None, ) -> float: - """Get the reference energy for the given BQM, computing and caching it if needed.""" + """Get the reference energy for the given BQM, computing and caching it + if needed. + + Args: + experiment: The experiment for which to get the reference energy. Used + to determine the path for caching and loading the reference energy. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: The sample set is not used in this observable, but is + included in the signature for consistency with other observables. + path: Optional path to load/save the reference energy. If not provided, + a default path will be generated based on the experiment and BQM. + + Returns: + The reference energy for the given BQM. + """ if path is not None: path = Path(path) else: - path = get_reference_energy_path(experiment, bqm=bqm) + path = get_reference_energy_path(bqm, experiment) if path.exists(): energy, sample, method_string = self.load(experiment, bqm, path) @@ -187,7 +254,7 @@ def load( if path is not None: path = Path(path) else: - path = get_reference_energy_path(experiment, bqm=bqm) + path = get_reference_energy_path(bqm, experiment) with open(path, "r") as f: method_string = f.readline().strip() @@ -223,17 +290,16 @@ def update( if new_energy < reference_energy: if path is None: - path = get_reference_energy_path(experiment, bqm=bqm) + path = get_reference_energy_path(bqm, experiment) self.save(path, new_energy, sample, reference_method_string) - print(f"Updated energy from {reference_energy} to {new_energy}.") else: - raise ValueError + raise ValueError("New energy is not better than reference energy, not updating.") def get_reference_energy_path( + bqm: dimod.BQM, 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. @@ -242,18 +308,17 @@ def get_reference_energy_path( pathstring, for example when ground-state energies depend on the specific chip. Args: + bqm: The BQM for which to get the reference energy path. experiment: The experiment for which to get the reference energy path. root: Optional root directory to use instead of the experiment's data root. - bqm: The BQM for which to get the reference energy path. + dummy_experiment_data_dict: A dictionary containing the keys ``run_index``, + ``num_random_instances``, and ``inst`` to use when no experiment is + provided. This allows for generation of dummy experiment data without + all the overhead, for running without an actual experiment. Returns: The path to the reference energy file. """ - if 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: diff --git a/dwave/experimental/lattice_utils/observable/triangular.py b/dwave/experimental/lattice_utils/observable/triangular.py index b6faac9..9c9509e 100644 --- a/dwave/experimental/lattice_utils/observable/triangular.py +++ b/dwave/experimental/lattice_utils/observable/triangular.py @@ -24,7 +24,7 @@ class TriangularOP(Observable): - """For calculating the order parameter of triangular lattices.""" + """A class for calculating the order parameter of triangular lattices.""" def evaluate( self, diff --git a/examples/example_1D_Ising_chain.py b/examples/example_1D_Ising_chain.py index f9ea39f..158fc2b 100644 --- a/examples/example_1D_Ising_chain.py +++ b/examples/example_1D_Ising_chain.py @@ -34,8 +34,8 @@ # Two samplers: an Advantage2 prototype and an Advantage system. samplers = [ - DWaveSampler(solver="Advantage2_system1"), - DWaveSampler(solver="Advantage_system4.1"), + #DWaveSampler(solver="Advantage2_system1"), + DWaveSampler(solver="Advantage_system4"), ] NUM_SPINS = 256 @@ -86,22 +86,27 @@ config = experiment.FastAnnealExperimentConfig( energy_scale=energy_scale, coupler_shim_step=0.05, - flux_bias_shim_step=1e-6, + flux_bias_shim_step=3e-6, ) - exp = experiment.Experiment(inst=inst, sampler=sampler, max_iterations=5, config=config) + #exp = experiment.Experiment(inst=inst, sampler=sampler, max_iterations=5, config=config) + exp = experiment.Experiment(inst=inst, sampler=sampler, 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. + + # Note that only running 5 iterations will give poor results for the + # kink-kink correlator, which is fairly sensitive and requires a + # converged shim, but this is just for demonstration purposes. 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) + while True: + done = exp.run_iteration(parameter_list, progress=True) if done: break @@ -119,7 +124,7 @@ fbshim = [] # flux bias shim kkc = [] # kink-kink correlator for param in parameter_list: - exp.apply_param(param) #TODO clean this up? + exp.apply_param(param) res = exp.load_results(num_iterations=1000) frust.append(np.array([np.mean(it["CouplerFrustration"]) for it in res])) @@ -136,8 +141,8 @@ 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) + fig.tight_layout() + fig.subplots_adjust(hspace=0.35, wspace=0.3, top=0.9) ax = axes[0, 0] ax.loglog() @@ -173,7 +178,7 @@ **point_style, ) - ax.set_title("Kink density (with ~t_a^{-1/2} guideline)") + 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]) @@ -200,37 +205,37 @@ ax = axes[1, 0] ax.plot(fbshim[0]) - ax.set_title(f"Flux bias shim, t_a={ANNEAL_TIMES[0]:.3f}μs") + 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_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_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_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_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_title(f"Coupler shim, $t_a$={ANNEAL_TIMES[-1]:.3f}μs") ax.set_xlabel("Iteration") ax.grid(which="both", alpha=0.3) @@ -248,6 +253,7 @@ 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) +fig2.tight_layout(rect=(0.05, 0.05, 1, 0.95), w_pad=4.0) for isampler, sampler in enumerate(samplers): @@ -303,6 +309,7 @@ time_ns = ANNEAL_TIMES[0] * 1000 title = f"1D chain kink-kink correlator, {dims}, {time_ns:.1f} ns" fig3.suptitle(title, fontsize=16) +fig3.tight_layout(rect=(0.05, 0.05, 1, 0.95), w_pad=4.0) for isampler, sampler in enumerate(samplers): diff --git a/examples/example_2D_geometric_frustration.py b/examples/example_2D_geometric_frustration.py index 2a68546..3957ead 100644 --- a/examples/example_2D_geometric_frustration.py +++ b/examples/example_2D_geometric_frustration.py @@ -125,8 +125,8 @@ 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) +fig.tight_layout() +fig.subplots_adjust(hspace=0.35, wspace=0.3, top=0.9, left=0.07, bottom=0.07) ax = axes[0, 0] ax.loglog() diff --git a/examples/mca_shim_AO_FB.py b/examples/mca_shim_AO_FB.py index 72a494a..547f4f3 100644 --- a/examples/mca_shim_AO_FB.py +++ b/examples/mca_shim_AO_FB.py @@ -25,6 +25,12 @@ import numpy as np from tqdm import tqdm +from pathlib import Path +import sys +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + import dimod from dwave.system import DWaveSampler from dwave.system.composites import ParallelEmbeddingComposite @@ -38,7 +44,6 @@ ) from dwave.experimental.shimming import shim_flux_biases - def _make_anneal_schedules( exp_feature_info: list, target_c: float = 0.37, From 06d43bf23e08b763d98cb77105aa21b844a20208 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 8 Jun 2026 11:45:01 -0700 Subject: [PATCH 08/13] Add progress bars and missing docstrings -added progress bars using TQDM -added docstrings to the evaluation methods in observable.py -misc cleanup --- dwave/experimental/lattice_utils/__init__.py | 2 +- .../lattice_utils/experiment/experiment.py | 141 +++++++++++++++--- .../lattice_utils/lattice/chain.py | 8 +- .../lattice_utils/lattice/embedded_lattice.py | 9 +- .../lattice_utils/lattice/lattice.py | 16 +- .../lattice_utils/lattice/optimize.py | 2 +- .../lattice_utils/lattice/orbits.py | 13 +- .../lattice_utils/lattice/triangular.py | 40 ++--- .../lattice_utils/observable/kinks.py | 5 +- .../lattice_utils/observable/observable.py | 95 ++++++++++-- .../lattice_utils/observable/triangular.py | 2 +- examples/example_1D_Ising_chain.py | 36 +++-- examples/example_2D_geometric_frustration.py | 4 +- examples/mca_shim_AO_FB.py | 7 +- 14 files changed, 284 insertions(+), 96 deletions(-) diff --git a/dwave/experimental/lattice_utils/__init__.py b/dwave/experimental/lattice_utils/__init__.py index d8cb8ad..fb1e903 100644 --- a/dwave/experimental/lattice_utils/__init__.py +++ b/dwave/experimental/lattice_utils/__init__.py @@ -12,6 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from . import experiment, lattice, observable +from dwave.experimental.lattice_utils import experiment, lattice, observable __all__ = ["experiment", "lattice", "observable"] diff --git a/dwave/experimental/lattice_utils/experiment/experiment.py b/dwave/experimental/lattice_utils/experiment/experiment.py index 331e81e..c2dae1c 100644 --- a/dwave/experimental/lattice_utils/experiment/experiment.py +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -24,6 +24,7 @@ import dimod import numpy as np +from tqdm.auto import tqdm from dwave.experimental.lattice_utils.lattice import Lattice from dwave.experimental.lattice_utils.observable import ( @@ -38,18 +39,10 @@ __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 +DW_TEAL = "#17bebb" +DW_BLUE = "#2a7de1" +DW_ORANGE = "#f37820" + @dataclass class ExperimentConfig: @@ -171,7 +164,12 @@ def apply_param(self, param: dict[str, float]) -> None: 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: + def run_iteration( + self, + parameter_list: list, + progress: bool = False, + scaling_factor: float = 1.0, + ) -> bool: """Run one experiment iteration for each parameter set in ``parameter_list``. For each parametrization, this method applies the parameters, builds the @@ -180,26 +178,36 @@ def run_iteration(self, parameter_list: list, **kwargs) -> bool: Args: parameter_list: List of parameter dictionaries to run. + progress: If true, displays a progress bar for waiting on results. + scaling_factor: A multiplicative factor to apply to the BQM before sampling. Returns: A boolean value corresponding to whether or not the experiment is finished. """ try: - self.inst._load_embeddings(self.sampler, **kwargs) + self.inst._load_embeddings(self.sampler) 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()})" + tqdm.write( + f"\n{type(self.inst).__name__}={self.inst.dimensions}, " + f"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 = {} + create_bar = self._make_progress_bar( + total=len(parameter_list), + desc="Creating sampler calls", + colour=DW_BLUE, + enabled=progress, + ) + for index, param in enumerate(parameter_list): self.apply_param(param) call_dict[index] = self._build_sampler_call() @@ -207,14 +215,29 @@ def run_iteration(self, parameter_list: list, **kwargs) -> bool: call_dict.pop(index) else: response_dict[index] = self.sampler.sample( - call_dict[index].bqm * kwargs.get("scaling_factor", 1.0), + call_dict[index].bqm * scaling_factor, **call_dict[index].sampler_params, ) + if create_bar is not None: + create_bar.update() + + if create_bar is not None: + create_bar.close() if len(call_dict) == 0: - print(f"***\n***\nFINISHED for all {len(parameter_list)} parameterizations.\n***\n***") + if progress: + tqdm.write( + f"***\n***\nFINISHED for all {len(parameter_list)} parameterizations.\n***\n***" + ) return True + wait_bar = self._make_progress_bar( + total=len(call_dict), + desc=" Awaiting/parsing data", + colour=DW_TEAL, + enabled=progress, + ) + # Get and manage all the results while response_dict: made_progress = False @@ -227,6 +250,8 @@ def run_iteration(self, parameter_list: list, **kwargs) -> bool: self._update_shim(call_dict[index], results) savedata = self._generate_data_to_save(call_dict[index], results) self._save_results(savedata, quiet=True) + if wait_bar is not None: + wait_bar.update() del response_dict[index] made_progress = True break @@ -234,6 +259,10 @@ def run_iteration(self, parameter_list: list, **kwargs) -> bool: if not made_progress: time.sleep(0.1) # Waiting for results to come in + if wait_bar is not None: + wait_bar.close() + + self._print_iteration_status(call_dict, len(parameter_list), enabled=progress) return False def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[str, Any]: @@ -273,6 +302,78 @@ def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[st return results + def _make_progress_bar( + self, + *, + total: int, + desc: str, + colour: str, + enabled: bool, + bar_format: str | None = None, + initial: int | float = 0, + ) -> tqdm | None: + """Create a tqdm progress bar with consistent formatting.""" + if not enabled: + return None + + if bar_format is None: + bar_width = min(100, max(total, 20)) + bar_format = f"{{desc}}: |{{bar:{bar_width}}}{{r_bar}}{{bar:-{bar_width}b}}" + + return tqdm( + total=total, + initial=initial, + desc=desc, + bar_format=bar_format, + colour=colour, + ) + + def _print_iteration_status( + self, + call_dict: dict[int, SamplerCall], + num_params: int, + enabled: bool, + ) -> None: + """Print a summary of the iteration status, including progress and iteration ranges.""" + if not enabled: + return + iteration_range = ( + f"Iteration range " + f"{min(call.shimdata['total_iterations'] for call in call_dict.values())}-" + f"{max(call.shimdata['total_iterations'] for call in call_dict.values())} " + ) + if self.max_iterations is None: + tqdm.write(" Total progress: " + iteration_range) + return + + total = num_params * self.max_iterations + progress_value = ( + sum(call.shimdata["total_iterations"] for call in call_dict.values()) + + (num_params - len(call_dict)) * self.max_iterations + ) + + progress_string = ( + f"{progress_value / total * 100:.1f}% " + f"Iteration range " + f"{min(call.shimdata['total_iterations'] for call in call_dict.values())}-" + f"{max(call.shimdata['total_iterations'] for call in call_dict.values())} " + f"of {self.max_iterations} " + f"({num_params - len(call_dict)} of {num_params} parameters finished)" + ) + + bar_width = min(100, max(num_params, 20)) + bar_format = f"{{desc}}: |{{bar:{bar_width}}}| {progress_string}" + + total_bar = self._make_progress_bar( + total=total, + desc=" Total progress", + bar_format=bar_format, + colour=DW_ORANGE, + enabled=enabled, + initial=progress_value, + ) + total_bar.close() + def _save_results( self, data_dict: dict[str, Any], diff --git a/dwave/experimental/lattice_utils/lattice/chain.py b/dwave/experimental/lattice_utils/lattice/chain.py index 246cad0..39391ab 100644 --- a/dwave/experimental/lattice_utils/lattice/chain.py +++ b/dwave/experimental/lattice_utils/lattice/chain.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections.abc import Iterator +from collections.abc import Generator from pathlib import Path from numpy.typing import NDArray @@ -27,7 +27,7 @@ class 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. + non-periodic (open chain) based on the ``periodic`` parameter. Args: dimensions: One-element tuple giving the number of spins in the chain. @@ -65,11 +65,11 @@ def __init__( coupler_orbits=coupler_orbits, ) - def generate_edges(self) -> Iterator[tuple[int, int]]: + def generate_edges(self) -> Generator[tuple[int, int]]: """Yield edges for a 1D chain lattice. Returns: - An iterator of tuples, where each tuple represents an edge between + A generator of tuples, where each tuple represents an edge between two spins in the chain. """ n = self.dimensions[0] diff --git a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py index a33b0ba..842ad46 100644 --- a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -15,7 +15,7 @@ from itertools import combinations, product from numbers import Integral -from collections.abc import Iterator, Hashable +from collections.abc import Generator, Hashable import dimod import numpy as np @@ -56,6 +56,7 @@ def __init__( 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"): @@ -76,7 +77,7 @@ def get_chain_connectivity( 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 + in the chains corresponding to u and v that should be connected. For a chain edge (u == v or v is None), this will return pairs of indices within the same chain. For a logical edge (u != v), this will return pairs of indices between the two chains. @@ -88,11 +89,11 @@ def get_chain_connectivity( # 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]]: + def generate_edges(self) -> Generator[tuple[Hashable, Hashable]]: """Yield physical edges for the embedded lattice. Returns: - An iterator of tuples, where each tuple represents an edge between + A generator of tuples, where each tuple represents an edge between two spins in the physical lattice. """ logical_bqm = self.logical_lattice.make_nominal_bqm() diff --git a/dwave/experimental/lattice_utils/lattice/lattice.py b/dwave/experimental/lattice_utils/lattice/lattice.py index a724923..63db7a3 100644 --- a/dwave/experimental/lattice_utils/lattice/lattice.py +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -14,7 +14,7 @@ import os from pathlib import Path -from collections.abc import Iterator, Hashable +from collections.abc import Generator, Hashable from abc import ABC, abstractmethod import warnings @@ -66,6 +66,12 @@ def __init__( self.data_root = Path(data_root) self.periodic = periodic if periodic is not None else tuple(False for _ in dimensions) + if len(self.periodic) != len(self.dimensions): + raise ValueError( + f"periodic and dimensions must have the same length: " + f"got {len(self.periodic)} and {len(self.dimensions)}." + ) + self.edge_list: list[tuple[Hashable, Hashable]] = list(self.generate_edges()) if not hasattr(self, "num_spins"): @@ -76,9 +82,8 @@ def __init__( self.initialize_orbits(qubit_orbits, coupler_orbits) @abstractmethod - def generate_edges(self) -> Iterator[tuple[Hashable, Hashable]]: + def generate_edges(self) -> Generator[tuple[Hashable, Hashable]]: """Yield the edges for this lattice.""" - raise NotImplementedError def embed_lattice( self, @@ -87,7 +92,7 @@ def embed_lattice( timeout: int = 10, max_number_of_embeddings: int | None = None, min_number_of_embeddings: int = 1, - exclude_qubits: list = [], + exclude_qubits: list | None = None, **kwargs, ) -> None: """Find or load embeddings onto the sampler graph. @@ -102,6 +107,9 @@ def embed_lattice( exclude_qubits: Qubits to remove from the sampler graph before searching for embeddings. """ + if exclude_qubits is None: + exclude_qubits = [] + graph_bqm = dimod.to_networkx_graph(self.make_nominal_bqm()) graph_sampler = sampler.to_networkx_graph() graph_sampler.remove_nodes_from(exclude_qubits) diff --git a/dwave/experimental/lattice_utils/lattice/optimize.py b/dwave/experimental/lattice_utils/lattice/optimize.py index 1a91f38..11943d8 100644 --- a/dwave/experimental/lattice_utils/lattice/optimize.py +++ b/dwave/experimental/lattice_utils/lattice/optimize.py @@ -88,7 +88,7 @@ def optimize_increasing_sa_sweeps( Returns: A tuple containing the best energy found, the corresponding sample as a - NumPy array, and a string indicating the optimization method used. + NumPy array, and a string indicating the optimization method used. """ sa = SimulatedAnnealingSampler() diff --git a/dwave/experimental/lattice_utils/lattice/orbits.py b/dwave/experimental/lattice_utils/lattice/orbits.py index 40852c3..cdcc812 100644 --- a/dwave/experimental/lattice_utils/lattice/orbits.py +++ b/dwave/experimental/lattice_utils/lattice/orbits.py @@ -13,6 +13,7 @@ # limitations under the License. from collections.abc import Hashable +from collections import defaultdict import dimod import networkx as nx @@ -37,7 +38,7 @@ def reindex(mapping: dict[Hashable, int]) -> dict[Hashable, int]: mapping: Dictionary whose values represent indices or labels. Returns: - A new dictionary with the same keys as `mapping` but with values reindexed + 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()))} @@ -113,9 +114,9 @@ def get_bqm_orbits( 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())} + mapping_h = defaultdict(list) + mapping_mp = defaultdict(list) + mapping_J = defaultdict(list) for p, q in bqm.linear.items(): mapping_h[q].append(f"hnode_{p}") @@ -213,8 +214,8 @@ def get_orbits(bqm: dimod.BQM, edge_list: list[int, int]) -> tuple[NDArray, NDAr """Provide a bqm and receive a set of usable orbits derived from the signed BQM. Args: - bqm: Ising model to analyze - edge_list + bqm: Ising model to analyze. + edge_list: List of edges from the original unsigned BQM. Returns: A tuple ``(qubit_orbits_array, coupler_orbits_array)`` where diff --git a/dwave/experimental/lattice_utils/lattice/triangular.py b/dwave/experimental/lattice_utils/lattice/triangular.py index 6ec42e9..b83622e 100644 --- a/dwave/experimental/lattice_utils/lattice/triangular.py +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections.abc import Iterator, Hashable +from collections.abc import Generator, Hashable from pathlib import Path import networkx as nx @@ -37,13 +37,16 @@ class Triangular(Lattice): 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. + 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. """ @@ -61,6 +64,7 @@ def __init__( ): 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] @@ -68,14 +72,6 @@ def __init__( 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." @@ -84,6 +80,14 @@ def __init__( raise ValueError( "For Triangular with periodic[1]=True, dimensions[1] must be divisible by 3." ) + super().__init__( + dimensions=dimensions, + periodic=periodic, + data_root=data_root, + orbit_type=orbit_type, + qubit_orbits=qubit_orbits, + coupler_orbits=coupler_orbits, + ) def coordinates(self, node: int) -> tuple[int, int]: """Return the coordinates of a node in the lattice given its index. @@ -121,7 +125,7 @@ def make_nominal_bqm(self) -> dimod.BQM: return bqm - def generate_edges(self) -> Iterator[tuple[int, int]]: + def generate_edges(self) -> Generator[tuple[int, int]]: """Yield edges for the triangular lattice and initialize coordinate attributes. y is the first dimension, x is the second. Edges are straight along @@ -129,7 +133,7 @@ def generate_edges(self) -> Iterator[tuple[int, int]]: not periodic. Returns: - An iterator of tuples, where each tuple represents an edge between + A generator of tuples, where each tuple represents an edge between two spins in the lattice. """ length_y, length_x = self.dimensions diff --git a/dwave/experimental/lattice_utils/observable/kinks.py b/dwave/experimental/lattice_utils/observable/kinks.py index 9f879b8..65495cc 100644 --- a/dwave/experimental/lattice_utils/observable/kinks.py +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -24,10 +24,7 @@ class KinkKinkCorrelator(Observable): - """For 1D chains.""" - - def __init__(self): - super().__init__() + """A class for computing the kink-kink correlator for 1D spin chains.""" def evaluate(self, experiment: Experiment, bqm: BQM, sample_set: SampleSet) -> NDArray: """Compute the kink-kink correlator for 1D spin chains. diff --git a/dwave/experimental/lattice_utils/observable/observable.py b/dwave/experimental/lattice_utils/observable/observable.py index 9e8c4b9..4cc9ebf 100644 --- a/dwave/experimental/lattice_utils/observable/observable.py +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -42,6 +42,7 @@ class Observable(ABC): 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__ @@ -52,7 +53,7 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> ObservableResult: - raise NotImplementedError + pass class QubitMagnetization(Observable): @@ -64,6 +65,16 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> NDArray: + """Return per-qubit mean spin values over the provided samples. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples used to compute the magnetization. + + Returns: + A numpy array containing the mean magnetization for each qubit. + """ sample_array = dimod.as_samples(sample_set)[0].astype(float) return np.mean(sample_array, axis=0) @@ -77,6 +88,16 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> NDArray: + """Return per-coupler pairwise spin correlations over the provided samples. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples used to compute the coupler correlations. + + Returns: + A numpy array containing the pairwise spin correlations for each coupler. + """ sample_array = dimod.as_samples(sample_set)[0].astype(float) if len(experiment.inst.edge_list) == 0: return np.empty(0, dtype=float) @@ -95,6 +116,16 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> NDArray: + """Return the mean coupler frustration over the provided samples. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples used to compute the mean coupler frustration. + + Returns: + A numpy array containing the mean coupler frustration for each edge. + """ sample_array = dimod.as_samples(sample_set)[0].astype(float) if len(experiment.inst.edge_list) == 0: return np.empty(0, dtype=float) @@ -120,11 +151,22 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> NDArray: + """Return signed sample energies from the sample set. + + Args: + experiment: Experiment context providing the ``energy_scale`` sign. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples containing energy data. + + Returns: + A numpy array containing the sample energies multiplied by the sign + of ``energy_scale``. + """ 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.""" + """Compute bitpacked spins.""" def evaluate( self, @@ -132,6 +174,17 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> tuple[NDArray, tuple[int, int]]: + """Return bitpacked spin samples and their original array shape. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: Samples containing the spin values to unpack. + + Returns: + A tuple containing the bitpacked spin array and the original array + shape. + """ sample_array = dimod.as_samples(sample_set)[0] # Bitpack solutions @@ -153,11 +206,25 @@ def evaluate( path: str | Path | None = None, inst: Lattice | None = None, ) -> float: - """Get the reference energy for the given BQM, computing and caching it if needed.""" + """Get the reference energy for the given BQM, computing and caching it + if needed. + + Args: + experiment: The experiment for which to get the reference energy. Used + to determine the path for caching and loading the reference energy. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: The sample set is not used in this observable, but is + included in the signature for consistency with other observables. + path: Optional path to load/save the reference energy. If not provided, + a default path will be generated based on the experiment and BQM. + + Returns: + The reference energy for the given BQM. + """ if path is not None: path = Path(path) else: - path = get_reference_energy_path(experiment, bqm=bqm) + path = get_reference_energy_path(bqm, experiment) if path.exists(): energy, sample, method_string = self.load(experiment, bqm, path) @@ -187,7 +254,7 @@ def load( if path is not None: path = Path(path) else: - path = get_reference_energy_path(experiment, bqm=bqm) + path = get_reference_energy_path(bqm, experiment) with open(path, "r") as f: method_string = f.readline().strip() @@ -223,17 +290,16 @@ def update( if new_energy < reference_energy: if path is None: - path = get_reference_energy_path(experiment, bqm=bqm) + path = get_reference_energy_path(bqm, experiment) self.save(path, new_energy, sample, reference_method_string) - print(f"Updated energy from {reference_energy} to {new_energy}.") else: - raise ValueError + raise ValueError("New energy is not better than reference energy, not updating.") def get_reference_energy_path( + bqm: dimod.BQM, 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. @@ -242,18 +308,17 @@ def get_reference_energy_path( pathstring, for example when ground-state energies depend on the specific chip. Args: + bqm: The BQM for which to get the reference energy path. experiment: The experiment for which to get the reference energy path. root: Optional root directory to use instead of the experiment's data root. - bqm: The BQM for which to get the reference energy path. + dummy_experiment_data_dict: A dictionary containing the keys ``run_index``, + ``num_random_instances``, and ``inst`` to use when no experiment is + provided. This allows for generation of dummy experiment data without + all the overhead, for running without an actual experiment. Returns: The path to the reference energy file. """ - if 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: diff --git a/dwave/experimental/lattice_utils/observable/triangular.py b/dwave/experimental/lattice_utils/observable/triangular.py index b6faac9..9c9509e 100644 --- a/dwave/experimental/lattice_utils/observable/triangular.py +++ b/dwave/experimental/lattice_utils/observable/triangular.py @@ -24,7 +24,7 @@ class TriangularOP(Observable): - """For calculating the order parameter of triangular lattices.""" + """A class for calculating the order parameter of triangular lattices.""" def evaluate( self, diff --git a/examples/example_1D_Ising_chain.py b/examples/example_1D_Ising_chain.py index f9ea39f..2ea72f9 100644 --- a/examples/example_1D_Ising_chain.py +++ b/examples/example_1D_Ising_chain.py @@ -34,8 +34,8 @@ # Two samplers: an Advantage2 prototype and an Advantage system. samplers = [ - DWaveSampler(solver="Advantage2_system1"), - DWaveSampler(solver="Advantage_system4.1"), + #DWaveSampler(solver="Advantage2_system1"), + DWaveSampler(solver="Advantage_system4"), ] NUM_SPINS = 256 @@ -86,7 +86,7 @@ config = experiment.FastAnnealExperimentConfig( energy_scale=energy_scale, coupler_shim_step=0.05, - flux_bias_shim_step=1e-6, + flux_bias_shim_step=3e-6, ) exp = experiment.Experiment(inst=inst, sampler=sampler, max_iterations=5, config=config) @@ -95,13 +95,17 @@ # 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. + + # Note that only running 5 iterations will give poor results for the + # kink-kink correlator, which is fairly sensitive and requires a + # converged shim, but this is just for demonstration purposes. 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) + while True: + done = exp.run_iteration(parameter_list, progress=True) if done: break @@ -119,7 +123,7 @@ fbshim = [] # flux bias shim kkc = [] # kink-kink correlator for param in parameter_list: - exp.apply_param(param) #TODO clean this up? + exp.apply_param(param) res = exp.load_results(num_iterations=1000) frust.append(np.array([np.mean(it["CouplerFrustration"]) for it in res])) @@ -136,8 +140,8 @@ 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) + fig.tight_layout() + fig.subplots_adjust(hspace=0.35, wspace=0.3, top=0.9) ax = axes[0, 0] ax.loglog() @@ -173,7 +177,7 @@ **point_style, ) - ax.set_title("Kink density (with ~t_a^{-1/2} guideline)") + 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]) @@ -200,37 +204,37 @@ ax = axes[1, 0] ax.plot(fbshim[0]) - ax.set_title(f"Flux bias shim, t_a={ANNEAL_TIMES[0]:.3f}μs") + 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_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_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_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_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_title(f"Coupler shim, $t_a$={ANNEAL_TIMES[-1]:.3f}μs") ax.set_xlabel("Iteration") ax.grid(which="both", alpha=0.3) @@ -248,6 +252,7 @@ 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) +fig2.tight_layout(rect=(0.05, 0.05, 1, 0.95), w_pad=4.0) for isampler, sampler in enumerate(samplers): @@ -303,6 +308,7 @@ time_ns = ANNEAL_TIMES[0] * 1000 title = f"1D chain kink-kink correlator, {dims}, {time_ns:.1f} ns" fig3.suptitle(title, fontsize=16) +fig3.tight_layout(rect=(0.05, 0.05, 1, 0.95), w_pad=4.0) for isampler, sampler in enumerate(samplers): diff --git a/examples/example_2D_geometric_frustration.py b/examples/example_2D_geometric_frustration.py index 2a68546..3957ead 100644 --- a/examples/example_2D_geometric_frustration.py +++ b/examples/example_2D_geometric_frustration.py @@ -125,8 +125,8 @@ 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) +fig.tight_layout() +fig.subplots_adjust(hspace=0.35, wspace=0.3, top=0.9, left=0.07, bottom=0.07) ax = axes[0, 0] ax.loglog() diff --git a/examples/mca_shim_AO_FB.py b/examples/mca_shim_AO_FB.py index 72a494a..547f4f3 100644 --- a/examples/mca_shim_AO_FB.py +++ b/examples/mca_shim_AO_FB.py @@ -25,6 +25,12 @@ import numpy as np from tqdm import tqdm +from pathlib import Path +import sys +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + import dimod from dwave.system import DWaveSampler from dwave.system.composites import ParallelEmbeddingComposite @@ -38,7 +44,6 @@ ) from dwave.experimental.shimming import shim_flux_biases - def _make_anneal_schedules( exp_feature_info: list, target_c: float = 0.37, From dab4d6dd0d841c69b87ce0ce621c18a11cfafe33 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 8 Jun 2026 12:05:07 -0700 Subject: [PATCH 09/13] Update default parameters of 1D example --- examples/example_1D_Ising_chain.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/example_1D_Ising_chain.py b/examples/example_1D_Ising_chain.py index 158fc2b..2ea72f9 100644 --- a/examples/example_1D_Ising_chain.py +++ b/examples/example_1D_Ising_chain.py @@ -88,8 +88,7 @@ coupler_shim_step=0.05, flux_bias_shim_step=3e-6, ) - #exp = experiment.Experiment(inst=inst, sampler=sampler, max_iterations=5, config=config) - exp = experiment.Experiment(inst=inst, sampler=sampler, config=config) + 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. From 80c39296803524eceb8f42ffee9ee027426c1b48 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 11 Jun 2026 11:56:35 -0700 Subject: [PATCH 10/13] Update tests and examples -fix failing tests due to updated method signatures with explicit arguments -rename energy_scale to signed_energy_scale --- .../lattice_utils/experiment/experiment.py | 53 +- .../lattice_utils/lattice/embedded_lattice.py | 32 +- .../lattice_utils/lattice/triangular.py | 20 +- .../lattice_utils/observable/kinks.py | 2 +- .../lattice_utils/observable/observable.py | 10 +- examples/example_1D_Ising_chain.py | 25 +- examples/example_2D_geometric_frustration.py | 31 +- tests/test_lattice_utils.py | 824 ++++++++++-------- 8 files changed, 542 insertions(+), 455 deletions(-) diff --git a/dwave/experimental/lattice_utils/experiment/experiment.py b/dwave/experimental/lattice_utils/experiment/experiment.py index c2dae1c..1e1eced 100644 --- a/dwave/experimental/lattice_utils/experiment/experiment.py +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -48,7 +48,7 @@ class ExperimentConfig: """Container for the parameters that define an experiment.""" - energy_scale: float = 1.0 + signed_energy_scale: float = 1.0 num_reads: int = 100 anneal_time: float = 1.0 num_random_instances: int | None = 1 @@ -78,7 +78,7 @@ def __init__( inst: Lattice, sampler: dimod.Sampler, max_iterations: int | None = None, - config: ExperimentConfig, + config: ExperimentConfig = ExperimentConfig(), ): self.inst = inst self.sampler = sampler @@ -190,12 +190,13 @@ def run_iteration( except FileNotFoundError as e: raise FileNotFoundError("No Embedding Found: ", e) from e - tqdm.write( - f"\n{type(self.inst).__name__}={self.inst.dimensions}, " - f"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()})" - ) + if progress: + tqdm.write( + f"\n{type(self.inst).__name__}={self.inst.dimensions}, " + f"J={self.param['signed_energy_scale']}, " + f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')} " + f"({self.inst._get_instance_pathstring()}/{self._get_solver_pathstring()})" + ) parameter_list = self._format_parameter_list(parameter_list) response_dict = {} @@ -378,7 +379,7 @@ def _save_results( self, data_dict: dict[str, Any], run_index: int | None = None, - quiet: bool = False, + quiet: bool = True, filename: str | None = None, ) -> None: """Save results to disk using LZMA-compressed pickle.""" @@ -453,12 +454,12 @@ def _get_parameter_pathstring(self) -> str: 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"] + signed_energy_scale = self.param["signed_energy_scale"] if "anneal_time" in self.param: - pathstring = f'energyscale{energy_scale:0.3}/atime{self.param["anneal_time"]:010.6f}us' + pathstring = f'energyscale{signed_energy_scale:0.3}/atime{self.param["anneal_time"]:010.6f}us' elif "anneal_schedule" in self.param: - pathstring = f'energyscale{energy_scale:0.3}/asched{self.param["anneal_schedule"]}' + pathstring = f'energyscale{signed_energy_scale:0.3}/asched{self.param["anneal_schedule"]}' else: raise ValueError @@ -671,7 +672,7 @@ def _update_coupler_shim( ) -> None: """Update relative coupler strength based on measured frustration.""" orbits = self.inst.coupler_orbits - energy_scale = self.param["energy_scale"] + signed_energy_scale = self.param["signed_energy_scale"] relative_coupler_strength = sampler_call.shimdata["relative_coupler_strength"] # Allow for zero step size, which will just truncate the shim. @@ -713,21 +714,21 @@ def _update_coupler_shim( violators = ( relative_coupler_strength[iemb, bin_edges] * nominal_values[bin_edges] - * energy_scale + * signed_energy_scale > 1 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - 0.99999 / nominal_values[bin_edges[violators]] / energy_scale + 0.99999 / nominal_values[bin_edges[violators]] / signed_energy_scale ) violators = ( relative_coupler_strength[iemb, bin_edges] * nominal_values[bin_edges] - * energy_scale + * signed_energy_scale < -2 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - -1.99999 / nominal_values[bin_edges[violators]] / energy_scale + -1.99999 / nominal_values[bin_edges[violators]] / signed_energy_scale ) # Renormalize each orbit after truncation @@ -748,24 +749,24 @@ def _update_coupler_shim( violators = ( relative_coupler_strength[iemb, bin_edges] * nominal_values[bin_edges] - * energy_scale + * signed_energy_scale > 1 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - 0.99999 / nominal_values[bin_edges[violators]] / energy_scale + 0.99999 / nominal_values[bin_edges[violators]] / signed_energy_scale ) violators = ( relative_coupler_strength[iemb, bin_edges] * nominal_values[bin_edges] - * energy_scale + * signed_energy_scale < -2 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - -1.99999 / nominal_values[bin_edges[violators]] / energy_scale + -1.99999 / nominal_values[bin_edges[violators]] / signed_energy_scale ) - Q = nominal_values * relative_coupler_strength * energy_scale + Q = nominal_values * relative_coupler_strength * signed_energy_scale Q_max = np.max(Q) Q_min = np.min(Q) if Q_max > 1 or Q_min < -2: @@ -776,7 +777,7 @@ def _update_coupler_shim( def _make_bqm(self, sampler_call: SamplerCall) -> dimod.BQM: """Construct a BQM for the current sampler call.""" - energy_scale = self.param["energy_scale"] + signed_energy_scale = self.param["signed_energy_scale"] bqm = dimod.BQM(vartype="SPIN") if not hasattr(self.inst, "embedding_list"): nominal_bqm = sampler_call.nominal_bqms[0] @@ -787,8 +788,8 @@ def _make_bqm(self, sampler_call: SamplerCall) -> dimod.BQM: 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) + for u, v in self.inst.edge_list: + bqm.add_quadratic(u, v, nominal_bqm.quadratic[u, v] * signed_energy_scale) return bqm @@ -805,7 +806,7 @@ def _make_bqm(self, sampler_call: SamplerCall) -> dimod.BQM: bias = ( nominal_bqm.quadratic[*edge] * relative_coupler_strength[iemb, iedge] - * energy_scale + * signed_energy_scale ) bqm.add_quadratic(emb[edge[0]], emb[edge[1]], bias) diff --git a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py index 842ad46..bb4a6fa 100644 --- a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -16,6 +16,7 @@ from itertools import combinations, product from numbers import Integral from collections.abc import Generator, Hashable +from pathlib import Path import dimod import numpy as np @@ -46,9 +47,16 @@ class EmbeddedLattice(Lattice): def __init__( self, + *, logical_lattice: Lattice, chain_nodes: dict[int, tuple[int, Integral]], - **kwargs, + dimensions: tuple[int, ...], + data_root: str | Path | None = None, + periodic: tuple[bool, ...] | None = None, + orbit_type: str = "singleton", + qubit_orbits: NDArray | None = None, + coupler_orbits: NDArray | None = None, + chain_strength: float = 2, ): if not isinstance(logical_lattice, Lattice): raise TypeError("logical_lattice must be a Lattice instance.") @@ -57,12 +65,26 @@ def __init__( 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 data_root is None: + data_root = logical_lattice.data_root + + self.chain_nodes = chain_nodes + self.chain_coupling = -chain_strength + if not hasattr(self, "num_spins"): self.num_spins = sum(len(c) for c in chain_nodes.values()) - kwargs.setdefault("periodic", self.logical_lattice.periodic) - super().__init__(**kwargs) + + if periodic is None: + periodic = self.logical_lattice.periodic + + super().__init__( + dimensions=dimensions, + data_root=data_root, + periodic=periodic, + orbit_type=orbit_type, + qubit_orbits=qubit_orbits, + coupler_orbits=coupler_orbits, + ) def get_chain_connectivity( self, diff --git a/dwave/experimental/lattice_utils/lattice/triangular.py b/dwave/experimental/lattice_utils/lattice/triangular.py index b83622e..6f0c712 100644 --- a/dwave/experimental/lattice_utils/lattice/triangular.py +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -56,7 +56,7 @@ def __init__( *, dimensions: tuple[int, int], periodic: tuple[bool, bool] = (True, False), - data_root: Path | None = None, + data_root: str | Path, orbit_type: str = "singleton", qubit_orbits: NDArray | None = None, coupler_orbits: NDArray | None = None, @@ -72,14 +72,6 @@ def __init__( 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 - 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." - ) super().__init__( dimensions=dimensions, periodic=periodic, @@ -88,6 +80,14 @@ def __init__( 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. @@ -210,7 +210,7 @@ def __init__( *, dimensions: tuple[int, int], periodic: tuple[bool, bool] = (True, False), - data_root: Path | None = None, + data_root: str | Path, orbit_type: str = "singleton", qubit_orbits: NDArray | None = None, coupler_orbits: NDArray | None = None, diff --git a/dwave/experimental/lattice_utils/observable/kinks.py b/dwave/experimental/lattice_utils/observable/kinks.py index 65495cc..1fb7e44 100644 --- a/dwave/experimental/lattice_utils/observable/kinks.py +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -40,7 +40,7 @@ def evaluate(self, experiment: Experiment, bqm: BQM, sample_set: SampleSet) -> N 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"]) + kink_mask = shifted_samples * samples == np.sign(experiment.param["signed_energy_scale"]) chain_length = kink_mask.shape[-1] kink_mask = np.reshape(kink_mask, (-1, chain_length)) kink_density = np.mean(kink_mask) diff --git a/dwave/experimental/lattice_utils/observable/observable.py b/dwave/experimental/lattice_utils/observable/observable.py index 4cc9ebf..94dbf9d 100644 --- a/dwave/experimental/lattice_utils/observable/observable.py +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -134,7 +134,7 @@ def evaluate( 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"]) + ) * np.sign(experiment.param["signed_energy_scale"]) return spin_product * coupler_signs / 2 + 1 / 2 @@ -142,7 +142,7 @@ def evaluate( class SampleEnergy(Observable): """Compute sample energies with respect to the nominal BQM. - Energies exclude the magnitude of ``energy_scale`` but include its sign. + Energies exclude the magnitude of ``signed_energy_scale`` but include its sign. """ def evaluate( @@ -154,15 +154,15 @@ def evaluate( """Return signed sample energies from the sample set. Args: - experiment: Experiment context providing the ``energy_scale`` sign. + experiment: Experiment context providing ``signed_energy_scale``. bqm: The binary quadratic model corresponding to the problem instance. sample_set: Samples containing energy data. Returns: A numpy array containing the sample energies multiplied by the sign - of ``energy_scale``. + of ``signed_energy_scale``. """ - return sample_set.data_vectors["energy"] * np.sign(experiment.param["energy_scale"]) + return sample_set.data_vectors["energy"] * np.sign(experiment.param["signed_energy_scale"]) class BitpackedSpins(Observable): diff --git a/examples/example_1D_Ising_chain.py b/examples/example_1D_Ising_chain.py index 2ea72f9..49e999e 100644 --- a/examples/example_1D_Ising_chain.py +++ b/examples/example_1D_Ising_chain.py @@ -82,9 +82,9 @@ # 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: + for signed_energy_scale in ENERGY_SCALES: config = experiment.FastAnnealExperimentConfig( - energy_scale=energy_scale, + signed_energy_scale=signed_energy_scale, coupler_shim_step=0.05, flux_bias_shim_step=3e-6, ) @@ -135,7 +135,10 @@ 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}" + dimensions = "x".join(str(dim) for dim in inst.dimensions) + signed_energy_scale = exp.param["signed_energy_scale"] + title = f"1D chain, {dimensions}, J={signed_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) @@ -245,8 +248,8 @@ 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) + kd_dict[sampler.solver.name, signed_energy_scale] = np.asarray(frust) + kkc_dict[sampler.solver.name, signed_energy_scale] = np.asarray(kkc) # Now plot the kink densities together, for a nice comparison. fig2, ax2 = plt.subplots(1, 2, figsize=(10, 8)) @@ -256,8 +259,8 @@ for isampler, sampler in enumerate(samplers): - for energy_scale in [-1.8, 0.1]: - M = kd_dict[sampler.solver.name, energy_scale] + for signed_energy_scale in [-1.8, 0.1]: + M = kd_dict[sampler.solver.name, signed_energy_scale] # Kink density plot theoryx = ANNEAL_TIMES[0] @@ -281,7 +284,7 @@ ax2[isampler].set_xlim([0.002, 9e1]) x = ANNEAL_TIMES - M = kd_dict[sampler.solver.name, energy_scale] + M = kd_dict[sampler.solver.name, signed_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]) @@ -312,9 +315,9 @@ 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]) + for signed_energy_scale in [-1.8, 0.1]: + magnetization = kkc_dict[sampler.solver.name, signed_energy_scale][0].T + kd = np.mean(kd_dict[sampler.solver.name, signed_energy_scale][0]) # Kink density plot ax3[isampler].grid(which="both", alpha=0.3) diff --git a/examples/example_2D_geometric_frustration.py b/examples/example_2D_geometric_frustration.py index 3957ead..9372712 100644 --- a/examples/example_2D_geometric_frustration.py +++ b/examples/example_2D_geometric_frustration.py @@ -67,34 +67,21 @@ 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, +config = experiment.FastAnnealExperimentConfig( + signed_energy_scale=0.8, coupler_shim_step=0.1, flux_bias_shim_step=5e-6, + readout_thermalization=100, + num_reads=100, ) - -exp.observables_to_collect = [ - observable.QubitMagnetization(), - observable.CouplerCorrelation(), - observable.CouplerFrustration(), - observable.SampleEnergy(), - observable.TriangularOP(), - observable.ReferenceEnergy(), - observable.BitpackedSpins(), -] +exp = experiment.Experiment(inst=inst, sampler=sampler, max_iterations=210, config=config) +exp.observables_to_collect.add(observable.TriangularOP()) # Make parameter list parameter_list = [{"anneal_time": rate} for rate in ANNEAL_TIMES] for _ in range(1000): - done = exp.run_iteration(parameter_list) + done = exp.run_iteration(parameter_list, progress=True) if done: break @@ -119,7 +106,7 @@ title = ( f"DimerizedTriangular, {'x'.join([str(dim) for dim in inst.dimensions])}, " - f"J={exp.param["energy_scale"]}, {sampler.solver.name}" + f"J={exp.param["signed_energy_scale"]}, {sampler.solver.name}" ) fig, axes = plt.subplots(3, 3, figsize=(16, 10)) fig.suptitle(title, fontsize=16) @@ -216,7 +203,7 @@ # 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"]}' +title = f'Triangular, global orbit, J={exp.param["signed_energy_scale"]}' fig2.suptitle(title, fontsize=16) M = m_dict[sampler.solver.name] diff --git a/tests/test_lattice_utils.py b/tests/test_lattice_utils.py index 96198e1..058c5a9 100644 --- a/tests/test_lattice_utils.py +++ b/tests/test_lattice_utils.py @@ -45,7 +45,10 @@ 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 +from dwave.experimental.lattice_utils.experiment.fast_anneal_experiment import ( + FastAnnealExperimentConfig, +) +from dwave.experimental.lattice_utils.experiment.experiment import ExperimentConfig def _make_triangular( @@ -55,12 +58,14 @@ def _make_triangular( orbit_type="singleton", halve_boundary_couplers=False, ): - return Triangular( - dimensions=(ly, lx), - periodic=periodic, - orbit_type=orbit_type, - halve_boundary_couplers=halve_boundary_couplers, - ) + with tempfile.TemporaryDirectory() as tmpdir: + return Triangular( + dimensions=(ly, lx), + periodic=periodic, + data_root=tmpdir, + orbit_type=orbit_type, + halve_boundary_couplers=halve_boundary_couplers, + ) def _make_mock_sampler(num_qubits=128, nodelist=None, solver_name="TestSolver"): @@ -100,12 +105,12 @@ def samples(self, sorted_by=None): def _make_mock_experiment( - inst, energy_scale=1.0, run_index=0, num_random_instances=1, extra_params=None + inst, signed_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.param = {"signed_energy_scale": signed_energy_scale, "num_random_instances": num_random_instances} exp.run_index = run_index if extra_params: exp.param.update(extra_params) @@ -113,16 +118,17 @@ def _make_mock_experiment( def _make_embedded_chain(chain_nodes): - return EmbeddedLattice( - logical_lattice=Chain( - dimensions=(len(chain_nodes),), + with tempfile.TemporaryDirectory() as tmpdir: + return EmbeddedLattice( + logical_lattice=Chain( + dimensions=(len(chain_nodes),), + periodic=(False,), + data_root=tmpdir, + ), + chain_nodes=chain_nodes, + dimensions=(sum(len(chain) for chain in chain_nodes.values()),), periodic=(False,), - ignore_embedding=True, - ), - chain_nodes=chain_nodes, - dimensions=(sum(len(chain) for chain in chain_nodes.values()),), - periodic=(False,), - ) + ) class TestUtils(unittest.TestCase): @@ -161,81 +167,98 @@ def test_confidence_interval_width(self): 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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + self.assertEqual(chain.num_spins, 6) + self.assertEqual(chain.num_edges, 6) + self.assertIn((5, 0), chain.edge_list) def test_non_periodic(self): - 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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(6,), periodic=(False,), data_root=tmpdir) + self.assertEqual(chain.num_spins, 6) + self.assertEqual(chain.num_edges, 5) + self.assertNotIn((5, 0), chain.edge_list) def test_single_node_periodic(self): - chain = Chain(dimensions=(1,), periodic=(True,)) - self.assertEqual(chain.num_spins, 1) - self.assertEqual(chain.num_edges, 0) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(1,), periodic=(True,), data_root=tmpdir) + 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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(2,), periodic=(True,), data_root=tmpdir) + self.assertEqual(chain.num_edges, 2) def test_geometry_name(self): - chain = Chain(dimensions=(6,), periodic=(True,)) - self.assertEqual(chain.geometry_name, "Chain") + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + 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]) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + 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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(5,), periodic=(False,), data_root=tmpdir) + 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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + 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) + with tempfile.TemporaryDirectory() as tmpdir: + bqm = Chain(dimensions=(3,), data_root=tmpdir).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)) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), orbit_type="singleton", data_root=tmpdir) + np.testing.assert_array_equal(chain.qubit_orbits, np.arange(4)) + np.testing.assert_array_equal(chain.coupler_orbits, np.arange(chain.num_edges)) def test_orbit_global(self): - 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)) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), orbit_type="global", data_root=tmpdir) + np.testing.assert_array_equal(chain.qubit_orbits, np.zeros(4, dtype=int)) + np.testing.assert_array_equal( + chain.coupler_orbits, np.zeros(chain.num_edges, dtype=int) + ) def test_orbit_explicit(self): - 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]) + with tempfile.TemporaryDirectory() as tmpdir: + 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]), + data_root=tmpdir, + ) + 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") + with tempfile.TemporaryDirectory() as tmpdir: + with self.assertRaises(ValueError): + Chain(dimensions=(4,), periodic=(True,), orbit_type="bogus", data_root=tmpdir) def test_get_path_invalid_kind(self): - chain = Chain(dimensions=(4,), periodic=(True,)) - with self.assertRaises(ValueError): - chain._get_path(None, "invalid") + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + with self.assertRaises(ValueError): + chain._get_path(None, "invalid") def test_standard_orbit_save_and_load(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -243,7 +266,7 @@ def test_standard_orbit_save_and_load(self): dimensions=(4,), periodic=(True,), orbit_type="standard", - lattice_data_root=Path(tmpdir), + data_root=tmpdir, ) self.assertIsNotNone(chain.qubit_orbits) self.assertIsNotNone(chain.coupler_orbits) @@ -252,47 +275,61 @@ def test_standard_orbit_save_and_load(self): dimensions=(4,), periodic=(True,), orbit_type="standard", - lattice_data_root=Path(tmpdir), + data_root=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() + with tempfile.TemporaryDirectory() as tmpdir: + base_chain = Chain( + dimensions=(2,), + periodic=(False,), + data_root=tmpdir, + ) + embedded_once = EmbeddedLattice( + logical_lattice=base_chain, + chain_nodes={0: (0, 1), 1: (2, 3)}, + dimensions=(4,), + periodic=(False,), + ) + with self.assertRaises(NotImplementedError): + EmbeddedLattice( + logical_lattice=embedded_once, + chain_nodes={0: (0, 1), 1: (2, 3), 2: (4, 5), 3: (6, 7)}, + dimensions=(8,), + periodic=(False,), + ) 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 tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = mock.MagicMock() + type(sampler).__name__ = "MockDWaveSampler" + sampler.to_networkx_graph.return_value = chain._make_networkx_graph() - with mock.patch( - "dwave.experimental.lattice_utils.lattice.lattice.find_multiple_embeddings", - return_value=[], - ): - with self.assertRaises(ValueError): - chain.embed_lattice(sampler, try_to_load=False, timeout=1) + 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)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = mock.MagicMock() type(sampler).__name__ = "MockDWaveSampler" sampler.to_networkx_graph.return_value = chain._make_networkx_graph() embeddings = np.array([[0, 1, 2, 3]]) - chain._save_embeddings(sampler, embeddings, data_root=Path(tmpdir)) + chain._save_embeddings(sampler, embeddings) - chain.embed_lattice(sampler, try_to_load=True, data_root=Path(tmpdir)) + chain.embed_lattice(sampler, try_to_load=True, data_root=tmpdir) np.testing.assert_array_equal(chain.embedding_list, embeddings) def test_embed_find_and_save(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), lattice_data_root=Path(tmpdir)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = mock.MagicMock() type(sampler).__name__ = "MockDWaveSampler" sampler.to_networkx_graph.return_value = chain._make_networkx_graph() @@ -302,18 +339,25 @@ def test_embed_find_and_save(self): "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)) + chain.embed_lattice(sampler, try_to_load=False, timeout=1, data_root=tmpdir) # Verify embedding was found and saved - emb_path = chain._get_path(Path(tmpdir), "embedding", sampler_name="MockDWaveSampler") + emb_path = chain._get_path("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") + with tempfile.TemporaryDirectory() as tmpdir: + tri = Triangular( + dimensions=(3, 3), + periodic=(True, False), + data_root=tmpdir, + orbit_type="singleton", + halve_boundary_couplers=False, + ) + self.assertEqual(tri.num_spins, 9) + self.assertGreater(tri.num_edges, 0) + self.assertEqual(tri.geometry_name, "Triangular") def test_coordinates(self): tri = _make_triangular(3, 3) @@ -340,44 +384,61 @@ def test_periodicity(self): 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) + with tempfile.TemporaryDirectory() as tmpdir: + dt = DimerizedTriangular( + dimensions=(3, 3), periodic=(True, False), orbit_type="singleton", data_root=tmpdir + ) + self.assertEqual(dt.geometry_name, "DimerizedTriangular") + self.assertIsNotNone(dt.logical_lattice) + self.assertEqual(dt.num_spins, 18) def test_chain_connectivity_self(self): - dt = DimerizedTriangular(dimensions=(3, 3), periodic=(True, False), orbit_type="singleton") - cc = dt.get_chain_connectivity(0) - self.assertEqual(cc, ((0, 1),)) + with tempfile.TemporaryDirectory() as tmpdir: + dt = DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + 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) + with tempfile.TemporaryDirectory() as tmpdir: + dt = DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + cases = [ + ((0,), ((0, 1),)), + ((0, 1), ((1, 0),)), + ((0, 3), ((1, 0),)), + ] + for args, expected in cases: + with self.subTest(args=args): + self.assertEqual(dt.get_chain_connectivity(*args), expected) class TestEmbeddedLattice(unittest.TestCase): def test_embed_sample(self): - 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]) + with tempfile.TemporaryDirectory() as tmpdir: + dt = DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + logical_sample = np.array([1, -1, 1, -1, 1, -1, 1, -1, 1]) + embedded = dt.embed_sample(logical_sample) + self.assertEqual(len(embedded), dt.num_spins) + # Each chain should have the same value + for spin, chain in dt.chain_nodes.items(): + for node in chain: + self.assertEqual(embedded[node], logical_sample[spin]) def test_unembed_sample(self): - 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])) + chain_nodes = {0: (0, 1, 2), 1: (3, 4, 5)} + with tempfile.TemporaryDirectory() as tmpdir: + embedded = EmbeddedLattice( + logical_lattice=Chain( + dimensions=(len(chain_nodes),), + periodic=(False,), + data_root=tmpdir, + ), + chain_nodes=chain_nodes, + dimensions=(sum(len(chain) for chain in chain_nodes.values()),), + periodic=(False,), + ) + physical_sample = np.array([1, 1, -1, -1, -1, 1]) + logical = embedded.unembed_sample(physical_sample) + np.testing.assert_array_equal(logical, np.array([1, -1])) def test_unembed_sample_breaks_ties_randomly(self): embedded = _make_embedded_chain({0: (0, 1), 1: (2, 3)}) @@ -436,77 +497,84 @@ def test_signed_bqm_symmetry(self): 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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + 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})) + with tempfile.TemporaryDirectory() as tmpdir: + dt = DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + 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]) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + 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) - ) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + 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) - ) + """All aligned (corr=1) -> frustration = 1.0""" + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + 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]) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + 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, signed_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)) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + 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,)) + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) bqm = chain.make_nominal_bqm() sample = np.array([1, -1, 1, -1]) obs = ReferenceEnergy() @@ -520,7 +588,7 @@ def test_reference_energy_save_load_roundtrip(self): 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)) + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) bqm = chain.make_nominal_bqm() obs = ReferenceEnergy() @@ -540,7 +608,7 @@ def test_reference_energy_evaluate_generates_and_caches(self): def test_reference_energy_update(self): with tempfile.TemporaryDirectory() as tmpdir: path = Path(tmpdir) / "ref.txt" - chain = Chain(dimensions=(4,), periodic=(False,)) + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) bqm = chain.make_nominal_bqm() obs = ReferenceEnergy() exp = _make_mock_experiment(chain) @@ -558,47 +626,47 @@ def test_reference_energy_update(self): 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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + bqm = chain.make_nominal_bqm() + exp = _make_mock_experiment(chain) - path = get_reference_energy_path(experiment=exp, bqm=bqm) - self.assertTrue(str(path).endswith(".txt")) + path = get_reference_energy_path(bqm, experiment=exp) + 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")) + # Via dummy data dict (experiment=None) + dummy = {"run_index": 0, "num_random_instances": 1, "inst": chain} + path2 = get_reference_energy_path(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)) + with tempfile.TemporaryDirectory() as tmpdir: + path3 = get_reference_energy_path(bqm, experiment=exp, 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)) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + 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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + 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): @@ -613,14 +681,15 @@ def test_uniform_state_vanishes(self): 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) + with tempfile.TemporaryDirectory() as tmpdir: + dt = DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + 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): @@ -648,83 +717,80 @@ def test_with_values(self): 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,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() - exp = Experiment(chain, sampler, results_root=tmpdir) - self.assertEqual(exp.experiment_results_root, Path(tmpdir).resolve()) + exp = Experiment(inst=chain, sampler=sampler) + self.assertEqual(exp.param["signed_energy_scale"], 1.0) + self.assertEqual(exp.param["num_reads"], 100) + self.assertIs(exp.inst, chain) def test_data_path_with_schedule(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() - exp = Experiment(chain, sampler, results_root=tmpdir, anneal_time=5.0) + exp = Experiment(inst=chain, sampler=sampler) 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)]}) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_schedule": [(0, 1), (5, 0.5)]}) self.assertIn("asched", str(exp.data_path)) def test_apply_param_unknown_sampler_raises(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() type(sampler).__name__ = "UnknownSampler" - exp = Experiment(chain, sampler, results_root=tmpdir) + exp = Experiment(inst=chain, sampler=sampler) with self.assertRaises(TypeError): - exp.apply_param({"energy_scale": 1.0, "anneal_time": 1.0}) + exp.apply_param({"signed_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,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() - exp = Experiment(chain, sampler, results_root=tmpdir) + exp = Experiment(inst=chain, sampler=sampler) del exp.param["anneal_time"] with self.assertRaises(ValueError): - exp.apply_param({"energy_scale": 1.0}) + exp.apply_param({"signed_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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = Experiment(inst=chain, sampler=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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler(num_qubits=128) + exp = Experiment(inst=chain, sampler=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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler(num_qubits=128) + fb = np.ones(128) * 0.01 + exp = Experiment(inst=chain, sampler=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,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() - exp = Experiment(chain, sampler, results_root=tmpdir) - exp.data_path = Path(tmpdir) + exp = Experiment(inst=chain, sampler=sampler) exp.run_index = 1 + exp.data_path = Path(tmpdir) shimdata = {"total_iterations": 5, "flux_biases": np.zeros(10)} data = {"shimdata": shimdata} @@ -737,11 +803,12 @@ def test_load_shim_from_file(self): def test_load_shim_empty_file(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() - exp = Experiment(chain, sampler, results_root=tmpdir) - exp.data_path = Path(tmpdir) + exp = Experiment(inst=chain, sampler=sampler) exp.run_index = 1 + exp.data_path = Path(tmpdir) + fn = Path(tmpdir) / "iter00000.pkl.lzma" fn.touch() @@ -750,11 +817,11 @@ def test_load_shim_empty_file(self): def test_load_shim_corrupted_file(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() - exp = Experiment(chain, sampler, results_root=tmpdir) - exp.data_path = Path(tmpdir) + exp = Experiment(inst=chain, sampler=sampler) exp.run_index = 1 + exp.data_path = Path(tmpdir) fn = Path(tmpdir) / "iter00000.pkl.lzma" fn.write_bytes(b"not a valid lzma file") @@ -763,9 +830,9 @@ def test_load_shim_corrupted_file(self): def test_load_shim_missing_file(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() - exp = Experiment(chain, sampler, results_root=tmpdir) + exp = Experiment(inst=chain, sampler=sampler) exp.data_path = Path(tmpdir) exp.run_index = 1 @@ -776,9 +843,9 @@ def test_load_shim_missing_file(self): def test_save_and_reload(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() - exp = Experiment(chain, sampler, results_root=tmpdir) + exp = Experiment(inst=chain, sampler=sampler) exp.data_path = Path(tmpdir) exp.run_index = 0 data = {"QubitMagnetization": np.zeros(4)} @@ -792,18 +859,18 @@ def test_save_and_reload(self): def test_save_with_filename_and_run_index_raises(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() - exp = Experiment(chain, sampler, results_root=tmpdir) + exp = Experiment(inst=chain, sampler=sampler) exp.data_path = Path(tmpdir) with self.assertRaises(ValueError): exp._save_results({}, run_index=0, filename="test.pkl.lzma") def test_save_custom_filename(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() - exp = Experiment(chain, sampler, results_root=tmpdir) + exp = Experiment(inst=chain, sampler=sampler) exp.data_path = Path(tmpdir) data = {"x": 1} exp._save_results(data, filename="custom.pkl.lzma") @@ -811,33 +878,33 @@ def test_save_custom_filename(self): def test_apply_param_sets_run_index_zero(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) 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}) + exp = Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) self.assertEqual(exp.run_index, 0) def test_apply_param_resumes_from_existing_iterations(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) 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}) + exp = Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) for i in range(3): fn = exp.data_path / f"iter{i:05d}.pkl.lzma" fn.parent.mkdir(parents=True, exist_ok=True) with lzma.open(fn, "wb") as f: pickle.dump({}, f) - exp.apply_param({"energy_scale": 1.0, "anneal_time": 1.0}) + exp.apply_param({"signed_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,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) 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}) + exp = Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) fn = exp.data_path / "iter00000.pkl.lzma" fn.parent.mkdir(parents=True, exist_ok=True) with lzma.open(fn, "wb") as f: @@ -848,25 +915,25 @@ def test_load_results_ignore_shim(self): def test_load_results_starting_iteration(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) 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}) + exp = Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) for i in range(10): fn = exp.data_path / f"iter{i:05d}.pkl.lzma" fn.parent.mkdir(parents=True, exist_ok=True) with lzma.open(fn, "wb") as f: pickle.dump({"value": i, "shimdata": {}}, f) - results = exp.load_results(num_iterations=3, starting_iteration=2) + results = exp.load_results(num_iterations=3, start_iteration=2) self.assertEqual(len(results), 3) def test_load_results_corrupted_lzma(self): with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,)) + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) 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}) + exp = Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) fn = exp.data_path / "iter00000.pkl.lzma" fn.parent.mkdir(parents=True, exist_ok=True) fn.write_bytes(b"corrupted data") @@ -875,44 +942,51 @@ def test_load_results_corrupted_lzma(self): 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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = Experiment(inst=chain, sampler=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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = Experiment( + inst=chain, + sampler=sampler, + config=ExperimentConfig(signed_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)), - } + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler() + exp = Experiment(inst=chain, sampler=sampler) + 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) @@ -920,18 +994,18 @@ def test_make_bqm_with_embeddings(self): 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 = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + exp = Experiment(inst=chain, sampler=_make_sync_sampler(), max_iterations=1) chain._load_embeddings = mock.MagicMock() - finished = exp.run_iteration([{"energy_scale": 1.0, "anneal_time": 1.0}]) + finished = exp.run_iteration([{"signed_energy_scale": 1.0, "anneal_time": 1.0}]) self.assertFalse(finished) result_files = list(exp.data_path.glob("iter*.pkl.lzma")) self.assertEqual(len(result_files), 1) + with lzma.open(result_files[0], "rb") as f: data = pickle.load(f) + self.assertIn("QubitMagnetization", data) self.assertIn("CouplerCorrelation", data) self.assertIn("shimdata", data) @@ -940,67 +1014,67 @@ def test_run_iteration_basic(self): 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)) + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + config = ExperimentConfig() exp = Experiment( - chain, _make_sync_sampler(), results_root=tmpdir, anneal_time=1.0, max_iterations=0 + inst=chain, sampler=_make_sync_sampler(), config=config, max_iterations=0 ) chain._load_embeddings = mock.MagicMock() - finished = exp.run_iteration([{"energy_scale": 1.0, "anneal_time": 1.0}]) + finished = exp.run_iteration([{"signed_energy_scale": 1.0, "anneal_time": 1.0}]) self.assertTrue(finished) self.assertEqual(list(exp.data_path.glob("iter*.pkl.lzma")), []) 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)) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler() + config = ExperimentConfig(flux_bias_shim_step=0.001) + exp = Experiment(inst=chain, sampler=sampler, config=config) + 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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler() + config = ExperimentConfig(coupler_shim_step=0.01) + exp = Experiment(inst=chain, sampler=sampler, config=config) - 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)) + 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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = Experiment(inst=chain, sampler=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) + with tempfile.TemporaryDirectory() as tmpdir: + chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = Experiment(inst=chain, sampler=sampler, config=FastAnnealExperimentConfig()) + self.assertTrue(exp.param.get("fast_anneal")) + self.assertEqual(exp.param["num_reads"], 100) if __name__ == "__main__": From 2a7e9c1b38f6d7b3075c96917c0d17a9bc98ebee Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 23 Jun 2026 14:17:52 -0700 Subject: [PATCH 11/13] Update examples and restructure tests -remove existing examples and replaced with a new 1D chain shimming example -restructure tests into separate files in a new folder -update optimize function to accept custom dimod solvers -expand docstring for ``embed_bqm`` method to explain difference from existing methods in dwave-system -promote ``optimize`` and ``make_networkx_graph`` to public methods --- .../lattice_utils/experiment/experiment.py | 86 +- .../lattice_utils/experiment/samplercall.py | 8 +- .../lattice_utils/lattice/chain.py | 6 + .../lattice_utils/lattice/embedded_lattice.py | 28 +- .../lattice_utils/lattice/lattice.py | 69 +- .../lattice_utils/lattice/optimize.py | 132 +- .../lattice_utils/lattice/triangular.py | 4 +- .../lattice_utils/observable/observable.py | 4 +- examples/example_1D_Ising_chain.py | 352 ------ examples/example_1D_Ising_chain_shim.py | 186 +++ examples/example_2D_geometric_frustration.py | 280 ----- .../add-lattice-utils-49e5efd79268e69d.yaml | 8 + tests/test_lattice_utils.py | 1081 ----------------- tests/test_lattice_utils/__init__.py | 0 tests/test_lattice_utils/_helpers.py | 103 ++ tests/test_lattice_utils/test_experiment.py | 466 +++++++ tests/test_lattice_utils/test_lattice.py | 464 +++++++ tests/test_lattice_utils/test_observable.py | 188 +++ tests/test_lattice_utils/test_utils.py | 61 + 19 files changed, 1684 insertions(+), 1842 deletions(-) delete mode 100644 examples/example_1D_Ising_chain.py create mode 100644 examples/example_1D_Ising_chain_shim.py delete mode 100644 examples/example_2D_geometric_frustration.py create mode 100644 releasenotes/notes/add-lattice-utils-49e5efd79268e69d.yaml delete mode 100644 tests/test_lattice_utils.py create mode 100644 tests/test_lattice_utils/__init__.py create mode 100644 tests/test_lattice_utils/_helpers.py create mode 100644 tests/test_lattice_utils/test_experiment.py create mode 100644 tests/test_lattice_utils/test_lattice.py create mode 100644 tests/test_lattice_utils/test_observable.py create mode 100644 tests/test_lattice_utils/test_utils.py diff --git a/dwave/experimental/lattice_utils/experiment/experiment.py b/dwave/experimental/lattice_utils/experiment/experiment.py index 1e1eced..ef4b1c2 100644 --- a/dwave/experimental/lattice_utils/experiment/experiment.py +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -77,11 +77,18 @@ def __init__( *, inst: Lattice, sampler: dimod.Sampler, + reference_energy_sampler: dimod.Sampler | None = None, + reference_energy_sampler_kwargs: dict[str, Any] | None = None, max_iterations: int | None = None, - config: ExperimentConfig = ExperimentConfig(), + config: ExperimentConfig | None = None, ): + if config is None: + config = ExperimentConfig() + self.inst = inst self.sampler = sampler + self.reference_energy_sampler = reference_energy_sampler + self.reference_energy_sampler_kwargs = reference_energy_sampler_kwargs self.param = dict(vars(config)) self.experiment_results_root = inst.data_root / "results" self.data_path = None @@ -158,6 +165,13 @@ def apply_param(self, param: dict[str, float]) -> None: param: Dictionary of parameter values to apply to the experiment. """ param = self._format_parameter_list([param])[0] + + # anneal_time and anneal_schedule are mutually exclusive + if "anneal_schedule" in param: + self.param.pop("anneal_time", None) + elif "anneal_time" in param: + self.param.pop("anneal_schedule", None) + for param_name, param_val in param.items(): self.param[param_name] = param_val @@ -270,7 +284,7 @@ def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[st """Parse a sampler response into per-embedding observable results. Args: - call: Sampler call metadata, inluding the nominal BQMs + call: Sampler call metadata, inluding the logical BQMs response: Raw sample set returned by the sampler. Returns: @@ -287,14 +301,14 @@ def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[st sample_set = {} for iemb, sample_array in enumerate(sample_arrays): sample_set[iemb] = dimod.SampleSet.from_samples_bqm( - sample_array, call.nominal_bqms[iemb] + sample_array, call.logical_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] + bqm = call.logical_bqms[iemb] obs_result = observable.evaluate(self, bqm, sample_set[iemb]) results[observable.name].append(obs_result) @@ -457,9 +471,13 @@ def _get_parameter_pathstring(self) -> str: signed_energy_scale = self.param["signed_energy_scale"] if "anneal_time" in self.param: - pathstring = f'energyscale{signed_energy_scale:0.3}/atime{self.param["anneal_time"]:010.6f}us' + pathstring = ( + f'energyscale{signed_energy_scale:0.3}/atime{self.param["anneal_time"]:010.6f}us' + ) elif "anneal_schedule" in self.param: - pathstring = f'energyscale{signed_energy_scale:0.3}/asched{self.param["anneal_schedule"]}' + pathstring = ( + f'energyscale{signed_energy_scale:0.3}/asched{self.param["anneal_schedule"]}' + ) else: raise ValueError @@ -480,14 +498,14 @@ def _get_relative_data_path(self) -> str: ] ) - def _make_nominal_bqms(self) -> list[dimod.BQM]: - """Make nominal BQMs (one per embedding) for the experiment.""" - nominal_bqm = self.inst.make_nominal_bqm() + def _make_logical_bqms(self) -> list[dimod.BQM]: + """Make logical BQMs for the experiment.""" + logical_bqm = self.inst.make_bqm() if not hasattr(self.inst, "embedding_list"): - return [nominal_bqm] + return [logical_bqm] - return [nominal_bqm] * len(self.inst.embedding_list) + return [logical_bqm] * len(self.inst.embedding_list) def _build_sampler_call(self) -> None | SamplerCall: """Build the sampler call using attributes of the experiment and instance. @@ -495,7 +513,7 @@ def _build_sampler_call(self) -> None | SamplerCall: Returns a SamplerCall. """ sampler_call = SamplerCall(run_index=self.run_index) - sampler_call.nominal_bqms = self._make_nominal_bqms() + sampler_call.logical_bqms = self._make_logical_bqms() sampler_call.shimdata = self._get_shimdata() # Here we can find out that we're finished. @@ -685,13 +703,13 @@ def _update_coupler_shim( normalization_basis = np.ones_like(orbits, dtype=bool) # Assume we have multiple embeddings of the same BQM. - bqms = sampler_call.nominal_bqms + bqms = sampler_call.logical_bqms if len(bqms) > 1 and any(bqm != bqms[0] for bqm in bqms[1:]): raise NotImplementedError("Case for distinct embedded BQMs not implemented yet.") bqm = bqms[0] - nominal_values = np.array([bqm.quadratic[edge] for edge in self.inst.edge_list]) - coupler_signs = np.sign(nominal_values) + logical_values = np.array([bqm.quadratic[edge] for edge in self.inst.edge_list]) + coupler_signs = np.sign(logical_values) for orbit_bin in range(max(orbits) + 1): bin_edges = np.argwhere(orbits == orbit_bin).ravel() if step_size != 0: @@ -708,27 +726,27 @@ def _update_coupler_shim( np.multiply(coupler_signs[bin_edges], excess) * self.param["coupler_damp"] ) - # New truncation method... previous is buggy when we mix signs of nominal values. + # New truncation method... previous is buggy when we mix signs of logical values. # Let's try being more explicit. for iemb in range(len(relative_coupler_strength)): violators = ( relative_coupler_strength[iemb, bin_edges] - * nominal_values[bin_edges] + * logical_values[bin_edges] * signed_energy_scale > 1 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - 0.99999 / nominal_values[bin_edges[violators]] / signed_energy_scale + 0.99999 / logical_values[bin_edges[violators]] / signed_energy_scale ) violators = ( relative_coupler_strength[iemb, bin_edges] - * nominal_values[bin_edges] + * logical_values[bin_edges] * signed_energy_scale < -2 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - -1.99999 / nominal_values[bin_edges[violators]] / signed_energy_scale + -1.99999 / logical_values[bin_edges[violators]] / signed_energy_scale ) # Renormalize each orbit after truncation @@ -743,30 +761,30 @@ def _update_coupler_shim( 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. + # New truncation method... previous is buggy when we mix signs of logical values. # Let's try being more explicit. for iemb in range(len(relative_coupler_strength)): violators = ( relative_coupler_strength[iemb, bin_edges] - * nominal_values[bin_edges] + * logical_values[bin_edges] * signed_energy_scale > 1 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - 0.99999 / nominal_values[bin_edges[violators]] / signed_energy_scale + 0.99999 / logical_values[bin_edges[violators]] / signed_energy_scale ) violators = ( relative_coupler_strength[iemb, bin_edges] - * nominal_values[bin_edges] + * logical_values[bin_edges] * signed_energy_scale < -2 ) relative_coupler_strength[iemb, bin_edges[violators]] = ( - -1.99999 / nominal_values[bin_edges[violators]] / signed_energy_scale + -1.99999 / logical_values[bin_edges[violators]] / signed_energy_scale ) - Q = nominal_values * relative_coupler_strength * signed_energy_scale + Q = logical_values * relative_coupler_strength * signed_energy_scale Q_max = np.max(Q) Q_min = np.min(Q) if Q_max > 1 or Q_min < -2: @@ -780,31 +798,31 @@ def _make_bqm(self, sampler_call: SamplerCall) -> dimod.BQM: signed_energy_scale = self.param["signed_energy_scale"] bqm = dimod.BQM(vartype="SPIN") if not hasattr(self.inst, "embedding_list"): - nominal_bqm = sampler_call.nominal_bqms[0] + logical_bqm = sampler_call.logical_bqms[0] for v in range(self.inst.num_spins): # Make sure variables appear in the correct order when dealing with software solvers bqm.add_variable(v) - if v in nominal_bqm.variables: - bqm.add_linear(v, nominal_bqm.linear[v]) + if v in logical_bqm.variables: + bqm.add_linear(v, logical_bqm.linear[v]) for u, v in self.inst.edge_list: - bqm.add_quadratic(u, v, nominal_bqm.quadratic[u, v] * signed_energy_scale) + bqm.add_quadratic(u, v, logical_bqm.quadratic[u, v] * signed_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() + logical_bqm = sampler_call.logical_bqms[iemb].copy() for v in range(self.inst.num_spins): # Don't touch degree-zero spins. Relevant to partial yield. - if nominal_bqm.degree(v) > 0: - bqm.add_linear(emb[v], nominal_bqm.linear[v]) + if logical_bqm.degree(v) > 0: + bqm.add_linear(emb[v], logical_bqm.linear[v]) for iedge, edge in enumerate(self.inst.edge_list): bias = ( - nominal_bqm.quadratic[*edge] + logical_bqm.quadratic[*edge] * relative_coupler_strength[iemb, iedge] * signed_energy_scale ) diff --git a/dwave/experimental/lattice_utils/experiment/samplercall.py b/dwave/experimental/lattice_utils/experiment/samplercall.py index cbd859b..81e2654 100644 --- a/dwave/experimental/lattice_utils/experiment/samplercall.py +++ b/dwave/experimental/lattice_utils/experiment/samplercall.py @@ -27,12 +27,12 @@ def __init__( self, run_index: int, shimdata: dict | None = None, - bqm: dimod.BQM | None = None, - nominal_bqms: list | None = None, + embedded_bqm: dimod.BQM | None = None, + logical_bqms: list | None = None, sampler_params: dict | None = None, ): self.run_index: int = run_index - self.bqm: dimod.BQM | None = bqm + self.bqm: dimod.BQM | None = embedded_bqm self.shimdata: dict = {} if shimdata is None else shimdata - self.nominal_bqms: list = [] if nominal_bqms is None else nominal_bqms + self.logical_bqms: list = [] if logical_bqms is None else logical_bqms self.sampler_params: dict = {} if sampler_params is None else sampler_params diff --git a/dwave/experimental/lattice_utils/lattice/chain.py b/dwave/experimental/lattice_utils/lattice/chain.py index 39391ab..c447d62 100644 --- a/dwave/experimental/lattice_utils/lattice/chain.py +++ b/dwave/experimental/lattice_utils/lattice/chain.py @@ -14,8 +14,10 @@ from collections.abc import Generator from pathlib import Path +from typing import Any from numpy.typing import NDArray +import dimod from dwave.experimental.lattice_utils.lattice import Lattice @@ -50,6 +52,8 @@ def __init__( orbit_type: str = "singleton", qubit_orbits: NDArray | None = None, coupler_orbits: NDArray | None = None, + reference_energy_sampler: dimod.Sampler | None = None, + reference_energy_sampler_kwargs: dict[str, Any] | None = None, ): self.geometry_name = "Chain" self.num_spins = dimensions[0] @@ -63,6 +67,8 @@ def __init__( orbit_type=orbit_type, qubit_orbits=qubit_orbits, coupler_orbits=coupler_orbits, + reference_energy_sampler=reference_energy_sampler, + reference_energy_sampler_kwargs=reference_energy_sampler_kwargs, ) def generate_edges(self) -> Generator[tuple[int, int]]: diff --git a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py index bb4a6fa..9263432 100644 --- a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -118,7 +118,7 @@ def generate_edges(self) -> Generator[tuple[Hashable, Hashable]]: A generator of tuples, where each tuple represents an edge between two spins in the physical lattice. """ - logical_bqm = self.logical_lattice.make_nominal_bqm() + logical_bqm = self.logical_lattice.make_bqm() # Now embed it. First make embedded spins and connect the chains. for v in logical_bqm.variables: @@ -132,25 +132,41 @@ def generate_edges(self) -> Generator[tuple[Hashable, Hashable]]: 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. + def make_bqm(self, **kwargs) -> dimod.BQM: + """Construct the physical BQM for this embedded lattice. + + Overrides the base class ``make_bqm`` for the Args: kwargs: Keyword arguments to pass to the logical lattice's - `make_nominal_bqm` method. + `make_bqm` method. Returns: - A dimod.BQM representing the embedded nominal BQM. + A dimod.BQM representing the embedded logical BQM. """ if hasattr(self, "fixed_seed"): self.logical_lattice.fixed_seed = self.fixed_seed kwargs.pop("seed", None) - return self.embed_bqm(self.logical_lattice.make_nominal_bqm(**kwargs)) + return self.embed_bqm(self.logical_lattice.make_bqm(**kwargs)) def embed_bqm(self, logical_bqm: dimod.BQM) -> dimod.BQM: """Embed a logical BQM onto the physical lattice. + This is a lattice-aware alternative to ``dwave.embedding.embed_bqm``. + The standard implementation treats each chain as an unodered set of + physical qubits and routes interactions across whatever target edges + happen to be available. Here, chains are ordered tuples and the + physical edges used for each logical interaction are chosen + deterministically by ``get_chain_connectivity``, so that the position of + a qubit within its chain carries geometric meaining (e.g. in dimerized + lattices, index 0 vs. 1 corresponds to a specific sublattice). This + allows for more structured embeddings that can be tailored to the + geometry of the logical lattice and the physics of the problem. + + Chain couplings are fixed at ``self.chain_coupling`` rather than + computed by a chain-strength heuristic. + Args: logical_bqm: A dimod.BQM representing the BQM defined on the logical variable space of the embedded lattice. diff --git a/dwave/experimental/lattice_utils/lattice/lattice.py b/dwave/experimental/lattice_utils/lattice/lattice.py index 63db7a3..3b714a1 100644 --- a/dwave/experimental/lattice_utils/lattice/lattice.py +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -17,6 +17,7 @@ from collections.abc import Generator, Hashable from abc import ABC, abstractmethod import warnings +from typing import Any import dimod from minorminer.utils.parallel_embeddings import find_multiple_embeddings @@ -61,9 +62,13 @@ def __init__( orbit_type: str = "singleton", qubit_orbits: NDArray | None = None, coupler_orbits: NDArray | None = None, + reference_energy_sampler: dimod.Sampler | None = None, + reference_energy_sampler_kwargs: dict[str, Any] | None = None, ): self.dimensions = dimensions self.data_root = Path(data_root) + self.reference_energy_sampler = reference_energy_sampler + self.reference_energy_sampler_kwargs = reference_energy_sampler_kwargs self.periodic = periodic if periodic is not None else tuple(False for _ in dimensions) if len(self.periodic) != len(self.dimensions): @@ -110,7 +115,7 @@ def embed_lattice( if exclude_qubits is None: exclude_qubits = [] - graph_bqm = dimod.to_networkx_graph(self.make_nominal_bqm()) + graph_bqm = dimod.to_networkx_graph(self.make_bqm()) graph_sampler = sampler.to_networkx_graph() graph_sampler.remove_nodes_from(exclude_qubits) @@ -119,7 +124,7 @@ def embed_lattice( self._load_embeddings(sampler) return except FileNotFoundError: - warnings.warn("No embedding file found.") + warnings.warn("No cached embedding file found, computing new embeddings.") embedding_dicts = find_multiple_embeddings( graph_bqm, @@ -128,18 +133,21 @@ def embed_lattice( embedder_kwargs={'timeout': timeout}, ) if not embedding_dicts: - raise ValueError("No embeddings found") + raise ValueError( + f"No embeddings found for {type(self).__name__}" + f"(dimensions={self.dimensions}) on " + f"{type(sampler).__name__}" + f" (target graph: {graph_sampler.number_of_nodes()} nodes, " + f"{graph_sampler.number_of_edges()} edges; timeout={timeout}s). " + f"Try increasing timeout, reducing dimensions, or relaxing exclude_qubits." + ) embeddings = np.stack([list(emb.values()) for emb in embedding_dicts]) if len(embeddings) >= min_number_of_embeddings and np.prod(embeddings.shape): self._save_embeddings(sampler, embeddings) - def make_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. + def make_bqm(self) -> dimod.BQM: + """Construct a default with BQM coupling strength values set to +1. Returns: A binary quadratic model representing the lattice with uniform @@ -195,7 +203,7 @@ def initialize_orbits( self._load_orbits() except FileNotFoundError: # calculating orbits - bqm = self.make_nominal_bqm() + bqm = self.make_bqm() self.qubit_orbits, self.coupler_orbits = get_orbits(bqm, self.edge_list) self._save_orbits() @@ -222,6 +230,34 @@ def initialize_orbits( 'Must be "global", "standard", "singleton", or "explicit".' ) + def make_networkx_graph(self) -> nx.Graph: + """Construct a NetworkX graph representation of the lattice. + + Returns: + A NetworkX graph where nodes correspond to spins and edges correspond to couplers. + """ + graph = nx.Graph() + for v in range(self.num_spins): + graph.add_node(v) + for u, v in self.edge_list: + graph.add_edge(u, v) + + return graph + + def optimize(self, bqm: dimod.BQM) -> tuple[float, NDArray, str]: + """Return the lowest energy sample by optimizing the BQM using the reference sampler. + + Returns: + A tuple containing the best energy found, the corresponding sample as a + NumPy array, and a string indicating the optimization method used. + """ + return optimize( + lattice=self, + bqm=bqm, + sampler=self.reference_energy_sampler, + sampler_kwargs=self.reference_energy_sampler_kwargs, + ) + def _get_path( self, kind: str, @@ -291,16 +327,3 @@ def _get_instance_pathstring(self) -> str: 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 index 11943d8..624735a 100644 --- a/dwave/experimental/lattice_utils/lattice/optimize.py +++ b/dwave/experimental/lattice_utils/lattice/optimize.py @@ -20,13 +20,59 @@ import dimod from numpy.typing import NDArray -__all__ = ['optimize', 'optimize_increasing_sa_sweeps'] +__all__ = ['optimize', 'ExponentialBackoffSimulatedAnnealingSampler'] + + +class ExponentialBackoffSimulatedAnnealingSampler(dimod.Sampler): + """SA sampler that doubles num_sweeps until energy stops improving or a cap is hit. + + Starts at ``min_num_sweeps`` and doubles after each round that improves the + best energy, stopping when no improvement is found or ``max_num_sweeps`` is + exceeded. + + Args: + max_num_sweeps: Upper bound on the number of sweeps per SA call. Once + ``num_sweeps`` exceeds this value, the backoff loop terminates. + min_num_sweeps: Initial number of sweeps for the first SA call. + """ + + properties = None + parameters = None + + def __init__(self, max_num_sweeps=1024, min_num_sweeps=256): + self.sampler = SimulatedAnnealingSampler() + self.max_num_sweeps = max_num_sweeps + self.min_num_sweeps = min_num_sweeps + self.properties = self.sampler.properties.copy() + self.parameters = self.sampler.parameters.copy() + + def sample(self, bqm, **parameters): + num_sweeps = parameters.pop("num_sweeps", self.min_num_sweeps) + num_reads = parameters.pop("num_reads", 256) + + best_energy = np.inf + best_sampleset = None + + while num_sweeps <= self.max_num_sweeps: + ss = self.sampler.sample(bqm, num_sweeps=num_sweeps, num_reads=num_reads, **parameters) + energy = ss.first.energy + + if energy < best_energy: + best_energy = energy + best_sampleset = ss + num_sweeps *= 2 + else: + break + + best_sampleset.info["num_sweeps_exit"] = num_sweeps + return best_sampleset def optimize( lattice: Lattice, bqm: dimod.BQM, - sa_kwargs: dict[str, Any] | None = None, + sampler: dimod.Sampler | None = None, + sampler_kwargs: dict[str, Any] | None = None, ) -> tuple[float, NDArray, str]: """Return the best sample found by optimizing the BQM using simulated annealing. @@ -44,72 +90,42 @@ def optimize( 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``. + sampler: A dimod Sampler to use for optimization of the reference energy. + If None, a default ExponentialBackoffSimulatedAnnealingSampler will + be used. + sampler_kwargs: Optional keyword arguments to pass to the provided + sampler, such as ``num_reads`` and ``num_sweeps`` in the case of a + SA sampler. Returns: A tuple containing the best energy found, the corresponding sample as a NumPy array, and a string indicating the optimization method used. """ - if sa_kwargs is None: - sa_kwargs = {} + if sampler_kwargs is None: + sampler_kwargs = {} + + if sampler is None: + sampler = ExponentialBackoffSimulatedAnnealingSampler() + + reference_energy = np.inf + reference_sample = None # If the lattice is embedded, we should optimize the logical lattice if hasattr(lattice, "logical_lattice"): _, logical_sample, _ = optimize( - lattice.logical_lattice, lattice.unembed_bqm(bqm), sa_kwargs=sa_kwargs + lattice.logical_lattice, + lattice.unembed_bqm(bqm), + sampler=sampler, + sampler_kwargs=sampler_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. + reference_sample = lattice.embed_sample(logical_sample) + reference_energy = bqm.energy(reference_sample) - 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. + sampleset = sampler.sample(bqm, **sampler_kwargs) + best = sampleset.first - 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 + if best.energy < reference_energy: + sample = np.array([best.sample[v] for v in bqm.variables]) + return best.energy, sample, type(sampler).__name__ - return reference_energy, reference_sample, "sa_exponential" + return reference_energy, reference_sample, type(sampler).__name__ diff --git a/dwave/experimental/lattice_utils/lattice/triangular.py b/dwave/experimental/lattice_utils/lattice/triangular.py index 6f0c712..06eaf5d 100644 --- a/dwave/experimental/lattice_utils/lattice/triangular.py +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -103,7 +103,7 @@ def coordinates(self, node: int) -> tuple[int, int]: length_y = self.dimensions[0] return node % length_y, node // length_y - def make_nominal_bqm(self) -> dimod.BQM: + def make_bqm(self) -> dimod.BQM: """Construct the nominal triangular lattice BQM. If ``halve_boundary_couplers`` is True, couplers that are on the boundary @@ -112,7 +112,7 @@ def make_nominal_bqm(self) -> dimod.BQM: Returns: A dimod.BQM representing the nominal triangular lattice. """ - graph = self._make_networkx_graph() + graph = self.make_networkx_graph() bqm = dimod.BQM(vartype="SPIN") for v in range(self.num_spins): diff --git a/dwave/experimental/lattice_utils/observable/observable.py b/dwave/experimental/lattice_utils/observable/observable.py index 94dbf9d..a267e95 100644 --- a/dwave/experimental/lattice_utils/observable/observable.py +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -232,9 +232,9 @@ def evaluate( # And if we can't load, we generate a reference sample. if experiment is not None: - energy, sample, method_string = experiment.inst._optimize(bqm) + energy, sample, method_string = experiment.inst.optimize(bqm) elif inst is not None: - energy, sample, method_string = inst._optimize(bqm) + energy, sample, method_string = inst.optimize(bqm) else: raise ValueError( "Must provide either an experiment or a lattice to compute reference energy." diff --git a/examples/example_1D_Ising_chain.py b/examples/example_1D_Ising_chain.py deleted file mode 100644 index 49e999e..0000000 --- a/examples/example_1D_Ising_chain.py +++ /dev/null @@ -1,352 +0,0 @@ -# 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"), -] - -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 signed_energy_scale in ENERGY_SCALES: - config = experiment.FastAnnealExperimentConfig( - signed_energy_scale=signed_energy_scale, - coupler_shim_step=0.05, - flux_bias_shim_step=3e-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. - - # Note that only running 5 iterations will give poor results for the - # kink-kink correlator, which is fairly sensitive and requires a - # converged shim, but this is just for demonstration purposes. - 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] - - while True: - done = exp.run_iteration(parameter_list, progress=True) - 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) - 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)) - ) - - dimensions = "x".join(str(dim) for dim in inst.dimensions) - signed_energy_scale = exp.param["signed_energy_scale"] - title = f"1D chain, {dimensions}, J={signed_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) - fig.tight_layout() - fig.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, signed_energy_scale] = np.asarray(frust) - kkc_dict[sampler.solver.name, signed_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) -fig2.tight_layout(rect=(0.05, 0.05, 1, 0.95), w_pad=4.0) - -for isampler, sampler in enumerate(samplers): - - for signed_energy_scale in [-1.8, 0.1]: - M = kd_dict[sampler.solver.name, signed_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, signed_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) -fig3.tight_layout(rect=(0.05, 0.05, 1, 0.95), w_pad=4.0) - -for isampler, sampler in enumerate(samplers): - - for signed_energy_scale in [-1.8, 0.1]: - magnetization = kkc_dict[sampler.solver.name, signed_energy_scale][0].T - kd = np.mean(kd_dict[sampler.solver.name, signed_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_1D_Ising_chain_shim.py b/examples/example_1D_Ising_chain_shim.py new file mode 100644 index 0000000..8119384 --- /dev/null +++ b/examples/example_1D_Ising_chain_shim.py @@ -0,0 +1,186 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shimming example for 1D Ising chain.""" + +from pathlib import Path +import 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 + +# Set up the parameters + +# Get a QPU sampler with Zephyr topology. +sampler = DWaveSampler(solver=dict(topology__type="zephyr")) + +NUM_SPINS = 256 +SIGNED_ENERGY_SCALE = -1.6 + +# We will simulate four orders of magnitude in anneal time. +# File format rounds to the nearest picosecond, so we will do so explicitly here. +ANNEAL_TIMES = np.round(np.geomspace(0.01, 100, 3), 6) + +errorbar_style = {"marker": "", "linestyle": "", "capsize": 2} +point_style = {"linestyle": "", "markersize": 5} +cm = plt.get_cmap("tab10") + +# 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] + +# Make a lattice instance for a periodic 256-spin chain, so we can embed it. +inst = lattice.Chain( + dimensions=(NUM_SPINS,), + data_root=data_root, + periodic=(True,), + orbit_type="standard", +) + +# Find parallel embeddings of the lattice heuristically. The embed_lattice +# function is heuristic and is run here with a default timeout (10s) and no +# tuning of any parameters. Larger and more complex lattices can take longer +# to embed. +inst.embed_lattice(sampler) + +# Time to make an experiment. The aim in this example is to demonstrate a coupler +# and flux-bias shim on a chain, at fixed energy scale and varying anneal time. +# Since the susceptibility to changes in the shim changes as a function of anneal +# time, we will run each anneal time using separate shim steps. +# We will also set the orbit_type to 'standard', allowing the use of graph +# automorphisms to determine symmetries in the system that can be exploited by +# shimming. In this case, all couplers are equivalent (they go in the same orbit) +# so the coupler shim will compel them all to have the same spin-spin correlation +# for a given parameterization. +flux_bias_shim_step = { + 0.01: 5e-6, + 1.0: 2e-6, + 100: 0.5e-6, +} +coupler_shim_step = { + 0.01: 0.05, + 1.0: 0.2, + 100: 1.0, +} +for anneal_time in ANNEAL_TIMES: + + config = experiment.FastAnnealExperimentConfig( + signed_energy_scale=SIGNED_ENERGY_SCALE, + coupler_shim_step=coupler_shim_step[anneal_time], + flux_bias_shim_step=flux_bias_shim_step[anneal_time], + ) + exp = experiment.Experiment(inst=inst, sampler=sampler, max_iterations=100, config=config) + + # Make parameter list. We will only vary anneal time. + for _ in range(120): + done = exp.run_iteration([{"anneal_time": anneal_time}], progress=True) + if done: + break + + +# We will make a dict for the data we want to analyze. For each anneal time +# we will load all iterations into a corresponding dict entry. +# Disjoint embeddings are given along a separate axis in the results, except +# flux biases and anneal offsets, which are given as a single +# array since they are indexed by physical qubit. + +mag = {} # average qubit magnetization +frust = {} # average coupler frustration (kink density) +cshim = {} # coupler shim +fbshim = {} # flux bias shim +for anneal_time in ANNEAL_TIMES: + exp.apply_param({"anneal_time": anneal_time}) + res = exp.load_results() + mag[anneal_time] = np.array([it["QubitMagnetization"] for it in res]) + frust[anneal_time] = np.array([it["CouplerFrustration"] for it in res]) + cshim[anneal_time] = np.array( + [it["shimdata"]["relative_coupler_strength"].ravel() for it in res] + ) + fbshim[anneal_time] = np.array([it["shimdata"]["flux_biases"] for it in res]) + +title = ( + f"1D chain shim, " + f"{'x'.join([str(dim) for dim in inst.dimensions])}, " + f"J={exp.param['signed_energy_scale']}, " + f"{sampler.solver.name}" +) +fig, axes = plt.subplots(3, 6, figsize=(16, 8), sharex="col", sharey="col") +fig.suptitle(title, fontsize=16) + +for iat, anneal_time in enumerate(ANNEAL_TIMES): + + # Plot std of qubit magnetization + ax = axes[iat, 0] + ax.plot(mag[anneal_time].std(axis=(1, 2)), label=r"Mag std") + ax.set_ylabel("Magnetization std") + + # Plot histograms of first and last iterations + ax = axes[iat, 1] + ax.hist(mag[anneal_time][:5].ravel(), label="First 5 iterations", alpha=0.5) + ax.hist(mag[anneal_time][-5:].ravel(), label="Last 5 iterations", alpha=0.5) + ax.set_yticks([]) + ax.set_ylabel("Frequency") + + # Plot std of coupler frustration + ax = axes[iat, 2] + ax.plot(frust[anneal_time].std(axis=(1, 2)), label=r"Frust std") + ax.set_ylabel("Frustration std") + + # Plot histograms of first and last iterations + ax = axes[iat, 3] + ax.hist(frust[anneal_time][:5].ravel(), label="First iteration", alpha=0.5) + ax.hist(frust[anneal_time][-5:].ravel(), label="Last iteration", alpha=0.5) + ax.set_yticks([]) + ax.set_ylabel("Frequency") + + # Plot flux biases + ax = axes[iat, 4] + ax.plot(fbshim[anneal_time], alpha=0.2) + ax.set_ylabel("Flux bias") + + # Plot coupler shim + ax = axes[iat, 5] + ax.plot(cshim[anneal_time], alpha=0.2) + ax.set_ylabel("Rel. cplr. strength") + + +axes[iat, 0].set_xlabel("Iteration") +axes[iat, 1].set_xlabel("Magnetization") +axes[iat, 1].legend() +axes[iat, 2].set_xlabel("Iteration") +axes[iat, 3].set_xlabel("Frustration") +axes[iat, 3].legend() +axes[iat, 4].set_xlabel("Iteration") +axes[iat, 5].set_xlabel("Iteration") + +axes[0, 0].set_title("Magnetization std") +axes[0, 1].set_title("Magnetization") +axes[0, 2].set_title("Frustration std") +axes[0, 3].set_title("Frustration") +axes[0, 4].set_title("Flux offset shim") +axes[0, 5].set_title("Coupler shim") + +fig.tight_layout() + +filename = title +for bad_symbol in "/: ;,": + filename = filename.replace(bad_symbol, "_") +fig.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 deleted file mode 100644 index 9372712..0000000 --- a/examples/example_2D_geometric_frustration.py +++ /dev/null @@ -1,280 +0,0 @@ -# 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) - -config = experiment.FastAnnealExperimentConfig( - signed_energy_scale=0.8, - coupler_shim_step=0.1, - flux_bias_shim_step=5e-6, - readout_thermalization=100, - num_reads=100, -) -exp = experiment.Experiment(inst=inst, sampler=sampler, max_iterations=210, config=config) -exp.observables_to_collect.add(observable.TriangularOP()) - -# Make parameter list -parameter_list = [{"anneal_time": rate} for rate in ANNEAL_TIMES] - -for _ in range(1000): - done = exp.run_iteration(parameter_list, progress=True) - 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["signed_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) -fig.tight_layout() -fig.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["signed_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/releasenotes/notes/add-lattice-utils-49e5efd79268e69d.yaml b/releasenotes/notes/add-lattice-utils-49e5efd79268e69d.yaml new file mode 100644 index 0000000..27fd3bb --- /dev/null +++ b/releasenotes/notes/add-lattice-utils-49e5efd79268e69d.yaml @@ -0,0 +1,8 @@ +--- +features: + - | + Add ``dwave.experimental.lattice_utils`` submodule with utilities for + constructing lattice graphs (``Chain``, ``Triangular``, + ``DimerizedTriangular``, ``EmbeddedLattice``), evaluating physics + observables on sample sets, and running shimmed Ising experiments on + QPU samplers. diff --git a/tests/test_lattice_utils.py b/tests/test_lattice_utils.py deleted file mode 100644 index 058c5a9..0000000 --- a/tests/test_lattice_utils.py +++ /dev/null @@ -1,1081 +0,0 @@ -# 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 ( - FastAnnealExperimentConfig, -) -from dwave.experimental.lattice_utils.experiment.experiment import ExperimentConfig - - -def _make_triangular( - ly=3, - lx=3, - periodic=(True, False), - orbit_type="singleton", - halve_boundary_couplers=False, -): - with tempfile.TemporaryDirectory() as tmpdir: - return Triangular( - dimensions=(ly, lx), - periodic=periodic, - data_root=tmpdir, - 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, signed_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 = {"signed_energy_scale": signed_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): - with tempfile.TemporaryDirectory() as tmpdir: - return EmbeddedLattice( - logical_lattice=Chain( - dimensions=(len(chain_nodes),), - periodic=(False,), - data_root=tmpdir, - ), - chain_nodes=chain_nodes, - dimensions=(sum(len(chain) for chain in chain_nodes.values()),), - periodic=(False,), - ) - - -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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) - self.assertEqual(chain.num_spins, 6) - self.assertEqual(chain.num_edges, 6) - self.assertIn((5, 0), chain.edge_list) - - def test_non_periodic(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(6,), periodic=(False,), data_root=tmpdir) - self.assertEqual(chain.num_spins, 6) - self.assertEqual(chain.num_edges, 5) - self.assertNotIn((5, 0), chain.edge_list) - - def test_single_node_periodic(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(1,), periodic=(True,), data_root=tmpdir) - self.assertEqual(chain.num_spins, 1) - self.assertEqual(chain.num_edges, 0) - - def test_two_node_periodic(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(2,), periodic=(True,), data_root=tmpdir) - self.assertEqual(chain.num_edges, 2) - - def test_geometry_name(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) - self.assertEqual(chain.geometry_name, "Chain") - - -class TestLattice(unittest.TestCase): - def test_default_periodic(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - self.assertFalse(chain.periodic[0]) - - def test_edge_list_sorted(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(5,), periodic=(False,), data_root=tmpdir) - for u, v in chain.edge_list: - self.assertLess(u, v) - - def test_bqm_structure(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - 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): - with tempfile.TemporaryDirectory() as tmpdir: - bqm = Chain(dimensions=(3,), data_root=tmpdir).make_nominal_bqm() - self.assertEqual(bqm.vartype, dimod.SPIN) - - def test_orbit_singleton(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), orbit_type="singleton", data_root=tmpdir) - np.testing.assert_array_equal(chain.qubit_orbits, np.arange(4)) - np.testing.assert_array_equal(chain.coupler_orbits, np.arange(chain.num_edges)) - - def test_orbit_global(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), orbit_type="global", data_root=tmpdir) - np.testing.assert_array_equal(chain.qubit_orbits, np.zeros(4, dtype=int)) - np.testing.assert_array_equal( - chain.coupler_orbits, np.zeros(chain.num_edges, dtype=int) - ) - - def test_orbit_explicit(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain( - dimensions=(4,), - periodic=(True,), - orbit_type="explicit", - qubit_orbits=np.array([0, 0, 1, 1]), - coupler_orbits=np.array([0, 0, 1, 1]), - data_root=tmpdir, - ) - np.testing.assert_array_equal(chain.qubit_orbits, [0, 0, 1, 1]) - - def test_unknown_orbit_type(self): - with tempfile.TemporaryDirectory() as tmpdir: - with self.assertRaises(ValueError): - Chain(dimensions=(4,), periodic=(True,), orbit_type="bogus", data_root=tmpdir) - - def test_get_path_invalid_kind(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - with self.assertRaises(ValueError): - chain._get_path(None, "invalid") - - def test_standard_orbit_save_and_load(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain( - dimensions=(4,), - periodic=(True,), - orbit_type="standard", - data_root=tmpdir, - ) - self.assertIsNotNone(chain.qubit_orbits) - self.assertIsNotNone(chain.coupler_orbits) - # Second instantiation should load from disk - chain2 = Chain( - dimensions=(4,), - periodic=(True,), - orbit_type="standard", - data_root=tmpdir, - ) - np.testing.assert_array_equal(chain.qubit_orbits, chain2.qubit_orbits) - - def test_nested_embedded_raises(self): - with tempfile.TemporaryDirectory() as tmpdir: - base_chain = Chain( - dimensions=(2,), - periodic=(False,), - data_root=tmpdir, - ) - embedded_once = EmbeddedLattice( - logical_lattice=base_chain, - chain_nodes={0: (0, 1), 1: (2, 3)}, - dimensions=(4,), - periodic=(False,), - ) - with self.assertRaises(NotImplementedError): - EmbeddedLattice( - logical_lattice=embedded_once, - chain_nodes={0: (0, 1), 1: (2, 3), 2: (4, 5), 3: (6, 7)}, - dimensions=(8,), - periodic=(False,), - ) - - def test_embed_no_embeddings_found(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = mock.MagicMock() - type(sampler).__name__ = "MockDWaveSampler" - sampler.to_networkx_graph.return_value = chain._make_networkx_graph() - - with mock.patch( - "dwave.experimental.lattice_utils.lattice.lattice.find_multiple_embeddings", - return_value=[], - ): - with self.assertRaises(ValueError): - chain.embed_lattice(sampler, try_to_load=False, timeout=1) - - def test_embed_load_existing(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = mock.MagicMock() - type(sampler).__name__ = "MockDWaveSampler" - sampler.to_networkx_graph.return_value = chain._make_networkx_graph() - - embeddings = np.array([[0, 1, 2, 3]]) - chain._save_embeddings(sampler, embeddings) - - chain.embed_lattice(sampler, try_to_load=True, data_root=tmpdir) - np.testing.assert_array_equal(chain.embedding_list, embeddings) - - def test_embed_find_and_save(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = mock.MagicMock() - type(sampler).__name__ = "MockDWaveSampler" - sampler.to_networkx_graph.return_value = chain._make_networkx_graph() - - emb_dict = {i: i for i in range(4)} - with mock.patch( - "dwave.experimental.lattice_utils.lattice.lattice.find_multiple_embeddings", - return_value=[emb_dict], - ): - chain.embed_lattice(sampler, try_to_load=False, timeout=1, data_root=tmpdir) - # Verify embedding was found and saved - emb_path = chain._get_path("embedding", sampler_name="MockDWaveSampler") - self.assertTrue(emb_path.exists()) - - -class TestTriangular(unittest.TestCase): - def test_basic_construction(self): - with tempfile.TemporaryDirectory() as tmpdir: - tri = Triangular( - dimensions=(3, 3), - periodic=(True, False), - data_root=tmpdir, - orbit_type="singleton", - halve_boundary_couplers=False, - ) - self.assertEqual(tri.num_spins, 9) - self.assertGreater(tri.num_edges, 0) - self.assertEqual(tri.geometry_name, "Triangular") - - def test_coordinates(self): - 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): - with tempfile.TemporaryDirectory() as tmpdir: - dt = DimerizedTriangular( - dimensions=(3, 3), periodic=(True, False), orbit_type="singleton", data_root=tmpdir - ) - self.assertEqual(dt.geometry_name, "DimerizedTriangular") - self.assertIsNotNone(dt.logical_lattice) - self.assertEqual(dt.num_spins, 18) - - def test_chain_connectivity_self(self): - with tempfile.TemporaryDirectory() as tmpdir: - dt = DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) - cc = dt.get_chain_connectivity(0) - self.assertEqual(cc, ((0, 1),)) - - def test_chain_connectivity_cases(self): - with tempfile.TemporaryDirectory() as tmpdir: - dt = DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) - cases = [ - ((0,), ((0, 1),)), - ((0, 1), ((1, 0),)), - ((0, 3), ((1, 0),)), - ] - for args, expected in cases: - with self.subTest(args=args): - self.assertEqual(dt.get_chain_connectivity(*args), expected) - - -class TestEmbeddedLattice(unittest.TestCase): - def test_embed_sample(self): - with tempfile.TemporaryDirectory() as tmpdir: - dt = DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) - logical_sample = np.array([1, -1, 1, -1, 1, -1, 1, -1, 1]) - embedded = dt.embed_sample(logical_sample) - self.assertEqual(len(embedded), dt.num_spins) - # Each chain should have the same value - for spin, chain in dt.chain_nodes.items(): - for node in chain: - self.assertEqual(embedded[node], logical_sample[spin]) - - def test_unembed_sample(self): - chain_nodes = {0: (0, 1, 2), 1: (3, 4, 5)} - with tempfile.TemporaryDirectory() as tmpdir: - embedded = EmbeddedLattice( - logical_lattice=Chain( - dimensions=(len(chain_nodes),), - periodic=(False,), - data_root=tmpdir, - ), - chain_nodes=chain_nodes, - dimensions=(sum(len(chain) for chain in chain_nodes.values()),), - periodic=(False,), - ) - physical_sample = np.array([1, 1, -1, -1, -1, 1]) - logical = embedded.unembed_sample(physical_sample) - np.testing.assert_array_equal(logical, np.array([1, -1])) - - def test_unembed_sample_breaks_ties_randomly(self): - 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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - 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): - with tempfile.TemporaryDirectory() as tmpdir: - dt = DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) - 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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - 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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - 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""" - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - 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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - 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, signed_energy_scale=1.0) - np.testing.assert_array_almost_equal(SampleEnergy().evaluate(exp_pos, bqm, ss), [3.0]) - - def test_bitpacked_spins(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - 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,), data_root=tmpdir) - 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,), data_root=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,), data_root=tmpdir) - 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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - bqm = chain.make_nominal_bqm() - exp = _make_mock_experiment(chain) - - path = get_reference_energy_path(bqm, experiment=exp) - 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, dummy_experiment_data_dict=dummy) - self.assertTrue(str(path2).endswith(".txt")) - - with tempfile.TemporaryDirectory() as tmpdir: - path3 = get_reference_energy_path(bqm, experiment=exp, root=tmpdir) - self.assertIn(tmpdir, str(path3)) - - -class TestKinks(unittest.TestCase): - def test_all_aligned(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) - 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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) - 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): - with tempfile.TemporaryDirectory() as tmpdir: - dt = DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) - 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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - self.assertEqual(exp.param["signed_energy_scale"], 1.0) - self.assertEqual(exp.param["num_reads"], 100) - self.assertIs(exp.inst, chain) - - def test_data_path_with_schedule(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.param["anneal_schedule"] = [(0, 1), (5, 0.5)] - del exp.param["anneal_time"] - exp.apply_param({"signed_energy_scale": 1.0, "anneal_schedule": [(0, 1), (5, 0.5)]}) - self.assertIn("asched", str(exp.data_path)) - - def test_apply_param_unknown_sampler_raises(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - type(sampler).__name__ = "UnknownSampler" - exp = Experiment(inst=chain, sampler=sampler) - with self.assertRaises(TypeError): - exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) - - def test_apply_param_no_anneal_or_schedule_raises(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - del exp.param["anneal_time"] - with self.assertRaises(ValueError): - exp.apply_param({"signed_energy_scale": 1.0}) - - def test_initial_shim_no_embeddings(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - chain.embedding_list = np.array([[0, 1, 2, 3]]) - sampler = _make_mock_sampler(num_qubits=128) - exp = Experiment(inst=chain, sampler=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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - chain.embedding_list = np.array([[0, 1, 2, 3]]) - sampler = _make_mock_sampler(num_qubits=128) - fb = np.ones(128) * 0.01 - exp = Experiment(inst=chain, sampler=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,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.run_index = 1 - exp.data_path = Path(tmpdir) - - 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,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.run_index = 1 - exp.data_path = Path(tmpdir) - - fn = Path(tmpdir) / "iter00000.pkl.lzma" - fn.touch() - - with self.assertRaises(FileNotFoundError): - exp._load_shim() - - def test_load_shim_corrupted_file(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.run_index = 1 - exp.data_path = Path(tmpdir) - fn = Path(tmpdir) / "iter00000.pkl.lzma" - fn.write_bytes(b"not a valid lzma file") - - with self.assertRaises(OSError): - exp._load_shim() - - def test_load_shim_missing_file(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - 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,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.data_path = Path(tmpdir) - exp.run_index = 0 - data = {"QubitMagnetization": np.zeros(4)} - exp._save_results(data) - fn = Path(tmpdir) / "iter00000.pkl.lzma" - self.assertTrue(fn.exists()) - - with lzma.open(fn, "rb") as f: - loaded = pickle.load(f) - np.testing.assert_array_equal(loaded["QubitMagnetization"], np.zeros(4)) - - def test_save_with_filename_and_run_index_raises(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.data_path = Path(tmpdir) - with self.assertRaises(ValueError): - exp._save_results({}, run_index=0, filename="test.pkl.lzma") - - def test_save_custom_filename(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.data_path = Path(tmpdir) - data = {"x": 1} - exp._save_results(data, filename="custom.pkl.lzma") - self.assertTrue((Path(tmpdir) / "custom.pkl.lzma").exists()) - - def test_apply_param_sets_run_index_zero(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) - self.assertEqual(exp.run_index, 0) - - def test_apply_param_resumes_from_existing_iterations(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) - for i in range(3): - fn = exp.data_path / f"iter{i:05d}.pkl.lzma" - fn.parent.mkdir(parents=True, exist_ok=True) - with lzma.open(fn, "wb") as f: - pickle.dump({}, f) - - exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) - self.assertEqual(exp.run_index, 3) - - def test_load_results_ignore_shim(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) - fn = exp.data_path / "iter00000.pkl.lzma" - fn.parent.mkdir(parents=True, exist_ok=True) - with lzma.open(fn, "wb") as f: - pickle.dump({"value": 0, "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,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) - for i in range(10): - fn = exp.data_path / f"iter{i:05d}.pkl.lzma" - fn.parent.mkdir(parents=True, exist_ok=True) - with lzma.open(fn, "wb") as f: - pickle.dump({"value": i, "shimdata": {}}, f) - - results = exp.load_results(num_iterations=3, start_iteration=2) - self.assertEqual(len(results), 3) - - def test_load_results_corrupted_lzma(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) - fn = exp.data_path / "iter00000.pkl.lzma" - fn.parent.mkdir(parents=True, exist_ok=True) - fn.write_bytes(b"corrupted data") - - with self.assertRaises(lzma.LZMAError): - exp.load_results(num_iterations=1) - - def test_generate_data_type_conversions(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment( - inst=chain, - sampler=sampler, - config=ExperimentConfig(signed_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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - chain.embedding_list = np.array([[0, 1, 2, 3]]) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - 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,), data_root=tmpdir) - exp = Experiment(inst=chain, sampler=_make_sync_sampler(), max_iterations=1) - chain._load_embeddings = mock.MagicMock() - finished = exp.run_iteration([{"signed_energy_scale": 1.0, "anneal_time": 1.0}]) - - self.assertFalse(finished) - result_files = list(exp.data_path.glob("iter*.pkl.lzma")) - self.assertEqual(len(result_files), 1) - - with lzma.open(result_files[0], "rb") as f: - data = pickle.load(f) - - self.assertIn("QubitMagnetization", data) - self.assertIn("CouplerCorrelation", data) - self.assertIn("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,), data_root=tmpdir) - config = ExperimentConfig() - exp = Experiment( - inst=chain, sampler=_make_sync_sampler(), config=config, max_iterations=0 - ) - chain._load_embeddings = mock.MagicMock() - finished = exp.run_iteration([{"signed_energy_scale": 1.0, "anneal_time": 1.0}]) - - self.assertTrue(finished) - self.assertEqual(list(exp.data_path.glob("iter*.pkl.lzma")), []) - - def test_flux_bias_shim_basic_update(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - chain.embedding_list = np.array([[0, 1, 2, 3]]) - sampler = _make_mock_sampler() - config = ExperimentConfig(flux_bias_shim_step=0.001) - exp = Experiment(inst=chain, sampler=sampler, config=config) - 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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) - chain.embedding_list = np.array([[0, 1, 2, 3]]) - sampler = _make_mock_sampler() - config = ExperimentConfig(coupler_shim_step=0.01) - exp = Experiment(inst=chain, sampler=sampler, config=config) - - 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): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler) - exp.already_initialized = False - shimdata = exp._get_shimdata() - self.assertEqual(shimdata["total_iterations"], 0) - - -class TestFastAnnealExperiment(unittest.TestCase): - def test_default_params(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) - sampler = _make_mock_sampler() - exp = Experiment(inst=chain, sampler=sampler, config=FastAnnealExperimentConfig()) - self.assertTrue(exp.param.get("fast_anneal")) - self.assertEqual(exp.param["num_reads"], 100) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_lattice_utils/__init__.py b/tests/test_lattice_utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_lattice_utils/_helpers.py b/tests/test_lattice_utils/_helpers.py new file mode 100644 index 0000000..8fe25ad --- /dev/null +++ b/tests/test_lattice_utils/_helpers.py @@ -0,0 +1,103 @@ +# 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. + +"""Shared test fixtures for the ``lattice_utils`` test package.""" + +from unittest import mock + +import dimod +import numpy as np + +from dwave.experimental.lattice_utils import lattice + +__all__ = [ + "_make_triangular", + "_make_mock_sampler", + "_make_mock_experiment", + "_make_embedded_chain", +] + + +def _make_triangular( + data_root, + ly=3, + lx=3, + periodic=(True, False), + orbit_type="singleton", + halve_boundary_couplers=False, +): + return lattice.Triangular( + dimensions=(ly, lx), + periodic=periodic, + data_root=data_root, + orbit_type=orbit_type, + halve_boundary_couplers=halve_boundary_couplers, + ) + + +def _make_mock_sampler( + num_qubits=128, + nodelist=None, + solver_name="TestSolver", + type_name="DWaveSampler", + *, + sync_response=False, +): + """Build a mock sampler resembling DWaveSampler. + + Production code detects the sampler via ``type(sampler).__name__`` and reads + ``nodelist``, ``properties["num_qubits"]``, and ``solver.name``. When + ``sync_response`` is True, ``sampler.sample(...)`` returns a mock response + mimicking the DWaveSampler async interface (``.done()`` -> True, + ``.samples()`` -> all-ones ndarray) used by ``run_iteration`` tests. + """ + sampler = mock.MagicMock(spec=dimod.Sampler) + type(sampler).__name__ = type_name + if nodelist is None: + nodelist = list(range(num_qubits)) + sampler.nodelist = nodelist + sampler.properties = {"num_qubits": num_qubits} + sampler.solver = mock.MagicMock() + sampler.solver.name = solver_name + if sync_response: + response = mock.MagicMock() + response.done.return_value = True + response.samples.return_value = np.ones((10, num_qubits), dtype=float) + sampler.sample.return_value = response + return sampler + + +def _make_mock_experiment(inst, signed_energy_scale=1.0, run_index=0, num_random_instances=1): + """Return a lightweight mock Experiment with .inst and .param.""" + exp = mock.MagicMock() + exp.inst = inst + exp.param = { + "signed_energy_scale": signed_energy_scale, + "num_random_instances": num_random_instances, + } + exp.run_index = run_index + return exp + + +def _make_embedded_chain(chain_nodes, data_root): + return lattice.EmbeddedLattice( + logical_lattice=lattice.Chain( + dimensions=(len(chain_nodes),), + periodic=(False,), + data_root=data_root, + ), + chain_nodes=chain_nodes, + dimensions=(sum(len(chain) for chain in chain_nodes.values()),), + periodic=(False,), + ) diff --git a/tests/test_lattice_utils/test_experiment.py b/tests/test_lattice_utils/test_experiment.py new file mode 100644 index 0000000..1302dcb --- /dev/null +++ b/tests/test_lattice_utils/test_experiment.py @@ -0,0 +1,466 @@ +# 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 import experiment, lattice +from tests.test_lattice_utils._helpers import _make_mock_sampler + + +class TestSamplerCall(unittest.TestCase): + def test_defaults(self): + sc = experiment.SamplerCall(run_index=0) + self.assertEqual(sc.run_index, 0) + self.assertIsNone(sc.bqm) + self.assertEqual(sc.shimdata, {}) + self.assertEqual(sc.logical_bqms, []) + self.assertEqual(sc.sampler_params, {}) + + def test_with_values(self): + bqm = dimod.BQM(vartype="SPIN") + sc = experiment.SamplerCall( + run_index=5, + embedded_bqm=bqm, + shimdata={"total_iterations": 1}, + logical_bqms=[bqm], + sampler_params={"num_reads": 100}, + ) + self.assertEqual(sc.run_index, 5) + self.assertIs(sc.bqm, bqm) + self.assertEqual(sc.shimdata["total_iterations"], 1) + + +class TestExperimentInit(unittest.TestCase): + def test_default_params(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + self.assertEqual(exp.param["signed_energy_scale"], 1.0) + self.assertEqual(exp.param["num_reads"], 100) + self.assertIs(exp.inst, chain) + + +class TestApplyParam(unittest.TestCase): + def test_data_path_with_schedule(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_schedule": [(0, 1), (5, 0.5)]}) + self.assertIn("asched", str(exp.data_path)) + + def test_apply_param_unknown_sampler_raises(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler(type_name="UnknownSampler") + exp = experiment.Experiment(inst=chain, sampler=sampler) + with self.assertRaises(TypeError): + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + + def test_apply_param_sets_run_index_zero(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + self.assertEqual(exp.run_index, 0) + + def test_apply_param_resumes_from_existing_iterations(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + for i in range(3): + fn = exp.data_path / f"iter{i:05d}.pkl.lzma" + fn.parent.mkdir(parents=True, exist_ok=True) + with lzma.open(fn, "wb") as f: + pickle.dump({}, f) + + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + self.assertEqual(exp.run_index, 3) + + +class TestShimdata(unittest.TestCase): + def test_initial_shim_no_embeddings(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.already_initialized = False + shimdata = exp._make_initial_shim() + self.assertEqual(shimdata["total_iterations"], 0) + self.assertNotIn("flux_biases", shimdata) + + def test_initial_shim_with_embeddings(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler(num_qubits=128) + exp = experiment.Experiment(inst=chain, sampler=sampler) + 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): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler(num_qubits=128) + fb = np.ones(128) * 0.01 + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.param["flux_biases"] = fb + 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 = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.run_index = 1 + exp.data_path = Path(tmpdir) + + 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 = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.run_index = 1 + exp.data_path = Path(tmpdir) + + fn = Path(tmpdir) / "iter00000.pkl.lzma" + fn.touch() + + with self.assertRaises(FileNotFoundError): + exp._load_shim() + + def test_load_shim_corrupted_file(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.run_index = 1 + exp.data_path = Path(tmpdir) + fn = Path(tmpdir) / "iter00000.pkl.lzma" + fn.write_bytes(b"not a valid lzma file") + + with self.assertRaises(OSError): + exp._load_shim() + + def test_load_shim_missing_file(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.data_path = Path(tmpdir) + exp.run_index = 1 + + with self.assertRaises(FileNotFoundError): + exp._load_shim() + + def test_get_shimdata_not_initialized(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.already_initialized = False + shimdata = exp._get_shimdata() + self.assertEqual(shimdata["total_iterations"], 0) + + +class TestCouplerShim(unittest.TestCase): + def test_coupler_shim_basic_update(self): + """rcs += step_size * (frust - mean(frust)) within each orbit bin. + + signed_energy_scale=0.5 keeps the effective coupler |J*rcs*scale| below the + truncation thresholds (>1 / <-2) so we can assert the raw update math. + """ + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + orbit_type="global", # all edges in one bin -> update is non-trivial + ) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler() + exp = experiment.Experiment( + inst=chain, + sampler=sampler, + config=experiment.ExperimentConfig(coupler_shim_step=0.01, signed_energy_scale=0.5), + ) + + sc = experiment.SamplerCall(run_index=0) + sc.logical_bqms = [chain.make_bqm()] + sc.shimdata = { + "total_iterations": 0, + "relative_coupler_strength": np.ones((1, chain.num_edges)), + } + # mean(frust) = 0.5, so delta = 0.01 * [-0.2, 0.0, 0.2] = [-0.002, 0.0, 0.002] + # Post-update mean(|rcs|) = 1.0 exactly -> renormalization is a no-op. + # Q = rcs * J(=1) * scale(=0.5) stays in [0.499, 0.501] -> no truncation. + results = {"CouplerFrustration": np.array([[0.3, 0.5, 0.7]])} + + exp._update_coupler_shim(sc, results) + + np.testing.assert_array_almost_equal( + sc.shimdata["relative_coupler_strength"], + np.array([[0.998, 1.0, 1.002]]), + ) + + def test_coupler_shim_singleton_orbits_is_noop(self): + """With singleton orbits, mean equals the single value, so rcs is unchanged.""" + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + orbit_type="singleton", + ) + chain.embedding_list = np.array([[0, 1, 2, 3]]) + sampler = _make_mock_sampler() + exp = experiment.Experiment( + inst=chain, + sampler=sampler, + config=experiment.ExperimentConfig(coupler_shim_step=0.01), + ) + + sc = experiment.SamplerCall(run_index=0) + sc.logical_bqms = [chain.make_bqm()] + rcs_before = np.ones((1, chain.num_edges)) + sc.shimdata = { + "total_iterations": 0, + "relative_coupler_strength": rcs_before.copy(), + } + results = {"CouplerFrustration": np.array([[0.3, 0.5, 0.7]])} + + exp._update_coupler_shim(sc, results) + + np.testing.assert_array_equal(sc.shimdata["relative_coupler_strength"], rcs_before) + + +class TestSaveLoadResults(unittest.TestCase): + def test_save_and_reload(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.data_path = Path(tmpdir) + exp.run_index = 0 + data = {"QubitMagnetization": np.zeros(4)} + exp._save_results(data) + fn = Path(tmpdir) / "iter00000.pkl.lzma" + self.assertTrue(fn.exists()) + + with lzma.open(fn, "rb") as f: + loaded = pickle.load(f) + np.testing.assert_array_equal(loaded["QubitMagnetization"], np.zeros(4)) + + def test_save_with_filename_and_run_index_raises(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.data_path = Path(tmpdir) + with self.assertRaises(ValueError): + exp._save_results({}, run_index=0, filename="test.pkl.lzma") + + def test_save_custom_filename(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.data_path = Path(tmpdir) + data = {"x": 1} + exp._save_results(data, filename="custom.pkl.lzma") + self.assertTrue((Path(tmpdir) / "custom.pkl.lzma").exists()) + + def test_load_results_ignore_shim(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + fn = exp.data_path / "iter00000.pkl.lzma" + fn.parent.mkdir(parents=True, exist_ok=True) + with lzma.open(fn, "wb") as f: + pickle.dump({"value": 0, "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 = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + for i in range(10): + fn = exp.data_path / f"iter{i:05d}.pkl.lzma" + fn.parent.mkdir(parents=True, exist_ok=True) + with lzma.open(fn, "wb") as f: + pickle.dump({"value": i, "shimdata": {}}, f) + + results = exp.load_results(num_iterations=3, start_iteration=2) + self.assertEqual(len(results), 3) + + def test_load_results_corrupted_lzma(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + exp.apply_param({"signed_energy_scale": 1.0, "anneal_time": 1.0}) + fn = exp.data_path / "iter00000.pkl.lzma" + fn.parent.mkdir(parents=True, exist_ok=True) + fn.write_bytes(b"corrupted data") + + with self.assertRaises(lzma.LZMAError): + exp.load_results(num_iterations=1) + + def test_generate_data_type_conversions(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment(inst=chain, sampler=sampler) + sc = experiment.SamplerCall(run_index=0) + sc.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) + + +class TestMakeBqm(unittest.TestCase): + def test_make_bqm_no_embeddings(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + sampler = _make_mock_sampler() + exp = experiment.Experiment( + inst=chain, + sampler=sampler, + config=experiment.ExperimentConfig(signed_energy_scale=0.5), + ) + sc = experiment.SamplerCall(run_index=0) + sc.logical_bqms = [chain.make_bqm()] + sc.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): + """Physical biases = logical_bias * relative_coupler_strength * signed_energy_scale.""" + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + chain.embedding_list = np.array([[10, 11, 12, 13]]) # offset to detect mapping + sampler = _make_mock_sampler() + exp = experiment.Experiment( + inst=chain, + sampler=sampler, + config=experiment.ExperimentConfig(signed_energy_scale=0.5), + ) + sc = experiment.SamplerCall(run_index=0) + sc.logical_bqms = [chain.make_bqm()] + sc.shimdata = { + "total_iterations": 0, + "relative_coupler_strength": np.full((1, chain.num_edges), 2.0), + } + + bqm = exp._make_bqm(sc) + + # Physical variables come from the embedding, not the logical indices. + self.assertEqual(set(bqm.variables), {10, 11, 12, 13}) + for u_log, v_log in chain.edge_list: + self.assertAlmostEqual(bqm.quadratic[(10 + u_log, 10 + v_log)], 1.0) + + self.assertEqual(len(bqm.quadratic), chain.num_edges) + + +class TestRunIteration(unittest.TestCase): + def test_run_iteration_basic(self): + """run_iteration() exercises the full pipeline: build call, sample, parse, shim, save.""" + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + exp = experiment.Experiment( + inst=chain, sampler=_make_mock_sampler(sync_response=True), max_iterations=1 + ) + chain._load_embeddings = mock.MagicMock() + finished = exp.run_iteration([{"signed_energy_scale": 1.0, "anneal_time": 1.0}]) + + self.assertFalse(finished) + result_files = list(exp.data_path.glob("iter*.pkl.lzma")) + self.assertEqual(len(result_files), 1) + + with lzma.open(result_files[0], "rb") as f: + data = pickle.load(f) + + self.assertIn("QubitMagnetization", data) + self.assertIn("CouplerCorrelation", data) + self.assertIn("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 = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + exp = experiment.Experiment( + inst=chain, + sampler=_make_mock_sampler(sync_response=True), + config=experiment.ExperimentConfig(), + max_iterations=0, + ) + chain._load_embeddings = mock.MagicMock() + finished = exp.run_iteration([{"signed_energy_scale": 1.0, "anneal_time": 1.0}]) + + self.assertTrue(finished) + self.assertEqual(list(exp.data_path.glob("iter*.pkl.lzma")), []) + + +class TestFastAnnealExperiment(unittest.TestCase): + def test_default_params(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = _make_mock_sampler() + config = experiment.FastAnnealExperimentConfig() + exp = experiment.Experiment(inst=chain, sampler=sampler, config=config) + self.assertTrue(exp.param.get("fast_anneal")) + self.assertEqual(exp.param["num_reads"], 100) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_lattice_utils/test_lattice.py b/tests/test_lattice_utils/test_lattice.py new file mode 100644 index 0000000..c85c546 --- /dev/null +++ b/tests/test_lattice_utils/test_lattice.py @@ -0,0 +1,464 @@ +# 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 unittest +from unittest import mock + +import dimod +import numpy as np +from dwave.samplers import SteepestDescentSolver + +from dwave.experimental.lattice_utils import lattice +from tests.test_lattice_utils._helpers import _make_embedded_chain, _make_triangular + + +class TestChain(unittest.TestCase): + def test_periodic(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + self.assertEqual(chain.num_spins, 6) + self.assertEqual(chain.num_edges, 6) + self.assertIn((5, 0), chain.edge_list) + + def test_non_periodic(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(6,), periodic=(False,), data_root=tmpdir) + self.assertEqual(chain.num_spins, 6) + self.assertEqual(chain.num_edges, 5) + self.assertNotIn((5, 0), chain.edge_list) + + def test_single_node_periodic(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(1,), periodic=(True,), data_root=tmpdir) + self.assertEqual(chain.num_spins, 1) + self.assertEqual(chain.num_edges, 0) + + def test_two_node_periodic(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(2,), periodic=(True,), data_root=tmpdir) + self.assertEqual(chain.num_edges, 2) + + def test_geometry_name(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + self.assertEqual(chain.geometry_name, "Chain") + + def test_default_periodic(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), data_root=tmpdir) + self.assertTrue(chain.periodic[0]) + + def test_edge_list_sorted(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(5,), periodic=(False,), data_root=tmpdir) + for u, v in chain.edge_list: + self.assertLess(u, v) + + def test_bqm_structure(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + self.assertEqual(len(bqm.variables), 4) + self.assertEqual(len(bqm.quadratic), 3) + for u, v in chain.edge_list: + self.assertAlmostEqual(bqm.quadratic[(u, v)], 1.0) + + def test_bqm_vartype(self): + with tempfile.TemporaryDirectory() as tmpdir: + bqm = lattice.Chain(dimensions=(3,), data_root=tmpdir).make_bqm() + self.assertEqual(bqm.vartype, dimod.SPIN) + + +class TestLattice(unittest.TestCase): + def test_orbit_singleton(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), orbit_type="singleton", data_root=tmpdir) + np.testing.assert_array_equal(chain.qubit_orbits, np.arange(4)) + np.testing.assert_array_equal(chain.coupler_orbits, np.arange(chain.num_edges)) + + def test_orbit_global(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), orbit_type="global", data_root=tmpdir) + np.testing.assert_array_equal(chain.qubit_orbits, np.zeros(4, dtype=int)) + np.testing.assert_array_equal( + chain.coupler_orbits, np.zeros(chain.num_edges, dtype=int) + ) + + def test_orbit_explicit(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(True,), + orbit_type="explicit", + qubit_orbits=np.array([0, 0, 1, 1]), + coupler_orbits=np.array([0, 0, 1, 1]), + data_root=tmpdir, + ) + np.testing.assert_array_equal(chain.qubit_orbits, [0, 0, 1, 1]) + + def test_unknown_orbit_type(self): + with tempfile.TemporaryDirectory() as tmpdir: + with self.assertRaises(ValueError): + lattice.Chain( + dimensions=(4,), periodic=(True,), orbit_type="bogus", data_root=tmpdir + ) + + def test_get_path_invalid_kind(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + with self.assertRaises(ValueError): + chain._get_path(None, "invalid") + + def test_standard_orbit_save_and_load(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(True,), + orbit_type="standard", + data_root=tmpdir, + ) + self.assertIsNotNone(chain.qubit_orbits) + self.assertIsNotNone(chain.coupler_orbits) + # Second instantiation should load from disk + chain2 = lattice.Chain( + dimensions=(4,), + periodic=(True,), + orbit_type="standard", + data_root=tmpdir, + ) + np.testing.assert_array_equal(chain.qubit_orbits, chain2.qubit_orbits) + + def test_embed_no_embeddings_found(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = mock.MagicMock() + type(sampler).__name__ = "MockDWaveSampler" + sampler.to_networkx_graph.return_value = chain.make_networkx_graph() + + with mock.patch( + "dwave.experimental.lattice_utils.lattice.lattice.find_multiple_embeddings", + return_value=[], + ): + with self.assertRaises(ValueError): + chain.embed_lattice(sampler, try_to_load=False, timeout=1) + + def test_embed_load_existing(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = mock.MagicMock() + type(sampler).__name__ = "MockDWaveSampler" + sampler.to_networkx_graph.return_value = chain.make_networkx_graph() + + embeddings = np.array([[0, 1, 2, 3]]) + chain._save_embeddings(sampler, embeddings) + + chain.embed_lattice(sampler, try_to_load=True, data_root=tmpdir) + np.testing.assert_array_equal(chain.embedding_list, embeddings) + + def test_embed_find_and_save(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) + sampler = mock.MagicMock() + type(sampler).__name__ = "MockDWaveSampler" + sampler.to_networkx_graph.return_value = chain.make_networkx_graph() + + emb_dict = {i: i for i in range(4)} + with mock.patch( + "dwave.experimental.lattice_utils.lattice.lattice.find_multiple_embeddings", + return_value=[emb_dict], + ): + chain.embed_lattice(sampler, try_to_load=False, timeout=1, data_root=tmpdir) + # Verify embedding was found and saved + emb_path = chain._get_path("embedding", sampler_name="MockDWaveSampler") + self.assertTrue(emb_path.exists()) + + +class TestLatticeOptimize(unittest.TestCase): + """Tests for the ``Lattice.optimize`` instance method (vs. the free function + ``lattice.optimize`` covered by ``TestOptimizeFunction``).""" + + def test_plain_lattice_private_optimize_default_sampler(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + reference_energy_sampler_kwargs={"num_sweeps": 256, "num_reads": 16}, + ) + bqm = chain.make_bqm() + energy, sample, method = chain.optimize(bqm) + + self.assertEqual(energy, -3) + self.assertEqual(bqm.energy(sample), energy) + self.assertEqual(method, "ExponentialBackoffSimulatedAnnealingSampler") + + def test_plain_lattice_private_optimize_custom_sampler(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + reference_energy_sampler=dimod.ExactSolver(), + ) + bqm = chain.make_bqm() + energy, sample, method = chain.optimize(bqm) + + self.assertEqual(energy, -3) + self.assertEqual(bqm.energy(sample), energy) + self.assertEqual(method, "ExactSolver") + + def test_optimize_with_custom_sampler_steepest_descent(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + reference_energy_sampler=SteepestDescentSolver(), + reference_energy_sampler_kwargs={"initial_states": np.array([[1, -1, 1, -1]])}, + ) + bqm = chain.make_bqm() + energy, sample, method = chain.optimize(bqm) + + self.assertEqual(energy, -3) + self.assertAlmostEqual(bqm.energy(sample), energy) + self.assertEqual(method, "SteepestDescentSolver") + + def test_optimize_with_custom_exponential_backoff_params(self): + with tempfile.TemporaryDirectory() as tmpdir: + sampler = lattice.ExponentialBackoffSimulatedAnnealingSampler( + max_num_sweeps=512, min_num_sweeps=64 + ) + chain = lattice.Chain( + dimensions=(4,), + periodic=(False,), + data_root=tmpdir, + reference_energy_sampler=sampler, + reference_energy_sampler_kwargs={"num_reads": 16}, + ) + bqm = chain.make_bqm() + energy, sample, method = chain.optimize(bqm) + + self.assertEqual(energy, -3) + self.assertAlmostEqual(bqm.energy(sample), energy) + self.assertEqual(method, "ExponentialBackoffSimulatedAnnealingSampler") + self.assertEqual(sampler.max_num_sweeps, 512) + self.assertEqual(sampler.min_num_sweeps, 64) + + +class TestTriangular(unittest.TestCase): + def test_basic_construction(self): + with tempfile.TemporaryDirectory() as tmpdir: + tri = lattice.Triangular( + dimensions=(3, 3), + periodic=(True, False), + data_root=tmpdir, + orbit_type="singleton", + halve_boundary_couplers=False, + ) + self.assertEqual(tri.num_spins, 9) + self.assertGreater(tri.num_edges, 0) + self.assertEqual(tri.geometry_name, "Triangular") + + def test_coordinates(self): + with tempfile.TemporaryDirectory() as tmpdir: + tri = _make_triangular(tmpdir, 3, 3) + y, x = tri.coordinates(0) + self.assertEqual(y, 0) + self.assertEqual(x, 0) + y, x = tri.coordinates(4) + self.assertEqual(y, 1) + self.assertEqual(x, 1) + + def test_halve_boundary_couplers(self): + with tempfile.TemporaryDirectory() as tmpdir: + tri = _make_triangular( + tmpdir, 3, 3, periodic=(False, False), halve_boundary_couplers=True + ) + bqm = tri.make_bqm() + graph = tri.make_networkx_graph() + for u, v in tri.edge_list: + expected = 1.0 if (graph.degree[u] == 6 or graph.degree[v] == 6) else 0.5 + self.assertAlmostEqual(bqm.quadratic[(u, v)], expected) + + def test_periodicity(self): + with tempfile.TemporaryDirectory() as tmpdir: + tri = _make_triangular(tmpdir, 3, 3, periodic=(False, True)) + self.assertFalse(tri.periodic[0]) + self.assertTrue(tri.periodic[1]) + + +class TestDimerizedTriangular(unittest.TestCase): + def test_basic_construction(self): + with tempfile.TemporaryDirectory() as tmpdir: + dt = lattice.DimerizedTriangular( + dimensions=(3, 3), periodic=(True, False), orbit_type="singleton", data_root=tmpdir + ) + self.assertEqual(dt.geometry_name, "DimerizedTriangular") + self.assertIsNotNone(dt.logical_lattice) + self.assertEqual(dt.num_spins, 18) + + def test_chain_connectivity_self(self): + with tempfile.TemporaryDirectory() as tmpdir: + dt = lattice.DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + cc = dt.get_chain_connectivity(0) + self.assertEqual(cc, ((0, 1),)) + + def test_chain_connectivity_cases(self): + with tempfile.TemporaryDirectory() as tmpdir: + dt = lattice.DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + cases = [ + ((0,), ((0, 1),)), + ((0, 1), ((1, 0),)), + ((0, 3), ((1, 0),)), + ] + for args, expected in cases: + with self.subTest(args=args): + self.assertEqual(dt.get_chain_connectivity(*args), expected) + + +class TestEmbeddedLattice(unittest.TestCase): + def test_embed_sample(self): + with tempfile.TemporaryDirectory() as tmpdir: + dt = lattice.DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + logical_sample = np.array([1, -1, 1, -1, 1, -1, 1, -1, 1]) + embedded = dt.embed_sample(logical_sample) + self.assertEqual(len(embedded), dt.num_spins) + # Each chain should have the same value + for spin, chain in dt.chain_nodes.items(): + for node in chain: + self.assertEqual(embedded[node], logical_sample[spin]) + + def test_unembed_sample(self): + chain_nodes = {0: (0, 1, 2), 1: (3, 4, 5)} + with tempfile.TemporaryDirectory() as tmpdir: + embedded = lattice.EmbeddedLattice( + logical_lattice=lattice.Chain( + dimensions=(len(chain_nodes),), + periodic=(False,), + data_root=tmpdir, + ), + chain_nodes=chain_nodes, + dimensions=(sum(len(chain) for chain in chain_nodes.values()),), + periodic=(False,), + ) + physical_sample = np.array([1, 1, -1, -1, -1, 1]) + logical = embedded.unembed_sample(physical_sample) + np.testing.assert_array_equal(logical, np.array([1, -1])) + + def test_unembed_sample_breaks_ties_randomly(self): + with tempfile.TemporaryDirectory() as tmpdir: + embedded = _make_embedded_chain({0: (0, 1), 1: (2, 3)}, tmpdir) + physical_sample = np.array([1, -1, 1, -1]) + with mock.patch( + "dwave.experimental.lattice_utils.lattice.embedded_lattice.np.random.rand", + side_effect=[0.9, 0.1], + ): + logical = embedded.unembed_sample(physical_sample) + np.testing.assert_array_equal(logical, np.array([1, -1])) + + def test_unembed_sampleset(self): + with tempfile.TemporaryDirectory() as tmpdir: + embedded = _make_embedded_chain({0: (0, 1), 1: (2, 3)}, tmpdir) + samples = np.array( + [ + [1, 1, -1, -1], + [1, -1, 1, -1], + ] + ) + ss = dimod.SampleSet.from_samples(samples, vartype=dimod.SPIN, energy=0) + with mock.patch( + "dwave.experimental.lattice_utils.lattice.embedded_lattice.np.random.rand", + return_value=np.array([[0.2, 0.9], [0.2, 0.1]]), + ): + result = embedded.unembed_sampleset(ss) + np.testing.assert_array_equal(dimod.as_samples(result)[0], np.array([[1, -1], [1, -1]])) + + def test_connectivity_generic_self(self): + chain_nodes = {0: (10, 11, 12), 1: (20, 21, 22)} + with tempfile.TemporaryDirectory() as tmpdir: + el = _make_embedded_chain(chain_nodes, tmpdir) + cc = lattice.EmbeddedLattice.get_chain_connectivity(el, 0) + self.assertEqual(cc, ((0, 1), (0, 2), (1, 2))) + self.assertEqual( + {tuple(chain_nodes[0][index] for index in edge) for edge in cc}, + {(10, 11), (10, 12), (11, 12)}, + ) + + def test_nested_embedded_raises(self): + with tempfile.TemporaryDirectory() as tmpdir: + base_chain = lattice.Chain( + dimensions=(2,), + periodic=(False,), + data_root=tmpdir, + ) + embedded_once = lattice.EmbeddedLattice( + logical_lattice=base_chain, + chain_nodes={0: (0, 1), 1: (2, 3)}, + dimensions=(4,), + periodic=(False,), + ) + with self.assertRaises(NotImplementedError): + lattice.EmbeddedLattice( + logical_lattice=embedded_once, + chain_nodes={0: (0, 1), 1: (2, 3), 2: (4, 5), 3: (6, 7)}, + dimensions=(8,), + periodic=(False,), + ) + + +class TestOrbits(unittest.TestCase): + def test_reindex_basic(self): + mapping = {"a": 5, "b": 5, "c": 10} + result = lattice.reindex(mapping) + self.assertEqual(result, {'a': 0, 'b': 0, 'c': 1}) + + def test_signed_bqm_symmetry(self): + bqm = dimod.BQM(vartype="SPIN") + bqm.add_variable(0, 0.5) + bqm.add_variable(1, -0.3) + bqm.add_quadratic(0, 1, 1.0) + signed = lattice.make_signed_bqm(bqm) + self.assertAlmostEqual(signed.linear["p0"], 0.5) + self.assertAlmostEqual(signed.linear["m0"], -0.5) + + +class TestOptimizeFunction(unittest.TestCase): + """Tests for the free function ``lattice.optimize`` (vs. the ``Lattice.optimize`` + instance method covered by ``TestLatticeOptimize``).""" + + def test_plain_lattice(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + energy, sample, _ = lattice.optimize( + chain, bqm, sampler_kwargs={"num_sweeps": 256, "num_reads": 16} + ) + self.assertEqual(energy, -3) + self.assertEqual(bqm.energy(sample), energy) + + def test_embedded_lattice(self): + with tempfile.TemporaryDirectory() as tmpdir: + dt = lattice.DimerizedTriangular(dimensions=(3, 3), data_root=tmpdir) + bqm = dt.make_bqm() + energy, sample, _ = lattice.optimize( + dt, bqm, sampler_kwargs={"num_sweeps": 256, "num_reads": 16} + ) + self.assertEqual(len(sample), dt.num_spins) + self.assertEqual(bqm.energy(sample), energy) + self.assertTrue(set(sample).issubset({-1, 1})) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_lattice_utils/test_observable.py b/tests/test_lattice_utils/test_observable.py new file mode 100644 index 0000000..330e830 --- /dev/null +++ b/tests/test_lattice_utils/test_observable.py @@ -0,0 +1,188 @@ +# 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 unittest +from pathlib import Path + +import dimod +import numpy as np + +from dwave.experimental.lattice_utils import lattice, observable +from tests.test_lattice_utils._helpers import _make_mock_experiment, _make_triangular + + +class TestQubitMagnetization(unittest.TestCase): + def test_qubit_magnetization(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + samples = np.array([[1, 1, -1, -1], [-1, -1, 1, 1]]) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(chain) + result = observable.QubitMagnetization().evaluate(exp, bqm, ss) + np.testing.assert_array_equal(result, [0.0, 0.0, 0.0, 0.0]) + + +class TestCouplerCorrelation(unittest.TestCase): + def test_coupler_correlation(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + exp = _make_mock_experiment(chain) + alt = np.tile([1, -1, 1, -1], (4, 1)) + ss_alt = dimod.SampleSet.from_samples_bqm(alt, bqm) + np.testing.assert_array_equal( + observable.CouplerCorrelation().evaluate(exp, bqm, ss_alt), + -np.ones(chain.num_edges), + ) + + +class TestCouplerFrustration(unittest.TestCase): + def test_coupler_frustration(self): + """All aligned (corr=1) -> frustration = 1.0""" + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + samples = np.ones((4, 4)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(chain) + np.testing.assert_array_almost_equal( + observable.CouplerFrustration().evaluate(exp, bqm, ss), np.ones(chain.num_edges) + ) + + +class TestSampleEnergy(unittest.TestCase): + def test_sample_energy(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + # All-ones: energy = sum of J for 3 edges = 3 + samples = np.ones((1, 4)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp_pos = _make_mock_experiment(chain, signed_energy_scale=1.0) + np.testing.assert_array_almost_equal( + observable.SampleEnergy().evaluate(exp_pos, bqm, ss), [3] + ) + + +class TestBitpackedSpins(unittest.TestCase): + def test_bitpacked_spins(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + samples = np.array([[1, -1, 1, -1], [-1, 1, -1, 1]]) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(chain) + packed, shape = observable.BitpackedSpins().evaluate(exp, bqm, ss) + self.assertEqual(shape, (2, 4)) + # Unpack and verify round-trip + unpacked = np.unpackbits(packed)[: shape[0] * shape[1]].reshape(shape) + np.testing.assert_array_equal(unpacked, np.equal(samples, 1)) + + +class TestReferenceEnergy(unittest.TestCase): + def test_reference_energy_save_load_roundtrip(self): + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "ref.txt" + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + sample = np.array([1, -1, 1, -1]) + obs = observable.ReferenceEnergy() + obs.save(path, -3, sample, "SA") + + exp = _make_mock_experiment(chain) + energy, loaded_sample, method = obs.load(exp, bqm, path) + self.assertEqual(energy, -3) + self.assertEqual(method, "SA") + np.testing.assert_array_equal(loaded_sample, sample) + + def test_reference_energy_evaluate_generates_and_caches(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + obs = observable.ReferenceEnergy() + + path1 = Path(tmpdir) / "ref_inst.txt" + energy1 = obs.evaluate(None, bqm, None, path=path1, inst=chain) + self.assertTrue(path1.exists()) + # Second call loads from cache — same value + energy1b = obs.evaluate(None, bqm, None, path=path1) + self.assertAlmostEqual(energy1, energy1b) + + exp = _make_mock_experiment(chain, run_index=0, num_random_instances=1) + path2 = Path(tmpdir) / "ref_exp.txt" + energy2 = obs.evaluate(exp, bqm, None, path=path2) + self.assertTrue(path2.exists()) + self.assertAlmostEqual(energy1, energy2) + + def test_reference_energy_update(self): + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "ref.txt" + chain = lattice.Chain(dimensions=(4,), periodic=(False,), data_root=tmpdir) + bqm = chain.make_bqm() + obs = observable.ReferenceEnergy() + exp = _make_mock_experiment(chain) + + bad_sample = np.ones(4) + obs.save(path, bqm.energy(bad_sample), bad_sample, "SA") + + better = np.array([1, -1, 1, -1]) + obs.update(exp, bqm, better, path=path) + energy, _, _ = obs.load(exp, bqm, path) + self.assertAlmostEqual(energy, bqm.energy(better)) + + # Attempting to update with a worse sample raises ValueError + with self.assertRaises(ValueError): + obs.update(exp, bqm, bad_sample, path=path) + + +class TestKinks(unittest.TestCase): + def test_all_aligned(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + bqm = chain.make_bqm() + samples = np.ones((10, 6)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(chain) + result = observable.KinkKinkCorrelator().evaluate(exp, bqm, ss) + np.testing.assert_array_equal(result, np.zeros(6)) + + def test_mixed_pattern(self): + with tempfile.TemporaryDirectory() as tmpdir: + chain = lattice.Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) + bqm = chain.make_bqm() + # [1,1,-1,-1,1,1]: kink at sites 2,4 (domain walls) + samples = np.tile([1, 1, -1, -1, 1, 1], (20, 1)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(chain) + result = observable.KinkKinkCorrelator().evaluate(exp, bqm, ss) + expected = np.array([0.0, -0.25, 0.125, -0.25, 0.125, -0.25]) + np.testing.assert_array_almost_equal(result, expected) + + +class TestTriangularOP(unittest.TestCase): + def test_uniform_state_vanishes(self): + with tempfile.TemporaryDirectory() as tmpdir: + tri = _make_triangular(tmpdir, 3, 3, periodic=(True, False)) + bqm = tri.make_bqm() + samples = np.ones((5, 9)) + ss = dimod.SampleSet.from_samples_bqm(samples, bqm) + exp = _make_mock_experiment(tri) + result = observable.TriangularOP().evaluate(exp, bqm, ss) + np.testing.assert_array_almost_equal(np.abs(result), np.zeros(5), decimal=10) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_lattice_utils/test_utils.py b/tests/test_lattice_utils/test_utils.py new file mode 100644 index 0000000..5458f3e --- /dev/null +++ b/tests/test_lattice_utils/test_utils.py @@ -0,0 +1,61 @@ +# 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 unittest + +import numpy as np + +from dwave.experimental.lattice_utils.utils import ( + bootstrap, + confidence_interval, + generate_bootstrap_indices, +) + + +class TestUtils(unittest.TestCase): + def test_bootstrap_all_nan_skipnan(self): + rng = np.random.default_rng(seed=0) + result = bootstrap(np.array([np.nan, np.nan]), rng, repetitions=5, skipnan=True) + self.assertEqual(len(result), 5) + for val in result: + self.assertTrue(np.isnan(val)) + + def test_bootstrap_skipnan_false(self): + rng = np.random.default_rng(seed=0) + result = bootstrap(np.array([1.0, 2.0, np.nan]), rng, repetitions=5, skipnan=False) + self.assertEqual(len(result), 5) + + def test_bootstrap_custom_function(self): + rng = np.random.default_rng(seed=0) + result = bootstrap(np.arange(20), rng, repetitions=10, bootstrap_function=np.mean) + self.assertEqual(len(result), 10) + + def test_generate_bootstrap_indices_correct_count(self): + rng = np.random.default_rng(seed=0) + indices = list(generate_bootstrap_indices(10, 5, rng)) + self.assertEqual(len(indices), 5) + for idx in indices: + self.assertEqual(len(idx), 10) + self.assertTrue(np.all(idx >= 0)) + self.assertTrue(np.all(idx < 10)) + + def test_confidence_interval_width(self): + arr = np.arange(1000) + _, low1, high1 = confidence_interval(arr, width=0.5) + _, low2, high2 = confidence_interval(arr, width=0.99) + self.assertGreater(low2 + high2, low1 + high1) + + +if __name__ == "__main__": + unittest.main() From 2b02ceee8a3ed2c1081d539cd4cf5a938c101b33 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 20 Jul 2026 12:09:50 -0700 Subject: [PATCH 12/13] Make TQDM optional -make TQDM an optional import -order imports -add some missing docstrings -removed subscript unpacking for Python 3.10 compatibility -updated licenses to 2026 --- dwave/experimental/lattice_utils/__init__.py | 2 +- .../lattice_utils/experiment/__init__.py | 4 +- .../lattice_utils/experiment/experiment.py | 134 +++++++++--------- .../experiment/fast_anneal_experiment.py | 2 +- .../lattice_utils/experiment/samplercall.py | 2 +- .../lattice_utils/lattice/__init__.py | 8 +- .../lattice_utils/lattice/chain.py | 8 +- .../lattice_utils/lattice/embedded_lattice.py | 7 +- .../lattice_utils/lattice/lattice.py | 35 ++--- .../lattice_utils/lattice/optimize.py | 7 +- .../lattice_utils/lattice/orbits.py | 4 +- .../lattice_utils/lattice/triangular.py | 6 +- .../lattice_utils/observable/__init__.py | 4 +- .../lattice_utils/observable/kinks.py | 5 +- .../lattice_utils/observable/observable.py | 65 ++++++--- .../lattice_utils/observable/triangular.py | 7 +- dwave/experimental/lattice_utils/utils.py | 2 +- .../1D_Ising_chain_shim.py} | 10 +- pyproject.toml | 3 + .../add-lattice-utils-49e5efd79268e69d.yaml | 16 ++- tests/__init__.py | 2 +- tests/test_lattice_utils/__init__.py | 13 ++ tests/test_lattice_utils/_helpers.py | 2 +- tests/test_lattice_utils/test_experiment.py | 2 +- tests/test_lattice_utils/test_lattice.py | 7 +- tests/test_lattice_utils/test_observable.py | 2 +- tests/test_lattice_utils/test_utils.py | 2 +- 27 files changed, 211 insertions(+), 150 deletions(-) rename examples/{example_1D_Ising_chain_shim.py => lattice_utils/1D_Ising_chain_shim.py} (96%) diff --git a/dwave/experimental/lattice_utils/__init__.py b/dwave/experimental/lattice_utils/__init__.py index fb1e903..bd91f26 100644 --- a/dwave/experimental/lattice_utils/__init__.py +++ b/dwave/experimental/lattice_utils/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/dwave/experimental/lattice_utils/experiment/__init__.py b/dwave/experimental/lattice_utils/experiment/__init__.py index 5f75b63..667811a 100644 --- a/dwave/experimental/lattice_utils/experiment/__init__.py +++ b/dwave/experimental/lattice_utils/experiment/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,5 +13,5 @@ # 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 * +from dwave.experimental.lattice_utils.experiment.samplercall import * diff --git a/dwave/experimental/lattice_utils/experiment/experiment.py b/dwave/experimental/lattice_utils/experiment/experiment.py index ef4b1c2..947c226 100644 --- a/dwave/experimental/lattice_utils/experiment/experiment.py +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,30 +12,37 @@ # See the License for the specific language governing permissions and # limitations under the License. -import tempfile +from __future__ import annotations + +import copy import lzma import os import pickle +import tempfile import time -from pathlib import Path +from dataclasses import dataclass from datetime import datetime +from pathlib import Path from typing import Any -from dataclasses import dataclass import dimod import numpy as np -from tqdm.auto import tqdm -from dwave.experimental.lattice_utils.lattice import Lattice +try: + from tqdm.auto import tqdm +except ImportError: + tqdm = None + +from dwave.experimental.lattice_utils.experiment.samplercall import SamplerCall +from dwave.experimental.lattice_utils.lattice.lattice import Lattice from dwave.experimental.lattice_utils.observable import ( - QubitMagnetization, + BitpackedSpins, CouplerCorrelation, CouplerFrustration, - SampleEnergy, - BitpackedSpins, + QubitMagnetization, ReferenceEnergy, + SampleEnergy, ) -from dwave.experimental.lattice_utils.experiment.samplercall import SamplerCall __all__ = ['Experiment', 'ExperimentConfig'] @@ -89,28 +96,27 @@ def __init__( self.sampler = sampler self.reference_energy_sampler = reference_energy_sampler self.reference_energy_sampler_kwargs = reference_energy_sampler_kwargs - self.param = dict(vars(config)) + self.param = vars(config).copy() self.experiment_results_root = inst.data_root / "results" self.data_path = None self.run_index = 0 self.config = config self.max_iterations = max_iterations self.already_initialized: bool = False - self.observables_to_collect = { + 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. @@ -122,7 +128,6 @@ def load_results( 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. @@ -144,12 +149,10 @@ def load_results( 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") + if ignore_shim: + result_fields.remove("shimdata") results.append({k: data[k] for k in result_fields}) @@ -172,8 +175,7 @@ def apply_param(self, param: dict[str, float]) -> None: elif "anneal_time" in param: self.param.pop("anneal_schedule", None) - for param_name, param_val in param.items(): - self.param[param_name] = param_val + self.param.update(param) self.data_path = self.experiment_results_root / self._get_relative_data_path() self.already_initialized = self._prepare_run_index() @@ -199,6 +201,9 @@ def run_iteration( A boolean value corresponding to whether or not the experiment is finished. """ + if progress and tqdm is None: + raise ImportError("Progress reporting requires the optional 'tqdm' dependency.") + try: self.inst._load_embeddings(self.sampler) except FileNotFoundError as e: @@ -216,11 +221,14 @@ def run_iteration( response_dict = {} call_dict = {} - create_bar = self._make_progress_bar( - total=len(parameter_list), - desc="Creating sampler calls", - colour=DW_BLUE, - enabled=progress, + create_bar = ( + self._make_progress_bar( + total=len(parameter_list), + desc="Creating sampler calls", + colour=DW_BLUE, + ) + if progress + else None ) for index, param in enumerate(parameter_list): @@ -239,18 +247,21 @@ def run_iteration( if create_bar is not None: create_bar.close() - if len(call_dict) == 0: + if not call_dict: if progress: tqdm.write( f"***\n***\nFINISHED for all {len(parameter_list)} parameterizations.\n***\n***" ) return True - wait_bar = self._make_progress_bar( - total=len(call_dict), - desc=" Awaiting/parsing data", - colour=DW_TEAL, - enabled=progress, + wait_bar = ( + self._make_progress_bar( + total=len(call_dict), + desc=" Awaiting/parsing data", + colour=DW_TEAL, + ) + if progress + else None ) # Get and manage all the results @@ -264,7 +275,7 @@ def run_iteration( 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) + self._save_results(savedata) if wait_bar is not None: wait_bar.update() del response_dict[index] @@ -277,7 +288,9 @@ def run_iteration( if wait_bar is not None: wait_bar.close() - self._print_iteration_status(call_dict, len(parameter_list), enabled=progress) + if progress: + self._print_iteration_status(call_dict, len(parameter_list)) + return False def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[str, Any]: @@ -305,14 +318,14 @@ def parse_results(self, call: SamplerCall, response: dimod.SampleSet) -> dict[st ) results = {} - for observable in set(self.observables_to_collect): + for observable in self.observables_to_collect: results[observable.name] = [] for iemb, sample_array in enumerate(sample_arrays): bqm = call.logical_bqms[iemb] obs_result = observable.evaluate(self, bqm, sample_set[iemb]) results[observable.name].append(obs_result) - if type(results[observable.name][0]) == np.ndarray: + if isinstance(results[observable.name][0], np.ndarray): results[observable.name] = np.asarray(results[observable.name]) return results @@ -323,14 +336,10 @@ def _make_progress_bar( total: int, desc: str, colour: str, - enabled: bool, bar_format: str | None = None, initial: int | float = 0, - ) -> tqdm | None: + ) -> tqdm: """Create a tqdm progress bar with consistent formatting.""" - if not enabled: - return None - if bar_format is None: bar_width = min(100, max(total, 20)) bar_format = f"{{desc}}: |{{bar:{bar_width}}}{{r_bar}}{{bar:-{bar_width}b}}" @@ -347,11 +356,8 @@ def _print_iteration_status( self, call_dict: dict[int, SamplerCall], num_params: int, - enabled: bool, ) -> None: """Print a summary of the iteration status, including progress and iteration ranges.""" - if not enabled: - return iteration_range = ( f"Iteration range " f"{min(call.shimdata['total_iterations'] for call in call_dict.values())}-" @@ -384,7 +390,6 @@ def _print_iteration_status( desc=" Total progress", bar_format=bar_format, colour=DW_ORANGE, - enabled=enabled, initial=progress_value, ) total_bar.close() @@ -393,7 +398,6 @@ def _save_results( self, data_dict: dict[str, Any], run_index: int | None = None, - quiet: bool = True, filename: str | None = None, ) -> None: """Save results to disk using LZMA-compressed pickle.""" @@ -403,7 +407,7 @@ def _save_results( filename = f"iter{run_index:05d}.pkl.lzma" else: if run_index is not None: - raise ValueError + raise ValueError("Cannot specify both filename and run_index.") # Write to a temp directory first to reduce disk write errors from killed jobs. with tempfile.TemporaryDirectory(dir=self.data_path) as tmp: @@ -412,14 +416,9 @@ def _save_results( 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] + return [str(fn) for fn in sorted(self.data_path.glob("iter*.pkl.lzma"))] def _get_next_run_index(self) -> tuple[int, bool]: """Get the next run index based on the existing files in the data path.""" @@ -479,7 +478,9 @@ def _get_parameter_pathstring(self) -> str: f'energyscale{signed_energy_scale:0.3}/asched{self.param["anneal_schedule"]}' ) else: - raise ValueError + raise ValueError( + "Parameter list must contain either 'anneal_time' or 'anneal_schedule'." + ) # Strip spaces and replace other unswanted symbols with underscores. pathstring = pathstring.replace(" ", "_") @@ -537,16 +538,14 @@ def _format_parameter_list( Some parameters can cause bugs if they are not appropriately formatted, rounded, etc. in accordance with filenames. """ - ret = parameter_list.copy() + ret_unique = [] + ret = copy.deepcopy(parameter_list) for entry in ret: if "anneal_time" in entry: entry["anneal_time"] = np.round(entry["anneal_time"], 6) if "anneal_schedule" in entry: entry["anneal_schedule"] = [tuple(np.round(p, 6)) for p in entry["anneal_schedule"]] - # 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) @@ -560,7 +559,7 @@ def _generate_data_to_save( """Construct a single dictionary containing results and shim data for saving.""" savedata = {} for key in results: - if type(results[key]) == np.ndarray: + if isinstance(results[key], np.ndarray): if results[key].dtype == "complex128": savedata[key] = results[key].astype(np.complex64) elif results[key].dtype == "float64": @@ -572,9 +571,9 @@ def _generate_data_to_save( savedata["shimdata"] = {} for key in sampler_call.shimdata: - if type(sampler_call.shimdata[key]) == np.ndarray: + if isinstance(sampler_call.shimdata[key], np.ndarray): savedata["shimdata"][key] = sampler_call.shimdata[key].astype(np.float32) - elif type(sampler_call.shimdata[key]) == int: + elif isinstance(sampler_call.shimdata[key], int): savedata["shimdata"][key] = sampler_call.shimdata[key] else: savedata["shimdata"][key] = sampler_call.shimdata[key].copy() @@ -601,7 +600,10 @@ def _make_sampler_params(self, **kwargs) -> dict[str, Any]: if self.param.get("fast_anneal", False): ret["fast_anneal"] = True - ret["annealing_time"] = self.param["anneal_time"] + if "anneal_schedule" in self.param: + ret["anneal_schedule"] = self.param["anneal_schedule"] + elif "anneal_time" in self.param: + ret["annealing_time"] = self.param["anneal_time"] return ret @@ -629,7 +631,7 @@ 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): + def _load_shim(self) -> dict[str, Any]: """Load shim data from the most recently completed iteration.""" filename = self._get_latest_iteration_filename() @@ -647,7 +649,7 @@ def _load_shim(self): except Exception as e: raise OSError("Failed to open file") from e - def _update_shim(self, sampler_call: SamplerCall, results: dict[str, Any]): + def _update_shim(self, sampler_call: SamplerCall, results: dict[str, Any]) -> None: """Update shim parameters according to shim data and parameters.""" if "flux_biases" in sampler_call.shimdata and self.param.get("flux_bias_shim_step", 0) != 0: self._update_flux_bias_shim(sampler_call, results) @@ -659,7 +661,7 @@ def _update_shim(self, sampler_call: SamplerCall, results: dict[str, Any]): sampler_call.shimdata["total_iterations"] += 1 - def _update_flux_bias_shim(self, sampler_call: SamplerCall, results: dict[str, Any]): + def _update_flux_bias_shim(self, sampler_call: SamplerCall, results: dict[str, Any]) -> None: """Update flux-bias shim values based on qubit magnetization.""" target_magnetization = self.param["target_magnetization"] qubit_magnetization = results["QubitMagnetization"] @@ -822,7 +824,7 @@ def _make_bqm(self, sampler_call: SamplerCall) -> dimod.BQM: for iedge, edge in enumerate(self.inst.edge_list): bias = ( - logical_bqm.quadratic[*edge] + logical_bqm.quadratic[tuple(edge)] * relative_coupler_strength[iemb, iedge] * signed_energy_scale ) diff --git a/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py index 59ca0b2..a34860d 100644 --- a/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py +++ b/dwave/experimental/lattice_utils/experiment/fast_anneal_experiment.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/dwave/experimental/lattice_utils/experiment/samplercall.py b/dwave/experimental/lattice_utils/experiment/samplercall.py index 81e2654..e02c9d5 100644 --- a/dwave/experimental/lattice_utils/experiment/samplercall.py +++ b/dwave/experimental/lattice_utils/experiment/samplercall.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/dwave/experimental/lattice_utils/lattice/__init__.py b/dwave/experimental/lattice_utils/lattice/__init__.py index 22a583d..1ff81e5 100644 --- a/dwave/experimental/lattice_utils/lattice/__init__.py +++ b/dwave/experimental/lattice_utils/lattice/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,9 +12,9 @@ # 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.embedded_lattice import * +from dwave.experimental.lattice_utils.lattice.lattice import * from dwave.experimental.lattice_utils.lattice.optimize import * -from dwave.experimental.lattice_utils.lattice.triangular import * from dwave.experimental.lattice_utils.lattice.orbits import * -from dwave.experimental.lattice_utils.lattice.embedded_lattice import * +from dwave.experimental.lattice_utils.lattice.triangular import * diff --git a/dwave/experimental/lattice_utils/lattice/chain.py b/dwave/experimental/lattice_utils/lattice/chain.py index c447d62..6636a78 100644 --- a/dwave/experimental/lattice_utils/lattice/chain.py +++ b/dwave/experimental/lattice_utils/lattice/chain.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -16,10 +16,10 @@ from pathlib import Path from typing import Any -from numpy.typing import NDArray import dimod +from numpy.typing import NDArray -from dwave.experimental.lattice_utils.lattice import Lattice +from dwave.experimental.lattice_utils.lattice.lattice import Lattice __all__ = ['Chain'] @@ -82,5 +82,5 @@ def generate_edges(self) -> Generator[tuple[int, int]]: for i in range(n - 1): yield (i, i + 1) - if self.periodic[0] and n > 1: + if self.periodic[0] and n > 2: yield (n - 1, 0) diff --git a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py index 9263432..02bf708 100644 --- a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,17 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. - +from collections.abc import Generator, Hashable from itertools import combinations, product from numbers import Integral -from collections.abc import Generator, Hashable from pathlib import Path import dimod import numpy as np from numpy.typing import NDArray -from dwave.experimental.lattice_utils.lattice import Lattice +from dwave.experimental.lattice_utils.lattice.lattice import Lattice __all__ = ['EmbeddedLattice'] diff --git a/dwave/experimental/lattice_utils/lattice/lattice.py b/dwave/experimental/lattice_utils/lattice/lattice.py index 3b714a1..2c670bc 100644 --- a/dwave/experimental/lattice_utils/lattice/lattice.py +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,20 +13,20 @@ # limitations under the License. import os -from pathlib import Path -from collections.abc import Generator, Hashable -from abc import ABC, abstractmethod import warnings +from abc import ABC, abstractmethod +from collections.abc import Generator, Hashable +from pathlib import Path from typing import Any import dimod -from minorminer.utils.parallel_embeddings import find_multiple_embeddings import networkx as nx import numpy as np +from minorminer.utils.parallel_embeddings import find_multiple_embeddings from numpy.typing import NDArray -from dwave.experimental.lattice_utils.lattice.orbits import get_orbits from dwave.experimental.lattice_utils.lattice.optimize import optimize +from dwave.experimental.lattice_utils.lattice.orbits import get_orbits __all__ = ['Lattice'] @@ -212,16 +212,19 @@ def initialize_orbits( 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)}." - ) + if qubit_orbits is None or coupler_orbits is None: + raise ValueError( + 'orbit_type "explicit" requires both qubit_orbits and coupler_orbits.' + ) + if len(qubit_orbits) != self.num_spins: + raise ValueError( + f"qubit_orbits must have length {self.num_spins}, got {len(qubit_orbits)}." + ) + if len(coupler_orbits) != self.num_edges: + raise ValueError( + f"coupler_orbits must have length {self.num_edges}, " + f"got {len(coupler_orbits)}." + ) self.qubit_orbits = qubit_orbits self.coupler_orbits = coupler_orbits else: diff --git a/dwave/experimental/lattice_utils/lattice/optimize.py b/dwave/experimental/lattice_utils/lattice/optimize.py index 624735a..c4a121e 100644 --- a/dwave/experimental/lattice_utils/lattice/optimize.py +++ b/dwave/experimental/lattice_utils/lattice/optimize.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,11 +13,12 @@ # limitations under the License. from __future__ import annotations + from typing import Any -from dwave.samplers import SimulatedAnnealingSampler -import numpy as np import dimod +import numpy as np +from dwave.samplers import SimulatedAnnealingSampler from numpy.typing import NDArray __all__ = ['optimize', 'ExponentialBackoffSimulatedAnnealingSampler'] diff --git a/dwave/experimental/lattice_utils/lattice/orbits.py b/dwave/experimental/lattice_utils/lattice/orbits.py index cdcc812..6fb9af9 100644 --- a/dwave/experimental/lattice_utils/lattice/orbits.py +++ b/dwave/experimental/lattice_utils/lattice/orbits.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from collections.abc import Hashable from collections import defaultdict +from collections.abc import Hashable import dimod import networkx as nx diff --git a/dwave/experimental/lattice_utils/lattice/triangular.py b/dwave/experimental/lattice_utils/lattice/triangular.py index 06eaf5d..0a6ac5c 100644 --- a/dwave/experimental/lattice_utils/lattice/triangular.py +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,13 +15,13 @@ from collections.abc import Generator, Hashable from pathlib import Path +import dimod import networkx as nx import numpy as np from numpy.typing import NDArray -import dimod -from dwave.experimental.lattice_utils.lattice.lattice import Lattice from dwave.experimental.lattice_utils.lattice.embedded_lattice import EmbeddedLattice +from dwave.experimental.lattice_utils.lattice.lattice import Lattice __all__ = ['Triangular', 'DimerizedTriangular'] diff --git a/dwave/experimental/lattice_utils/observable/__init__.py b/dwave/experimental/lattice_utils/observable/__init__.py index cb3ef4f..8f32897 100644 --- a/dwave/experimental/lattice_utils/observable/__init__.py +++ b/dwave/experimental/lattice_utils/observable/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,6 +12,6 @@ # 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.observable import * from dwave.experimental.lattice_utils.observable.triangular import * diff --git a/dwave/experimental/lattice_utils/observable/kinks.py b/dwave/experimental/lattice_utils/observable/kinks.py index 1fb7e44..665d5f3 100644 --- a/dwave/experimental/lattice_utils/observable/kinks.py +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,9 +13,10 @@ # limitations under the License. from __future__ import annotations + import dimod -from dimod import BQM, SampleSet import numpy as np +from dimod import BQM, SampleSet from numpy.typing import NDArray from dwave.experimental.lattice_utils.observable.observable import Observable diff --git a/dwave/experimental/lattice_utils/observable/observable.py b/dwave/experimental/lattice_utils/observable/observable.py index a267e95..4fb91a4 100644 --- a/dwave/experimental/lattice_utils/observable/observable.py +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,15 +13,16 @@ # limitations under the License. from __future__ import annotations + from abc import ABC, abstractmethod from pathlib import Path from typing import Any, TypeAlias +import dimod import numpy as np from numpy.typing import NDArray -import dimod -from dwave.experimental.lattice_utils.lattice import Lattice +from dwave.experimental.lattice_utils.lattice.lattice import Lattice __all__ = [ 'Observable', @@ -53,7 +54,13 @@ def evaluate( bqm: dimod.BQM, sample_set: dimod.SampleSet, ) -> ObservableResult: - pass + """Compute the observable from the provided sample set. + + Args: + experiment: Experiment object containing the context for this observable. + bqm: The binary quadratic model corresponding to the problem instance. + sample_set: The samples used to compute the observable. + """ class QubitMagnetization(Observable): @@ -99,7 +106,7 @@ def evaluate( A numpy array containing the pairwise spin correlations for each coupler. """ sample_array = dimod.as_samples(sample_set)[0].astype(float) - if len(experiment.inst.edge_list) == 0: + if not experiment.inst.edge_list: return np.empty(0, dtype=float) row, col = np.asarray(experiment.inst.edge_list).T @@ -127,7 +134,7 @@ def evaluate( A numpy array containing the mean coupler frustration for each edge. """ sample_array = dimod.as_samples(sample_set)[0].astype(float) - if len(experiment.inst.edge_list) == 0: + if not experiment.inst.edge_list: return np.empty(0, dtype=float) row, col = np.asarray(experiment.inst.edge_list).T @@ -216,7 +223,7 @@ def evaluate( sample_set: The sample set is not used in this observable, but is included in the signature for consistency with other observables. path: Optional path to load/save the reference energy. If not provided, - a default path will be generated based on the experiment and BQM. + a default path will be generated based on the experiment and BQM. Returns: The reference energy for the given BQM. @@ -250,7 +257,18 @@ def load( bqm: dimod.BQM, path: str | Path | None = None, ) -> tuple[float, NDArray, str]: - """Load and get the full data tuple, not just the energy.""" + """Load and get the full data tuple, not just the energy. + + Args: + experiment: The experiment for which to load the reference energy. + bqm: The binary quadratic model corresponding to the problem instance. + path: Optional path to load the reference energy. If not provided, + a default path will be generated based on the experiment and BQM. + + Returns: + A tuple containing the reference energy, the corresponding sample + as a NumPy array, and a string indicating the optimization method used. + """ if path is not None: path = Path(path) else: @@ -265,7 +283,15 @@ def load( 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.""" + """Save the reference energy to disk. + + Args: + path: Path to save the reference energy file. + energy: The reference energy to save. + sample: The corresponding sample to save. + method_string: A string indicating the optimization method used to + obtain the reference energy. + """ path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) np.savetxt(path, sample, fmt="%d", header=f"{method_string}\n{energy}", comments="") @@ -274,13 +300,20 @@ def update( self, experiment: Experiment, bqm: dimod.BQM, - sample, + sample: NDArray, path: str | Path | None = None, ) -> None: """Update the cached reference energy if the provided sample improves it. Use this when you get an energy that is lower than the reference energy. We want to keep the old method string unless it is specified. + + Args: + experiment: The experiment for which to update the reference energy. + bqm: The binary quadratic model corresponding to the problem instance. + sample: The new sample that may improve the reference energy. + path: Optional path to load/save the reference energy. If not provided, + a default path will be generated based on the experiment and BQM. """ if path is not None: path = Path(path) @@ -288,13 +321,13 @@ def update( 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(bqm, experiment) - self.save(path, new_energy, sample, reference_method_string) - else: + if new_energy >= reference_energy: raise ValueError("New energy is not better than reference energy, not updating.") + if path is None: + path = get_reference_energy_path(bqm, experiment) + self.save(path, new_energy, sample, reference_method_string) + def get_reference_energy_path( bqm: dimod.BQM, @@ -308,7 +341,7 @@ def get_reference_energy_path( pathstring, for example when ground-state energies depend on the specific chip. Args: - bqm: The BQM for which to get the reference energy path. + bqm: The binary quadratic model for which to get the reference energy path. experiment: The experiment for which to get the reference energy path. root: Optional root directory to use instead of the experiment's data root. dummy_experiment_data_dict: A dictionary containing the keys ``run_index``, diff --git a/dwave/experimental/lattice_utils/observable/triangular.py b/dwave/experimental/lattice_utils/observable/triangular.py index 9c9509e..541fbfa 100644 --- a/dwave/experimental/lattice_utils/observable/triangular.py +++ b/dwave/experimental/lattice_utils/observable/triangular.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,10 +13,11 @@ # limitations under the License. from __future__ import annotations -import numpy as np -from numpy.typing import NDArray + import dimod +import numpy as np from dimod import BQM +from numpy.typing import NDArray from dwave.experimental.lattice_utils.observable.observable import Observable diff --git a/dwave/experimental/lattice_utils/utils.py b/dwave/experimental/lattice_utils/utils.py index 94341ab..b191cd7 100644 --- a/dwave/experimental/lattice_utils/utils.py +++ b/dwave/experimental/lattice_utils/utils.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/example_1D_Ising_chain_shim.py b/examples/lattice_utils/1D_Ising_chain_shim.py similarity index 96% rename from examples/example_1D_Ising_chain_shim.py rename to examples/lattice_utils/1D_Ising_chain_shim.py index 8119384..d771323 100644 --- a/examples/example_1D_Ising_chain_shim.py +++ b/examples/lattice_utils/1D_Ising_chain_shim.py @@ -14,16 +14,14 @@ """Shimming example for 1D Ising chain.""" -from pathlib import Path import os +from pathlib import Path -from dwave.system import DWaveSampler import matplotlib.pyplot as plt -from matplotlib.colors import to_rgb import numpy as np +from dwave.system import DWaveSampler -from dwave.experimental.lattice_utils import lattice, experiment, observable -from dwave.experimental.lattice_utils.utils import bootstrap, confidence_interval +from dwave.experimental.lattice_utils import lattice, experiment # Set up the parameters @@ -90,7 +88,7 @@ # Make parameter list. We will only vary anneal time. for _ in range(120): - done = exp.run_iteration([{"anneal_time": anneal_time}], progress=True) + done = exp.run_iteration([{"anneal_time": anneal_time}]) if done: break diff --git a/pyproject.toml b/pyproject.toml index 57a587a..1b8a959 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,9 @@ dependencies = [ # https://github.com/python/cpython/issues/142214 requires-python = ">=3.10, !=3.14.1" +[project.optional-dependencies] +progress = ["tqdm>=4,<5"] + [project.urls] Issues = "https://github.com/dwavesystems/dwave-experimental/issues" Repository = "https://github.com/dwavesystems/dwave-experimental.git" diff --git a/releasenotes/notes/add-lattice-utils-49e5efd79268e69d.yaml b/releasenotes/notes/add-lattice-utils-49e5efd79268e69d.yaml index 27fd3bb..71d728b 100644 --- a/releasenotes/notes/add-lattice-utils-49e5efd79268e69d.yaml +++ b/releasenotes/notes/add-lattice-utils-49e5efd79268e69d.yaml @@ -2,7 +2,19 @@ features: - | Add ``dwave.experimental.lattice_utils`` submodule with utilities for - constructing lattice graphs (``Chain``, ``Triangular``, - ``DimerizedTriangular``, ``EmbeddedLattice``), evaluating physics + constructing lattice graphs, evaluating physics observables on sample sets, and running shimmed Ising experiments on QPU samplers. + - | + Add lattice graph constructors ``Chain``, ``Triangular``, and + ``DimerizedTriangular`` for building common lattice geometries. + - | + Add the ``EmbeddedLattice`` class for embedding logical lattices onto physical + hardware graphs using ordered chains. + - | + Add the ``Observable`` base class and several observables (e.g. qubit + magnetization, coupler correlation, and coupler frustration) for evaluating + physics quantities on sample sets. + - | + Add the ``Experiment`` class for running shimmed Ising experiments on QPU + samplers and collecting observables from the resulting sample sets. diff --git a/tests/__init__.py b/tests/__init__.py index de79690..0002aae 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/test_lattice_utils/__init__.py b/tests/test_lattice_utils/__init__.py index e69de29..0002aae 100644 --- a/tests/test_lattice_utils/__init__.py +++ b/tests/test_lattice_utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 D-Wave +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/test_lattice_utils/_helpers.py b/tests/test_lattice_utils/_helpers.py index 8fe25ad..680efbf 100644 --- a/tests/test_lattice_utils/_helpers.py +++ b/tests/test_lattice_utils/_helpers.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/test_lattice_utils/test_experiment.py b/tests/test_lattice_utils/test_experiment.py index 1302dcb..971993c 100644 --- a/tests/test_lattice_utils/test_experiment.py +++ b/tests/test_lattice_utils/test_experiment.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/test_lattice_utils/test_lattice.py b/tests/test_lattice_utils/test_lattice.py index c85c546..3afdeb6 100644 --- a/tests/test_lattice_utils/test_lattice.py +++ b/tests/test_lattice_utils/test_lattice.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -45,11 +45,6 @@ def test_single_node_periodic(self): self.assertEqual(chain.num_spins, 1) self.assertEqual(chain.num_edges, 0) - def test_two_node_periodic(self): - with tempfile.TemporaryDirectory() as tmpdir: - chain = lattice.Chain(dimensions=(2,), periodic=(True,), data_root=tmpdir) - self.assertEqual(chain.num_edges, 2) - def test_geometry_name(self): with tempfile.TemporaryDirectory() as tmpdir: chain = lattice.Chain(dimensions=(6,), periodic=(True,), data_root=tmpdir) diff --git a/tests/test_lattice_utils/test_observable.py b/tests/test_lattice_utils/test_observable.py index 330e830..2aeea18 100644 --- a/tests/test_lattice_utils/test_observable.py +++ b/tests/test_lattice_utils/test_observable.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/test_lattice_utils/test_utils.py b/tests/test_lattice_utils/test_utils.py index 5458f3e..1b821ad 100644 --- a/tests/test_lattice_utils/test_utils.py +++ b/tests/test_lattice_utils/test_utils.py @@ -1,4 +1,4 @@ -# Copyright 2025 D-Wave +# Copyright 2026 D-Wave # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. From b34c98453654a085aaa1078334b4949636bfcfef Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 13 Aug 2026 10:27:18 -0700 Subject: [PATCH 13/13] Fix type hints -import Experiment and Lattice types for proper type hints -fix typos -SamplerCall is now a proper data class --- .../lattice_utils/experiment/experiment.py | 99 +++++++++++-------- .../lattice_utils/experiment/samplercall.py | 20 ++-- .../lattice_utils/lattice/embedded_lattice.py | 1 + .../lattice_utils/lattice/lattice.py | 10 +- .../lattice_utils/lattice/optimize.py | 5 +- .../lattice_utils/lattice/triangular.py | 2 +- .../lattice_utils/observable/kinks.py | 5 + .../lattice_utils/observable/observable.py | 7 +- .../lattice_utils/observable/triangular.py | 5 + examples/lattice_utils/1D_Ising_chain_shim.py | 14 +-- tests/test_lattice_utils/test_experiment.py | 60 +++++------ 11 files changed, 127 insertions(+), 101 deletions(-) diff --git a/dwave/experimental/lattice_utils/experiment/experiment.py b/dwave/experimental/lattice_utils/experiment/experiment.py index 947c226..524c66a 100644 --- a/dwave/experimental/lattice_utils/experiment/experiment.py +++ b/dwave/experimental/lattice_utils/experiment/experiment.py @@ -23,7 +23,7 @@ from dataclasses import dataclass from datetime import datetime from pathlib import Path -from typing import Any +from typing import Any, TYPE_CHECKING import dimod import numpy as np @@ -33,6 +33,9 @@ except ImportError: tqdm = None +if TYPE_CHECKING: + from tqdm.std import tqdm as tqdm_bar + from dwave.experimental.lattice_utils.experiment.samplercall import SamplerCall from dwave.experimental.lattice_utils.lattice.lattice import Lattice from dwave.experimental.lattice_utils.observable import ( @@ -128,7 +131,7 @@ def load_results( results are loaded. result_fields: Subset of fields to extract from each result file. If ``None``, all fields present in the first result file are used. - ignore_shim: If true, the ``shimdata`` field is removed from the + ignore_shim: If true, the ``shim_data`` field is removed from the returned results. Returns: @@ -152,7 +155,7 @@ def load_results( if result_fields is None: result_fields = list(data.keys()) if ignore_shim: - result_fields.remove("shimdata") + result_fields.remove("shim_data") results.append({k: data[k] for k in result_fields}) @@ -338,8 +341,15 @@ def _make_progress_bar( colour: str, bar_format: str | None = None, initial: int | float = 0, - ) -> tqdm: - """Create a tqdm progress bar with consistent formatting.""" + ) -> tqdm_bar: + """Create a tqdm progress bar with consistent formatting. + + The optional ``tqdm`` dependency is required to call this method. If + ``tqdm`` is not installed, an ImportError will be raised. + """ + if tqdm is None: + raise ImportError("Progress reporting requires the optional 'tqdm' dependency.") + if bar_format is None: bar_width = min(100, max(total, 20)) bar_format = f"{{desc}}: |{{bar:{bar_width}}}{{r_bar}}{{bar:-{bar_width}b}}" @@ -357,11 +367,18 @@ def _print_iteration_status( call_dict: dict[int, SamplerCall], num_params: int, ) -> None: - """Print a summary of the iteration status, including progress and iteration ranges.""" + """Print a summary of the iteration status, including progress and iteration ranges. + + The optional ``tqdm`` dependency is required to call this method. If + ``tqdm`` is not installed, an ImportError will be raised. + """ + if tqdm is None: + raise ImportError("Progress reporting requires the optional 'tqdm' dependency.") + iteration_range = ( f"Iteration range " - f"{min(call.shimdata['total_iterations'] for call in call_dict.values())}-" - f"{max(call.shimdata['total_iterations'] for call in call_dict.values())} " + f"{min(call.shim_data['total_iterations'] for call in call_dict.values())}-" + f"{max(call.shim_data['total_iterations'] for call in call_dict.values())} " ) if self.max_iterations is None: tqdm.write(" Total progress: " + iteration_range) @@ -369,15 +386,15 @@ def _print_iteration_status( total = num_params * self.max_iterations progress_value = ( - sum(call.shimdata["total_iterations"] for call in call_dict.values()) + sum(call.shim_data["total_iterations"] for call in call_dict.values()) + (num_params - len(call_dict)) * self.max_iterations ) progress_string = ( f"{progress_value / total * 100:.1f}% " f"Iteration range " - f"{min(call.shimdata['total_iterations'] for call in call_dict.values())}-" - f"{max(call.shimdata['total_iterations'] for call in call_dict.values())} " + f"{min(call.shim_data['total_iterations'] for call in call_dict.values())}-" + f"{max(call.shim_data['total_iterations'] for call in call_dict.values())} " f"of {self.max_iterations} " f"({num_params - len(call_dict)} of {num_params} parameters finished)" ) @@ -515,17 +532,17 @@ def _build_sampler_call(self) -> None | SamplerCall: """ sampler_call = SamplerCall(run_index=self.run_index) sampler_call.logical_bqms = self._make_logical_bqms() - sampler_call.shimdata = self._get_shimdata() + sampler_call.shim_data = self._get_shim_data() # Here we can find out that we're finished. if ( self.max_iterations is not None - and sampler_call.shimdata["total_iterations"] >= self.max_iterations + and sampler_call.shim_data["total_iterations"] >= self.max_iterations ): return None sampler_call.bqm = self._make_bqm(sampler_call) - sampler_call.sampler_params = self._make_sampler_params(shimdata=sampler_call.shimdata) + sampler_call.sampler_params = self._make_sampler_params(shim_data=sampler_call.shim_data) return sampler_call @@ -569,14 +586,14 @@ def _generate_data_to_save( else: savedata[key] = results[key].copy() - savedata["shimdata"] = {} - for key in sampler_call.shimdata: - if isinstance(sampler_call.shimdata[key], np.ndarray): - savedata["shimdata"][key] = sampler_call.shimdata[key].astype(np.float32) - elif isinstance(sampler_call.shimdata[key], int): - savedata["shimdata"][key] = sampler_call.shimdata[key] + savedata["shim_data"] = {} + for key in sampler_call.shim_data: + if isinstance(sampler_call.shim_data[key], np.ndarray): + savedata["shim_data"][key] = sampler_call.shim_data[key].astype(np.float32) + elif isinstance(sampler_call.shim_data[key], int): + savedata["shim_data"][key] = sampler_call.shim_data[key] else: - savedata["shimdata"][key] = sampler_call.shimdata[key].copy() + savedata["shim_data"][key] = sampler_call.shim_data[key].copy() return savedata @@ -591,11 +608,11 @@ def _make_sampler_params(self, **kwargs) -> dict[str, Any]: "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 "shim_data" in kwargs: + if "flux_biases" in kwargs["shim_data"]: + ret["flux_biases"] = list(kwargs["shim_data"]["flux_biases"]) + if "anneal_offsets" in kwargs["shim_data"]: + ret["anneal_offsets"] = list(kwargs["shim_data"]["anneal_offsets"]) if self.param.get("fast_anneal", False): ret["fast_anneal"] = True @@ -607,7 +624,7 @@ def _make_sampler_params(self, **kwargs) -> dict[str, Any]: return ret - def _get_shimdata(self) -> dict[str, Any]: + def _get_shim_data(self) -> dict[str, Any]: """Load shim data if possible, otherwise make an initial shim.""" if self.already_initialized: return self._load_shim() @@ -615,17 +632,17 @@ def _get_shimdata(self) -> dict[str, Any]: 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} + shim_data = {"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)) + shim_data["flux_biases"] = np.zeros(self.sampler.properties["num_qubits"]) + shim_data["anneal_offsets"] = np.zeros(self.sampler.properties["num_qubits"]) + shim_data["relative_coupler_strength"] = np.ones((num_embeddings, self.inst.num_edges)) if self.param.get("flux_biases", None) is not None: - shimdata["flux_biases"] = self.param.get("flux_biases") + shim_data["flux_biases"] = self.param.get("flux_biases") - return shimdata + return shim_data def _get_latest_iteration_filename(self) -> Path: """Return the filename of the most recently completed iteration.""" @@ -642,8 +659,8 @@ def _load_shim(self) -> dict[str, Any]: try: with lzma.open(filename, "rb") as f: data = pickle.load(f) - shimdata = data["shimdata"] - return shimdata + shim_data = data["shim_data"] + return shim_data except FileNotFoundError as e: raise FileNotFoundError(f"{filename} does not exist") from e except Exception as e: @@ -651,21 +668,21 @@ def _load_shim(self) -> dict[str, Any]: def _update_shim(self, sampler_call: SamplerCall, results: dict[str, Any]) -> None: """Update shim parameters according to shim data and parameters.""" - if "flux_biases" in sampler_call.shimdata and self.param.get("flux_bias_shim_step", 0) != 0: + if "flux_biases" in sampler_call.shim_data and self.param.get("flux_bias_shim_step", 0) != 0: self._update_flux_bias_shim(sampler_call, results) if ( - "relative_coupler_strength" in sampler_call.shimdata + "relative_coupler_strength" in sampler_call.shim_data and self.param.get("coupler_shim_step", 0) != 0 ): self._update_coupler_shim(sampler_call, results) - sampler_call.shimdata["total_iterations"] += 1 + sampler_call.shim_data["total_iterations"] += 1 def _update_flux_bias_shim(self, sampler_call: SamplerCall, results: dict[str, Any]) -> None: """Update flux-bias shim values based on qubit magnetization.""" target_magnetization = self.param["target_magnetization"] qubit_magnetization = results["QubitMagnetization"] - flux_biases = sampler_call.shimdata["flux_biases"] + flux_biases = sampler_call.shim_data["flux_biases"] shim_step = self.param["flux_bias_shim_step"] steps = shim_step * (qubit_magnetization.ravel() - target_magnetization) @@ -693,7 +710,7 @@ def _update_coupler_shim( """Update relative coupler strength based on measured frustration.""" orbits = self.inst.coupler_orbits signed_energy_scale = self.param["signed_energy_scale"] - relative_coupler_strength = sampler_call.shimdata["relative_coupler_strength"] + relative_coupler_strength = sampler_call.shim_data["relative_coupler_strength"] # Allow for zero step size, which will just truncate the shim. if step_size is None: @@ -813,7 +830,7 @@ def _make_bqm(self, sampler_call: SamplerCall) -> dimod.BQM: return bqm - relative_coupler_strength = sampler_call.shimdata["relative_coupler_strength"] + relative_coupler_strength = sampler_call.shim_data["relative_coupler_strength"] for iemb, emb in enumerate(self.inst.embedding_list): logical_bqm = sampler_call.logical_bqms[iemb].copy() diff --git a/dwave/experimental/lattice_utils/experiment/samplercall.py b/dwave/experimental/lattice_utils/experiment/samplercall.py index e02c9d5..b8cffb9 100644 --- a/dwave/experimental/lattice_utils/experiment/samplercall.py +++ b/dwave/experimental/lattice_utils/experiment/samplercall.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dataclasses import dataclass +from dataclasses import dataclass, field import dimod @@ -23,16 +23,8 @@ class SamplerCall: """Data class for managing asynchronous sampler calls.""" - def __init__( - self, - run_index: int, - shimdata: dict | None = None, - embedded_bqm: dimod.BQM | None = None, - logical_bqms: list | None = None, - sampler_params: dict | None = None, - ): - self.run_index: int = run_index - self.bqm: dimod.BQM | None = embedded_bqm - self.shimdata: dict = {} if shimdata is None else shimdata - self.logical_bqms: list = [] if logical_bqms is None else logical_bqms - self.sampler_params: dict = {} if sampler_params is None else sampler_params + run_index: int + bqm: dimod.BQM | None = None + shim_data: dict = field(default_factory=dict) + logical_bqms: list = field(default_factory=list) + sampler_params: dict = field(default_factory=dict) diff --git a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py index 02bf708..9ad8a71 100644 --- a/dwave/experimental/lattice_utils/lattice/embedded_lattice.py +++ b/dwave/experimental/lattice_utils/lattice/embedded_lattice.py @@ -96,6 +96,7 @@ def get_chain_connectivity( u: The first node in the logical edge. v: The second node in the logical edge. If None, this is treated as a chain edge (u == v). + Returns: A tuple of tuples, where each inner tuple represents a pair of indices in the chains corresponding to u and v that should be connected. For diff --git a/dwave/experimental/lattice_utils/lattice/lattice.py b/dwave/experimental/lattice_utils/lattice/lattice.py index 2c670bc..e6784eb 100644 --- a/dwave/experimental/lattice_utils/lattice/lattice.py +++ b/dwave/experimental/lattice_utils/lattice/lattice.py @@ -34,7 +34,7 @@ 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, + Subclasses are responsible for defining the lattice geometry itself. In particular, a subclass must: - Implement the ``generate_edges`` method, which yields the edges of the lattice as pairs @@ -76,12 +76,10 @@ def __init__( f"periodic and dimensions must have the same length: " f"got {len(self.periodic)} and {len(self.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.edge_list: list[tuple[Hashable, Hashable]] = list(self.generate_edges()) self.num_edges: int = len(self.edge_list) self.orbit_type: str = orbit_type self.initialize_orbits(qubit_orbits, coupler_orbits) @@ -267,7 +265,7 @@ def _get_path( sampler_name: str | None = None, extra_subdir: str | Path | None = None, ) -> Path: - """Construct a standarized file path for embedding or orbit data.""" + """Construct a standardized file path for embedding or orbit data.""" if kind not in {"embedding", "orbits"}: raise ValueError("kind must be provided as either `embedding` or `orbits`") @@ -297,7 +295,7 @@ def _save_embeddings(self, sampler: dimod.Sampler, embeddings: NDArray) -> None: os.makedirs(cache_filename.parent, exist_ok=True) np.savetxt(cache_filename, embeddings, fmt="%d") - def _load_embeddings(self, sampler: str) -> None: + def _load_embeddings(self, sampler: dimod.Sampler) -> None: """Load embedding data.""" filename = self._make_filename("embedding", sampler=sampler) self.embedding_list = np.atleast_2d(np.loadtxt(filename, dtype=int)) diff --git a/dwave/experimental/lattice_utils/lattice/optimize.py b/dwave/experimental/lattice_utils/lattice/optimize.py index c4a121e..574ce7a 100644 --- a/dwave/experimental/lattice_utils/lattice/optimize.py +++ b/dwave/experimental/lattice_utils/lattice/optimize.py @@ -14,13 +14,16 @@ from __future__ import annotations -from typing import Any +from typing import Any, TYPE_CHECKING import dimod import numpy as np from dwave.samplers import SimulatedAnnealingSampler from numpy.typing import NDArray +if TYPE_CHECKING: + from dwave.experimental.lattice_utils.experiment.experiment import Lattice + __all__ = ['optimize', 'ExponentialBackoffSimulatedAnnealingSampler'] diff --git a/dwave/experimental/lattice_utils/lattice/triangular.py b/dwave/experimental/lattice_utils/lattice/triangular.py index 0a6ac5c..1046341 100644 --- a/dwave/experimental/lattice_utils/lattice/triangular.py +++ b/dwave/experimental/lattice_utils/lattice/triangular.py @@ -262,7 +262,7 @@ def get_chain_connectivity( 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 + in the chains corresponding to u and v that should be connected. For a chain edge (u == v or v is None), this will return pairs of indices within the same chain. For a logical edge (u != v), this will return pairs of indices between the two chains. diff --git a/dwave/experimental/lattice_utils/observable/kinks.py b/dwave/experimental/lattice_utils/observable/kinks.py index 665d5f3..3f55717 100644 --- a/dwave/experimental/lattice_utils/observable/kinks.py +++ b/dwave/experimental/lattice_utils/observable/kinks.py @@ -14,6 +14,8 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import dimod import numpy as np from dimod import BQM, SampleSet @@ -21,6 +23,9 @@ from dwave.experimental.lattice_utils.observable.observable import Observable +if TYPE_CHECKING: + from dwave.experimental.lattice_utils.experiment.experiment import Experiment + __all__ = ['KinkKinkCorrelator'] diff --git a/dwave/experimental/lattice_utils/observable/observable.py b/dwave/experimental/lattice_utils/observable/observable.py index 4fb91a4..fed8475 100644 --- a/dwave/experimental/lattice_utils/observable/observable.py +++ b/dwave/experimental/lattice_utils/observable/observable.py @@ -16,7 +16,7 @@ from abc import ABC, abstractmethod from pathlib import Path -from typing import Any, TypeAlias +from typing import Any, TypeAlias, TYPE_CHECKING import dimod import numpy as np @@ -24,6 +24,9 @@ from dwave.experimental.lattice_utils.lattice.lattice import Lattice +if TYPE_CHECKING: + from dwave.experimental.lattice_utils.experiment.experiment import Experiment + __all__ = [ 'Observable', 'QubitMagnetization', @@ -353,6 +356,8 @@ def get_reference_energy_path( The path to the reference energy file. """ if experiment is None: + if dummy_experiment_data_dict is None: + raise ValueError("Provide either 'experiment' or 'dummy_experiment_data_dict'.") experiment_data_dict = dummy_experiment_data_dict else: experiment_data_dict = { diff --git a/dwave/experimental/lattice_utils/observable/triangular.py b/dwave/experimental/lattice_utils/observable/triangular.py index 541fbfa..a24f7b4 100644 --- a/dwave/experimental/lattice_utils/observable/triangular.py +++ b/dwave/experimental/lattice_utils/observable/triangular.py @@ -14,6 +14,8 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import dimod import numpy as np from dimod import BQM @@ -21,6 +23,9 @@ from dwave.experimental.lattice_utils.observable.observable import Observable +if TYPE_CHECKING: + from dwave.experimental.lattice_utils.experiment.experiment import Experiment + __all__ = ['TriangularOP'] diff --git a/examples/lattice_utils/1D_Ising_chain_shim.py b/examples/lattice_utils/1D_Ising_chain_shim.py index d771323..95fd885 100644 --- a/examples/lattice_utils/1D_Ising_chain_shim.py +++ b/examples/lattice_utils/1D_Ising_chain_shim.py @@ -14,7 +14,6 @@ """Shimming example for 1D Ising chain.""" -import os from pathlib import Path import matplotlib.pyplot as plt @@ -39,9 +38,6 @@ point_style = {"linestyle": "", "markersize": 5} cm = plt.get_cmap("tab10") -# 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] # Make a lattice instance for a periodic 256-spin chain, so we can embed it. @@ -109,9 +105,9 @@ mag[anneal_time] = np.array([it["QubitMagnetization"] for it in res]) frust[anneal_time] = np.array([it["CouplerFrustration"] for it in res]) cshim[anneal_time] = np.array( - [it["shimdata"]["relative_coupler_strength"].ravel() for it in res] + [it["shim_data"]["relative_coupler_strength"].ravel() for it in res] ) - fbshim[anneal_time] = np.array([it["shimdata"]["flux_biases"] for it in res]) + fbshim[anneal_time] = np.array([it["shim_data"]["flux_biases"] for it in res]) title = ( f"1D chain shim, " @@ -180,5 +176,9 @@ filename = title for bad_symbol in "/: ;,": filename = filename.replace(bad_symbol, "_") -fig.savefig(Path(os.getcwd()) / "figures" / f"{filename}.png") + +# Create a folder to save figures in if it doesn't already exist. +(data_root / "figures").mkdir(exist_ok=True) + +fig.savefig(data_root / "figures" / f"{filename}.png") plt.show() diff --git a/tests/test_lattice_utils/test_experiment.py b/tests/test_lattice_utils/test_experiment.py index 971993c..69a0427 100644 --- a/tests/test_lattice_utils/test_experiment.py +++ b/tests/test_lattice_utils/test_experiment.py @@ -31,7 +31,7 @@ def test_defaults(self): sc = experiment.SamplerCall(run_index=0) self.assertEqual(sc.run_index, 0) self.assertIsNone(sc.bqm) - self.assertEqual(sc.shimdata, {}) + self.assertEqual(sc.shim_data, {}) self.assertEqual(sc.logical_bqms, []) self.assertEqual(sc.sampler_params, {}) @@ -39,14 +39,14 @@ def test_with_values(self): bqm = dimod.BQM(vartype="SPIN") sc = experiment.SamplerCall( run_index=5, - embedded_bqm=bqm, - shimdata={"total_iterations": 1}, + bqm=bqm, + shim_data={"total_iterations": 1}, logical_bqms=[bqm], sampler_params={"num_reads": 100}, ) self.assertEqual(sc.run_index, 5) self.assertIs(sc.bqm, bqm) - self.assertEqual(sc.shimdata["total_iterations"], 1) + self.assertEqual(sc.shim_data["total_iterations"], 1) class TestExperimentInit(unittest.TestCase): @@ -108,9 +108,9 @@ def test_initial_shim_no_embeddings(self): sampler = _make_mock_sampler() exp = experiment.Experiment(inst=chain, sampler=sampler) exp.already_initialized = False - shimdata = exp._make_initial_shim() - self.assertEqual(shimdata["total_iterations"], 0) - self.assertNotIn("flux_biases", shimdata) + shim_data = exp._make_initial_shim() + self.assertEqual(shim_data["total_iterations"], 0) + self.assertNotIn("flux_biases", shim_data) def test_initial_shim_with_embeddings(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -118,9 +118,9 @@ def test_initial_shim_with_embeddings(self): chain.embedding_list = np.array([[0, 1, 2, 3]]) sampler = _make_mock_sampler(num_qubits=128) exp = experiment.Experiment(inst=chain, sampler=sampler) - shimdata = exp._make_initial_shim() - self.assertIn("flux_biases", shimdata) - self.assertEqual(len(shimdata["flux_biases"]), 128) + shim_data = exp._make_initial_shim() + self.assertIn("flux_biases", shim_data) + self.assertEqual(len(shim_data["flux_biases"]), 128) def test_initial_shim_with_preset_flux_biases(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -130,8 +130,8 @@ def test_initial_shim_with_preset_flux_biases(self): fb = np.ones(128) * 0.01 exp = experiment.Experiment(inst=chain, sampler=sampler) exp.param["flux_biases"] = fb - shimdata = exp._make_initial_shim() - np.testing.assert_array_almost_equal(shimdata["flux_biases"], fb) + shim_data = exp._make_initial_shim() + np.testing.assert_array_almost_equal(shim_data["flux_biases"], fb) def test_load_shim_from_file(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -141,8 +141,8 @@ def test_load_shim_from_file(self): exp.run_index = 1 exp.data_path = Path(tmpdir) - shimdata = {"total_iterations": 5, "flux_biases": np.zeros(10)} - data = {"shimdata": shimdata} + shim_data = {"total_iterations": 5, "flux_biases": np.zeros(10)} + data = {"shim_data": shim_data} fn = Path(tmpdir) / "iter00000.pkl.lzma" with lzma.open(fn, "wb") as f: pickle.dump(data, f) @@ -188,14 +188,14 @@ def test_load_shim_missing_file(self): with self.assertRaises(FileNotFoundError): exp._load_shim() - def test_get_shimdata_not_initialized(self): + def test_get_shim_data_not_initialized(self): with tempfile.TemporaryDirectory() as tmpdir: chain = lattice.Chain(dimensions=(4,), periodic=(True,), data_root=tmpdir) sampler = _make_mock_sampler() exp = experiment.Experiment(inst=chain, sampler=sampler) exp.already_initialized = False - shimdata = exp._get_shimdata() - self.assertEqual(shimdata["total_iterations"], 0) + shim_data = exp._get_shim_data() + self.assertEqual(shim_data["total_iterations"], 0) class TestCouplerShim(unittest.TestCase): @@ -222,7 +222,7 @@ def test_coupler_shim_basic_update(self): sc = experiment.SamplerCall(run_index=0) sc.logical_bqms = [chain.make_bqm()] - sc.shimdata = { + sc.shim_data = { "total_iterations": 0, "relative_coupler_strength": np.ones((1, chain.num_edges)), } @@ -234,7 +234,7 @@ def test_coupler_shim_basic_update(self): exp._update_coupler_shim(sc, results) np.testing.assert_array_almost_equal( - sc.shimdata["relative_coupler_strength"], + sc.shim_data["relative_coupler_strength"], np.array([[0.998, 1.0, 1.002]]), ) @@ -258,7 +258,7 @@ def test_coupler_shim_singleton_orbits_is_noop(self): sc = experiment.SamplerCall(run_index=0) sc.logical_bqms = [chain.make_bqm()] rcs_before = np.ones((1, chain.num_edges)) - sc.shimdata = { + sc.shim_data = { "total_iterations": 0, "relative_coupler_strength": rcs_before.copy(), } @@ -266,7 +266,7 @@ def test_coupler_shim_singleton_orbits_is_noop(self): exp._update_coupler_shim(sc, results) - np.testing.assert_array_equal(sc.shimdata["relative_coupler_strength"], rcs_before) + np.testing.assert_array_equal(sc.shim_data["relative_coupler_strength"], rcs_before) class TestSaveLoadResults(unittest.TestCase): @@ -314,10 +314,10 @@ def test_load_results_ignore_shim(self): 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) + pickle.dump({"value": 0, "shim_data": {}}, f) results = exp.load_results(num_iterations=1, ignore_shim=True) - self.assertNotIn("shimdata", results[0]) + self.assertNotIn("shim_data", results[0]) def test_load_results_starting_iteration(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -329,7 +329,7 @@ def test_load_results_starting_iteration(self): 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) + pickle.dump({"value": i, "shim_data": {}}, f) results = exp.load_results(num_iterations=3, start_iteration=2) self.assertEqual(len(results), 3) @@ -353,7 +353,7 @@ def test_generate_data_type_conversions(self): sampler = _make_mock_sampler() exp = experiment.Experiment(inst=chain, sampler=sampler) sc = experiment.SamplerCall(run_index=0) - sc.shimdata = {"total_iterations": 1, "flux_biases": np.zeros(4)} + sc.shim_data = {"total_iterations": 1, "flux_biases": np.zeros(4)} results = { "QubitMagnetization": np.array([0.1, 0.2, 0.3, 0.4]), @@ -363,7 +363,7 @@ def test_generate_data_type_conversions(self): 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) + self.assertEqual(savedata["shim_data"]["total_iterations"], 1) class TestMakeBqm(unittest.TestCase): @@ -378,7 +378,7 @@ def test_make_bqm_no_embeddings(self): ) sc = experiment.SamplerCall(run_index=0) sc.logical_bqms = [chain.make_bqm()] - sc.shimdata = {"total_iterations": 0} + sc.shim_data = {"total_iterations": 0} bqm = exp._make_bqm(sc) for u, v in chain.edge_list: self.assertAlmostEqual(bqm.quadratic[(u, v)], 0.5) @@ -396,7 +396,7 @@ def test_make_bqm_with_embeddings(self): ) sc = experiment.SamplerCall(run_index=0) sc.logical_bqms = [chain.make_bqm()] - sc.shimdata = { + sc.shim_data = { "total_iterations": 0, "relative_coupler_strength": np.full((1, chain.num_edges), 2.0), } @@ -431,8 +431,8 @@ def test_run_iteration_basic(self): self.assertIn("QubitMagnetization", data) self.assertIn("CouplerCorrelation", data) - self.assertIn("shimdata", data) - self.assertEqual(data["shimdata"]["total_iterations"], 1) + self.assertIn("shim_data", data) + self.assertEqual(data["shim_data"]["total_iterations"], 1) def test_run_iteration_returns_true_when_finished(self): """run_iteration() returns True when max_iterations already reached."""