Functions for directly calculating the Rabi frequencies from experime…

…ntal parameters

atomic_physics/beam.py

@@ -27,17 +27,17 @@ class Beam:
     intensity radius.
 
     The beam polarisation is defined using the three angles phi, gamma and theta.
-    Phi is the angle between the beam and the B-field, such that if they are 
+    Phi is the angle between the beam and the B-field, such that if they are
     parallel, phi=0. The beam direction and the B-field then define a plane: we
-    decompose the polarisation vector into two unit vectors, one in this plane 
+    decompose the polarisation vector into two unit vectors, one in this plane
     (k1) and one perpendicular to it (k2). The beam polarisation is then defined
-    as 
+    as
         e = cos(gamma) k1 + exp(i theta) sin(gamma) k2
     Gamma = 0 gives linear polarisation that has maximal projection along the B-field.
-    Gamma = +- pi/4, theta = 0 gives linear polarisation along +-45 deg to the 
+    Gamma = +- pi/4, theta = 0 gives linear polarisation along +-45 deg to the
         plane.
     Gamma = pi/4, theta = +-pi/2 give circular polarisations.
-    
+
     :param wavelength: The wavelength of the light in m.
     :param waist_radius: The (minimum) waist radius of the beam in m.
     :param power: The power in the beam in W.
@@ -46,15 +46,18 @@ class Beam:
         and the plane described by the B-field and beam vector.
     :param theta: The phase angle in rad between the polarisation component in the
         plane and out of the plane."""
+
     wavelength: float
     waist_radius: float
     power: float
     phi: float
     gamma: float
     theta: float
     I_peak: float = field(init=False)
-    E_field: float = field(init=False)  # The magnitude of the E-field at the beam centre
-    rayleigh_range: float = field(init=False) 
+    E_field: float = field(
+        init=False
+    )  # The magnitude of the E-field at the beam centre
+    rayleigh_range: float = field(init=False)
     omega: float = field(init=False)
 
     def __post_init__(self):
@@ -64,32 +67,32 @@ def __post_init__(self):
         self.omega = consts.c / self.wavelength * 2 * np.pi
 
     def beam_direction_cartesian(self):
-        """Return the unit vector along the beam direction in Cartesian 
-            coordinates [x, y, z]"""
+        """Return the unit vector along the beam direction in Cartesian
+        coordinates [x, y, z]"""
         return [np.sin(self.phi), 0, np.cos(self.phi)]
 
     def polarisation_cartesian(self):
         """Return a list of the Cartesian components of the (complex) polarisation
-            vector [x, y, z]"""
+        vector [x, y, z]"""
         return [
             np.cos(self.phi) * np.cos(self.gamma),
             cmath.exp(1j * self.theta) * np.sin(self.gamma),
             -np.sin(self.phi) * np.cos(self.gamma),
-            ]
-    
+        ]
+
     def polarisation_rank1_polar(self):
         """
         Returns the rank-1 tensor components of the polarisation vector that drive
-        dipole transitions, i.e. the sigma_m, sigma_p and pi components of the 
+        dipole transitions, i.e. the sigma_m, sigma_p and pi components of the
         polarisation vector. Output is a list of the form [sigma_m, pi, sigma_p]
         in polar form.
         """
         return [cmath.polar(q) for q in self.polarisation_rank1()]
-    
+
     def polarisation_rank1(self):
         """
         Returns the rank-1 tensor components of the polarisation vector that drive
-        dipole transitions, i.e. the sigma_m, sigma_p and pi components of the 
+        dipole transitions, i.e. the sigma_m, sigma_p and pi components of the
         polarisation vector. Output is a list of the form [sigma_m, pi, sigma_p]
         in rectangular form (i.e. re + j im).
         """
@@ -98,24 +101,24 @@ def polarisation_rank1(self):
         pi = np.sum(vect * C1_0)
         sigma_p = np.sum(vect * C1_P1)
         return [sigma_m, pi, sigma_p]
-    
+
     def polarisation_rank2_polar(self):
         """
-        Returns the rank-2 tensor components of the polarisation tensor, which 
+        Returns the rank-2 tensor components of the polarisation tensor, which
         drive quadrupole transitions. Output is a list containing the five components
         in increasing order from q_m2 to q_p2 in polar form.
         """
         return [cmath.polar(q) for q in self.polarisation_rank2()]
-    
+
     def polarisation_rank2(self):
         """
-        Returns the rank-2 tensor components of the polarisation tensor, which 
+        Returns the rank-2 tensor components of the polarisation tensor, which
         drive quadrupole transitions. Output is a list containing the five components
         in increasing order from q_m2 to q_p2 in rectangular form (i.e. re + j im).
         """
         beam_vect = np.array(self.beam_direction_cartesian())
         pol_vect = np.array(self.polarisation_cartesian())
-        
+
         q_m2 = pol_vect.T @ C2_M2 @ beam_vect
         q_m1 = pol_vect.T @ C2_M1 @ beam_vect
         q_0 = pol_vect.T @ C2_0 @ beam_vect
@@ -134,11 +137,10 @@ def print_rank1(self, title=None):
         print(f"pi: {pi:.3f}")
         print(f"sigma_p: {sigma_p:.3f}")
 
-
     def print_rank2(self, title=None):
         """Pretty-print the rank-2 tensor components of the beam."""
         if title is None:
             title = ""
         print(title)
         for idx, p in enumerate(self.polarisation_rank1_rect()):
-            print(f"q_{idx-2}: {p:.3f}")
+            print(f"q_{idx-2}: {p:.3f}")

atomic_physics/common.py

@@ -32,6 +32,12 @@
      atom.delta)
 """
 
+State = namedtuple("State", ["level", "F", "M"])
+State.__doc__ = """Convenient holder for storing atomic state data.
+    :param level: the Level object that the state is in.
+    :param F: total angular momentum number.
+    :param M: z-projection of F."""
+
 
 class LevelData:
     """Stored atomic structure information about a single level."""
@@ -131,6 +137,11 @@ def __init__(
 
         self.levels = levels
         self.transitions = transitions
+        # Store the transitions indexed by the upper and lower levels it connects,
+        # useful fro reverse searching
+        self.reverse_transitions = {
+            (t.lower, t.upper): t for t in self.transitions.values()
+        }
 
         # ordered in terms of increasing state energies
         self.num_states = None  # Total number of electronic states
@@ -422,9 +433,7 @@ def calc_Epole(self):
         since they are somewhat more convenient to work with. The two are
         related by constants and power of the transition frequencies.
 
-        To do: define what we mean by a matrix element, and how the amplitudes
-        stored in ePole are related to the matrix elements and to the Rabi
-        frequencies.
+        For converting to Rabi frequencies, use the functions in atomic_physics.rabi.
 
         To do: double check the signs of the amplitudes.
         """

atomic_physics/rabi.py

@@ -0,0 +1,231 @@
+from .common import Atom, State, Level
+import numpy as np
+from .beam import Beam
+import scipy.constants as consts
+
+__doc__ = """
+    Calculate experimentally observable Rabi frequencies and light shifts (AC
+    Stark shifts) in experimentally useful units. Essentially just does the relevant
+    unit conversion from atom.ePole and takes into account the beam power and
+    geometry.
+    """
+
+
+def scattering_amplitude_to_dipole_matrix_element_prefactor(wavelength: float):
+    """
+    Get the prefactor to convert a scattering amplitude (the scale that
+    is used in atom.ePole) to a matrix element for a dipole transition.
+
+    Taken from https://doi.org/10.1007/s003400050373
+
+    :param wavelength: The wavelength of the transition.
+    """
+    return np.sqrt(
+        3 * consts.epsilon_0 * consts.hbar * wavelength**3 / (8 * np.pi**2)
+    )
+
+
+def scattering_amplitude_to_quadrupole_matrix_element_prefactor(wavelength: float):
+    """
+    Get the prefactor to convert a scattering amplitude (the scale that
+    is used in atom.ePole) to a matrix element for a quadrupole transition.
+
+    Taken from https://doi.org/10.1007/s003400050373
+
+    :param wavelength: The wavelength of the transition.
+    """
+    return np.sqrt(
+        15 * consts.epsilon_0 * consts.hbar * wavelength**3 / (8 * np.pi**2)
+    )
+
+
+def laser_rabi_omega(state1: State, state2: State, atom: Atom, beam: Beam):
+    """
+    Return the (complex-valued) on-resonance Rabi frequency of the transition
+    between state1 and state2 when driven by a laser beam in rad/s. The
+    transition can be a dipole (E1) or quadrupole (E2) one.
+
+    NOTE: The Rabi frequency returned has a phase arises from the Wigner3J
+    symbols and the complex laser polarisation. This is useful to keep for
+    more complex phase-dependent interactions such as light shifts and Raman
+    transitions.
+
+    TODO: The quadrupole Rabi frequency disagreed with the experimentally measured
+    one in ABaQuS by a factor of ~2.3 in Dec 2024. Whether this is due to not fully
+    known experimental params or errors in the maths is not clear. Check quadrupole
+    Rabi frequency.
+
+    :param state1: One of the atomic states in the transition.
+    :param state2: The other atomic state in the transition.
+    :param atom: The Atom object. ePole must have been pre-calculated.
+    :param beam: Beam object containing the parameters of the driving
+        beam.
+    """
+
+    assert (
+        atom.ePole is not None
+    ), "atom.ePole has not been calculated. Has the magnetic field been set?"
+    idx1 = atom.index(state1.level, M=state1.M, F=state1.F)
+    idx2 = atom.index(state2.level, M=state2.M, F=state2.F)
+    if idx2 > idx1:
+        idx_u = idx2
+        idx_l = idx1
+        state_u = state2
+        state_l = state1
+    else:
+        idx_u = idx1
+        idx_l = idx2
+        state_u = state1
+        state_l = state2
+
+    q = state_u.M - state_l.M
+
+    assert (state_l.level, state_u.level,) in atom.reverse_transitions.keys(), (
+        "No dipole or quadrupole transition between"
+        + f" {state_l.level} and {state_u.level} exists."
+    )
+    transition = atom.reverse_transitions[(state_l.level, state_u.level)]
+    transition_freq = transition.freq + atom.delta(idx_l, idx_u)
+    wavelength = 2 * np.pi * consts.c / transition_freq
+
+    if abs(state_u.level.L - state_l.level.L) == 1:
+        C = scattering_amplitude_to_dipole_matrix_element_prefactor(wavelength)
+        if abs(q) > 1:
+            return 0.0
+        polarisation_coeff = beam.polarisation_rank1()[int(q + 1)]
+    elif abs(state_u.level.L - state_l.level.L) == 2:
+        C = scattering_amplitude_to_quadrupole_matrix_element_prefactor(wavelength)
+        if abs(q) > 2:
+            return 0.0
+        polarisation_coeff = beam.polarisation_rank2()[int(q + 2)]
+    else:
+        raise ValueError(
+            "Rabi frequencies for lasers can only be calculated for dipole or"
+            + " quadrupole transitions. A transition with delta_L ="
+            + f" {abs(state_u.level.L - state_l.level.L)} was provided."
+        )
+
+    prefactor = C * beam.E_field / consts.hbar * polarisation_coeff
+    return prefactor * atom.ePole[idx_l, idx_u]
+
+
+def raman_rabi_omega(
+    state1: State,
+    state2: State,
+    atom: Atom,
+    beam_red: Beam,
+    beam_blue: Beam,
+    virtual_levels: list[Level],
+):
+    """
+    Return the (complex-valued) Rabi frequency  in rad/s of the transition between
+    state1 and state2 when driven by a Raman transition. Assumes that these are
+    within the same level.
+
+    NOTE: The Rabi frequency returned has a phase arises from the Wigner3J
+    symbols and the complex laser polarisation. This is useful to keep for
+    more complex phase-dependent interactions such as light shifts.
+
+    :param state1: One of the atomic states in the transition.
+    :param state2: The other atomic state in the transition.
+    :param atom: The Atom object. ePole must have been pre-calculated.
+    :param beam_red: Beam object containing the parameters of the
+        red-detuned driving beam.
+    :param beam_red: Beam object containing the parameters of the
+        blue-detuned driving beam.
+    :param levels: List of all levels through which the virtual transition
+        can occur (e.g. for transitions between two states in an S1/2 level,
+        these would be the P1/2 and P3/2 levels). Only dipole transitions are
+        considered for this.
+    """
+
+    assert (
+        state1.level == state2.level
+    ), "Raman transitions between states in different levels not supported"
+
+    # It's important to determine which beam "drives" which virtual transition
+    # out of state1 <-> virt. state and state2 <-> virt. state if the two beams
+    # have different polarisations.
+    # The red-detuned beam drives the transition from the higher of state1,2 and
+    # the blue-detuned one drives the other one.
+    idx_1 = atom.index(state1.level, M=state1.M, F=state1.F)
+    idx_2 = atom.index(state2.level, M=state2.M, F=state2.F)
+    if atom.E[idx_2] > atom.E[idx_1]:
+        state_high = state2
+        state_low = state1
+    else:
+        state_high = state1
+        state_low = state2
+
+    I = atom.I
+    rabi_omega = 0
+
+    for v_level in virtual_levels:
+        # Check whether the virtual level is above or below the level that each
+        # of the two states are in, to correctly determine the sign of delta_m
+        # for each arm.
+        if (state_high.level, v_level) in atom.reverse_transitions.keys():
+            v_transition = atom.reverse_transitions[(state_high.level, v_level)]
+        else:
+            v_transition = atom.reverse_transitions[(v_level, state_high.level)]
+
+        v_transition_delta = beam_red.omega - v_transition.freq
+
+        v_J = v_level.J
+        for F_int in np.arange(abs(I - v_J), I + v_J + 1):
+            for m_int in np.arange(-F_int, F_int + 1):
+                state_int = State(level=v_level, M=m_int, F=F_int)
+                rabi_omega_red = laser_rabi_omega(state_high, state_int, atom, beam_red)
+                rabi_omega_blue = laser_rabi_omega(
+                    state_low, state_int, atom, beam_blue
+                )
+                rabi_omega += rabi_omega_red * rabi_omega_blue
+
+        rabi_omega *= 1 / (4 * v_transition_delta)
+    return rabi_omega
+
+
+def far_detuned_light_shift(
+    state: State,
+    atom: Atom,
+    beam: Beam,
+    levels: list[Level],
+):
+    """
+    Return the (complex) light shift frequency in rad/s on a state due to very
+    far off-resonant transitions to a set of levels (which may occur when driving
+    a Raman or quadrupole transition, for example).
+
+    :param state: The State to calculate the light-shift for.
+    :param beam: The Beam object that causes the light shift.
+    :param atom: The Atom object. ePole must have been pre-calculated.
+    :param levels: List of atomic levels that induce a light shift on ```state```.
+    """
+    state_level = state.level
+    I = atom.I
+    light_shift = 0
+    # Since the transition is very off-resonant, say that all the transitions have the
+    # same detuning
+    laser_omega = beam.omega
+    for aux_level in levels:
+        transition = atom.reverse_transitions.get((state_level, aux_level), None)
+        transition = atom.reverse_transitions.get((aux_level), transition)
+
+        # Correct for sign if the state happens to be the top state
+        if state_level == transition.lower:
+            sign = 1
+        else:
+            sign = -1
+        # Detuning has the sign such that a laser with lower frequency than the
+        # transition pushes the state energy down, which occurs when detuning < 0
+        detuning = laser_omega - transition.freq
+        aux_J = aux_level.J
+        for F_aux in np.arange(abs(I - aux_J), I + aux_J + 1):
+            for m_aux in np.arange(-F_aux, F_aux + 1):
+                delta_m = m_aux - state.M
+                if abs(delta_m) <= 1:
+                    state2 = State(level=aux_level, F=F_aux, M=m_aux)
+                    rabi = laser_rabi_omega(state, state2, atom, beam)
+                    light_shift += sign * abs(rabi) ** 2 / (4 * detuning)
+
+    return light_shift