Line data    Source code

       1             : # SPDX-FileCopyrightText: 2024 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4             : """Test the mapping between kets and states."""
       5             : 
       6           1 : from typing import TYPE_CHECKING
       7             : 
       8           1 : import numpy as np
       9           1 : import pairinteraction.real as pi
      10           1 : from scipy.optimize import linear_sum_assignment
      11             : 
      12             : if TYPE_CHECKING:
      13             :     from pairinteraction.units import NDArray
      14             : 
      15             : 
      16           1 : def test_mapping() -> None:
      17             :     """Test generation of a mapping."""
      18             :     # Get the eigenbasis of the Hamiltonian describing an atom in an electric field
      19           1 :     basis = pi.BasisAtom("Rb", n=(58, 62), l=(0, 2))
      20           1 :     system = pi.SystemAtom(basis).set_electric_field([0, 0, 2.5], unit="V/cm")
      21           1 :     system.diagonalize(diagonalizer="eigen", sort_by_energy=True)
      22           1 :     eigenbasis = system.get_eigenbasis()
      23             : 
      24           1 :     assert eigenbasis.number_of_states == eigenbasis.number_of_kets
      25             : 
      26             :     # Obtain the mapping
      27           1 :     state_indices = [eigenbasis.get_corresponding_state_index(ket) for ket in eigenbasis.kets]
      28             : 
      29             :     # Calculate the mapping from the coefficient matrix using scipy
      30           1 :     coefficient_matrix = np.square(np.abs(eigenbasis.get_coefficients().todense()))
      31             :     rows: NDArray
      32             :     cols: NDArray
      33           1 :     rows, cols = linear_sum_assignment(-coefficient_matrix)
      34             : 
      35           1 :     sorter = np.argsort(rows)
      36           1 :     rows = rows[sorter]
      37           1 :     cols = cols[sorter]
      38             : 
      39             :     # Because we have chosen the electric field to be weak enough to avoid strong mixing of states,
      40             :     # the mapping obtained by pairinteraction's heuristic should be the same as the optimal mapping
      41             :     # obtained by scipy's linear_sum_assignment
      42           1 :     np.testing.assert_array_equal(rows, np.arange(eigenbasis.number_of_kets))
      43           1 :     np.testing.assert_array_equal(cols, state_indices)  # TODO