# -*- coding: utf-8 -*-
import json
import logging
import os
from typing import Any
import numpy as np
from maxwellbloch import ob_atom
[docs]
class OBSolve(object):
"""Time-domain master equation solver for a single spatial point.
Wraps :func:`qutip.mesolve` to evolve an atomic density matrix described
by an :class:`~maxwellbloch.ob_atom.OBAtom` over a user-defined time
grid. Supports Doppler broadening via velocity classes and can save/load
results to avoid recomputation.
Args:
atom: Dict (or empty dict for defaults) describing the atomic system;
passed to :class:`~maxwellbloch.ob_atom.OBAtom`.
t_min: Start time of the simulation.
t_max: End time of the simulation.
t_steps: Number of time steps.
method: Solver method — ``'mesolve'`` (Lindblad master equation) is
the only supported value.
opts: Dict of QuTiP solver options passed to ``qutip.Options``.
savefile: Path prefix for caching solved results (``'.qu'`` extension
appended automatically). Pass ``None`` to disable caching.
"""
def __init__(
self,
atom: dict | None = None,
t_min: float = 0.0,
t_max: float = 1.0,
t_steps: int = 100,
method: str = "mesolve",
opts: dict | None = None,
savefile: str | None = None,
) -> None:
if atom is None:
atom = {}
self._build_atom(atom)
self._build_tlist(t_min=t_min, t_max=t_max, t_steps=t_steps)
self.method = method
self.build_opts(opts)
self._build_savefile(savefile)
def __repr__(self):
return (
"OBSolve(atom={0}, "
+ "t_min={1}, "
+ "t_max={2}, "
+ "t_steps={3}, "
+ "method={4}, "
+ "opts={5})"
).format(
self.atom, self.t_min, self.t_max, self.t_steps, self.method, self.opts
)
def _build_atom(self, atom_dict: dict) -> ob_atom.OBAtom:
self.atom = ob_atom.OBAtom(**atom_dict)
return self.atom
def _build_tlist(self, t_min: float, t_max: float, t_steps: int) -> np.ndarray:
from numpy import linspace
self.t_min = t_min
self.t_max = t_max
self.t_steps = t_steps
self.tlist = linspace(t_min, t_max, t_steps + 1)
return self.tlist
[docs]
def build_opts(self, opts: dict | None = None) -> dict:
"""Build the options dict to be passed into the QuTiP solver.
Any option available to the QuTiP solver is available here, we
provide defaults for solving the optical Bloch equations. See [0]
for details of all the available options.
Notes:
- For a stiff problem, it may help to set 'method' to
'bdf' instead of 'adams'.
- If the solver times out, try more 'nsteps', though
this will take longer.
- To speed up the solver, reduce atol and rtol, though this
will reduce accuracy.
Warning:
There is no validation checking here. If you pass in an option
which is not known to QuTiP it will throw an exception.
[0]: http://qutip.org/docs/4.2/guide/dynamics/dynamics-options.html
"""
self.opts = {
"atol": 1e-8,
"rtol": 1e-6,
"method": "adams",
"nsteps": 1000,
"max_step": self.t_step(),
}
if opts:
self.opts.update(opts) # Update any specified in the parameter
return self.opts
[docs]
def set_field_rabi_freq_t_func(self, field_idx: int, t_func: Any) -> None:
"""Set the Rabi frequency time function of a field to a new time
function. This is useful when you want to set a custom function,
not one available in t_funcs.py
Args:
field_idx: The field for which to set the Rabi frequency t_func
t_func: Rabi frequency as a function of time, f(t, args)
"""
self.atom.fields[field_idx].rabi_freq_t_func = t_func
self.atom.build_operators() # Rebuild so H_Omega is updated
[docs]
def set_field_rabi_freq_t_args(self, field_idx: int, t_args: dict) -> None:
"""Set the Rabi frequency time function arguments. To be used with
set_field_rabi_freq_t_func
Args:
field_idx: The field for which to set the Rabi frequency t_args
t_args: A dict representing the args to go with the t_func.
"""
self.atom.fields[field_idx].rabi_freq_t_args = t_args
[docs]
def obsolve(
self,
e_ops: list | None = None,
opts: dict | None = None,
recalc: bool = True,
show_pbar: bool = False,
save: bool = True,
) -> np.ndarray:
# When we're calling from MBSolve, we don't want to save each step.
# So we pass in save=False.
if save:
savefile = self.savefile
else:
savefile = None
# Choosing to overwrite opts at the solve stage.
if opts:
self.build_opts(opts)
options = self.opts
if self.method == "mesolve":
self.atom.mesolve(
self.tlist,
e_ops=e_ops,
options=options,
recalc=recalc,
savefile=savefile,
show_pbar=show_pbar,
)
return self.atom.states_t() # self.atom.result
[docs]
def states_t(self) -> np.ndarray:
return self.atom.states_t()
[docs]
def t_step(self) -> float:
return (self.t_max - self.t_min) / self.t_steps
def _build_savefile(self, savefile: str | None) -> None:
self.savefile = savefile
[docs]
def savefile_exists(self) -> bool:
"""Returns true if savefile (with appended extension .qu) exists."""
return os.path.isfile(str(self.savefile) + ".qu")
### JSON Serialise and Deserialise
[docs]
def get_json_dict(self) -> dict[str, Any]:
json_dict = {
"atom": self.atom.get_json_dict(),
"t_min": self.t_min,
"t_max": self.t_max,
"t_steps": self.t_steps,
"method": self.method,
"opts": self.opts,
}
return json_dict
[docs]
def to_json_str(self) -> str:
return json.dumps(self.get_json_dict(), sort_keys=True)
[docs]
def to_json(self, file_path: str) -> None:
with open(file_path, "w") as fp:
json.dump(
self.get_json_dict(), fp=fp, indent=2, separators=None, sort_keys=True
)
[docs]
@classmethod
def from_json_str(cls, json_str: str) -> "OBSolve":
"""Initialise OBSolve from JSON string."""
json_dict = json.loads(json_str)
return cls(**json_dict)
[docs]
@classmethod
def from_json(cls, file_path: str) -> "OBSolve":
"""Initialise OBSolve from JSON file."""
with open(file_path) as json_file:
json_dict = json.load(json_file)
# This needs to be here to get the savefile name from json
if "savefile" not in json_dict:
logging.debug("The savefile JSON element is missing.")
savefile = os.path.splitext(file_path)[0]
json_dict["savefile"] = savefile
elif not json_dict["savefile"]:
logging.debug("The savefile JSON element is null.")
return cls(**json_dict)