def test_works(sess, traj):
rmsd = tftraj.rmsd_op.load()
inds = [5, 19, 234, 235]
target = np.array(traj.xyz[inds])
target = tf.Variable(target)
prmsd, _ = rmsd.pairwise_msd(traj.xyz, target)
sess.run(tf.global_variables_initializer())
result = sess.run(prmsd)
print(result.shape)
grad = tf.gradients(prmsd, [target])[0]
grad_result = sess.run(grad)
print(grad_result.shape)
print(grad_result)
def test_vs_tensorflow_target(sess, traj):
rmsd = tftraj.rmsd_op.load()
inds = [5, 19, 234, 235]
target = np.array(traj.xyz[inds])
target = tf.Variable(target)
sess.run(tf.global_variables_initializer())
rmsd_from_op, _ = rmsd.pairwise_msd(tf.constant(traj.xyz), target)
grad_from_op = tf.gradients(rmsd_from_op, [target])[0]
grad_from_op = sess.run(grad_from_op)
rmsd_from_tf = tftraj.rmsd.pairwise_msd(tf.constant(traj.xyz), target)
grad_from_tf = tf.gradients(rmsd_from_tf, [target])[0]
grad_from_tf = sess.run(grad_from_tf)
np.testing.assert_almost_equal(grad_from_op, grad_from_tf, decimal=3)
def test_vs_tensorflow_both_input(sess, traj):
rmsd = tftraj.rmsd_op.load()
inds = [5, 19, 234, 235]
target = np.array(traj.xyz[inds])
traj = tf.Variable(traj.xyz)
target = tf.Variable(target)
sess.run(tf.global_variables_initializer())
rmsd_from_op, _ = rmsd.pairwise_msd(traj, target)
grad_from_op = tf.gradients(rmsd_from_op, [traj, target])
grad_from_op = sess.run(grad_from_op)
rmsd_from_tf = tftraj.rmsd.pairwise_msd(traj, target)
grad_from_tf = tf.gradients(rmsd_from_tf, [traj, target])
grad_from_tf = sess.run(grad_from_tf)
assert len(grad_from_tf) == 2
assert len(grad_from_op) == 2
for i in range(2):
print(i)
np.testing.assert_almost_equal(grad_from_op[i],
grad_from_tf[i],
decimal=3)
import tftraj.rmsd_op
import tftraj.rmsd
results = {}
sess = tf.Session()
traj = md.load(
['fs_peptide/trajectory-{}.xtc'.format(i + 1) for i in range(28)],
top='fs_peptide/fs-peptide.pdb')
traj = traj[::100]
traj_xyz = np.array(traj.xyz)
traj_target = traj[::100]
traj_target_xyz = np.array(traj_target.xyz)
print(len(traj_xyz), len(traj_target_xyz))
rmsd = tftraj.rmsd_op.load()
prmsd, _ = rmsd.pairwise_msd(traj_xyz, traj_target_xyz)
results['tf-cpu'] = timeit.timeit(
'sess.run(prmsd)', number=30, globals=globals()) / 30
results['mdtraj'] = timeit.timeit(
'[md.rmsd(traj, traj_target, i) ** 2 for i in range(traj_target.n_frames)]',
number=30,
globals=globals()) / 30
tfnative = tftraj.rmsd.pairwise_msd(tf.constant(traj_xyz),
tf.constant(traj_target_xyz))
results['tf-native'] = timeit.timeit('sess.run(tfnative)',
number=1,
globals=globals())
print("{:10s} {:7s}".format("Algo", "time/ms"))