Skills › Scientific Agent Skills › qutip › convergence_sweep.py
convergence_sweep.py
358 lines · skills/qutip/scripts/convergence_sweep.py at 5f2bc03
1#!/usr/bin/env python3
2"""Run bounded deterministic or trajectory convergence sweeps."""
4from __future__ import annotations
6import argparse
7from typing import Any
9from _common import (
10 DEFAULT_SEED,
11 MAX_SWEEP_RUNS,
12 MAX_TIME_POINTS,
13 MAX_TOTAL_SWEEP_TRAJECTORIES,
14 MAX_TRAJECTORIES,
15 QUTIP_VERSION,
16 CliError,
17 add_output_arguments,
18 emit_json,
19 finite_float,
20 load_qutip,
21 parse_csv_ints,
22 run_cli,
23)
24from two_level_simulation import SimulationConfig, run_simulation
27def _parse_rtols(value: str) -> list[float]:
28 pieces = [piece.strip() for piece in value.split(",")]
29 if not pieces or any(not piece for piece in pieces):
30 raise CliError("rtols must be a comma-separated number list")
31 if len(pieces) > MAX_SWEEP_RUNS:
32 raise CliError(f"rtols may contain at most {MAX_SWEEP_RUNS} values")
33 try:
34 values = [float(piece) for piece in pieces]
35 except ValueError as exc:
36 raise CliError("rtols must contain only numbers") from exc
37 checked = [
38 finite_float(
39 item,
40 name="rtol",
41 minimum=1.0e-12,
42 maximum=1.0e-2,
43 )
44 for item in values
45 ]
46 if any(left <= right for left, right in zip(checked, checked[1:])):
47 raise CliError("rtols must be strictly decreasing from coarse to fine")
48 return checked
51def _base_config(
52 *,
53 solver: str,
54 args: argparse.Namespace,
55 time_points: int,
56 trajectories: int,
57 rtol: float,
58) -> SimulationConfig:
59 return SimulationConfig(
60 solver=solver,
61 initial_state="excited",
62 omega=finite_float(
63 args.omega,
64 name="omega",
65 minimum=-100_000.0,
66 maximum=100_000.0,
67 ),
68 drive=finite_float(
69 args.drive,
70 name="drive",
71 minimum=-100_000.0,
72 maximum=100_000.0,
73 ),
74 decay_rate=finite_float(
75 args.decay_rate,
76 name="decay_rate",
77 minimum=0.0,
78 maximum=10_000.0,
79 ),
80 dephasing_rate=finite_float(
81 args.dephasing_rate,
82 name="dephasing_rate",
83 minimum=0.0,
84 maximum=10_000.0,
85 ),
86 t_final=finite_float(
87 args.t_final,
88 name="t_final",
89 minimum=0.0,
90 maximum=100_000.0,
91 minimum_inclusive=False,
92 ),
93 time_points=time_points,
94 trajectories=trajectories,
95 seed=int(args.seed),
96 method=str(args.method),
97 atol=max(1.0e-14, rtol * 0.01),
98 rtol=rtol,
99 )
102def run_sweep(
103 args: argparse.Namespace,
104 qutip_module: Any | None = None,
105) -> dict[str, Any]:
106 """Execute the requested bounded convergence sweep."""
108 import numpy as np
110 qutip = qutip_module or load_qutip()
111 acceptance = finite_float(
112 args.acceptance,
113 name="acceptance",
114 minimum=1.0e-12,
115 maximum=0.5,
116 )
117 seed = int(args.seed)
118 if seed < 0 or seed > 2**63 - 1:
119 raise CliError("seed must be from 0 through 2^63-1")
121 reports: list[dict[str, Any]] = []
122 run_summaries: list[dict[str, Any]] = []
124 if args.mode == "deterministic":
125 points = parse_csv_ints(
126 args.time_points,
127 name="time_points",
128 minimum=2,
129 maximum=MAX_TIME_POINTS,
130 )
131 rtols = _parse_rtols(args.rtols)
132 if len(points) < 2:
133 raise CliError("deterministic convergence requires at least two runs")
134 if len(points) != len(rtols):
135 raise CliError("time_points and rtols must contain the same number of values")
136 for index, (count, rtol) in enumerate(zip(points, rtols)):
137 config = _base_config(
138 solver="mesolve",
139 args=args,
140 time_points=count,
141 trajectories=1,
142 rtol=rtol,
143 )
144 report = run_simulation(config, qutip)
145 reports.append(report)
146 run_summaries.append(
147 {
148 "index": index,
149 "time_points": count,
150 "atol": config.atol,
151 "rtol": rtol,
152 "final_excited_population": float(
153 report["expectations"]["excited_population"][-1]
154 ),
155 "analytic_max_abs_error": report["analytic_reference"][
156 "max_abs_error"
157 ],
158 "final_state_valid": report["checks"]["final_state"].get(
159 "valid_within_tolerance", False
160 ),
161 "solver_stats": report["solver"]["stats"],
162 }
163 )
165 reference_times = np.asarray(reports[-1]["times"], dtype=float)
166 reference_values = np.asarray(
167 reports[-1]["expectations"]["excited_population"],
168 dtype=float,
169 )
170 comparisons = []
171 for index, report in enumerate(reports[:-1]):
172 times = np.asarray(report["times"], dtype=float)
173 values = np.asarray(
174 report["expectations"]["excited_population"],
175 dtype=float,
176 )
177 reference_on_grid = np.interp(times, reference_times, reference_values)
178 difference = float(np.max(np.abs(values - reference_on_grid)))
179 comparisons.append(
180 {
181 "run_index": index,
182 "reference_run_index": len(reports) - 1,
183 "max_abs_population_difference": difference,
184 "within_acceptance": difference <= acceptance,
185 }
186 )
187 status = (
188 "converged_at_requested_threshold"
189 if comparisons and all(item["within_acceptance"] for item in comparisons)
190 else "not_converged_at_requested_threshold"
191 )
192 design = {
193 "varied": ["time_points", "rtol", "atol"],
194 "paired_randomness": None,
195 }
196 else:
197 counts = parse_csv_ints(
198 args.trajectory_counts,
199 name="trajectory_counts",
200 minimum=2,
201 maximum=MAX_TRAJECTORIES,
202 )
203 if len(counts) < 2:
204 raise CliError("Monte Carlo convergence requires at least two runs")
205 if sum(counts) > MAX_TOTAL_SWEEP_TRAJECTORIES:
206 raise CliError(
207 f"sum of trajectory_counts exceeds {MAX_TOTAL_SWEEP_TRAJECTORIES}"
208 )
209 time_points = int(args.mc_time_points)
210 if not 2 <= time_points <= MAX_TIME_POINTS:
211 raise CliError(
212 f"mc_time_points must be from 2 through {MAX_TIME_POINTS}"
213 )
214 rtol = finite_float(
215 args.mc_rtol,
216 name="mc_rtol",
217 minimum=1.0e-12,
218 maximum=1.0e-2,
219 )
220 for index, count in enumerate(counts):
221 config = _base_config(
222 solver="mcsolve",
223 args=args,
224 time_points=time_points,
225 trajectories=count,
226 rtol=rtol,
227 )
228 if config.decay_rate == 0.0 and config.dephasing_rate == 0.0:
229 raise CliError(
230 "monte-carlo mode requires a nonzero decay or dephasing rate"
231 )
232 report = run_simulation(config, qutip)
233 reports.append(report)
234 standard_error = report["trajectory_statistics"][
235 "excited_population_standard_error_estimate"
236 ]
237 run_summaries.append(
238 {
239 "index": index,
240 "trajectories_requested": count,
241 "trajectories_run": report["trajectory_statistics"][
242 "trajectories_run"
243 ],
244 "final_excited_population": float(
245 report["expectations"]["excited_population"][-1]
246 ),
247 "final_standard_error_estimate": float(standard_error[-1]),
248 "analytic_max_abs_error": report["analytic_reference"][
249 "max_abs_error"
250 ],
251 "seed_count": len(
252 report["trajectory_statistics"]["seeds"] or []
253 ),
254 }
255 )
256 reference_value = run_summaries[-1]["final_excited_population"]
257 comparisons = []
258 for summary in run_summaries[:-1]:
259 difference = abs(summary["final_excited_population"] - reference_value)
260 uncertainty_scale = max(
261 summary["final_standard_error_estimate"],
262 run_summaries[-1]["final_standard_error_estimate"],
263 )
264 comparisons.append(
265 {
266 "run_index": summary["index"],
267 "reference_run_index": len(run_summaries) - 1,
268 "absolute_final_population_difference": difference,
269 "acceptance_threshold": acceptance,
270 "within_absolute_acceptance": difference <= acceptance,
271 "difference_over_larger_standard_error": (
272 difference / uncertainty_scale
273 if uncertainty_scale > 0.0
274 else None
275 ),
276 }
277 )
278 status = "sampling_diagnostic_complete"
279 design = {
280 "varied": ["ntraj"],
281 "paired_randomness": (
282 "The same base seed is intentionally used so QuTiP spawns "
283 "comparable trajectory prefixes; these are not independent replicates."
284 ),
285 }
287 return {
288 "report_type": "qutip.convergence_sweep",
289 "schema_version": 1,
290 "qutip_version": QUTIP_VERSION,
291 "mode": args.mode,
292 "acceptance": acceptance,
293 "design": design,
294 "runs": run_summaries,
295 "comparisons": comparisons,
296 "status": status,
297 "limits": {
298 "maximum_runs": MAX_SWEEP_RUNS,
299 "maximum_total_trajectories": MAX_TOTAL_SWEEP_TRAJECTORIES,
300 "maximum_time_points_per_run": MAX_TIME_POINTS,
301 },
302 "interpretation": (
303 "A synthetic sweep demonstrates numerical checks; repeat convergence "
304 "for every cutoff and reported observable in the scientific model."
305 ),
306 }
309def build_parser() -> argparse.ArgumentParser:
310 parser = argparse.ArgumentParser(
311 description=(
312 "Run a bounded synthetic QuTiP convergence sweep over deterministic "
313 "tolerances/grids or Monte Carlo trajectory counts."
314 )
315 )
316 parser.add_argument(
317 "--mode",
318 choices=("deterministic", "monte-carlo"),
319 default="deterministic",
320 )
321 parser.add_argument("--omega", type=float, default=1.0)
322 parser.add_argument("--drive", type=float, default=0.0)
323 parser.add_argument("--decay-rate", type=float, default=0.2)
324 parser.add_argument("--dephasing-rate", type=float, default=0.0)
325 parser.add_argument("--t-final", type=float, default=10.0)
326 parser.add_argument("--method", choices=("adams", "bdf", "lsoda"), default="adams")
327 parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
328 parser.add_argument("--acceptance", type=float, default=1.0e-5)
329 parser.add_argument(
330 "--time-points",
331 default="101,201,401",
332 help="deterministic mode: increasing comma-separated output-grid sizes",
333 )
334 parser.add_argument(
335 "--rtols",
336 default="1e-5,1e-7,1e-9",
337 help="deterministic mode: decreasing comma-separated relative tolerances",
338 )
339 parser.add_argument(
340 "--trajectory-counts",
341 default="50,100,200",
342 help="Monte Carlo mode: increasing comma-separated trajectory counts",
343 )
344 parser.add_argument("--mc-time-points", type=int, default=101)
345 parser.add_argument("--mc-rtol", type=float, default=1.0e-8)
346 add_output_arguments(parser)
347 return parser
350def main(argv: list[str] | None = None) -> int:
351 args = build_parser().parse_args(argv)
352 report = run_sweep(args)
353 emit_json(report, output=args.output, force=args.force)
354 return 0
357if __name__ == "__main__":
358 raise SystemExit(run_cli(main))