def test_two_representations_correlation():
x = np.linspace(-100, 100, 10001)
x1 = 0.0
x2 = 3.0
sigma1 = 1.0
sigma2 = 5.0
psi1 = __gauss(x, x1, sigma1)
psi2 = __gauss(x, x2, sigma2)
p, psi1_p = wavefunction.momentum_representation(x, psi1)
_, psi2_p = wavefunction.momentum_representation(x, psi2)
corr1 = wavefunction.correlation(x, psi1, psi2)
corr2 = wavefunction.correlation(p, psi1_p, psi2_p)
np.testing.assert_almost_equal(corr1, corr2)
def test_correlation():
x = np.linspace(-100, 100, 10001)
psi1 = __gauss(x, 0.0, 2.0)
phase_mult = np.exp(0.25j * np.pi)
psi2 = psi1 * phase_mult
np.testing.assert_allclose(wavefunction.norm(x, psi1),
wavefunction.correlation(x, psi1, psi1))
np.testing.assert_allclose(wavefunction.norm(x, psi2),
wavefunction.correlation(x, psi2, psi2))
np.testing.assert_allclose(
wavefunction.norm(x, psi1) * phase_mult,
wavefunction.correlation(x, psi1, psi2))
np.testing.assert_allclose(
wavefunction.norm(x, psi1) * np.conj(phase_mult),
wavefunction.correlation(x, psi2, psi1))
def test_coulomb_potential():
x = np.linspace(-30, 30, 201)
xx, yy, zz = np.meshgrid(x, x, x, indexing='ij')
r = (xx, yy, zz)
levels = [(1, 0, 0), (2, 0, 0), (2, 1, 1), (2, 1, -1), (3, 2, -1)]
v = potential.CoulombPotential()
psis = []
for l in levels:
e = v.get_eigenenergy(*l)
np.testing.assert_allclose(e, -0.5 / l[0]**2)
psi = v.get_eigenfunction(*l)(*r)
psis.append(psi)
np.testing.assert_array_almost_equal(
wavefunction.norm(r, psi),
1.0,
decimal=3,
err_msg=f"Wavefunction norm for level {l} is not unity")
from itertools import combinations
for psi1, psi2 in combinations(psis, 2):
np.testing.assert_allclose(wavefunction.correlation(r, psi1, psi2),
0.0,
atol=0.001,
err_msg=f"Non-orthogonal wavefunctions")
def test_quadratic_orthogonality():
frequencies = [0.1, 1.0, 7.5]
x = np.linspace(-50, 50, 40000)
levels = range(10)
for f in frequencies:
v = potential.QuadraticPotential1D(f)
for l1 in levels:
for l2 in levels[l1 + 1:]:
psi1 = v.get_eigenfunction(l1)(x)
psi2 = v.get_eigenfunction(l2)(x)
np.testing.assert_almost_equal(
wavefunction.correlation(x, psi1, psi2),
0.0,
err_msg=
f'Functions for levels {l1} and {l2} are not orthogonal '
f'for frequency {f}')
def test_square_potential_orthogonality():
from itertools import combinations
widths = [1.0, 2.0]
depths = [10.0, 5.0]
x = np.linspace(-15, 15, 3000)
for V0, a in zip(depths, widths):
v = potential.SquarePotential1D(V0, a)
assert v.get_depth() == V0, f"Depth {v.get_depth()} is not {V0}"
assert v.get_width() == a, f"Width {v.get_width()} is not {a}"
psis = [
v.get_eigenfunction(n)(x) for n in range(v.get_number_of_levels())
]
for psi1, psi2 in combinations(psis, 2):
np.testing.assert_almost_equal(
wavefunction.correlation(x, psi1, psi2),
0.0,
err_msg="Non-orthogonal eigenfunctions for V0={V0}, a={a}")
def test_eigen_quadratic():
freqs = [0.5, 1.0, 5.0]
x = np.linspace(-200, 200, 10000)
levels = range(10)
for f in freqs:
v = potential.QuadraticPotential1D(f)
es, psis = eigen.calculate_eigenstates(x, v.get_potential(), 10, 0.0)
for i, psi in enumerate(psis):
np.testing.assert_almost_equal(wavefunction.norm(x, psi), 1.0, err_msg=f'Eigenfunction {i} norm is not 1')
e0s = np.array([v.get_eigenenergy(l) for l in levels])
psi0s = [v.get_eigenfunction(l)(x) for l in levels]
np.testing.assert_allclose(e0s, es, rtol=0.02, err_msg=f"Energy spectra are different for frequency {f}")
assert len(psis) == len(psi0s), f"Incorrect number of eigenfunctions {len(psis)}, expected {len(psi0s)}"
for i, (psi0, psi) in enumerate(zip(psi0s, psis)):
corr = np.abs(wavefunction.correlation(x, psi0, psi))
np.testing.assert_almost_equal(corr, 1.0, decimal=3, err_msg=f'Function {i} is incorrect')
def test_eigen_square():
widths = [4.0, 6.0]
depths = [2., 5.]
x = np.linspace(-100, 100, 10000)
for a, V0 in zip(widths, depths):
v = potential.SquarePotential1D(V0, a)
levels = v.get_number_of_levels()
es, psis = eigen.calculate_eigenstates(x, v.get_potential(), levels, -V0)
for i, psi in enumerate(psis):
np.testing.assert_almost_equal(wavefunction.norm(x, psi), 1.0, err_msg=f'Eigenfunction {i} norm is not 1')
e0s = np.array([v.get_eigenenergy(l) for l in range(levels)])
psi0s = [v.get_eigenfunction(l)(x) for l in range(levels)]
np.testing.assert_allclose(e0s, es, atol=0.02, rtol=0.05,
err_msg=f"Energy spectra are different for a={a}, V0={V0}")
assert len(psis) == len(psi0s), f"Incorrect number of eigenfunctions {len(psis)}, expected {len(psi0s)}"
for i, (psi0, psi) in enumerate(zip(psi0s, psis)):
corr = np.abs(wavefunction.correlation(x, psi0, psi))
np.testing.assert_almost_equal(corr, 1.0, decimal=3, err_msg=f'Function {i} is incorrect')