Skills › Scientific Agent Skills › qutip › two_level_simulation.py
two_level_simulation.py
402 lines · skills/qutip/scripts/two_level_simulation.py at 5f2bc03
1#!/usr/bin/env python3
2"""Run a bounded synthetic two-level open-system simulation."""
4from __future__ import annotations
6import argparse
7import math
8from dataclasses import asdict, dataclass
9from typing import Any
11from _common import (
12 DEFAULT_SEED,
13 MAX_ABS_FREQUENCY,
14 MAX_RATE,
15 MAX_TIME,
16 MAX_TIME_POINTS,
17 MAX_TRAJECTORIES,
18 QUTIP_VERSION,
19 CliError,
20 add_output_arguments,
21 bounded_int,
22 emit_json,
23 finite_float,
24 load_qutip,
25 run_cli,
26 to_jsonable,
27)
30METHODS = {"adams", "bdf", "lsoda", "dop853", "vern7", "vern9"}
31INITIAL_STATES = {"excited", "ground", "plus"}
34@dataclass(frozen=True)
35class SimulationConfig:
36 solver: str
37 initial_state: str
38 omega: float
39 drive: float
40 decay_rate: float
41 dephasing_rate: float
42 t_final: float
43 time_points: int
44 trajectories: int
45 seed: int
46 method: str
47 atol: float
48 rtol: float
50 @classmethod
51 def from_namespace(cls, args: argparse.Namespace) -> "SimulationConfig":
52 solver = str(args.solver)
53 if solver not in {"mesolve", "mcsolve"}:
54 raise CliError("solver must be mesolve or mcsolve")
55 initial_state = str(args.initial_state)
56 if initial_state not in INITIAL_STATES:
57 raise CliError(
58 f"initial_state must be one of: {', '.join(sorted(INITIAL_STATES))}"
59 )
60 method = str(args.method)
61 if method not in METHODS:
62 raise CliError(f"method must be one of: {', '.join(sorted(METHODS))}")
63 config = cls(
64 solver=solver,
65 initial_state=initial_state,
66 omega=finite_float(
67 args.omega,
68 name="omega",
69 minimum=-MAX_ABS_FREQUENCY,
70 maximum=MAX_ABS_FREQUENCY,
71 ),
72 drive=finite_float(
73 args.drive,
74 name="drive",
75 minimum=-MAX_ABS_FREQUENCY,
76 maximum=MAX_ABS_FREQUENCY,
77 ),
78 decay_rate=finite_float(
79 args.decay_rate,
80 name="decay_rate",
81 minimum=0.0,
82 maximum=MAX_RATE,
83 ),
84 dephasing_rate=finite_float(
85 args.dephasing_rate,
86 name="dephasing_rate",
87 minimum=0.0,
88 maximum=MAX_RATE,
89 ),
90 t_final=finite_float(
91 args.t_final,
92 name="t_final",
93 minimum=0.0,
94 maximum=MAX_TIME,
95 minimum_inclusive=False,
96 ),
97 time_points=bounded_int(
98 args.time_points,
99 name="time_points",
100 minimum=2,
101 maximum=MAX_TIME_POINTS,
102 ),
103 trajectories=bounded_int(
104 args.trajectories,
105 name="trajectories",
106 minimum=1,
107 maximum=MAX_TRAJECTORIES,
108 ),
109 seed=bounded_int(
110 args.seed,
111 name="seed",
112 minimum=0,
113 maximum=2**63 - 1,
114 ),
115 method=method,
116 atol=finite_float(
117 args.atol,
118 name="atol",
119 minimum=1.0e-14,
120 maximum=1.0e-2,
121 ),
122 rtol=finite_float(
123 args.rtol,
124 name="rtol",
125 minimum=1.0e-14,
126 maximum=1.0e-2,
127 ),
128 )
129 if config.atol > config.rtol:
130 raise CliError("atol must be no greater than rtol")
131 if (
132 config.solver == "mcsolve"
133 and config.decay_rate == 0.0
134 and config.dephasing_rate == 0.0
135 ):
136 raise CliError("mcsolve requires at least one nonzero collapse rate")
137 return config
140def _initial_state(qutip: Any, label: str) -> Any:
141 excited = qutip.basis(2, 0)
142 ground = qutip.basis(2, 1)
143 if label == "excited":
144 return excited
145 if label == "ground":
146 return ground
147 return (excited + ground).unit()
150def _state_audit(state: Any, tolerance: float) -> dict[str, Any]:
151 if state is None:
152 return {"available": False}
153 if state.isket:
154 norm = float(state.norm())
155 return {
156 "available": True,
157 "representation": "ket",
158 "norm": norm,
159 "normalization_error": abs(norm - 1.0),
160 "valid_within_tolerance": abs(norm - 1.0) <= tolerance,
161 }
162 trace = complex(state.tr())
163 hermitian = bool(state.isherm)
164 minimum_eigenvalue = (
165 min(float(value) for value in state.eigenenergies())
166 if hermitian
167 else None
168 )
169 valid = (
170 hermitian
171 and abs(trace - 1.0) <= tolerance
172 and minimum_eigenvalue is not None
173 and minimum_eigenvalue >= -tolerance
174 )
175 return {
176 "available": True,
177 "representation": "density_matrix",
178 "is_hermitian": hermitian,
179 "trace": trace,
180 "trace_error": abs(trace - 1.0),
181 "minimum_eigenvalue": minimum_eigenvalue,
182 "valid_within_tolerance": valid,
183 }
186def _seed_manifest(seeds: Any) -> list[Any] | None:
187 if seeds is None:
188 return None
189 manifest: list[Any] = []
190 for seed in seeds:
191 if isinstance(seed, int):
192 manifest.append(seed)
193 continue
194 manifest.append(
195 {
196 "entropy": to_jsonable(getattr(seed, "entropy", None)),
197 "spawn_key": list(getattr(seed, "spawn_key", ())),
198 }
199 )
200 return manifest
203def run_simulation(
204 config: SimulationConfig,
205 qutip_module: Any | None = None,
206) -> dict[str, Any]:
207 """Execute one bounded model and return a strict-JSON-ready report."""
209 qutip = qutip_module or load_qutip()
210 import numpy as np
212 times = np.linspace(0.0, config.t_final, config.time_points)
213 excited = qutip.basis(2, 0)
214 psi0 = _initial_state(qutip, config.initial_state)
215 H = 0.5 * config.omega * qutip.sigmaz() + 0.5 * config.drive * qutip.sigmax()
216 c_ops = []
217 if config.decay_rate > 0.0:
218 c_ops.append(np.sqrt(config.decay_rate) * qutip.sigmam())
219 if config.dephasing_rate > 0.0:
220 c_ops.append(
221 np.sqrt(config.dephasing_rate / 2.0) * qutip.sigmaz()
222 )
223 e_ops = [excited.proj(), qutip.sigmax(), qutip.sigmay(), qutip.sigmaz()]
224 options = {
225 "method": config.method,
226 "atol": config.atol,
227 "rtol": config.rtol,
228 "store_final_state": True,
229 "progress_bar": "",
230 }
232 if config.solver == "mesolve":
233 # Default output normalization can conceal norm/trace integration drift.
234 options["normalize_output"] = False
235 result = qutip.mesolve(
236 H,
237 psi0,
238 times,
239 c_ops=c_ops,
240 e_ops=e_ops,
241 options=options,
242 )
243 trajectories_run = None
244 trajectory_std = None
245 standard_error = None
246 seed_manifest = None
247 else:
248 result = qutip.mcsolve(
249 H,
250 psi0,
251 times,
252 c_ops,
253 e_ops=e_ops,
254 ntraj=config.trajectories,
255 seeds=config.seed,
256 options={**options, "keep_runs_results": False},
257 )
258 trajectories_run = int(
259 getattr(result, "num_trajectories", config.trajectories)
260 )
261 trajectory_std = np.asarray(result.std_expect[0], dtype=float)
262 standard_error = trajectory_std / math.sqrt(trajectories_run)
263 seed_manifest = _seed_manifest(getattr(result, "seeds", None))
265 population = np.asarray(result.expect[0], dtype=float)
266 sigma_x = np.asarray(result.expect[1], dtype=float)
267 sigma_y = np.asarray(result.expect[2], dtype=float)
268 sigma_z = np.asarray(result.expect[3], dtype=float)
270 initial_population = {
271 "excited": 1.0,
272 "ground": 0.0,
273 "plus": 0.5,
274 }[config.initial_state]
275 analytic_applicable = config.drive == 0.0
276 analytic_population = (
277 initial_population * np.exp(-config.decay_rate * times)
278 if analytic_applicable
279 else None
280 )
281 max_abs_analytic_error = (
282 float(np.max(np.abs(population - analytic_population)))
283 if analytic_population is not None
284 else None
285 )
286 lower_violation = max(0.0, float(-np.min(population)))
287 upper_violation = max(0.0, float(np.max(population) - 1.0))
288 probability_violation = max(lower_violation, upper_violation)
289 invariant_tolerance = max(10.0 * config.atol, 10.0 * config.rtol)
290 final_state_audit = _state_audit(
291 getattr(result, "final_state", None),
292 invariant_tolerance,
293 )
295 return {
296 "report_type": "qutip.two_level_simulation",
297 "schema_version": 1,
298 "qutip_version": QUTIP_VERSION,
299 "unit_convention": "hbar=1; angular-frequency and reciprocal-time units",
300 "configuration": asdict(config),
301 "model": {
302 "hilbert_dimension": 2,
303 "basis_labels": "excited=basis(2,0); ground=basis(2,1); fixed sigma_z basis, not dressed energy order",
304 "hamiltonian": "0.5 * omega * sigma_z + 0.5 * drive * sigma_x",
305 "amplitude_decay": "sqrt(decay_rate) * sigma_minus",
306 "pure_dephasing": "sqrt(dephasing_rate / 2) * sigma_z",
307 "initial_population": initial_population,
308 "assumptions": [
309 "finite two-level model",
310 "time-independent Hamiltonian",
311 "Markovian Lindblad channels",
312 "dephasing_rate denotes off-diagonal coherence decay",
313 ],
314 },
315 "times": times,
316 "expectations": {
317 "excited_population": population,
318 "sigma_x": sigma_x,
319 "sigma_y": sigma_y,
320 "sigma_z": sigma_z,
321 },
322 "analytic_reference": {
323 "applicable": analytic_applicable,
324 "excited_population": analytic_population,
325 "max_abs_error": max_abs_analytic_error,
326 "reason_if_not_applicable": (
327 None
328 if analytic_applicable
329 else "nonzero transverse drive changes the population equation"
330 ),
331 },
332 "trajectory_statistics": {
333 "uncertainty_note": (
334 "Independent-trajectory spread/sqrt(n); one trajectory or no "
335 "observed rare jumps cannot establish sampling precision."
336 if config.solver == "mcsolve" else None
337 ),
338 "trajectories_requested": (
339 config.trajectories if config.solver == "mcsolve" else None
340 ),
341 "trajectories_run": trajectories_run,
342 "seeds": seed_manifest,
343 "excited_population_std": trajectory_std,
344 "excited_population_standard_error_estimate": standard_error,
345 },
346 "checks": {
347 "time_grid_strictly_increasing": True,
348 "maximum_population_bound_violation": probability_violation,
349 "population_within_tolerance": probability_violation
350 <= invariant_tolerance,
351 "final_state": final_state_audit,
352 },
353 "solver": {
354 "name": getattr(result, "solver", config.solver),
355 "options_requested": options,
356 "stats": getattr(result, "stats", {}),
357 },
358 }
361def build_parser() -> argparse.ArgumentParser:
362 parser = argparse.ArgumentParser(
363 description=(
364 "Run a bounded two-level Lindblad or quantum-jump simulation and "
365 "emit strict JSON with analytic and physical checks."
366 )
367 )
368 parser.add_argument("--solver", choices=("mesolve", "mcsolve"), default="mesolve")
369 parser.add_argument(
370 "--initial-state",
371 choices=tuple(sorted(INITIAL_STATES)),
372 default="excited",
373 )
374 parser.add_argument("--omega", type=float, default=1.0)
375 parser.add_argument("--drive", type=float, default=0.0)
376 parser.add_argument("--decay-rate", type=float, default=0.2)
377 parser.add_argument("--dephasing-rate", type=float, default=0.0)
378 parser.add_argument("--t-final", type=float, default=10.0)
379 parser.add_argument("--time-points", type=int, default=201)
380 parser.add_argument("--trajectories", type=int, default=400)
381 parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
382 parser.add_argument("--method", choices=tuple(sorted(METHODS)), default="adams")
383 parser.add_argument("--atol", type=float, default=1.0e-10)
384 parser.add_argument("--rtol", type=float, default=1.0e-8)
385 add_output_arguments(parser)
386 return parser
389def main(argv: list[str] | None = None) -> int:
390 args = build_parser().parse_args(argv)
391 config = SimulationConfig.from_namespace(args)
392 report = run_simulation(config)
393 emit_json(report, output=args.output, force=args.force)
394 checks = report["checks"]
395 return 0 if (
396 checks["population_within_tolerance"]
397 and checks["final_state"].get("valid_within_tolerance", False)
398 ) else 1
401if __name__ == "__main__":
402 raise SystemExit(run_cli(main))