biopython_wrapper.py 文件源码

def calculate_rms_with_rotran(atoms1: np.ndarray, atoms2: np.ndarray, rotran: tuple) -> Tuple[float, float, Tuple]:
    # apply rotran on atoms2
    atoms2 = np.dot(atoms2, rotran[0]) + rotran[1]

    # fix atoms1 in place, use KDTree for distances
    kd_tree_a = KDTree(atoms1)
    kd_tree_b = KDTree(atoms2)
    # get nearest neighbours for atoms2 in atoms1 using KDTree, and also the other way around
    distances, indexes_a = kd_tree_b.query(atoms1)
    _, indexes_b = kd_tree_a.query(atoms2)

    # pick only the pairs that are both closest to each other
    closest_indexes = np.where(indexes_b[indexes_a] == list(range(len(indexes_a))))
    smoothing_rotran = calculate_rotran(atoms1[closest_indexes], atoms2[indexes_a[closest_indexes]])
    atoms2 = np.dot(atoms2, smoothing_rotran[0]) + smoothing_rotran[1]

    kd_tree_b = KDTree(atoms2)
    distances, indexes_a = kd_tree_b.query(atoms1)
    _, indexes_b = kd_tree_a.query(atoms2)

    # pick only the pairs that are both closest to each other
    closest_indexes = np.where(indexes_b[indexes_a] == list(range(len(indexes_a))))
    distances = distances[closest_indexes]

    # return RMSD ( sqrt(1/N * SUM(distance(xi, yi)^2)) )
    post_rms = math.sqrt(np.sum(np.power(distances, 2)) / float(len(distances)))

    psi = 100.0 * np.sum(distances <= 4.0) / float(min(len(atoms1), len(atoms2)))

    return post_rms, psi, smoothing_rotran