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 str(state) == "1.00 |Rb:60,P_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 = ket.to_state()
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_negating_state(basis: BasisAtom, pi_module: PairinteractionModule) -> None:
114 : """Negating a state negates its coefficients, and subtracting states is equivalent to adding the negated state."""
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 1 : state1 = basis.get_corresponding_state(ket1)
118 1 : state2 = basis.get_corresponding_state(ket2)
119 1 : assert pytest.approx((-state1).get_coefficients()) == -state1.get_coefficients() # NOSONAR
120 1 : difference = (state1 - state2).normalize()
121 1 : negated_sum = (state1 + (-state2)).normalize()
122 1 : assert pytest.approx(negated_sum.get_coefficients()) == difference.get_coefficients() # NOSONAR
123 :
124 :
125 1 : def test_add_states_with_different_bases_merges(pi_module: PairinteractionModule) -> None:
126 : """Adding states with different bases merges the bases and re-expresses the coefficients."""
127 1 : ket1 = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
128 1 : ket2 = pi_module.KetAtom("Rb", n=60, l=1, j=1.5, m=0.5)
129 :
130 : # Both states start from a minimal single-ket basis, so their bases differ.
131 1 : state1 = ket1.to_state()
132 1 : state2 = ket2.to_state()
133 1 : assert state1.number_of_kets == 1
134 1 : assert state2.number_of_kets == 1
135 :
136 1 : combined = (state1 + 2 * state2).normalize()
137 : # The merged basis contains both kets.
138 1 : assert combined.number_of_kets == 2
139 1 : assert pytest.approx(combined.get_overlap(ket1)) == 1 / 5 # NOSONAR
140 1 : assert pytest.approx(combined.get_overlap(ket2)) == 4 / 5 # NOSONAR
141 :
142 : # Adding a single-ket state onto a full-basis state merges into the full basis.
143 1 : energy_min = ket1.get_energy(unit="GHz") - 100
144 1 : energy_max = ket1.get_energy(unit="GHz") + 100
145 1 : full_basis = pi_module.BasisAtom("Rb", n=(58, 62), l=(0, 2), energy=(energy_min, energy_max), energy_unit="GHz")
146 1 : state_full = full_basis.get_corresponding_state(ket1)
147 1 : merged = (state_full + ket2.to_state()).normalize()
148 1 : assert merged.number_of_kets == full_basis.number_of_kets
149 1 : assert pytest.approx(merged.get_overlap(ket1)) == 0.5 # NOSONAR
150 1 : assert pytest.approx(merged.get_overlap(ket2)) == 0.5 # NOSONAR
151 :
152 :
153 1 : def test_state_from_coefficients_and_kets(pi_module: PairinteractionModule) -> None:
154 : """A state can be created from an arbitrary coefficient vector and a list of kets."""
155 1 : ket1 = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
156 1 : ket2 = pi_module.KetAtom("Rb", n=60, l=1, j=1.5, m=0.5)
157 :
158 1 : state = pi_module.StateAtom([1, 2], [ket1, ket2]).normalize()
159 1 : assert state.number_of_kets == 2
160 1 : assert not state.is_canonical
161 1 : assert pytest.approx(state.get_overlap(ket1)) == 1 / 5 # NOSONAR
162 1 : assert pytest.approx(state.get_overlap(ket2)) == 4 / 5 # NOSONAR
163 :
164 :
165 1 : def test_state_from_coefficients_and_kets_with_basis(basis: BasisAtom, pi_module: PairinteractionModule) -> None:
166 : """Passing a basis expresses the state in the larger Hilbert space of that basis."""
167 1 : ket1 = pi_module.KetAtom("Rb", n=60, l=0, j=0.5, m=0.5)
168 1 : ket2 = pi_module.KetAtom("Rb", n=60, l=1, j=1.5, m=0.5)
169 :
170 1 : state = pi_module.StateAtom([1, 2], [ket1, ket2], basis=basis).normalize()
171 1 : assert state.number_of_kets == basis.number_of_kets
172 1 : assert pytest.approx(state.get_overlap(ket1)) == 1 / 5 # NOSONAR
173 1 : assert pytest.approx(state.get_overlap(ket2)) == 4 / 5 # NOSONAR
174 :
175 : # a ket outside of the basis is rejected, and duplicate kets are rejected as well
176 1 : ket_outside = pi_module.KetAtom("Rb", n=65, l=0, j=0.5, m=0.5)
177 1 : with pytest.raises(ValueError, match="not part of the given basis"):
178 1 : pi_module.StateAtom([1], [ket_outside], basis=basis)
179 1 : with pytest.raises(ValueError, match="must be unique"):
180 1 : pi_module.StateAtom([1, 2], [ket1, ket1], basis=basis)
|
