Line data    Source code

       1             : # SPDX-FileCopyrightText: 2024 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           1 : from __future__ import annotations
       5             : 
       6           1 : import contextlib
       7           1 : import logging
       8           1 : from pathlib import Path
       9           1 : from typing import TYPE_CHECKING, Protocol
      10             : 
      11           1 : import numpy as np
      12           1 : import pytest
      13             : 
      14             : if TYPE_CHECKING:
      15             :     from collections.abc import Callable, Iterator
      16             : 
      17             :     import pairinteraction as pi
      18             :     import pint
      19             :     from pairinteraction.units import NDArray
      20             : 
      21             : 
      22           1 : REFERENCE_PATHS = {
      23             :     "stark_map": Path(__file__).parent.parent / "data" / "reference_stark_map",
      24             :     "pair_potential": Path(__file__).parent.parent / "data" / "reference_pair_potential",
      25             : }
      26             : 
      27           1 : SHRUNK_DATABASE_PATH = (Path(__file__).parent.parent / "data" / "database").resolve()
      28             : 
      29             : 
      30           1 : def is_shrunk_database_used() -> bool:
      31             :     """Check whether the shrunk database that comes with the repository is used.
      32             : 
      33             :     The shrunk database only contains states within a narrow range of the effective principal quantum number, a
      34             :     few low-lying states, and only small values of the quantum number l_ryd. Thus, results that depend on the
      35             :     completeness of the database, like lifetimes, are far off and must not be checked against reference values.
      36             :     To check these values as well, run the tests with `pytest --database-dir "" --download-missing`.
      37             :     """
      38           1 :     from pairinteraction import Database
      39             : 
      40           1 :     return Path(Database.get_global_database().database_dir).resolve() == SHRUNK_DATABASE_PATH
      41             : 
      42             : 
      43           1 : def skip_value_check_if_shrunk_database() -> None:
      44             :     """Skip the rest of the current test if the results cannot be checked against reference values.
      45             : 
      46             :     Call this directly before the checks of the resulting values, see :func:`is_shrunk_database_used`. The code
      47             :     before the call is still executed, only the checks are skipped and the test is reported as skipped.
      48             :     """
      49           1 :     if is_shrunk_database_used():
      50           1 :         pytest.skip(
      51             :             "the shrunk database is used, run with `pytest --database-dir '' --download-missing` to include all checks"
      52             :         )
      53             : 
      54             : 
      55           1 : def compare_eigensystem_to_reference(
      56             :     reference_path: Path,
      57             :     eigenenergies: NDArray,
      58             :     overlaps: NDArray | None = None,
      59             :     eigenvectors: NDArray | None = None,
      60             :     kets: list[str] | None = None,
      61             : ) -> None:
      62           1 :     n_systems, n_kets = eigenenergies.shape
      63           1 :     np.testing.assert_allclose(eigenenergies, np.loadtxt(reference_path / "eigenenergies.txt"))
      64             : 
      65           1 :     if overlaps is not None:
      66             :         # Ensure that the overlaps sum up to one
      67           1 :         np.testing.assert_allclose(np.sum(overlaps, axis=1), np.ones(n_systems))
      68           1 :         np.testing.assert_allclose(overlaps, np.loadtxt(reference_path / "overlaps.txt"), atol=1e-7)
      69             : 
      70           1 :     if kets is not None:
      71           1 :         np.testing.assert_equal(kets, np.loadtxt(reference_path / "kets.txt", dtype=str, delimiter="\t"))
      72             : 
      73           1 :     if eigenvectors is not None:
      74             :         # Because of degeneracies, checking the eigenvectors against reference data is complicated.
      75             :         # Thus, we only check their normalization and orthogonality.
      76           1 :         cumulative_norm = (np.array(eigenvectors) * np.array(eigenvectors).conj()).sum(axis=1)
      77           1 :         np.testing.assert_allclose(cumulative_norm, n_kets * np.ones(n_systems))
      78             : 
      79             : 
      80           1 : @contextlib.contextmanager
      81           1 : def no_log_propagation(logger: logging.Logger | str) -> Iterator[None]:
      82             :     """Context manager to temporarily disable log propagation for a given logger."""
      83           1 :     if isinstance(logger, str):
      84           1 :         logger = logging.getLogger(logger)
      85           1 :     old_value = logger.propagate
      86           1 :     try:
      87           1 :         logger.propagate = False
      88           1 :         yield
      89             :     finally:
      90           1 :         logger.propagate = old_value
      91             : 
      92             : 
      93           1 : class PairinteractionModule(Protocol):
      94           1 :     ureg: pint.UnitRegistry
      95           1 :     Database: type[pi.Database]
      96           1 :     KetAtom: type[pi.KetAtom]
      97           1 :     StateAtom: type[pi.StateAtom]
      98           1 :     BasisAtom: type[pi.BasisAtom]
      99           1 :     SystemAtom: type[pi.SystemAtom]
     100           1 :     KetPair: type[pi.KetPair]
     101           1 :     BasisPair: type[pi.BasisPair]
     102           1 :     SystemPair: type[pi.SystemPair]
     103           1 :     EffectiveSystemPair: type[pi.EffectiveSystemPair]
     104           1 :     C3: type[pi.C3]
     105           1 :     C6: type[pi.C6]
     106           1 :     diagonalize: Callable[..., None]