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]