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