Line data    Source code

       1             : # SPDX-FileCopyrightText: 2024 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3           1 : from __future__ import annotations
       4             : 
       5           1 : import logging
       6           1 : from typing import TYPE_CHECKING, ClassVar
       7             : 
       8           1 : from pairinteraction import _backend
       9           1 : from pairinteraction.custom_logging import _flush_pending_logs
      10             : 
      11             : if TYPE_CHECKING:
      12             :     import os
      13             :     from pathlib import Path
      14             : 
      15             : 
      16           1 : logger = logging.getLogger(__name__)
      17             : 
      18             : 
      19           1 : class Database:
      20             :     """Class for handling the databases for the PairInteraction package.
      21             : 
      22             :     The Databases are used to store the atomic states, their energies, and their matrix elements to other states.
      23             :     The databases are stored in the user's cache directory by default, but can be stored in a different directory.
      24             :     When running PairInteraction for the first time, the databases have to be downloaded from the internet
      25             :     (e.g. by explicitly passing `download_missing=True` to the constructor).
      26             :     Once the databases are downloaded, the user usually does not have to interact with the Database class directly.
      27             : 
      28             :     """
      29             : 
      30           1 :     _global_database: ClassVar[Database | None] = None
      31             : 
      32           1 :     def __init__(
      33             :         self,
      34             :         download_missing: bool = False,
      35             :         use_cache: bool = True,
      36             :         database_dir: str | os.PathLike[str] = "",
      37             :     ) -> None:
      38             :         """Create a new database instance with the given parameters.
      39             : 
      40             :         Args:
      41             :             download_missing: Whether to download missing databases if needed. Default False.
      42             :             use_cache: Whether to load the Wigner 3j symbols table into memory. Default True.
      43             :             database_dir: The directory where the databases are stored.
      44             :                 Default "", i.e. use the default directory (the user's cache directory).
      45             : 
      46             :         """
      47           1 :         self._cpp = _backend.Database(download_missing, use_cache, database_dir)
      48           1 :         _flush_pending_logs()  # call it manually because constructors of classes from nanobind cannot be decorated
      49             : 
      50           1 :     @classmethod
      51           1 :     def _from_cpp_object(cls, cpp_obj: _backend.Database) -> Database:
      52             :         """Create a Database instance from a C++ Database object.
      53             : 
      54             :         This is used internally to convert C++ objects returned by the C++ API to Python objects.
      55             :         """
      56           1 :         obj = cls.__new__(cls)
      57           1 :         obj._cpp = cpp_obj
      58           1 :         return obj
      59             : 
      60           1 :     @classmethod
      61           1 :     def get_global_database(cls) -> Database:
      62             :         """Return the global database instance if it was initialized, otherwise None."""
      63           1 :         return cls._global_database  # type: ignore [return-value]
      64             : 
      65           1 :     @classmethod
      66           1 :     def initialize_global_database(
      67             :         cls,
      68             :         download_missing: bool = False,
      69             :         use_cache: bool = True,
      70             :         database_dir: str | os.PathLike[str] = "",
      71             :     ) -> None:
      72             :         """Initialize the global database with the given parameters.
      73             : 
      74             :         The arguments are the same as for the constructor of this class.
      75             :         """
      76           1 :         db = cls(download_missing, use_cache, database_dir)
      77           1 :         if cls._global_database is None:
      78           1 :             cls._global_database = db
      79           0 :         elif (
      80             :             cls._global_database.download_missing == db.download_missing
      81             :             and cls._global_database.use_cache == db.use_cache
      82             :             and cls._global_database.database_dir == db.database_dir
      83             :         ):
      84           0 :             pass  # already initialized with the same parameters
      85             :         else:
      86           0 :             raise ValueError(
      87             :                 "Global database was already initialized with different parameters. "
      88             :                 "The global database is automatically initialized when needed. "
      89             :                 "If you explicitly want to initialize the global database, do this at the beginning of your script."
      90             :             )
      91             : 
      92           1 :     @property
      93           1 :     def download_missing(self) -> bool:
      94             :         """Whether to download missing databases if needed."""
      95           1 :         return self._cpp.get_download_missing()
      96             : 
      97           1 :     @property
      98           1 :     def use_cache(self) -> bool:
      99             :         """Whether to load the Wigner 3j symbols table into memory."""
     100           1 :         return self._cpp.get_use_cache()
     101             : 
     102           1 :     @property
     103           1 :     def database_dir(self) -> Path:
     104             :         """The directory where the databases are stored."""
     105           1 :         return self._cpp.get_database_dir()
     106             : 
     107           1 :     def get_versions_info(self) -> str:
     108             :         """Return a formatted table of local and remote database versions."""
     109           0 :         return self._cpp.get_versions_info()
     110             : 
     111             : 
     112           1 : def print_database_info(
     113             :     download_missing: bool = True,
     114             :     use_cache: bool = True,
     115             :     database_dir: str | os.PathLike[str] = "",
     116             : ) -> None:
     117             :     """Print an overview of local and remote database tables."""
     118           0 :     db = Database(download_missing, use_cache, database_dir)
     119           0 :     print(f"Database directory: {db.database_dir}")
     120           0 :     print(db.get_versions_info())