Python torch.argsort() Examples

def _run_encoder(self, batch):
        src, src_lengths = batch.src if isinstance(batch.src, tuple) \
                           else (batch.src, None)

        # added by @memray, resort examples in batch by lengths first
        sort_idx = torch.argsort(src_lengths, descending=True)
        sorted_src = src[:,sort_idx,:]
        sorted_src_lengths = src_lengths[sort_idx]
        unsort_idx = torch.argsort(sort_idx)

        enc_states, memory_bank, src_lengths = self.model.encoder(
            sorted_src, sorted_src_lengths)

        enc_states = enc_states[:, unsort_idx, :]
        memory_bank = memory_bank[:, unsort_idx, :]
        src_lengths = src_lengths[unsort_idx]

        if src_lengths is None:
            assert not isinstance(memory_bank, tuple), \
                'Ensemble decoding only supported for text data'
            src_lengths = torch.Tensor(batch.batch_size) \
                               .type_as(memory_bank) \
                               .long() \
                               .fill_(memory_bank.size(0))
        return src, enc_states, memory_bank, src_lengths 

def _get_augmented_pareto_front_indices(self) -> Tensor:
        r"""Get indices of augmented pareto front."""
        pf_idx = torch.argsort(self._pareto_Y, dim=0)
        return torch.cat(
            [
                torch.zeros(
                    1, self.num_outcomes, dtype=torch.long, device=self.Y.device
                ),
                # Add 1 because index zero is used for the ideal point
                pf_idx + 1,
                torch.full(
                    torch.Size([1, self.num_outcomes]),
                    self._pareto_Y.shape[0] + 1,
                    dtype=torch.long,
                    device=self.Y.device,
                ),
            ],
            dim=0,
        ) 

def order_points(pts):
    pts_reorder = []

    for idx, pt in enumerate(pts):
        idx = torch.argsort(pt[:, 0])
        xSorted = pt[idx, :]
        leftMost = xSorted[:2, :]
        rightMost = xSorted[2:, :]

        leftMost = leftMost[torch.argsort(leftMost[:, 1]), :]
        (tl, bl) = leftMost

        D = torch.cdist(tl[np.newaxis], rightMost)[0]
        (br, tr) = rightMost[torch.argsort(D, descending=True), :]
        pts_reorder.append(torch.stack([tl, tr, br, bl]))

    return torch.stack([p for p in pts_reorder]) 

def get_mask(self, base_mask, weight, num_prune, wrapper, wrapper_idx):
        assert wrapper_idx is not None
        activations = self.pruner.collected_activation[wrapper_idx]
        if len(activations) < self.statistics_batch_num:
            return None
        mean_activation = self._cal_mean_activation(activations)
        prune_indices = torch.argsort(mean_activation)[:num_prune]
        for idx in prune_indices:
            base_mask['weight_mask'][idx] = 0.
            if base_mask['bias_mask'] is not None:
                base_mask['bias_mask'][idx] = 0.

        if len(activations) >= self.statistics_batch_num and self.pruner.hook_id in self.pruner._fwd_hook_handles:
            self.pruner.remove_activation_collector(self.pruner.hook_id)

        return base_mask 

def compute(self, pred: torch.Tensor, target: torch.Tensor) \
            -> torch.Tensor:
        """Computes the recall @ k.

        Parameters
        ----------
        pred: Tensor
            input logits of shape (B x N)
        target: LongTensor
            target tensor of shape (B) or (B x N)

        Returns
        -------
        recall: torch.Tensor
            single label recall, of shape (B)

        """
        # If 2-dimensional, select the highest score in each row
        if len(target.size()) == 2:
            target = target.argmax(dim=1)

        ranked_scores = torch.argsort(pred, dim=1)[:, -self.top_k:]
        recalled = torch.sum((target.unsqueeze(1) == ranked_scores).float(), dim=1)
        return recalled.mean() 

def new_predict(query, neighbors, neighbor_labels):
    '''
    query: p
    neighbors: n x p
    neighbor_labels: n (int)
    '''
    diffs = (query.unsqueeze(1) - neighbors.unsqueeze(0))
    squared_diffs = diffs ** 2
    norms = squared_diffs.sum(-1)
    indices = torch.argsort(norms, dim=-1)
    labels = neighbor_labels.take(indices[:, :k])
    prediction = [majority(l.tolist()) for l in labels]
    return torch.Tensor(prediction).to(device=gpu).long() 

def _inverse_permutation(self):
        return torch.argsort(self._permutation) 

def soft_rank_sampling(loc, covariance_matrix=None, inds_style=True, descending=True):
    '''
    :param loc: mean of the distribution
    :param covariance_matrix: positive-definite covariance matrix
    :param inds_style: true means the indice leading to the ltr_adhoc
    :return:
    '''
    m = MultivariateNormal(loc, covariance_matrix)
    vals = m.sample()
    if inds_style:
        sorted_inds = torch.argsort(vals, descending=descending)
        return sorted_inds
    else:
        vals 

def get_mask(self, base_mask, weight, num_prune, wrapper, wrapper_idx):
        if self.pruner.iterations < self.pruner.statistics_batch_num:
            return None

        if wrapper.contribution is None:
            return None

        prune_indices = torch.argsort(wrapper.contribution)[:num_prune]
        for idx in prune_indices:
            base_mask['weight_mask'][idx] = 0.
            if base_mask['bias_mask'] is not None:
                base_mask['bias_mask'][idx] = 0.
        return base_mask 

def get_mask(self, base_mask, weight, num_prune, wrapper, wrapper_idx):
        assert wrapper_idx is not None
        activations = self.pruner.collected_activation[wrapper_idx]
        if len(activations) < self.statistics_batch_num:
            return None
        apoz = self._calc_apoz(activations)
        prune_indices = torch.argsort(apoz, descending=True)[:num_prune]
        for idx in prune_indices:
            base_mask['weight_mask'][idx] = 0.
            if base_mask['bias_mask'] is not None:
                base_mask['bias_mask'][idx] = 0.
        return base_mask

        if len(activations) >= self.statistics_batch_num and self.pruner.hook_id in self.pruner._fwd_hook_handles:
            self.pruner.remove_activation_collector(self.pruner.hook_id) 

def auc(x: torch.Tensor, y: torch.Tensor, reorder: bool = True) -> torch.Tensor:
    """
    Computes Area Under the Curve (AUC) using the trapezoidal rule

    Args:
        x: x-coordinates
        y: y-coordinates
        reorder: reorder coordinates, so they are increasing.

    Return:
        Tensor containing AUC score (float)

    Example:

        >>> x = torch.tensor([0, 1, 2, 3])
        >>> y = torch.tensor([0, 1, 2, 2])
        >>> auc(x, y)
        tensor(4.)
    """
    direction = 1.

    if reorder:
        # can't use lexsort here since it is not implemented for torch
        order = torch.argsort(x)
        x, y = x[order], y[order]
    else:
        dx = x[1:] - x[:-1]
        if (dx < 0).any():
            if (dx, 0).all():
                direction = -1.
            else:
                raise ValueError("Reordering is not turned on, and "
                                 "the x array is not increasing: %s" % x)

    return direction * torch.trapz(y, x) 

def add(self, errs, poses, pose_ids):
        e = torch.cat((self.min_errs, errs))
        a = torch.argsort(e)[:self.max_keep]
        self.min_errs = e[a]
        self.min_poses = torch.cat((self.min_poses,poses))[a]
        self.min_pose_ids = torch.cat((self.min_pose_ids, pose_ids))[a] 

def arg_shuffle_ties(vec, descending=True):
    ''' the same as shuffle_ties, but return the corresponding indice '''
    if len(vec.size()) > 1:
        raise NotImplementedError
    else:
        length = vec.size()[0]
        perm = torch.randperm(length)
        sorted_shuffled_vec_inds = torch.argsort(vec[perm], descending=descending)
        shuffle_ties_inds = perm[sorted_shuffled_vec_inds]
        return shuffle_ties_inds 

def _eval_mrr(self, y_pred_pos, y_pred_neg, type_info):
        """
            compute mrr
            y_pred_neg is an array with shape (batch size, num_entities_neg).
            y_pred_pos is an array with shape (batch size, )
        """


        if type_info == 'torch':
            y_pred = torch.cat([y_pred_pos.view(-1,1), y_pred_neg], dim = 1)
            argsort = torch.argsort(y_pred, dim = 1, descending = True)
            ranking_list = (argsort == 0).nonzero()
            ranking_list = ranking_list[:, 1] + 1
            hits1_list = (ranking_list <= 1).to(torch.float)
            hits3_list = (ranking_list <= 3).to(torch.float)
            hits10_list = (ranking_list <= 10).to(torch.float)
            mrr_list = 1./ranking_list.to(torch.float)

            return {"hits@1_list": hits1_list, 
                     "hits@3_list": hits3_list,
                     "hits@10_list": hits10_list,
                     "mrr_list": mrr_list}

        else:
            y_pred = np.concatenate([y_pred_pos.reshape(-1,1), y_pred_neg], axis = 1)
            argsort = np.argsort(-y_pred, axis = 1)
            ranking_list = (argsort == 0).nonzero()
            ranking_list = ranking_list[1] + 1
            hits1_list = (ranking_list <= 1).astype(np.float32)
            hits3_list = (ranking_list <= 3).astype(np.float32)
            hits10_list = (ranking_list <= 10).astype(np.float32)
            mrr_list = 1./ranking_list.astype(np.float32)

            return {"hits@1_list": hits1_list, 
                     "hits@3_list": hits3_list,
                     "hits@10_list": hits10_list,
                     "mrr_list": mrr_list} 

def _update_pareto_Y(self) -> bool:
        r"""Update the non-dominated front."""
        # is_non_dominated assumes maximization
        non_dominated_mask = is_non_dominated(-self.Y)
        pf = self.Y[non_dominated_mask]
        # sort by first objective
        new_pareto_Y = pf[torch.argsort(pf[:, 0])]
        if not hasattr(self, "_pareto_Y") or not torch.equal(
            new_pareto_Y, self._pareto_Y
        ):
            self.register_buffer("_pareto_Y", new_pareto_Y)
            return True
        return False 

def find_local_max(self, box_prob):
        B, C, H, W = box_prob.shape
        max_prob, idx = F.max_pool2d_with_indices(box_prob, 3, 1, 1, return_indices=True)
        max_prob = max_prob[0, 0]
        box_prob = box_prob[0, 0]
        is_local_max = torch.nonzero(box_prob == max_prob)
        y, x = is_local_max[:, 0], is_local_max[:, 1]
        idx = torch.argsort(-box_prob[y, x])
        k = self.scatter_topk
        y = y[idx[:k]]
        x = x[idx[:k]]
        return y.cpu().numpy(), x.cpu().numpy(), box_prob[y, x] 

def find_local_max(self, box_prob):
        B, C, H, W = box_prob.shape
        max_prob, idx = F.max_pool2d_with_indices(box_prob, 3, 1, 1, return_indices=True)
        max_prob = max_prob[0, 0]
        box_prob = box_prob[0, 0]
        is_local_max = torch.nonzero(box_prob == max_prob)
        y, x = is_local_max[:, 0], is_local_max[:, 1]
        idx = torch.argsort(-box_prob[y, x])
        k = self.scatter_topk
        y = y[idx[:k]]
        x = x[idx[:k]]
        return y.cpu().numpy(), x.cpu().numpy(), box_prob[y, x] 

def forward(self, wav):
        batch, streams, channels, time = wav.size()
        device = wav.device

        if self.training:
            group_size = self.group_size or batch
            if batch % group_size != 0:
                raise ValueError(f"Batch size {batch} must be divisible by group size {group_size}")
            groups = batch // group_size
            wav = wav.view(groups, group_size, streams, channels, time)
            permutations = th.argsort(th.rand(groups, group_size, streams, 1, 1, device=device),
                                      dim=1)
            wav = wav.gather(1, permutations.expand(-1, -1, -1, channels, time))
            wav = wav.view(batch, streams, channels, time)
        return wav 

def __init__(self, dim, amortized='none', **kwargs):
        super(ReverseFlow, self).__init__()
        self.permute = torch.arange(dim - 1, -1, -1)
        self.inverse = torch.argsort(self.permute) 

def test_compute_edge_score_tanh():
    edge_index = torch.tensor([[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5],
                               [1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2, 5, 4]])
    raw = torch.randn(edge_index.size(1))
    e = EdgePooling.compute_edge_score_tanh(raw, edge_index, 6)
    assert torch.all(e >= -1) and torch.all(e <= 1)
    assert torch.all(torch.argsort(raw) == torch.argsort(e)) 

def test_compute_edge_score_sigmoid():
    edge_index = torch.tensor([[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5],
                               [1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2, 5, 4]])
    raw = torch.randn(edge_index.size(1))
    e = EdgePooling.compute_edge_score_sigmoid(raw, edge_index, 6)
    assert torch.all(e >= 0) and torch.all(e <= 1)
    assert torch.all(torch.argsort(raw) == torch.argsort(e)) 

def build_part_with_score_torch(score_threshold, local_max_radius, scores):
    lmd = 2 * local_max_radius + 1
    max_vals = F.max_pool2d(scores, lmd, stride=1, padding=1)
    max_loc = (scores == max_vals) & (scores >= score_threshold)
    max_loc_idx = max_loc.nonzero()
    scores_vec = scores[max_loc]
    sort_idx = torch.argsort(scores_vec, descending=True)
    return scores_vec[sort_idx], max_loc_idx[sort_idx]


# FIXME leaving here as reference for now
# def build_part_with_score_fast(score_threshold, local_max_radius, scores):
#     parts = []
#     num_keypoints = scores.shape[0]
#     lmd = 2 * local_max_radius + 1
#
#     # NOTE it seems faster to iterate over the keypoints and perform maximum_filter
#     # on each subarray vs doing the op on the full score array with size=(lmd, lmd, 1)
#     for keypoint_id in range(num_keypoints):
#         kp_scores = scores[keypoint_id, :, :].copy()
#         kp_scores[kp_scores < score_threshold] = 0.
#         max_vals = ndi.maximum_filter(kp_scores, size=lmd, mode='constant')
#         max_loc = np.logical_and(kp_scores == max_vals, kp_scores > 0)
#         max_loc_idx = max_loc.nonzero()
#         for y, x in zip(*max_loc_idx):
#             parts.append((
#                 scores[keypoint_id, y, x],
#                 keypoint_id,
#                 np.array((y, x))
#             ))
#
#    return parts 

def argsort(input, dim, descending):
    return th.argsort(input, dim=dim, descending=descending) 

def get_rank(batch_score, dim=0):
    rank = torch.argsort(batch_score, dim=dim)
    rank = torch.argsort(rank, dim=dim)
    rank = (rank * -1) + batch_score.size(dim)
    rank = rank.float()
    rank = rank / batch_score.size(dim)

    return rank 

def __getitem__(self, index):
        rand_seq = get_rand_seq(self.seq_len, self.dist)
        zipp_sort_ind = zip(np.argsort(rand_seq)[::-1], range(self.seq_len))

        ranks = [((y[1] + 1) / float(self.seq_len)) for y in sorted(zipp_sort_ind, key=lambda x: x[0])]

        return torch.FloatTensor(rand_seq), torch.FloatTensor(ranks) 

def get_rank_single(batch_score):
        rank = torch.argsort(batch_score, dim=0)
        rank = torch.argsort(rank, dim=0)
        rank = (rank * -1) + batch_score.size(0)
        rank = rank.float()
        rank = rank / batch_score.size(0)

        return rank 

def energy_spectrum(vel):
    """
    Compute energy spectrum given a velocity field
    :param vel: tensor of shape (N, 3, res, res, res)
    :return spec: tensor of shape(N, res/2)
    :return k: tensor of shape (res/2,), frequencies corresponding to spec
    """
    device = vel.device
    res = vel.shape[-2:]

    assert(res[0] == res[1])
    r = res[0]
    k_end = int(r/2)
    vel_ = pad_rfft3(vel, onesided=False) # (N, 3, res, res, res, 2)
    uu_ = (torch.norm(vel_, dim=-1) / r**3)**2
    e_ = torch.sum(uu_, dim=1)  # (N, res, res, res)
    k = fftfreqs(res).to(device) # (3, res, res, res)
    rad = torch.norm(k, dim=0) # (res, res, res)
    k_bin = torch.arange(k_end, device=device).float()+1
    bins = torch.zeros(k_end+1).to(device)
    bins[1:-1] = (k_bin[1:]+k_bin[:-1])/2
    bins[-1] = k_bin[-1]
    bins = bins.unsqueeze(0)
    bins[1:] += 1e-3
    inds = searchsorted(bins, rad.flatten().unsqueeze(0)).squeeze().int()
    # bincount = torch.histc(inds.cpu(), bins=bins.shape[1]+1).to(device)
    bincount = torch.bincount(inds)
    asort = torch.argsort(inds.squeeze())
    sorted_e_ = e_.view(e_.shape[0], -1)[:, asort]
    csum_e_ = torch.cumsum(sorted_e_, dim=1)
    binloc = torch.cumsum(bincount, dim=0).long()-1
    spec_ = csum_e_[:,binloc[1:]] - csum_e_[:,binloc[:-1]]
    spec_ = spec_[:, :-1]
    spec_ = spec_ * 2 * np.pi * (k_bin.float()**2) / bincount[1:-1].float()
    return spec_, k_bin


##################### COMPUTE STATS ########################### 

def word_shuffle(x, l, shuffle_len):
    if not shuffle_len:
        return x

    noise = torch.rand(x.size(), dtype=torch.float).to(x.device)
    pos_idx = torch.arange(x.size(1)).unsqueeze(0).expand_as(x).to(x.device)
    pad_mask = (pos_idx >= l.unsqueeze(1)).float()

    scores = pos_idx.float() + ((1 - pad_mask) * noise + pad_mask) * shuffle_len
    x2 = x.clone()
    x2 = x2.gather(1, scores.argsort(1))

    return x2 

def word_dropout_new(x, l, unk_drop_fac, rand_drop_fac, drop_prob, vocab):
    if not unk_drop_fac and not rand_drop_fac:
        return x

    assert unk_drop_fac + rand_drop_fac <= 1

    batch_size = x.size(0)
    unk_idx = vocab.stoi['<unk>']
    unk_drop_idx = int(batch_size * unk_drop_fac)
    rand_drop_idx = int(batch_size * rand_drop_fac)

    shuffle_idx = torch.argsort(torch.rand(batch_size))
    orignal_idx = torch.argsort(shuffle_idx)

    x2 = x.clone()
    x2 = x2[shuffle_idx]
    
    if unk_drop_idx:
        unk_dropout_(x2[:unk_drop_idx], l[:unk_drop_idx], drop_prob, unk_idx)

    if rand_drop_idx:
        rand_dropout_(x2[-rand_drop_idx:], l[-rand_drop_idx:], drop_prob, len(vocab))

    x2 = x2[orignal_idx]

    return x2 

def recommend(self, u):
        """Return recommended item list given users.

        Args:
            u(torch.LongTensor): tensor stored user indexes. [batch_size,]

        Returns:
            pred(torch.LongTensor): recommended item list sorted by preference. [batch_size, item_size]
        """
        u = self.W[u, :]
        x_ui = torch.mm(u, self.H.t())
        pred = torch.argsort(x_ui, dim=1)
        return pred