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 : from typing import TYPE_CHECKING
       7              : 
       8            1 : import numpy as np
       9            1 : import pytest
      10              : 
      11            1 : from .utils import REFERENCE_PATHS, compare_eigensystem_to_reference
      12              : 
      13              : if TYPE_CHECKING:
      14              :     from .utils import PairinteractionModule
      15              : 
      16              : 
      17            1 : def test_pair_potential(pi_module: PairinteractionModule, generate_reference: bool) -> None:
      18              :     """Test calculating a pair potential."""
      19              :     # Create a single-atom system
      20            1 :     basis = pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2))
      21            1 :     print(f"Number of single-atom basis states: {basis.number_of_states}")
      22              : 
      23            1 :     system = pi_module.SystemAtom(basis)
      24              : 
      25              :     # Create two-atom systems for different interatomic distances
      26            1 :     ket = pi_module.KetAtom("Rb", n=60, l=0, m=0.5)
      27            1 :     delta_energy = 3  # GHz
      28            1 :     min_energy = 2 * ket.get_energy(unit="GHz") - delta_energy
      29            1 :     max_energy = 2 * ket.get_energy(unit="GHz") + delta_energy
      30              : 
      31            1 :     basis_pair = pi_module.BasisPair([system, system], energy=(min_energy, max_energy), energy_unit="GHz", m=(1, 1))
      32            1 :     print(f"Number of two-atom basis states: {basis_pair.number_of_states}")
      33              : 
      34            1 :     distances = np.linspace(1, 5, 5)
      35            1 :     system_pairs = [pi_module.SystemPair(basis_pair).set_distance(d, unit="micrometer") for d in distances]
      36              : 
      37              :     # Diagonalize the systems in parallel
      38            1 :     pi_module.diagonalize(system_pairs, diagonalizer="eigen")
      39              : 
      40              :     # Get the overlap with |ket, ket>
      41            1 :     overlaps = np.array([system.get_eigenbasis().get_overlaps([ket, ket]) for system in system_pairs])
      42              : 
      43              :     # Ensure that the overlaps sum up to one
      44            1 :     np.testing.assert_allclose(np.sum(overlaps, axis=1), np.ones(5))
      45              : 
      46              :     # Compare to reference data
      47            1 :     kets = [ket.get_label("raw") for ket in basis_pair.kets]
      48            1 :     eigenenergies = np.array([system.get_eigenenergies(unit="GHz") for system in system_pairs])
      49            1 :     eigenvectors = np.array([system.get_eigenbasis().get_coefficients().todense().A1 for system in system_pairs])
      50              : 
      51            1 :     reference_path = REFERENCE_PATHS["pair_potential"]
      52            1 :     if generate_reference:
      53            1 :         reference_path.mkdir(parents=True, exist_ok=True)
      54            1 :         np.savetxt(reference_path / "kets.txt", kets, fmt="%s", delimiter="\t")
      55            1 :         np.savetxt(reference_path / "eigenenergies.txt", eigenenergies)
      56            1 :         np.savetxt(reference_path / "overlaps.txt", overlaps)
      57            1 :         pytest.skip("Reference data generated, skipping comparison test")
      58              : 
      59            1 :     compare_eigensystem_to_reference(reference_path, eigenenergies, overlaps, eigenvectors, kets)