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 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)._set_green_tensor_interpolator(gti)
      53             : 
      54             :     # Create a reference system using the build in dipole-dipole interaction
      55           1 :     system_pairs_reference = (
      56             :         pi_module.SystemPair(basis_pair).set_interaction_order(3).set_distance(distance_mum, unit="micrometer")
      57             :     )
      58             : 
      59             :     # Check that the two systems are equivalent
      60           1 :     np.testing.assert_allclose(
      61             :         system_pairs.get_hamiltonian(unit="GHz").data, system_pairs_reference.get_hamiltonian(unit="GHz").data
      62             :     )
      63             : 
      64             : 
      65           1 : @pytest.mark.parametrize("distance_mum", [1, 2, 11])
      66           1 : def test_omega_dependent_green_tensor_interpolator(
      67             :     green_tensor_interpolator_class: type[GreenTensorInterpolator],
      68             :     distance_mum: float,
      69             : ) -> None:
      70             :     """Test the interpolation for different values of the transition energy."""
      71             :     # Define an simple linear omega-dependent green tensor interpolator
      72             :     # note that at least four entries are needed for the applied spline interpolation.
      73           1 :     gti = green_tensor_interpolator_class()
      74           1 :     transition_energies = [ureg.Quantity(i, "planck_constant * GHz") for i in np.linspace(1, 5, 20)]
      75           1 :     tensors = [
      76             :         np.array([[1, 0, 0], [0, 1, 0], [0, 0, -2]]) * i / (freq.m**2 * distance_mum**3)
      77             :         for (i, freq) in enumerate(transition_energies, start=1)
      78             :     ]
      79           1 :     tensors_pint = [ureg.Quantity(tensor, "1/micrometer") for tensor in tensors]
      80           1 :     gti.set_list(1, 1, tensors_pint, transition_energies)
      81             : 
      82             :     # Check the interpolation
      83           1 :     tensors_spherical = [
      84             :         np.array([[1, 0, 0], [0, -2, 0], [0, 0, 1]]) * i / (freq.m**2 * distance_mum**3)
      85             :         for (i, freq) in enumerate(transition_energies, start=1)
      86             :     ]
      87             : 
      88             :     # the interpolation near the edges is bad, so we only check the middle
      89           1 :     for idx in range(3, len(transition_energies) - 3):
      90           1 :         tensor = gti.get(1, 1, transition_energies[idx], unit="1/micrometer", coordinates="spherical")
      91           1 :         np.testing.assert_allclose(tensor, tensors_spherical[idx])
      92             : 
      93           1 :     for ind in range(3, len(transition_energies) - 5):
      94           1 :         ind1, ind2 = ind, ind + 1
      95           1 :         transition_energy = (transition_energies[ind1] + transition_energies[ind2]) / 2
      96           1 :         reference_tensor = (tensors_spherical[ind1] + tensors_spherical[ind2]) / 2
      97             : 
      98           1 :         tensor = gti.get(1, 1, transition_energy, unit="1/micrometer", coordinates="spherical")
      99           1 :         np.testing.assert_allclose(tensor, reference_tensor, rtol=2e-2)