Line data    Source code

       1             : # SPDX-FileCopyrightText: 2024 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4             : """Test the Stark map calculation."""
       5             : 
       6           1 : from pathlib import Path
       7           1 : from typing import TYPE_CHECKING, Optional
       8             : 
       9           1 : import numpy as np
      10           1 : import pairinteraction.real as pi
      11           1 : import pytest
      12             : 
      13             : if TYPE_CHECKING:
      14             :     from pairinteraction.units import NDArray
      15             : 
      16           1 : reference_kets_file = Path(__file__).parent.parent / "data/reference_stark_map/kets.txt"
      17           1 : reference_eigenenergies_file = Path(__file__).parent.parent / "data/reference_stark_map/eigenenergies.txt"
      18           1 : reference_overlaps_file = Path(__file__).parent.parent / "data/reference_stark_map/overlaps.txt"
      19             : 
      20             : 
      21           1 : def test_starkmap(generate_reference: bool) -> None:
      22             :     """Test calculating a Stark map."""
      23             :     # Create a basis
      24           1 :     ket = pi.KetAtom("Rb", n=60, l=0, m=0.5)
      25           1 :     basis = pi.BasisAtom("Rb", n=(58, 62), l=(0, 2))
      26           1 :     print(f"Number of basis states: {basis.number_of_states}")
      27             : 
      28           1 :     electric_fields = np.linspace(0, 10, 11)
      29             :     # Create systems for different values of the electric field
      30           1 :     systems = [pi.SystemAtom(basis).set_electric_field([0, 0, e], unit="V/cm") for e in electric_fields]
      31             : 
      32             :     # Diagonalize the systems in parallel
      33           1 :     pi.diagonalize(systems, diagonalizer="eigen", sort_by_energy=True)
      34             : 
      35             :     # Get the overlap with |ket>
      36           1 :     overlaps = np.array([system.basis.get_overlaps(ket) for system in systems])
      37             : 
      38             :     # Compare to reference data
      39           1 :     kets = [repr(ket) for ket in systems[0].basis.kets]
      40           1 :     eigenenergies = np.array([system.get_eigenenergies(unit="GHz") for system in systems])
      41           1 :     eigenvectors = np.array([system.get_eigenbasis().get_coefficients().todense().A1 for system in systems])
      42             : 
      43           1 :     if generate_reference:
      44           1 :         reference_kets_file.parent.mkdir(parents=True, exist_ok=True)
      45           1 :         np.savetxt(reference_kets_file, kets, fmt="%s", delimiter="\t")
      46           1 :         np.savetxt(reference_eigenenergies_file, eigenenergies)
      47           1 :         np.savetxt(reference_overlaps_file, overlaps)
      48           1 :         pytest.skip("Reference data generated, skipping comparison test")
      49             : 
      50           1 :     compare_starkmap_to_reference(eigenenergies, overlaps, eigenvectors, kets)
      51             : 
      52             : 
      53           1 : def compare_starkmap_to_reference(
      54             :     eigenenergies: "NDArray",
      55             :     overlaps: Optional["NDArray"] = None,
      56             :     eigenvectors: Optional["NDArray"] = None,
      57             :     kets: Optional[list[str]] = None,
      58             : ) -> None:
      59           1 :     np.testing.assert_allclose(eigenenergies, np.loadtxt(reference_eigenenergies_file))
      60             : 
      61           1 :     if overlaps is not None:
      62             :         # Ensure that the overlaps sum up to one
      63           1 :         np.testing.assert_allclose(np.sum(overlaps, axis=1), np.ones(len(eigenenergies)))
      64           1 :         np.testing.assert_allclose(overlaps, np.loadtxt(reference_overlaps_file), atol=1e-10)
      65             : 
      66           1 :     if kets is not None:
      67           1 :         np.testing.assert_equal(kets, np.loadtxt(reference_kets_file, dtype=str, delimiter="\t"))
      68             : 
      69           1 :     if eigenvectors is not None:
      70             :         # Because of degeneracies, checking the eigenvectors against reference data is complicated.
      71             :         # Thus, we only check their normalization and orthogonality.
      72           1 :         cumulative_norm = (np.array(eigenvectors) * np.array(eigenvectors).conj()).sum(axis=1)
      73           1 :         np.testing.assert_allclose(cumulative_norm, 90 * np.ones(len(eigenenergies)))