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