improve __init__

Author: johannes-moegerle

src/rydstate/__init__.py

@@ -1,5 +1,10 @@
-from rydstate import angular, radial, species
-from rydstate.basis import BasisSQDTAlkali, BasisSQDTAlkalineFJ, BasisSQDTAlkalineJJ, BasisSQDTAlkalineLS
+from rydstate import angular, basis, radial, rydberg, species
+from rydstate.basis import (
+    BasisSQDTAlkali,
+    BasisSQDTAlkalineFJ,
+    BasisSQDTAlkalineJJ,
+    BasisSQDTAlkalineLS,
+)
 from rydstate.rydberg import (
     RydbergStateSQDT,
     RydbergStateSQDTAlkali,
@@ -20,7 +25,9 @@
     "RydbergStateSQDTAlkalineJJ",
     "RydbergStateSQDTAlkalineLS",
     "angular",
+    "basis",
     "radial",
+    "rydberg",
     "species",
     "ureg",
 ]

src/rydstate/angular/__init__.py

@@ -1,3 +1,4 @@
+from rydstate.angular import utils
 from rydstate.angular.angular_ket import AngularKetFJ, AngularKetJJ, AngularKetLS
 from rydstate.angular.angular_state import AngularState
 
@@ -6,4 +7,5 @@
     "AngularKetJJ",
     "AngularKetLS",
     "AngularState",
+    "utils",
 ]

src/rydstate/basis/__init__.py

@@ -1,3 +1,4 @@
+from rydstate.basis.basis_base import BasisBase
 from rydstate.basis.basis_sqdt import BasisSQDTAlkali, BasisSQDTAlkalineFJ, BasisSQDTAlkalineJJ, BasisSQDTAlkalineLS
 
-__all__ = ["BasisSQDTAlkali", "BasisSQDTAlkalineFJ", "BasisSQDTAlkalineJJ", "BasisSQDTAlkalineLS"]
+__all__ = ["BasisBase", "BasisSQDTAlkali", "BasisSQDTAlkalineFJ", "BasisSQDTAlkalineJJ", "BasisSQDTAlkalineLS"]

src/rydstate/rydberg/__init__.py

@@ -1,3 +1,4 @@
+from rydstate.rydberg.rydberg_base import RydbergStateBase
 from rydstate.rydberg.rydberg_sqdt import (
     RydbergStateSQDT,
     RydbergStateSQDTAlkali,
@@ -7,6 +8,7 @@
 )
 
 __all__ = [
+    "RydbergStateBase",
     "RydbergStateSQDT",
     "RydbergStateSQDTAlkali",
     "RydbergStateSQDTAlkalineFJ",