Generate mqdt db

src/rydstate/basis/basis_base.py

@@ -28,12 +28,6 @@ def __init__(self, species: str) -> None:
     def __len__(self) -> int:
         return len(self.states)
 
-    def copy(self) -> Self:
-        new_basis = self.__class__.__new__(self.__class__)
-        new_basis.__dict__ = self.__dict__.copy()
-        new_basis.states = list(self.states)
-        return new_basis
-
     @overload
     def filter_states(
         self, qn: str, value: tuple[float, float], *, delta: float = 1e-10, keep_unknown: bool = False

src/rydstate/cli.py

@@ -12,7 +12,7 @@
 logger = logging.getLogger(__name__)
 
 
-def main() -> None:
+def main() -> None:  # noqa: C901, PLR0912, PLR0915
     """Entry point for the generate_database script."""
     parser = argparse.ArgumentParser(
         description="Generate a database, containing energies and matrix elements, for a given species.",
@@ -25,7 +25,6 @@ def main() -> None:
         default=None,
         type=int,
         help="The minimal principal quantum number n for the states to be included in the database. "
-        "This is used for species, where the low lying states do not converge nicely, so we exclude those states. "
         "Default 1 will start with the ground state configuration of the specific species (e.g. n=5 for Rb).",
     )
     parser.add_argument(
@@ -35,16 +34,29 @@ def main() -> None:
         help="The maximum principal quantum number n for the states to be included in the database.",
     )
     parser.add_argument(
-        "--max-delta-n",
+        "--nu-min",
+        default=None,
+        type=int,
+        help="The minimal effective principal quantum number nu for the states to be included in the database. "
+        "Default 0 will include all low lying states.",
+    )
+    parser.add_argument(
+        "--nu-max",
+        default=None,
+        type=int,
+        help="The maximum effective principal quantum number nu for the states to be included in the database.",
+    )
+    parser.add_argument(
+        "--max-delta-nu",
         default=float("inf"),
         type=float,
-        help="The maximum difference in principal quantum number n for matrix elements to be calculated.",
+        help="The maximum difference in effective principal quantum number nu for matrix elements to be calculated.",
     )
     parser.add_argument(
-        "--all-n-up-to",
+        "--all-nu-up-to",
         default=float("inf"),
         type=float,
-        help="Calculate all matrix elements where at least one state has principal quantum number n "
+        help="Calculate all matrix elements where at least one state has effective principal quantum number nu "
         "smaller than or equal to this value.",
     )
     parser.add_argument(
@@ -81,11 +93,14 @@ def main() -> None:
         if (
             args.n_min is not None
             or args.n_max is not None
-            or args.max_delta_n != float("inf")
-            or args.all_n_up_to != float("inf")
+            or args.nu_min is not None
+            or args.nu_max is not None
+            or args.max_delta_nu != float("inf")
+            or args.all_nu_up_to != float("inf")
         ):
             parser.error(
-                "--n-min, --n-max, --max-delta-n, and --all-n-up-to are only valid when generating a species database."
+                "--n-min, --n-max, --nu-min, --nu-max, --max-delta-nu, and --all-nu-up-to are only valid "
+                "when generating a species database."
             )
     elif args.f_max is not None:
         parser.error("--f-max is only valid when generating the misc database.")
@@ -108,10 +123,25 @@ def main() -> None:
             parser.error("--f-max is required when generating the misc database.")
         create_tables_for_misc(f_max=args.f_max, kappa_max=3)
     else:
-        n_min = args.n_min if args.n_min is not None else 1
-        if args.n_max is None:
-            parser.error("--n-max is required when generating a species database.")
-        create_tables_for_one_species(args.species, n_min, args.n_max, args.max_delta_n, args.all_n_up_to)
+        if args.n_max is None and args.nu_max is None:
+            parser.error("At least one of --n-max or --nu-max must be provided.")
+
+        if args.n_min is None and args.n_max is None:
+            n = None
+        else:
+            n_min = args.n_min if args.n_min is not None else 1
+            n_max = args.n_max if args.n_max is not None else int(args.nu_max) + 10
+            n = (n_min, n_max)
+
+        if args.nu_min is None and args.nu_max is None:
+            nu = None
+        else:
+            nu_min = args.nu_min if args.nu_min is not None else 0
+            nu_max = args.nu_max if args.nu_max is not None else args.n_max
+            nu = (nu_min, nu_max)
+        create_tables_for_one_species(
+            args.species, n=n, nu=nu, max_delta_nu=args.max_delta_nu, all_nu_up_to=args.all_nu_up_to
+        )
     logger.info("Time taken: %.2f seconds", time.perf_counter() - time_start)
 
 

src/rydstate/generate_database/generate_database.py

@@ -4,19 +4,16 @@
 import sqlite3
 from importlib.resources import files
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 import pandas as pd
 
 from rydstate import __version__
 from rydstate.angular.wigner_symbols import calc_wigner_3j
+from rydstate.basis.basis_mqdt import BasisMQDT
 from rydstate.basis.basis_sqdt import BasisSQDT
-from rydstate.generate_database.generate_matrix_elements_table import (
-    _calc_radial_matrix_element_cached,
-    calc_reduced_angular_matrix_element_cached,
-    generate_matrix_elements_tables,
-    get_radial_state_cached,
-)
+from rydstate.generate_database.generate_matrix_elements_table import generate_matrix_elements_tables
 from rydstate.generate_database.generate_misc_table import generate_wigner_table
 from rydstate.generate_database.generate_states_table import generate_states_table
 
@@ -27,25 +24,39 @@
 
 
 def create_tables_for_one_species(
-    species: str,
-    n_min: int,
-    n_max: int,
-    max_delta_n: float = np.inf,
-    all_n_up_to: float = np.inf,
+    species_specifier: str,
+    n: tuple[int, int] | None = None,
+    nu: tuple[float, float] | None = None,
+    max_delta_nu: float = np.inf,
+    all_nu_up_to: float = np.inf,
 ) -> None:
     """Create the database tables for a given species in the current directory."""
-    logger.info("Start creating database for %s", species)
-    logger.info("n-min=%d, n-max=%d", n_min, n_max)
-    logger.info("max_delta_n=%s, all_n_up_to=%s", max_delta_n, all_n_up_to)
+    logger.info("Start creating database for %s", species_specifier)
+    logger.info("n-range=%s", n)
+    logger.info("nu-range=%s", nu)
+    logger.info("max_delta_nu=%s, all_nu_up_to=%s", max_delta_nu, all_nu_up_to)
     logger.info("rydstate.__version__=%s", __version__)
 
     # create the database and populate the states and matrix elements tables
     db_file = Path("database.db")
     with sqlite3.connect(db_file) as conn:
         conn.executescript(DATABASE_SQL_FILE.read_text(encoding="utf-8"))
-        basis = BasisSQDT(species, n=(n_min, n_max), coupling_scheme="LS")
+        basis: BasisSQDT[Any] | BasisMQDT
+        species = species_specifier.removesuffix("_mqdt").removesuffix("_sqdt")
+        if species_specifier.endswith("_mqdt"):
+            if nu is None:
+                raise ValueError("nu must be provided for MQDT basis")
+            basis = BasisMQDT(species, nu=nu)
+            if n is not None:
+                basis.filter_states("n", n)
+        else:
+            if n is None:
+                raise ValueError("n must be provided for SQDT basis")
+            basis = BasisSQDT(species, n=n, coupling_scheme="LS")
+            if nu is not None:
+                basis.filter_states("nu", nu)
         generate_states_table(basis, conn)
-        generate_matrix_elements_tables(basis, conn, max_delta_n, all_n_up_to)
+        generate_matrix_elements_tables(basis, conn, max_delta_nu, all_nu_up_to, free_memory=True)
     logger.info("Size of %s: %.6f megabytes", db_file, db_file.stat().st_size * 1e-6)
 
     # convert the tables to parquet files
@@ -67,12 +78,6 @@ def create_tables_for_one_species(
                 with Path("log").open("a") as buf:
                     table.info(buf=buf)
 
-    logger.info(
-        "calc_reduced_angular_matrix_element_cached: %s", calc_reduced_angular_matrix_element_cached.cache_info()
-    )
-    logger.info("_calc_radial_matrix_element_cached: %s", _calc_radial_matrix_element_cached.cache_info())
-    logger.info("get_radial_state_cached: %s", get_radial_state_cached.cache_info())
-
 
 def create_tables_for_misc(f_max: float, kappa_max: int = 3) -> None:
     """Create misc databases, i.e. the wigner table in the current directory."""