Hotfix release v2.1.3

docs/source/changelog.rst

@@ -1,6 +1,28 @@
 Changelog
 =========
 
+v2.1.3
+------
+
+Improvements
+++++++++++++
+
+- Add `kmag_detuning_correction` keyword argument to `Cell.add_single_coupling`.
+  Default behavior is to only use the atomic transition frequency to calculate the k-vector magnitude
+  for a coupling with `kunit` provided,
+  but this leads to inaccurate doppler averages if the coupling detuning is large relative to Doppler shifts.
+  This new argument allows passing in a detuning to be applied to the calculation of the k-vector magnitude
+  for a coupling. 
+
+Bug Fixes
++++++++++
+
+- Fix issue where `Solution.coupling_rabi` would fail for a coupling group
+  that contains non-dipole-allowed transitions (as common in `Cell` with sublevel structure).
+- Fix incorrect `Sensor.dm_basis` labels. E.g. `01_real`, `01_imag` is actually `10_real`, `10_imag`.
+  For real terms, change is purely cosmetic, but imaginary terms there is a sign difference.
+  This change brings `dm_basis` labels to agree with the actual basis definitions (ie `Solution.rho_ij`).
+
 v2.1.2
 ------
 

docs/source/dev/tests.rst

@@ -67,6 +67,8 @@ The markers we use are
     - Marks a test that incorporates Doppler averaging.
   * - experiments
     - Marks a test that represents a full experiment.
+  * - solution
+    - Marks a test of the `Solution` object and methods.
   * - util
     - Marks a test of the ancillary utilities.
   * - structure

docs/source/examples/Analytical Doppler Demo.ipynb

@@ -239,12 +239,12 @@
       "Rel(diff) max 1.137e+03, mean 0.784, std 22.4\n",
       "Not close matrix elements 593 out of 9045 total\n",
       "Not close solutions 399 out of 603 total\n",
-      "\tdm element 01_real-[1] has 220 misses: Abs-diff (max, mean, diff) 1.797e-04, 5.971e-05, 4.712e-05\n",
-      "\tdm element 02_real-[2] has 100 misses: Abs-diff (max, mean, diff) 1.615e-04, 3.366e-05, 3.979e-05\n",
-      "\tdm element 03_real-[7] has 112 misses: Abs-diff (max, mean, diff) 1.604e-04, 4.080e-05, 3.453e-05\n",
-      "\tdm element 01_imag-[9] has 73 misses: Abs-diff (max, mean, diff) 1.082e-04, 3.516e-05, 2.708e-05\n",
+      "\tdm element 10_real-[1] has 220 misses: Abs-diff (max, mean, diff) 1.797e-04, 5.971e-05, 4.712e-05\n",
+      "\tdm element 20_real-[2] has 100 misses: Abs-diff (max, mean, diff) 1.615e-04, 3.366e-05, 3.979e-05\n",
+      "\tdm element 30_real-[7] has 112 misses: Abs-diff (max, mean, diff) 1.604e-04, 4.080e-05, 3.453e-05\n",
+      "\tdm element 10_imag-[9] has 73 misses: Abs-diff (max, mean, diff) 1.082e-04, 3.516e-05, 2.708e-05\n",
       "\tdm element 11_real-[11] has 49 misses: Abs-diff (max, mean, diff) 1.221e-04, 2.376e-05, 2.820e-05\n",
-      "\tdm element 12_real-[14] has 39 misses: Abs-diff (max, mean, diff) 9.864e-05, 2.241e-05, 2.532e-05\n"
+      "\tdm element 21_real-[14] has 39 misses: Abs-diff (max, mean, diff) 9.864e-05, 2.241e-05, 2.532e-05\n"
      ]
     }
    ],

docs/source/intro_nbs/Cell_Basics.ipynb

@@ -940,7 +940,7 @@
    "metadata": {},
    "source": [
     "### 3.5 Doppler Shifts\n",
-    "Recall that when we want to define a system with doppler broadening in a `Sensor`, we define the `vP` parameter either in the constructor or by accessing the attribute directly after construction, and then specifying the `kvec` parameter in couplings. In a `Cell`, however, the magnitude of the `kvec` parameter is fixed by the transition frequency, and the most probable speed is already set by the temperature. So when adding a coupling with doppler broadening in a `Cell`, we set only the newly-introduced `kunit` when couplings are defined. This is a tuple representing an `(x, y, z)` unit vector representing the direction of the plane wave of the field. Other than that, solving with doppler broadening in a `Cell` is identical to doing so in `Sensor`, and we can otherwise do everything the same way.\n",
+    "Recall that when we want to define a system with doppler broadening in a `Sensor`, we define the `vP` parameter either in the constructor or by accessing the attribute directly after construction, and then specifying the `kvec` parameter in couplings. In a `Cell`, however, the magnitude of the `kvec` parameter is calculated from the transition frequency by default, and the most probable speed is already set by the temperature. So when adding a coupling with doppler broadening in a `Cell`, we set only the newly-introduced `kunit` when couplings are defined. This is a tuple representing an `(x, y, z)` unit vector representing the direction of the plane wave of the field. Other than that, solving with doppler broadening in a `Cell` is identical to doing so in `Sensor`, and we can otherwise do everything the same way.\n",
     "\n",
     "Note that for consistency with `Sensor`, the same `kvec` attribute is added to the graph edge, it is just computed automatically from `kunit`."
    ]
@@ -988,6 +988,16 @@
     "print(rq.solve_steady_state(RbCell_dop, doppler=True).rho)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "4dd4e58f",
+   "metadata": {},
+   "source": [
+    "While using the atomic transition frequency to calculate the magnitude of the k-vector works well in most cases, it leads to inaccuracy in Doppler-averaged calculations if the field detuning is large relative to typical Doppler shifts. In this case, you can use the `kmag_detuning_correction` argument for a coupling to provide the average detuning relative to the atomic transition to be used when calculating the k-vector magnitude.\n",
+    "\n",
+    "Note that passing this argument does not change the `detuning` argument, so this average shift should be accounted for in both places (ie `detuning = laser_scan + offsetFreq` and `kmag_detuning_correction=offsetFreq`)."
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "49f242a2",

docs/source/writeups/observables.rst

@@ -37,7 +37,7 @@ but written in terms of atomic constants :math:`\kappa` and :math:`\eta` (see [2
 We validate that rydiqule calculates these observable quantities in a manner consistent with
 a canonical reference, namely Steck's Quantum and Atom Optics notes [1].
 The unit tests in 
-`test_experiments.py <https://github.com/qtc-umd/rydiqule/blob/main/tests/test_experiments.py>`_
+`test_solution.py <https://github.com/qtc-umd/rydiqule/blob/main/tests/test_solution.py>`_
 test that rydiqule and Steck align,
 but assume that the density matrix element :math:`\rho_{eg}` is correct.
 Validity of density matrix elements is checked in numerous other tests.
@@ -49,7 +49,7 @@ as is done in Steck's Quantum and Atom Optics notes, Eq. 5.273 [1].
 Another way to calculate the scattering rate is using the imaginary term of the susceptibility
 corresponding to the probing field coupling.
 This is the method Rydiqule uses, and can be found in Eq 6.73 of [1]. 
-This test is implemented in `test_experiments.py <https://github.com/qtc-umd/rydiqule/blob/main/tests/test_experiments.py>`_
+This test is implemented in `test_solution.py <https://github.com/qtc-umd/rydiqule/blob/main/tests/test_solution.py>`_
 as the `test_OD_with_steck` unit test.
 
 Running just these observable and experiment unit tests can be done by installing the `pytest` dependencies (see :doc:`Unit Tests <../dev/tests>`),

pyproject.toml

@@ -117,7 +117,8 @@ markers = [
     "steady_state: Marks a test as using the steady-state solver.",
     "time: Marks a test as using the time solver.",
     "doppler: Marks a test as incorporating Doppler averaging.",
-    "experiments: Marks test of experiments.py",
+    "experiments: Marks test that represents a full experiment",
+    "solution: Marks test of Solution object and methods",
     "util: Marks a test of the ancillary utilities.",
     "structure: Marks a test of the definition of the atomic system.",
     "exception: Marks a test of error handling.",

src/rydiqule/cell.py

@@ -391,7 +391,7 @@ def kappa(self) -> float:
             warnings.warn(msg)
 
         # extract q from coupling group, this ensures we find a dipole allowed edge with q defined
-        manifold_subgraph = self._coupling_subgraph(self.probe_tuple)
+        manifold_subgraph = self.coupling_subgraph(self.probe_tuple)
         for (_, _, d) in manifold_subgraph.edges(data=True):
             if 'q' in d:
                 q = d['q']
@@ -489,7 +489,7 @@ def eta(self) -> float:
             warnings.warn(msg)
 
         # extract q from coupling group, this ensures we find a dipole allowed edge with q defined
-        manifold_subgraph = self._coupling_subgraph(self.probe_tuple)
+        manifold_subgraph = self.coupling_subgraph(self.probe_tuple)
         for (_, _, d) in manifold_subgraph.edges(data=True):
             if 'q' in d:
                 q = d['q']
@@ -611,6 +611,7 @@ def add_single_coupling(
             e_field: Optional[ScannableParameter] = None, beam_power: Optional[float] = None,
             beam_waist: Optional[float] = None, coherent_cc: Optional[float]=None,
             q: Literal[-1, 0, 1] = 0,
+            kmag_detuning_correction: Optional[float] = None,
             **extra_kwargs) -> None:
         """
         Overload of :meth:`~.Sensor.add_single_coupling`, which allows for alternate specifications
@@ -710,6 +711,13 @@ def add_single_coupling(
             Coupling polarization in spherical basis.
             Valid values are -1, 0, 1 for :math:`-\\sigma`, linear, :math:`+\\sigma`.
             Default is 0 for linear.
+        kmag_detuning_correction: float, optional
+            Detuning to use when calculating the magnitude of the k-vector.
+            By default, `Cell` uses the transition frequency to define the k-vector.
+            For large detuned couplings, this can lead to inaccurate results.
+            Detuning should be given in units of Mrad/s.
+        **extra_kwargs:
+            Keyword arguments that are passed directly to :meth:`Sensor.add_single_coupling`.
 
         Raises
         ------
@@ -824,9 +832,9 @@ def add_single_coupling(
         not_nlj = (state1.stype != 'NLJ' and state2.stype != 'NLJ')
 
         # check that tuple energy convention matches atomic properties
-        freq_diff = 2*np.pi*self.atom.get_transition_frequency(state1, state2)*1e-6
-        det_sign = np.sign(freq_diff)
-        if np.sign(freq_diff) != 1:
+        ang_freq_diff = 2*np.pi*self.atom.get_transition_frequency(state1, state2)*1e-6
+        det_sign = np.sign(ang_freq_diff)
+        if np.sign(ang_freq_diff) != 1:
             if det_sign > 0:
                 msg = ' higher energy, but it is actually lower. '
             else:
@@ -915,7 +923,14 @@ def add_single_coupling(
                 it is calculated based on atomic properties."""
                 raise RydiquleError(msg)
             else:
-                transition_frequency = freq_diff
+                transition_frequency = ang_freq_diff
+
+        # define coupling angular frequency to use for calculating magnitude of k-vector
+        # by default, uses the atomic transition frequency
+        if kmag_detuning_correction is not None:
+            ang_freq_mean = abs(ang_freq_diff) + kmag_detuning_correction
+        else:
+            ang_freq_mean = abs(ang_freq_diff)
 
         if 'kvec' in extra_kwargs:
             raise RydiquleError("Cell couplings no longer accept 'kvec' as a parameter. " +
@@ -927,9 +942,9 @@ def add_single_coupling(
             # doppler not requested for this coupling, pass default along
             kvec = kunit
         elif np.isclose(k_norm_sq, 1.0):
-            # apply standard doppler shift
-            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
+            # pass full k-vector, accounting for detuning correction if provided
+            kmag = ang_freq_mean / scipy.constants.c # already scaled to Mrad/m
+            kvec = kmag*np.asarray(kunit)
         else:
             raise RydiquleError(f'Coupling {states} has un-normalized |kunit|={math.sqrt(k_norm_sq):.2f}!=1')
 

src/rydiqule/sensor.py

@@ -9,7 +9,8 @@
 import itertools
 
 from .sensor_utils import (ScannableParameter, CouplingDict, State, States, StateSpec, StateSpecs, TimeFunc,
-                           match_states, _squeeze_dims, expand_statespec, state_tuple_to_str, process_scannable_parameter)
+                           match_states, _squeeze_dims, expand_statespec, state_tuple_to_str, process_scannable_parameter,
+                           coupling_subgraph, nx_edges_with)
 from .exceptions import RydiquleError, CouplingNotAllowedError
 from .exceptions import RWAWarning, PopulationNotConservedWarning, RydiquleWarning, debug_state
 
@@ -3019,18 +3020,8 @@ def couplings_with(self, *keys: str,
         >>> print(s.couplings_with("detuning"))
         {(0, 1): {'rabi_frequency': 2, 'detuning': 1, 'phase': 0, 'kvec': (0, 0, 0), 'coherent_cc': 1.0, 'label': '(0,1)'}}
         """
-        def notAll(x):
-            return not all(x)
 
-        def notAny(x):
-            return not any(x)
-
-        methods = {"any":any, "all":all, "not any": notAny, "not all": notAll}
-
-        return {s:p
-                for s,p in self.couplings.edges.items()
-                if methods[method]([k in p for k in keys])
-                }
+        return nx_edges_with(self.couplings, *keys, method=method)
 
 
     def states_with_spec(self, statespec: StateSpec) -> List[State]:
@@ -3068,7 +3059,7 @@ def states_with_spec(self, statespec: StateSpec) -> List[State]:
         return match_states(statespec, self.states)
 
 
-    def _coupling_subgraph(self, coupling: StateSpecs) -> nx.Graph:
+    def coupling_subgraph(self, coupling: StateSpecs) -> nx.Graph:
         """
         Returns a subgraph view of the couplings graph corresponding to `coupling`.
 
@@ -3083,12 +3074,7 @@ def _coupling_subgraph(self, coupling: StateSpecs) -> nx.Graph:
             View of the corresponding subgraph
         """
 
-        states1 = self.states_with_spec(coupling[0])
-        states2 = self.states_with_spec(coupling[1])
-
-        subgraph = self.couplings.subgraph(states1 + states2)
-
-        return subgraph
+        return coupling_subgraph(coupling, self.couplings)
 
 
     def get_couplings(self) -> Dict[States, CouplingDict]:
@@ -3257,11 +3243,11 @@ def dm_basis(self) -> np.ndarray:
         --------
         >>> s = rq.Sensor(3)
         >>> print(s.dm_basis())
-        ['01_real' '02_real' '01_imag' '11_real' '12_real' '02_imag' '12_imag'
+        ['10_real' '20_real' '10_imag' '11_real' '21_real' '20_imag' '21_imag'
          '22_real']
 
         """
-        dm_basis = [f'{j:d}{i:d}_imag' if i > j else f'{i:d}{j:d}_real'
+        dm_basis = [f'{i:d}{j:d}_imag' if i > j else f'{j:d}{i: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)

src/rydiqule/sensor_solution.py

@@ -11,7 +11,6 @@
 import networkx as nx
 from scipy.constants import c
 import warnings
-import itertools
 
 # have to import this way to prevent circular imports
 from rydiqule import sensor_utils
@@ -219,7 +218,8 @@ def coupling_rabi(self, coupling: sensor_utils.StateSpecs) -> ComplexResult:
         Raises
         ------
         ValueError
-            If `coupling` is a group with mismatched rabi frequencies, either in shape or value.
+            If `coupling` is a group with mismatched rabi frequencies, either in shape or value,
+            or coupling has no defined rabi frequencies.
 
         Examples
         --------
@@ -257,11 +257,15 @@ def coupling_rabi(self, coupling: sensor_utils.StateSpecs) -> ComplexResult:
         """
 
         #check all rabis are match, we do this to avoid accidental misuse of the function
-        states1 = sensor_utils.match_states(coupling[0], list(self.couplings.nodes))
-        states2 = sensor_utils.match_states(coupling[1], list(self.couplings.nodes))
+        coupling_subgraph = sensor_utils.coupling_subgraph(coupling, self.couplings)
+
+        nonzero_couplings = sensor_utils.nx_edges_with(coupling_subgraph, 'rabi_frequency')
+
+        if not len(nonzero_couplings):
+            raise ValueError(f'Coupling {coupling} has no edges with defined rabi frequencies')
 
-        rabis = [self.couplings.edges[s1,s2].get("rabi_frequency")
-                 for s1,s2 in itertools.product(states1, states2)]
+        rabis = [edge.get("rabi_frequency")
+                 for edge in nonzero_couplings.values()]
 
         #check all are the same shape
         #cast to array if another data type
@@ -371,10 +375,10 @@ def get_solution_element(self, idx: int) -> Result:
         0.00131399
 
         """
-        basis_size = np.sqrt(self.rho.shape[-1] + 1)
         try:
             sol_slice = self.rho.take(indices=idx, axis=-1)
         except IndexError as err:
+            basis_size = np.sqrt(self.rho.shape[-1] + 1)
             raise RydiquleError(f"No element with given index for {basis_size}-level system") from err
 
         return sol_slice