<reponame>JiajunBao/neural-dimension-reduction
import torch
from tqdm.auto import tqdm
from torch.nn import functional as F


def nearest_neighbors(x, top_k, device):
    """
    calculate the nearest neighbors of x, return the
    :param x: for matrix to calculate nearest neighbor
    :param top_k: number of the nearest neighbor to be returned
    :param device: device used during computation
    :return:
            ground_min_dist_square: torch.tensor (n, ) distance to the nearest neighbor
            topk_neighbors: torch.tensor (n, top_k) the index of the top-k nearest neighbors;
    """
    batch_size = 2000
    x_to_device = x.to(device)
    if x_to_device.shape[0] * x_to_device.shape[1] < batch_size * 200:  # direct computes the whole matrix
        dist = torch.cdist(x1=x_to_device, x2=x_to_device, p=2)  # (n, n)
        sorted_dist, indices = torch.sort(dist, dim=1, descending=False)
        ground_min_dist_square = sorted_dist[:, 1]  # the 0-th column is the distance to oneself
        topk_neighbors = indices[:, 1:1 + top_k]
        topk_dists = sorted_dist[:, 1:1 + top_k]
    else:  # calculate the nearest neighbors in batches
        num_iter = x_to_device.shape[0] // batch_size + 1
        topk_neighbors_list = list()
        ground_min_dist_square_list = list()
        sorted_dist_list = list()
        for i in tqdm(torch.arange(num_iter), desc='computing nearest neighbors'):
            batch_x = x_to_device[i * batch_size: (i + 1) * batch_size, :]
            dist = torch.cdist(x1=batch_x, x2=x_to_device, p=2)  # (n, n)
            sorted_dist, indices = torch.sort(dist, dim=1, descending=False)
            batch_ground_min_dist_square = sorted_dist[:, 1]  # the 0-th column is the distance to oneself
            batch_topk_neighbors = indices[:, 1:1 + top_k]
            topk_neighbors_list.append(batch_topk_neighbors.cpu())
            ground_min_dist_square_list.append(batch_ground_min_dist_square.cpu())
            sorted_dist_list.append(sorted_dist[:, 1:1 + top_k].cpu())
        ground_min_dist_square = torch.cat(ground_min_dist_square_list, dim=0)
        topk_neighbors = torch.cat(topk_neighbors_list, dim=0)
        topk_dists = torch.cat(sorted_dist_list, dim=0)
    return ground_min_dist_square.cpu(), topk_neighbors.cpu(), topk_dists.cpu()


def euclidean_softmax_similarity(vec_i, vec_j, ground_min_dist_square_i=None, two_eps_square=1):
    """
    calculate inverse similarity for inputs:
        1 / ((d_{in}(x_i, x_j))^2 / d_i^2 + eps)
    :param vec_i:  torch.tensor of the shape (n, m)
                xi
    :param vec_j: torch.tensor of the shape (n, m)
                xj
    :param two_eps_square: 2 * (epsilon)^2
    :param ground_min_dist_square_i:
    :return: q: qij in formula (P20-3) torch.tensor of the shape (n, m)
    """

    din = (vec_i.unsqueeze(dim=1) - vec_j).square().sum(dim=2)  # (n, m)
    sim_j_given_i = F.softmin(din / two_eps_square, dim=1)  # (n, m)
    return sim_j_given_i


def kl_div_loss(input_similarity, output_similarity):
    """
    calculate the sum of kl divergence and mse loss
    :param input_similarity: input similarity
    :param output_similarity: output similarity
    :return: torch.tensor of the shape (,)
    """
    return torch.sum(input_similarity * torch.log(input_similarity / output_similarity))


def precomputing(x, top_k, device):
    """
    compute ground true nearest neighbors
    :param x:
    :param top_k: top-k neighbors that are considered
    :param device: device used during computation
    :return: anchor_idx: each point has m points as anchors (in the case, we pick m near neighbors of x as anchors)
             input_similarity: input_similarity
    """
    ground_min_dist_square, anchor_idx, topk_dists = nearest_neighbors(x, top_k, device)

    xi = x
    xj = x[anchor_idx, :]
    input_similarity = euclidean_softmax_similarity(xi, xj, ground_min_dist_square)
    return anchor_idx, input_similarity, ground_min_dist_square, topk_dists

import torch
from tqdm.auto import tqdm
from torch.nn import functional as F


def nearest_neighbors(x, top_k, device):
    """
    calculate the nearest neighbors of x, return the
    :param x: for matrix to calculate nearest neighbor
    :param top_k: number of the nearest neighbor to be returned
    :param device: device used during computation
    :return:
            ground_min_dist_square: torch.tensor (n, ) distance to the nearest neighbor
            topk_neighbors: torch.tensor (n, top_k) the index of the top-k nearest neighbors;
    """
    batch_size = 2000
    x_to_device = x.to(device)
    if x_to_device.shape[0] * x_to_device.shape[1] < batch_size * 200:  # direct computes the whole matrix
        dist = torch.cdist(x1=x_to_device, x2=x_to_device, p=2)  # (n, n)
        sorted_dist, indices = torch.sort(dist, dim=1, descending=False)
        ground_min_dist_square = sorted_dist[:, 1]  # the 0-th column is the distance to oneself
        topk_neighbors = indices[:, 1:1 + top_k]
        topk_dists = sorted_dist[:, 1:1 + top_k]
    else:  # calculate the nearest neighbors in batches
        num_iter = x_to_device.shape[0] // batch_size + 1
        topk_neighbors_list = list()
        ground_min_dist_square_list = list()
        sorted_dist_list = list()
        for i in tqdm(torch.arange(num_iter), desc='computing nearest neighbors'):
            batch_x = x_to_device[i * batch_size: (i + 1) * batch_size, :]
            dist = torch.cdist(x1=batch_x, x2=x_to_device, p=2)  # (n, n)
            sorted_dist, indices = torch.sort(dist, dim=1, descending=False)
            batch_ground_min_dist_square = sorted_dist[:, 1]  # the 0-th column is the distance to oneself
            batch_topk_neighbors = indices[:, 1:1 + top_k]
            topk_neighbors_list.append(batch_topk_neighbors.cpu())
            ground_min_dist_square_list.append(batch_ground_min_dist_square.cpu())
            sorted_dist_list.append(sorted_dist[:, 1:1 + top_k].cpu())
        ground_min_dist_square = torch.cat(ground_min_dist_square_list, dim=0)
        topk_neighbors = torch.cat(topk_neighbors_list, dim=0)
        topk_dists = torch.cat(sorted_dist_list, dim=0)
    return ground_min_dist_square.cpu(), topk_neighbors.cpu(), topk_dists.cpu()


def euclidean_softmax_similarity(vec_i, vec_j, ground_min_dist_square_i=None, two_eps_square=1):
    """
    calculate inverse similarity for inputs:
        1 / ((d_{in}(x_i, x_j))^2 / d_i^2 + eps)
    :param vec_i:  torch.tensor of the shape (n, m)
                xi
    :param vec_j: torch.tensor of the shape (n, m)
                xj
    :param two_eps_square: 2 * (epsilon)^2
    :param ground_min_dist_square_i:
    :return: q: qij in formula (P20-3) torch.tensor of the shape (n, m)
    """

    din = (vec_i.unsqueeze(dim=1) - vec_j).square().sum(dim=2)  # (n, m)
    sim_j_given_i = F.softmin(din / two_eps_square, dim=1)  # (n, m)
    return sim_j_given_i


def kl_div_loss(input_similarity, output_similarity):
    """
    calculate the sum of kl divergence and mse loss
    :param input_similarity: input similarity
    :param output_similarity: output similarity
    :return: torch.tensor of the shape (,)
    """
    return torch.sum(input_similarity * torch.log(input_similarity / output_similarity))


def precomputing(x, top_k, device):
    """
    compute ground true nearest neighbors
    :param x:
    :param top_k: top-k neighbors that are considered
    :param device: device used during computation
    :return: anchor_idx: each point has m points as anchors (in the case, we pick m near neighbors of x as anchors)
             input_similarity: input_similarity
    """
    ground_min_dist_square, anchor_idx, topk_dists = nearest_neighbors(x, top_k, device)

    xi = x
    xj = x[anchor_idx, :]
    input_similarity = euclidean_softmax_similarity(xi, xj, ground_min_dist_square)
    return anchor_idx, input_similarity, ground_min_dist_square, topk_dists
calculate the nearest neighbors of x, return the     :param x: for matrix to calculate nearest neighbor     :param top_k: number of the nearest neighbor to be returned     :param device: device used during computation     :return:             ground_min_dist_square: torch.tensor (n, ) distance to the nearest neighbor             topk_neighbors: torch.tensor (n, top_k) the index of the top-k nearest neighbors; # direct computes the whole matrix # (n, n) # the 0-th column is the distance to oneself # calculate the nearest neighbors in batches # (n, n) # the 0-th column is the distance to oneself calculate inverse similarity for inputs:         1 / ((d_{in}(x_i, x_j))^2 / d_i^2 + eps)     :param vec_i:  torch.tensor of the shape (n, m)                 xi     :param vec_j: torch.tensor of the shape (n, m)                 xj     :param two_eps_square: 2 * (epsilon)^2     :param ground_min_dist_square_i:     :return: q: qij in formula (P20-3) torch.tensor of the shape (n, m) # (n, m) # (n, m) calculate the sum of kl divergence and mse loss     :param input_similarity: input similarity     :param output_similarity: output similarity     :return: torch.tensor of the shape (,) compute ground true nearest neighbors     :param x:     :param top_k: top-k neighbors that are considered     :param device: device used during computation     :return: anchor_idx: each point has m points as anchors (in the case, we pick m near neighbors of x as anchors)              input_similarity: input_similarity