def test_ferminet(self, hidden_units, after_det):
"""Check that FermiNet is actually antisymmetric."""
atoms = [
system.Atom(symbol='H', coords=(0, 0, -1.0)),
system.Atom(symbol='H', coords=(0, 0, 1.0)),
]
ne = (3, 2)
x1 = tf.random_normal([3, 3 * sum(ne)])
xs = tf.split(x1, sum(ne), axis=1)
# swap indices to test antisymmetry
x2 = tf.concat([xs[1], xs[0]] + xs[2:], axis=1)
x3 = tf.concat(
xs[:ne[0]] + [xs[ne[0] + 1], xs[ne[0]]] + xs[ne[0] + 2:], axis=1)
ferminet = networks.FermiNet(
atoms=atoms,
nelectrons=ne,
slater_dets=4,
hidden_units=hidden_units,
after_det=after_det)
y1 = ferminet(x1)
y2 = ferminet(x2)
y3 = ferminet(x3)
with tf.train.MonitoredSession() as session:
out1, out2, out3 = session.run([y1, y2, y3])
np.testing.assert_allclose(out1, -out2, rtol=4.e-5, atol=1.e-6)
np.testing.assert_allclose(out1, -out3, rtol=4.e-5, atol=1.e-6)
def test_ferminet_pretrain(self, hidden_units, after_det):
"""Check that FermiNet pretraining runs."""
atoms = [
system.Atom(symbol='Li', coords=(0, 0, -1.0)),
system.Atom(symbol='Li', coords=(0, 0, 1.0)),
]
ne = (4, 2)
x = tf.random_normal([1, 10, 3 * sum(ne)])
strategy = tf.distribute.get_strategy()
with strategy.scope():
ferminet = networks.FermiNet(
atoms=atoms,
nelectrons=ne,
slater_dets=4,
hidden_units=hidden_units,
after_det=after_det,
pretrain_iterations=10)
# Test Hartree fock pretraining - no change of position.
hf_approx = scf.Scf(atoms, nelectrons=ne)
hf_approx.run()
pretrain_op_hf = networks.pretrain_hartree_fock(ferminet, x, strategy,
hf_approx)
self.assertEqual(ferminet.pretrain_iterations, 10)
with tf.train.MonitoredSession() as session:
for _ in range(ferminet.pretrain_iterations):
session.run(pretrain_op_hf)
def test_ferminet_size(self, size):
atoms = [
system.Atom(symbol='H', coords=(0, 0, -1.0)),
system.Atom(symbol='H', coords=(0, 0, 1.0)),
]
ne = (size, size)
nout = 1
batch_size = 3
ndet = 4
hidden_units = [[8, 8], [8, 8]]
after_det = [8, 8, nout]
x = tf.random_normal([batch_size, 3 * sum(ne)])
ferminet = networks.FermiNet(
atoms=atoms,
nelectrons=ne,
slater_dets=ndet,
hidden_units=hidden_units,
after_det=after_det)
y = ferminet(x)
for i in range(len(ne)):
self.assertEqual(ferminet._dets[i].shape.as_list(), [batch_size, ndet])
self.assertEqual(ferminet._orbitals[i].shape.as_list(),
[batch_size, ndet, size, size])
self.assertEqual(y.shape.as_list(), [batch_size, nout])
def create_o2_hf(bond_length):
molecule = [
system.Atom('O', (0, 0, 0)),
system.Atom('O', (0, 0, bond_length))
]
spin = 2
oxygen_atomic_number = 8
nelectrons = [oxygen_atomic_number + spin, oxygen_atomic_number - spin]
hf = scf.Scf(molecule=molecule, nelectrons=nelectrons, restricted=False)
return hf
def setUp(self):
super(HamiltonianMoleculeTest, self).setUp()
self.rh1 = np.array([0, 0, 0], dtype=np.float32)
self.rh2 = np.array([0, 0, 200], dtype=np.float32)
self.rhh = self.rh2 - self.rh1
self.dhh = 1.0 / np.sqrt(np.dot(self.rhh, self.rhh))
self.xs = (2 * np.random.random((6, 3)) - 1).astype(np.float32)
self.xs = [self.xs + self.rh1, self.xs + self.rh2]
self.atoms = [
system.Atom(symbol='H', coords=self.rh1),
system.Atom(symbol='H', coords=self.rh2)
]
def test_tf_eval_hf(self):
# Check we get consistent answers between multiple calls to eval_mos and
# a single tf_eval_hf call.
molecule = [system.Atom('O', (0, 0, 0))]
nelectrons = (5, 3)
hf = scf.Scf(molecule=molecule,
nelectrons=nelectrons,
restricted=False)
hf.run()
batch = 100
xs = [np.random.randn(batch, 3) for _ in range(sum(nelectrons))]
mos = []
for i, x in enumerate(xs):
ispin = 0 if i < nelectrons[0] else 1
orbitals = hf.eval_mos(x)[ispin]
# Select occupied orbitals via Aufbau.
mos.append(orbitals[:, :nelectrons[ispin]])
np_mos = (np.stack(mos[:nelectrons[0]],
axis=1), np.stack(mos[nelectrons[0]:], axis=1))
tf_xs = tf.constant(np.stack(xs, axis=1))
tf_mos = hf.tf_eval_hf(tf_xs, deriv=True)
with tf.train.MonitoredSession() as session:
tf_mos_ = session.run(tf_mos)
for i, (np_mos_mat, tf_mos_mat) in enumerate(zip(np_mos, tf_mos_)):
self.assertEqual(np_mos_mat.shape, tf_mos_mat.shape)
self.assertEqual(np_mos_mat.shape,
(batch, nelectrons[i], nelectrons[i]))
np.testing.assert_allclose(np_mos_mat, tf_mos_mat)
def test_tf_eval_slog_wavefuncs(self):
# Check TensorFlow evaluation runs and gives correct shapes.
molecule = [system.Atom('O', (0, 0, 0))]
nelectrons = (5, 3)
total_electrons = sum(nelectrons)
num_spatial_dim = 3
hf = scf.Scf(molecule=molecule,
nelectrons=nelectrons,
restricted=False)
hf.run()
batch = 100
rng = np.random.RandomState(1)
flat_positions_np = rng.randn(batch, total_electrons * num_spatial_dim)
flat_positions_tf = tf.constant(flat_positions_np)
for method in [
hf.tf_eval_slog_slater_determinant,
hf.tf_eval_slog_hartree_product
]:
slog_wavefunc, signs = method(flat_positions_tf)
with tf.train.MonitoredSession() as session:
slog_wavefunc_, signs_ = session.run([slog_wavefunc, signs])
self.assertEqual(slog_wavefunc_.shape, (batch, 1))
self.assertEqual(signs_.shape, (batch, 1))
hartree_product = hf.tf_eval_hartree_product(flat_positions_tf)
with tf.train.MonitoredSession() as session:
hartree_product_ = session.run(hartree_product)
np.testing.assert_allclose(hartree_product_,
np.exp(slog_wavefunc_) * signs_)
def test_hydrogen_atom(self, dtype):
atoms = [system.Atom(symbol='H', coords=(0, 0, 0))]
def net(x):
psi = _hydrogen(x)
return (tf.log(tf.abs(psi)), tf.sign(psi))
e = _run_mcmc(atoms, 1, net, dtype=dtype)
np.testing.assert_allclose(e, -np.ones_like(e) / 2)
def test_helium_atom(self):
atoms = [system.Atom(symbol='He', coords=(0, 0, 0))]
def net(x):
psi = _helium(x)
return (tf.log(tf.abs(psi)), tf.sign(psi))
e = _run_mcmc(atoms, 2, net, steps=500)
np.testing.assert_allclose(e[100:].mean(), -2.901188, atol=5.e-3)
def test_r12_features_atom1(self):
atoms = [system.Atom(symbol='H', coords=(0, 0, 0))]
one = np.ones((1, 3))
xs = tf.constant(one, dtype=tf.float32)
xs12 = hamiltonian.r12_features(xs, atoms, 1, flatten=True)
with tf.train.MonitoredSession() as session:
xs12_ = session.run(xs12)
# Should add |x|.
x_test = np.concatenate([[[np.linalg.norm(one[0])]], one], axis=1)
np.testing.assert_allclose(xs12_, x_test)
def test_hydrogen_atom(self):
xs = tf.random_uniform((6, 3), minval=-1.0, maxval=1.0)
atoms = [system.Atom(symbol='H', coords=(0, 0, 0))]
kin, pot = hamiltonian.operators(atoms, 1, 0.0)
logpsi = tf.log(tf.abs(_hydrogen(xs)))
e_loc = kin(logpsi, xs) + pot(xs)
with tf.train.MonitoredSession() as session:
e_loc_ = session.run(e_loc)
# Wavefunction is the ground state eigenfunction. Hence H\Psi = -1/2 \Psi.
np.testing.assert_allclose(e_loc_, -0.5, rtol=1.e-6, atol=1.e-7)
def test_ferminet_mask(self):
"""Check that FermiNet with a decaying mask on the output works."""
atoms = [
system.Atom(symbol='H', coords=(0, 0, -1.0)),
system.Atom(symbol='H', coords=(0, 0, 1.0)),
]
ne = (3, 2)
hidden_units = [[8, 8], [8, 8]]
after_det = [8, 8, 2]
x = tf.random_normal([3, 3 * sum(ne)])
ferminet = networks.FermiNet(
atoms=atoms,
nelectrons=ne,
slater_dets=4,
hidden_units=hidden_units,
after_det=after_det,
envelope=True)
y = ferminet(x)
with tf.train.MonitoredSession() as session:
session.run(y)
def test_r12_features_molecule(self):
atoms = [
system.Atom(symbol='H', coords=(0, 0, 0)),
system.Atom(symbol='H', coords=(1, 0, 0))
]
one = np.concatenate([np.ones((1, 3)), 0.5 * np.ones((1, 3))], axis=1)
xs = tf.constant(one, dtype=tf.float32)
xs12 = hamiltonian.r12_features(xs, atoms, 2, flatten=True)
with tf.train.MonitoredSession() as session:
xs12_ = session.run(xs12)
# Should add |x_11|, |x_21|, |x_12|, |x_22|, |x1-x2|
norm = np.linalg.norm
x_test = np.concatenate([[[
norm(one[0, :3]),
norm(one[0, :3] - atoms[1].coords),
norm(one[0, 3:]),
norm(one[0, 3:] - atoms[1].coords),
norm(one[0, :3] - one[0, 3:]),
]], one],
axis=1)
np.testing.assert_allclose(xs12_, x_test, rtol=1e-6)
def test_assign_electrons(self, molecule, electrons, electrons_per_atom):
molecule = [system.Atom(v[0], v[1:]) for v in molecule]
e_means = train.assign_electrons(molecule, electrons)
expected_e_means = self._expected_result(molecule, electrons_per_atom)
np.testing.assert_allclose(e_means, expected_e_means)
e_means = train.assign_electrons(molecule[::-1], electrons)
if any(atom.symbol != molecule[0].symbol for atom in molecule):
expected_e_means = self._expected_result(molecule[::-1],
electrons_per_atom[::-1])
else:
expected_e_means = self._expected_result(molecule[::-1],
electrons_per_atom)
np.testing.assert_allclose(e_means, expected_e_means)
def test_helium_atom(self):
x1 = np.linspace(-1, 1, 6, dtype=np.float32)
x2 = np.zeros(6, dtype=np.float32) + 0.25
x12 = np.array([x1] + [x2 for _ in range(5)]).T
atoms = [system.Atom(symbol='He', coords=(0, 0, 0))]
hpsi_test = np.array([
-0.25170374, -0.43359983, -0.61270618, -1.56245542, -0.39977553,
-0.2300467
])
log_helium = _wfn_to_log(_helium)
h = hamiltonian.exact_hamiltonian(atoms, 2)
xs = tf.constant(x12)
_, hpsi = h(log_helium, xs)
with tf.train.MonitoredSession() as session:
hpsi_ = session.run(hpsi)
xs = tf.reverse(xs, axis=[1])
_, hpsi2 = h(log_helium, xs)
with tf.train.MonitoredSession() as session:
hpsi2_ = session.run(hpsi2)
np.testing.assert_allclose(hpsi_[:, 0], hpsi_test, rtol=1.e-6)
np.testing.assert_allclose(hpsi_, hpsi2_, rtol=1.e-6)
class ScfTest(tf.test.TestCase, parameterized.TestCase):
def setUp(self):
super(ScfTest, self).setUp()
# disable use of temp directory in pyscf.
# Test calculations are small enough to fit in RAM and we don't need
# checkpoint files.
pyscf.lib.param.TMPDIR = None
@parameterized.parameters(
{
'molecule': [system.Atom('He', (0, 0, 0))],
'nelectrons': (1, 1)
},
{
'molecule': [system.Atom('N', (0, 0, 0))],
'nelectrons': (5, 2)
},
{
'molecule': [system.Atom('N', (0, 0, 0))],
'nelectrons': (5, 3)
},
{
'molecule': [system.Atom('N', (0, 0, 0))],
'nelectrons': (4, 2)
},
{
'molecule': [system.Atom('O', (0, 0, 0))],
'nelectrons': (5, 3),
'restricted': False,
},
{
'molecule':
[system.Atom('N', (0, 0, 0)),
system.Atom('N', (0, 0, 1.4))],
'nelectrons': (7, 7)
},
{
'molecule':
[system.Atom('O', (0, 0, 0)),
system.Atom('O', (0, 0, 1.4))],
'nelectrons': (9, 7),
'restricted':
False,
},
)
def test_scf_interface(self,
molecule: List[system.Atom],
nelectrons: Tuple[int, int],
restricted: bool = True):
"""Tests SCF interface to a pyscf calculation.
pyscf has its own tests so only check that we can run calculations over
atoms and simple diatomics using the interface in ferminet.scf.
Args:
molecule: List of system.Atom objects giving atoms in the molecule.
nelectrons: Tuple containing number of alpha and beta electrons.
restricted: If true, run a restricted Hartree-Fock calculation, otherwise
run an unrestricted Hartree-Fock calculation.
"""
npts = 100
xs = np.random.randn(npts, 3)
hf = scf.Scf(molecule=molecule,
nelectrons=nelectrons,
restricted=restricted)
hf.run()
mo_vals = hf.eval_mos(xs)
self.assertLen(mo_vals,
2) # alpha-spin orbitals and beta-spin orbitals.
for spin_mo_vals in mo_vals:
# Evalute npts points on M orbitals/functions - (npts, M) array.
self.assertEqual(spin_mo_vals.shape, (npts, hf._mol.nao_nr()))
@parameterized.parameters(np.float32, np.float64)
def test_tf_eval_mos(self, dtype):
"""Tests tensorflow op for evaluating Hartree-Fock orbitals.
Args:
dtype: numpy type to use for position vectors.
"""
xs = np.random.randn(100, 3).astype(dtype)
hf = create_o2_hf(1.4)
hf.run()
# Evaluate MOs at positions xs directly in pyscf.
# hf.eval_mos returns np.float64 arrays always. Cast to dtype for
# like-for-like comparison.
mo_vals = [
mo_val.astype(dtype) for mo_val in hf.eval_mos(xs, deriv=True)
]
# Evaluate MOs as a tensorflow op.
tf_xs = tf.constant(xs)
tf_mo_vals = hf.tf_eval_mos(tf_xs, deriv=True)
tf_dmo_vals = tf.gradients(tf.square(tf_mo_vals), tf_xs)[0]
tf_dmo_vals_shape = tf.shape(tf_dmo_vals)
with tf.train.MonitoredSession() as session:
tf_mo_vals_, tf_dmo_vals_ = session.run([tf_mo_vals, tf_dmo_vals])
tf_dmo_vals_shape_ = session.run(tf_dmo_vals_shape)
tf_mo_vals_ = np.split(tf_mo_vals_, 2, axis=-1)
for np_val, tf_val in zip(mo_vals, tf_mo_vals_):
self.assertEqual(np_val.dtype, tf_val.dtype)
np.testing.assert_allclose(np_val[0], tf_val)
np.testing.assert_array_equal(tf_dmo_vals_shape_, xs.shape)
np.testing.assert_array_equal(tf_dmo_vals_shape_, tf_dmo_vals_.shape)
# Compare analytic derivative of orbital^2 with tensorflow backprop.
# Sum over the orbital index and spin channel because of definition of
# tensorflow gradients.
np_deriv = sum(
np.sum(2 * np_val[1:] * np.expand_dims(np_val[0], 0), axis=-1)
for np_val in mo_vals)
# Tensorflow returns (N,3) but pyscf returns (3, N, M).
np_deriv = np.transpose(np_deriv)
self.assertEqual(np_deriv.dtype, tf_dmo_vals_.dtype)
if dtype == np.float32:
atol = 1e-6
rtol = 1e-6
else:
atol = 0
rtol = 1.e-7
np.testing.assert_allclose(np_deriv,
tf_dmo_vals_,
atol=atol,
rtol=rtol)
def test_tf_eval_mos_deriv(self):
hf = create_o2_hf(1.4)
hf.run()
xs = tf.random_normal((100, 3))
tf_mos_derivs = hf.tf_eval_mos(xs, deriv=True)
tf_mos = hf.tf_eval_mos(xs, deriv=False)
with tf.train.MonitoredSession() as session:
tf_mos_, tf_mos_derivs_ = session.run([tf_mos, tf_mos_derivs])
np.testing.assert_allclose(tf_mos_, tf_mos_derivs_)
def test_tf_eval_hf(self):
# Check we get consistent answers between multiple calls to eval_mos and
# a single tf_eval_hf call.
molecule = [system.Atom('O', (0, 0, 0))]
nelectrons = (5, 3)
hf = scf.Scf(molecule=molecule,
nelectrons=nelectrons,
restricted=False)
hf.run()
batch = 100
xs = [np.random.randn(batch, 3) for _ in range(sum(nelectrons))]
mos = []
for i, x in enumerate(xs):
ispin = 0 if i < nelectrons[0] else 1
orbitals = hf.eval_mos(x)[ispin]
# Select occupied orbitals via Aufbau.
mos.append(orbitals[:, :nelectrons[ispin]])
np_mos = (np.stack(mos[:nelectrons[0]],
axis=1), np.stack(mos[nelectrons[0]:], axis=1))
tf_xs = tf.constant(np.stack(xs, axis=1))
tf_mos = hf.tf_eval_hf(tf_xs, deriv=True)
with tf.train.MonitoredSession() as session:
tf_mos_ = session.run(tf_mos)
for i, (np_mos_mat, tf_mos_mat) in enumerate(zip(np_mos, tf_mos_)):
self.assertEqual(np_mos_mat.shape, tf_mos_mat.shape)
self.assertEqual(np_mos_mat.shape,
(batch, nelectrons[i], nelectrons[i]))
np.testing.assert_allclose(np_mos_mat, tf_mos_mat)
def test_tf_eval_slog_wavefuncs(self):
# Check TensorFlow evaluation runs and gives correct shapes.
molecule = [system.Atom('O', (0, 0, 0))]
nelectrons = (5, 3)
total_electrons = sum(nelectrons)
num_spatial_dim = 3
hf = scf.Scf(molecule=molecule,
nelectrons=nelectrons,
restricted=False)
hf.run()
batch = 100
rng = np.random.RandomState(1)
flat_positions_np = rng.randn(batch, total_electrons * num_spatial_dim)
flat_positions_tf = tf.constant(flat_positions_np)
for method in [
hf.tf_eval_slog_slater_determinant,
hf.tf_eval_slog_hartree_product
]:
slog_wavefunc, signs = method(flat_positions_tf)
with tf.train.MonitoredSession() as session:
slog_wavefunc_, signs_ = session.run([slog_wavefunc, signs])
self.assertEqual(slog_wavefunc_.shape, (batch, 1))
self.assertEqual(signs_.shape, (batch, 1))
hartree_product = hf.tf_eval_hartree_product(flat_positions_tf)
with tf.train.MonitoredSession() as session:
hartree_product_ = session.run(hartree_product)
np.testing.assert_allclose(hartree_product_,
np.exp(slog_wavefunc_) * signs_)
def test_atom_units(self):
system.Atom(symbol='H', coords=[1, 2, 3], units='bohr')
system.Atom(symbol='H', coords=[1, 2, 3], units='angstrom')
with self.assertRaises(ValueError):
system.Atom(symbol='H', coords=[1, 2, 3], units='dummy')
def test_atom_coords(self):
xs = np.random.uniform(size=3)
atom = system.Atom(symbol='H', coords=xs, units='angstrom')
np.testing.assert_allclose(atom.coords / xs, [units.BOHR_ANGSTROM] * 3)
np.testing.assert_allclose(atom.coords_angstrom, xs)
class ScfTest(parameterized.TestCase):
def setUp(self):
super(ScfTest, self).setUp()
# disable use of temp directory in pyscf.
# Test calculations are small enough to fit in RAM and we don't need
# checkpoint files.
pyscf.lib.param.TMPDIR = None
@parameterized.parameters(
{
'molecule': [system.Atom('He', (0, 0, 0))],
'nelectrons': (1, 1)
},
{
'molecule': [system.Atom('N', (0, 0, 0))],
'nelectrons': (5, 2)
},
{
'molecule': [system.Atom('N', (0, 0, 0))],
'nelectrons': (5, 3)
},
{
'molecule': [system.Atom('N', (0, 0, 0))],
'nelectrons': (4, 2)
},
{
'molecule': [system.Atom('O', (0, 0, 0))],
'nelectrons': (5, 3),
'restricted': False,
},
{
'molecule':
[system.Atom('N', (0, 0, 0)),
system.Atom('N', (0, 0, 1.4))],
'nelectrons': (7, 7)
},
{
'molecule':
[system.Atom('O', (0, 0, 0)),
system.Atom('O', (0, 0, 1.4))],
'nelectrons': (9, 7),
'restricted':
False,
},
)
def test_scf_interface(self,
molecule: List[system.Atom],
nelectrons: Tuple[int, int],
restricted: bool = True):
"""Tests SCF interface to a pyscf calculation.
pyscf has its own tests so only check that we can run calculations over
atoms and simple diatomics using the interface in ferminet.scf.
Args:
molecule: List of system.Atom objects giving atoms in the molecule.
nelectrons: Tuple containing number of alpha and beta electrons.
restricted: If true, run a restricted Hartree-Fock calculation, otherwise
run an unrestricted Hartree-Fock calculation.
"""
npts = 100
xs = np.random.randn(npts, 3)
hf = scf.Scf(molecule=molecule,
nelectrons=nelectrons,
restricted=restricted)
hf.run()
mo_vals = hf.eval_mos(xs)
self.assertLen(mo_vals,
2) # alpha-spin orbitals and beta-spin orbitals.
for spin_mo_vals in mo_vals:
# Evalute npts points on M orbitals/functions - (npts, M) array.
self.assertEqual(spin_mo_vals.shape, (npts, hf._mol.nao_nr()))
