def test_get_xc_potential(self):
grids = jnp.linspace(-5, 5, 10001)
# We use the form of 3d LDA exchange functional as an example. So the
# correlation contribution is 0.
# exchange energy = -0.73855 \int n^(4 / 3) dx
# exchange potential should be -0.73855 * (4 / 3) n^(1 / 3)
# by taking functional derivative on exchange energy.
xc_energy_density_fn = lambda density: -0.73855 * density**(1 / 3)
density = jnp.exp(-(grids - 1)**2)
np.testing.assert_allclose(
scf.get_xc_potential(density,
xc_energy_density_fn=xc_energy_density_fn,
grids=grids),
-0.73855 * (4 / 3) * density**(1 / 3))
def _kohn_sham_iteration(density, external_potential, grids, num_electrons,
xc_energy_density_fn, interaction_fn,
enforce_reflection_symmetry):
"""One iteration of Kohn-Sham calculation."""
# NOTE(leeley): Since num_electrons in KohnShamState need to specify as
# static argument in jit function, this function can not directly take
# KohnShamState as input arguments. The related attributes in KohnShamState
# are used as input arguments for this helper function.
if enforce_reflection_symmetry:
xc_energy_density_fn = _flip_and_average_on_center_fn(
xc_energy_density_fn)
hartree_potential = scf.get_hartree_potential(
density=density, grids=grids, interaction_fn=interaction_fn)
xc_potential = scf.get_xc_potential(
density=density,
xc_energy_density_fn=xc_energy_density_fn,
grids=grids)
ks_potential = hartree_potential + xc_potential + external_potential
xc_energy_density = xc_energy_density_fn(density)
# Solve Kohn-Sham equation.
density, total_eigen_energies, gap = scf.solve_noninteracting_system(
external_potential=ks_potential,
num_electrons=num_electrons,
grids=grids)
total_energy = (
# kinetic energy = total_eigen_energies - external_potential_energy
total_eigen_energies - scf.get_external_potential_energy(
external_potential=ks_potential, density=density, grids=grids)
# Hartree energy
+ scf.get_hartree_energy(
density=density, grids=grids, interaction_fn=interaction_fn)
# xc energy
+ scf.get_xc_energy(density=density,
xc_energy_density_fn=xc_energy_density_fn,
grids=grids)
# external energy
+ scf.get_external_potential_energy(
external_potential=external_potential,
density=density,
grids=grids))
if enforce_reflection_symmetry:
density = _flip_and_average_on_center(density)
return (density, total_energy, hartree_potential, xc_potential,
xc_energy_density, gap)
def test_get_xc_potential_hartree(self):
grids = jnp.linspace(-5, 5, 10001)
density = utils.gaussian(grids=grids, center=1., sigma=1.)
def half_hartree_potential(density):
return 0.5 * scf.get_hartree_potential(
density=density,
grids=grids,
interaction_fn=utils.exponential_coulomb)
np.testing.assert_allclose(
scf.get_xc_potential(
density=density,
xc_energy_density_fn=half_hartree_potential,
grids=grids),
scf.get_hartree_potential(
density, grids=grids, interaction_fn=utils.exponential_coulomb))