OneScience-Group/SurfDock

1# reference from DeepDock nature machine intellience paper2# import numpy as np3import torch4 5def compute_euclidean_distances_matrix(X, Y):6    # Based on: https://medium.com/@souravdey/l2-distance-matrix-vectorization-trick-26aa3247ac6c7    # (X-Y)^2 = X^2 + Y^2 -2XY8    X = X.double()9    Y = Y.double()10    dists = -2 * torch.bmm(X, Y.permute(0, 2, 1)) + torch.sum(Y**2,    axis=-1).unsqueeze(1) + torch.sum(X**2, axis=-1).unsqueeze(-1)11    return dists**0.512def compute_euclidean_distances_matrix_TopN( X, Y,B, N_l,topN = 1):13    X = X.double()14    Y = Y.double()15    dists = -2 * torch.bmm(X, Y.permute(0, 2, 1)) + torch.sum(Y**2,    axis=-1).unsqueeze(1) + torch.sum(X**2, axis=-1).unsqueeze(-1)16    dists = torch.nan_to_num((dists**0.5).view(B, N_l,-1,24),10000).sort(axis=-1)[0][:,:,:,:topN]17    dist_topN = []18    for i in range(topN):19        dist_topN.append(dists[:,:,:,i])20    return dist_topN21