Merge pull request #36 from QTC-UMD/lazy_imports

Lazy imports for leveldiagram and arc dependencies

Author: dihm

docs/source/changelog.rst

@@ -7,6 +7,12 @@ Latest
 Improvements
 ++++++++++++
 
+- `leveldiagram` and `arc` dependencies are now lazy loaded at first use
+  (`draw_diagram` and `Cell` instantiation, respectively).
+  This improves `import rydiqule as rq` time by about a second
+  and ensures second-order dependencies such as matplotlib and sympy are not
+  imported unless explicitly needed.
+
 
 Bug Fixes
 +++++++++

src/rydiqule/arc_utils.py

@@ -5,8 +5,11 @@
 from __future__ import annotations
 from typing import TYPE_CHECKING, Literal, Dict, Callable, Tuple
 
-from arc import C_e, C_c, C_h
 from scipy.constants import pi, hbar, epsilon_0, physical_constants
+from scipy.constants import c as C_c
+from scipy.constants import h as C_h
+from scipy.constants import e as C_e
+
 import numpy as np
 
 from math import sqrt

src/rydiqule/atom_utils.py

@@ -2,28 +2,48 @@
 Utilities for interacting with atomic parameters and ARC.
 """
 
-import arc.alkali_atom_data as arc_atoms
 from scipy.constants import epsilon_0, hbar, c
 import numpy as np
-import re 
+import re
 from .sensor_utils import expand_statespec
 
-from typing import Union, Optional, Literal, Tuple, List, Dict, Callable, NamedTuple
+from typing import Union, Optional, Literal, Tuple, List, Dict, Callable, NamedTuple, TYPE_CHECKING
 
-from .exceptions import RydiquleError
+from .exceptions import RydiquleError, AtomError
+
+if TYPE_CHECKING:
+    import arc.alkali_atom_data
 
 ATOMS = {
-    'H': arc_atoms.Hydrogen,
-    'Li6': arc_atoms.Lithium6, 'Li7': arc_atoms.Lithium7,
-    'Na': arc_atoms.Sodium,
-    'K39': arc_atoms.Potassium39, 'K40': arc_atoms.Potassium40, 'K41': arc_atoms.Potassium41,
-    'Rb85': arc_atoms.Rubidium85, 'Rb87': arc_atoms.Rubidium87,
-    'Cs': arc_atoms.Caesium
+    '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,

src/rydiqule/cell.py

@@ -14,9 +14,9 @@
 from .sensor import Sensor
 from .sensor_utils import scale_dipole
 from .sensor_utils import ScannableParameter, TimeFunc
-from .atom_utils import ATOMS, calc_kappa, calc_eta, expand_qnums, validate_qnums, A_QState, ground_state
+from .atom_utils import calc_kappa, calc_eta, expand_qnums, validate_qnums, A_QState, ground_state, _load_arc_atom
 from .arc_utils import RQ_AlkaliAtom
-from .exceptions import RydiquleError, AtomError, CouplingNotAllowedError
+from .exceptions import RydiquleError, CouplingNotAllowedError
 from .exceptions import RydiquleWarning, debug_state
 
 from typing import Literal, Optional, Sequence, List, Tuple, Callable, Union, Dict
@@ -165,11 +165,9 @@ def __init__(self, atom_flag: AtomFlags, atomic_states: List[A_QState],
         (5, 0, 0.5, f=2.0, m_f=2.0)
 
         """
-        if atom_flag not in ATOMS.keys():
-            raise AtomError(f"Atom flag must be one of {ATOMS.keys()}")
 
         self.atom_flag = atom_flag
-        self.atom = RQ_AlkaliAtom(ATOMS[atom_flag]())
+        self.atom = RQ_AlkaliAtom(_load_arc_atom(atom_flag))
         self.I = self.atom.arc_atom.I
 
         #prepare states by expanding statespec into list and initialize graph

src/rydiqule/sensor_utils.py

@@ -10,12 +10,12 @@
 from .exceptions import RydiquleError
 import scipy.constants
 from scipy.constants import hbar, e
-from leveldiagram import LD
 
 from typing import Dict, Tuple, Union, List, Callable, TYPE_CHECKING
 if TYPE_CHECKING:
     # only import when type checking, avoid circular import
     from .sensor import Sensor
+    from leveldiagram import LD
 
 a0 = scipy.constants.physical_constants["Bohr radius"][0]
 
@@ -713,7 +713,7 @@ def scale_dipole(dipole: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
     return dipole
 
 
-def draw_diagram(sensor: "Sensor", include_dephasing: bool = True) -> LD:
+def draw_diagram(sensor: "Sensor", include_dephasing: bool = True) -> "LD":
     """
     Draw a matplotlib plot that shows the energy level diagram, couplings, and dephasing paths.
 
@@ -745,6 +745,9 @@ def draw_diagram(sensor: "Sensor", include_dephasing: bool = True) -> LD:
         Diagram handle
 
     """
+    # lazy import of leveldiagram to avoid matplotlib import when not needed
+    from leveldiagram import LD
+
     rq_g = sensor.couplings.copy()
     sensor_states = sensor.states