Skip to content

sweeps

Sweeping tools.

run_jax_sweep(params, sweep_params, metrics_func, fixed_kwargs=None, is_parallel=False)

Evaluate a pure metric over sweep points with jax.vmap.

Source code in jaxquantum/devices/analysis/sweeps.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def run_jax_sweep(
    params,
    sweep_params,
    metrics_func,
    fixed_kwargs=None,
    is_parallel=False,
):
    """Evaluate a pure metric over sweep points with ``jax.vmap``."""
    if not sweep_params:
        raise ValueError("sweep_params must not be empty")
    keys = tuple(sweep_params)
    values = tuple(jnp.asarray(sweep_params[key]) for key in keys)
    if is_parallel:
        length = values[0].shape[0]
        if not all(value.shape[0] == length for value in values):
            raise ValueError("parallel sweep parameters must have equal lengths")
        points = values
    else:
        points = tuple(
            grid.reshape(-1) for grid in jnp.meshgrid(*values, indexing="ij")
        )
    fixed_kwargs = dict(fixed_kwargs or {})

    def evaluate(*point):
        current = dict(params)
        current.update(zip(keys, point))
        return metrics_func(current, **fixed_kwargs)

    return jax.vmap(evaluate)(*points)

run_sweep(params, sweep_params, metrics_func, fixed_kwargs=None, data=None, is_parallel=False, save_file=None, data_save_mode='end', return_errors=False)

Run a sweep over a single parameter, or multiple parameters.

Parameters:

Name Type Description Default
params dict

The base parameters to sweep over.

 required
sweep_params dict

The parameters to sweep over. key: The parameter name. value: The list of values to sweep over.

 required
metrics_func function

The function to evaluate the metrics.

 required
fixed_kwargs dict

Fixed keyword arguments for metrics_func.

 None
data dict

The data to append to. Defaults to None.

 None
is_parallel bool

Whether to sweep through the sweep_params lists in parallel or through their cartesian product. Defaults to False.

 False
save_file str

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).

 None
data_save_mode str

The mode to save the data. Defaults to None. Options are: "no" - don't save data "end" - save data at the end of the sweep "during" - save data during and at the end of the sweep

 'end'

Returns: dict: The data after the sweep.

Source code in jaxquantum/devices/analysis/sweeps.py 
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def run_sweep(
    params,
    sweep_params,
    metrics_func,
    fixed_kwargs=None,
    data=None,
    is_parallel=False,
    save_file=None,
    data_save_mode="end",
    return_errors=False,
):
    """Run a sweep over a single parameter, or multiple parameters.

    Args:
        params (dict): The base parameters to sweep over.
        sweep_params (dict): The parameters to sweep over.
            key: The parameter name.
            value: The list of values to sweep over.
        metrics_func (function): The function to evaluate the metrics.
        fixed_kwargs (dict, optional): Fixed keyword arguments for metrics_func.
        data (dict, optional): The data to append to. Defaults to None.
        is_parallel (bool, optional): Whether to sweep through the sweep_params lists in parallel or through their cartesian product. Defaults to False.
        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).
        data_save_mode (str, optional): The mode to save the data. Defaults to None.
            Options are:
                "no" - don't save data
                "end" - save data at the end of the sweep
                "during" - save data during and at the end of the sweep
    Returns:
        dict: The data after the sweep.
    """
    if data is None:
        data = {}
        run = -1
    else:
        run = max(data.keys())

    assert data_save_mode in ["no", "end", "during"], "Invalid data_save_mode."

    if data_save_mode in ["during", "end"]:
        if isinstance(save_file, str):
            print("Saving data to: ", save_file)
            dirname = os.path.dirname(save_file)
            if not os.path.exists(dirname):
                os.makedirs(dirname)
        else:
            save_file = save_file or NamedTemporaryFile()  # noqa: SIM115
            print("Saving data to a temporary file: ", save_file.name)

    fixed_kwargs = fixed_kwargs or {}

    if is_parallel:
        sweep_length = len(next(iter(sweep_params.values())))
        assert all(len(vals) == sweep_length for vals in sweep_params.values()), (
            "Parallel sweep parameters must have the same length."
        )

        errors = []
        try:
            for j in tqdm(range(sweep_length)):
                run += 1
                data[run] = {}
                data[run]["params"] = deepcopy(params)
                sweep_point_info = {
                    "labels": [],
                    "values": [],
                    "indices": [],
                }
                for key, vals in sweep_params.items():
                    data[run]["params"][key] = vals[j]
                    sweep_point_info["labels"].append(key)
                    sweep_point_info["values"].append(vals[j])
                    sweep_point_info["indices"].append(j)
                data[run]["results"] = metrics_func(data[run]["params"], **fixed_kwargs)
                data[run]["sweep_point_info"] = sweep_point_info
                if data_save_mode == "during":
                    jnp.savez(
                        save_file, data=data, sweep_params=sweep_params, params=params
                    )
        except Exception as e:  # noqa: BLE001
            errors.append(str(e))
            print("Error during run: ", errors[-1])

        try:
            if data_save_mode in ["during", "end"]:
                jnp.savez(
                    save_file,
                    data=data,
                    sweep_params=sweep_params,
                    params=params,
                    error=None,
                )
        except Exception as e:  # noqa: BLE001
            errors.append(str(e))
            print("Error during saving: ", errors[-1])

        if return_errors:
            return data, errors
        else:
            return data
    else:
        # Product Sweep
        sweep_points = list(itertools.product(*list(sweep_params.values())))
        sweep_points_indxs = list(
            itertools.product(
                *[list(range(len(vals))) for vals in list(sweep_params.values())]
            )
        )
        sweep_point_labels = list(sweep_params.keys())

        errors = []

        try:
            with tqdm(total=len(sweep_points)) as pbar:
                for j, sweep_point in enumerate(sweep_points):
                    run += 1
                    data[run] = {}
                    data[run]["params"] = deepcopy(params)
                    sweep_point_info = {
                        "labels": [],
                        "values": [],
                        "indices": [],
                    }
                    for i, key in enumerate(sweep_point_labels):
                        data[run]["params"][key] = sweep_point[i]
                        sweep_point_info["labels"].append(key)
                        sweep_point_info["values"].append(sweep_point[i])
                        sweep_point_info["indices"].append(sweep_points_indxs[j][i])
                    data[run]["results"] = metrics_func(
                        data[run]["params"], **fixed_kwargs
                    )
                    data[run]["sweep_point_info"] = sweep_point_info
                    pbar.update(1)
                    if data_save_mode == "during":
                        jnp.savez(
                            save_file,
                            data=data,
                            sweep_params=sweep_params,
                            params=params,
                        )
        except Exception as e:  # noqa: BLE001
            errors.append(str(e))
            print("Error during run: ", errors[-1])

        try:
            if data_save_mode in ["during", "end"]:
                jnp.savez(
                    save_file, data=data, sweep_params=sweep_params, params=params
                )
        except Exception as e:  # noqa: BLE001
            errors.append(str(e))
            print("Error during saving: ", errors[-1])

        if return_errors:
            return data, errors
        else:
            return data