diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000..3de9f8d --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,95 @@ +cff-version: 1.2.0 +message: If you use this software, please cite it using the preferred article citation. +title: RydIQule +abstract: >- + The Rydberg Interactive Quantum module is a modeling library designed to simulate the response of Rydberg atoms to arbitrary input RF waveforms. + It functions as a general master equation solver for quantum systems based on the semi-classical density matrix method. +authors: + - family-names: Miller + given-names: Benjamin N + orcid: 'https://orcid.org/0000-0003-0017-1355' + - family-names: Meyer + given-names: David H + orcid: 'https://orcid.org/0000-0003-2452-2017' + - family-names: Virtanen + given-names: Teemu + - family-names: O'Brien + given-names: Christopher M + orcid: 'https://orcid.org/0000-0003-2974-0531' + - family-names: Cox + given-names: Kevin C + orcid: 'https://orcid.org/0000-0001-5049-3999' + +version: 1.1.0 +date-released: "2023-10-11" +license: Apache-2.0 +repository-code: "https://github.com/QTC-UMD/rydiqule" +url: "https://doi.org/10.1016/j.cpc.2023.108952" + +identifiers: + - description: Journal article describing the software + doi: 10.1016/j.cpc.2023.108952 + +preferred-citation: + type: article + title: "RydIQule: A Graph-based paradigm for modeling Rydberg and atomic sensors" + authors: + - family-names: Miller + given-names: Benjamin N + - family-names: Meyer + given-names: David H + orcid: 'https://orcid.org/0000-0003-2452-2017' + - family-names: Virtanen + given-names: Teemu + - family-names: O'Brien + given-names: Christopher M + - family-names: Cox + given-names: Kevin C + doi: "10.1016/j.cpc.2023.108952" + journal: "Computer Physics Communications" + month: 1 + year: 2024 + start: 108952 # page number + issue: 294 + +references: + - type: software + title: ARC (Alkali.ne Rydberg Calculator) + authors: + - family-names: Šibalić + given-names: Nikola + repository-code: "https://github.com/nikolasibalic" + url: "http://arc-alkali-rydberg-calculator.readthedocs.io/" + - type: software + title: NetworkX + repository-code: "https://github.com/networkx/networkx" + url: "https://networkx.org" + - type: software + title: NumPy + repository-code: "https://github.com/numpy/numpy" + url: "https://numpy.org" + - type: software + title: SciPy + repository-code: "https://github.com/scipy/scipy" + url: "https://scipy.org" + - type: software + title: numbakit-ode + authors: + - family-names: Grecco + given-names: Hernan E + repository-code: "https://github.com/hgrecco/numbakit-ode" + url: "https://numbakit-ode.readthedocs.io/en/latest/index.html" + - type: software + title: CyRK + authors: + - family-names: Renaud + given-names: Joe P + repository-code: "https://github.com/jrenaud90/CyRK" + - type: software + title: leveldiagram + authors: + - family-names: Meyer + given-names: David H + repository-code: "https://github.com/dihm/leveldiagram" + url: "https://leveldiagram.readthedocs.io/en/latest" + \ No newline at end of file diff --git a/README.md b/README.md index efe1ac2..478f4a0 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ rydiqule -The Rydberg Interactive Quantum module is a modelling library designed to simulate -the response of a Rydberg atoms to arbitrary input RF waveforms. +The Rydberg Interactive Quantum module is a modeling library designed to simulate +the response of Rydberg atoms to arbitrary input RF waveforms. It also functions as a general master equation solver based on the semi-classical density matrix method. [![PyPI](https://img.shields.io/pypi/v/rydiqule.svg)](https://pypi.org/project/rydiqule) @@ -9,6 +9,13 @@ It also functions as a general master equation solver based on the semi-classica [![License](https://img.shields.io/pypi/l/rydiqule.svg)](https://github.com/QTC-UMD/rydiqule/raw/main/LICENSE) [![Docs](https://readthedocs.org/projects/rydiqule/badge/?version=latest)](https://rydiqule.readthedocs.io/en/latest) +### Please cite as + +B. N Miller, D. H. Meyer, T. Virtanen, C. M O'Brien, and K. C. Cox, +RydIQule: A Graph-based paradigm for modeling Rydberg and atomic sensors, +*Computer Physics Communications*, **294**, 108952 (2024) +[https://doi.org/10.1016/j.cpc.2023.108952](https://doi.org/10.1016/j.cpc.2023.108952) + ## Installation Presently, installation is done via pip. @@ -43,8 +50,8 @@ Now use pip to install rydiqule and remaining dependencies. ```shell # for normal installation (rydiqule) ~/> pip install rydiqule -# for editable installation, so source can be modified locally -(rydiqule) ~/> pip install -e rydiqule +# for editable installation of cloned repo, so source can be modified locally +(rydiqule) ~/> pip install -e . ``` ### Pure pip installation @@ -55,9 +62,10 @@ pip install rydiqule ``` This command will use pip to install all necessary dependencies. -To install in an editable way (which allows edits of the source code), run: +To install in an editable way (which allows edits of the source code), +run the following from the root directory of the cloned repository: ```shell -pip install -e rydiqule +pip install -e . ``` ### Confirm installation @@ -70,8 +78,8 @@ Proper installation can be confirmed by executing the following commands in a py Rydiqule ================ -Rydiqule Version: 1.0.0 -Installation Path: C:\Users\naqsL\Miniconda3\envs\rydiqule\lib\site-packages\rydiqule +Rydiqule Version: 1.1.0 +Installation Path: ~\Miniconda3\envs\rydiqule\lib\site-packages\rydiqule Dependencies ================ @@ -81,7 +89,7 @@ SciPy Version: 1.10.1 Matplotlib Version: 3.7.1 ARC Version: 3.3.0 Python Version: 3.9.16 -Python Install Path: C:\Users\naqsL\Miniconda3\envs\rydiqule +Python Install Path: ~\Miniconda3\envs\rydiqule Platform Info: Windows (AMD64) CPU Count: 12 Total System Memory: 128 GB @@ -90,16 +98,17 @@ Total System Memory: 128 GB ### Updating an existing installation Upgrading an existing installation is simple. -Simply run the pip installation commands described above with the update flag. +Simply run the pip installation commands described above. +Optionally, include the update flag to greedily update dependencies as well. ```shell pip install -U rydiqule ``` This command will also install any new dependencies that are required. If using an editable install, simply replacing the files in the same directory is sufficient. -Though it is recommended to also run the appropriate pip update command as well. +Though it is recommended to also run the appropriate pip update command as well to capture updated dependencies. ```shell -pip install -U -e rydiqule +pip install -U -e . ``` ### Dependencies diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 084d6cd..1ebc5be 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -1,6 +1,30 @@ Changelog ========= +v1.2.0 +------ + +Improvements +++++++++++++ + +- Level diagrams now use `Sensor.get_rotating_frames` to provide better plotting of energy ordering of levels. +- Level diagrams now allow for optional control of plotting parameters by manually specifying `ld_kw` options on nodes and edges. +- Added the ability to specify energy level shifts (additional Hamiltonian digonal terms) not accounted for by the coupling infrastructure. + + +Bug Fixes ++++++++++ + +- `Sensor.make_real` now returns correct sized `const` array when ground is not removed. +- Many updates to type hints to improve their accuracy. + +Deprecations +++++++++++++ + +- Remove `Solution._variable_parameters` in favor of property checking the observable parameters. +- Renamed `Sensor.basis()` and `Solution.basis` to `Sensor.dm_basis()` and `Solution.dm_basis` + to disambiguate physical basis from computational basis. + v1.1.0 ------ diff --git a/docs/source/index.rst b/docs/source/index.rst index 80cb029..2f8c43b 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -25,6 +25,31 @@ For more details, see the :doc:`overview`. For detailed usage examples, see the :doc:`_intro_nbs/Introduction_To_Rydiqule/Introduction_To_Rydiqule` Jupyter notebook. +If you use rydiqule in your work, please cite as + +.. raw:: html + +
+ B. N. Miller, et. al., RydIQule: A Graph-based paradigm for modeling Rydberg and atomic sensors, Computer Physics Communications 294, 108952 (2024). arXiv:2307.15673. + +.. code-block:: bibtex + + @article{rydiqule_2024, + author = {Miller, B. N. and Meyer, D. H. and Virtanen, T. and O'Brien, C. M. and Cox, K. C.}, + title = {RydIQule: A Graph-based paradigm for modeling Rydberg and atomic sensors}, + journal = {Computer Physics Communications}, + volume = {294}, + pages = {108952}, + year = {2024}, + doi = {10.1016/j.cpc.2023.108952}, + url = {https://doi.org/10.1016/j.cpc.2023.108952}, + eprint = {https://doi.org/10.1016/j.cpc.2023.108952} + } + +.. raw:: html + +
+ .. toctree:: :maxdepth: 2 :hidden: diff --git a/docs/source/installation.rst b/docs/source/installation.rst index be34098..b08724f 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -40,8 +40,8 @@ Now use pip to install rydiqule and remaining dependencies. # for normal installation (rydiqule) ~/Rydiqule> pip install rydiqule - # for editable installation, so source can be modified locally - (rydiqule) ~/Rydiqule> pip install -e rydiqule + # for editable installation of cloned repo, so source can be modified locally + (rydiqule) ~/Rydiqule> pip install -e . Pure pip installation --------------------- @@ -54,11 +54,12 @@ To install normally, run: This command will use pip to install all necessary dependencies. -To install in an editable way (which allows edits of the source code), run: +To install in an editable way (which allows edits of the source code), +run the following from the root directory of the cloned repository: .. code-block:: shell - pip install -e rydiqule + pip install -e . Confirm installation -------------------- @@ -73,8 +74,8 @@ Proper installation can be confirmed by executing the following commands in a py Rydiqule ================ - Rydiqule Version: 1.0.0 - Installation Path: C:\Users\naqsL\Miniconda3\envs\rydiqule\lib\site-packages\rydiqule + Rydiqule Version: 1.1.0 + Installation Path: ~\Miniconda3\envs\rydiqule\lib\site-packages\rydiqule Dependencies ================ @@ -84,7 +85,7 @@ Proper installation can be confirmed by executing the following commands in a py Matplotlib Version: 3.7.1 ARC Version: 3.3.0 Python Version: 3.9.16 - Python Install Path: C:\Users\naqsL\Miniconda3\envs\rydiqule + Python Install Path: ~\Miniconda3\envs\rydiqule Platform Info: Windows (AMD64) CPU Count: 12 Total System Memory: 128 GB @@ -93,7 +94,8 @@ Updating an existing installation --------------------------------- Upgrading an existing installation is simple. -Simply run the pip installation commands described above with the update flag. +Simply run the pip installation commands described above. +Optionally, include the update flag to greedily update dependencies as well. .. code-block:: shell @@ -102,11 +104,11 @@ Simply run the pip installation commands described above with the update flag. This command will also install any new dependencies that are required. If using an editable install, simply replacing the files in the same directory is sufficient. -Though it is recommended to also run the appropriate pip update command as well. +Though it is recommended to also run the appropriate pip update command as well to capture updated depedencies. .. code-block:: shell - pip install -U -e rydiqule + pip install -U -e . Dependencies diff --git a/mypy.ini b/mypy.ini index 6cec3ff..a178035 100644 --- a/mypy.ini +++ b/mypy.ini @@ -7,6 +7,7 @@ warn_redundant_casts = True warn_unused_ignores = True warn_unreachable = True show_error_codes = True +allow_redefinition = True files = src/**/*.py plugins = numpy.typing.mypy_plugin @@ -32,7 +33,13 @@ ignore_missing_imports = True [mypy-psutil.*] ignore_missing_imports = True -# Module options +[mypy-CyRK.*] +ignore_missing_imports = True + +[mypy-networkx.*] +ignore_missing_imports = True -[mypy-rydiqule.energy_diagram] -ignore_errors = False +[mypy-leveldiagram.*] +ignore_missing_imports = True + +# Module options diff --git a/readthedocs.yaml b/readthedocs.yaml index 2355f6e..84a7a5a 100644 --- a/readthedocs.yaml +++ b/readthedocs.yaml @@ -8,7 +8,7 @@ version: 2 build: os: ubuntu-22.04 tools: - python: "3.8" + python: "3.11" # Build documentation in the docs/ directory with Sphinx sphinx: diff --git a/src/rydiqule/__init__.py b/src/rydiqule/__init__.py index 9a9cbc7..1420c04 100644 --- a/src/rydiqule/__init__.py +++ b/src/rydiqule/__init__.py @@ -15,4 +15,4 @@ from .slicing.slicing import compute_grid, matrix_slice, memory_size, get_slice_num, get_slice_num_t -__version__ = '1.1.0' +__version__ = '1.2.0' diff --git a/src/rydiqule/cell.py b/src/rydiqule/cell.py index 47ebe28..54b2a98 100644 --- a/src/rydiqule/cell.py +++ b/src/rydiqule/cell.py @@ -10,8 +10,9 @@ from scipy.constants import Boltzmann, e # rydiqule imports -from .sensor import Sensor, ScannableParameter +from .sensor import Sensor from .sensor_utils import scale_dipole +from .sensor_utils import ScannableParameter, QState, States from .atom_utils import ATOMS, calc_kappa, calc_eta from typing import Literal, Optional, Sequence, List, Tuple, Callable, Union @@ -19,8 +20,6 @@ a0 = scipy.constants.physical_constants["Bohr radius"][0] AtomFlags = Literal['H', 'Li6', 'Li7', 'Na', 'K39', 'K40', 'K41', 'Rb85', 'Rb87', 'Cs'] -QState = Sequence # TODO: consider using named tuples here -States = Tuple[int, ...] class Cell(Sensor): @@ -118,10 +117,10 @@ def __init__(self, atom_flag: AtomFlags, *atomic_states: QState, def set_experiment_values(self, probe_tuple: Tuple[int,int], probe_freq: float, - cell_length: float, - beam_area: float, - eta: float, - kappa: float): + kappa: float, + eta: Optional[float] = None, + cell_length: Optional[float] = None, + beam_area: Optional[float] = None): """`Sensor` specific method. Do not use with `Cell`. This function does not do anything as Cell automatically handles @@ -328,6 +327,37 @@ def kappa(self): return kappa + @kappa.setter + def kappa(self, value: float): + """Setter for the kappa attribute. + + Updates the self._kappa class attribute. + + Parameters + ---------- + value : float + The floating-point value to set as the eta parameter for the system. + """ + self._kappa = value + + + @kappa.deleter + def kappa(self): + """Setter for the kappa attribute. + + Removes the self._kappa class attribute. + + Raises + ------ + AttributeError: + If kappa has not been set. + """ + try: + del self._kappa + except AttributeError: + raise AttributeError("The \"kappa\" attribute has not been set") + + @property def eta(self): """Get the eta value for the system. @@ -368,22 +398,8 @@ def eta(self): eta = calc_eta(omega_rad, dipole_moment, self.beam_area) return eta - - - @kappa.setter - def kappa(self, value: float): - """Setter for the kappa attribute. - - Updates the self._kappa class attribute. - - Parameters - ---------- - value : float - The floating-point value to set as the eta parameter for the system. - """ - self._kappa = value - + @eta.setter def eta(self, value): """Setter for the eta attribute. @@ -398,23 +414,6 @@ def eta(self, value): self._eta = value - @kappa.deleter - def kappa(self): - """Setter for the kappa attribute. - - Removes the self._kappa class attribute. - - Raises - ------ - AttributeError: - If kappa has not been set. - """ - try: - del self._kappa - except AttributeError: - raise AttributeError("The \"kappa\" attribute has not been set") - - @eta.deleter def eta(self): """Deleter for the eta attribute. @@ -430,7 +429,8 @@ def eta(self): del self._eta except AttributeError: raise AttributeError("The \"eta\" attribute has not been set") - + + @property def probe_freq(self): """Get the probe transition frequency, in rad/s @@ -513,7 +513,7 @@ def add_coupling( self, states: States, rabi_frequency: Optional[ScannableParameter] = None, detuning: Optional[ScannableParameter] = None, transition_frequency: Optional[float] = None, - phase: Optional[ScannableParameter] = 0, kvec: Tuple[float,float,float] = (0,0,0), + phase: ScannableParameter = 0, kvec: Tuple[float,float,float] = (0,0,0), time_dependence: Optional[Callable[[float],float]] = None, label: Optional[str] = None, e_field: Optional[ScannableParameter] = None, beam_power: Optional[float] = None, beam_waist: Optional[float] = None, @@ -974,8 +974,10 @@ def get_qnums(self, state: Union[QState, int]): """ if isinstance(state, int): try: - state = self.couplings.nodes[state]["qnums"] + qnums = self.couplings.nodes[state]["qnums"] except KeyError: raise ValueError(f"Cell basis is smaller that state {state}") + else: + qnums = state - return self._validate_qnums(state) + return self._validate_qnums(qnums) diff --git a/src/rydiqule/doppler_utils.py b/src/rydiqule/doppler_utils.py index 5b9c65f..4433353 100644 --- a/src/rydiqule/doppler_utils.py +++ b/src/rydiqule/doppler_utils.py @@ -34,7 +34,7 @@ class DirectMethod(TypedDict, total=False): doppler_velocities: Union[np.ndarray, Sequence] -MeshMethod = Union[UniformMethod, SplitMethod, DirectMethod] +MeshMethod = Union[UniformMethod, IsoPopMethod, SplitMethod, DirectMethod] def get_doppler_equations(base_eoms: np.ndarray, diff --git a/src/rydiqule/experiments.py b/src/rydiqule/experiments.py index fc72f0e..1370456 100644 --- a/src/rydiqule/experiments.py +++ b/src/rydiqule/experiments.py @@ -8,7 +8,7 @@ import numpy as np import warnings -from typing import Tuple +from typing import Tuple, List def get_transmission_coef(*args, **kwargs): @@ -66,7 +66,7 @@ def get_snr(sensor: Sensor, phase_quadrature: bool = False, diff_nearest: bool = False, **kwargs - ) -> Tuple[np.ndarray, Tuple[np.ndarray, ...]]: + ) -> Tuple[np.ndarray, List[np.ndarray]]: """ Calculate a Sensor's signal-to-noise ratio in standard deviation, in a 1Hz bandwidth, @@ -161,7 +161,6 @@ def get_snr(sensor: Sensor, 'results may not be as expected.')) full_sols = solve_steady_state(sensor, **kwargs) - full_sols._variables_defined('cell_length', 'eta', 'kappa') rhos_ij = full_sols.rho_ij(*full_sols.probe_tuple) _ = full_sols.get_OD() diff --git a/src/rydiqule/rydiqule_utils.py b/src/rydiqule/rydiqule_utils.py index 71fac38..e1576b6 100644 --- a/src/rydiqule/rydiqule_utils.py +++ b/src/rydiqule/rydiqule_utils.py @@ -49,7 +49,9 @@ def about(obscure_paths: bool = True): """ home = Path.home() - rydiqule_install_path = Path(inspect.getsourcefile(rydiqule)).parent + install_path = inspect.getsourcefile(rydiqule) + assert install_path is not None + rydiqule_install_path = Path(install_path).parent try: ryd_path = '~' / rydiqule_install_path.relative_to(home) except ValueError: diff --git a/src/rydiqule/sensor.py b/src/rydiqule/sensor.py index d697f4b..ec8a8f4 100644 --- a/src/rydiqule/sensor.py +++ b/src/rydiqule/sensor.py @@ -3,25 +3,22 @@ """ import numpy as np -import networkx as nx # type: ignore +import networkx as nx import re import warnings import itertools from .sensor_utils import _combine_parameter_labels +from .sensor_utils import ScannableParameter, CouplingDict, State, States, TimeFunc -from typing import List, Tuple, Dict, Literal, Callable, Optional, Union - -ScannableParameter = Union[float, List[float], np.ndarray] - -CouplingDict = Dict -States = Tuple[int, ...] +from typing import List, Tuple, Dict, Literal, Callable, Optional, Union, Sequence, cast BASE_SCANNABLE_KEYS = ["detuning", "rabi_frequency", "phase", - "transition_frequency"] + "transition_frequency", + "e_shift"] """Reference list of all coherent coupling keys that support rydiqules stacking convention. Note that all decoherence keys (keys beginning with `gamma_`) are supported, but handled separately. """ @@ -60,7 +57,7 @@ class Sensor(): Must be specified when using :class:`Sensor`. Automatically calculated when using :class:`Cell`.""" - probe_tuple: Optional[Tuple[int, int]] = None + probe_tuple: Optional[States] = None """Coupling edge that corresponds to the probing field. Defaults to `(0,1)` in :class:`Cell`.""" @@ -73,29 +70,80 @@ class Sensor(): beam_area: Optional[float] = None """Cross-sectional area of the probing beam, in square meters.""" - - def __init__(self, basis_size: int, *couplings: CouplingDict) -> None: + def __init__(self, basis: Union[int, List[Union[int, str]]], + *couplings: CouplingDict) -> None: """ - Initializes the Sensor of the specified basis size. + Initializes the Sensor with the specified basis . + + Can be specified as either an integer number of states (which will automatically + label the states `[0,...,basis_size]`) or list of state labels. Parameters - ---- - basis_size : int - Number of states in the basis. + --------- + basis: int or list of int, str + The specification of the basis size and labelling for a new `Sensor`. Can be specified + by either a integer or a list. If specified as an integer `n`, the created `Sensor` + will have `n` states labelled as `0,...n`. In the case of a list, a number of states + equal to the length of the list will be created in the sensor, indexed by the integer + or string values of the nodes. *couplings : tuple(dict) Couplings dictionaries to pass to :meth:`~.add_couplings` on sensor construction. Raises ------ TypeError - If `basis_size` is not an integer. + If `basis` is not an integer or iterable. + TypeError + If any of the state label specifications of basis are the wrong type. + + Examples + -------- + Providing an integer will define a sensor with the given basis size, labelled with + ascending integers. + + >>> s = rq.Sensor(3) + >>> print(s.states) + [0, 1, 2] + + States can also be defined with a list of integers: + + >>> s = rq.Sensor([0, 1, 2]) + >>> print(s.states) + [0, 1, 2] + + States can also be strings + + >>> s = rq.Sensor(['g', 'e1', 'e2']) + >>> print(s.states) + ['g', 'e1', 'e2'] + + Using `None` in a list will default to using the integer correspoinding to the state: + + >>> s = rq.Sensor(['g', None, None]) + >>> print(s.states) + ['g', 1, 2] """ - if not isinstance(basis_size, int): - raise TypeError("basis_size must be an integer") - self.valid_parameters = BASE_EDGE_KEYS.copy() + #if its an int, expand to variables + if isinstance(basis, int): + basis = list(range(basis)) + + try: + valid_node = [isinstance(n, (str,int)) for n in basis] + except TypeError: + raise TypeError("'basis' must be an integer or iterable") + + #ensure node types are valid + if not all(valid_node): + raise ValueError("Nodes in \'basis\' must be integers or strings") + + #ensure unique labels + if len(basis) != len(set(basis)): + raise ValueError("All state labels must be unique") + + self.valid_parameters = BASE_EDGE_KEYS.copy() self.couplings = nx.DiGraph() - self.couplings.add_nodes_from(range(basis_size)) + self.couplings.add_nodes_from(basis) if len(couplings) > 0: self.add_couplings(*couplings) @@ -103,7 +151,6 @@ def __init__(self, basis_size: int, *couplings: CouplingDict) -> None: self._zipped_parameters: Dict = {} - def set_experiment_values(self, probe_tuple: Tuple[int,int], probe_freq: float, kappa: float, @@ -111,7 +158,7 @@ def set_experiment_values(self, probe_tuple: Tuple[int,int], cell_length: Optional[float] = None, beam_area: Optional[float] = None, ): - """Sets attributes needed for observable calculations + """Sets attributes needed for observable calculations. Parameters ---------- @@ -141,7 +188,7 @@ def set_experiment_values(self, probe_tuple: Tuple[int,int], @property def basis_size(self): - """Property to return the number of nodes on the Sensor graph + """Property to return the number of nodes on the Sensor graph. Returns ------- @@ -151,50 +198,77 @@ def basis_size(self): return len(self.couplings) - def label_states(self, labels: dict): - """Add labels to the node of the graph. Does not add functionality immediatly, just - makes it easier to track which nodes correspond to specific states as the number - of states increases. + @property + def states(self): + """Property which gets a list of labels for the sensor in the order defined in + :meth:`~.Sensor.__init__`. This is also the order corresponding the rows and columns + in the system Hamiltonian and decoherence matrix. + + Returns + ------- + list + List of states of the system defined the constructor, in the order corresponding to + rows and columns of the Hamiltonian. + """ + + return list(self.couplings.nodes()) + + + def add_energy_shift(self, state: State, shift: ScannableParameter): + """Add an energy shift to a state. + + First perfoms validation that the provided `state` is actually a node in the graph, then + adds the shift specified by `shift` to a self-loop edge keyed with `"e_shift"`. This value + will be added to the corresponding diagonal term when the hamiltonian is generated. If + the provided node Parameters ---------- - labels : dict - A dictionary defining which labels should be applied to which node. The keys of the - dictionary must be integer values corresponding to the numbered nodes of the sensor - graph, and the values must be string values. Specific nodes can be accessed directly - or through regex patterns in the :meth:`~.Sensor.states_with_label` function. + state : str or int + The label corresponding to the atomic state to which the shift will be added. + shift : float or list-like of float + The magnitude of the energy shift, in Mrad/s Raises ------ - ValueError - If `labels` cannot be interpreted as a dictionary. KeyError - If any of the keys in `labels` are not integer nodes of the graph. - ValueError - If any of the labels are not a string. + If the supplied `state` is not in the system. + """ + if not self.couplings.has_node(state): + raise KeyError(f"state {state} is not a node on the graph") + + self._remove_edge_data((state, state), kind="coherent") + self.couplings.add_edge(state, state, e_shift=shift) + + + def add_energy_shifts(self, shifts:dict): + """Add multiple energy shifts to different nodes. + + Shifts are specified with the `shifts` dictionary, which is keyed with states and + has values corresponding to the energy shift applied to the state in Mrad/s. Error + handling and validation is done with the :meth:`~.Sensor.add_energy_shift` function. + + Parameters + ---------- + shifts : dict + Dictionary keyed with states with values corresponding to the energy shift, in Mrad/s, + of the corresponding state. """ + try: - labels_items = labels.items() + shifts_items = shifts.items() except AttributeError: - raise ValueError("The labels argument must be a dictionary") + raise ValueError("Shifts parameters must be a dictionary-like object") - nodes_list = list(self.couplings.nodes()) + for state, shift in shifts_items: + self.add_energy_shift(state, shift) - for state, label in labels_items: - if state not in nodes_list: - raise KeyError(f"{state} is not a node of the sensor graph.") - if not isinstance(label, str): - raise ValueError("States can only be labelled with a string") - - self.couplings.nodes[state]["label"] = label - - def add_coupling( self, states: States, rabi_frequency: Optional[ScannableParameter] = None, detuning: Optional[ScannableParameter] = None, transition_frequency: Optional[float] = None, - phase: Optional[ScannableParameter] = 0, + phase: ScannableParameter = 0, kvec: Tuple[float,float,float] = (0,0,0), time_dependence: Optional[Callable[[float],float]] = None, label: Optional[str] = None, @@ -209,12 +283,11 @@ def add_coupling( Parameters ---------- - states : tuple of ints of length 2 - The pair of states of the sensor which the state couples - Must be a tuple of intergers of length 2, and the integers must - be unique indicies within the basis. - Tuple order indicates which state to has higher energy: - namely the second state is always assumed to have higher energy. + states : tuple of ints or strings of length 2 + The pair of states of the sensor which the state couples. Must be a tuple + of length 2, where each element is a string or integer corresponding to a state in the + Sensor as defined in the constructor. Tuple order indicates which state to has higher + energy; the second state is always assumed to have higher energy. rabi_frequency : float or complex, or list-like of float or complex The rabi frequency of the field being added. Defined in units of Mrad/s. List-like values will invoke Rydiqule's stacking convention when relevant quantities are calculated. @@ -223,9 +296,10 @@ def add_coupling( units of Mrad/s. List-like values will invoke Rydiqule's stacking convention when relevant quantities are calculated. If specified, the coupling is treated with the rotating-wave approximation rather than in the lab frame, and `transition_frequency` is ignored if present. + A positive number always indicates a blue detuning, and a negative number indicates a blue + detuning. transition_frequency : float or list-like of floats, optional - The transition frequency between a particular pair of states. - Must be a positive number. + The transition frequency between a particular pair of states. Must be a positive number. List-like values will invoke Rydiqule's stacking convention when relevant quantities are calculated. Only used directly in calculations if `detuning` is `None`, ignored otherwise. Note that on its own, it only defines the spacing between two energy levels and not the @@ -273,6 +347,9 @@ def add_coupling( >>> s = rq.Sensor(2) >>> s.add_coupling(states=(0,1), detuning=1, rabi_frequency=2) + >>> s = rq.Sensor(['g', 'e']) + >>> s.add_coupling(('g', 'e'), detuning=1, rabi_frequency=1) + >>> s = rq.Sensor(2) >>> s.add_coupling(states=(0,1), detuning=np.linspace(-10, 10, 101), rabi_frequency=2, label="laser") @@ -284,8 +361,8 @@ def add_coupling( >>> kp = 250*np.array([1,0,0]) >>> s.add_coupling(states=(0,1), detuning=1, rabi_frequency=2, kvec=kp) - """ # noqa - + """ + #ensure states are unique if states[0] == states[1]: raise ValueError(f'{states}: Coherent coupling must couple different states.') @@ -300,19 +377,19 @@ def add_coupling( " or \'transition_frequency\' for a coupling without the approximation," " but not both.") - if (transition_frequency is not None) and (transition_frequency < 0): - raise ValueError(f"{states}: \'transition_frequency\' must be positive.") - - if detuning is None and transition_frequency > 5000 and not suppress_rwa: - msg = (f"{states}: Not using the rotating wave approximation" - " for large transition frequencies can result in " - "prohibitively long computation times. Specify detuning to use " - "the rotating wave approximation or pass \"suppress_rwa_warn=True\" " - "to add_coupling() to suppress this warning.") - - with warnings.catch_warnings(): - warnings.simplefilter("always") - warnings.warn(msg, stacklevel=2) + if transition_frequency is not None: + if transition_frequency < 0: + raise ValueError(f"{states}: \'transition_frequency\' must be positive.") + elif transition_frequency > 5000 and not suppress_rwa: + msg = (f"{states}: Not using the rotating wave approximation" + " for large transition frequencies can result in " + "prohibitively long computation times. Specify detuning to use " + "the rotating wave approximation or pass \"suppress_rwa_warn=True\" " + "to add_coupling() to suppress this warning.") + + with warnings.catch_warnings(): + warnings.simplefilter("always") + warnings.warn(msg, stacklevel=2) self._add_coupling(**field_params, **extra_kwargs) @@ -397,7 +474,7 @@ def _add_coupling(self, states: States, **field_params) -> None: self.couplings.add_edge(*states, **field_params) - def zip_parameters(self, *parameters: str, zip_label:str=None) -> None: + def zip_parameters(self, *parameters: str, zip_label: Optional[str]=None): """ Define 2 scannable parameters as "zipped" so they are scanned in parallel. @@ -411,10 +488,11 @@ def zip_parameters(self, *parameters: str, zip_label:str=None) -> None: ---------- parameters : str Parameter labels to scan together. Parameter labels are strings of the form - `", "`, such as `"(0,1)_detuning"`. + `"_"`, such as `"(0,1)_detuning"`. Must be at least 2 labels to zip. Note that couplings are specified in the :meth:`~.Sensor.add_coupling` function. If unspecified in this function, the pair of states in the coupling cast to a string will be used. + zip_label : optional, str String label shorthand for the zipped parameters. The label for the axis of these @@ -441,12 +519,17 @@ def zip_parameters(self, *parameters: str, zip_label:str=None) -> None: .. note:: This function should be called last after all Sensor couplings and dephasings have been added. Changing a coupling that has already been zipped removes it from - the `self.zipped_parameters` list. + the `self.zipped_parameters` list. .. note:: Modifying the `Sensor.zipped_parameters` attribute directly can break some functionality and should be avoided. Use this function or :meth:`~.Sensor.unzip_parameters` instead. + .. note:: + When defining the zip strings for states labelled with strings, be sure to additional + `'` or `"` characters on either side of the labels, as demonstrated in the second + example below. + Examples -------- >>> s = rq.Sensor(3) @@ -457,6 +540,16 @@ def zip_parameters(self, *parameters: str, zip_label:str=None) -> None: >>> print(s._zipped_parameters) #NOT modifying directly {'demo_zip': ['(1,2)_detuning', 'probe_detuning']} + Make sure to add the appropriate additional string markings when the states are strings. + + >>> s = rq.Sensor(['g', 'e1', 'e2']) + >>> det = np.linspace(-1,1,11) + >>> s.add_coupling(states=('g','e1'), detuning=det, rabi_frequency=1, label="probe") + >>> s.add_coupling(states=('e1','e2'), detuning=det, rabi_frequency=1) + >>> s.zip_parameters("probe_detuning", "('e1','e2')_detuning", zip_label="demo_zip") + >>> print(s._zipped_parameters) #NOT modifying directly + {'demo_zip': ["('e1','e2')_detuning", 'probe_detuning']} + """ #give a dummy label if not provided if zip_label is None: @@ -589,13 +682,14 @@ def unzip_parameters(self, zip_label, verbose=True): print(f"No label matching {zip_label}, no action taken") - def add_decoherence(self, states: Tuple[int, ...], gamma: ScannableParameter, - label: Optional[str] = None) -> None: + def add_decoherence(self, states: States, gamma: ScannableParameter, + label: Optional[str] = None): """ Add decoherent coupling to the graph between two states. If `gamma` is list-like, the sensor will scan over the values, - solving the system for each different gamma. + solving the system for each different gamma, identically to the + scannable parameters in coherent couplings. Parameters ---------- @@ -659,7 +753,8 @@ def add_decoherence(self, states: Tuple[int, ...], gamma: ScannableParameter, def add_transit_broadening(self, gamma_transit: ScannableParameter, - repop: Dict[int, float] = {0: 1.0}) -> None: + repop: Union[None, Dict[State, float]]= None, + label: str = "transit") -> None: """ Adds transit broadening by adding a decoherence from each node to ground. @@ -679,7 +774,8 @@ def add_transit_broadening(self, gamma_transit: ScannableParameter, The keys represent the state labels. The values represent the fractional amount that goes to that state. If the sum of values does not equal 1, population will not be conserved. - Default is to repopulate everything into state 0. + Default is to repopulate everything into the ground state (either state 0 + or the first state in the basis passed to the :meth:`~.Sensor.__init__` method). Warns ----- @@ -697,7 +793,20 @@ def add_transit_broadening(self, gamma_transit: ScannableParameter, [0.1 0. 0. ] [0.1 0. 0. ]] - """ # noqa + >>> s = rq.Sensor(['g', 'e1', 'e2']) + >>> repop = {'g':0.75, 'e1': 0.25} + >>> s.add_transit_broadening(0.2, repop=repop) + >>> print(s.decoherence_matrix()) + [[0.15 0.05 0. ] + [0.15 0.05 0. ] + [0.15 0.05 0. ]] + + """ + + if repop is None: + ground_state = self.states[0] + repop = {ground_state: 1.0} + if isinstance(gamma_transit, list): # needed for multiplying branching ratios gamma_transit = np.array(gamma_transit) @@ -707,18 +816,18 @@ def add_transit_broadening(self, gamma_transit: ScannableParameter, ' Population will not be conserved.')) for t, br in repop.items(): - for i in range(self.basis_size): + for i in self.states: self.add_decoherence((i, t), gamma=gamma_transit*br, label="transit") if isinstance(gamma_transit, (list,np.ndarray)): # need to zip together all the transit rates axes = self.axis_labels() transit_axes = [s for s in axes if s.endswith('gamma_transit')] - self.zip_parameters(*transit_axes) + self.zip_parameters(*transit_axes, zip_label=label) def add_self_broadening(self, node: int, gamma: ScannableParameter, - label: Optional[str] = "self") -> None: + label: str = "self"): """ Specify self-broadening (such as collisional broadening) of a level. @@ -797,11 +906,23 @@ def decoherence_matrix(self) -> np.ndarray: >>> gamma = np.linspace(0,0.5, 11) >>> s.add_decoherence((1,0), gamma) >>> print(s.decoherence_matrix().shape) - (11,3,3) - - """ # noqa + (11,3,3) + + Defining decoherences between states labelled with string values works just like coherent couplings: + + >>> s = rq.Sensor(['g', 'e1', 'e2']) + >>> s.add_decoherence(('e1', 'g'), 0.1) + >>> s.add_decoherence(('e2', 'g'),0.1) + >>> s.decoherence_matrix() + array([[0. , 0. , 0. ], + [0.1, 0. , 0. ], + [0.1, 0. , 0. ]]) + + """ base_couplings = self.couplings.copy() + int_states = {state: i for (i, state) in enumerate(self.states)} + stack_shape = self._stack_shape() for states, param, arr in self.variable_parameters(apply_mesh=True): self.couplings.edges[states][param] = arr @@ -818,7 +939,8 @@ def decoherence_matrix(self) -> np.ndarray: for label in decoherence_labels: for states, f in self.couplings_with(label).items(): - idx = (..., *states) + states_n = tuple([int_states[s] for s in states]) + idx = (...,*states_n) gamma_matrix[idx] += f[label] self.couplings = base_couplings @@ -838,6 +960,8 @@ def axis_labels(self, collapse: bool=True, full_labels:bool=False) -> List[str]: will be combined into a single label, as this is how :meth:`~.Sensor.get_hamiltonian` treats these axes. + The ordering of axis labels is + Returns ------- list of str @@ -869,6 +993,20 @@ def axis_labels(self, collapse: bool=True, full_labels:bool=False) -> List[str]: (11, 11, 3, 3) ['blue_detuning', '(1, 2)_detuning'] + The ordering of labels may change if string state names are used. The ordering + is determined by the output of the :meth:`~.Sensor.variable_parameters` method. + See method documentation for more detail. + + >>> s = rq.Sensor(['g', 'e1', 'e2']) + >>> det = np.linspace(-10, 10, 11) + >>> blue = {"states":('g','e1'), "rabi_frequency":1, "detuning":det, "label":"blue"} + >>> red = {"states":('e1','e2'), "rabi_frequency":3, "detuning":det} + >>> s.add_couplings(blue, red) + >>> print(s.get_hamiltonian().shape) + >>> print(s.axis_labels()) + (11, 11, 3, 3) + ["('e1','e2')_detuning", 'blue_detuning'] + Zipping parameters combines their corresponding labels, since their Hamiltonians now lie on a single axis of the stack. Here the axis of length 7 (axis 0) corresponds to the rabi frequencies and the axis of shape 11 (axis 1) corresponds to the zipped detunings @@ -884,7 +1022,7 @@ def axis_labels(self, collapse: bool=True, full_labels:bool=False) -> List[str]: ['(0,1)_rabi_frequency', 'detunings'] ['(0,1)_rabi_frequency', '(0,1)_detuning|(1,2)_detuning'] - """ # noqa + """ key_labels = [] item_labels = [] @@ -914,14 +1052,18 @@ def axis_labels(self, collapse: bool=True, full_labels:bool=False) -> List[str]: else: return key_labels - def variable_parameters(self, apply_mesh:bool = False): + + def variable_parameters(self, apply_mesh:bool = False + ) -> List[Tuple[States, str, np.ndarray]]: """ Property to retrieve the values of parameters that were stored on the graph as arrays. Values are returned as a list of tuples in the standard order of pythons default sorting, applied first to the tuple indicating states and then to the key of the parameter itself. This means that couplings are sorted first by lower state, then by upper state, then - alphabetically by the name of the parameter. + alphabetically by the name of the parameter.To determine order, all state labels treated + as their integer position in the basis as determined by ordering in the constructor + :meth:`~.Sensor.__init__`. Returns ------- @@ -942,10 +1084,30 @@ def variable_parameters(self, apply_mesh:bool = False): (0, 1): detuning=[-1. 0.5 2. ] (0, 1): rabi_frequency=[-1. 0.5 2. ] (1, 2): rabi_frequency=[-1. 0.5 2. ] + + The order is important; in the unzipped case, it will sort as though all state labels + were cast to strings, meaning integers will always be treated as first. + >>> s = rq.Sensor([None, 'e1', 'e2']) + >>> det1 = np.linspace(-1, 1, 3) + >>> det2 = np.linspace(-1, 1, 5) + >>> blue = {"states":(0,'e1'), "rabi_frequency":1, "detuning":det1} + >>> red = {"states":('e1','e2'), "rabi_frequency":3, "detuning":det2} + >>> s.add_couplings(blue, red) + >>> for states, key, value in s.variable_parameters(): + ... print(f"{states}: {key}={value}") + >>> print(f"Axis Labels: {s.axis_labels()}") + ('g', 1): detuning=[-1. 0. 1.] + (1, 2): detuning=[-1. -0.5 0. 0.5 1. ] + Axis Labels: ["('g',1)_detuning", '(1,2)_detuning'] + """ l=[] - for states in sorted(self.couplings.edges): + #function to compare mixed tuples + def state_order(values): + return [self.states.index(x) for x in values] + + for states in sorted(self.couplings.edges, key=state_order): edge_data = self.couplings.edges.get(states) @@ -954,7 +1116,7 @@ def variable_parameters(self, apply_mesh:bool = False): continue if isinstance(value, (list,np.ndarray)): - l.append((states, key, value)) + l.append((states, key, np.asarray(value))) if apply_mesh: vals = [val for _,_,val in l] @@ -965,16 +1127,19 @@ def variable_parameters(self, apply_mesh:bool = False): return l - def get_parameter_mesh(self) -> Tuple[np.ndarray, ...]: + def get_parameter_mesh(self) -> List[np.ndarray]: """ Returns the parameter mesh of the sensor. The parameter mesh is the flattened grid of variable parameters in all the couplings of a sensor. - Wraps `numpy.meshgrid` by passing arguments as the variable parameters in a - sensor. The `indexing` argument is always `"ij"` for matrix indexing. - For full documention of all arguments, see `numpy.meshgrid` documentation. - Argument documentation copied from `numpy.meshgrid`. + Wraps `numpy.meshgrid` with the `indexing` argument + always `"ij"` for matrix indexing. + + Returns + ------- + list of numpy.ndarray + list of mesh grids for every variable parameter Examples -------- @@ -1013,7 +1178,11 @@ def get_hamiltonian(self) -> np.ndarray: If any parameters have been zipped with the :meth:`~.Sensor._zip_parameters` method, those parameters will share an axis in the final hamiltonian stack. In this case, if axis N1 and N2 above are the same shape and zipped, the final - Hamiltonian will be of shape `(N1,...,Nm, n, n)` + Hamiltonian will be of shape `(N1,...,Nm, n, n)`. + + In the case where the basis of the `Sensor` was explicitly defined with a list + of states, the ordering of rows and coulumns in the hamiltonian corresponds to the + ordering of states passed in the basis. See rydiqule's conventions for matrix stacking for more details. @@ -1054,9 +1223,23 @@ def get_hamiltonian(self) -> np.ndarray: >>> print(H.shape) (11, 3, 3) + If the basis is provided as a list of string labels, the ordering of Hamiltonian rows + And columns will correspond to the order of states provided. + + >>> s = rq.Sensor(['g', 'e1', 'e2']) + >>> s.add_coupling(('g', 'e1'), detuning=1, rabi_frequency=1) + >>> s.add_coupling(('e1', 'e2'), detuning=1.5, rabi_frequency=1) + >>> print(s.get_hamiltonian()) + [[ 0. +0.j 0.5+0.j 0. +0.j] + [ 0.5-0.j -1. +0.j 0.5+0.j] + [ 0. +0.j 0.5-0.j -2.5+0.j]] + """ base_couplings = self.couplings.copy() + #dictionary of state ordering used to determine indeces in ham + int_states = {state: i for (i, state) in enumerate(self.states)} + stack_shape = self._stack_shape(time_dependence='steady') for states, param, arr in self.variable_parameters(apply_mesh=True): @@ -1075,14 +1258,24 @@ def get_hamiltonian(self) -> np.ndarray: if 'rabi_frequency' not in f: continue - idx = (...,*states) - conj_idx = (...,*states[::-1]) + #convert the state label to an index of position in ham + states_n = tuple([int_states[s] for s in states]) + idx = (...,*states_n) + conj_idx = (...,*states_n[::-1]) # factor of 1/2 accounts for implicit rotating wave approximation hamiltonian[idx] = (f['rabi_frequency']*np.cos(f['phase']) + 1j*f['rabi_frequency']*np.sin(f['phase']))/2 hamiltonian[conj_idx] = np.conj(hamiltonian[idx]) + #add the numerical diagonal shifts + for state1, state2, shift in nx.selfloop_edges(self.couplings, data='e_shift', default=0): + + shift_state = int_states[state1] + idx = (..., shift_state, shift_state) + + hamiltonian[idx] = shift + self.couplings = base_couplings return hamiltonian @@ -1132,7 +1325,7 @@ def _collapse_mesh(self, mesh): return mesh_copy - def get_time_hamiltonians(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + def get_time_hamiltonians(self) -> Tuple[np.ndarray, List[np.ndarray], List[np.ndarray]]: """ Get the hamiltonians for the time solver. @@ -1141,6 +1334,10 @@ def get_time_hamiltonians(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: The time dependent hamiltonians give 2 terms, the hamiltonian corresponding to the real part of the coupling and the hamiltonian corresponding to the imaginary part. + In the case where the basis of the `Sensor` was explicitly defined with a list + of states, the ordering of rows and coulumns in the hamiltonian corresponds to the + ordering of states passed in the basis. + Returns ------- hamiltonian_base : np.ndarray @@ -1176,6 +1373,22 @@ def get_time_hamiltonians(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: [array([[0.+0.j , 0.+0.5j], [0.-0.5j, 0.+0.j ]])] + If the basis is passed as a list, rows and columns are in the order specified: + + >>> s = rq.Sensor(['g', 'e']) + >>> step = lambda t: 0. if t<1 else 1. + >>> s.add_coupling(states=('g','e'), detuning=1, rabi_frequency=1, time_dependence=step) + >>> H_base, H_time_real, H_time_imaginary = s.get_time_hamiltonians() + >>> print(H_base) + >>> print(H_time_real) + >>> print(H_time_imaginary) + [[ 0.+0.j 0.+0.j] + [ 0.+0.j -1.+0.j]] + [array([[0. +0.j, 0.5+0.j], + [0.5+0.j, 0. +0.j]])] + [array([[0.+0.j , 0.+0.5j], + [0.-0.5j, 0.+0.j ]])] + """ hamiltonian_base = self.get_hamiltonian() dipole_matrices = self.get_time_couplings() @@ -1183,7 +1396,7 @@ def get_time_hamiltonians(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: return hamiltonian_base, *dipole_matrices - def get_time_couplings(self) -> Tuple[np.ndarray, np.ndarray]: + def get_time_couplings(self) -> Tuple[List[np.ndarray], List[np.ndarray]]: """ Returns the list of matrices of all couplings in the system defined with a `time_dependence` key. @@ -1250,12 +1463,15 @@ def get_time_couplings(self) -> Tuple[np.ndarray, np.ndarray]: (11, 1, 3, 3) (11, 1, 3, 3) - """ # noqa + """ #save to re-add later base_couplings = self.couplings.copy() stack_shape = self._stack_shape(time_dependence='time') + #dictionary of state ordering used to determine indeces in ham + int_states = {state: i for (i, state) in enumerate(self.states)} + for states, param, arr in self.variable_parameters(apply_mesh=True): self.couplings.edges[states][param] = arr @@ -1273,8 +1489,10 @@ def get_time_couplings(self) -> Tuple[np.ndarray, np.ndarray]: time_hamiltonian = np.zeros(hamiltonian_shape, dtype='complex') time_hamiltonian_i = np.zeros(hamiltonian_shape, dtype='complex') - idx = (...,*states) - conj_idx = (...,*states[::-1]) + #convert state label to an index of position in the ham + states_n = tuple([int_states[s] for s in states]) + idx = (...,*states_n) + conj_idx = (...,*states_n[::-1]) #ignores numpy casting to real warning that we already account for with warnings.catch_warnings(): @@ -1300,7 +1518,7 @@ def get_time_couplings(self) -> Tuple[np.ndarray, np.ndarray]: return matrix_list, matrix_list_i - def get_time_dependence(self) -> List[Callable[[float], float]]: + def get_time_dependence(self) -> List[TimeFunc]: """ Function which returns a list of the `time_dependence` functions. @@ -1324,7 +1542,7 @@ def get_time_dependence(self) -> List[Callable[[float], float]]: >>> print(s.get_time_dependence()) [ at 0x7fb310edd9d0>, at 0x7fb37c0c81f0>] - """ # noqa + """ time_dependence = [] for (i,j),f in self.couplings_with("time_dependence").items(): time_dependence.append(f['time_dependence']) @@ -1370,20 +1588,22 @@ def get_hamiltonian_diagonal(self, values: dict, no_stack: bool=False) -> np.nda dtype=np.complex128) subgraphs = self.get_rotating_frames() + int_states = {state: i for (i, state) in enumerate(self.states)} #ref list of states/idxs - for g, paths in subgraphs.items(): + for paths in subgraphs.values(): - for i, path in paths.items(): + for base_node, path in paths.items(): + # print(base_node, path) term = 0 for j in range(1, len(path)): - + n, sign = path[j] + n_prev, _ = path[j-1] # get the jth couplings along the path # remove frame signs - field = tuple(np.abs(path[j-1:j+1])) + field = (n_prev, n) # get the sign from the rotating frame - sign = np.sign(path[j]) if sign < 0: # Since it is getting an existing edge from the undirected graph, @@ -1392,7 +1612,7 @@ def get_hamiltonian_diagonal(self, values: dict, no_stack: bool=False) -> np.nda # sum to the cumulative term along the path from ground term = term + values.get(field,0)*sign - + i = int_states[base_node] diag[..., i] = term return diag @@ -1425,22 +1645,29 @@ def get_rotating_frames(self) -> dict: coherent_graph = self.couplings.edge_subgraph(coherent_edges) connected_levels = nx.weakly_connected_components(coherent_graph) - subgraphs = {coherent_graph.subgraph(ls):None for ls in connected_levels} + subgraphs: dict = {coherent_graph.subgraph(ls):{} for ls in connected_levels} for g in subgraphs: # min sets lowest state in graph as "ground" - paths = nx.shortest_path(nx.to_undirected(g), source=min(g.nodes)) + source_node = list(g.nodes)[0] + paths = nx.shortest_path(nx.to_undirected(g), source=source_node) - for i, path in paths.items(): + path_and_sign = {} + for node, path in paths.items(): + + path_sign = [1 for _ in path] for j in range(1, len(path)): # get the jth couplings along the path and assume the sign as positive - field = tuple(path[j-1:j+1]) + field = (path[j-1], path[j]) if field not in coherent_graph.edges: # switch the sign if the arrow points in the opposite direction # This corresponds to moving to a lower energy state - path[j] *= -1 + path_sign[j] = -1 + # print("path", path) + # print("sign", path_sign) + path_and_sign[node] = [ps for ps in zip(path, path_sign)] - subgraphs[g] = paths + subgraphs[g] = path_and_sign return subgraphs @@ -1508,7 +1735,7 @@ def get_value_dictionary(self, key: str) -> dict: return {states:params[key] for states, params in couplings_with_key.items()} - def set_gamma_matrix(self, gamma_matrix: np.ndarray) -> None: + def set_gamma_matrix(self, gamma_matrix: np.ndarray): """ Set the decoherence matrix for the system. @@ -1556,7 +1783,7 @@ def set_gamma_matrix(self, gamma_matrix: np.ndarray) -> None: f'must be {(self.basis_size,self.basis_size)}')) for states, gamma in np.ndenumerate(gamma_matrix): - + states = cast(Tuple[int, int], states) # cast for mypy #remove existing decoherence data if self.couplings.has_edge(*states): @@ -1662,7 +1889,7 @@ def couplings_with(self, *keys: str, >>> s.set_gamma_matrix(gamma) >>> print(s.couplings_with("detuning")) {(0, 1): {'rabi_frequency': 2, 'detuning': 1, 'phase': 0, 'kvec': (0, 0, 0), 'no_rwa_warning': False, 'label': '(0,1)'}} - """ # noqa + """ def notAll(x): return not all(x) @@ -1677,7 +1904,7 @@ def notAny(x): } - def states_with_label(self, label): + def states_with_label(self, label: str) -> List[State]: """Return a dict of all states with a label matching a given regular expression (regex) pattern. The dictionary will be consist of keys which are string labels applied to states with the :meth:`~.Sensor.label_states` function, and values which are the corresponding @@ -1692,9 +1919,8 @@ def states_with_label(self, label): Returns ------- - dict - Dictionary consisting of string labels for keys and the corresponding integer node - numbers as the values. + list + List of all labels of states in the sensor which match the provided regex pattern. Raises ------ @@ -1708,7 +1934,7 @@ def states_with_label(self, label): >>> s.add_coupling((1,2), detuning=2, rabi_requency=2) >>> s.label_states({0:"g", 1:"e1", 2:"e2"}) >>> print(s.states_with_label("e[12]")) - {'e1': 1, 'e2': 2} + ['e1', 'e2'] """ try: @@ -1716,13 +1942,7 @@ def states_with_label(self, label): except TypeError: raise ValueError("label must be a regex string.") - out_dict = {} - - for (n, l) in self.couplings.nodes(data="label"): - if re_match.match(l) is not None: - out_dict[l] = n - - return out_dict + return [l for l in self.states if isinstance(l, str) and re_match.match(l) is not None] def get_couplings(self) -> Dict[States, CouplingDict]: @@ -1794,7 +2014,7 @@ def spatial_dim(self) -> int: return np.sum(k_vector_dim) - def _states_valid(self, states: States) -> States: + def _states_valid(self, states: Sequence) -> States: """ Confirms that the provided states are in a valid format. @@ -1834,23 +2054,21 @@ def _states_valid(self, states: States) -> States: raise ValueError( f'A field must couple exactly 2 states, but {len(tpl)} are specified in {states}') for i in tpl: - if not isinstance(i, int): - raise TypeError('State specifications must be an integer') - if i >= self.basis_size: - raise ValueError((f'State specification {i} is outside ' - f'the basis size {self.basis_size:d}')) + if i not in self.states: + raise ValueError((f'State specification {i} is not a state in the basis')) - return tpl + return cast(States, tpl) # cast for mypy - def _stack_shape(self, time_dependence="all") -> Tuple[int, ...]: + def _stack_shape(self, time_dependence: Literal['steady', 'time', 'all']="all" + ) -> Tuple[int, ...]: """ Internal function to get the shape of the tuple preceding the two hamiltonian axes in :meth:`~.get_hamiltonian()` """ if len(self.variable_parameters()) == 0: - return [] + return () #make sure time_dependence is valid if time_dependence not in ["steady", "time", "all"]: @@ -1886,7 +2104,7 @@ def _stack_shape(self, time_dependence="all") -> Tuple[int, ...]: return final_shape - def basis(self) -> np.ndarray: + def dm_basis(self) -> np.ndarray: """ Generate basis labels of density matrix components. @@ -1909,10 +2127,10 @@ def basis(self) -> np.ndarray: '22_real'] """ - basis = [f'{j:d}{i:d}_imag' if i > j else f'{i:d}{j:d}_real' - for i in range(self.basis_size) - for j in range(self.basis_size)][1:] # indexing removes ground state label - return np.array(basis) + dm_basis = [f'{j:d}{i:d}_imag' if i > j else f'{i:d}{j:d}_real' + for i in range(self.basis_size) + for j in range(self.basis_size)][1:] # indexing removes ground state label + return np.array(dm_basis) def _remove_edge_data(self, states: States, kind: str): @@ -1969,7 +2187,7 @@ def _remove_edge_data(self, states: States, kind: str): self.couplings.remove_edge(*states) - def _coupling_with_label(self, label: str) -> dict: + def _coupling_with_label(self, label: str) -> States: """ Helper function to return the pair of states corresponding to a particular label string. For internal use. @@ -1982,7 +2200,7 @@ def _coupling_with_label(self, label: str) -> dict: raise ValueError(f"No coupling with label {label}") - def get_coupling_rabi(self,coupling_tuple: Tuple[int, ...] = (0, 1) + def get_coupling_rabi(self, coupling_tuple: States = (0, 1) ) -> Union[float, np.ndarray]: """ Helper function that returns the Rabi frequency of the coupling diff --git a/src/rydiqule/sensor_solution.py b/src/rydiqule/sensor_solution.py index d55d64b..bfcce01 100644 --- a/src/rydiqule/sensor_solution.py +++ b/src/rydiqule/sensor_solution.py @@ -3,7 +3,7 @@ Adds essential keys with "None" entries """ from __future__ import annotations -from typing import Optional, Tuple, Union +from typing import Optional, Union import copy @@ -14,6 +14,8 @@ # have to import this way to prevent circular imports from rydiqule import sensor_utils +Result = Union[float, np.ndarray] +ComplexResult = Union[complex, np.ndarray] class Solution(dict): """ @@ -26,25 +28,25 @@ class Solution(dict): # common attributes rho: np.ndarray """numpy.ndarray : Solutions returned by the solver.""" - eta: Optional[float] + _eta: Optional[float] """float, optional : Eta constant from the Cell. Not generally defined when using a Sensor.""" - kappa: Optional[float] + _kappa: Optional[float] """float, optional : Kappa constant from the Cell. Not generally defined when using a Sensor.""" - probe_tuple: Optional[Tuple[int, int]] + _probe_tuple: Optional[sensor_utils.States] """Coupling edge corresponding to probing field. Not generally defined when using a Sensor.""" - probe_freq: Optional[float] + _probe_freq: Optional[float] """Probing transition frequency, in rad/s. Not generally defined when using a Sensor.""" - probe_rabi: Optional[Union[float, np.ndarray]] + _probe_rabi: Optional[Union[float, np.ndarray]] """Probe Rabi frequency, in Mrad/s. Not generally defined when using a Sensor.""" - cell_length: Optional[float] + _cell_length: Optional[float] """Optical path length of the medium, in meters. Not generally defined when using a Sensor.""" - beam_area: Optional[float] + _beam_area: Optional[float] """Cross-sectional area of the probing beam, in square meters. Not generally defined when using a Sensor.""" @@ -58,7 +60,7 @@ class Solution(dict): If doppler averaging but not summing, doppler classes in internal units are added.""" rq_version: str """str : Version of rydiqule that created the Solution.""" - basis: list[str] + dm_basis: np.ndarray """list of str: The list of density matrix elements in the order they appear in the solution. See :meth:`Sensor.basis` for details.""" @@ -79,7 +81,79 @@ def __init__(self, **kwargs): super().__init__(**kwargs) self.__dict__ = self - def rho_ij(self, i: int, j: int) -> np.ndarray: + + # property attributes + @property + def eta(self) -> float: + """Eta constant from the Cell. + Not generally defined when using a Sensor.""" + if self._eta is None: + raise AttributeError("eta not defined. " + "Please define in Sensor and redo calculation") + return self._eta + + + @property + def kappa(self) -> float: + """Kappa constant from the Cell. + Not generally defined when using a Sensor.""" + if self._kappa is None: + raise AttributeError("kappa not defined. " + "Please define in Sensor and redo calculation") + return self._kappa + + + @property + def probe_tuple(self) -> sensor_utils.States: + """Coupling edge corresponding to probing field. + Not generally defined when using a Sensor.""" + if self._probe_tuple is None: + raise AttributeError("probe_tuple not defined. " + "Please define in Sensor and redo calculation") + return self._probe_tuple + + + @property + def probe_freq(self) -> float: + """Probing transition frequency, in rad/s. + Not generally defined when using a Sensor.""" + if self._probe_freq is None: + raise AttributeError("probe_freq not defined. " + "Please define in Sensor and redo calculation") + return self._probe_freq + + + @property + def probe_rabi(self) -> Union[complex, np.ndarray]: + """Probe Rabi frequency, in Mrad/s. + Not generally defined when using a Sensor.""" + if self._probe_rabi is None: + raise AttributeError("probe_rabi not defined. " + "Please define in Sensor and redo calculation") + return self._probe_rabi + + + @property + def cell_length(self) -> float: + """Optical path length of the medium, in meters. + Not generally defined when using a Sensor.""" + if self._cell_length is None: + raise AttributeError("cell_length not defined. " + "Please define in Sensor and redo calculation") + return self._cell_length + + + @property + def beam_area(self) -> float: + """Cross-sectional area of the probing beam, in square meters. + Not generally defined when using a Sensor.""" + if self._beam_area is None: + raise AttributeError("beam_area not defined. " + "Please define in Sensor and redo calculation") + return self._beam_area + + + def rho_ij(self, i: int, j: int) -> ComplexResult: """ Gets the i,j element(s) of the density matrix solutions. @@ -94,13 +168,13 @@ def rho_ij(self, i: int, j: int) -> np.ndarray: Returns ------- - numpy.ndarray - `[i,j]` elments of the density matrix + complex or numpy.ndarray + `[i,j]` elment(s) of the density matrix """ return sensor_utils.get_rho_ij(self.rho, i, j) - def get_solution_element(self, idx: int) -> np.ndarray: + def get_solution_element(self, idx: int) -> Result: """ Return a slice of an n_dimensional matrix of solutions of shape (...,n^2-1), where n is the basis size of the quantum system. @@ -112,7 +186,7 @@ def get_solution_element(self, idx: int) -> np.ndarray: Returns ------- - numpy.ndarray + float or numpy.ndarray Slice of solutions corresponding to index idx. For example, if sols has shape (..., n^2-1), sol_slice will have shape (...). @@ -141,7 +215,7 @@ def get_solution_element(self, idx: int) -> np.ndarray: return sol_slice - def get_susceptibility(self) -> np.ndarray: + def get_susceptibility(self) -> ComplexResult: """ Return the atomic susceptibility on the probe transition. @@ -149,7 +223,7 @@ def get_susceptibility(self) -> np.ndarray: Returns ------- - numpy.ndarray + complex or numpy.ndarray Susceptibility of the density matrix solution. Examples @@ -164,9 +238,6 @@ def get_susceptibility(self) -> np.ndarray: (1.9046090082907774e-05+0.0003680924230367812j) """ - #transition_frequency: Optional[float] = None - - self._variables_defined('probe_rabi', 'kappa', 'probe_freq') probe_rabi = self.probe_rabi*1e6 #rad/s kappa = self.kappa @@ -178,7 +249,7 @@ def get_susceptibility(self) -> np.ndarray: # See Steck for last factor of 2. Comes from Steck QO Notes page return kappa * (rho_probe * 2*c) / (probe_freq * (probe_rabi/2)) - def get_OD(self) -> np.ndarray: + def get_OD(self) -> Result: """ Calculates the optical depth from the solution. This equation comes from Steck's Quantum Optics Notes Eq. 6.74. @@ -189,7 +260,7 @@ def get_OD(self) -> np.ndarray: Returns ------- - OD: numpy.ndarray + OD: float or numpy.ndarray Optical depth of the sample Warns @@ -213,7 +284,6 @@ def get_OD(self) -> np.ndarray: Integrated results are likely invalid. """ - self._variables_defined('probe_freq', 'cell_length') probe_wavelength = np.abs(c/(self.probe_freq/(2*np.pi))) #meters probe_wavevector = 2*np.pi/probe_wavelength #1/meters @@ -226,7 +296,7 @@ def get_OD(self) -> np.ndarray: return OD - def get_transmission_coef(self) -> np.ndarray: + def get_transmission_coef(self) -> Result: """ Extract the transmission term from a solution. @@ -234,7 +304,7 @@ def get_transmission_coef(self) -> np.ndarray: Returns ------- - numpy.ndarray + float or numpy.ndarray Numerical value of the probe absorption in fractional units (P_out/P_in). @@ -254,7 +324,7 @@ def get_transmission_coef(self) -> np.ndarray: transmission_coef = np.exp(-OD) return transmission_coef - def get_phase_shift(self) -> np.ndarray: + def get_phase_shift(self) -> Result: """ Extract the phase shift from a solution. @@ -262,7 +332,7 @@ def get_phase_shift(self) -> np.ndarray: Returns ------- - numpy.ndarray + float or numpy.ndarray Probe phase in radians. Examples @@ -276,7 +346,6 @@ def get_phase_shift(self) -> np.ndarray: (3,) 80.52949114644437 """ - self._variables_defined('probe_rabi', 'kappa', 'cell_length') # reverse probe tuple order to get correct sign of imag # not actually necessary for this function since only need real part @@ -292,28 +361,4 @@ def copy(self): return copy.copy(self) def deepcopy(self): - return copy.deepcopy(self) - - def _variables_defined(self, *var_names): - """Check if provided experimental parameter variables are defined. - - Parameters - ---------- - var_names: str - Parameter names to check if not None - - Raises - ------ - AttributeError - If one or more parameters are not defined. - """ - - undef_vars = [] - for var in var_names: - if self[var] is None: - undef_vars.append(var) - - if len(undef_vars): - raise AttributeError(f'Required parameters not defined: {undef_vars}\n' - 'Please define in Sensor and redo calculation.') - + return copy.deepcopy(self) diff --git a/src/rydiqule/sensor_utils.py b/src/rydiqule/sensor_utils.py index afc160b..e6877b1 100644 --- a/src/rydiqule/sensor_utils.py +++ b/src/rydiqule/sensor_utils.py @@ -10,13 +10,23 @@ from scipy.constants import hbar, e from leveldiagram import LD -from typing import Dict, Tuple, Union, TYPE_CHECKING +from typing import Dict, Tuple, Union, Sequence, List, Callable, TYPE_CHECKING if TYPE_CHECKING: # only import when type checking, avoid circular import from .sensor import Sensor a0 = scipy.constants.physical_constants["Bohr radius"][0] +# put composite types of Sensor/Cell/Solution here +ScannableParameter = Union[float, List[float], np.ndarray] + +CouplingDict = Dict +State = Union[int, str] +States = Tuple[State, State] +QState = Sequence # TODO: consider using named tuples here + +TimeFunc = Callable[[float], complex] + RHO: Dict[int, np.ndarray] = {} U: Dict[int, np.ndarray] = {} B: Dict[int, int] = {} @@ -276,7 +286,7 @@ def make_real(equations: np.ndarray, constant: np.ndarray, # transform to the real basis new_eqns = u@(equations@u_inv) - new_const = 0 + new_const = np.zeros(equations.shape[:-1]) if ground_removed: new_const = np.einsum('ij,...j', u, constant) @@ -385,7 +395,8 @@ def _get_rho(n: int) -> np.ndarray: return rho -def get_rho_ij(sols: Union[np.ndarray, Solution], i: int, j: int) -> np.ndarray: +def get_rho_ij(sols: Union[np.ndarray, Solution], + i: int, j: int) -> Union[complex, np.ndarray]: """ For a given density matrix solution, retrieve a specific element of the density matrix. @@ -534,14 +545,25 @@ def draw_diagram(sensor: "Sensor", include_dephasing: bool = True) -> LD: rq_g = sensor.couplings.copy() # level settings - if isinstance(sensor, Cell): - for lev, vals in rq_g.nodes.items(): - ld_kw = {'left_text':vals['qnums']} - rq_g.nodes[lev]['ld_kw'] = ld_kw + # use rotating frames to send levels up or down relative to others + frames = sensor.get_rotating_frames() + # get flattend dictionary of all state paths + # also calculates energy based on sum of frame signs in path + energies = {k: np.sum(np.sign(v)[1:]) # ignore first state so all subgraphs start at 0 + for d in list(frames.values()) # get list of subgraph path dicts + for k, v in d.items()} + for lev, vals in rq_g.nodes.items(): + ld_kw = vals.get('ld_kw', {}) + ld_kw['energy'] = energies[lev] + + if isinstance(sensor, Cell): + ld_kw['left_text'] = vals['qnums'] + + rq_g.nodes[lev]['ld_kw'] = ld_kw # coupling settings for edge, vals in rq_g.edges.items(): - ld_kw = {} + ld_kw = vals.get('ld_kw', {}) if 'dipole_moment' in vals: ld_kw['linestyle'] = 'dashed' elif 'rabi_frequency' in vals: @@ -570,7 +592,7 @@ def draw_diagram(sensor: "Sensor", include_dephasing: bool = True) -> LD: idxs = np.argwhere(gamma_matrix != 0.0)[::-1,:] for idx in idxs: - ld_kw = {} + ld_kw = rq_g.edges[tuple(idx)].get('ld_kw', {}) ld_kw['wavy'] = True ld_kw['deflect'] = True diff --git a/src/rydiqule/slicing/slicing.py b/src/rydiqule/slicing/slicing.py index 2cfa690..bea6a4c 100644 --- a/src/rydiqule/slicing/slicing.py +++ b/src/rydiqule/slicing/slicing.py @@ -8,10 +8,10 @@ import numpy as np -from typing import Tuple, Iterator, Union, List +from typing import Tuple, Union, Optional, Iterable -def compute_grid(stack_shape: np.ndarray, n_slices: int): +def compute_grid(stack_shape: Tuple[int, ...], n_slices: int): """Calculate the bin edges to break a given stack shape into at least a certain number of pieces Works by iterating first over a number of slices per axis (N=1,2,3), @@ -26,7 +26,7 @@ def compute_grid(stack_shape: np.ndarray, n_slices: int): Parameters ---------- - stack_shape : np.ndarray + stack_shape : tuple of int The shape of the stack to be sliced. Does not include Hamiltonian or matrix equation dimensions, so for a hamiltonain stack of shape `(*l,n,n)`, `stack_shape` will be `*l`. @@ -68,10 +68,11 @@ def compute_grid(stack_shape: np.ndarray, n_slices: int): return [np.linspace(0,stack_shape[i], n_ax_slices[i]+1, dtype=int) for i in range(total_axes)] - +# in python 3.11+ return should be Iterator[tuple[tuple[slice, ...], *tuple[np.ndarray]]] def matrix_slice(*matrices: np.ndarray, - edges:Union[List, None]=None, - n_slices:Union[int, None]=None) -> Iterator[Tuple[tuple, np.ndarray]]: + edges: Optional[Iterable] = None, + n_slices: Optional[int] = None + ): # type: (...) -> Iterator[Tuple[Tuple[slice, ...], Unpack[Tuple[np.ndarray,...]]]] """ Generator that returns parts of a stack of matrices. @@ -89,17 +90,19 @@ def matrix_slice(*matrices: np.ndarray, broadcast by numpy's broadcasting rules, with the additional restriction that all matrices must have the same number of dimensions, even if some dimensions are of size 1. For example, matrices of sizes `(10,1,4,4)` and `(1,20,4,4)` can be sliced together. - edges: list(iterable) + edges: iterable, optional The values along each axis that define the edges of bins on an n-dimensional grid. For example, to slice a grid of hamiltonians with stack_shape `(10,10)` into 4 pieces, `edges` could be defined as `[0,5,10]` for each of the 2 stack axes. - n_slices : int or None + n_slices : int, optional The number of slices into which to break the matrix stack. Ignored if the `edges` parameter is not `None`. Must be specified as an integer value if `edges` is `None`, ignored otherwise. Defaults to None. Yields ------ + tuple of slices + slicing for each corresponding matrix numpy.ndarray Slice of hamiltonian stack @@ -141,11 +144,13 @@ def matrix_slice(*matrices: np.ndarray, return #Build the grid if only the number of slices is specified. if edges is None: - if n_slices in (1,None): + if n_slices is None or n_slices == 1: yield (), *matrices return - elif n_slices>1: + elif n_slices > 1: edges = compute_grid(stack_shape, n_slices) + else: + raise ValueError("n_slices must be positive int if specified") #check that the bins slice each axis appropriately for i,e in enumerate(edges): diff --git a/src/rydiqule/solvers.py b/src/rydiqule/solvers.py index 2a2eb59..bdeccf6 100644 --- a/src/rydiqule/solvers.py +++ b/src/rydiqule/solvers.py @@ -150,7 +150,7 @@ def solve_steady_state( spatial_dim = sensor.spatial_dim() # initialize doppler-related quantities - doppler_axis_shape: Iterable[int] = () + doppler_axis_shape: Tuple[int, ...] = () dop_classes = None doppler_shifts = None doppler_axes: Iterable[slice] = () @@ -203,17 +203,17 @@ def solve_steady_state( # save results to Solution object solution.rho = sols # specific to observable calculations - solution.eta = sensor.eta - solution.kappa = sensor.kappa - solution.cell_length = sensor.cell_length - solution.beam_area = sensor.beam_area - solution.probe_freq = sensor.probe_freq - solution.probe_tuple = sensor.probe_tuple + solution._eta = sensor.eta + solution._kappa = sensor.kappa + solution._cell_length = sensor.cell_length + solution._beam_area = sensor.beam_area + solution._probe_freq = sensor.probe_freq + solution._probe_tuple = sensor.probe_tuple if sensor.probe_tuple is not None: probe_rabi = sensor.get_coupling_rabi(sensor.probe_tuple) - solution.probe_rabi = probe_rabi + solution._probe_rabi = probe_rabi else: - solution.probe_rabi = None + solution._probe_rabi = None solution.couplings = sensor.get_couplings() solution.axis_labels = ([f'doppler_{i:d}' for i in range(spatial_dim) if not sum_doppler] @@ -221,8 +221,8 @@ def solve_steady_state( + ["density_matrix"]) solution.axis_values = ([dop_classes for i in range(spatial_dim) if not sum_doppler] + [val for _,_,val in sensor.variable_parameters()] - + [sensor.basis()]) - solution.basis = sensor.basis() + + [sensor.dm_basis()]) + solution.dm_basis = sensor.dm_basis() solution.rq_version = version("rydiqule") solution.doppler_classes = dop_classes diff --git a/src/rydiqule/stack_solvers/cyrk_solver.py b/src/rydiqule/stack_solvers/cyrk_solver.py index 1c150f1..09647bf 100644 --- a/src/rydiqule/stack_solvers/cyrk_solver.py +++ b/src/rydiqule/stack_solvers/cyrk_solver.py @@ -10,13 +10,14 @@ cyrk_available = False cyrk_import_error = e -from typing import Sequence, Callable, Union, Tuple +from typing import Sequence, Callable +from rydiqule.sensor_utils import TimeFunc def cyrk_solve(eoms_base: np.ndarray, const_base: np.ndarray, eom_time_r: np.ndarray, const_r: np.ndarray, eom_time_i: np.ndarray, const_i: np.ndarray, - time_inputs: Sequence[Callable[[float], Union[float, complex]]], + time_inputs: Sequence[TimeFunc], t_eval: np.ndarray, init_cond: np.ndarray, **kwargs ) -> np.ndarray: @@ -83,7 +84,7 @@ def cyrk_solve(eoms_base: np.ndarray, const_base: np.ndarray, OverflowError: If system size exceeds cyrk backend limit of 65535 equations. If we see this error a lot, consider getting CyRK project to increase it by changing type of `y_size` from unisgned short. - """ # noqa + """ if not cyrk_available: raise ImportError('CyRK backend not installed') from cyrk_import_error @@ -121,7 +122,7 @@ def cyrk_solve(eoms_base: np.ndarray, const_base: np.ndarray, def _derEqns(obes_base: np.ndarray, const_base: np.ndarray, obes_time_r: np.ndarray, const_r: np.ndarray, obes_time_i: np.ndarray, const_i: np.ndarray, - time_inputs: Tuple[Callable[[float], Union[float, complex]]] + time_inputs: Sequence[TimeFunc] ) -> Callable[[float, np.ndarray, np.ndarray], None]: """ Function to build the callable passed to CyRK's cyrk_ode cython solver. diff --git a/src/rydiqule/stack_solvers/nbkode_solver.py b/src/rydiqule/stack_solvers/nbkode_solver.py index 5820161..97da2d6 100644 --- a/src/rydiqule/stack_solvers/nbkode_solver.py +++ b/src/rydiqule/stack_solvers/nbkode_solver.py @@ -10,13 +10,14 @@ nbkode_available = False nbkode_import_error = e -from typing import Sequence, Tuple, Callable, Union +from typing import Sequence, Callable +from rydiqule.sensor_utils import TimeFunc def nbkode_solve(eoms_base: np.ndarray, const_base: np.ndarray, eom_time_r: np.ndarray, const_r: np.ndarray, eom_time_i: np.ndarray, const_i: np.ndarray, - time_inputs: Sequence[Callable[[float], Union[float, complex]]], + time_inputs: Sequence[TimeFunc], t_eval: np.ndarray, init_cond: np.ndarray, **kwargs ) -> np.ndarray: """ @@ -76,7 +77,7 @@ def nbkode_solve(eoms_base: np.ndarray, const_base: np.ndarray, numpy.ndarray The matrix solution of shape `(*l,n,n_t)` representing the density matrix of the system at each time t. - """ # noqa + """ if not nbkode_available: raise ImportError('numbakit-ode backend not installed') from nbkode_import_error @@ -111,7 +112,7 @@ def nbkode_solve(eoms_base: np.ndarray, const_base: np.ndarray, def _derEqns(obes_base: np.ndarray, const_base: np.ndarray, obes_time_r: np.ndarray, const_r: np.ndarray, obes_time_i: np.ndarray, const_i: np.ndarray, - time_inputs: Tuple[Callable[[float], Union[float, complex]]] + time_inputs: Sequence[TimeFunc] ) -> Callable[[float, np.ndarray], np.ndarray]: """ Function to build the callable passed to nbkode numba solver. diff --git a/src/rydiqule/stack_solvers/nbrk_solver.py b/src/rydiqule/stack_solvers/nbrk_solver.py index 0848662..73b7dbc 100644 --- a/src/rydiqule/stack_solvers/nbrk_solver.py +++ b/src/rydiqule/stack_solvers/nbrk_solver.py @@ -10,13 +10,14 @@ nbrk_available = False nbrk_import_error = e -from typing import Sequence, Callable, Union, Tuple +from typing import Sequence, Callable +from rydiqule.sensor_utils import TimeFunc def nbrk_solve(eoms_base: np.ndarray, const_base: np.ndarray, eom_time_r: np.ndarray, const_r: np.ndarray, eom_time_i: np.ndarray, const_i: np.ndarray, - time_inputs: Sequence[Callable[[float], Union[float, complex]]], + time_inputs: Sequence[TimeFunc], t_eval: np.ndarray, init_cond: np.ndarray, **kwargs ) -> np.ndarray: """ @@ -76,7 +77,7 @@ def nbrk_solve(eoms_base: np.ndarray, const_base: np.ndarray, numpy.ndarray The matrix solution of shape `(*l,n,n_t)` representing the density matrix of the system at each time t. - """ # noqa + """ if not nbrk_available: raise ImportError('CyRK backend not installed') from nbrk_import_error @@ -109,7 +110,7 @@ def nbrk_solve(eoms_base: np.ndarray, const_base: np.ndarray, def _derEqns(obes_base: np.ndarray, const_base: np.ndarray, obes_time_r: np.ndarray, const_r: np.ndarray, obes_time_i: np.ndarray, const_i: np.ndarray, - time_inputs: Tuple[Callable[[float], Union[float, complex]]] + time_inputs: Sequence[TimeFunc] ) -> Callable[[float, np.ndarray], np.ndarray]: """ Function to build the callable passed to CyRK's nbrk_ode numba solver. diff --git a/src/rydiqule/stack_solvers/scipy_solver.py b/src/rydiqule/stack_solvers/scipy_solver.py index 5411296..43cd2ee 100644 --- a/src/rydiqule/stack_solvers/scipy_solver.py +++ b/src/rydiqule/stack_solvers/scipy_solver.py @@ -2,12 +2,13 @@ from scipy.integrate import solve_ivp -from typing import Sequence, Callable, Union, Literal +from typing import Sequence, Callable, Literal +from rydiqule.sensor_utils import TimeFunc def scipy_solve(eoms_base: np.ndarray, const: np.ndarray, eom_time_r: np.ndarray, const_r: np.ndarray, eom_time_i: np.ndarray, const_i: np.ndarray, - time_inputs: Sequence[Callable[[float], Union[float, complex]]], + time_inputs: Sequence[TimeFunc], t_eval: np.ndarray, init_cond: np.ndarray, eqns: Literal["loop", "comp"] = "loop", **kwargs @@ -72,7 +73,7 @@ def scipy_solve(eoms_base: np.ndarray, const: np.ndarray, numpy.ndarray The matrix solution of shape `(*l,n,n_t)` representing the density matrix of the system at each time t. - """ # noqa + """ _derEqns = {"loop": _derEqns_loop, "comp": _derEqn_comp} @@ -102,7 +103,7 @@ def scipy_solve(eoms_base: np.ndarray, const: np.ndarray, def _derEqns_loop(obes_base: np.ndarray, const_base: np.ndarray, obes_time_r: np.ndarray, const_r: np.ndarray, obes_time_i: np.ndarray, const_i: np.ndarray, - time_inputs: Sequence[Callable[[float], Union[float, complex]]], + time_inputs: Sequence[TimeFunc], ) -> Callable[[float, np.ndarray], np.ndarray]: """ Function to build the callable passed to scipy's solve_ivp in :func:`~.scipy_solve`. @@ -138,7 +139,7 @@ def func(t: float, A_flat: np.ndarray) -> np.ndarray: def _derEqn_comp(obes_base: np.ndarray, const: np.ndarray, obes_time_r: np.ndarray, const_r: np.ndarray, obes_time_i: np.ndarray, const_i: np.ndarray, - time_inputs: Sequence[Callable[[float], Union[float, complex]]], + time_inputs: Sequence[TimeFunc], ) -> Callable[[float, np.ndarray], np.ndarray]: """ Function to build the callable passed to scipy's solve_ivp in :func:`~.scipy_solve`. diff --git a/src/rydiqule/timesolvers.py b/src/rydiqule/timesolvers.py index e5607ce..ff211dd 100644 --- a/src/rydiqule/timesolvers.py +++ b/src/rydiqule/timesolvers.py @@ -5,7 +5,7 @@ import numpy as np from importlib.metadata import version -from .sensor_utils import _hamiltonian_term, generate_eom, make_real +from .sensor_utils import _hamiltonian_term, generate_eom, make_real, TimeFunc from .sensor_solution import Solution from .sensor import Sensor # only needing for type hinting from .slicing.slicing import matrix_slice, get_slice_num_t @@ -17,12 +17,12 @@ from rydiqule.stack_solvers.cyrk_solver import cyrk_solve from rydiqule.stack_solvers.nbrk_solver import nbrk_solve -from typing import Optional, Tuple, Iterable, List, Callable, Literal, Union +from typing import Optional, Tuple, List, Callable, Literal, Union, Dict -stack_solvers = {"scipy": scipy_solve, - "nbkode": nbkode_solve, - "cyrk": cyrk_solve, - "nbrk": nbrk_solve} +stack_solvers: Dict[str, Callable] = {"scipy": scipy_solve, + "nbkode": nbkode_solve, + "cyrk": cyrk_solve, + "nbrk": nbrk_solve} solver_type = Union[Callable, Literal['scipy', 'nbkode', 'cyrk', 'nbrk']] @@ -145,7 +145,7 @@ def solve_time(sensor: Sensor, end_time: float, num_pts: int, f"{solver} is not a built-in solver." f" Supported built-in solvers are {list(stack_solvers.keys())}") - if not isinstance(solver, Callable): + if not callable(solver): raise ValueError("Solvers must be callable functions") solution = Solution() @@ -160,11 +160,11 @@ def solve_time(sensor: Sensor, end_time: float, num_pts: int, t_eval = np.linspace(*time_range, num=num_pts, dtype=np.float64) # initialize doppler-related quantities - doppler_axis_shape: Iterable[int] = () + doppler_axis_shape: Tuple[int, ...] = () dop_classes = None doppler_shifts = None - out_doppler_axes: Iterable[slice] = () - doppler_axes: Iterable[slice] = () + out_doppler_axes: Tuple[slice, ...] = () + doppler_axes: Tuple[slice, ...] = () # update doppler-related values if doppler: @@ -198,10 +198,10 @@ def solve_time(sensor: Sensor, end_time: float, num_pts: int, n_slices_true = sum(1 for _ in matrix_slice(gamma, n_slices=n_slices)) - for i, (idx, H, G, *time_hams), in enumerate(matrix_slice(hamiltonians, gamma, - *hamiltonians_time, - *hamiltonians_time_i, - n_slices=n_slices)): + for i, (idx, H, G, *time_hams) in enumerate(matrix_slice(hamiltonians, gamma, + *hamiltonians_time, + *hamiltonians_time_i, + n_slices=n_slices)): if n_slices_true > 1: print(f"Solving equation slice {i+1}/{n_slices_true}", end='\r') @@ -222,25 +222,25 @@ def solve_time(sensor: Sensor, end_time: float, num_pts: int, # save results to the Solution solution.rho = sols # specific to observable calculations - solution.eta = sensor.eta - solution.kappa = sensor.kappa - solution.cell_length = sensor.cell_length - solution.beam_area = sensor.beam_area - solution.probe_freq = sensor.probe_freq - solution.probe_tuple = sensor.probe_tuple + solution._eta = sensor.eta + solution._kappa = sensor.kappa + solution._cell_length = sensor.cell_length + solution._beam_area = sensor.beam_area + solution._probe_freq = sensor.probe_freq + solution._probe_tuple = sensor.probe_tuple if sensor.probe_tuple is not None: probe_rabi = sensor.get_coupling_rabi(sensor.probe_tuple) - solution.probe_rabi = probe_rabi + solution._probe_rabi = probe_rabi else: - solution.probe_rabi = None + solution._probe_rabi = None solution.axis_labels = ([f'doppler_{i:d}' for i in range(spatial_dim) if not sum_doppler] + sensor.axis_labels() + ["time", "density_matrix"]) solution.axis_values = ([dop_classes for i in range(spatial_dim) if not sum_doppler] + [val for _,_,val in sensor.variable_parameters()] - + [t_eval, sensor.basis()]) - solution.basis = sensor.basis() + + [t_eval, sensor.dm_basis()]) + solution.dm_basis = sensor.dm_basis() solution.rq_version = version("rydiqule") solution.doppler_classes = dop_classes # time solver specific @@ -252,7 +252,7 @@ def solve_time(sensor: Sensor, end_time: float, num_pts: int, def _solve_hamiltonian_stack(hamiltonians_base: np.ndarray, hamiltonians_time: np.ndarray, hamiltonians_time_i: np.ndarray, gamma_matrix: np.ndarray, - time_functions: List[Callable[[float], float]], t_eval: np.ndarray, + time_functions: List[TimeFunc], t_eval: np.ndarray, init_cond: np.ndarray, solver, doppler: bool = False, dop_classes: Optional[np.ndarray] = None, sum_doppler: bool = True, weight_doppler: bool = True, @@ -299,7 +299,7 @@ def _solve_hamiltonian_stack(hamiltonians_base: np.ndarray, hamiltonians_time: n def solve_eom_stack(eoms_base: np.ndarray, const: np.ndarray, eom_time_r: np.ndarray, const_r: np.ndarray, eom_time_i: np.ndarray, const_i: np.ndarray, - time_inputs: List[Callable[[float],float]], + time_inputs: List[TimeFunc], t_eval: np.ndarray, init_cond: np.ndarray, solver, **kwargs ) -> np.ndarray: diff --git a/tests/test_experiments.py b/tests/test_experiments.py index 92ccecc..6e31f3d 100644 --- a/tests/test_experiments.py +++ b/tests/test_experiments.py @@ -80,7 +80,7 @@ def test_phase_shift_with_steck(): e_field=1e-7 - rb_cell = rq.Sensor(basis_size = 2) + rb_cell = rq.Sensor(2) rb_cell.set_experiment_values(cell_length=cell_length, probe_freq = probe_freq, beam_area=beam_area, probe_tuple = (0,1), kappa = kappa) rb_cell.set_gamma_matrix(np.array([[0.23229296, 0.00000000e+00], diff --git a/tests/test_sensor.py b/tests/test_sensor.py index 1d315e4..d87c710 100644 --- a/tests/test_sensor.py +++ b/tests/test_sensor.py @@ -89,15 +89,28 @@ def test_states_validation(): # too many labels passed s._states_valid((0,1,2)) - with pytest.raises(TypeError, match='must be an integer'): - # non-integer inputs - s._states_valid('01') - - with pytest.raises(ValueError, match='outside the basis size'): + with pytest.raises(ValueError, match='not a state in the basis'): # outside basis s._states_valid((0,2)) +@pytest.mark.exception +def test_basis_exception(): + '''Assures that invalid bases are correctly caught''' + #bad state type + with pytest.raises(ValueError, match="must be integers or strings"): + s = rq.Sensor(['g', 'e', None]) + + #double state labels + with pytest.raises(ValueError, match="All state labels must be unique"): + s = rq.Sensor(['g', 'e', 'e']) + + #bad basis type + with pytest.raises(TypeError, match="must be an integer or iterable"): + l = lambda x: x + s = rq.Sensor(l) + + @pytest.mark.exception def test_zip_param_validation(): '''Assures that incorrect usage of Sensor.zip_parameters is caught.''' diff --git a/tests/test_steady_state_ham_gen.py b/tests/test_steady_state_ham_gen.py index 690e702..e2231b1 100644 --- a/tests/test_steady_state_ham_gen.py +++ b/tests/test_steady_state_ham_gen.py @@ -22,6 +22,20 @@ def test_ladder_5_level(): np.testing.assert_allclose(ham, expected_ham, err_msg="5 Level ladder failed hamilonian generation") +@pytest.mark.steady_state +def test_mixed_state_ladder(): + + f1 = {'states':('g',1), 'rabi_frequency':0, 'detuning':1} + f2 = {'states':(1,2), 'rabi_frequency':0, 'detuning':2} + + s = Sensor(['g', 1, 2]) + s.add_couplings(f1, f2) + + expected_ham = np.diag([0, -1, -1-2]) + ham = s.get_hamiltonian() + + np.testing.assert_allclose(ham, expected_ham, + err_msg="5 Level ladder failed hamilonian generation") @pytest.mark.steady_state def test_lambda_4_level(): @@ -30,14 +44,24 @@ def test_lambda_4_level(): f2 = {'states':(1,3), 'rabi_frequency':0, 'detuning':2} f3 = {'states':(2,3), 'rabi_frequency':0, 'detuning':4} + s_f1 = {'states':('a','b'), 'rabi_frequency':0, 'detuning':1} + s_f2 = {'states':('b','d'), 'rabi_frequency':0, 'detuning':2} + s_f3 = {'states':('c','d'), 'rabi_frequency':0, 'detuning':4} + sensor = Sensor(4) sensor.add_couplings(f1, f2, f3) + sensor_str = Sensor(['a', 'b', 'c', 'd']) + sensor_str.add_couplings(s_f1, s_f2, s_f3) + expected_ham = np.diag([0, -1, -1-2+4, -1-2]) ham = sensor.get_hamiltonian() + ham_str = sensor.get_hamiltonian() np.testing.assert_allclose(ham, expected_ham, err_msg="4 Level lambda failed hamilonian generation") + np.testing.assert_allclose(ham_str, expected_ham, + err_msg="4 Level lambda failed ham generation with str labels") @pytest.mark.steady_state @@ -79,3 +103,17 @@ def test_v_5_cell(): np.testing.assert_allclose(ham, expected_ham, err_msg="5 Level V cell failed hamilonian generation") + +@pytest.mark.steady_state +def test_e_shift(): + f1 = {'states':('g',1), 'rabi_frequency':0, 'detuning':1} + f2 = {'states':(1,2), 'rabi_frequency':0, 'detuning':2} + + s = Sensor(['g', 1, 2]) + s.add_couplings(f1, f2) + s.add_energy_shift('g', .5) + s.add_energy_shift(1, 1) + s.add_energy_shift(2, 1.5) + + expected_ham = np.diag([0, -1, -1-2]) + np.diag([.5, 1, 1.5]) + ham = s.get_hamiltonian() \ No newline at end of file