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