Line data Source code
1 : # SPDX-FileCopyrightText: 2024 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 1 : from pairinteraction.units import ureg
11 :
12 : if TYPE_CHECKING:
13 : from pairinteraction.green_tensor.green_tensor_interpolator import GreenTensorInterpolator
14 :
15 : from .utils import PairinteractionModule
16 :
17 :
18 1 : @pytest.mark.parametrize("distance_mum", [1, 2, 11])
19 1 : def test_static_green_tensor_interpolator(
20 : pi_module: PairinteractionModule,
21 : green_tensor_interpolator_class: type[GreenTensorInterpolator],
22 : distance_mum: float,
23 : ) -> None:
24 : """Test calculating a pair potential using a user-defined static green tensor interpolator."""
25 : # Create a single-atom system
26 1 : basis = pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2))
27 1 : system = pi_module.SystemAtom(basis)
28 :
29 : # Create two-atom basis
30 1 : ket = pi_module.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 1 : basis_pair = pi_module.BasisPair([system, system], energy=(min_energy, max_energy), energy_unit="GHz", m=(1, 1))
35 :
36 : # Create a system using a user-defined green tensor interpolator for dipole-dipole interaction
37 1 : gti = green_tensor_interpolator_class()
38 1 : omega = ureg.Quantity(1, "Hz")
39 1 : distance_au = ureg.Quantity(distance_mum, "micrometer").to("bohr").m
40 :
41 1 : tensor = np.array([[1, 0, 0], [0, 1, 0], [0, 0, -2]]) / (distance_au**3)
42 1 : tensor_unit = "hartree / (e * bohr)^2"
43 1 : gti.set_constant(1, 1, tensor, tensor_unit)
44 :
45 1 : tensor_spherical = gti.get(1, 1, omega, unit=tensor_unit, scaled=True, coordinates="spherical")
46 1 : tensor_spherical_ref = np.array([[1, 0, 0], [0, -2, 0], [0, 0, 1]]) / distance_au**3
47 1 : np.testing.assert_allclose(tensor_spherical, tensor_spherical_ref)
48 :
49 1 : tensor_spherical = gti.get(1, 1, ureg.Quantity(2.5, "GHz"), unit=tensor_unit, scaled=True, coordinates="spherical")
50 1 : np.testing.assert_allclose(tensor_spherical, tensor_spherical_ref)
51 :
52 1 : system_pairs = pi_module.SystemPair(basis_pair)
53 1 : system_pairs._cpp.set_green_tensor_interpolator(gti._cpp)
54 :
55 : # Create a reference system using the build in dipole-dipole interaction
56 1 : system_pairs_reference = (
57 : pi_module.SystemPair(basis_pair).set_interaction_order(3).set_distance(distance_mum, unit="micrometer")
58 : )
59 :
60 : # Check that the two systems are equivalent
61 1 : np.testing.assert_allclose(
62 : system_pairs.get_hamiltonian(unit="GHz").data, system_pairs_reference.get_hamiltonian(unit="GHz").data
63 : )
64 :
65 :
66 1 : @pytest.mark.parametrize("distance_mum", [1, 2, 11])
67 1 : def test_omega_dependent_green_tensor_interpolator(
68 : green_tensor_interpolator_class: type[GreenTensorInterpolator],
69 : distance_mum: float,
70 : ) -> None:
71 : """Test the interpolation for different values of the transition energy."""
72 : # Define an simple linear omega-dependent green tensor interpolator
73 : # note that at least four entries are needed for the applied spline interpolation.
74 1 : gti = green_tensor_interpolator_class()
75 1 : transition_energies = [ureg.Quantity(i, "planck_constant * GHz") for i in np.linspace(1, 5, 20)]
76 1 : tensors = [
77 : np.array([[1, 0, 0], [0, 1, 0], [0, 0, -2]]) * i / (freq.m**2 * distance_mum**3)
78 : for (i, freq) in enumerate(transition_energies, start=1)
79 : ]
80 1 : tensors_pint = [ureg.Quantity(tensor, "1/micrometer") for tensor in tensors]
81 1 : gti.set_list(1, 1, tensors_pint, transition_energies)
82 :
83 : # Check the interpolation
84 1 : tensors_spherical = [
85 : np.array([[1, 0, 0], [0, -2, 0], [0, 0, 1]]) * i / (freq.m**2 * distance_mum**3)
86 : for (i, freq) in enumerate(transition_energies, start=1)
87 : ]
88 :
89 : # the interpolation near the edges is bad, so we only check the middle
90 1 : for idx in range(3, len(transition_energies) - 3):
91 1 : tensor = gti.get(1, 1, transition_energies[idx], unit="1/micrometer", coordinates="spherical")
92 1 : np.testing.assert_allclose(tensor, tensors_spherical[idx])
93 :
94 1 : for ind in range(3, len(transition_energies) - 5):
95 1 : ind1, ind2 = ind, ind + 1
96 1 : transition_energy = (transition_energies[ind1] + transition_energies[ind2]) / 2
97 1 : reference_tensor = (tensors_spherical[ind1] + tensors_spherical[ind2]) / 2
98 :
99 1 : tensor = gti.get(1, 1, transition_energy, unit="1/micrometer", coordinates="spherical")
100 1 : np.testing.assert_allclose(tensor, reference_tensor, rtol=2e-2)
101 :
102 :
103 1 : def test_system_atom_green_tensor_interpolator(
104 : pi_module: PairinteractionModule,
105 : green_tensor_interpolator_class: type[GreenTensorInterpolator],
106 : ) -> None:
107 : """Test self interaction in a single-atom system from a user-defined Green tensor interpolator."""
108 1 : ket1 = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
109 1 : ket2 = pi_module.KetAtom("Rb", n=60, l=1, j=0.5, m=0.5)
110 :
111 1 : basis = pi_module.BasisAtom("Rb", n=(0, 0), additional_kets=[ket1, ket2])
112 1 : system = pi_module.SystemAtom(basis)
113 1 : reference_hamiltonian = system.get_hamiltonian("hartree").toarray()
114 :
115 1 : dipole = ket1.get_matrix_element(ket2, "electric_dipole", q=0, unit="e * a0")
116 1 : green_tensor_zz = 2.5e-12
117 :
118 1 : gti = green_tensor_interpolator_class()
119 1 : gti.set_constant(
120 : 1,
121 : 1,
122 : np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, green_tensor_zz]]),
123 : "hartree / (e * a0)^2",
124 : )
125 :
126 1 : system._cpp.set_green_tensor_interpolator(gti._cpp)
127 1 : hamiltonian = system.get_hamiltonian("hartree").toarray()
128 1 : expected_shift = green_tensor_zz * abs(dipole) ** 2 / 2
129 1 : np.testing.assert_allclose(hamiltonian, reference_hamiltonian + expected_shift * np.eye(2))
|
