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            1 : import pytest
      10              : 
      11            1 : from .utils import REFERENCE_PATHS, compare_eigensystem_to_reference
      12              : 
      13              : if TYPE_CHECKING:
      14              :     from .utils import PairinteractionModule
      15              : 
      16              : 
      17            1 : def test_starkmap(pi_module: PairinteractionModule, generate_reference: bool) -> None:
      18              :     """Test calculating a Stark map."""
      19              :     # Create a basis
      20            1 :     ket = pi_module.KetAtom("Rb", n=60, l=0, m=0.5)
      21            1 :     basis = pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2))
      22            1 :     print(f"Number of basis states: {basis.number_of_states}")
      23              : 
      24            1 :     electric_fields = np.linspace(0, 10, 11)
      25              :     # Create systems for different values of the electric field
      26            1 :     systems = [pi_module.SystemAtom(basis).set_electric_field([0, 0, e], unit="V/cm") for e in electric_fields]
      27              : 
      28              :     # Diagonalize the systems in parallel
      29            1 :     pi_module.diagonalize(systems, diagonalizer="eigen")
      30              : 
      31              :     # Get the overlap with |ket>
      32            1 :     overlaps = np.array([system.basis.get_overlaps(ket) for system in systems])
      33              : 
      34              :     # Compare to reference data
      35            1 :     kets = [ket.get_label("raw") for ket in systems[0].basis.kets]
      36            1 :     eigenenergies = np.array([system.get_eigenenergies(unit="GHz") for system in systems])
      37            1 :     eigenvectors = np.array([system.get_eigenbasis().get_coefficients().todense().A1 for system in systems])
      38              : 
      39            1 :     reference_path = REFERENCE_PATHS["stark_map"]
      40            1 :     if generate_reference:
      41            1 :         reference_path.mkdir(parents=True, exist_ok=True)
      42            1 :         np.savetxt(reference_path / "kets.txt", kets, fmt="%s", delimiter="\t")
      43            1 :         np.savetxt(reference_path / "eigenenergies.txt", eigenenergies)
      44            1 :         np.savetxt(reference_path / "overlaps.txt", overlaps)
      45            1 :         pytest.skip("Reference data generated, skipping comparison test")
      46              : 
      47            1 :     compare_eigensystem_to_reference(reference_path, eigenenergies, overlaps, eigenvectors, kets)