From 7dac34b1536f15c14e4619dcdc7325b207808558 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Fri, 13 Mar 2026 13:41:22 -0700 Subject: [PATCH 1/6] 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 2/6] 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 3/6] 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 4/6] 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 5/6] 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 6/6] 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()