convert field_insensitive_point to use F/M_F instead of state indices

.github/workflows/docs.yml

@@ -20,7 +20,7 @@ concurrency:
 
 jobs:
   build:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     steps:
       - name: Checkout
         uses: actions/checkout@v2
@@ -51,7 +51,7 @@ jobs:
     environment:
       name: github-pages
       url: ${{ steps.deployment.outputs.page_url }}
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     needs: build
     steps:
       - name: Deploy to GitHub Pages

.github/workflows/package.yml

@@ -9,7 +9,7 @@ on:
 jobs:
   build_and_upload_wheel:
     name: Build and upload wheel
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v3
       - run: pipx install poetry==1.3.2

.github/workflows/test.yml

@@ -9,7 +9,7 @@ on:
 
 jobs:
   test:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     steps:
       - name: Checkout
         uses: actions/checkout@v2

atomic_physics/examples/clock_qubits.py

@@ -8,12 +8,19 @@
 print("Field-independent points:")
 for M3 in range(-3, +3 + 1):
     for q in [-1, 0, 1]:
-        ion = Ca43(1e-4)
-        F4 = ion.get_state_for_F(ca43.ground_level, F=4, M_F=M3 - q)
-        F3 = ion.get_state_for_F(ca43.ground_level, F=3, M_F=M3)
-        B0 = field_insensitive_point(Ca43, (F4, F3))
+        B0 = field_insensitive_point(
+            Ca43,
+            level_0=ca43.ground_level,
+            F_0=4,
+            M_F_0=M3 - q,
+            level_1=ca43.ground_level,
+            F_1=3,
+            M_F_1=M3,
+        )
         if B0 is not None:
             ion = Ca43(B0)
+            F4 = ion.get_state_for_F(ca43.ground_level, F=4, M_F=M3 - q)
+            F3 = ion.get_state_for_F(ca43.ground_level, F=3, M_F=M3)
             f0 = ion.get_transition_frequency_for_states((F4, F3))
             d2fdB2 = d2f_dB2(atom_factory=Ca43, magnetic_field=B0, states=(F4, F3))
             print(

atomic_physics/tests/test_ca43.py

@@ -19,20 +19,29 @@ def test_s12_d52_clock(self):
         S1/2 4,4 -> D5/2 4,3 transition at 3.38 G and 4.96 G with [1].
         """
         Ca43 = ca43.Ca43.filter_levels(level_filter=(ca43.S12, ca43.D52))
-        ion = Ca43(magnetic_field=3e-4)
-
-        l_index = ion.get_state_for_F(ca43.S12, F=4, M_F=+4)
-        u_index = ion.get_state_for_F(ca43.D52, F=4, M_F=+3)
-
         # Start with the field-insensitive transition at around 3G
         B0 = field_insensitive_point(
-            Ca43, (l_index, u_index), magnetic_field_guess=3.38e-4
+            Ca43,
+            level_0=ca43.S12,
+            F_0=4,
+            M_F_0=+4,
+            level_1=ca43.D52,
+            F_1=4,
+            M_F_1=+3,
+            magnetic_field_guess=3.38e-4,
         )
         np.testing.assert_allclose(B0, 3.38e-4, atol=1e-6)
 
         # Field-insensitive transition at around 5G
         B0 = field_insensitive_point(
-            Ca43, (l_index, u_index), magnetic_field_guess=5e-4
+            Ca43,
+            level_0=ca43.S12,
+            F_0=4,
+            M_F_0=+4,
+            level_1=ca43.D52,
+            F_1=4,
+            M_F_1=+3,
+            magnetic_field_guess=5e-4,
         )
         np.testing.assert_allclose(B0, 4.96e-4, atol=1e-6)
 
@@ -41,14 +50,14 @@ def test_s12_40_31_clock(self):
         transition at 146.094G to [1].
         """
         Ca43 = ca43.Ca43.filter_levels(level_filter=(ca43.S12,))
-        ion = Ca43(magnetic_field=146.0942e-4)
-
-        s12_40_index = ion.get_state_for_F(ca43.S12, F=4, M_F=0)
-        s12_31_index = ion.get_state_for_F(ca43.S12, F=3, M_F=+1)
-
         B0 = field_insensitive_point(
             Ca43,
-            (s12_40_index, s12_31_index),
+            level_0=ca43.S12,
+            F_0=4,
+            M_F_0=0,
+            level_1=ca43.S12,
+            F_1=3,
+            M_F_1=+1,
             magnetic_field_guess=140e-4,
         )
         np.testing.assert_allclose(B0, 146.0942e-4, atol=1e-8)
@@ -58,14 +67,14 @@ def test_s12_41_31_clock(self):
         transition at 288G to [1].
         """
         Ca43 = ca43.Ca43.filter_levels(level_filter=(ca43.S12,))
-        ion = Ca43(magnetic_field=287.7827e-4)
-
-        s12_41_index = ion.get_state_for_F(ca43.S12, F=4, M_F=+1)
-        s12_31_index = ion.get_state_for_F(ca43.S12, F=3, M_F=+1)
-
         B0 = field_insensitive_point(
             Ca43,
-            (s12_41_index, s12_31_index),
+            level_0=ca43.S12,
+            F_0=4,
+            M_F_0=+1,
+            level_1=ca43.S12,
+            F_1=3,
+            M_F_1=+1,
             magnetic_field_guess=288e-4,
         )
         np.testing.assert_allclose(B0, 287.7827e-4, atol=1e-8)

atomic_physics/tests/test_mg25.py

@@ -19,13 +19,15 @@ def test_s12_clock(self):
         212.8 G to [1].
         """
         Mg25 = mg25.Mg25.filter_levels(level_filter=(mg25.S12,))
-        ion = Mg25(magnetic_field=212.8e-4)
-
-        l_index = ion.get_state_for_F(mg25.S12, F=3, M_F=+1)
-        u_index = ion.get_state_for_F(mg25.S12, F=2, M_F=+1)
-
         B0 = field_insensitive_point(
-            Mg25, (l_index, u_index), magnetic_field_guess=212.8e-4
+            Mg25,
+            level_0=mg25.S12,
+            F_0=3,
+            M_F_0=+1,
+            level_1=mg25.S12,
+            F_1=2,
+            M_F_1=+1,
+            magnetic_field_guess=212.8e-4,
         )
         np.testing.assert_allclose(B0, 212.8e-4, atol=1e-5)
 

atomic_physics/tests/test_utils.py

@@ -108,14 +108,15 @@ def test_d2f_db2(self):
 
     def test_field_insensitive_point(self):
         """Tests for ``utils.field_insensitive_point``."""
-        ion = ca43.Ca43(magnetic_field=146.0942e-4)
         np.testing.assert_allclose(
             field_insensitive_point(
                 atom_factory=ca43.Ca43,
-                states=(
-                    ion.get_state_for_F(ca43.S12, F=4, M_F=0),
-                    ion.get_state_for_F(ca43.S12, F=3, M_F=+1),
-                ),
+                level_0=ca43.S12,
+                F_0=4,
+                M_F_0=0,
+                level_1=ca43.S12,
+                F_1=3,
+                M_F_1=+1,
                 magnetic_field_guess=10e-4,
             ),
             146.0942e-4,

atomic_physics/utils.py

@@ -2,7 +2,7 @@
 import scipy.constants as consts
 import scipy.optimize as opt
 
-from atomic_physics.core import Atom, AtomFactory, RFDrive, Transition
+from atomic_physics.core import Atom, AtomFactory, Level, RFDrive, Transition
 from atomic_physics.polarization import (
     cartesian_to_spherical,
     dM_for_transition,
@@ -58,29 +58,53 @@ def d2f_dB2(
 
 def field_insensitive_point(
     atom_factory: AtomFactory,
-    states: tuple[int, int],
+    level_0: Level,
+    F_0: float,
+    M_F_0: float,
+    level_1: Level,
+    F_1: float,
+    M_F_1: float,
     magnetic_field_guess: float = 1e-4,
     eps: float = 1e-4,
 ) -> float | None:
     """Returns the magnetic field at which the frequency of a transition
     between two states in the same level becomes first-order field independent.
 
+    Since the energy ordering of states can change with magnetic field, we label states
+    by ``F`` and ``M_F`` instead of using state indices.
+
     :param atom_factory: factory class for the atom of interest.
-    :param states: tuple of indices involved in this transition.
+    :param level_0: level the first state involved in the transition lies in.
+    :param F_0: value of ``F`` for the first state involved in the transition.
+    :param M_F_0: value of ``M_F`` for the first state involved in the transition.
+    :param level_1: level the second state involved in the transition lies in.
+    :param F_1: value of ``F`` for the second state involved in the transition.
+    :param M_F_1: value of ``M_F`` for the second state involved in the transition.
     :param magnetic_field_guess: Initial guess for the magnetic field insensitive point
         (T). This is used both as a seed for the root finding algorithm and as a scale
         factor to help numerical accuracy.
     :param eps: step size as a fraction of ``magnetic_field_guess`` to use when
         calculating numerical derivatives.
     :return: the field-independent point (T) or ``None`` if none found.
     """
-    res = opt.root(
-        lambda magnetic_field: df_dB(
+
+    def opt_fun(x):
+        magnetic_field = max(x, 10 * eps) * magnetic_field_guess
+        atom = atom_factory(magnetic_field=magnetic_field)
+        states = (
+            atom.get_state_for_F(level=level_0, F=F_0, M_F=M_F_0),
+            atom.get_state_for_F(level=level_1, F=F_1, M_F=M_F_1),
+        )
+
+        return df_dB(
             atom_factory,
-            max(magnetic_field, 10 * eps) * magnetic_field_guess,
+            magnetic_field,
             states,
             eps=eps * magnetic_field_guess,
-        ),
+        )
+
+    res = opt.root(
+        opt_fun,
         x0=1,
         options={"xtol": 1e-4, "eps": eps},
     )

docs/changes.rst

@@ -95,4 +95,5 @@ API refactor:
   equations, I've started moving us over to Jones vectors.
 * Add ``operators.expectation_value`` helper method.
 * Add ``Atom.levels`` field.
-* Add new ``Atom.get_states_for_level`` method.
+* Add new ``Atom.get_states_for_level`` method.
+* ``utils.field_insensitive_point`` now works with values of ``F`` and ``M_F`` instead of state indices.