| # reference from DeepDock nature machine intellience paper | |
| # import numpy as np | |
| import torch | |
| def compute_euclidean_distances_matrix(X, Y): | |
| # Based on: https://medium.com/@souravdey/l2-distance-matrix-vectorization-trick-26aa3247ac6c | |
| # (X-Y)^2 = X^2 + Y^2 -2XY | |
| X = X.double() | |
| Y = Y.double() | |
| 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) | |
| return dists**0.5 | |
| def compute_euclidean_distances_matrix_TopN( X, Y,B, N_l,topN = 1): | |
| X = X.double() | |
| Y = Y.double() | |
| 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) | |
| dists = torch.nan_to_num((dists**0.5).view(B, N_l,-1,24),10000).sort(axis=-1)[0][:,:,:,:topN] | |
| dist_topN = [] | |
| for i in range(topN): | |
| dist_topN.append(dists[:,:,:,i]) | |
| return dist_topN | |