def test_kohn_sham_neural_xc_density_mse_converge_tolerance(
self, density_mse_converge_tolerance, expected_converged):
init_fn, xc_energy_density_fn = neural_xc.local_density_approximation(
stax.serial(stax.Dense(8), stax.Elu, stax.Dense(1)))
params_init = init_fn(rng=random.PRNGKey(0))
states = jit_scf.kohn_sham(
locations=self.locations,
nuclear_charges=self.nuclear_charges,
num_electrons=self.num_electrons,
num_iterations=3,
grids=self.grids,
xc_energy_density_fn=tree_util.Partial(xc_energy_density_fn,
params=params_init),
interaction_fn=utils.exponential_coulomb,
initial_density=self.num_electrons *
utils.gaussian(grids=self.grids, center=0., sigma=0.5),
density_mse_converge_tolerance=density_mse_converge_tolerance)
np.testing.assert_array_equal(states.converged, expected_converged)
for single_state in scf.state_iterator(states):
self._test_state(
single_state,
self._create_testing_external_potential(
utils.exponential_coulomb))
def test_wrap_network_with_self_interaction_layer_large_num_electrons(
self):
grids = jnp.linspace(-5, 5, 9, dtype=jnp.float32)
density = 100. * utils.gaussian(grids=grids, center=1.,
sigma=1.).astype(jnp.float32)
reshaped_density = density[jnp.newaxis, :, jnp.newaxis]
inner_network_init_fn, inner_network_apply_fn = neural_xc.build_unet(
num_filters_list=[2, 4], core_num_filters=4, activation='swish')
init_fn, apply_fn = neural_xc.wrap_network_with_self_interaction_layer(
network=(inner_network_init_fn, inner_network_apply_fn),
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))
self.assertEqual(
apply_fn(init_params, reshaped_density).shape, (1, 9, 1))
np.testing.assert_allclose(
apply_fn(init_params, reshaped_density),
inner_network_apply_fn(
# The initial parameters of the inner network.
init_params[1][1],
reshaped_density))
def test_self_interaction_weight(self, density_integral, expected_weight):
grids = jnp.linspace(-5, 5, 11)
self.assertAlmostEqual(
neural_xc.self_interaction_weight(
reshaped_density=density_integral *
utils.gaussian(grids=grids, center=1.,
sigma=1.)[jnp.newaxis, :, jnp.newaxis],
dx=utils.get_dx(grids),
width=1.), expected_weight)
def test_kohn_sham_iteration_neural_xc_density_loss_gradient_symmetry(
self):
# The network only has one layer.
# The initial params contains weights with shape (1, 1) and bias (1,).
init_fn, xc_energy_density_fn = neural_xc.local_density_approximation(
stax.serial(stax.Dense(1)))
init_params = init_fn(rng=random.PRNGKey(0))
initial_state = self._create_testing_initial_state(
utils.exponential_coulomb)
target_density = (
utils.gaussian(grids=self.grids, center=-0.5, sigma=1.) +
utils.gaussian(grids=self.grids, center=0.5, sigma=1.))
spec, flatten_init_params = np_utils.flatten(init_params)
def loss(flatten_params, initial_state, target_density):
state = scf.kohn_sham_iteration(
state=initial_state,
num_electrons=self.num_electrons,
xc_energy_density_fn=tree_util.Partial(
xc_energy_density_fn,
params=np_utils.unflatten(spec, flatten_params)),
interaction_fn=utils.exponential_coulomb,
enforce_reflection_symmetry=True)
return jnp.sum(jnp.abs(state.density -
target_density)) * utils.get_dx(self.grids)
grad_fn = jax.grad(loss)
params_grad = grad_fn(flatten_init_params,
initial_state=initial_state,
target_density=target_density)
# Check gradient values.
np.testing.assert_allclose(params_grad, [0.2013181, 0.], atol=1e-7)
# Check whether the gradient values match the numerical gradient.
np.testing.assert_allclose(optimize.approx_fprime(
xk=flatten_init_params,
f=functools.partial(loss,
initial_state=initial_state,
target_density=target_density),
epsilon=1e-9),
params_grad,
atol=1e-4)
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 _create_testing_initial_state(self, interaction_fn):
locations = jnp.array([-0.5, 0.5])
nuclear_charges = jnp.array([1, 1])
return scf.KohnShamState(
density=self.num_electrons *
utils.gaussian(grids=self.grids, center=0., sigma=1.),
# Set initial energy as inf, the actual value is not used in Kohn-Sham
# calculation.
total_energy=jnp.inf,
locations=locations,
nuclear_charges=nuclear_charges,
external_potential=utils.get_atomic_chain_potential(
grids=self.grids,
locations=locations,
nuclear_charges=nuclear_charges,
interaction_fn=interaction_fn),
grids=self.grids,
num_electrons=self.num_electrons)
def test_get_hartree_energy(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_energy = 0.
for x_0, n_0 in zip(grids, density):
for x_1, n_1 in zip(grids, density):
expected_hartree_energy += 0.5 * n_0 * n_1 * interaction_fn(
x_0 - x_1) * dx**2
self.assertAlmostEqual(
float(
scf.get_hartree_energy(density=density,
grids=grids,
interaction_fn=interaction_fn)),
float(expected_hartree_energy))
def test_self_interaction_layer_large_num_electrons(self):
grids = jnp.linspace(-5, 5, 11)
density = 100. * utils.gaussian(grids=grids, center=1., sigma=1.)
reshaped_density = density[jnp.newaxis, :, jnp.newaxis]
features = np.random.rand(*reshaped_density.shape)
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 output is completely the features (second input).
apply_fn(init_params, (reshaped_density, features)),
features)
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, dtype=jnp.float32)
density = utils.gaussian(grids=grids, center=1.,
sigma=1.).astype(jnp.float32)
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 setUp(self):
super(GlobalFunctionalTest, self).setUp()
self.grids = jnp.linspace(-5, 5, 17)
self.density = utils.gaussian(grids=self.grids, center=1., sigma=1.)
def test_gaussian(self, center, sigma, expected_max_value):
gaussian = utils.gaussian(
grids=jnp.linspace(-10, 10, 201), center=center, sigma=sigma)
self.assertAlmostEqual(float(jnp.sum(gaussian) * 0.1), 1, places=5)
self.assertAlmostEqual(float(jnp.amax(gaussian)), expected_max_value)
