Remove manual frame trafo

src/rydstate/angular/angular_ket.py

@@ -800,6 +800,15 @@ def _calc_prefactor_of_operator_in_coupled_scheme(
         prefactor = calc_prefactor_of_operator_in_coupled_scheme(f1, f2, f_tot, i1, i2, i_tot, kappa, operator_acts_on)
         return prefactor * self._calc_prefactor_of_operator_in_coupled_scheme(other, qn_combined, kappa)
 
+    def get_core_ket(self) -> CoreKet:
+        """Get the core ket corresponding to this ket, j_c and f_c might be unknown dependent on the coupling scheme."""
+        j_c = self.get_qn("j_c", allow_unknown=True)
+        f_c = self.get_qn("f_c", allow_unknown=True)
+        label = self.label
+        if label is None and (is_unknown(f_c) or is_unknown(j_c)):
+            label = Unknown
+        return CoreKet(i_c=self.i_c, s_c=self.s_c, l_c=self.l_c, j_c=j_c, f_c=f_c, label=label)
+
 
 class AngularKetLS(AngularKetBase[GenericT_Unknown], Generic[GenericT_Unknown]):
     """Spin ket in LS coupling."""
@@ -1128,7 +1137,3 @@ 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 get_core_ket(self) -> CoreKet:
-        """Get the core ket corresponding to this FJ ket."""
-        return CoreKet(i_c=self.i_c, s_c=self.s_c, l_c=self.l_c, j_c=self.j_c, f_c=self.f_c, label=self.label)

src/rydstate/angular/utils.py

@@ -130,7 +130,7 @@ def try_trivial_spin_addition(
 ) -> float | Unknown:
     """Try to determine s_tot from s_1 and s_2 if it is not given.
 
-    If s_tot is None and cannot be uniquely determined from s_1 and s_2, raise an error.
+    If s_tot is None and cannot be uniquely determined from s_1 and s_2, return Unknown.
     Otherwise return s_tot or the trivial sum s_1 + s_2.
     """
     if s_tot is not None and not is_unknown(s_tot):

src/rydstate/basis/basis_mqdt.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -131,7 +131,7 @@ def _init_states(
         self.states.sort(key=lambda state: state.nu)
 
 
-def get_mqdt_states_from_fmodel(  # noqa: C901
+def get_mqdt_states_from_fmodel(  # noqa: C901, PLR0912
     model: FModel,
     nu_range: tuple[float, float],
     m_range: tuple[float, float] | None | NotSet,
@@ -171,6 +171,15 @@ def get_mqdt_states_from_fmodel(  # noqa: C901
         )
         return []
 
+    coefficients_fj: list[float] = []
+    angular_kets_fj: list[AngularKetFJ[Any]] = []
+    number_kets_fj: list[int] = [0] * len(model.outer_channels)
+    for i, angular_ket in enumerate(model.outer_channels):
+        for coeff_fj, ket_fj in angular_ket.to_state("FJ"):
+            coefficients_fj.append(coeff_fj)
+            angular_kets_fj.append(ket_fj)
+            number_kets_fj[i] += 1
+
     states: list[RydbergStateMQDT] = []
     for nu in nu_list:
         det_mmat = model.calc_det_m_matrix(nu)
@@ -204,16 +213,24 @@ def get_mqdt_states_from_fmodel(  # noqa: C901
                 radial = RadialDummy(1.0, nui)
             radial_kets.append(radial)
 
+        radial_kets_fj = [
+            radial for radial, number in zip(radial_kets, number_kets_fj, strict=True) for _ in range(number)
+        ]
+        coefficients_all = [
+            coeff for coeff, number in zip(coefficients, number_kets_fj, strict=True) for _ in range(number)
+        ]
+        coefficients_all = np.array(coefficients_all) * np.array(coefficients_fj)
+
         energy_au = model.calc_energy_au(nu)
         for m in get_m_range(model.f_tot, m_range):
             rydberg_kets = [
                 RydbergKet(model.species, angular_ket.replace_m(m), radial_ket)
-                for angular_ket, radial_ket in zip(model.outer_channels, radial_kets, strict=True)
+                for angular_ket, radial_ket in zip(angular_kets_fj, radial_kets_fj, strict=True)
             ]
             states.append(
                 RydbergStateMQDT(
                     model.species,
-                    coefficients,
+                    coefficients_all,
                     rydberg_kets,
                     nu=nu,
                     energy_au=energy_au,

src/rydstate/rydberg_state/rydberg_base.py

@@ -69,8 +69,9 @@ def __init__(
             raise ValueError("RydbergState must be initialized with at least one state.")
         if len(coefficients) != len(rydberg_kets):
             raise ValueError("Length of coefficients and rydberg_kets must be the same.")
-        if len(set(rydberg_kets)) != len(rydberg_kets):
-            raise ValueError("RydbergState initialized with duplicate rydberg_kets.")
+        angular_kets = [rydberg_ket.angular for rydberg_ket in rydberg_kets]
+        if len(set(angular_kets)) != len(angular_kets):
+            raise ValueError("RydbergState initialized with duplicate angular kets.")
 
         if abs(self.norm - 1) > 1e-10:
             raise ValueError(
@@ -137,14 +138,16 @@ def to_coupling_scheme(self, coupling_scheme: CouplingScheme) -> RydbergState:
         return RydbergState(self.species, coefficients, rydberg_kets, nu=self.nu, energy_au=self._energy_au)
 
     def _free_memory(self) -> None:
-        """Release the cached radial and angular data to reduce memory usage.
+        """Release the cached radial, angular and core state data to reduce memory usage.
 
-        This drops the references to the (potentially large) radial wavefunctions of the rydberg kets.
+        This drops the references to the (potentially large) radial wavefunctions of the rydberg kets
+        (including the radial wavefunctions of the cached core states).
         After calling this, matrix elements, overlaps and expectation values can no longer be calculated for this state.
         """
         for rydberg_ket in self.rydberg_kets:
             rydberg_ket.__dict__.pop("radial", None)
             rydberg_ket.__dict__.pop("angular", None)
+            rydberg_ket.__dict__.pop("core_state", None)
         self.__dict__.pop("rydberg_kets", None)
 
     @property

src/rydstate/rydberg_state/rydberg_ket.py

@@ -2,6 +2,7 @@
 
 import logging
 import math
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, Literal, overload
 
 from rydstate.angular.angular_ket import AngularKetLS
@@ -13,6 +14,7 @@
 if TYPE_CHECKING:
     from rydstate.angular.angular_ket import AngularKetBase
     from rydstate.radial.radial_base import Radial
+    from rydstate.rydberg_state.rydberg_sqdt import RydbergStateSQDT
     from rydstate.units import MatrixElementOperator, PintFloat
 
 
@@ -165,46 +167,47 @@ def _calc_core_radial_matrix_element_au(self, other: RydbergKet, k_radial: int)
 
         The core electron is treated as the low-lying valence electron of the corresponding singly charged SQDT ion
         (e.g. Yb174_ion for Yb174):
-        the Rydberg electron is ignored and the radial matrix element is calculated between the two ion states,
+        the Rydberg electron is ignored and the radial matrix element is calculated between the two core states,
         where the principal quantum number of each core electron is given by the lowest allowed shell of the ion
         for the given l_c.
         """
+        if self.core_state is None or other.core_state is None:
+            return 0.0
+
+        return self.core_state.radial.calc_matrix_element(other.core_state.radial, k_radial, unit="a.u.")
+
+    @cached_property
+    def core_state(self) -> RydbergStateSQDT[Any] | None:
+        """Get the corresponding ion state of the Rydberg ket."""
         from rydstate.rydberg_state.rydberg_sqdt import RydbergStateSQDT  # noqa: PLC0415
 
-        species = self.species
-        ion_species = f"{species}_ion"
+        ion_species = f"{self.species}_ion"
         try:
             ion_sqdt = get_sqdt(ion_species)
         except ValueError:
             logger.warning(
-                "No SQDT data available for the ion species of %s "
-                "returning 0 for the dipole matrix element core contribution.",
-                species,
+                "No SQDT data available for the ion species of %s, "
+                "thus we cannot calculate the ion state and its matrix elements.",
+                self.species,
             )
-            return 0.0
+            return None
 
-        kets = {"self": self, "other": other}
-        ion_states: dict[str, RydbergStateSQDT[Any]] = {}
-        for ket_name, ket in kets.items():
-            l_c = ket.angular.l_c
-            j_c = ket.angular.get_qn("j_c", allow_unknown=True)
-            f_c = ket.angular.get_qn("f_c", allow_unknown=True)
-            if is_unknown(l_c) or is_unknown(j_c) or is_unknown(f_c):
-                return 0.0
+        l_c = self.angular.l_c
+        j_c = self.angular.get_qn("j_c", allow_unknown=True)
+        f_c = self.angular.get_qn("f_c", allow_unknown=True)
+        if is_unknown(l_c) or is_unknown(j_c) or is_unknown(f_c):
+            return None
 
-            angular_ket = AngularKetLS(l_r=l_c, j_tot=j_c, f_tot=f_c, species=ion_species)
+        angular_ket = AngularKetLS(l_r=l_c, j_tot=j_c, f_tot=f_c, species=ion_species)
 
-            # TODO: we should probably also store n_c for the core angular ket in the future
-            # for now, it is correct to assume that the core electron is
-            # in the lowest allowed shell of the ion for the given l_c
-            for n_c in range(l_c + 1, l_c + 15):
-                if ion_sqdt.is_allowed_shell(n_c, l_c, 0.5):
-                    ion_states[ket_name] = RydbergStateSQDT(ion_species, n_c, angular_ket=angular_ket, sqdt=ion_sqdt)
-                    break
-            else:  # no break
-                raise ValueError(f"No allowed shell found for ion species {ion_species} with l_c={l_c}.")
+        # TODO: we should probably also store n_c for the core angular ket in the future
+        # for now, it is correct to assume that the core electron is
+        # in the lowest allowed shell of the ion for the given l_c
+        for n_c in range(l_c + 1, l_c + 15):
+            if ion_sqdt.is_allowed_shell(n_c, l_c, 0.5):
+                return RydbergStateSQDT(ion_species, n_c, angular_ket=angular_ket, sqdt=ion_sqdt)
 
-        return ion_states["self"].radial.calc_matrix_element(ion_states["other"].radial, k_radial, unit="a.u.")
+        raise ValueError(f"No allowed shell found for ion species {ion_species} with l_c={l_c}.")
 
     @overload
     def calc_matrix_element(

src/rydstate/species/fmodel.py

@@ -13,7 +13,7 @@
 if TYPE_CHECKING:
     from types import ModuleType
 
-    from rydstate.angular.angular_ket import AngularKetBase, AngularKetFJ
+    from rydstate.angular.angular_ket import AngularKetBase, AngularKetFJ, AngularKetJJ, AngularKetLS
     from rydstate.angular.utils import AllKnown
     from rydstate.species.mqdt import MQDT
     from rydstate.species.utils import RydbergRitzParameters
@@ -40,7 +40,7 @@ class FModel:
     inner_channels: ClassVar[list[AngularKetBase[Any]]]
     """List of inner channels in the MQDT model."""
 
-    outer_channels: ClassVar[list[AngularKetFJ[Any]]]
+    outer_channels: ClassVar[list[AngularKetFJ[Any] | AngularKetJJ[Any] | AngularKetLS[Any]]]
     """List of outer channels in the MQDT model."""
 
     eigen_quantum_defects: ClassVar[list[RydbergRitzParameters]]
@@ -52,12 +52,6 @@ class FModel:
     Each entry is a tuple (i_idx, j_idx, params) where i_idx and j_idx are the indices of the involved channels
     and params are the parameters for the energy dependence of the angle (constant or polynomial coefficients)."""
 
-    manual_frame_transformation_outer_inner: ClassVar[NDArray | None] = None
-    """Optional manually specified frame transformation matrix Q mapping inner to outer channels.
-    This is mainly needed for models with unknown quantum numbers,
-    where the frame transformation cannot (yet) be computed from Wigner coefficients.
-    """
-
     def __init__(self, mqdt: MQDT) -> None:
         self.mqdt = mqdt
         self.element_properties = get_element_properties(self.species)
@@ -176,9 +170,6 @@ def calc_frame_transformation_outer_inner(self) -> NDArray:
             Unitary transformation matrix Q (n_outer, n_inner).
 
         """
-        if self.manual_frame_transformation_outer_inner is not None:
-            return self.manual_frame_transformation_outer_inner
-
         n = len(self.inner_channels)
         u = np.zeros((n, n))
 

src/rydstate/species/mqdt.py

@@ -92,7 +92,11 @@ def reference_ionization_energy_au(self) -> float:
 
     def get_mqdt_models(self, outer_channel: AngularKetFJ[Any]) -> list[FModel]:
         """Return a list of MQDT models for the outer_channel."""
-        models = [model for model in self.models if any(ket == outer_channel for ket in model.outer_channels)]
+        models = [
+            model
+            for model in self.models
+            if any(abs(outer_channel.calc_reduced_overlap(ket)) > 0 for ket in model.outer_channels)
+        ]
         if len(models) == 0:
             models = [FModelSQDT(self.species, outer_channel, mqdt=self)]
         return models