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