Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4             : """Test atom-ion interaction."""
       5             : 
       6           1 : import numpy as np
       7           1 : import pytest
       8             : 
       9           1 : import pairinteraction.complex as pi
      10             : 
      11             : 
      12           1 : def test_ion_z() -> None:
      13             :     """Test the calculation of energy shifts in the field on an ion positioned along z."""
      14             :     # Create a basis
      15           1 :     ket = pi.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
      16           1 :     basis = pi.BasisAtom("Rb", n=(ket.n - 2, ket.n + 2), l=(0, ket.l + 2), m=(ket.m, ket.m))
      17             : 
      18             :     # Create systems for different distances to the ion
      19           1 :     distance = 3
      20           1 :     system_z = (
      21             :         pi.SystemAtom(basis)
      22             :         .set_ion_interaction_order(3)
      23             :         .set_ion_charge(1, unit="e")
      24             :         .set_ion_distance_vector([0, 0, distance], unit="um")
      25             :     )
      26             : 
      27             :     # Diagonalize the system
      28           1 :     system_z = system_z.diagonalize(diagonalizer="eigen", sort_by_energy=True)
      29             : 
      30             :     # Ensure that values are correct for the system where the ion is closest to the atom
      31           1 :     eigenenergies = system_z.get_eigenenergies(unit="GHz")
      32           1 :     overlaps = system_z.basis.get_overlaps(ket)
      33           1 :     idx = np.argmax(overlaps)
      34           1 :     assert pytest.approx(overlaps[idx], rel=1e-6) == 0.8754278674619744  # NOSONAR
      35           1 :     assert pytest.approx(eigenenergies[idx] - ket.get_energy(unit="GHz"), rel=1e-12) == -0.2531503909267485  # NOSONAR
      36             : 
      37             : 
      38           1 : def test_ion_x() -> None:
      39             :     """Test the calculation of energy shifts in the field on an ion positioned along x."""
      40             :     # Create a basis
      41           1 :     ket = pi.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
      42           1 :     basis = pi.BasisAtom("Rb", n=(ket.n - 2, ket.n + 2), l=(0, ket.l + 2))
      43             : 
      44             :     # Create systems for different distances to the ion
      45           1 :     distance = 3
      46           1 :     system_x = (
      47             :         pi.SystemAtom(basis)
      48             :         .set_ion_interaction_order(3)
      49             :         .set_ion_charge(1, unit="e")
      50             :         .set_ion_distance_vector([distance, 0, 0], unit="um")
      51             :     )
      52             : 
      53             :     # Diagonalize the system
      54           1 :     system_x = system_x.diagonalize(diagonalizer="eigen", sort_by_energy=True)
      55             : 
      56             :     # Ensure that values are correct for the system where the ion is closest to the atom
      57           1 :     eigenenergies = system_x.get_eigenenergies(unit="GHz")
      58           1 :     overlaps = system_x.basis.get_overlaps(ket)
      59           1 :     idx = np.argmax(overlaps)
      60             :     # Note that we must use a large relative tolerance for the overlaps because the calculated value
      61             :     # is very sensitive on the actual method that is used by eigen to diagonalize the system (whether eigen
      62             :     # is using its own implementation, mkl on a Intel CPU, mkl on a AMD CPU, or lapack). This is because of
      63             :     # eigenstates belonging to different degenerate Zeeman sublevels.
      64           1 :     assert pytest.approx(overlaps[idx], rel=0.2) == 0.8754157131466024  # NOSONAR
      65           1 :     assert pytest.approx(eigenenergies[idx] - ket.get_energy(unit="GHz"), rel=1e-12) == -0.2531503909267485  # NOSONAR
      66             : 
      67             : 
      68           1 : def test_ion_angle_dependence() -> None:
      69             :     """Test the calculation of energy shifts in the field on an ion for different angles."""
      70             :     # Create a basis
      71           1 :     basis = pi.BasisAtom("Rb", n=(58, 62), l=(0, 2))
      72             : 
      73             :     # Create systems for different distances to the ion
      74           1 :     distances = np.linspace(3, 10, 5)
      75           1 :     systems_x = [
      76             :         pi.SystemAtom(basis)
      77             :         .set_ion_interaction_order(3)
      78             :         .set_ion_charge(1, unit="e")
      79             :         .set_ion_distance_vector([d, 0, 0], unit="um")
      80             :         for d in distances
      81             :     ]
      82           1 :     systems_y = [
      83             :         pi.SystemAtom(basis)
      84             :         .set_ion_interaction_order(3)
      85             :         .set_ion_charge(1, unit="e")
      86             :         .set_ion_distance_vector([0, d, 0], unit="um")
      87             :         for d in distances
      88             :     ]
      89           1 :     systems_z = [
      90             :         pi.SystemAtom(basis)
      91             :         .set_ion_interaction_order(3)
      92             :         .set_ion_charge(1, unit="e")
      93             :         .set_ion_distance_vector([0, 0, d], unit="um")
      94             :         for d in distances
      95             :     ]
      96             : 
      97             :     # Diagonalize the systems in parallel
      98           1 :     pi.diagonalize(systems_x + systems_y + systems_z, diagonalizer="eigen", sort_by_energy=True)
      99             : 
     100             :     # Ensure that all eigenenergies are the same
     101           1 :     for system_x, system_y, system_z in zip(systems_x, systems_y, systems_z):
     102           1 :         np.testing.assert_allclose(system_x.get_eigenenergies(unit="GHz"), system_y.get_eigenenergies(unit="GHz"))
     103           1 :         np.testing.assert_allclose(system_x.get_eigenenergies(unit="GHz"), system_z.get_eigenenergies(unit="GHz"))