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             : 
      11             : if TYPE_CHECKING:
      12             :     from .utils import PairinteractionModule
      13             : 
      14             : 
      15           1 : @pytest.mark.parametrize("distance_mum", [1, 2, 11])
      16           1 : def test_static_green_tensor(pi_module: PairinteractionModule, distance_mum: float) -> None:
      17             :     """Test calculating a pair potential using a user-defined static green tensor."""
      18             :     # Create a single-atom system
      19           1 :     basis = pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2))
      20           1 :     system = pi_module.SystemAtom(basis)
      21             : 
      22             :     # Create two-atom basis
      23           1 :     ket = pi_module.KetAtom("Rb", n=60, l=0, m=0.5)
      24           1 :     delta_energy = 3  # GHz
      25           1 :     min_energy = 2 * ket.get_energy(unit="GHz") - delta_energy
      26           1 :     max_energy = 2 * ket.get_energy(unit="GHz") + delta_energy
      27           1 :     basis_pair = pi_module.BasisPair([system, system], energy=(min_energy, max_energy), energy_unit="GHz", m=(1, 1))
      28             : 
      29             :     # Create a system using a user-defined green tensor for dipole-dipole interaction
      30           1 :     gt = pi_module.GreenTensor()
      31           1 :     tensor = np.array([[1, 0, 0], [0, 1, 0], [0, 0, -2]]) / distance_mum**3
      32           1 :     tensor_unit = "hartree / (e^2 micrometer^3)"
      33           1 :     gt.set_from_cartesian(1, 1, tensor, tensor_unit)
      34             : 
      35           1 :     tensor_spherical = np.array([[1, 0, 0], [0, -2, 0], [0, 0, 1]]) / distance_mum**3
      36           1 :     np.testing.assert_allclose(gt.get_spherical(1, 1, unit=tensor_unit), tensor_spherical)
      37           1 :     np.testing.assert_allclose(gt.get_spherical(1, 1, omega=2.5, omega_unit="GHz", unit=tensor_unit), tensor_spherical)
      38             : 
      39           1 :     system_pairs = pi_module.SystemPair(basis_pair).set_green_tensor(gt)
      40             : 
      41             :     # Create a reference system using the build in dipole-dipole interaction
      42           1 :     system_pairs_reference = (
      43             :         pi_module.SystemPair(basis_pair).set_interaction_order(3).set_distance(distance_mum, unit="micrometer")
      44             :     )
      45             : 
      46             :     # Check that the two systems are equivalent
      47           1 :     np.testing.assert_allclose(
      48             :         system_pairs.get_hamiltonian(unit="GHz").data, system_pairs_reference.get_hamiltonian(unit="GHz").data
      49             :     )
      50             : 
      51             : 
      52           1 : @pytest.mark.parametrize("distance_mum", [1, 2, 11])
      53           1 : def test_omega_dependent_green_tensor(pi_module: PairinteractionModule, distance_mum: float) -> None:
      54             :     """Test the interpolation for different values of omega."""
      55             :     # Define an simple linear omega-dependent green tensor
      56             :     # note that at least four entries are needed for the applied spline interpolation.
      57           1 :     gt = pi_module.GreenTensor()
      58           1 :     omegas = [1, 2, 3, 4]  # GHz
      59           1 :     tensors = [np.array([[1, 0, 0], [0, 1, 0], [0, 0, -2]]) * i / distance_mum**3 for i in range(1, 5)]
      60           1 :     tensor_unit = "hartree / (e^2 micrometer^3)"
      61           1 :     gt.set_from_cartesian(1, 1, tensors, tensor_unit, omegas, omegas_unit="GHz")
      62             : 
      63             :     # Check the interpolation
      64           1 :     tensors_spherical = [np.array([[1, 0, 0], [0, -2, 0], [0, 0, 1]]) * i / distance_mum**3 for i in range(1, 5)]
      65             : 
      66           1 :     for idx in range(1, len(omegas) - 1):  # the interpolation near the edges is bad, so we only check the middle
      67           1 :         tensor = gt.get_spherical(1, 1, omega=omegas[idx], omega_unit="GHz", unit=tensor_unit)
      68           1 :         np.testing.assert_allclose(tensor, tensors_spherical[idx])
      69             : 
      70           1 :     tensor = gt.get_spherical(1, 1, omega=2.5, omega_unit="GHz", unit=tensor_unit)
      71           1 :     reference_tensor = (tensors_spherical[1] + tensors_spherical[2]) / 2
      72           1 :     np.testing.assert_allclose(tensor, reference_tensor, rtol=2e-2)