Line data    Source code

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