Line data    Source code

       1             : # SPDX-FileCopyrightText: 2024 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4             : """Test the calculation of matrix elements."""
       5             : 
       6           1 : import numpy as np
       7           1 : import pairinteraction.real as pi
       8             : 
       9             : 
      10           1 : def test_energy() -> None:
      11             :     """Test calculating energies of ket states."""
      12             :     # Energy of unperturbed state
      13           1 :     ket = pi.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
      14           1 :     energy_unperturbed = ket.get_energy(unit="GHz")
      15             : 
      16           1 :     assert np.isclose(energy_unperturbed, ket.get_energy("GHz"))
      17             : 
      18             :     # Energy of Stark shifted state
      19           1 :     basis = pi.BasisAtom("Rb", n=(58, 62), l=(0, 2), m=(0.5, 0.5))
      20             : 
      21           1 :     system = pi.SystemAtom(basis).set_electric_field([0, 0, 1], unit="V/cm").diagonalize(diagonalizer="eigen")
      22             : 
      23           1 :     energy_perturbed = system.get_corresponding_energy(ket, unit="GHz")
      24             : 
      25           1 :     shift = energy_perturbed - energy_unperturbed
      26           1 :     print(f"Energy shift: {shift} GHz")
      27             : 
      28           1 :     assert shift < 0
      29             : 
      30             : 
      31           1 : def test_electric_dipole_matrix_element() -> None:
      32             :     """Test calculating dipole matrix elements."""
      33             :     # The dipole element between dipole-coupled states should be non-zero
      34           1 :     ket_initial = pi.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
      35           1 :     ket_final = pi.KetAtom("Rb", n=60, l=1, j=0.5, m=0.5)
      36             : 
      37           1 :     dipole = ket_initial.get_matrix_element(ket_final, "electric_dipole", 0, unit="e a0")
      38           1 :     assert dipole.real > 0
      39           1 :     assert dipole.imag == 0
      40             : 
      41             :     # The dipole element between the same, unperturbed state should be zero
      42           1 :     dipole = ket_initial.get_matrix_element(ket_initial, "electric_dipole", 0, unit="e a0")
      43           1 :     assert np.isclose(dipole, 0)
      44             : 
      45             :     # Stark effect induces a permanent dipole moment
      46           1 :     basis = pi.BasisAtom("Rb", n=(58, 62), l=(0, 2))
      47             : 
      48           1 :     system = (
      49             :         pi.SystemAtom(basis)
      50             :         .set_electric_field([1, 0, 1], unit="V/cm")
      51             :         .diagonalize(diagonalizer="eigen", sort_by_energy=True)
      52             :     )
      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)