def generate_reference_ensemble():
intg.reset(sess)
rs = ReservoirSampler(reference_generator(), reservoir_size)
rs.sample_all()
ensemble = []
for x, t in rs.R:
ensemble.append(x)
stacked_ensemble = np.stack(ensemble, axis=0)
ensemble_ph = tf.placeholder(shape=(reservoir_size, num_atoms, 3),
dtype=np.float64)
ref_d2ij_op = observable.sorted_squared_distances(ensemble_ph)
return ref_d2ij_op, stacked_ensemble, ensemble_ph
def test_dense_sorted_dij(self):
"""
Testing sorted distance matrix observables implemented correctly
"""
x0 = np.array([
[-0.0070, -0.0100, 0.0000],
[-0.1604, 0.4921, 0.0000],
[0.5175, 0.0128, 0.0000],
],
dtype=np.float64)
all_xs = []
n_confs = 5
for p in range(n_confs):
rand_dx = np.random.rand(3, 3) / 10
all_xs.append(x0 + rand_dx)
stacked_dijs = []
for x in all_xs:
# loop over the atoms
dij = np.zeros(shape=(3, 3), dtype=np.float64)
for a_idx, a in enumerate(x):
for b_idx, b in enumerate(x):
dij[a_idx][b_idx] = np.sum(np.power(a - b, 2), axis=-1)
stacked_dijs.append(dij)
stacked_dijs = np.stack(stacked_dijs, axis=-1)
reference_sorted_dijs = np.sort(stacked_dijs, axis=-1)
x_ph = tf.placeholder(shape=(n_confs, 3, 3), dtype=np.float64)
test_sorted_dijs_op = observable.sorted_squared_distances(x_ph)
dOdx_op = tf.gradients(test_sorted_dijs_op, x_ph)
sess = tf.Session()
test_sorted_dijs, test_grads = sess.run(
[test_sorted_dijs_op, dOdx_op],
feed_dict={x_ph: np.array(all_xs, dtype=np.float64)})
# test that distances are equivalent
np.testing.assert_almost_equal(reference_sorted_dijs,
test_sorted_dijs,
decimal=15)
assert test_grads is not None
assert not np.any(np.isnan(test_grads))
def test_optimize_ensemble(self):
"""
Testing that we can optimize ensembles with different integrator and ff settings relative to a canonical one.
"""
# 1. Generate a water ensemble by doing 5000 steps of MD starting from an idealized geometry.
# 2. Reservoir sample along the trajectory.
# 3. Generate a loss function used sum of square distances.
# 4. Compute parameter derivatives.
x_opt = np.array([
[-0.0070, -0.0100, 0.0000],
[-0.1604, 0.4921, 0.0000],
[0.5175, 0.0128, 0.0000],
],
dtype=np.float64) # idealized geometry
x_opt.setflags(write=False)
masses = np.array([8.0, 1.0, 1.0], dtype=np.float64)
num_atoms = len(masses)
ideal_bond = 0.52
ideal_angle = 1.81
bond_params = [
tf.get_variable("OH_kb",
shape=tuple(),
dtype=tf.float64,
initializer=tf.constant_initializer(100.0)),
tf.get_variable("OH_b0",
shape=tuple(),
dtype=tf.float64,
initializer=tf.constant_initializer(ideal_bond)),
]
hb = bonded_force.HarmonicBondForce(
params=bond_params,
bond_idxs=np.array([[0, 1], [0, 2]], dtype=np.int32),
param_idxs=np.array([[0, 1], [0, 1]], dtype=np.int32))
angle_params = [
tf.get_variable("HOH_ka",
shape=tuple(),
dtype=tf.float64,
initializer=tf.constant_initializer(75.0)),
tf.get_variable("HOH_a0",
shape=tuple(),
dtype=tf.float64,
initializer=tf.constant_initializer(ideal_angle)),
]
ha = bonded_force.HarmonicAngleForce(
params=angle_params,
angle_idxs=np.array([[1, 0, 2]], dtype=np.int32),
param_idxs=np.array([[0, 1]], dtype=np.int32))
# standard MD used to generate a canonical ensemble
friction = 1.0
dt = 0.0025
temp = 300.0
num_steps = 20000
x_ph = tf.placeholder(name="input_geom",
dtype=tf.float64,
shape=(num_atoms, 3))
intg = integrator.LangevinIntegrator(masses, x_ph, [hb, ha], dt,
friction, temp)
reservoir_size = 200
ref_dx_op, _ = intg.step_op(inference=True)
sess = tf.Session()
sess.run(tf.initializers.global_variables())
def reference_generator():
x = np.copy(
x_opt
) # (ytz): Do not remove this copy else you'll spend hours tracking down bugs.
for step in range(num_steps):
if step % 100 == 0:
print(step)
dx_val = sess.run(ref_dx_op, feed_dict={x_ph: x})
x += dx_val
# this copy is important here, otherwise we're just adding the same reference
# since += modifies the object in_place
yield np.copy(x), step # wip: decouple
def generate_reference_ensemble():
intg.reset(sess)
rs = ReservoirSampler(reference_generator(), reservoir_size)
rs.sample_all()
ensemble = []
for x, t in rs.R:
ensemble.append(x)
stacked_ensemble = np.stack(ensemble, axis=0)
ensemble_ph = tf.placeholder(shape=(reservoir_size, num_atoms, 3),
dtype=np.float64)
ref_d2ij_op = observable.sorted_squared_distances(ensemble_ph)
return ref_d2ij_op, stacked_ensemble, ensemble_ph
a1, inp1, ensemble_ph1 = generate_reference_ensemble()
a2, inp2, ensemble_ph2 = generate_reference_ensemble()
# compute MSE of two _identical_ simulations except for the RNG
loss = tf.reduce_sum(tf.pow(a1 - a2, 2)) / reservoir_size # 0.003
mutual_MSE = sess.run(loss,
feed_dict={
ensemble_ph1: inp1,
ensemble_ph2: inp2
})
print("Optimal MSE", mutual_MSE)
# on average, two identical ensembles should yield the same result
assert mutual_MSE < 0.05
# Use completely different integration parameters
friction = 10.0
dt = 0.05
temp = 100
num_steps = 500 # we need to make the num_steps variable and only stop once we've converged
x_ph = tf.placeholder(name="input_geom",
dtype=tf.float64,
shape=(num_atoms, 3))
with tf.variable_scope("bad"):
bad_intg = integrator.LangevinIntegrator(masses,
x_ph, [hb, ha],
dt,
friction,
temp,
buffer_size=None)
bad_dx_op, bad_dxdp_op = bad_intg.step_op()
d0_ph = tf.placeholder(shape=tuple(), dtype=tf.float64)
d1_ph = tf.placeholder(shape=tuple(), dtype=tf.float64)
d2_ph = tf.placeholder(shape=tuple(), dtype=tf.float64)
d3_ph = tf.placeholder(shape=tuple(), dtype=tf.float64)
grads_and_vars = []
for dp, var in zip([d0_ph, d1_ph, d2_ph, d3_ph],
bond_params + angle_params):
grads_and_vars.append((dp, var))
# param_optimizer = tf.train.AdamOptimizer(0.02)
# param_optimizer = tf.train.AdamOptimizer(0.1) # unstable
param_optimizer = tf.train.RMSPropOptimizer(0.1)
train_op = param_optimizer.apply_gradients(grads_and_vars)
sess.run(tf.initializers.global_variables())
# it turns out the force constants don't really matter a whole lot in this case, but the
# ideal lengths/angles do matter.
# Use completely different forcefield parameters, changing bond constants and lengths.
# See if we can recover the original parameters again.
sess.run([
tf.assign(bond_params[0], 60),
tf.assign(bond_params[1], 1.3),
tf.assign(angle_params[0], 43),
tf.assign(angle_params[1], 2.1)
])
print("Starting params", sess.run(bond_params + angle_params))
for epoch in range(10000):
def sub_optimal_generator():
x = np.copy(x_opt)
for step in range(num_steps):
dx_val, dxdp_val = sess.run([bad_dx_op, bad_dxdp_op],
feed_dict={x_ph: x})
x += dx_val
yield np.copy(x), dxdp_val, step # wip: decouple
bad_intg.reset(sess)
rs = ReservoirSampler(sub_optimal_generator(), reservoir_size)
rs.sample_all()
bad_ensemble = []
bad_ensemble_grads = []
for x, dxdp, t in rs.R:
bad_ensemble.append(x)
bad_ensemble_grads.append(dxdp)
stacked_bad_ensemble = np.stack(bad_ensemble, axis=0)
stacked_bad_ensemble_grads = np.stack(bad_ensemble_grads, axis=0)
# compute ensemble's average angle and bond length
bad_ensemble_ph = tf.placeholder(shape=(reservoir_size, num_atoms,
3),
dtype=np.float64)
ref_d2ij_op = observable.sorted_squared_distances(bad_ensemble_ph)
# b1, bnp1, bbad_ensemble_ph1 = generate_sub_optimal_bad_ensemble()
loss_op = tf.reduce_sum(tf.pow(a1 - ref_d2ij_op,
2)) / reservoir_size # 0.003
dLdx_op = tf.gradients(loss_op, bad_ensemble_ph)
loss, dLdx_val = sess.run([loss_op, dLdx_op],
feed_dict={
ensemble_ph1: inp1,
bad_ensemble_ph: stacked_bad_ensemble
}) # MSE 12.953122852970827
dLdx_dxdp = np.multiply(np.expand_dims(dLdx_val[0], 1),
stacked_bad_ensemble_grads)
reduced_dLdp = np.sum(dLdx_dxdp, axis=tuple([0, 2, 3]))
sess.run(train_op,
feed_dict={
d0_ph: reduced_dLdp[0],
d1_ph: reduced_dLdp[1],
d2_ph: reduced_dLdp[2],
d3_ph: reduced_dLdp[3],
})
if loss < mutual_MSE * 5:
# succesfully converged (should take about 600 epochs)
return
print("loss", loss, "epoch", epoch, "current params",
sess.run(bond_params + angle_params))
assert 0
def test_mol(smiles):
nrgs1, offsets1, intg1, context1, init_x_1, total_params_1 = initialize_system(
dt=0.001,
temperature=20,
forcefield_file='forcefield/smirnoff99Frosst.offxml')
# generate the observable
print("generating observable")
ksize = 500
confs1, _ = run_once(nrgs1, context1, intg1, init_x_1, 40000,
total_params_1, ksize)
x1 = tf.convert_to_tensor(confs1)
obs1_rij = observable.sorted_squared_distances(x1)
# train this secondary system
print("starting training...")
bond_learning_rate = np.array([[0.1, 0.0001]])
angle_learning_rate = np.array([[0.01, 0.001]])
torsion_learning_rate = np.array([[0.01, 0.001, 0.0]])
lj_learning_rate = np.array([[0.000, 0.000]])
nrgs0, offsets0, intg0, context0, init_x_0, total_params_0 = initialize_system(
dt=0.001,
temperature=20,
forcefield_file='forcefield/smirnoff99Frosst_perturbed.offxml')
for epoch in range(5000):
print("starting epoch", epoch)
confs0, dxdp0 = run_once(nrgs0, context0, intg0, init_x_0, 40000,
total_params_0, ksize)
x0 = tf.convert_to_tensor(confs0)
obs0_rij = observable.sorted_squared_distances(x0)
loss = tf.sqrt(tf.reduce_sum(tf.pow(obs0_rij - obs1_rij, 2)) /
ksize) # RMSE
x0_grads = tf.gradients(loss, x0)[0]
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session()
np_loss, x0g = sess.run([loss, x0_grads])
# np_loss = sess.run(loss)
print("------------------LOSS", np_loss)
# print("x0g shape", x0g.shape)
# assert 0
x0g = np.expand_dims(x0g, 1) # [B, 1, N, 3]
res = np.multiply(x0g, dxdp0) # dL/dx * dx/dp [B, P, N, 3]
dLdp = np.sum(res, axis=(0, 2, 3))
for dparams, nrg in zip(np.split(dLdp, offsets0)[1:], nrgs0):
if isinstance(nrg, custom_ops.HarmonicBondGPU_double):
cp = nrg.get_params()
dp = bond_learning_rate * dparams.reshape((-1, 2))
print("BOND PARAMS", cp)
print("BOND CONSTANTS, LENGTHS", dp)
nrg.set_params(cp - dp.reshape(-1))
elif isinstance(nrg, custom_ops.HarmonicAngleGPU_double):
dp = angle_learning_rate * dparams.reshape((-1, 2))
print("ANGLE CONSTANTS, ANGLES", dp)
elif isinstance(nrg, custom_ops.PeriodicTorsionGPU_double):
dp = torsion_learning_rate * dparams.reshape((-1, 3))
print("TORSION CONSTANTS, PERIODS, PHASES", dp)
elif isinstance(nrg, custom_ops.LennardJonesGPU_double):
dp = lj_learning_rate * dparams.reshape((-1, 2))
print("LJ SIG, EPS", dp)
else:
assert 0
# nrg.set_params(cp - dp.reshape(-1))
# nrg.set_params(cp)
assert 0
for step in range(num_steps):
if step % 500 == 0:
# print(step)
coords = np.array(intg.get_coordinates()).reshape((-1, 3))
print(coords)
center = np.sum(coords, axis=0) / coords.shape[0]
dto = np.sqrt(np.sum(np.power(center - origin, 2)))
velocities = np.array(intg.get_velocities()).reshape((-1, 3))
net_mass = np.sum(masses)
# nv = np.sum(np.expand_dims(np.array(masses), axis=-1)*velocities, axis=0)
nv = np.sum(np.expand_dims(np.array(masses), axis=-1) * velocities,
axis=0)
# assert 0
cc_bond_length = np.sqrt(
np.sum(np.power(coords[0, :] - coords[1, :], 2)))
write_xyz(ofs, mol, np.array(coords) * 10)
dxdp = np.array(intg.get_dxdp()).reshape(
(total_params, num_atoms, 3))
amax, amin = np.amax(dxdp), np.amin(dxdp)
print(step, "\tdto\t", dto, "\tnv\t", nv, "\tcc_bond_length\t",
cc_bond_length, "\tamax/amin", amax, "\t", amin)
segments = np.split(dxdp, offsets)[1:]
for grads, force in zip(segments, nrgs):
# print(force)
if isinstance(force, custom_ops.HarmonicBondGPU_double):
grads = grads.reshape((-1, 2, num_atoms, 3))
print("Bond Constants:", np.amax(grads[:, 0, :, :]),
np.amin(grads[:, 0, :, :]))
print("Bond Lengths:", np.amax(grads[:, 1, :, :]),
np.amin(grads[:, 1, :, :]))
# print(grads[:, 1, :, :])
elif isinstance(force, custom_ops.HarmonicAngleGPU_double):
grads = grads.reshape((-1, 2, num_atoms, 3))
print("Angle Constants:", np.amax(grads[:, 0, :, :]),
np.amin(grads[:, 0, :, :]))
# print(grads[:, 0, :, :])
print("Angle Radians:", np.amax(grads[:, 1, :, :]),
np.amin(grads[:, 1, :, :]))
# print(grads[:, 1, :, :])
elif isinstance(force, custom_ops.PeriodicTorsionGPU_double):
grads = grads.reshape((-1, 3, num_atoms, 3))
print("Torsion Constants:", np.amax(grads[:, 0, :, :]),
np.amin(grads[:, 0, :, :]))
print("Torsion Phase:", np.amax(grads[:, 1, :, :]),
np.amin(grads[:, 1, :, :]))
print("Torsion Periods:", np.amax(grads[:, 2, :, :]),
np.amin(grads[:, 2, :, :]))
elif isinstance(force, custom_ops.LennardJonesGPU_double):
grads = grads.reshape((-1, 2, num_atoms, 3))
print("LJ sigma:", np.amax(grads[:, 0, :, :]),
np.amin(grads[:, 0, :, :]))
print("LJ epsilon:", np.amax(grads[:, 1, :, :]),
np.amin(grads[:, 1, :, :]))
else:
assert 0
if np.any(np.isnan(dxdp)):
assert 0
context.step()