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])