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 = (
      53             :         pi_module.SystemAtom(basis)
      54             :         .set_electric_field([1, 0, 1], unit="V/cm")
      55             :         .diagonalize(diagonalizer="eigen", sort_by_energy=True)
      56             :     )
      57             : 
      58           1 :     state = system.basis.get_corresponding_state(ket_initial)
      59             : 
      60           1 :     dipole_zero = state.get_matrix_element(state, "electric_dipole", 0, unit="e a0")
      61           1 :     dipole_plus = state.get_matrix_element(state, "electric_dipole", 1, unit="e a0")
      62           1 :     dipole_minus = state.get_matrix_element(state, "electric_dipole", -1, unit="e a0")
      63             : 
      64           1 :     dipole_z = dipole_zero
      65           1 :     dipole_x = (dipole_minus - dipole_plus) / np.sqrt(2)
      66           1 :     dipole_y = 1j * (dipole_plus + dipole_minus) / np.sqrt(2)
      67             : 
      68           1 :     assert dipole_z.real > 0
      69           1 :     assert dipole_z.imag == 0
      70           1 :     assert np.isclose(dipole_z, dipole_x)
      71           1 :     assert np.isclose(dipole_y, 0)