Restructure rydberg state

README.md

@@ -83,7 +83,7 @@ This package relies on quantum defects provided by the community. Consider citin
 ## Using custom quantum defects
 To use custom quantum defects (or quantum defects for a new species), you can simply create a subclass of `rydstate.species.species_object.SpeciesObject` (e.g. `class CustomRubidium(SpeciesObject):`) with a custom species name (e.g. `name = "Custom_Rb"`).
 Then, similarly to `rydstate.species.rubidium.py` you can define the quantum defects (and model potential parameters, ...) for your species.
-Finally, you can use the custom species by simply calling `rydstate.RydbergStateAlkali("Custom_Rb", n=50, l=0, j=1/2, m=1/2)` (the code will look for all subclasses of `SpeciesObject` until it finds one with the species name "Custom_Rb").
+Finally, you can use the custom species by simply calling `rydstate.RydbergStateSQDTAlkali("Custom_Rb", n=50, l=0, j=1/2, m=1/2)` (the code will look for all subclasses of `SpeciesObject` until it finds one with the species name "Custom_Rb").
 
 
 ## License

docs/examples.rst

@@ -2,10 +2,10 @@ Examples
 ========
 
 
-RadialState
------------
+RadialKet
+---------
 
-Some examples demonstrating the usage of the RadialState class, which uses the Numerov method for solving the radial Schrödinger equation.
+Some examples demonstrating the usage of the RadialKet class, which uses the Numerov method for solving the radial Schrödinger equation.
 
 .. nbgallery::
    examples/radial/hydrogen_wavefunction

docs/examples/benchmark/benchmark_njit.ipynb

@@ -17,7 +17,7 @@
     "\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali\n",
+    "from rydstate import RydbergStateSQDTAlkali\n",
     "\n",
     "test_cases: list[tuple[str, int, int, bool]] = [\n",
     "    # species, n, l, use_njit\n",
@@ -28,7 +28,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -41,21 +41,21 @@
     "    \"\"\"\n",
     "    # run the integration once to compile the numba function\n",
     "    species, n, l, use_njit = test_cases[0]\n",
-    "    state = RydbergStateAlkali(species, n, l, j=l + 0.5)\n",
+    "    state = RydbergStateSQDTAlkali(species, n, l, j=l + 0.5)\n",
     "    state.radial.create_wavefunction(_use_njit=True)\n",
     "\n",
     "    results = []\n",
     "    for species, n, l, use_njit in test_cases:\n",
     "        # Setup the test function\n",
     "        stmt = (\n",
-    "            \"state = RydbergStateAlkali(species, n, l, j=l+0.5)\\n\"\n",
+    "            \"state = RydbergStateSQDTAlkali(species, n, l, j=l+0.5)\\n\"\n",
     "            \"state.radial.create_grid(dz=1e-3)\\n\"\n",
     "            \"state.radial.create_wavefunction(_use_njit=use_njit)\"\n",
     "        )\n",
     "\n",
     "        # Time the integration multiple times and take average/std\n",
     "        globals_dict = {\n",
-    "            \"RydbergStateAlkali\": RydbergStateAlkali,\n",
+    "            \"RydbergStateSQDTAlkali\": RydbergStateSQDTAlkali,\n",
     "            \"species\": species,\n",
     "            \"n\": n,\n",
     "            \"l\": l,\n",

docs/examples/comparisons/compare_model_potentials.ipynb

@@ -9,13 +9,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali, RydbergStateAlkalineLS"
+    "from rydstate import RydbergStateSQDTAlkali, RydbergStateSQDTAlkalineLS"
    ]
   },
   {
@@ -41,15 +41,15 @@
     }
    ],
    "source": [
-    "state = RydbergStateAlkali(\"Rb\", n=40, l=0, j=0.5)\n",
+    "state = RydbergStateSQDTAlkali(\"Rb\", n=40, l=0, j=0.5)\n",
     "\n",
-    "states: dict[str, RydbergStateAlkali] = {}\n",
+    "states: dict[str, RydbergStateSQDTAlkali] = {}\n",
     "\n",
     "\n",
-    "states[\"model_potential_marinescu_1993\"] = RydbergStateAlkali(state.species, n=state.n, l=state.l, j=state.j)\n",
+    "states[\"model_potential_marinescu_1993\"] = RydbergStateSQDTAlkali(state.species, n=state.n, l=state.l, j=state.j)\n",
     "states[\"model_potential_marinescu_1993\"].radial.create_model(potential_type=\"model_potential_marinescu_1993\")\n",
     "\n",
-    "states[\"model_potential_fei_2009\"] = RydbergStateAlkali(state.species, n=state.n, l=state.l, j=state.j)\n",
+    "states[\"model_potential_fei_2009\"] = RydbergStateSQDTAlkali(state.species, n=state.n, l=state.l, j=state.j)\n",
     "states[\"model_potential_fei_2009\"].radial.create_model(potential_type=\"model_potential_fei_2009\")\n",
     "\n",
     "for label, state in states.items():\n",
@@ -122,17 +122,17 @@
     }
    ],
    "source": [
-    "state = RydbergStateAlkalineLS(\"Sr88\", n=8, l=0, j_tot=0, s_tot=0)\n",
+    "state = RydbergStateSQDTAlkalineLS(\"Sr88\", n=8, l=0, j_tot=0, s_tot=0)\n",
     "\n",
-    "states: dict[str, RydbergStateAlkalineLS] = {}\n",
+    "states: dict[str, RydbergStateSQDTAlkalineLS] = {}\n",
     "\n",
     "\n",
-    "states[\"model_potential_marinescu_1993\"] = RydbergStateAlkalineLS(\n",
+    "states[\"model_potential_marinescu_1993\"] = RydbergStateSQDTAlkalineLS(\n",
     "    state.species, n=state.n, l=state.l, j_tot=state.j_tot, s_tot=state.s_tot\n",
     ")\n",
     "states[\"model_potential_marinescu_1993\"].radial.create_model(potential_type=\"model_potential_marinescu_1993\")\n",
     "\n",
-    "states[\"model_potential_fei_2009\"] = RydbergStateAlkalineLS(\n",
+    "states[\"model_potential_fei_2009\"] = RydbergStateSQDTAlkalineLS(\n",
     "    state.species, n=state.n, l=state.l, j_tot=state.j_tot, s_tot=state.s_tot\n",
     ")\n",
     "states[\"model_potential_fei_2009\"].radial.create_model(potential_type=\"model_potential_fei_2009\")\n",

docs/examples/comparisons/compare_whittaker.ipynb

@@ -21,7 +21,7 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate import RydbergStateAlkali\n",
+    "from rydstate import RydbergStateSQDTAlkali\n",
     "\n",
     "logging.basicConfig(level=logging.INFO, format=\"%(levelname)s %(filename)s: %(message)s\")"
    ]
@@ -41,7 +41,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "states: dict[str, RydbergStateAlkali] = {}"
+    "states: dict[str, RydbergStateSQDTAlkali] = {}"
    ]
   },
   {
@@ -62,15 +62,15 @@
     }
    ],
    "source": [
-    "state = RydbergStateAlkali(\"Rb\", n=21, l=0, j=0.5)\n",
+    "state = RydbergStateSQDTAlkali(\"Rb\", n=21, l=0, j=0.5)\n",
     "\n",
     "state.radial.create_model(potential_type=\"model_potential_marinescu_1993\")\n",
     "state.radial.create_wavefunction(\"numerov\")\n",
     "states[\"Numerov with Model Potentials\"] = state\n",
     "\n",
     "# Using Numerov without model potentials will lead to some warnings,\n",
     "# since the resulting wavefunction does not pass all heuristic checks\n",
-    "state_without_mp = RydbergStateAlkali(state.species, state.n, state.l, state.j)\n",
+    "state_without_mp = RydbergStateSQDTAlkali(state.species, state.n, state.l, state.j)\n",
     "state_without_mp.radial.create_model(potential_type=\"coulomb\")\n",
     "state_without_mp.radial.create_wavefunction(\"numerov\")\n",
     "states[\"Numerov without Model Potentials\"] = state_without_mp"
@@ -91,7 +91,7 @@
     }
    ],
    "source": [
-    "state_whittaker = RydbergStateAlkali(state.species, state.n, state.l, state.j)\n",
+    "state_whittaker = RydbergStateSQDTAlkali(state.species, state.n, state.l, state.j)\n",
     "state_whittaker.radial.create_grid(x_min=state.radial.grid.x_min, x_max=state.radial.grid.x_max)\n",
     "state_whittaker.radial.create_wavefunction(\"whittaker\")\n",
     "states[\"Whittaker\"] = state_whittaker"
@@ -188,15 +188,15 @@
     }
    ],
    "source": [
-    "state1 = RydbergStateAlkali(\"Rb\", n=10, l=0, j=0.5)\n",
-    "state2 = RydbergStateAlkali(\"Rb\", n=9, l=1, j=1.5)\n",
+    "state1 = RydbergStateSQDTAlkali(\"Rb\", n=10, l=0, j=0.5)\n",
+    "state2 = RydbergStateSQDTAlkali(\"Rb\", n=9, l=1, j=1.5)\n",
     "\n",
     "dipole_me = state1.radial.calc_matrix_element(state2.radial, 1)\n",
     "print(f\"Numerov with model potentials: {dipole_me}\", flush=True)\n",
     "\n",
-    "_state1 = RydbergStateAlkali(state1.species, state1.n, state1.l, state1.j)\n",
+    "_state1 = RydbergStateSQDTAlkali(state1.species, state1.n, state1.l, state1.j)\n",
     "_state1.radial.create_model(potential_type=\"coulomb\")\n",
-    "_state2 = RydbergStateAlkali(state2.species, state2.n, state2.l, state2.j)\n",
+    "_state2 = RydbergStateSQDTAlkali(state2.species, state2.n, state2.l, state2.j)\n",
     "_state2.radial.create_model(potential_type=\"coulomb\")\n",
     "\n",
     "dipole_me = _state1.radial.calc_matrix_element(_state2.radial, 1)\n",
@@ -206,10 +206,10 @@
     "# to avoid integrating over the diverging peak at the origin (see plots above)\n",
     "xmin1, xmax1 = _state1.radial.grid.x_min, _state1.radial.grid.x_max\n",
     "xmin2, xmax2 = _state2.radial.grid.x_min, _state2.radial.grid.x_max\n",
-    "_state1 = RydbergStateAlkali(state1.species, state1.n, state1.l, state1.j)\n",
+    "_state1 = RydbergStateSQDTAlkali(state1.species, state1.n, state1.l, state1.j)\n",
     "_state1.radial.create_grid(x_min=xmin1, x_max=xmax1)\n",
     "_state1.radial.create_wavefunction(\"whittaker\")\n",
-    "_state2 = RydbergStateAlkali(state2.species, state2.n, state2.l, state2.j)\n",
+    "_state2 = RydbergStateSQDTAlkali(state2.species, state2.n, state2.l, state2.j)\n",
     "_state2.radial.create_grid(x_min=xmin2, x_max=xmax2)\n",
     "_state2.radial.create_wavefunction(\"whittaker\")\n",
     "\n",

docs/examples/comparisons/compare_z_min_cutoff.ipynb

@@ -9,13 +9,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali"
+    "from rydstate import RydbergStateSQDTAlkali"
    ]
   },
   {
@@ -56,7 +56,7 @@
     "z_i_dict = {\"hydrogen\": [], \"classical\": [], \"rydstate cutoff\": []}\n",
     "for qn in qn_list:\n",
     "    print(f\"n={qn[0]}\", end=\"\\r\")\n",
-    "    state = RydbergStateAlkali(\"Rb\", n=qn[0], l=qn[1], j=qn[2])\n",
+    "    state = RydbergStateSQDTAlkali(\"Rb\", n=qn[0], l=qn[1], j=qn[2])\n",
     "\n",
     "    hydrogen_z_i = state.radial.model.calc_hydrogen_turning_point_z(state.n, state.l)\n",
     "    z_i_dict[\"hydrogen\"].append(hydrogen_z_i)\n",

docs/examples/dipole_matrix_elements.ipynb

@@ -9,13 +9,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import numpy as np\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali"
+    "from rydstate import RydbergStateSQDTAlkali"
    ]
   },
   {
@@ -34,8 +34,8 @@
     }
    ],
    "source": [
-    "state_i = RydbergStateAlkali(\"Rb\", 60, 2, j=3 / 2, m=1 / 2)\n",
-    "state_f = RydbergStateAlkali(\"Rb\", 60, 3, j=5 / 2, m=1 / 2)\n",
+    "state_i = RydbergStateSQDTAlkali(\"Rb\", 60, 2, j=3 / 2, m=1 / 2)\n",
+    "state_f = RydbergStateSQDTAlkali(\"Rb\", 60, 3, j=5 / 2, m=1 / 2)\n",
     "\n",
     "kappa = 1\n",
     "radial = state_i.radial.calc_matrix_element(state_f.radial, k_radial=1)\n",

docs/examples/radial/hydrogen_wavefunction.ipynb

@@ -9,7 +9,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -18,7 +18,7 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate.radial import RadialState"
+    "from rydstate.radial import RadialKet"
    ]
   },
   {
@@ -27,7 +27,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "state = RadialState(\"H_textbook\", nu=10, l_r=5)\n",
+    "state = RadialKet(\"H_textbook\", nu=10, l_r=5)\n",
     "state.set_n_for_sanity_check(10)\n",
     "state.create_model()\n",
     "state.create_grid(dz=1e-2)\n",

docs/examples/radial/rubidium_wavefunction.ipynb

@@ -18,7 +18,7 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
-    "from rydstate.rydberg_state import RydbergStateAlkali\n",
+    "from rydstate import RydbergStateSQDTAlkali\n",
     "\n",
     "logging.basicConfig(level=logging.INFO, format=\"%(levelname)s %(filename)s: %(message)s\")\n",
     "logging.getLogger(\"rydstate\").setLevel(logging.DEBUG)"
@@ -30,7 +30,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "state = RydbergStateAlkali(\"Rb\", n=130, l=129, j=129.5)\n",
+    "state = RydbergStateSQDTAlkali(\"Rb\", n=130, l=129, j=129.5)\n",
     "state.radial.create_wavefunction()\n",
     "\n",
     "turning_points = {\n",
@@ -45,7 +45,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "hydrogen = RydbergStateAlkali(\"H_textbook\", n=state.n, l=state.l, j=state.j)\n",
+    "hydrogen = RydbergStateSQDTAlkali(\"H_textbook\", n=state.n, l=state.l, j=state.j)\n",
     "hydrogen.radial.create_model()\n",
     "hydrogen.radial.create_wavefunction()"
    ]

src/rydstate/__init__.py

@@ -1,11 +1,17 @@
 from rydstate import angular, radial, species
-from rydstate.rydberg_state import RydbergStateAlkali, RydbergStateAlkalineJJ, RydbergStateAlkalineLS
+from rydstate.rydberg import (
+    RydbergStateSQDT,
+    RydbergStateSQDTAlkali,
+    RydbergStateSQDTAlkalineJJ,
+    RydbergStateSQDTAlkalineLS,
+)
 from rydstate.units import ureg
 
 __all__ = [
-    "RydbergStateAlkali",
-    "RydbergStateAlkalineJJ",
-    "RydbergStateAlkalineLS",
+    "RydbergStateSQDT",
+    "RydbergStateSQDTAlkali",
+    "RydbergStateSQDTAlkalineJJ",
+    "RydbergStateSQDTAlkalineLS",
     "angular",
     "radial",
     "species",

src/rydstate/angular/angular_ket.py

@@ -76,7 +76,7 @@ class AngularKetBase(ABC):
     """Total atom angular quantum number (including nuclear, core electron and rydberg electron contributions)."""
     m: float | None
     """Magnetic quantum number, which is the projection of `f_tot` onto the quantization axis.
-    If None, only reduced matrix elements can be calculated
+    If None, only reduced matrix elements can be calculated.
     """
 
     def __init__(
@@ -708,3 +708,53 @@ def sanity_check(self, msgs: list[str] | None = None) -> None:
             msgs.append(f"{self.f_c=}, {self.j_r=}, {self.f_tot=} don't satisfy spin addition rule.")
 
         super().sanity_check(msgs)
+
+
+def quantum_numbers_to_angular_ket(
+    species: str | SpeciesObject,
+    s_c: float | None = None,
+    l_c: int = 0,
+    j_c: float | None = None,
+    f_c: float | None = None,
+    s_r: float = 0.5,
+    l_r: int | None = None,
+    j_r: float | None = None,
+    s_tot: float | None = None,
+    l_tot: int | None = None,
+    j_tot: float | None = None,
+    f_tot: float | None = None,
+    m: float | None = None,
+) -> AngularKetBase:
+    r"""Return an AngularKet object in the corresponding coupling scheme from the given quantum numbers.
+
+    Args:
+        species: Atomic species.
+        s_c: Spin quantum number of the core electron (0 for Alkali, 0.5 for divalent atoms).
+        l_c: Orbital angular momentum quantum number of the core electron.
+        j_c: Total angular momentum quantum number of the core electron.
+        f_c: Total angular momentum quantum number of the core (core electron + nucleus).
+        s_r: Spin quantum number of the rydberg electron (always 0.5).
+        l_r: Orbital angular momentum quantum number of the rydberg electron.
+        j_r: Total angular momentum quantum number of the rydberg electron.
+        s_tot: Total spin quantum number of all electrons.
+        l_tot: Total orbital angular momentum quantum number of all electrons.
+        j_tot: Total angular momentum quantum number of all electrons.
+        f_tot: Total angular momentum quantum number of the atom (rydberg electron + core).
+        m: Total magnetic quantum number.
+          Optional, only needed for concrete angular matrix elements.
+
+    """
+    if all(qn is None for qn in [j_c, f_c, j_r]):
+        return AngularKetLS(
+            s_c=s_c, l_c=l_c, s_r=s_r, l_r=l_r, s_tot=s_tot, l_tot=l_tot, j_tot=j_tot, f_tot=f_tot, m=m, species=species
+        )
+    if all(qn is None for qn in [s_tot, l_tot, f_c]):
+        return AngularKetJJ(
+            s_c=s_c, l_c=l_c, j_c=j_c, s_r=s_r, l_r=l_r, j_r=j_r, j_tot=j_tot, f_tot=f_tot, m=m, species=species
+        )
+    if all(qn is None for qn in [s_tot, l_tot, j_tot]):
+        return AngularKetFJ(
+            s_c=s_c, l_c=l_c, j_c=j_c, f_c=f_c, s_r=s_r, l_r=l_r, j_r=j_r, f_tot=f_tot, m=m, species=species
+        )
+
+    raise ValueError("Invalid combination of angular quantum numbers provided.")