def apply_fn(params, inputs, **kwargs): # pylint: disable=missing-docstring
del kwargs
width, = params
reshaped_density, features = inputs
beta = self_interaction_weight(
reshaped_density=reshaped_density, dx=dx, width=width)
hartree = -0.5 * scf.get_hartree_potential(
density=reshaped_density.reshape(-1),
grids=grids,
interaction_fn=interaction_fn).reshape(reshaped_density.shape)
return hartree * beta + features * (1 - beta)
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))
def test_get_hartree_potential(self, interaction_fn):
grids = jnp.linspace(-5, 5, 11)
dx = utils.get_dx(grids)
density = utils.gaussian(grids=grids, center=1., sigma=1.)
# Compute the expected Hartree energy by nested for loops.
expected_hartree_potential = np.zeros_like(grids)
for i, x_0 in enumerate(grids):
for x_1, n_1 in zip(grids, density):
expected_hartree_potential[i] += np.sum(
n_1 * interaction_fn(x_0 - x_1)) * dx
np.testing.assert_allclose(
scf.get_hartree_potential(
density=density, grids=grids, interaction_fn=interaction_fn),
expected_hartree_potential)
def test_self_interaction_layer_one_electron(self):
grids = jnp.linspace(-5, 5, 11)
density = utils.gaussian(grids=grids, center=1., sigma=1.)
reshaped_density = density[jnp.newaxis, :, jnp.newaxis]
init_fn, apply_fn = neural_xc.self_interaction_layer(
grids=grids, interaction_fn=utils.exponential_coulomb)
output_shape, init_params = init_fn(
random.PRNGKey(0), input_shape=((-1, 11, 1), (-1, 11, 1)))
self.assertEqual(output_shape, (-1, 11, 1))
self.assertAlmostEqual(init_params, (1.,))
np.testing.assert_allclose(
# The features (second input) is not used for one electron.
apply_fn(
init_params, (reshaped_density, jnp.ones_like(reshaped_density))),
-0.5 * scf.get_hartree_potential(
density=density,
grids=grids,
interaction_fn=utils.exponential_coulomb)[
jnp.newaxis, :, jnp.newaxis])
def test_wrap_network_with_self_interaction_layer_one_electron(self):
grids = jnp.linspace(-5, 5, 9)
density = utils.gaussian(grids=grids, center=1., sigma=1.)
reshaped_density = density[jnp.newaxis, :, jnp.newaxis]
init_fn, apply_fn = neural_xc.wrap_network_with_self_interaction_layer(
network=neural_xc.build_unet(num_filters_list=[2, 4],
core_num_filters=4,
activation='swish'),
grids=grids,
interaction_fn=utils.exponential_coulomb)
output_shape, init_params = init_fn(random.PRNGKey(0),
input_shape=((-1, 9, 1)))
self.assertEqual(output_shape, (-1, 9, 1))
np.testing.assert_allclose(
apply_fn(init_params, reshaped_density),
-0.5 * scf.get_hartree_potential(
density=density,
grids=grids,
interaction_fn=utils.exponential_coulomb)[jnp.newaxis, :,
jnp.newaxis])
def half_hartree_potential(density):
return 0.5 * scf.get_hartree_potential(
density=density,
grids=grids,
interaction_fn=utils.exponential_coulomb)
