def test_momentum_representation_planewave():
x = np.linspace(-100, 100, 10001)
l = (x[-1] - x[0] + x[1] - x[0])
p0s = [10 * np.pi / l, 100 * np.pi / l, 1000 * np.pi / l]
for p0 in p0s:
psi = np.exp(1.0j * p0 * x) / np.sqrt(l)
np.testing.assert_array_almost_equal(wavefunction.norm(x, psi),
1.0,
decimal=4)
p, psi_p = wavefunction.momentum_representation(x, psi)
np.testing.assert_equal(np.sort(p),
p,
err_msg="Array of p is not sorted")
np.testing.assert_array_almost_equal(
wavefunction.norm(p, psi_p),
1.0,
decimal=4,
err_msg="The norm is not conserved in momentum representation")
p_max = p[np.argmax(np.abs(psi_p))]
np.testing.assert_allclose(
p0,
p_max,
err_msg=f"Spectral maximum {p_max} is not equal to expected {p0}")
def test_euler_delta():
depths = [0.2, 1.0, 5.0]
for d in depths:
v = potential.DeltaPotential1D(d)
e0 = v.get_eigenenergy()
tmax = -2 * np.pi / e0
dt = tmax / 500000
s = solver.EulerSolver(20 / d, 0.1 / d, dt, v)
psi0 = v.get_eigenfunction()
psi0_value = psi0(s.x)
np.testing.assert_almost_equal(wavefunction.norm(s.x, psi0_value),
1,
decimal=2)
psi = s.execute(tmax, psi0=psi0)
np.testing.assert_almost_equal(
wavefunction.norm(s.x, psi),
wavefunction.norm(s.x, psi0_value),
decimal=3,
err_msg=f"Norm not conserved for depth {d}")
np.testing.assert_allclose(np.abs(psi), np.abs(psi0_value), atol=0.06)
def test_coordinate_representation_planewave():
p = np.linspace(-100, 100, 10001)
l = (p[-1] - p[0] + p[1] - p[0])
x0s = [10 * np.pi / l, 100 * np.pi / l, 1000 * np.pi / l]
for x0 in x0s:
psi_p = np.exp(-1.0j * p * x0) / np.sqrt(l)
np.testing.assert_array_almost_equal(wavefunction.norm(p, psi_p),
1.0,
decimal=4)
x, psi = wavefunction.coordinate_representation(p, psi_p)
np.testing.assert_equal(np.sort(x),
x,
err_msg="Array of x is not sorted")
np.testing.assert_array_almost_equal(
wavefunction.norm(x, psi),
1.0,
decimal=4,
err_msg="The norm is not conserved in coordinate representation")
x_max = x[np.argmax(np.abs(psi))]
np.testing.assert_allclose(
x0,
x_max,
err_msg=f"Spectral maximum {x_max} is not equal to expected {x0}")
def test_cn_delta():
depths = [0.2, 1.0, 5.0]
for d in depths:
v = potential.DeltaPotential1D(d)
e0 = v.get_eigenenergy()
tmax = -2 * np.pi / e0
dt = tmax / 500
s = solver.CrankNicolsonSolver(20 / d, 0.1 / d, dt, v)
psi0 = v.get_eigenfunction()
psi0_value = psi0(s.x)
np.testing.assert_almost_equal(wavefunction.norm(s.x, psi0_value),
1,
decimal=2)
times = [tmax, tmax / 2, tmax / 4]
psi_expected = [
psi0_value, -psi0_value, psi0_value * np.exp(0.5j * np.pi)
]
for t, psi1 in zip(times, psi_expected):
psi = s.execute(t, psi0=psi0)
np.testing.assert_almost_equal(
wavefunction.norm(s.x, psi),
wavefunction.norm(s.x, psi0_value),
decimal=7,
err_msg=f"Norm not conserved for depth {d}")
np.testing.assert_allclose(np.abs(psi),
np.abs(psi0_value),
atol=0.02)
np.testing.assert_allclose(psi, psi1, atol=0.05)
def test_norm_1d():
phases = np.linspace(0, 2 * np.pi, 20)
x = np.linspace(-100, 100, 10001)
for phase in phases:
psi = np.exp(1j * phase) * np.power(2 / np.pi, 0.25) * np.exp(-x**2)
norm = wavefunction.norm(x, psi)
assert np.imag(norm) == 0.0, "Norm is not real"
np.testing.assert_almost_equal(norm, 1.0)
psi *= 4.0
norm = wavefunction.norm(x, psi)
np.testing.assert_almost_equal(norm, 16.0)
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 calculate_eigenstates(x, potential, n=1, energy_guess=0.0):
"""
Calculates `n` eigenstates in the potential V(x).
:param x: array of x coordinates
:param potential: function V(x)
:param n: number of eigenstates to calculate
:param energy_guess: the expected level of energy of the first eigenstate
:return: a tuple of eigenenergies (as a np.array) and eigenfunctions (as a list of np.array-s)
"""
from scipy import sparse
from scipy.sparse.linalg import eigsh
import wavefunction
dim = x.shape[0]
dx = x[1] - x[0]
v_arr = potential(x)
A = sparse.dok_matrix((dim, dim))
for i in range(dim):
A[i, i] = 1.0 / dx**2 + v_arr[i]
A[i, i - 1] = -0.5 / dx**2
for i in range(dim - 1):
A[i, i + 1] = -0.5 / dx**2
A[dim - 1, 0] = -0.5 / dx**2
A = sparse.csc_matrix(A)
w, vectors = eigsh(A, k=n, sigma=energy_guess)
psis0 = [v / _np.sqrt(wavefunction.norm(x, v)) for v in vectors.T]
return w, psis0
def test_quadratic_potential():
frequencies = [0.1, 1.0, 7.5]
x = np.linspace(-50, 50, 40000)
levels = range(20)
for f in frequencies:
v = potential.QuadraticPotential1D(f)
assert (v.get_frequency() == f)
with pytest.raises(Exception):
v.get_number_of_levels()
for l in levels:
e = v.get_eigenenergy(l)
np.testing.assert_almost_equal(e, (2 * l + 1) * f * 0.5)
psi = v.get_eigenfunction(l)
psi_value = psi(x)
np.testing.assert_almost_equal(
wavefunction.norm(x, psi_value),
1.0,
err_msg=f'Norm is not 1 for frequency {f}, level {l}')
psi0_value = v.get_eigenfunction()(x)
psi0_expected = (f / np.pi)**0.25 * np.exp(-0.5 * f * x**2)
np.testing.assert_allclose(psi0_value, psi0_expected)
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_square_potential():
widths = [1.0, 0.5, 2.0]
depths = [1.0, 5.0, 10.0]
x = np.linspace(-150, 150, 3000)
expected_levels = {
(1.0, 1.0): 1,
(0.5, 1.0): 1,
(2.0, 1.0): 2,
(1.0, 5.0): 3,
(0.5, 5.0): 2,
(2.0, 5.0): 5,
(1.0, 10.0): 3,
(0.5, 10.0): 2,
(2.0, 10.0): 6
}
from itertools import product
for V0, a in product(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}"
assert v.get_delta_depth(
) == 0.0, f"Delta depth {v.get_delta_depth()} is non-zero"
max_levels = v.get_number_of_levels()
assert max_levels == expected_levels[
a,
V0], f"Max levels {max_levels}, expected {expected_levels[a, V0]}"
with pytest.raises(ValueError):
v.get_eigenenergy(max_levels)
with pytest.raises(ValueError):
v.get_eigenfunction(max_levels)
assert v.get_potential()(0.0) == -V0
assert v.get_potential()(2 * a) == 0.0
assert v.get_potential()(-2 * a) == 0.0
energies = np.array([v.get_eigenenergy(i) for i in range(max_levels)])
np.testing.assert_equal(
energies,
np.sort(energies),
err_msg=f"Energies aren't sorted for V0={V0}, a={a}")
assert np.all(energies < 0.0), f"Positive energies for V0={V0}, a={a}"
assert np.all(energies > -V0), f"Too low energies for V0={V0}, a={a}"
for i in range(max_levels):
psi = v.get_eigenfunction(i)(x)
np.testing.assert_almost_equal(
wavefunction.norm(x, psi),
1.0,
decimal=4,
err_msg=
f"Eigenfunction {i} norm is incorrect for V0={V0}, a={a}")
def _test_quadratic(solver_class):
frequencies = [0.1, 1.0, 7.5]
levels = range(10)
for f in frequencies:
v = potential.QuadraticPotential1D(f)
for l in levels:
e = v.get_eigenenergy(l)
tmax = 2 * np.pi / e
dt = tmax / 100
s = solver_class(20 / f, 0.05 / f, dt, v)
psi0 = v.get_eigenfunction(l)
psi0_value = psi0(s.x)
np.testing.assert_almost_equal(wavefunction.norm(s.x, psi0_value),
1,
decimal=2)
times = [tmax, tmax / 2, tmax / 4]
psi_expected = [
psi0_value, -psi0_value, psi0_value * np.exp(-0.5j * np.pi)
]
for t, psi1 in zip(times, psi_expected):
psi = s.execute(t, psi0=psi0)
np.testing.assert_almost_equal(
wavefunction.norm(s.x, psi),
wavefunction.norm(s.x, psi0_value),
decimal=7,
err_msg=f"Norm not conserved for frequency {f}, level {l}")
np.testing.assert_allclose(np.abs(psi),
np.abs(psi0_value),
atol=0.02)
np.testing.assert_allclose(psi, psi1, atol=0.05)
def test_momentum_representation_gauss():
x = np.linspace(-100, 100, 10001)
x0s = [0.0, -10.5, 5.0]
sigmas = [1.0, 5.0, 0.1]
for sigma in sigmas:
for x0 in x0s:
psi = __gauss(x, x0, sigma)
np.testing.assert_array_almost_equal(wavefunction.norm(x, psi),
1.0)
p, psi_p = wavefunction.momentum_representation(x, psi)
np.testing.assert_equal(np.sort(p),
p,
err_msg="Array of p is not sorted")
np.testing.assert_array_almost_equal(
wavefunction.norm(p, psi_p),
1.0,
err_msg="The norm is not conserved in momentum representation")
psi_p_expected = __gauss(p, 0.0, 2 / sigma) * np.exp(
-1.0j * p * x0)
np.testing.assert_allclose(psi_p, psi_p_expected, atol=0.02)
def test_coordinate_representation_gauss():
p = np.linspace(-100, 100, 10001)
p0s = [0.0, -10.5, 5.0]
sigmas = [1.0, 5.0, 0.1]
for sigma in sigmas:
for p0 in p0s:
psi_p = __gauss(p, p0, 2 / sigma)
np.testing.assert_array_almost_equal(wavefunction.norm(p, psi_p),
1.0)
x, psi = wavefunction.coordinate_representation(p, psi_p)
np.testing.assert_equal(np.sort(x),
x,
err_msg="Array of x is not sorted")
np.testing.assert_array_almost_equal(
wavefunction.norm(x, psi),
1.0,
err_msg="The norm is not conserved in coordinate representation"
)
psi_expected = __gauss(x, 0.0, sigma) * np.exp(1.0j * p0 * x)
np.testing.assert_allclose(psi, psi_expected, atol=0.02)
def test_norm_multidimensional():
x = np.linspace(0, 10, 200)
y = np.linspace(0, 20, 100)
z = np.linspace(0, 30, 50)
xx, yy, zz = np.meshgrid(x, y, z, indexing='ij')
f = np.ones(xx.shape) / np.sqrt(6000)
norm = wavefunction.norm((xx, yy, zz), f)
assert np.imag(norm) == 0.0, "Norm is not real"
np.testing.assert_almost_equal(norm, 1)
phases = np.linspace(0, 2 * np.pi, 3)
x = np.linspace(-10, 10, 101)
for phase in phases:
xx, yy, zz = np.meshgrid(x, x, x, indexing='ij')
psi = np.exp(1j * phase) * np.power(
2 / np.pi, 0.75) * np.exp(-xx**2 - yy**2 - zz**2)
norm = wavefunction.norm((xx, yy, zz), psi)
assert np.imag(norm) == 0.0, "Norm is not real"
np.testing.assert_almost_equal(norm, 1.0)
psi *= 4.0
norm = wavefunction.norm((xx, yy, zz), psi)
np.testing.assert_almost_equal(norm, 16.0)
def test_delta_potential():
x = np.linspace(-50, 50, 40000)
depths = np.linspace(0.1, 10, 10)
for d in depths:
v = potential.DeltaPotential1D(d)
assert (v.get_delta_depth() == d)
assert (v.get_number_of_levels() == 1)
with pytest.raises(ValueError):
v.get_eigenenergy(random.randint(1, 100))
with pytest.raises(ValueError):
v.get_eigenfunction(random.randint(1, 100))
psi = v.get_eigenfunction()(x)
np.testing.assert_almost_equal(wavefunction.norm(x, psi),
1.0,
decimal=4,
err_msg=f"Norm is not 1 for depth {d}")
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')
