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 : if TYPE_CHECKING:
12 : from pairinteraction import BasisAtom, StateAtom
13 :
14 : from .utils import PairinteractionModule
15 :
16 :
17 1 : @pytest.fixture
18 1 : def basis(pi_module: PairinteractionModule) -> BasisAtom:
19 : """Create a test basis with a few states around Rb 60S."""
20 1 : ket = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
21 1 : energy_min = ket.get_energy(unit="GHz") - 100
22 1 : energy_max = ket.get_energy(unit="GHz") + 100
23 1 : return pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2), energy=(energy_min, energy_max), energy_unit="GHz")
24 :
25 :
26 1 : @pytest.fixture
27 1 : def state(pi_module: PairinteractionModule, basis: BasisAtom) -> StateAtom:
28 : """Create a test state."""
29 1 : ket = pi_module.KetAtom("Rb", n=60, l=1, j=1.5, m=-0.5)
30 1 : return basis.get_corresponding_state(ket)
31 :
32 :
33 1 : def test_state_creation(state: StateAtom) -> None:
34 : """Test basic properties of created state."""
35 1 : assert state.species == "Rb"
36 1 : assert state.number_of_kets == 80
37 1 : assert len(state.kets) == state.number_of_kets
38 1 : assert all(x in str(state) for x in ["StateAtom", "60", "S", "3/2", "-1/2"])
39 1 : assert state.is_canonical
40 :
41 :
42 1 : def test_coefficients(state: StateAtom) -> None:
43 : """Test coefficient matrix properties."""
44 1 : coeffs = state.get_coefficients()
45 1 : assert coeffs.shape == (state.number_of_kets,)
46 1 : assert np.count_nonzero(coeffs) == 1
47 1 : assert pytest.approx(coeffs.sum()) == 1.0 # NOSONAR
48 :
49 :
50 1 : def test_get_amplitude_and_overlap(state: StateAtom) -> None:
51 : """Test amplitude and overlap calculations."""
52 : # Test with ket
53 1 : test_ket = state.get_corresponding_ket()
54 1 : amplitude = state.get_amplitude(test_ket)
55 1 : assert np.isscalar(amplitude)
56 1 : assert pytest.approx(amplitude) == 1.0 # NOSONAR
57 1 : overlap = state.get_overlap(test_ket)
58 1 : assert np.isscalar(overlap)
59 1 : assert pytest.approx(overlap) == 1.0 # NOSONAR
60 :
61 : # Test with state
62 1 : amplitude = state.get_amplitude(state)
63 1 : assert np.isscalar(amplitude)
64 1 : assert pytest.approx(amplitude) == 1.0 # NOSONAR
65 1 : overlap = state.get_overlap(state)
66 1 : assert np.isscalar(overlap)
67 1 : assert pytest.approx(overlap) == 1.0 # NOSONAR
68 :
69 :
70 1 : def test_get_matrix_element(pi_module: PairinteractionModule, basis: BasisAtom) -> None:
71 : """Test matrix element calculations."""
72 1 : ket1 = pi_module.KetAtom("Rb", n=60, l=1, j=1.5, m=-0.5)
73 1 : ket2 = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
74 1 : state1 = basis.get_corresponding_state(ket1)
75 1 : state2 = basis.get_corresponding_state(ket2)
76 :
77 : # Test with ket
78 1 : element_dipole_ket = state1.get_matrix_element(ket2, "electric_dipole", q=1, unit="e * a0")
79 1 : assert np.isscalar(element_dipole_ket)
80 1 : assert element_dipole_ket != 0
81 :
82 : # Test with state
83 1 : element_dipole_state = state1.get_matrix_element(state2, "electric_dipole", q=1, unit="e * a0")
84 1 : assert np.isscalar(element_dipole_state)
85 1 : assert pytest.approx(element_dipole_ket) == element_dipole_state # NOSONAR
86 1 : assert state1.get_matrix_element(state1, "electric_dipole", q=1, unit="e * a0") == 0
87 1 : assert state1.get_matrix_element(state2, "electric_dipole", q=0, unit="e * a0") == 0
88 :
89 :
90 1 : def test_state_without_basis(pi_module: PairinteractionModule) -> None:
91 : """A StateAtom created without an explicit basis uses a minimal single-ket basis."""
92 1 : ket = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
93 1 : state = pi_module.StateAtom(ket)
94 1 : assert state.number_of_kets == 1
95 1 : assert state.is_canonical
96 1 : assert state.get_corresponding_ket() == ket
97 1 : assert pytest.approx(state.get_overlap(ket)) == 1.0 # NOSONAR
98 :
99 :
100 1 : def test_add_states_with_same_basis(basis: BasisAtom, pi_module: PairinteractionModule) -> None:
101 : """Adding states expressed in the same basis keeps that basis."""
102 1 : ket1 = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
103 1 : ket2 = pi_module.KetAtom("Rb", n=60, l=1, j=1.5, m=0.5)
104 1 : state1 = basis.get_corresponding_state(ket1)
105 1 : state2 = basis.get_corresponding_state(ket2)
106 :
107 1 : combined = (state1 + state2).normalize()
108 1 : assert combined.number_of_kets == basis.number_of_kets
109 1 : assert pytest.approx(combined.get_overlap(ket1)) == 0.5 # NOSONAR
110 1 : assert pytest.approx(combined.get_overlap(ket2)) == 0.5 # NOSONAR
111 :
112 :
113 1 : def test_add_states_with_different_bases_merges(pi_module: PairinteractionModule) -> None:
114 : """Adding states with different bases merges the bases and re-expresses the coefficients."""
115 1 : ket1 = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
116 1 : ket2 = pi_module.KetAtom("Rb", n=60, l=1, j=1.5, m=0.5)
117 :
118 : # Both states start from a minimal single-ket basis, so their bases differ.
119 1 : state1 = pi_module.StateAtom(ket1)
120 1 : state2 = pi_module.StateAtom(ket2)
121 1 : assert state1.number_of_kets == 1
122 1 : assert state2.number_of_kets == 1
123 :
124 1 : combined = (state1 + 2 * state2).normalize()
125 : # The merged basis contains both kets.
126 1 : assert combined.number_of_kets == 2
127 1 : assert pytest.approx(combined.get_overlap(ket1)) == 1 / 5 # NOSONAR
128 1 : assert pytest.approx(combined.get_overlap(ket2)) == 4 / 5 # NOSONAR
129 :
130 : # Adding a single-ket state onto a full-basis state merges into the full basis.
131 1 : energy_min = ket1.get_energy(unit="GHz") - 100
132 1 : energy_max = ket1.get_energy(unit="GHz") + 100
133 1 : full_basis = pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2), energy=(energy_min, energy_max), energy_unit="GHz")
134 1 : state_full = full_basis.get_corresponding_state(ket1)
135 1 : merged = (state_full + pi_module.StateAtom(ket2)).normalize()
136 1 : assert merged.number_of_kets == full_basis.number_of_kets
137 1 : assert pytest.approx(merged.get_overlap(ket1)) == 0.5 # NOSONAR
138 1 : assert pytest.approx(merged.get_overlap(ket2)) == 0.5 # NOSONAR
139 :
140 :
141 1 : def test_error_handling(state: StateAtom) -> None:
142 : """Test error cases."""
143 1 : with pytest.raises(TypeError):
144 1 : state.get_amplitude("not a ket") # type: ignore [arg-type]
145 :
146 1 : with pytest.raises(TypeError):
147 1 : state.get_overlap("not a ket") # type: ignore [arg-type]
148 :
149 1 : with pytest.raises(TypeError):
150 1 : state.get_matrix_element("not a ket", "energy", 0) # type: ignore [call-overload]
|
