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)