Line data    Source code

       1             : # SPDX-FileCopyrightText: 2024 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4             : """Test the pair potential calculation."""
       5             : 
       6           1 : from pathlib import Path
       7           1 : from typing import TYPE_CHECKING, Optional
       8             : 
       9           1 : import numpy as np
      10           1 : import pairinteraction.real as pi
      11           1 : import pytest
      12             : 
      13             : if TYPE_CHECKING:
      14             :     from pairinteraction.units import NDArray
      15             : 
      16           1 : reference_kets_file = Path(__file__).parent.parent / "data/reference_pair_potential/kets.txt"
      17           1 : reference_eigenenergies_file = Path(__file__).parent.parent / "data/reference_pair_potential/eigenenergies.txt"
      18           1 : reference_overlaps_file = Path(__file__).parent.parent / "data/reference_pair_potential/overlaps.txt"
      19             : 
      20             : 
      21           1 : def test_pair_potential(generate_reference: bool) -> None:
      22             :     """Test calculating a pair potential."""
      23             :     # Create a single-atom system
      24           1 :     basis = pi.BasisAtom("Rb", n=(58, 62), l=(0, 2))
      25           1 :     print(f"Number of single-atom basis states: {basis.number_of_states}")
      26             : 
      27           1 :     system = pi.SystemAtom(basis)
      28             : 
      29             :     # Create two-atom systems for different interatomic distances
      30           1 :     ket = pi.KetAtom("Rb", n=60, l=0, m=0.5)
      31           1 :     delta_energy = 3  # GHz
      32           1 :     min_energy = 2 * ket.get_energy(unit="GHz") - delta_energy
      33           1 :     max_energy = 2 * ket.get_energy(unit="GHz") + delta_energy
      34             : 
      35           1 :     basis_pair = pi.BasisPair([system, system], energy=(min_energy, max_energy), energy_unit="GHz", m=(1, 1))
      36           1 :     print(f"Number of two-atom basis states: {basis_pair.number_of_states}")
      37             : 
      38           1 :     distances = np.linspace(1, 5, 5)
      39           1 :     system_pairs = [pi.SystemPair(basis_pair).set_distance(d, unit="micrometer") for d in distances]
      40             : 
      41             :     # Diagonalize the systems in parallel
      42           1 :     pi.diagonalize(system_pairs, diagonalizer="eigen", sort_by_energy=True)
      43             : 
      44             :     # Get the overlap with |ket, ket>
      45           1 :     overlaps = np.array([system.get_eigenbasis().get_overlaps([ket, ket]) for system in system_pairs])
      46             : 
      47             :     # Ensure that the overlaps sum up to one
      48           1 :     np.testing.assert_allclose(np.sum(overlaps, axis=1), np.ones(5))
      49             : 
      50             :     # Compare to reference data
      51           1 :     kets = [repr(ket) for ket in basis_pair.kets]
      52           1 :     eigenenergies = np.array([system.get_eigenenergies(unit="GHz") for system in system_pairs])
      53           1 :     eigenvectors = np.array([system.get_eigenbasis().get_coefficients().todense().A1 for system in system_pairs])
      54             : 
      55           1 :     if generate_reference:
      56           1 :         reference_kets_file.parent.mkdir(parents=True, exist_ok=True)
      57           1 :         np.savetxt(reference_kets_file, kets, fmt="%s", delimiter="\t")
      58           1 :         np.savetxt(reference_eigenenergies_file, eigenenergies)
      59           1 :         np.savetxt(reference_overlaps_file, overlaps)
      60           1 :         pytest.skip("Reference data generated, skipping comparison test")
      61             : 
      62           1 :     compare_pair_potential_to_reference(eigenenergies, overlaps, eigenvectors, kets)
      63             : 
      64             : 
      65           1 : def compare_pair_potential_to_reference(
      66             :     eigenenergies: "NDArray",
      67             :     overlaps: Optional["NDArray"] = None,
      68             :     eigenvectors: Optional["NDArray"] = None,
      69             :     kets: Optional[list[str]] = None,
      70             : ) -> None:
      71           1 :     np.testing.assert_allclose(eigenenergies, np.loadtxt(reference_eigenenergies_file))
      72             : 
      73           1 :     if overlaps is not None:
      74           1 :         np.testing.assert_allclose(overlaps, np.loadtxt(reference_overlaps_file), atol=1e-10)
      75             : 
      76           1 :     if kets is not None:
      77           1 :         np.testing.assert_equal(kets, np.loadtxt(reference_kets_file, dtype=str, delimiter="\t"))
      78             : 
      79           1 :     if eigenvectors is not None:
      80             :         # Because of degeneracies, checking the eigenvectors against reference data is complicated.
      81             :         # Thus, we only check their normalization and orthogonality.
      82           1 :         cumulative_norm = (np.array(eigenvectors) * np.array(eigenvectors).conj()).sum(axis=1)
      83           1 :         np.testing.assert_allclose(cumulative_norm, 19 * np.ones(5))