update cli generate database

Author: johannes-moegerle

docs/api/generate_database.rst

@@ -6,5 +6,6 @@ Generate Database module
 .. autosummary::
     :toctree: ../_autosummary/
 
-    generate_database.create_tables_for_one_species
+    generate_database.create_tables_for_sqdt
+    generate_database.create_tables_for_mqdt
     generate_database.create_tables_for_misc

src/rydstate/cli.py

@@ -8,143 +8,231 @@
 import time
 from pathlib import Path
 
-from rydstate.generate_database.generate_database import create_tables_for_misc, create_tables_for_one_species
+from rydstate.generate_database.generate_database import (
+    create_tables_for_misc,
+    create_tables_for_mqdt,
+    create_tables_for_sqdt,
+)
 
 logger = logging.getLogger(__name__)
 
 
-def main() -> None:  # noqa: C901, PLR0912, PLR0915
+def main() -> None:
     """Entry point for the generate_database script."""
+    args = build_parser().parse_args()
+
+    directory = prepare_directory(args)
+    configure_logging(args.log_level, directory, args.warnings_as_exceptions)
+
+    time_start = time.perf_counter()
+    if args.mode == "misc":
+        create_tables_for_misc(f_max=args.f_max, kappa_max=3)
+    elif args.mode == "sqdt":
+        create_tables_for_sqdt(
+            args.species,
+            n=(args.n_min, args.n_max),
+            f_tot=get_f_tot_range(args),
+            max_delta_nu=args.max_delta_nu,
+            all_nu_up_to=args.all_nu_up_to,
+        )
+    elif args.mode == "mqdt":
+        create_tables_for_mqdt(
+            args.species,
+            nu=(args.nu_min, args.nu_max),
+            f_tot=get_f_tot_range(args),
+            max_delta_nu=args.max_delta_nu,
+            all_nu_up_to=args.all_nu_up_to,
+        )
+    else:
+        raise ValueError(f"Unknown mode: {args.mode}")
+
+    logger.info("Time taken: %.2f seconds", time.perf_counter() - time_start)
+    log_memory_usage()
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the argument parser with the sqdt, mqdt and misc subcommands."""
     parser = argparse.ArgumentParser(
         description="Generate a database, containing energies and matrix elements, for a given species.",
         formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog=("Example:\n  generate_database Rb --log-level INFO\n"),
+        epilog=(
+            "Examples:\n"
+            "  generate_database sqdt Rb --n-max 60\n"
+            "  generate_database mqdt Yb174 --nu-max 60\n"
+            "  generate_database misc --f-max 10\n"
+        ),
     )
-    parser.add_argument("species", help="The species name to generate the database for.")
-    parser.add_argument(
-        "--n-min",
+    subparsers = parser.add_subparsers(dest="mode", required=True, title="modes")
+
+    # arguments shared by all modes
+    common = argparse.ArgumentParser(add_help=False)
+    common.add_argument(
+        "--directory",
         default=None,
-        type=int,
-        help="The minimal principal quantum number n for the states to be included in the database. "
-        "Default 1 will start with the ground state configuration of the specific species (e.g. n=5 for Rb).",
+        type=str,
+        help="The directory where the database will be saved. Default database/<species> for the sqdt mode, "
+        "database/<species>_mqdt for the mqdt mode and database/misc for the misc mode.",
     )
-    parser.add_argument(
-        "--n-max",
-        default=None,
-        type=int,
-        help="The maximum principal quantum number n for the states to be included in the database.",
+    common.add_argument(
+        "--log-level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        help="set the logging level (default: INFO)",
     )
-    parser.add_argument(
-        "--nu-min",
+    common.add_argument(
+        "--warnings-as-exceptions",
+        action="store_true",
+        help="Treat warnings in rydstate as exceptions.",
+    )
+    common.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Delete the database folder if it exists and create a new one.",
+    )
+
+    # arguments shared by the species modes (sqdt and mqdt)
+    species_common = argparse.ArgumentParser(add_help=False)
+    species_common.add_argument("species", help="The species name to generate the database for (e.g. Rb).")
+    species_common.add_argument(
+        "--f-tot-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.",
+        type=float,
+        help="The minimal total angular momentum quantum number f_tot for the states to be included in the database. "
+        "Default 0 will include all states with small f_tot.",
     )
-    parser.add_argument(
-        "--nu-max",
+    species_common.add_argument(
+        "--f-tot-max",
         default=None,
-        type=int,
-        help="The maximum effective principal quantum number nu for the states to be included in the database.",
+        type=float,
+        help="The maximum total angular momentum quantum number f_tot for the states to be included in the database. "
+        "Default inf will include all states with large f_tot.",
     )
-    parser.add_argument(
+    species_common.add_argument(
         "--max-delta-nu",
         default=float("inf"),
         type=float,
         help="The maximum difference in effective principal quantum number nu for matrix elements to be calculated.",
     )
-    parser.add_argument(
+    species_common.add_argument(
         "--all-nu-up-to",
         default=float("inf"),
         type=float,
         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(
-        "--f-max",
-        default=None,
+
+    sqdt_parser = subparsers.add_parser(
+        "sqdt",
+        parents=[species_common, common],
+        help="Generate the database for a species using single-channel quantum defect theory.",
+        description="Generate the database for a species using single-channel quantum defect theory. "
+        "The basis is defined via the n-range.",
+    )
+    sqdt_parser.add_argument(
+        "--n-min",
+        default=1,
+        type=int,
+        help="The minimal principal quantum number n for the states to be included in the database. "
+        "Default 1 will start with the ground state configuration of the specific species (e.g. n=5 for Rb).",
+    )
+    sqdt_parser.add_argument(
+        "--n-max",
+        required=True,
+        type=int,
+        help="The maximum principal quantum number n for the states to be included in the database.",
+    )
+
+    mqdt_parser = subparsers.add_parser(
+        "mqdt",
+        parents=[species_common, common],
+        help="Generate the database for a species using multi-channel quantum defect theory.",
+        description="Generate the database for a species using multi-channel quantum defect theory. "
+        "The basis is defined via the nu-range.",
+    )
+    mqdt_parser.add_argument(
+        "--nu-min",
+        default=0,
         type=float,
-        help="The maximum angular momentum quantum number f for misc database tables.",
+        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(
-        "--directory",
-        default=None,
-        type=str,
-        help="The directory where the database will be saved.",
+    mqdt_parser.add_argument(
+        "--nu-max",
+        required=True,
+        type=float,
+        help="The maximum effective principal quantum number nu for the states to be included in the database.",
     )
-    parser.add_argument(
-        "--log-level",
-        default="INFO",
-        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
-        help="set the logging level (default: INFO)",
+
+    misc_parser = subparsers.add_parser(
+        "misc",
+        parents=[common],
+        help="Generate the misc database tables, which do not depend on a species.",
+        description="Generate the misc database tables, which do not depend on a species.",
     )
-    parser.add_argument(
-        "--warnings-as-exceptions",
-        action="store_true",
-        help="Treat warnings in rydstate as exceptions.",
+    misc_parser.add_argument(
+        "--f-max",
+        required=True,
+        type=float,
+        help="The maximum angular momentum quantum number f for misc database tables.",
     )
-    parser.add_argument(
-        "--overwrite",
-        action="store_true",
-        help="Delete the species folder if it exists and create a new one.",
-    )
-
-    args = parser.parse_args()
-    if args.species == "misc":
-        if (
-            args.n_min is not None
-            or args.n_max is not None
-            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, --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.")
-
-    directory = Path(args.directory) if args.directory is not None else Path("database") / args.species
+
+    return parser
+
+
+def prepare_directory(args: argparse.Namespace) -> Path:
+    """Create the (empty) database directory and change into it."""
+    if args.directory is not None:
+        directory = Path(args.directory)
+    elif args.mode == "misc":
+        directory = Path("database") / "misc"
+    elif args.mode == "sqdt":
+        directory = Path("database") / args.species
+    elif args.mode == "mqdt":
+        directory = Path("database") / f"{args.species}_mqdt"
+    else:
+        raise ValueError(f"Unknown mode: {args.mode}")
     directory = directory.resolve()
+
     if directory.exists():
-        if args.overwrite:
-            shutil.rmtree(directory)
-        else:
+        if not args.overwrite:
             raise FileExistsError(f"The folder {directory} already exists. Use --overwrite to overwrite it.")
+        check_is_generated_database(directory)
+        shutil.rmtree(directory)
     directory.mkdir(parents=True)
     os.chdir(directory)
+    return directory
 
-    configure_logging(args.log_level, directory, args.warnings_as_exceptions)
 
-    time_start = time.perf_counter()
-    if args.species == "misc":
-        if args.f_max is 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:
-        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
+def check_is_generated_database(directory: Path) -> None:
+    """Raise if the directory contains anything but the files written by a previous run.
+
+    Since --overwrite deletes the whole directory tree, only ever delete a directory that looks like a
+    previously generated database, i.e. that solely contains the log file and the parquet tables.
+    """
+    unexpected = sorted(
+        entry.name
+        for entry in directory.iterdir()
+        if not (entry.is_file() and (entry.name == "log" or entry.suffix == ".parquet"))
+    )
+    if unexpected:
+        listed = ", ".join(unexpected[:5]) + (", ..." if len(unexpected) > 5 else "")
+        raise FileExistsError(
+            f"Refusing to overwrite the folder {directory}, since it does not look like a generated database "
+            f"(it contains {listed}). Delete it manually if this is really what you want."
         )
-    logger.info("Time taken: %.2f seconds", time.perf_counter() - time_start)
-    log_memory_usage()
+
+
+def get_f_tot_range(args: argparse.Namespace) -> tuple[float, float] | None:
+    """Get the (f_tot_min, f_tot_max) range from the parsed arguments.
+
+    Returns None if neither the minimum nor the maximum is given, i.e. all f_tot values are included.
+    """
+    if args.f_tot_min is None and args.f_tot_max is None:
+        return None
+    return (
+        args.f_tot_min if args.f_tot_min is not None else 0,
+        args.f_tot_max if args.f_tot_max is not None else float("inf"),
+    )
 
 
 def log_memory_usage() -> None:

src/rydstate/generate_database/__init__.py

@@ -1,3 +1,7 @@
-from rydstate.generate_database.generate_database import create_tables_for_misc, create_tables_for_one_species
+from rydstate.generate_database.generate_database import (
+    create_tables_for_misc,
+    create_tables_for_mqdt,
+    create_tables_for_sqdt,
+)
 
-__all__ = ["create_tables_for_misc", "create_tables_for_one_species"]
+__all__ = ["create_tables_for_misc", "create_tables_for_mqdt", "create_tables_for_sqdt"]