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", sort_by_energy=True)
      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)