Line data    Source code

       1             : # SPDX-FileCopyrightText: 2024 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4             : """Test energy range restriction in the diagonalization."""
       5             : 
       6           1 : import numpy as np
       7             : 
       8           1 : import pairinteraction.real as pi
       9             : 
      10             : 
      11           1 : def test_energy_range() -> None:
      12             :     """Test restricting the energy range in the diagonalization."""
      13           1 :     ket = pi.KetAtom("Rb", n=60, l=0, m=0.5)
      14           1 :     pair_energy = 2 * ket.get_energy(unit="GHz")
      15           1 :     distances = np.linspace(1, 4, 10)
      16             : 
      17             :     # Create a single-atom system
      18           1 :     basis = pi.BasisAtom("Rb", n=(58, 62), l=(0, 1))
      19           1 :     system = pi.SystemAtom(basis)
      20             : 
      21             :     # Create two-atom basis
      22           1 :     basis_pair = pi.BasisPair(
      23             :         [system, system], energy=(pair_energy - 10, pair_energy + 10), energy_unit="GHz", m=(1, 1)
      24             :     )
      25             : 
      26             :     # Diagonalize the systems for different distances in parallel and get all eigenenergies
      27           1 :     system_pairs = [pi.SystemPair(basis_pair).set_distance(d, unit="micrometer") for d in distances]
      28           1 :     pi.diagonalize(system_pairs, diagonalizer="eigen", sort_by_energy=True)
      29           1 :     eigenenergies_all = [system.get_eigenenergies(unit="GHz") for system in system_pairs]
      30             : 
      31             :     # Diagonalize the systems for different distances in parallel and get only the eigenenergies in an energy range
      32           1 :     system_pairs = [pi.SystemPair(basis_pair).set_distance(d, unit="micrometer") for d in distances]
      33           1 :     pi.diagonalize(
      34             :         system_pairs,
      35             :         diagonalizer="eigen",
      36             :         sort_by_energy=True,
      37             :         energy_range=(pair_energy - 5, pair_energy + 5),
      38             :         energy_unit="GHz",
      39             :     )
      40           1 :     eigenenergies_restricted = [system.get_eigenenergies(unit="GHz") for system in system_pairs]
      41             : 
      42             :     # Check the result
      43           1 :     eigenenergies_all_restricted = [
      44             :         eigenenergies[(eigenenergies < pair_energy + 5) & (eigenenergies > pair_energy - 5)]
      45             :         for eigenenergies in eigenenergies_all
      46             :     ]
      47           1 :     for e1, e2 in zip(eigenenergies_restricted, eigenenergies_all_restricted):
      48           1 :         np.testing.assert_allclose(e1, e2)