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