Coverage for jaxquantum/devices/analysis/sweeps.py: 24%
103 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1"""Sweeping tools."""
3import itertools
4import os
5from copy import deepcopy
6from tempfile import NamedTemporaryFile
8import jax
9import jax.numpy as jnp
10from tqdm import tqdm
13def run_jax_sweep(
14 params,
15 sweep_params,
16 metrics_func,
17 fixed_kwargs=None,
18 is_parallel=False,
19):
20 """Evaluate a pure metric over sweep points with ``jax.vmap``."""
21 if not sweep_params:
22 raise ValueError("sweep_params must not be empty")
23 keys = tuple(sweep_params)
24 values = tuple(jnp.asarray(sweep_params[key]) for key in keys)
25 if is_parallel:
26 length = values[0].shape[0]
27 if not all(value.shape[0] == length for value in values):
28 raise ValueError("parallel sweep parameters must have equal lengths")
29 points = values
30 else:
31 points = tuple(
32 grid.reshape(-1) for grid in jnp.meshgrid(*values, indexing="ij")
33 )
34 fixed_kwargs = dict(fixed_kwargs or {})
36 def evaluate(*point):
37 current = dict(params)
38 current.update(zip(keys, point))
39 return metrics_func(current, **fixed_kwargs)
41 return jax.vmap(evaluate)(*points)
44def run_sweep(
45 params,
46 sweep_params,
47 metrics_func,
48 fixed_kwargs=None,
49 data=None,
50 is_parallel=False,
51 save_file=None,
52 data_save_mode="end",
53 return_errors=False,
54):
55 """Run a sweep over a single parameter, or multiple parameters.
57 Args:
58 params (dict): The base parameters to sweep over.
59 sweep_params (dict): The parameters to sweep over.
60 key: The parameter name.
61 value: The list of values to sweep over.
62 metrics_func (function): The function to evaluate the metrics.
63 fixed_kwargs (dict, optional): Fixed keyword arguments for metrics_func.
64 data (dict, optional): The data to append to. Defaults to None.
65 is_parallel (bool, optional): Whether to sweep through the sweep_params lists in parallel or through their cartesian product. Defaults to False.
66 save_file (str, optional): The file to save the data to. Defaults to None, in which case data is saved to a temporary file, which will be deleted upon closing (e.g. during garbage collection).
67 data_save_mode (str, optional): The mode to save the data. Defaults to None.
68 Options are:
69 "no" - don't save data
70 "end" - save data at the end of the sweep
71 "during" - save data during and at the end of the sweep
72 Returns:
73 dict: The data after the sweep.
74 """
75 if data is None:
76 data = {}
77 run = -1
78 else:
79 run = max(data.keys())
81 assert data_save_mode in ["no", "end", "during"], "Invalid data_save_mode."
83 if data_save_mode in ["during", "end"]:
84 if isinstance(save_file, str):
85 print("Saving data to: ", save_file)
86 dirname = os.path.dirname(save_file)
87 if not os.path.exists(dirname):
88 os.makedirs(dirname)
89 else:
90 save_file = save_file or NamedTemporaryFile() # noqa: SIM115
91 print("Saving data to a temporary file: ", save_file.name)
93 fixed_kwargs = fixed_kwargs or {}
95 if is_parallel:
96 sweep_length = len(next(iter(sweep_params.values())))
97 assert all(len(vals) == sweep_length for vals in sweep_params.values()), (
98 "Parallel sweep parameters must have the same length."
99 )
101 errors = []
102 try:
103 for j in tqdm(range(sweep_length)):
104 run += 1
105 data[run] = {}
106 data[run]["params"] = deepcopy(params)
107 sweep_point_info = {
108 "labels": [],
109 "values": [],
110 "indices": [],
111 }
112 for key, vals in sweep_params.items():
113 data[run]["params"][key] = vals[j]
114 sweep_point_info["labels"].append(key)
115 sweep_point_info["values"].append(vals[j])
116 sweep_point_info["indices"].append(j)
117 data[run]["results"] = metrics_func(data[run]["params"], **fixed_kwargs)
118 data[run]["sweep_point_info"] = sweep_point_info
119 if data_save_mode == "during":
120 jnp.savez(
121 save_file, data=data, sweep_params=sweep_params, params=params
122 )
123 except Exception as e: # noqa: BLE001
124 errors.append(str(e))
125 print("Error during run: ", errors[-1])
127 try:
128 if data_save_mode in ["during", "end"]:
129 jnp.savez(
130 save_file,
131 data=data,
132 sweep_params=sweep_params,
133 params=params,
134 error=None,
135 )
136 except Exception as e: # noqa: BLE001
137 errors.append(str(e))
138 print("Error during saving: ", errors[-1])
140 if return_errors:
141 return data, errors
142 else:
143 return data
144 else:
145 # Product Sweep
146 sweep_points = list(itertools.product(*list(sweep_params.values())))
147 sweep_points_indxs = list(
148 itertools.product(
149 *[list(range(len(vals))) for vals in list(sweep_params.values())]
150 )
151 )
152 sweep_point_labels = list(sweep_params.keys())
154 errors = []
156 try:
157 with tqdm(total=len(sweep_points)) as pbar:
158 for j, sweep_point in enumerate(sweep_points):
159 run += 1
160 data[run] = {}
161 data[run]["params"] = deepcopy(params)
162 sweep_point_info = {
163 "labels": [],
164 "values": [],
165 "indices": [],
166 }
167 for i, key in enumerate(sweep_point_labels):
168 data[run]["params"][key] = sweep_point[i]
169 sweep_point_info["labels"].append(key)
170 sweep_point_info["values"].append(sweep_point[i])
171 sweep_point_info["indices"].append(sweep_points_indxs[j][i])
172 data[run]["results"] = metrics_func(
173 data[run]["params"], **fixed_kwargs
174 )
175 data[run]["sweep_point_info"] = sweep_point_info
176 pbar.update(1)
177 if data_save_mode == "during":
178 jnp.savez(
179 save_file,
180 data=data,
181 sweep_params=sweep_params,
182 params=params,
183 )
184 except Exception as e: # noqa: BLE001
185 errors.append(str(e))
186 print("Error during run: ", errors[-1])
188 try:
189 if data_save_mode in ["during", "end"]:
190 jnp.savez(
191 save_file, data=data, sweep_params=sweep_params, params=params
192 )
193 except Exception as e: # noqa: BLE001
194 errors.append(str(e))
195 print("Error during saving: ", errors[-1])
197 if return_errors:
198 return data, errors
199 else:
200 return data