def test_harmonic_oscillator(self, num_electrons, expected_total_energy,
expected_kinetic_energy):
grids = np.linspace(-5, 5, 501)
potential_fn = functools.partial(single_electron.harmonic_oscillator,
k=1)
solver = single_electron.SparseEigenSolver(grids=grids,
potential_fn=potential_fn,
num_electrons=num_electrons)
solver.solve_ground_state()
# Integrate the density over the grid points.
norm = np.sum(solver.density) * solver.dx
np.testing.assert_allclose(norm, num_electrons)
# i-th eigen-energy of the quantum harmonic oscillator is
# 0.5 + (i - 1).
# The total energy for num_electrons states is
# (0.5 + 0.5 + num_electrons - 1) / 2 * num_electrons
# = num_electrons ** 2 / 2
np.testing.assert_allclose(solver.total_energy,
expected_total_energy,
atol=1e-3)
# Kinetic energy should equal to half of the total energy.
# num_electrons ** 2 / 4
np.testing.assert_allclose(solver.kinetic_energy,
expected_kinetic_energy,
atol=1e-3)
def test_additional_levels_negative(self):
with self.assertRaisesRegexp(
ValueError,
'additional_levels is expected to be non-negative, but got -1'):
single_electron.SparseEigenSolver(
grids=np.linspace(-5, 5, 501),
potential_fn=functools.partial(
single_electron.harmonic_oscillator, k=1),
num_electrons=3,
additional_levels=-1)
def test_additional_levels_too_large(self):
with self.assertRaisesRegexp(
ValueError,
r'additional_levels is expected to be smaller than '
r'num_grids - num_electrons \(498\), but got 499'):
single_electron.SparseEigenSolver(
grids=np.linspace(-5, 5, 501),
potential_fn=functools.partial(
single_electron.harmonic_oscillator, k=1),
num_electrons=3,
additional_levels=499)
def test_poschl_teller(
self, num_electrons, lam, expected_energy, atol):
# NOTE(leeley): Default additional_levels cannot converge to the correct
# eigenstate, even for the first eigenstate. I set the additional_levels
# to 20 so it can converge to the correct eigenstate and reach the same
# accuracy of the EigenSolver.
solver = single_electron.SparseEigenSolver(
grids=np.linspace(-10, 10, 1001),
potential_fn=functools.partial(single_electron.poschl_teller, lam=lam),
num_electrons=num_electrons,
additional_levels=20)
solver.solve_ground_state()
self.assertAlmostEqual(np.sum(solver.density) * solver.dx, num_electrons)
np.testing.assert_allclose(solver.total_energy, expected_energy, atol=atol)
def test_gaussian_dips(self):
grids = np.linspace(-5, 5, 501)
coeff = np.array([1., 1.])
sigma = np.array([1., 1.])
mu = np.array([-2., 2.])
potential_fn = functools.partial(single_electron.gaussian_dips,
coeff=coeff,
sigma=sigma,
mu=mu)
solver = single_electron.SparseEigenSolver(grids, potential_fn)
solver.solve_ground_state()
# Integrate the density over the grid points.
norm = np.sum(solver.density) * solver.dx
# Exact kinetic energy from von Weizsacker functional.
t_vw = single_electron.vw_grid(solver.density, solver.dx)
np.testing.assert_allclose(norm, 1.)
# Kinetic energy should equal to the exact solution from
# von Weizsacker kinetic energy functional.
np.testing.assert_allclose(solver.kinetic_energy, t_vw, atol=1e-4)