atoms.py 文件源码

def site_rdf(distances, cutoff, step, width, eps=1e-5,
             use_mean=False, lower_cutoff=None):
    with tf.variable_scope('srdf'):
        if lower_cutoff is None:
            vrange = cutoff
        else:
            vrange = cutoff - lower_cutoff
        distances = tf.expand_dims(distances, -1)
        n_centers = np.ceil(vrange / step)
        gap = vrange - n_centers * step
        n_centers = int(n_centers)

        if lower_cutoff is None:
            centers = tf.linspace(0., cutoff - gap, n_centers)
        else:
            centers = tf.linspace(lower_cutoff + 0.5 * gap, cutoff - 0.5 * gap,
                                  n_centers)
        centers = tf.reshape(centers, (1, 1, 1, -1))

        gamma = -0.5 / width / step ** 2

        rdf = tf.exp(gamma * (distances - centers) ** 2)

        mask = tf.cast(distances >= eps, tf.float32)
        rdf *= mask
        rdf = tf.reduce_sum(rdf, 2)
        if use_mean:
            N = tf.reduce_sum(mask, 2)
            N = tf.maximum(N, 1)
            rdf /= N

        new_shape = [None, None, n_centers]
        rdf.set_shape(new_shape)

    return rdf