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