"""
Utilities for interacting with atomic parameters and ARC.
"""
from scipy.constants import epsilon_0, hbar, c
import numpy as np
import math
import re
from .sensor_utils import expand_statespec
from typing import Union, Optional, Literal, Tuple, List, Dict, Callable, NamedTuple, TYPE_CHECKING
from .exceptions import RydiquleError, AtomError
if TYPE_CHECKING:
import arc.alkali_atom_data
ATOMS = {
'H': 'Hydrogen',
'Li6': 'Lithium6', 'Li7': 'Lithium7',
'Na': 'Sodium',
'K39': 'Potassium39', 'K40': 'Potassium40', 'K41': 'Potassium41',
'Rb85': 'Rubidium85', 'Rb87': 'Rubidium87',
'Cs': 'Caesium'
}
"""
Alkali atoms defined by ARC that can be used with :class:`~.Cell`.
"""
def _load_arc_atom(atom_flag: str) -> 'arc.alkali_atom_data.AlkaliAtom':
"""
Function that lazy loads ARC atoms from :external+arc:module:`~arc.alkali_atom_data`
Returns
-------
arc.alkali_atom_data.AlkaliAtom
ARC alkali atom class associated with the provided atom_flag.
"""
import arc.alkali_atom_data as arc_atoms
if atom_flag not in ATOMS.keys():
raise AtomError(f"Atom flag must be one of {ATOMS.keys()}")
return getattr(arc_atoms, ATOMS[atom_flag])() # instantiate the class here
ground_n = {
"H": 1,
"Li": 2,
"Na": 3,
"K": 4,
"Rb": 5,
"Cs": 6
}
splitting_qnums: Dict[Optional[str], tuple] = {
None: (None, None, None),
"fs": ("all", None, None),
"hfs": (None, "all", "all")
}
# optional keys: (m_j, f, m_f)
# (n,l,j are required so not included here
STATE_TYPES = {
(False, False, False): "NLJ",
(True, False, False): "FS",
(False, True, True): "HFS"
}
QSpec = Union[float, List[float], Literal['all']]
[docs]
class A_QState(NamedTuple):
"""
Named tuple class designed to represent the quantum numbers a state spec of an alkali
atom. `n`, `l`, and `j` quantum numbers are required, with optional `m_j`, `f`, and `m_f`.
"""
n: int
l: int
j: float
m_j: Optional[QSpec] = None
f: Optional[QSpec] = None
m_f: Optional[QSpec] = None
def __str__(self):
"""Compact string representation of an `A_QState` that removes labels for
n, l, and j as well as any of m_j, f, and m_f that are `None`.
Also removes `"A_QState"` from the print output.
Returns
-------
str
String representation of state. Prunes redundant/bulky parts of a standard NamedTuple
output.
"""
return self.__repr__().replace("n=","").replace("l=","").replace("j=","",1)
def __repr__(self):
"""Overload of the standard `__repr__` function which removes "A_QState" from the front of
the string and removes all the `None` values. This doesn't change any of the values of the
state, just prunes the output string to a more printout-friendly format.
Returns
-------
str
Pruned representation of A_QState.
"""
return '(' + ', '.join(f'{name}={val!r}' for name, val in zip(self._fields, self) if val is not None) + ')'
@property
def qnums(self) -> Tuple[QSpec, ...]:
return tuple(i for i in self if i is not None)
@property
def stype(self) -> str:
try:
return STATE_TYPES[tuple(i is not None for i in self[3:])] # pyright: ignore[reportArgumentType]
except KeyError:
raise RydiquleError(f'{self} has unrecognized type')
[docs]
class QState(NamedTuple):
"""
Named tuple class designed to represent the quantum numbers in the state of an alkali
atom. `n`, `l`, and `j` quantum numbers are required, with optional `m_j`, `f`, and `m_f`.
"""
n: int
l: int
j: float
m_j: Optional[float] = None
f: Optional[int] = None
m_f: Optional[int] = None
def __str__(self):
"""Compact string representation of an `A_QState` that removes labels for
n, l, and j as well as any of m_j, f, and m_f that are `None`.
Also removes `"A_QState"` from the print output.
Returns
-------
str
String representation of state. Prunes redundant/bulky parts of a standard NamedTuple
output.
"""
return self.__repr__().replace("n=","").replace("l=","").replace("j=","",1)
def __repr__(self):
"""Overload of the standard `__repr__` function which removes "A_QState" from the front of
the string and removes all the `None` values. This doesn't change any of the values of the
state, just prunes the output string to a more printout-friendly format.
Returns
-------
str
Pruned representation of A_QState.
"""
return '(' + ', '.join(f'{name}={val!r}' for name, val in zip(self._fields, self) if val is not None) + ')'
@property
def qnums(self) -> Tuple[float, ...]:
"""Return a basic `tuple` representation of an `A_QState` with all `None` values removed.
Returns
-------
tuple
Quantum numbers which are not `None`.
"""
return tuple(i for i in self if i is not None)
@property
def stype(self) -> str:
"""Type of state. One of "NLJ", "FS", "HFS"
Returns
-------
str
String representing state type.
"""
try:
return STATE_TYPES[tuple(i is not None for i in self[3:])] # pyright: ignore[reportArgumentType]
except KeyError:
raise RydiquleError(f'{self} has unrecognized type')
#STATE_TYPES but inverted and including nlj numbers (we reference often enough to precompute)
QNUMS_IN_STATE_SPEC = {state_type: (True, True, True) + qnums for qnums, state_type in STATE_TYPES.items()}
[docs]
def ground_state(n: Union[int, str],
splitting:Literal[None, "fs", "hfs"]=None,
expand:bool=False) -> Union[A_QState, List[A_QState]]:
"""
Retrieve `A_QState` for the ground state of an atom or principle quantum number.
Optionally, include fine structure splitting or hyperfine splitting, which will include all
`m_j` or all `f` and `m_f` values respectively. By default, specifying splitting does not return
a list of states, but rather the associated specification with `"all"` in the appropriate
place. The `expand` keyword argument can be used to modify this behavior, returning a list.
Parameters
----------
n: int or str
Either the string flag of the atom or the principle quantum number n of
an atom. If string, must begin with ['H', 'Li', 'Na', 'K',
'Rb', 'Cs'].
splitting: {None, 'fs', 'hfs'}, optional
Type of splitting. Must be one of `None`, `"fs"`, or `"hfs"`, corresponding to the
inclusion of `(n,l,j)` only, `m_j`, or both `f` and `m_f` respectively.
expand: boolean, optional
For states with splitting, whether to return them as a list of all states. If `False`,
return is a single state specification with `"all"` for the appropriate quantum numbers.
If `True`, a list of all individual states is returned.
Raises
------
RydiquleError
If `n` is not a valid atom string or integer value
ValueError
If `splitting` is not one of {None, "fs", "hfs"}.
Returns
-------
A_QState:
`A_QState` corresponding to the ground state of the provided atom with the provided
splitting, or list of `A_QState`s if `expand` is `True`.
Examples
--------
The simplest use is to return the nlj quantum numbers for a particular atom's ground state.
Principle quantum number and string atom flags can be used interchangeably.
>>> atom = "Rb85"
>>> print(rq.ground_state(atom))
(5, 0, 0.5)
>>> print(rq.ground_state(5))
(5, 0, 0.5)
This function also can return states with splitting, either as a list of states or as a
manifold specification.
>>> print(rq.ground_state(atom, splitting="fs"))
(5, 0, 0.5, m_j='all')
>>> print(rq.ground_state(5, splitting="fs", expand=True))
[(n=5, l=0, j=0.5, m_j=-0.5), (n=5, l=0, j=0.5, m_j=0.5)]
>>> print(rq.ground_state(atom, splitting="hfs"))
(5, 0, 0.5, f='all', m_f='all')
"""
if isinstance(n, str):
# strip out isotope numbers, if any
atom = re.sub(r'\d+', '', n)
if atom in ground_n.keys():
n = ground_n[atom]
else:
raise RydiquleError(f"For string value of n, must use one of {list(ground_n.keys())}")
if not isinstance(n, int):
raise RydiquleError(f"n must be int or str, but found type {type(n)}.")
try:
g_state = A_QState(n, 0, 1/2, *splitting_qnums[splitting])
except KeyError:
raise ValueError(f"Invalid splitting {splitting}")
if expand:
return expand_qnums([g_state])
return g_state
[docs]
def D1_excited(n: Union[int, str],
splitting:Literal[None, "fs", "hfs"]=None,
expand:bool=False) -> Union[A_QState, List[A_QState]]:
"""
Retrieve `A_QState` of the excited state of the D1 line of an atom or principle quantum number.
Optionally, include fine structure splitting or hyperfine splitting, which will include all
`m_j` or all `f` and `m_f` values respectively. By default, specifying splitting does not return
a list of states, but rather the associated specification with `"all"` in the appropriate
place. The `expand` keyword argument can be used to modify this behavior, returning a list.
place.
Parameters
----------
n: int or str
Either the string flag of the atom or the principle quantum number n of
an atom. If string, must begin with ['H', 'Li', 'Na', 'K',
'Rb', 'Cs'].
splitting: {None, 'fs', 'hfs'}
Type of splitting. Must be one of `None`, `"fs"`, or `"hfs"`, corresponding to the
inclusion of `(n,l,j)` only, `m_j`, or both `f` and `m_f` respectively.
expand: boolean, optional
For states with splitting, whether to return them as a list of all states. If `False`,
return is a single state specification with `"all"` for the appropriate quantum numbers.
If `True`, a list of all individual states is returned.
Raises
------
RydiquleError
If `n` is not a valid atom string or integer value
ValueError
If `splitting` is not one of {None, "fs", "hfs"}.
Returns
-------
A_QState or list of A_QState:
`A_QState` corresponding to the D1 excited state of the provided atom with the provided
splitting, or list of `A_QState`s if `expand` is `True`.
Examples
--------
The simplest use is to return the nlj quantum numbers for a particular atom's excited state of
the D1 transition. Principle quantum number and string atom flags can be used interchangeably.
>>> atom = "Rb85"
>>> print(rq.D1_excited(atom))
(5, 1, 0.5)
>>> print(rq.D1_excited(5))
(5, 1, 0.5)
This function also can return states with splitting, either as a list of states or as a
manifold specification.
>>> print(rq.D1_excited(atom, splitting="fs"))
(5, 1, 0.5, m_j='all')
>>> print(rq.D1_excited(5, splitting="fs", expand=True))
[(n=5, l=1, j=0.5, m_j=-0.5), (n=5, l=1, j=0.5, m_j=0.5)]
>>> print(rq.D1_excited(atom, splitting="hfs"))
(5, 1, 0.5, f='all', m_f='all')
"""
if isinstance(n, str):
# strip out isotope numbers, if any
atom = re.sub(r'\d+', '', n)
if atom in ground_n.keys():
n = ground_n[atom]
else:
raise RydiquleError(f"For string value of n, must use one of {list(ground_n.keys())}")
if not isinstance(n, int):
raise RydiquleError(f"n must be int or str, but found type {type(n)}.")
try:
e_state = A_QState(n, 1, 1/2, *splitting_qnums[splitting])
except KeyError:
raise ValueError(f"Invalid splitting {splitting}")
if expand:
return expand_qnums([e_state])
return e_state
[docs]
def D2_excited(n: Union[int, str],
splitting:Literal[None, "fs", "hfs"]=None,
expand:bool=False) -> Union[A_QState, List[A_QState]]:
"""
Retrieve `A_QState` of the excited state of the D2 line of an atom or principle quantum number.
Optionally, include fine structure splitting or hyperfine splitting, which will include all
`m_j` or all `f` and `m_f` values respectively. By default, specifying splitting does not return
a list of states, but rather the associated specification with `"all"` in the appropriate
place. The `expand` keyword argument can be used to modify this behavior, returning a list.
Parameters
----------
n: int or str
Either the string flag of the atom or the principle quantum number n of
an atom. If string, must begin with ['H', 'Li', 'Na', 'K',
'Rb', 'Cs'].
splitting: {None, 'fs', 'hfs'}
Type of splitting. Must be one of `None`, `"fs"`, or `"hfs"`, corresponding to the
inclusion of `(n,l,j)` only, `m_j`, or both `f` and `m_f` respectively.
expand: boolean, optional
For states with splitting, whether to return them as a list of all states. If `False`,
return is a single state specification with `"all"` for the appropriate quantum numbers.
If `True`, a list of all individual states is returned.
Raises
------
RydiquleError
If `n` is not a valid atom string or integer value
ValueError
If `splitting` is not one of {None, "fs", "hfs"}.
Returns
-------
A_QState or list of A_QState:
`A_QState` corresponding to the D2 excited state of the provided atom with the provided
splitting, or list of `A_QState`s if `expand` is `True`.
Examples
--------
The simplest use is to return the nlj quantum numbers for a particular atom's excited state of
the D2 transition. Principle quantum number and string atom flags can be used interchangeably.
>>> atom = "Rb85"
>>> print(rq.D2_excited(atom))
(5, 1, 1.5)
>>> print(rq.D2_excited(5))
(5, 1, 1.5)
This function also can return states with splitting, either as a list of states or as a
manifold specification.
>>> print(rq.D2_excited(atom, splitting="fs"))
(5, 1, 1.5, m_j='all')
>>> print(rq.D2_excited(5, splitting="fs", expand=True))
[(n=5, l=1, j=1.5, m_j=-1.5),
(n=5, l=1, j=1.5, m_j=-0.5),
(n=5, l=1, j=1.5, m_j=0.5),
(n=5, l=1, j=1.5, m_j=1.5)]
>>> print(rq.D2_excited(atom, splitting="hfs"))
(5, 1, 1.5, f='all', m_f='all')
"""
if isinstance(n, str):
# strip out isotope numbers, if any
atom = re.sub(r'\d+', '', n)
if atom in ground_n.keys():
n = ground_n[atom]
else:
raise RydiquleError(f"For string value of n, must use one of {list(ground_n.keys())}")
if not isinstance(n, int):
raise RydiquleError(f"n must be int or str, but found type {type(n)}.")
try:
e_state = A_QState(n, 1, 3/2, *splitting_qnums[splitting])
except KeyError:
raise ValueError(f"Invalid splitting {splitting}")
if expand:
return expand_qnums([e_state])
return e_state
[docs]
def D1_states(n: Union[int, str],
splitting:Literal[None, "fs", "hfs"]=None,
g_splitting:Literal[None, "fs", "hfs"]=None,
e_splitting:Literal[None, "fs", "hfs"]=None,
expand:bool=False
) -> List[Union[A_QState, List[A_QState]]]:
"""Return the ground and excited states for the D1 line of a rydberg atom.
States are returned as a pair of `A_QStates` with the provided splitting according to
the :func:`~.atom_utils.rydberg_ground` and :func:`~.atom_utils.D1_excited` functions with
the provided splitting values passed through. When splitting is `None`, the `g_splitting`
and `e_splitting` values are passed to `rydberg_ground` and `D1_excited` respectively.
Otherwise, the value of `splitting` is passed to both and `g_splitting` and `e_splitting`
are ignored.
By default, specifying splitting does not return a list of all states, but rather the associated
specifications with `"all"` in the appropriate place. The `expand` keyword argument can be
used to modify this behavior, returning a full list.
Parameters
----------
n : int or str
Either the string flag of the atom or the principle quantum number n of
an atom. If string, must begin with ['H', 'Li', 'Na', 'K', 'Rb', 'Cs'].
splitting : None, "fs", "hfs, optional
Type of splitting for both states. Must be one of `None`, `"fs"`, or `"hfs"`,
corresponding to the inclusion of `(n,l,j)` only, `m_j`, or both `f` and `m_f`
respectively.
g_splitting : None, "fs", "hfs, optional
Type of splitting for both states. Must be one of `None`, `"fs"`, or `"hfs"`,
corresponding to the inclusion of `(n,l,j)` only, `m_j`, or both `f` and `m_f`
respectively. Ignored if `splitting` is specified.
e_splitting : None, "fs", "hfs, optional
Type of splitting for both states. Must be one of `None`, `"fs"`, or `"hfs"`,
corresponding to the inclusion of `(n,l,j)` only, `m_j`, or both `f` and `m_f`
respectively. Ignored if `splitting` is specified.
Returns
-------
list of A_QState
Ground and D1 excited state specifications of the provided atom or principal quantum number.
Examples
--------
The basic use of this function is to return the A_QStates associated with the states of the D1
transition of a particular Rydberg atom. String flags and principle quantum numbers can be used
interchangeably.
>>> atom = "Rb85"
>>> print(rq.D1_states(atom))
[(n=5, l=0, j=0.5), (n=5, l=1, j=0.5)]
>>> print(rq.D1_states(5))
[(n=5, l=0, j=0.5), (n=5, l=1, j=0.5)]
Furthermore, splitting can be specified either for each state individually, or just for one of
the states using the optional `splitting`, `g_splitting`, or `e_splitting` argument.
>>> print(rq.D1_states(5, splitting="fs"))
[(n=5, l=0, j=0.5, m_j='all'), (n=5, l=1, j=0.5, m_j='all')]
>>> print(rq.D1_states(5, splitting="fs", expand=True))
[(n=5, l=0, j=0.5, m_j=-0.5), (n=5, l=0, j=0.5, m_j=0.5), (n=5, l=1, j=0.5, m_j=-0.5), (n=5, l=1, j=0.5, m_j=0.5)]
>>> print(rq.D1_states(5, g_splitting="fs"))
[(n=5, l=0, j=0.5, m_j='all'), (n=5, l=1, j=0.5)]
"""
if splitting is not None:
g_splitting = splitting
e_splitting = splitting
g_states = ground_state(n, splitting=g_splitting, expand=expand)
e_states = D1_excited(n, splitting=e_splitting, expand=expand)
if expand:
# fully expanded list of all states
assert isinstance(g_states, list) and isinstance(e_states, list)
return [*g_states, *e_states]
else:
# 2-element list of specs
return [g_states, e_states]
[docs]
def D2_states(n: Union[int, str],
splitting:Literal[None, "fs", "hfs"]=None,
g_splitting:Literal[None, "fs", "hfs"]=None,
e_splitting:Literal[None, "fs", "hfs"]=None,
expand: bool=False) -> List[Union[A_QState, List[A_QState]]]:
"""Return the ground and excited states for the D1 line of a rydberg atom.
States are returned as a pair of `A_QStates` with the provided splitting according to
the :func:`~.atom_utils.rydberg_ground` and :func:`~.atom_utils.D2_excited` functions with
the provided splitting values passed through. When splitting is `None`, the `g_splitting`
and `e_splitting` values are passed to `rydberg_ground` and `D2_excited` respectively.
Otherwise, the value of `splitting` is passed to both and `g_splitting` and `e_splitting`
are ignored.
By default, specifying splitting does not return a list of all states, but rather the associated
specifications with `"all"` in the appropriate place. The `expand` keyword argument can be
used to modify this behavior, returning a full list.
Parameters
----------
n : int or str
Either the string flag of the atom or the principle quantum number n of
an atom. If string, must begin with ['H', 'Li', 'Na', 'K', 'Rb', 'Cs'].
splitting : None, "fs", "hfs, optional
Type of splitting for both states. Must be one of `None`, `"fs"`, or `"hfs"`,
corresponding to the inclusion of `(n,l,j)` only, `m_j`, or both `f` and `m_f`
respectively.
g_splitting : None, "fs", "hfs, optional
Type of splitting for both states. Must be one of `None`, `"fs"`, or `"hfs"`,
corresponding to the inclusion of `(n,l,j)` only, `m_j`, or both `f` and `m_f`
respectively. Ignored if `splitting` is specified.
e_splitting : None, "fs", "hfs, optional
Type of splitting for both states. Must be one of `None`, `"fs"`, or `"hfs"`,
corresponding to the inclusion of `(n,l,j)` only, `m_j`, or both `f` and `m_f`
respectively. Ignored if `splitting` is specified.
Returns
-------
list of A_QState
Ground and D2 excited state specifications of the provided atom or principal quantum number.
Examples
--------
The basic use of this function is to return the A_QStates associated with the states of the D2
transition of a particular Rydberg atom. String flags and principle quantum numbers can be used
interchangeably.
>>> atom = "Rb85"
>>> print(rq.D2_states(atom))
[(n=5, l=0, j=0.5), (n=5, l=1, j=1.5)]
>>> print(rq.D2_states(5))
[(n=5, l=0, j=0.5), (n=5, l=1, j=1.5)]
Furthermore, splitting can be specified either for each state individually, or just for one of
the states using the optional `splitting`, `g_splitting`, or `e_splitting` argument.
>>> print(rq.D1_states(5, splitting="fs"))
[(n=5, l=0, j=0.5, m_j='all'), (n=5, l=1, j=0.5, m_j='all')]
>>> print(rq.D2_states(5, splitting="fs", expand=True))
[(n=5, l=0, j=0.5, m_j=-0.5),
(n=5, l=0, j=0.5, m_j=0.5),
(n=5, l=1, j=1.5, m_j=-1.5),
(n=5, l=1, j=1.5, m_j=-0.5),
(n=5, l=1, j=1.5, m_j=0.5),
(n=5, l=1, j=1.5, m_j=1.5)]
>>> print(rq.D1_states(5, g_splitting="fs"))
[(n=5, l=0, j=0.5, m_j='all'), (n=5, l=1, j=0.5)]
"""
if splitting is not None:
g_splitting = splitting
e_splitting = splitting
g_states = ground_state(n, splitting=g_splitting, expand=expand)
e_states = D2_excited(n, splitting=e_splitting, expand=expand)
if expand:
# fully expanded list of all states
assert isinstance(g_states, list) and isinstance(e_states, list)
return [*g_states, *e_states]
else:
# 2-element list of specs
return [g_states, e_states]
[docs]
def calc_kappa(omega: float, dipole_moment: float, density: float) -> float:
"""
Calculates the kappa constant needed for observable calculations.
The value is computed with the following formula Eq. 5 of
Meyer et. al. PRA 104, 043103 (2021)
.. math::
\\kappa = \\frac{\\omega n \\mu^2}{2c \\epsilon_0 \\hbar}
Where :math:`\\omega` is the probing frequency, :math:`\\mu` is the dipole moment,
:math:`n` is atomic cloud density, :math:`c` is the speed of light, :math:`\\epsilon_0`
is the dielectric constant, and :math:`\\hbar` is the reduced Plank constant.
Parameters
----------
omega: float
Atomic transition frequency, in rad/s
dipole_moment: float
Dipole moment of the atomic transition, in C*m
density: float
The atomic number density, in m^(-3)
Returns
-------
float
The value of kappa, in (rad/s)/m
"""
kappa = (omega*density*dipole_moment**2)/(2.0*c*epsilon_0*hbar)
return kappa
[docs]
def calc_eta(omega: float, dipole_moment:float, beam_area: float) -> float:
"""
Calculates the eta constant needed from some experiment calculations
The value is computed with the following formula Eq. 7 of
Meyer et. al. PRA 104, 043103 (2021)
.. math::
\\eta = \\sqrt{\\frac{\\omega \\mu^2}{2 c \\epsilon_0 \\hbar A}}
Where :math:`\\omega` is the probing frequency, :math:`\\mu` is the dipole moment,
:math:`A` is the beam area, :math:`c` is the speed of light, :math:`\\epsilon_0`
is the dielectric constant, and :math:`\\hbar` is the reduced Plank constant.
Parameters
----------
omega: float
The atomic transition frequency, in rad/s
dipole_moment: float
The atomic transition dipole moment, in C*m
beam_area: float
The cross-sectional area of the beam, in m^2
Returns
-------
float
The value of eta, in root(Hz)
"""
eta = math.sqrt((omega*dipole_moment**2)/(2.0*c*epsilon_0*hbar*beam_area))
return eta
[docs]
def expand_qnums(qstates: List[A_QState], I: Optional[float] = None,
) -> List[A_QState]:
"""Expand all list-like A_QStates in a list.
List-like quantum numbers are defined either with a list of quantum numbers or the string
"all". In the "all" case, that quantum number will be expanded into all physically allowed
values of that quantum number given the preceding numbers.
Iterates through the list, expanding each A_QState specification into a list of all states
matching that specification. For each state specification in the list, quantum numbers are
expanded from left to right. The final list of A_QStates will respect the ordering of the
initial states by ordering the states corresponding to each specification by
n, l, j, m_j, f, and finally m_f
Parameters
----------
qstates : list of A_QState
List of atomic quantum states specifications to be expanded.
I : Union[float,None], optional
Nuclear spin for the isotope of the atom. Used to calculate the f quantum number
when relevant, by default None
Returns
-------
list of A_QState
List of all atomic states corresponding to all the specifications in the given list.
Notes
-----
..note::
While this function can expand arbitrary states, it should be noted that the resulting
lists of states can be quite long. If they are to be used as the states of a
:class:`~.Cell`, these long state lists can dramatically increase computation time, and
it is often worth ensuring that tracking hyperfine states individually is absolutely
necessary.
Examples
--------
A basic piece of functionality for this function is as a shorthand for all states in a
given manifold.
>>> D1_ground = A_QState(5,0,0.5, f="all")
>>> D1_excited = A_QState(5,0,0.5,f="all")
>>> #manifolds for rubidium 87 (I=3/2)
>>> print(rq.expand_qnums([D1_ground], I=3/2))
[(n=5, l=0, j=0.5, f=1.0), (n=5, l=0, j=0.5, f=2.0)]
>>> print(rq.expand_qnums([D1_excited], I=3/2))
[(n=5, l=0, j=0.5, f=1.0), (n=5, l=0, j=0.5, f=2.0)]
>>> #manifolds for rubidium 85 (I=5/2)
>>> print(rq.expand_qnums([D1_ground], I=5/2))
[(n=5, l=0, j=0.5, f=2.0), (n=5, l=0, j=0.5, f=3.0)]
>>> print(rq.expand_qnums([D1_excited], I=5/2))
[(n=5, l=0, j=0.5, f=2.0), (n=5, l=0, j=0.5, f=3.0)]
Note that while this function is capable of getting large numbers of states, the resulting
lists can be quite cumbersome, and be substantially slower if used in a calculation, especially
for high angular momentum states.
>>> state = A_QState(7, 2, 2.5, f="all", m_f="all")
>>> states_all = rq.expand_qnums([state], I=7/2)
>>> print(len(states_all))
48
"""
has_len = [any([hasattr(qn, "__len__") for qn in state])
for state in qstates]
if not any(has_len):
return qstates
else:
expanded_qnums = [
expand_single_qnum(state, I=I)
for state in qstates
]
list_qnums = sum(expanded_qnums, start=[])
return(expand_qnums(list_qnums, I=I))
[docs]
def validate_qnums(qstate:A_QState, I: Optional[float]=None):
"""Validate that the provided named_tuple is a valid rydberg atomic state
Parameters
----------
qstate : A_QState
Named tuple to check, should have fields `("n","l","j","m_j","f","m_f")`
I : Union[None,float], optional
Nuclear spin of the rydberg atom of which this is a state. If `None`, all f
values are invalid automatically. Defaults to `None`
Raises
------
ValueError
If the tuple representing the state does not have 6 elements
AssertionError
If the states of the state are not physically allowed
"""
#confirm qstate is the correct type of state
try:
(n,l,j,m_j,f,m_f) = (
qstate.n,
qstate.l,
qstate.j,
qstate.m_j,
qstate.f,
qstate.m_f
)
assert len(qstate) == 6
except (AttributeError, ValueError, AssertionError):
raise ValueError("Atomic states must be represented with a rq.A_QState namedtuple")
none_qnums = tuple(i is not None for i in qstate[3:])
if none_qnums not in STATE_TYPES:
raise ValueError(f"State {qstate} is not a valid combination of quantum numbers")
#validate (n,l) int, j half int
assert int(n)==n, f"invalid n quantum number {n}."
assert (int(l)==l) and (l < n), f"invalid l quantum number {l}."
assert j==l+1/2 or j==abs(l-1/2), f"invalid j quantum number {j}"
#test m_j, f, m_f are allowed values
if m_j is not None:
valid_mj = get_valid_mj(qstate, I=I)
assert m_j in valid_mj, f"m_j must be one of {valid_mj}"
if f is not None:
valid_f = get_valid_f(qstate, I=I)
assert f in valid_f, f"f must be one of {valid_f}"
if m_f is not None:
valid_mf = get_valid_mf(qstate, I=I)
assert m_f in valid_mf, f"m_f must be one of {valid_mf}"
[docs]
def expand_single_qnum(qstate: A_QState, I: Optional[float] = None, wildcard: str = "all"
) -> List[A_QState]:
"""Generates a list of all valid states given a particular quantum number to be expanded.
For a given `A_Qstate` spec with one or more tuple elements specified as either a list or
the "all" string, returns a list of all valid state specifications matching that state
specification with the first list or string element only expanded. If multiple elements of the
statespec are specified with a list or string, only the first one is expanded. This function
is intended as a helper function for a single quantum number, and is not designed to be
used at the top-level
The the case that the element to be expanded is a list, the list returned will have a single
state specification corresponding to each element of that list, and allowed quantum number
rules will not be enforced. In the case that the element to be expanded is the "all" string,
all valid values of that particular quantum number will be used. Note
that only the `m_j`, `f`, and `m_f` quantum numbers can be expanded in this way.
Parameters
----------
qstate : A_QState
NamedTuple with fields `(n, l, j, m_j, f, m_f)` representing the quantum numbers of
the state. Each, element must be either a float, list of floats, or the "all" string.
Only `m_j`, `f`, and `m_f` may be specified with a "all".
I : float, optional
The nuclear spin of the rydberg atom. Only used for calculations of valid `f` quantum
numbers, defaults to 0.0
Returns
-------
List of A_QState
List of all possible quantum states matching the provided specification. Only the
first list-like quantum number will be expanded.
Raises
------
RydiquleError
If there is a string specification besides "all" in the provided state
Examples
--------
A simple m_j expansion of the D1 states for Rubidium
>>> D1_g = A_QState(5,0,0.5, m_j="all")
>>> D1_e = A_QState(5,1,0.5, m_j="all")
>>> print(rq.atom_utils.expand_single_qnum(D1_g))
[(n=5, l=0, j=0.5, m_j=-0.5), (n=5, l=0, j=0.5, m_j=0.5)]
>>> print(rq.atom_utils.expand_single_qnum(D1_e))
[(n=5, l=1, j=0.5, m_j=-0.5), (n=5, l=1, j=0.5, m_j=0.5)]
We can also expand into nuclear spin-coupled (f) states. Note that in this case we must
provide the nuclear spin I to use this function on its own (in this case I=7/2 for Rb87)
>>> D1_g = A_QState(5,0,0.5, f="all")
>>> D1_e = A_QState(5,1,0.5, f="all")
>>> print(rq.atom_utils.expand_single_qnum(D1_g, I=7/2))
[(n=5, l=0, j=0.5, f=3.0), (n=5, l=0, j=0.5, f=4.0)]
>>> print(rq.atom_utils.expand_single_qnum(D1_e, I=7/2))
[(n=5, l=1, j=0.5, f=3.0), (n=5, l=1, j=0.5, f=4.0)]
While we can provide the "all" flag to expand to all states, we may want to use a specific
subset of states if, for example, selection rules limit the states at play. In this case,
one can pass a list for a given quantum number.
>>> state = A_QState(5, 2 ,2.5, m_j=[-2.5, -1.5, -0.5])
>>> print(rq.atom_utils.expand_single_qnum(state))
[(n=5, l=2, j=2.5, m_j=-2.5), (n=5, l=2, j=2.5, m_j=-1.5), (n=5, l=2, j=2.5, m_j=-0.5)]
"""
list_idxs = [i for i,qn in enumerate(qstate)
if isinstance(qn, (list, str))]
if len(list_idxs) == 0:
return [qstate]
idx = list_idxs[0]
if isinstance(qstate[idx], str):
if not qstate[idx] == "all":
raise RydiquleError("String must be 'all' to specify all possible quantum numbers.")
qstate = A_QState(*(qn if i != idx
else valid_qnum_fns[i](qstate,I=I)
for i, qn in enumerate(qstate)
))
return expand_statespec(qstate)
[docs]
def match_A_QState(qstate: A_QState, compare_list=[], I: Optional[float] =None
) -> List[A_QState]:
"""Function to return all states in a list matching the provided pattern.
States are considered a match for `qstate` if they are an element of the list returned
by calling the :func:`~.expand_qnums` on `qstate`.
Parameters
----------
qstate : A_QState
The state against which elements of the list are compared.
compare_list : list, optional
List of states to test. Any matching the pattern provided by `qstate` will be in the
returned list, by default []
I : float, optional
Nuclear spin I of the atom containing the provided states, by default None
Returns
-------
List of A_QState
List of all the elements matching the pattern defined by `qstate`
"""
all_qstates = expand_qnums([qstate], I=I)
return [state for state in all_qstates if state in compare_list]
#functions to compute valid quantum numbers for rydberg atoms. While not
#underscored, these functions are not strictly designed as user-facing
[docs]
def get_valid_j(state: A_QState, I:Optional[float]=None) -> List[float]:
"""Return the valid values of j for given other quantum numbers.
For a given quantum state with principal and orbital quantum numbers
:math:`(n,l)`, the valid values of j are given by
.. math:: j = |l - \\frac{1}{2}|, l + \\frac{1}{2}
Note that if both values are the same, a list of length 1 is returned.
"""
L_qnum = state[1]
if not isinstance(L_qnum, (int, float)):
raise RydiquleError(f"Invalid J quantum number type {type(L_qnum)}.")
return list(set(L_qnum + s for s in [-.5,.5]))
[docs]
def get_valid_mj(state: A_QState, I:Optional[float]=None) -> List[float]:
"""Return the valid values of m_J for given other quantum numbers.
For a given quantum state with principal, orbital, and total quantum numbers
:math:`(n,L,J)`, the valid values of m_J are given by
.. math:: m_J = -J, -J+1, -J+2, ... , J-2, J-1, J
"""
J_qnum = state[2]
if not isinstance(J_qnum, (int, float)):
raise RydiquleError(f"Invalid J quantum number type {type(J_qnum)}.")
return np.arange(-1*J_qnum,J_qnum + 1).tolist()
[docs]
def get_valid_f(state: A_QState, I: Optional[float]=None) -> List[float]:
"""Return the valid values of f for given other quantum numbers.
For a given quantum state with principal, orbital, and spin-orbit quantum numbers
:math:`(n,L,J)` and nuclear quantum number :math:`I`, the valid values of m_f are given by
.. math:: f = |I-J|, |I-J|+1, ..., I+J
"""
J_qnum=state[2]
if not isinstance(J_qnum, (int, float)) or not isinstance(I, (int, float)):
raise ValueError(f"Invalid I,J quantum number types {(type(I),type(J_qnum))}.")
return np.arange(abs(J_qnum - I), J_qnum + I + 1).tolist()
[docs]
def get_valid_mf(state: A_QState, I: Optional[float]=None) -> List[float]:
"""Return the valid values of m_f for given other quantum numbers.
For a given quantum state with principal, orbital, and total quantum numbers
:math:`(n,L,J,f)`, the valid values of m_f are given by
.. math:: m_f = -f, -f+1, -f+2, ... , f-2, f-1, f
"""
f_qnum = state[4]
if not isinstance(f_qnum, (float, int)):
raise RydiquleError(f"Invalid f quantum number type {type(f_qnum)}.")
return np.arange(-1*f_qnum,f_qnum + 1).tolist()
valid_qnum_fns: Dict[int, Callable] = {
2: get_valid_j,
3: get_valid_mj,
4: get_valid_f,
5: get_valid_mf
}