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")
      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              :         energy_range=(pair_energy - 5, pair_energy + 5),
      43              :         energy_range_unit="GHz",
      44              :     )
      45            1 :     eigenenergies_restricted = [system.get_eigenenergies(unit="GHz") for system in system_pairs]
      46              : 
      47              :     # Check the result
      48            1 :     eigenenergies_all_restricted = [
      49              :         eigenenergies[(eigenenergies < pair_energy + 5) & (eigenenergies > pair_energy - 5)]
      50              :         for eigenenergies in eigenenergies_all
      51              :     ]
      52            1 :     for e1, e2 in zip(eigenenergies_restricted, eigenenergies_all_restricted, strict=True):
      53            1 :         np.testing.assert_allclose(e1, e2)