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