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              : 
      10              : if TYPE_CHECKING:
      11              :     from .utils import PairinteractionModule
      12              : 
      13              : 
      14            1 : def test_energy(pi_module: PairinteractionModule) -> None:
      15              :     """Test calculating energies of ket states."""
      16              :     # Energy of unperturbed state
      17            1 :     ket = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
      18            1 :     energy_unperturbed = ket.get_energy(unit="GHz")
      19              : 
      20            1 :     assert np.isclose(energy_unperturbed, ket.get_energy("GHz"))
      21              : 
      22              :     # Energy of Stark shifted state
      23            1 :     basis = pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2), m=(0.5, 0.5))
      24              : 
      25            1 :     system = pi_module.SystemAtom(basis).set_electric_field([0, 0, 1], unit="V/cm").diagonalize(diagonalizer="eigen")
      26              : 
      27            1 :     energy_perturbed = system.get_corresponding_energy(ket, unit="GHz")
      28              : 
      29            1 :     shift = energy_perturbed - energy_unperturbed
      30            1 :     print(f"Energy shift: {shift} GHz")
      31              : 
      32            1 :     assert shift < 0
      33              : 
      34              : 
      35            1 : def test_electric_dipole_matrix_element(pi_module: PairinteractionModule) -> None:
      36              :     """Test calculating dipole matrix elements."""
      37              :     # The dipole element between dipole-coupled states should be non-zero
      38            1 :     ket_initial = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
      39            1 :     ket_final = pi_module.KetAtom("Rb", n=60, l=1, j=0.5, m=0.5)
      40              : 
      41            1 :     dipole = ket_initial.get_matrix_element(ket_final, "electric_dipole", 0, unit="e a0")
      42            1 :     assert dipole.real > 0
      43            1 :     assert dipole.imag == 0
      44              : 
      45              :     # The dipole element between the same, unperturbed state should be zero
      46            1 :     dipole = ket_initial.get_matrix_element(ket_initial, "electric_dipole", 0, unit="e a0")
      47            1 :     assert np.isclose(dipole, 0)
      48              : 
      49              :     # Stark effect induces a permanent dipole moment
      50            1 :     basis = pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2))
      51              : 
      52            1 :     system = pi_module.SystemAtom(basis).set_electric_field([1, 0, 1], unit="V/cm").diagonalize(diagonalizer="eigen")
      53              : 
      54            1 :     state = system.basis.get_corresponding_state(ket_initial)
      55              : 
      56            1 :     dipole_zero = state.get_matrix_element(state, "electric_dipole", 0, unit="e a0")
      57            1 :     dipole_plus = state.get_matrix_element(state, "electric_dipole", 1, unit="e a0")
      58            1 :     dipole_minus = state.get_matrix_element(state, "electric_dipole", -1, unit="e a0")
      59              : 
      60            1 :     dipole_z = dipole_zero
      61            1 :     dipole_x = (dipole_minus - dipole_plus) / np.sqrt(2)
      62            1 :     dipole_y = 1j * (dipole_plus + dipole_minus) / np.sqrt(2)
      63              : 
      64            1 :     assert dipole_z.real > 0
      65            1 :     assert dipole_z.imag == 0
      66            1 :     assert np.isclose(dipole_z, dipole_x)
      67            1 :     assert np.isclose(dipole_y, 0)