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