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