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