Merge pull request #38 from QTC-UMD/cell_fixes

Multiple bug fixes in Cell

Author: dihm

docs/source/changelog.rst

@@ -23,6 +23,10 @@ Bug Fixes
   It is expected that other plots and values will differ by small amounts when using `Cell` and ARC v3.9.0
 - Fixed bug where `solve_doppler_analytic` ignored parameter time-dependence. 
   It now follows the behavior of other steady-state solvers by constructing the Hamiltonian and equations of motion at `t=0`.
+- Ensure that `Cell.eta`, `Cell.kappa`, and `Cell.probe_freq` properly cache their values as intended.
+- Fix bug where `Cell` would not add transit broadening automatically, as intended.
+- Fix errors in `gamma_mismatch` calculations when coupling groups are used.
+  Also fix issue that prevented `gamma_mismatch='all'` from working if only 1 dephasing is present that needs modification.
 
 Deprecations
 ++++++++++++

src/rydiqule/atom_utils.py

@@ -4,6 +4,7 @@
 
 from scipy.constants import epsilon_0, hbar, c
 import numpy as np
+import math
 import re
 from .sensor_utils import expand_statespec
 
@@ -667,7 +668,7 @@ def calc_eta(omega: float, dipole_moment:float, beam_area: float) -> float:
         The value of eta, in root(Hz)
     """
 
-    eta = np.sqrt((omega*dipole_moment**2)/(2.0*c*epsilon_0*hbar*beam_area))
+    eta = math.sqrt((omega*dipole_moment**2)/(2.0*c*epsilon_0*hbar*beam_area))
 
     return eta
 
@@ -793,7 +794,7 @@ def validate_qnums(qstate:A_QState, I: Optional[float]=None):
     #validate (n,l) int, j half int
     assert int(n)==n, f"invalid n quantum number {n}."
     assert (int(l)==l) and (l < n), f"invalid l quantum number {l}."
-    assert j==l+1/2 or j==np.abs(l-1/2), f"invalid j quantum number {j}"
+    assert j==l+1/2 or j==abs(l-1/2), f"invalid j quantum number {j}"
 
     #test m_j, f, m_f are allowed values
     if m_j is not None:
@@ -967,7 +968,7 @@ def get_valid_f(state: A_QState, I: Optional[float]=None) -> List[float]:
     J_qnum=state[2]
     if not isinstance(J_qnum, (int, float)) or not isinstance(I, (int, float)):
         raise ValueError(f"Invalid I,J quantum number types {(type(I),type(J_qnum))}.")
-    return np.arange(np.abs(J_qnum - I), J_qnum + I + 1).tolist()
+    return np.arange(abs(J_qnum - I), J_qnum + I + 1).tolist()
 
 
 def get_valid_mf(state: A_QState, I: Optional[float]=None) -> List[float]:

src/rydiqule/cell.py

@@ -5,7 +5,8 @@
 import scipy
 import numpy as np
 import warnings
-import itertools 
+import itertools
+import math
 
 import scipy.constants
 from scipy.constants import Boltzmann, e
@@ -185,25 +186,26 @@ def __init__(self, atom_flag: AtomFlags, atomic_states: List[A_QState],
         self.temp = temp  # K
         self.beam_area = beam_area
         self.density = self.atom.arc_atom.getNumberDensity(self.temp)
-        self.atom_mass = self.atom.arc_atom.mass
+        self.atom_mass: float = self.atom.arc_atom.mass
         if beam_diam is None:
-            self.beam_diameter = 2.0*np.sqrt(beam_area/np.pi)
+            self.beam_diameter = 2.0*math.sqrt(beam_area/np.pi)
         else:
             self.beam_diameter = beam_diam
 
         if gamma_transit is None:
-            gamma_transit = 1E-6*np.sqrt(8*Boltzmann*self.temp/(self.atom_mass*np.pi)
-                                         )/(self.beam_diameter/2*np.sqrt(2*np.log(2)))
+            gamma_transit = 1E-6*math.sqrt(8*Boltzmann*self.temp/(self.atom_mass*np.pi)
+                                         )/(self.beam_diameter/2*math.sqrt(2*math.log(2)))
         self.gamma_transit = gamma_transit
 
         # most probable speed for a 3D Maxwell-Boltzmann distribution
         # used when defining doppler averaging
-        self.vP = np.sqrt(2*Boltzmann*self.temp/self.atom_mass)
+        self.vP = math.sqrt(2*Boltzmann*self.temp/self.atom_mass)
 
         self._add_state_energies()
         self._add_state_lifetimes()
         self._add_decoherence_rates()
         self._add_gamma_mismatches(gamma_mismatch)
+        self.add_transit_broadening(gamma_transit)
 
 
     def set_experiment_values(self,
@@ -341,7 +343,7 @@ def level_ordering(self) -> List[A_QState]:
 
 
     @property
-    def kappa(self):
+    def kappa(self) -> float:
         """Property to calculate the kappa value of the system. 
 
         The value is computed with the following formula Eq. 5 of
@@ -369,7 +371,7 @@ def kappa(self):
             return self._kappa
 
         if self.probe_tuple is None:
-            raise RydiquleError("Cell.probe_tuple not set. Either set manually or add at least one coupling before calculation.")
+            raise RydiquleError("Cell.probe_tuple not set. Add at least one coupling before calculation.")
 
         ground_manifold = self.states_with_spec(self.probe_tuple[0])
         excited_manifold = self.states_with_spec(self.probe_tuple[1])
@@ -394,6 +396,7 @@ def kappa(self):
         dipole_moment = self.atom.get_dipole_matrix_element(probe_g_nlj, probe_e_nlj, q=q)*a0*e
 
         kappa = calc_kappa(omega_rad, dipole_moment, self.density)
+        self._kappa = kappa
 
         return kappa
 
@@ -433,7 +436,7 @@ def kappa(self):
 
 
     @property
-    def eta(self):
+    def eta(self) -> float:
         """Get the eta value for the system.
 
         The value is computed with the following formula Eq. 7 of
@@ -459,7 +462,7 @@ def eta(self):
         if hasattr(self, "_eta"):
             return self._eta
         if self.probe_tuple is None:
-            raise RydiquleError("Cell.probe_tuple not set. Either set manually or add at least one coupling before calculation.")
+            raise RydiquleError("Cell.probe_tuple not set. Add at least one coupling before calculation.")
 
         ground_manifold = self.states_with_spec(self.probe_tuple[0])
         excited_manifold = self.states_with_spec(self.probe_tuple[1])
@@ -483,12 +486,13 @@ def eta(self):
         omega_rad = self.atom.get_transition_frequency(probe_g_nlj, probe_e_nlj)*2.0*np.pi
         dipole_moment = self.atom.get_dipole_matrix_element(probe_g_nlj, probe_e_nlj, q=q)*a0*e
         eta = calc_eta(omega_rad, dipole_moment, self.beam_area)
+        self._eta = eta
 
         return eta
 
 
     @eta.setter
-    def eta(self, value):
+    def eta(self, value: float):
         """Setter for the eta attribute.
 
         Updates the self._eta class attribute.
@@ -519,7 +523,7 @@ def eta(self):
 
 
     @property
-    def probe_freq(self):
+    def probe_freq(self) -> float:
         """Get the probe transition frequency, in rad/s.
 
         Note that for :class:`~.Cell`, probing transition frequency is calculated using only
@@ -535,6 +539,8 @@ def probe_freq(self):
 
         if hasattr(self, '_probe_freq'):
             return self._probe_freq
+        if self.probe_tuple is None:
+            raise RydiquleError("Cell.probe_tuple not set. Add at least one coupling before calculation.")
 
         probe_lower_manifold = self.states_with_spec(self.probe_tuple[0])
         probe_upper_manifold = self.states_with_spec(self.probe_tuple[1])
@@ -544,11 +550,14 @@ def probe_freq(self):
 
         energy_lower = self.atom.get_state_energy(A_QState(n1, l1, j1), s=0.5)*2*np.pi
         energy_upper = self.atom.get_state_energy(A_QState(n2, l2, j2), s=0.5)*2*np.pi
+
+        probe_freq = abs(energy_upper - energy_lower)
+        self._probe_freq = probe_freq
 
-        return np.abs(energy_upper - energy_lower)
+        return probe_freq
 
     @probe_freq.setter
-    def probe_freq(self, value):
+    def probe_freq(self, value: float):
         """Setter for the probe_freq attribute.
 
         Updates the self._probe_freq class attribute.
@@ -737,7 +746,8 @@ def add_single_coupling(
         Coherent Couplings: 
             ((5, 0, 0.5),(5, 1, 1.5)): {rabi_frequency: 2.0, detuning: 1.0, phase: 0, kvec: (0, 0, 0), label: probe, coherent_cc: 1, dipole_moment: 2.44, q: 0}
         Decoherent Couplings:
-            ((5, 1, 1.5),(5, 0, 0.5)): {gamma_transition: 38.11316}
+            ((5, 0, 0.5),(5, 0, 0.5)): {gamma_transit: 0.40697}
+            ((5, 1, 1.5),(5, 0, 0.5)): {gamma_transition: 38.11316, gamma_transit: 0.40697}
         Energy Shifts:
             None
         
@@ -755,7 +765,8 @@ def add_single_coupling(
         Coherent Couplings: 
             ((5, 0, 0.5),(5, 1, 1.5)): {rabi_frequency: 2.0, detuning: 1.0, phase: 0, kvec: (0, 0, 0), label: probe, coherent_cc: 1, dipole_moment: 2.44, q: 0}
         Decoherent Couplings:
-            ((5, 1, 1.5),(5, 0, 0.5)): {gamma_transition: 38.11}
+            ((5, 0, 0.5),(5, 0, 0.5)): {gamma_transit: 0.40696}
+            ((5, 1, 1.5),(5, 0, 0.5)): {gamma_transition: 38.113, gamma_transit: 0.40696}
         Energy Shifts:
             None
 
@@ -772,7 +783,8 @@ def add_single_coupling(
         Coherent Couplings: 
             ((5, 0, 0.5),(5, 1, 1.5)): {rabi_frequency: 1.177, detuning: 1.0, phase: 0, kvec: (0, 0, 0), label: probe, coherent_cc: 1, dipole_moment: 2.44, q: 0}
         Decoherent Couplings:
-            ((5, 1, 1.5),(5, 0, 0.5)): {gamma_transition: 38.11}
+            ((5, 0, 0.5),(5, 0, 0.5)): {gamma_transit: 0.40696}
+            ((5, 1, 1.5),(5, 0, 0.5)): {gamma_transition: 38.11, gamma_transit: 0.40696}
         Energy Shifts:
             None
 
@@ -787,7 +799,8 @@ def add_single_coupling(
         Coherent Couplings: 
             ((5, 0, 0.5),(5, 1, 1.5)): {rabi_frequency: 4.3, detuning: 1.0, phase: 0, kvec: (0, 0, 0), label: probe, coherent_cc: 1, dipole_moment: 2.44, q: 0}
         Decoherent Couplings:
-            ((5, 1, 1.5),(5, 0, 0.5)): {gamma_transition: 38.11}
+            ((5, 0, 0.5),(5, 0, 0.5)): {gamma_transit: 0.4117}
+            ((5, 1, 1.5),(5, 0, 0.5)): {gamma_transition: 38.113, gamma_transit: 0.4117}
         Energy Shifts:
             None
 
@@ -904,7 +917,7 @@ def add_single_coupling(
             lam = abs(self.atom.get_transition_wavelength(state1, state2)) # in m
             kvec = 2*np.pi/lam*np.asarray(kunit)*1e-6 # scaled to Mrad/m
         else:
-            raise RydiquleError(f'Coupling {states} has un-normalized |kunit|={np.sqrt(k_norm_sq):.2f}!=1')
+            raise RydiquleError(f'Coupling {states} has un-normalized |kunit|={math.sqrt(k_norm_sq):.2f}!=1')
 
         super().add_single_coupling(states=states,
                                     rabi_frequency=passed_rabi,
@@ -1066,7 +1079,7 @@ def _add_gamma_mismatches(self, method:Union[str, dict]="ground"):
         `gamma_transition` values on edges leaving that state. However, it is not always
         desirable to account for all states in this way for simplicity or computational
         complexity reasons. This function allows the :class:`~.Cell` to account for any
-        differences in these values that arise as a result of excluding physicalstates from a
+        differences in these values that arise as a result of excluding physical states from a
         :class:`~.Cell`. There are multiple ways to resolve these discrepancies, specified by
         the `method` argument, which is detailed in the `Parameters` section.
 
@@ -1155,8 +1168,9 @@ def _add_gamma_mismatch_to_ground(self, state: A_QState):
             m_j = None if g.m_j is None else "all"
             (f, m_f) = (None, None) if g.f is None else ("all","all")
             ground_manifold = A_QState(g.n, g.l, g.j, m_j=m_j, f=f, m_f=m_f)
+            degeneracy = len(self.states_with_spec(ground_manifold))
 
-            self.add_decoherence((state, ground_manifold), gamma=(lifetime-transition_total), label="mismatch")
+            self.add_decoherence((state, ground_manifold), gamma=(lifetime-transition_total)/degeneracy, label="mismatch")
 
 
     def _add_gamma_mismatch_to_all(self, state:A_QState):
@@ -1180,18 +1194,21 @@ def _add_gamma_mismatch_to_all(self, state:A_QState):
 
         transition_total = sum(e[2] for e in out_edges)
 
-        #if they dom't match, we add a decoherence to the entire ground state manifold
-        #that matches what remains
+        #if they dom't match, we proportionally increase existing decoherences to make up the difference
         if not np.isclose(transition_total, lifetime):
-            #construct the dictionary of coupling coefficients coefficients
-            cc = {
-                (s1, s2):gamma/transition_total for s1, s2, gamma in out_edges
-                if gamma
-            }
-
+            gamma_total_mismatch = lifetime-transition_total
             out_states_list = [s2 for _, s2, _ in out_edges]
-            self.add_decoherence_group([state], out_states_list, gamma = transition_total, coupling_coefficients=cc, label="mismatch")
-
+            if len(out_states_list) > 1:
+                #construct the dictionary of coupling coefficients
+                cc = {
+                    (s1, s2):gamma/transition_total for s1, s2, gamma in out_edges
+                    if gamma
+                }
+
+                self.add_decoherence_group([state], out_states_list, gamma = gamma_total_mismatch, coupling_coefficients=cc, label="mismatch")
+            else:
+                self.add_single_decoherence((state, out_states_list[0]), gamma_total_mismatch,
+                                            label='mismatch')
 
     def _validate_input_states(self, atomic_states: List[A_QState]):
         """Helper function to check that input states are compatible and defined"""

src/rydiqule/experiments.py

@@ -90,8 +90,8 @@ def get_snr(sensor: Sensor,
     >>> c.add_coupling(states=(g,e), rabi_frequency=np.linspace(1e-6, 1, 5), detuning=1, label="probe")
     >>> snr, mesh = rq.get_snr(c, 'probe_rabi_frequency')
     >>> print(snr) 
-    [       0.       13947396.7 27887614.4 41813486.6
-     55717871.1]
+    [       0.       13654034.1 27301261.5 40934886.9
+     54548137.6]
     >>> print(mesh) # doctest: +SKIP
     [array([0.      , 0.25    , 0.499999, 0.749999, 0.999999])]
 

src/rydiqule/sensor_solution.py

@@ -352,7 +352,7 @@ def get_solution_element(self, idx: int) -> Result:
         (3,)
         >>> rho_01_im = sols.get_solution_element(0)
         >>> print(rho_01_im)
-        0.0013711656
+        0.00131399
 
         """
         basis_size = np.sqrt(self.rho.shape[-1] + 1)
@@ -382,7 +382,11 @@ def coupling_coefficient_matrix(self, coupling: sensor_utils.StateSpecs) -> np.n
         >>> sol = rq.solve_steady_state(my_cell)
         >>> for e in sol.couplings.edges(data="coherent_cc"):
         ...     print(*e)
+        (5, 0, 0.5, m_j=-0.5) (5, 0, 0.5, m_j=-0.5) None
+        (5, 0, 0.5, m_j=-0.5) (5, 0, 0.5, m_j=0.5) None
         (5, 0, 0.5, m_j=-0.5) (5, 1, 0.5, m_j=-0.5) -0.816496580927726
+        (5, 0, 0.5, m_j=0.5) (5, 0, 0.5, m_j=-0.5) None
+        (5, 0, 0.5, m_j=0.5) (5, 0, 0.5, m_j=0.5) None
         (5, 0, 0.5, m_j=0.5) (5, 1, 0.5, m_j=0.5) 0.816496580927726
         (5, 1, 0.5, m_j=-0.5) (5, 0, 0.5, m_j=-0.5) None
         (5, 1, 0.5, m_j=-0.5) (5, 0, 0.5, m_j=0.5) None
@@ -495,7 +499,7 @@ def get_susceptibility(self) -> ComplexResult:
         (3,)
         >>> sus = sols.get_susceptibility()
         >>> print(sus)
-        (1.9734254e-05+0.000376067j)
+        (1.891136e-05+0.00036817j)
 
         """
         probe_rabi = self.probe_rabi*1e6 #rad/s
@@ -539,11 +543,11 @@ def get_OD(self) -> Result:
         (3,)
         >>> OD = sols.get_OD()
         >>> print(OD)
-        0.3028422896
+        0.2964844
 
         """
 
-        probe_wavelength = np.abs(c/(self.probe_freq/(2*np.pi))) #meters
+        probe_wavelength = abs(c/(self.probe_freq/(2*np.pi))) #meters
         probe_wavevector = 2*np.pi/probe_wavelength #1/meters
         OD = self.get_susceptibility().imag*self.cell_length*probe_wavevector
         if np.any(OD > 1.0):
@@ -577,7 +581,7 @@ def get_transmission_coef(self) -> Result:
         (3,)
         >>> t = sols.get_transmission_coef()
         >>> print(t)
-        0.73871559029
+        0.743427
 
         """
         OD = self.get_OD()
@@ -605,7 +609,7 @@ def get_phase_shift(self) -> Result:
         >>> print(sols.rho.shape)
         (3,)
         >>> print(sols.get_phase_shift())
-        80.5295218956
+        80.5294887
         
         """
         # reverse probe tuple order to get correct sign of imag

src/rydiqule/sensor_utils.py

@@ -405,9 +405,9 @@ def convert_to_full_dm(dm: np.ndarray) -> np.ndarray:
     >>> c.add_coupling(states=(g, e), rabi_frequency=1, detuning=1)
     >>> sols = rq.solve_steady_state(c)
     >>> print(sols.rho)
-    [0.001371 0.02613  0.000686]
+    [0.001313 0.02558  0.000664]
     >>> print(rq.sensor_utils.convert_to_full_dm(sols.rho))
-    [9.993144e-01 1.371165e-03 2.612972e-02 6.855828e-04]
+    [9.993359e-01 1.31399e-03 2.55811e-02 6.64016e-04]
 
     """
     b, r = divmod(math.sqrt(dm.shape[-1]+1), 1.0)
@@ -461,10 +461,10 @@ def convert_dm_to_complex(dm: np.ndarray) -> np.ndarray:
     >>> c.add_coupling(states=(g, e), rabi_frequency=1, detuning=1)
     >>> sols = rq.solve_steady_state(c)
     >>> print(sols.rho)
-    [0.001371 0.02613  0.000686]
+    [0.001313 0.02558  0.000664]
     >>> print(rq.sensor_utils.convert_dm_to_complex(sols.rho))
-    [[9.993144e-01+0.j      1.371166e-03+0.02613j]
-     [1.371166e-03-0.02613j 6.855828e-04+0.j     ]]
+    [[9.993359e-01+0.j      1.31399e-03+0.025581j]
+     [1.313990e-03-0.025581j 6.64016e-04+0.j     ]]
 
     """
     b, r = divmod(math.sqrt(dm.shape[-1]+1), 1.0)