def interatomic_distances(positions, cell, pbc, cutoff):
with tf.variable_scope('distance'):
# calculate heights
# account for zero cell in case of no pbc
c = tf.reduce_sum(tf.cast(pbc, tf.int32)) > 0
icell = tf.cond(c, lambda: tf.matrix_inverse(cell),
lambda: tf.eye(3))
height = 1. / tf.sqrt(tf.reduce_sum(tf.square(icell), 0))
extent = tf.where(tf.cast(pbc, tf.bool),
tf.cast(tf.floor(cutoff / height), tf.int32),
tf.cast(tf.zeros_like(height), tf.int32))
n_reps = tf.reduce_prod(2 * extent + 1)
# replicate atoms
r = tf.range(-extent[0], extent[0] + 1)
v0 = tf.expand_dims(r, 1)
v0 = tf.tile(v0,
tf.stack(((2 * extent[1] + 1) * (2 * extent[2] + 1), 1)))
v0 = tf.reshape(v0, tf.stack((n_reps, 1)))
r = tf.range(-extent[1], extent[1] + 1)
v1 = tf.expand_dims(r, 1)
v1 = tf.tile(v1, tf.stack((2 * extent[2] + 1, 2 * extent[0] + 1)))
v1 = tf.reshape(v1, tf.stack((n_reps, 1)))
v2 = tf.expand_dims(tf.range(-extent[2], extent[2] + 1), 1)
v2 = tf.tile(v2,
tf.stack((1, (2 * extent[0] + 1) * (2 * extent[1] + 1))))
v2 = tf.reshape(v2, tf.stack((n_reps, 1)))
v = tf.cast(tf.concat((v0, v1, v2), axis=1), tf.float32)
offset = tf.matmul(v, cell)
offset = tf.expand_dims(offset, 0)
# add axes
positions = tf.expand_dims(positions, 1)
rpos = positions + offset
rpos = tf.expand_dims(rpos, 0)
positions = tf.expand_dims(positions, 1)
euclid_dist = tf.sqrt(
tf.reduce_sum(tf.square(positions - rpos),
reduction_indices=3))
return euclid_dist
评论列表
文章目录