text:
"""Cross-modal retrieval evaluation wrapper.

reference code:

- https://github.com/KevinMusgrave/pytorch-metric-learning/blob/0b575b556fe339c2a62043d0ff0efe7fe85107bc/src/pytorch_metric_learning/utils/accuracy_calculator.py#L45
"""
from functools import partial

import numpy as np
import torch
import torch.nn as nn
from tqdm import tqdm

from utils.tensor_utils import to_numpy

def batch(iterable, batch_size=1):
    """a batch generator
    """
    n_items = len(iterable)
    for batch_idx in range(0, n_items, batch_size):
        yield iterable[batch_idx:min(batch_idx + batch_size, n_items)]

def recall_at_k(knn_labels, gt_labels, k):
    curr_knn_labels = knn_labels[:, :k]
    accuracy_per_sample = np.any(curr_knn_labels == gt_labels, axis=1)

    return np.mean(accuracy_per_sample)

def precision_at_k(knn_labels, gt_labels, k):
    curr_knn_labels = knn_labels[:, :k]
    accuracy_per_sample = np.sum(curr_knn_labels == gt_labels, axis=1) / k
    return np.mean(accuracy_per_sample)

def get_relevance_mask(shape, gt_labels, label_counts, embeddings_come_from_same_source):
    relevance_mask = np.zeros(shape=shape, dtype=np.int)
    for k, max_column in label_counts.items():
        matching_rows = np.where(gt_labels == k)[0]
        if embeddings_come_from_same_source:
            max_column = max_column - 1
        relevance_mask[matching_rows, :max_column] = 1
    return relevance_mask

def r_precision(knn_labels, gt_labels, embeddings_come_from_same_source, label_counts):
    relevance_mask = get_relevance_mask(
        knn_labels.shape, gt_labels, label_counts, embeddings_come_from_same_source
    )
    matches_per_row = np.sum(
        (knn_labels == gt_labels) * relevance_mask.astype(bool), axis=1
    )
    max_possible_matches_per_row = np.sum(relevance_mask, axis=1)
    accuracy_per_sample = matches_per_row / max_possible_matches_per_row
    return accuracy_per_sample

class ParallelMatMulModule(nn.Module):
    def set_g_features(self, g_features, g_sigmas=None):
        self._g_features = g_features
        self.g_features = None

    def forward(self, q_features, n_embeddings=1, reduction=None,
                embeddings_come_from_same_source=False,
                q_indices=None, q_sigmas=None):
        if self.g_features is None:
            self.g_features = self._g_features.to(q_features.device)
        sims = q_features.mm(self.g_features)

        if n_embeddings > 1:
            sims = sims.view(int(len(q_features) / n_embeddings),
                             n_embeddings,
                             int(self.g_features.size()[-1] / n_embeddings),
                             n_embeddings)
            sims = sims.permute(0, 1, 3, 2)

            if reduction == 'sum':
                sims = torch.sum(torch.sum(sims, axis=1), axis=1)
            elif reduction == 'max':
                sims = torch.max(torch.max(sims, axis=1)[0], axis=1)[0]

        if embeddings_come_from_same_source:
            # ignore self
            for idx, qidx in enumerate(q_indices):
                sims[idx, qidx] = -np.inf
        sims, pred_ranks = (-sims).sort()
        return sims, pred_ranks

class MatchingProbModule(nn.Module):
    def __init__(self, match_prob_fn):
        super().__init__()
        self.match_prob_fn = match_prob_fn

    def set_g_features(self, g_features, g_sigmas=None):
        self._g_features = g_features
        self.g_features = None

    def forward(self, q_features, n_embeddings=1, reduction=None,
                embeddings_come_from_same_source=False,
                q_indices=None, q_sigmas=None):
        if self.g_features is None:
            self.g_features = self._g_features.to(q_features.device)
        sims = torch.zeros(len(q_features), len(self.g_features))
        for idx, q_feature in enumerate(q_features):
            _sim = self.match_prob_fn(q_feature.unsqueeze(0), self.g_features, None, None)
            sims[idx] = _sim
        sims, pred_ranks = (-sims).sort()
        return sims, pred_ranks

def elk_dist(mu1, mu2, log_sigma1, log_sigma2):
    sum_sigma_sq = torch.exp(2 * log_sigma1) + torch.exp(2 * log_sigma2)
    dist = (mu1 - mu2) ** 2
    dist = dist.squeeze(1)
    elk = dist / (sum_sigma_sq) + torch.log(sum_sigma_sq)
    return -0.5 * torch.sum(elk, dim=1)

def kl_dist(mu1, mu2, log_sigma1, log_sigma2):
    dist = (mu1 - mu2) ** 2
    dist = dist.squeeze(1)

    var2 = torch.exp(log_sigma2 * 2)
    dist = dist / var2 + 2 * (log_sigma1 - log_sigma2) + torch.exp(log_sigma1 * 2) / var2
    return -torch.sum(dist, dim=1)

def reverse_kl_dist(mu1, mu2, log_sigma1, log_sigma2):
    return kl_dist(mu2, mu1, log_sigma2, log_sigma1)

def js_dist(mu1, mu2, log_sigma1, log_sigma2):
    dist = (mu1 - mu2) ** 2
    dist = dist.squeeze(1)

    var1 = torch.exp(log_sigma1 * 2)
    var2 = torch.exp(log_sigma2 * 2)
    dist = (dist + var1) / var2 + (dist + var2) / var1
    return -torch.sum(dist, dim=1)

def bhattacharyya_dist(mu1, mu2, log_sigma1, log_sigma2):
    dist = (mu1 - mu2) ** 2
    dist = dist.squeeze(1)
    sigma1 = torch.exp(log_sigma1)
    sigma2 = torch.exp(log_sigma2)

    dist = dist / (torch.exp(log_sigma1 * 2) + torch.exp(log_sigma2 * 2))
    dist = dist + 2 * torch.log(sigma1 / sigma2 + sigma2 / sigma1)
    ddd = 2 * torch.log(torch.ones(1) * 2)
    dist = dist.float() - ddd.to(dist.device).float()
    dist = dist / 4
    return -torch.sum(dist, dim=1)

def wasserstein_dist(mu1, mu2, log_sigma1, log_sigma2):
    dist = (mu1 - mu2) ** 2
    dist = dist.squeeze(1)
    dist = dist + (torch.exp(log_sigma1) - torch.exp(log_sigma2)) ** 2
    return -torch.sum(dist, dim=1)

def l2_dist(mu1, mu2, log_sigma1, log_sigma2):
    dist = (mu1 - mu2) ** 2
    dist = dist.squeeze(1)
    return -torch.sum(dist, dim=1)

def min_dist(mu1, mu2, log_sigma1, log_sigma2):
    dist = (mu1 - mu2) ** 2
    dist = dist.squeeze(1)
    return -torch.min(dist, dim=1)[0]

def max_dist(mu1, mu2, log_sigma1, log_sigma2):
    dist = (mu1 - mu2) ** 2
    dist = dist.squeeze(1)
    return -torch.max(dist, dim=1)[0]

class HybridDistanceModule(nn.Module):
    def __init__(self, dist_fn_name):
        super().__init__()
        print(dist_fn_name)
        if dist_fn_name == 'elk':
            self.dist_fn = elk_dist
        elif dist_fn_name == 'l2':
            self.dist_fn = l2_dist
        elif dist_fn_name == 'min':
            self.dist_fn = min_dist
        elif dist_fn_name == 'max':
            self.dist_fn = max_dist
        elif dist_fn_name == 'wasserstein':
            self.dist_fn = wasserstein_dist
        elif dist_fn_name == 'kl':
            self.dist_fn = kl_dist
        elif dist_fn_name == 'reverse_kl':
            self.dist_fn = reverse_kl_dist
        elif dist_fn_name == 'js':
            self.dist_fn = js_dist
        elif dist_fn_name == 'bhattacharyya':
            self.dist_fn = bhattacharyya_dist
        else:
            raise ValueError(dist_fn_name)

    def set_g_features(self, g_features, g_sigmas=None):
        self._g_features = g_features
        self.g_features = None
        self._g_sigmas = torch.from_numpy(g_sigmas)
        self.g_sigmas = None

    def forward(self, q_features, n_embeddings=1, reduction=None,
                embeddings_come_from_same_source=False,
                q_indices=None, q_sigmas=None):
        if self.g_features is None:
            self.g_features = self._g_features.to(q_features.device)
            self.g_sigmas = self._g_sigmas.to(q_features.device)

        sims = torch.zeros(len(q_features), len(self.g_features))

        q_sigmas = torch.from_numpy(q_sigmas).to(q_features.device)
        for idx, (q_feature, q_sigma) in enumerate(zip(q_features, q_sigmas)):
            _sim = self.dist_fn(q_feature.unsqueeze(0), self.g_features, q_sigma.unsqueeze(0), self.g_sigmas)
            sims[idx] = _sim

        sims, pred_ranks = (-sims).sort()
        return sims, pred_ranks

class CUBEvaluator(object):
    """Evaluator wrapper

    Args:
        eval_method (str): distance function to use, should be in
            ('elk', 'l2', 'min', 'max', 'wasserstein', 'kl',
            'reverse_kl', 'js', 'bhattacharyya', 'matmul', 'matching_prob')
        n_crossfolds (int): default crossfold setting (-1 | 5)
    """
    def __init__(self,
                 eval_method='matmul',
                 extract_device='cuda',
                 eval_device='cuda',
                 verbose=False):
        self.eval_method = eval_method
        self.extract_device = extract_device
        self.eval_device = eval_device
        self.logger = None

        self.pbar = partial(tqdm, disable=not verbose)

    def set_model(self, model):
        """set model
        """
        self.model = model

        if isinstance(self.model, nn.DataParallel):
            self.n_embeddings = self.model.module.n_embeddings
            self.feat_size = self.model.module.embed_dim
        else:
            self.n_embeddings = self.model.n_embeddings
            self.feat_size = self.model.embed_dim

    def set_criterion(self, criterion):
        """set criterion (for matching prob retrieval)
        """
        self.criterion = criterion

    def set_logger(self, logger):
        """set logger
        """
        self.logger = logger

    @torch.no_grad()
    def extract_features(self, dataloader):
        """Extract image and caption features using the given model.

        Args:
            model (nn.Module): a model to extract features.
            dataloader (data.Dataloader): the target dataloader to feature extraction.
        """
        self.model.eval()
        self.model.to(self.extract_device)

        num_images = dataloader.dataset.n_images
        num_captions = len(dataloader.dataset)
        num_classes = len(dataloader.dataset.class_to_indices)

        image_classes = np.zeros(num_images)
        caption_classes = np.zeros(num_captions)

        image_features = np.zeros((num_images, self.n_embeddings, self.feat_size))
        caption_features = np.zeros((num_captions, self.n_embeddings, self.feat_size))

        image_sigmas = np.zeros((num_images, self.feat_size))
        caption_sigmas = np.zeros((num_captions, self.feat_size))

        print(f'num_images = {num_images}, num_captions = {num_captions}, num_classes = {num_classes}')

        image_paths = []
        raw_captions = []

        cur_image_idx = 0
        cur_caption_idx = 0
        seen_image_ids = set()

        dataset_meta = dataloader.dataset.targets

        for images, captions, caption_lens, class_ids, annotation_ids in self.pbar(dataloader):
            images = images.to(self.extract_device)
            captions = captions.to(self.extract_device)
            caption_lens = caption_lens.to(self.extract_device)

            output = self.model(images, captions, caption_lens)
            _image_features = output['image_features']
            _caption_features = output['caption_features']

            if output.get('image_logsigma') is not None:
                _image_sigmas = output['image_logsigma']
                _caption_sigmas = output['caption_logsigma']

            for idx, (class_id, ann_id) in enumerate(zip(class_ids, annotation_ids)):
                image_id, raw_caption = dataset_meta[ann_id]
                if image_id not in seen_image_ids:
                    image_classes[cur_image_idx] = class_id
                    image_features[cur_image_idx] = to_numpy(_image_features[idx])
                    if output.get('image_logsigma') is not None:
                        image_sigmas[cur_image_idx] = to_numpy(_image_sigmas[idx])
                    cur_image_idx += 1
                    seen_image_ids.add(image_id)
                    image_paths.append(image_id)

                raw_captions.append(raw_caption)
                caption_classes[cur_caption_idx] = class_id
                caption_features[cur_caption_idx] = to_numpy(_caption_features[idx])
                if output.get('image_logsigma') is not None:
                    caption_sigmas[cur_caption_idx] = to_numpy(_caption_sigmas[idx])
                cur_caption_idx += 1

        if cur_image_idx != num_images:
            raise RuntimeError('unexpected error, {} != {}'.format(cur_image_idx, num_images))
        if cur_caption_idx != num_captions:
            raise RuntimeError('unexpected error, {}, {}'.format(cur_caption_idx, num_captions))
        if set(image_classes) != set(caption_classes):
            raise RuntimeError('unexpected error, I({}) != C({})'.format(set(image_classes), set(caption_classes)))

        image_features = torch.from_numpy(image_features)
        caption_features = torch.from_numpy(caption_features)
        image_classes = torch.from_numpy(image_classes)
        caption_classes = torch.from_numpy(caption_classes)

        return {
            'image_features': image_features,
            'caption_features': caption_features,
            'image_sigmas': image_sigmas,
            'caption_sigmas': caption_sigmas,
            'image_ids': image_paths,
            'caption_ids': raw_captions,
            'image_classes': image_classes,
            'caption_classes': caption_classes,
        }

    @torch.no_grad()
    def retrieve(self, q_features, g_features,
                 q_ids, g_ids,
                 q_classes, g_classes,
                 topk=10, reduction='sum', batch_size=256):
        if len(q_features) != len(q_ids):
            raise RuntimeError('length mismatch {}, {}'.format(q_features.shape,
                                                               q_ids.shape))
        if len(g_features) != len(g_ids):
            raise RuntimeError('length mismatch {}, {}'.format(g_features.shape,
                                                               g_ids.shape))
        if isinstance(q_ids, list) or isinstance(g_ids, list):
            q_ids = np.array(q_ids)
            g_ids = np.array(g_ids)
        n_queries = len(q_ids)
        n_galleries = len(g_ids)

        if self.eval_method == 'matmul':
            pmm = ParallelMatMulModule()
            g_features = g_features.view(n_galleries * self.n_embeddings, -1).t()
        elif self.eval_method == 'matching_prob':
            pmm = MatchingProbModule(self.criterion.match_prob)
        pmm.set_g_features(g_features)

        q_features = q_features.to(self.eval_device)

        retrieved_items = {}
        retrieved_scores = {}
        query_meta = {}

        for q_indices in self.pbar(batch(range(n_queries), batch_size=batch_size)):
            q_indices = np.array(q_indices)

            if self.eval_method != 'matching_prob':
                _q_feature = q_features[q_indices, :]
                _q_feature = _q_feature.view(len(q_indices) * self.n_embeddings, -1)
            else:
                _q_feature = q_features[q_indices, :, :]
            sims, pred_ranks = pmm(_q_feature, n_embeddings=self.n_embeddings, reduction=reduction)

            _, sorted_indices = pred_ranks.sort()
            for idx, sorted_db_index in enumerate(sorted_indices):
                _, _sorted_idx = sorted_db_index.sort()
                qidx = q_ids[q_indices[idx]]
                retrieved_items[qidx] = g_ids[to_numpy(_sorted_idx)[0][:topk]]
                retrieved_scores[qidx] = sims[idx][:topk].cpu().numpy()
                query_meta[qidx] = {'class': int(q_classes[q_indices[idx]].item()),
                                    'retrieved_classes': [int(item) for item in
                                                          g_classes[to_numpy(_sorted_idx)[0][:topk]].tolist()]}

        return retrieved_items, retrieved_scores, query_meta

    def get_scores(self, knn_labels, gt_labels, best_pred_ranks,
                   label_counts, embeddings_come_from_same_source,
                   sigmas=None, dismatch_probs=None):
        # https://github.com/KevinMusgrave/pytorch-metric-learning/blob/0b575b556fe339c2a62043d0ff0efe7fe85107bc/src/pytorch_metric_learning/utils/accuracy_calculator.py#L45
        relevance_mask = get_relevance_mask(
            knn_labels.shape, gt_labels, label_counts, embeddings_come_from_same_source
        )
        num_samples, num_k = knn_labels.shape
        equality = (knn_labels == gt_labels) * relevance_mask.astype(bool)
        cumulative_correct = np.cumsum(equality, axis=1)
        k_idx = np.tile(np.arange(1, num_k + 1), (num_samples, 1))
        precision_at_ks = (cumulative_correct * equality) / k_idx
        summed_precision_per_row = np.sum(precision_at_ks * relevance_mask, axis=1)
        max_possible_matches_per_row = np.sum(relevance_mask, axis=1)
        mean_average_precision_at_r = np.mean(summed_precision_per_row / max_possible_matches_per_row)

        pr_per_sample = r_precision(knn_labels, gt_labels, embeddings_come_from_same_source, label_counts)
        precision_r = np.mean(pr_per_sample)

        sigma_vs_map = None
        dismatchprob_vs_map = None

        recall_1 = recall_at_k(knn_labels, gt_labels, 1)
        recall_5 = recall_at_k(knn_labels, gt_labels, 5)
        recall_10 = recall_at_k(knn_labels, gt_labels, 10)

        precision_5 = precision_at_k(knn_labels, gt_labels, 5)
        precision_10 = precision_at_k(knn_labels, gt_labels, 10)
        precision_50 = precision_at_k(knn_labels, gt_labels, 50)
        precision_100 = precision_at_k(knn_labels, gt_labels, 100)

        scores = {
            'recall_1': recall_1,
            'recall_5': recall_5,
            'recall_10': recall_10,
            'rsum': recall_1 + recall_5 + recall_10,
            'precision_5': precision_5,
            'precision_10': precision_10,
            'precision_50': precision_50,
            'precision_100': precision_100,
            'r_precision': precision_r,
            'map_r': mean_average_precision_at_r,
        }

        if sigma_vs_map:
            scores['sigma_vs_map'] = sigma_vs_map
            scores['dismatchprob_vs_map'] = dismatchprob_vs_map

        if best_pred_ranks is not None:
            medr = np.floor(np.median(best_pred_ranks)) + 1
            meanr = np.mean(best_pred_ranks) + 1
            scores['medr'] = medr
            scores['meanr'] = meanr
        return scores

    @torch.no_grad()
    def evaluate_recall(self, q_features, g_features, q_labels, g_labels,
                        q_ids=None, g_ids=None,
                        q_classes=None, g_classes=None,
                        label_counts=None,
                        embeddings_come_from_same_source=False,
                        q_sigmas=None, g_sigmas=None,
                        reduction='sum', batch_size=256
                        ):
        """Evaluate recall

        Args:
            q_features (tensor): N_q x d query features
            g_features (tensor): N_g x d gallery features
            q_labels (tensor): N query labels
            g_labels (tensor): N gallery labels
        """
        if len(q_features) != len(q_labels):
            raise RuntimeError('length mismatch {}, {}'.format(q_features.shape,
                                                               q_labels.shape))
        if len(g_features) != len(g_labels):
            raise RuntimeError('length mismatch {}, {}'.format(g_features.shape,
                                                               g_labels.shape))
        n_queries = len(q_labels)
        n_galleries = len(g_labels)
        best_pred_ranks = np.zeros(n_queries)

        if self.eval_method == 'matmul':
            pmm = ParallelMatMulModule()
            g_features = g_features.view(n_galleries * self.n_embeddings, -1).t()
        elif self.eval_method == 'matching_prob':
            pmm = MatchingProbModule(self.criterion.match_prob)
        else:
            pmm = HybridDistanceModule(self.eval_method)
        pmm.set_g_features(g_features, g_sigmas)

        q_features = q_features.to(self.eval_device)
        sigmas = None
        dismatch_probs = None
        if q_sigmas is not None and len(q_sigmas.shape) == 2:
            sigmas = []
            dismatch_probs = []

        knn_labels = []
        gt_labels = []
        for q_indices in self.pbar(batch(range(n_queries), batch_size=batch_size)):
            q_indices = np.array(q_indices)

            if self.eval_method != 'matching_prob':
                _q_feature = q_features[q_indices, :]
                _q_feature = _q_feature.view(len(q_indices) * self.n_embeddings, -1)
            else:
                _q_feature = q_features[q_indices, :, :]

            if q_sigmas is not None:
                _q_sigmas = q_sigmas[q_indices]
            else:
                _q_sigmas = None
            _, pred_ranks = pmm(_q_feature, n_embeddings=self.n_embeddings,
                                reduction=reduction, embeddings_come_from_same_source=embeddings_come_from_same_source,
                                q_indices=q_indices, q_sigmas=_q_sigmas)

            _, sorted_indices = pred_ranks.sort()
            for idx, sorted_db_index in enumerate(sorted_indices):
                _, _sorted_idx = sorted_db_index.sort()
                q_class = int(q_classes[q_indices[idx]].item())
                r_classes = [int(item) for item in
                             g_classes[to_numpy(_sorted_idx)[0]].tolist()]
                gt_labels.append(q_class)
                knn_labels.append(r_classes)

            for idx, q_idx in enumerate(q_indices):
                pos_indices = np.where(g_labels == q_labels[q_idx])[0]
                _pred_ranks = [torch.where(pred_ranks[idx] == pos_idx)[0][0].item() for pos_idx in pos_indices]
                best_pred_ranks[q_idx] = min(_pred_ranks)

        knn_labels = np.array(knn_labels)
        gt_labels = np.array(gt_labels)
        gt_labels = gt_labels[:, None]

        scores = self.get_scores(knn_labels, gt_labels, best_pred_ranks,
                                 label_counts, embeddings_come_from_same_source,
                                 sigmas, dismatch_probs)

        return scores

    @torch.no_grad()
    def evaluate(self, dataloader, n_crossfolds=None,
                 reduction='sum', eval_batch_size=256,
                 report_sigma=False,
                 key=None):
        """evaluate image-to-caption and caption-to-image retrieval tasks.
        """
        scores = {}

        if self.logger:
            self.logger.log('extracting features...')

        extracted_features = self.extract_features(dataloader)

        image_features = extracted_features['image_features']
        caption_features = extracted_features['caption_features']
        image_sigmas = extracted_features['image_sigmas']
        caption_sigmas = extracted_features['caption_sigmas']
        image_classes = extracted_features['image_classes']
        caption_classes = extracted_features['caption_classes']

        scores['image_sigma'] = np.mean(image_sigmas)
        scores['caption_sigma'] = np.mean(caption_sigmas)

        q_sigmas = None
        g_sigmas = None

        if report_sigma:
            q_sigmas = image_sigmas
            g_sigmas = caption_sigmas

        if self.logger:
            self.logger.log('evaluating i2t...')
        label_counts = {k: len(v) for k, v in dataloader.dataset.class_to_indices.items()}
        scores['i2t'] = self.evaluate_recall(image_features,
                                             caption_features,
                                             image_classes,
                                             caption_classes,
                                             image_classes,
                                             caption_classes,
                                             image_classes,
                                             caption_classes,
                                             label_counts,
                                             reduction=reduction,
                                             q_sigmas=q_sigmas,
                                             g_sigmas=g_sigmas,
                                             batch_size=eval_batch_size)

        if report_sigma:
            q_sigmas = caption_sigmas
            g_sigmas = image_sigmas

        if self.logger:
            self.logger.log('evaluating t2i...')
        label_counts = {k: len(v) for k, v in dataloader.dataset.class_to_img_indices.items()}
        scores['t2i'] = self.evaluate_recall(caption_features,
                                             image_features,
                                             caption_classes,
                                             image_classes,
                                             caption_classes,
                                             image_classes,
                                             caption_classes,
                                             image_classes,
                                             label_counts,
                                             reduction=reduction,
                                             q_sigmas=q_sigmas,
                                             g_sigmas=g_sigmas,
                                             batch_size=eval_batch_size)
        for key in ('rsum', 'medr', 'meanr'):
            scores[key] = scores['i2t'][key] + scores['t2i'][key]
        return scores

max_stars_repo_name: naver-ai/pcme

max_issues_repo_name: naver-ai/pcme

max_forks_repo_name: naver-ai/pcme


text:
import numpy as np
from sklearn.utils import shuffle

def prepareInput(data, groupSize=int(8e4)):
    data = inferWeight(data)
    groups, weights = createGroupsRandom(data, groupSize)
    return groups, weights

# alpha and beta are knowledge in some framework, i.e. RLL.
# we don't include them in NeuCrowd
def inferWeight(data, alpha=None, beta=None):
    votes = data[:, 1]
    maxVote = max(votes)
    weights = []
    for i in range(votes.shape[0]):
        v = votes[i]

        if v >= (1+maxVote)/2:
            if alpha is None or beta is None:
                weights.append(float(v/maxVote))
            else:
                weights.append(float((v+alpha)/(maxVote+alpha+beta)))
        else:
            if alpha is None or beta is None:
                weights.append(1-float(v/maxVote))
            else:
                weights.append(float((maxVote-v+alpha)/(maxVote+alpha+beta)))
    data[:, 1] = weights
    return data

def splitFeatureWeight(x):
    return x[:,1:], x[:,0]

def createGroupsRandom(data, groupSize=int(8e4)):
    positive = data[np.where(data[:, 0] == 1)]
    negative = data[np.where(data[:, 0] == 0)]
    posNum = positive.shape[0]
    negNum = negative.shape[0]

    idx = np.random.randint(low=0, high=posNum, size=groupSize)
    query = np.array([positive[i, 1:] for i in idx])
    posDoc = shuffle(query)

    idx = np.random.randint(low=0, high=negNum, size=groupSize)
    negDoc0 = np.array([negative[i, 1:] for i in idx])
    negDoc1 = shuffle(negDoc0)
    negDoc2 = shuffle(negDoc0)

    query, queryDocW = splitFeatureWeight(query)
    posDoc, posDocW = splitFeatureWeight(posDoc)
    negDoc0, negDoc0W = splitFeatureWeight(negDoc0)
    negDoc1, negDoc1W = splitFeatureWeight(negDoc1)
    negDoc2, negDoc2W = splitFeatureWeight(negDoc2)

    groups = (query, posDoc, negDoc0, negDoc1, negDoc2)
    weights = (queryDocW, posDocW, negDoc0W, negDoc1W, negDoc2W)
    return groups, weights

ext: py

reason: import numpy


text:
#!/usr/bin/env python
# Copyright (c) 2020 - The Procedural Generation for Gazebo authors
# For information on the respective copyright owner see the NOTICE file
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import sys
import unittest
import shutil
import numpy as np
from pcg_gazebo.utils import generate_random_string
from pcg_gazebo.path import Path
from pcg_gazebo.parsers import parse_sdf
from pcg_gazebo.simulation import add_custom_gazebo_resource_path, \
    SimulationModel
from pcg_gazebo.simulation.properties import Heightmap
from pcg_gazebo.generators import HeightmapGenerator
from pcg_gazebo.generators.biomes import Biome, WhittakerBiome

CUR_DIR = os.path.dirname(os.path.abspath(__file__))

class TestHeightmap(unittest.TestCase):
    def test_parse_heightmap_model(self):
        add_custom_gazebo_resource_path(
            os.path.join(CUR_DIR, '..', 'examples', 'models'))

        # Parse the SDF file
        sdf = parse_sdf(
            os.path.join(
                CUR_DIR, '..', 'examples', 'models',
                'pcg_winding_valley_heightmap', 'model.sdf'))
        self.assertIsNotNone(sdf)
        self.assertIsNotNone(sdf.models)
        self.assertEqual(len(sdf.models[0].links[0].visuals), 1)
        self.assertEqual(len(sdf.models[0].links[0].collisions), 1)

        sdf_heightmap = sdf.models[0].links[0].visuals[0].geometry.heightmap
        self.assertIsNotNone(sdf_heightmap)

        uri_path_prefix = \
            'model://pcg_winding_valley_heightmap/materials/textures'
        self.assertIn(
            uri_path_prefix,
            sdf_heightmap.uri.value)
        self.assertEqual(len(sdf_heightmap.textures), 3)
        for texture in sdf_heightmap.textures:
            self.assertIn(uri_path_prefix, texture.diffuse.value)
            self.assertIn(uri_path_prefix, texture.normal.value)
            self.assertEqual(texture.size.value, [10])
        self.assertEqual(len(sdf_heightmap.blends), 2)
        for blend in sdf_heightmap.blends:
            self.assertGreater(blend.min_height.value, 0)
            self.assertGreater(blend.fade_dist.value, 0)
        self.assertEqual(sdf_heightmap.size.value, [1000, 1000, 25])
        self.assertEqual(sdf_heightmap.pos.value, [0, 0, -4])

        sdf_heightmap = \
            sdf.models[0].links[0].collisions[0].geometry.heightmap
        self.assertIsNotNone(sdf_heightmap)

        self.assertIn(
            uri_path_prefix,
            sdf_heightmap.uri.value)
        self.assertEqual(sdf_heightmap.size.value, [1000, 1000, 25])
        self.assertEqual(sdf_heightmap.pos.value, [0, 0, -4])

    def test_from_sdf(self):
        add_custom_gazebo_resource_path(
            os.path.join(CUR_DIR, '..', 'examples', 'models'))

        # Parse the SDF file
        sdf = parse_sdf(
            os.path.join(
                CUR_DIR, '..', 'examples', 'models',
                'pcg_winding_valley_heightmap', 'model.sdf'))
        self.assertIsNotNone(sdf)

        # Convert SDF heightmap element to simulation heightmap element
        sdf_heightmap = sdf.models[0].links[0].visuals[0].geometry.heightmap
        heightmap = Heightmap.from_sdf(sdf_heightmap)
        self.assertIsNotNone(heightmap)

        output_sdf = heightmap.to_sdf()
        self.assertIsNotNone(output_sdf)

        sdf_heightmap = \
            sdf.models[0].links[0].collisions[0].geometry.heightmap
        heightmap = Heightmap.from_sdf(sdf_heightmap)
        self.assertIsNotNone(heightmap)

        output_sdf = heightmap.to_sdf()
        self.assertIsNotNone(output_sdf)

        self.assertEqual(sdf_heightmap.size.value, output_sdf.size.value)
        self.assertEqual(sdf_heightmap.pos.value, output_sdf.pos.value)

        # Convert model to SimulationModel element
        model = SimulationModel.from_sdf(sdf.models[0])
        self.assertIsNotNone(model)
        output_sdf = model.to_sdf()
        self.assertIsNotNone(output_sdf)

        self.assertIsNotNone(output_sdf)
        self.assertEqual(len(output_sdf.links[0].visuals), 1)
        self.assertEqual(len(output_sdf.links[0].collisions), 1)

        sdf_heightmap = output_sdf.links[0].visuals[0].geometry.heightmap
        self.assertIsNotNone(sdf_heightmap)

        uri_path_prefix = \
            'model://pcg_winding_valley_heightmap/materials/textures'
        self.assertIn(
            uri_path_prefix,
            sdf_heightmap.uri.value)
        self.assertEqual(len(sdf_heightmap.textures), 3)
        for texture in sdf_heightmap.textures:
            self.assertIn(uri_path_prefix, texture.diffuse.value)
            self.assertIn(uri_path_prefix, texture.normal.value)
            self.assertEqual(texture.size.value, [10])
        self.assertEqual(len(sdf_heightmap.blends), 2)
        for blend in sdf_heightmap.blends:
            self.assertGreater(blend.min_height.value, 0)
            self.assertGreater(blend.fade_dist.value, 0)
        self.assertEqual(sdf_heightmap.size.value, [1000, 1000, 25])
        self.assertEqual(sdf_heightmap.pos.value, [0, 0, -4])

        sdf_heightmap = \
            output_sdf.links[0].collisions[0].geometry.heightmap
        self.assertIsNotNone(sdf_heightmap)

        self.assertIn(
            uri_path_prefix,
            sdf_heightmap.uri.value)
        self.assertEqual(sdf_heightmap.size.value, [1000, 1000, 25])
        self.assertEqual(sdf_heightmap.pos.value, [0, 0, -4])

    def test_init_heightmap_generator_from_sdf(self):
        add_custom_gazebo_resource_path(
            os.path.join(CUR_DIR, '..', 'examples', 'models'))

        # Parse the SDF file
        sdf = parse_sdf(
            os.path.join(
                CUR_DIR, '..', 'examples', 'models',
                'pcg_winding_valley_heightmap', 'model.sdf'))
        self.assertIsNotNone(sdf)

        sdf_heightmap = sdf.models[0].links[0].visuals[0].geometry.heightmap
        hg = HeightmapGenerator.from_sdf(sdf_heightmap)
        self.assertIsNotNone(hg)
        self.assertIsNotNone(hg.biome)

        self.assertEqual(hg.biome.n_elevation_zones, 3)
        self.assertEqual(hg.biome.n_moisture_zones, 1)

        for i in range(hg.biome.n_elevation_zones):
            tag = hg.biome.get_biome(i, 0)
            self.assertIsNotNone(tag)
            self.assertIsNotNone(hg.biome.get_diffuse(tag))
            self.assertIsNotNone(hg.biome.get_normal(tag))

    def test_heightmap_generator_to_sdf(self):
        hg = HeightmapGenerator()
        model = hg.as_model()

        self.assertIsNotNone(model)
        self.assertEqual(model.n_links, 1)
        self.assertIn('link', model.links)

        self.assertEqual(len(model.links['link'].visuals), 1)
        self.assertEqual(
            model.links['link'].visuals[0].geometry.geo_type,
            'heightmap')
        self.assertEqual(len(model.links['link'].collisions), 1)
        self.assertEqual(
            model.links['link'].collisions[0].geometry.geo_type,
            'heightmap')

        model_sdf = model.to_sdf()
        self.assertIsNotNone(model_sdf)
        heightmap = model_sdf.links[0].visuals[0].geometry.heightmap
        self.assertIsNotNone(heightmap.uri)
        self.assertIn('.pcg', heightmap.uri.value)
        self.assertIn('file://', heightmap.uri.value)
        self.assertIn('/materials/textures/', heightmap.uri.value)

        image_path = Path(heightmap.uri.value)
        self.assertTrue(os.path.isfile(image_path.absolute_uri))

        # Clean up files
        os.remove(image_path.absolute_uri)

    def test_heightmap_generator_to_gazebo_model(self):
        hg = HeightmapGenerator()
        hg.biome = WhittakerBiome()
        model_name = generate_random_string(10)
        output_dir = os.path.join(
            os.path.expanduser('~'), '.gazebo', 'models')
        if not os.path.isdir(output_dir):
            os.makedirs(output_dir)

        model = hg.as_model()
        model.name = model_name

        model.to_gazebo_model(
            output_dir=output_dir,
            copy_resources=True)

        self.assertTrue(
            os.path.isdir(os.path.join(output_dir, model_name)))

        sdf = parse_sdf(os.path.join(output_dir, model_name, 'model.sdf'))
        self.assertIsNotNone(sdf)
        self.assertIsNotNone(sdf.models[0])

        heightmap = \
            sdf.models[0].links[0].visuals[0].geometry.heightmap
        self.assertIsNotNone(heightmap.uri)
        self.assertIn(
            'model://{}'.format(model_name), heightmap.uri.value)
        self.assertIn('/materials/textures/', heightmap.uri.value)
        self.assertTrue(heightmap.uri.value.endswith('.png'))

        heightmap = \
            sdf.models[0].links[0].collisions[0].geometry.heightmap
        self.assertIsNotNone(heightmap.uri)
        self.assertIn(
            'model://{}'.format(model_name), heightmap.uri.value)
        self.assertIn('/materials/textures/', heightmap.uri.value)
        self.assertTrue(heightmap.uri.value.endswith('.png'))

        shutil.rmtree(os.path.join(output_dir, model_name))

    def test_store_perlin_noise_map(self):
        if sys.version_info[0] < 3:
            return
        hg = HeightmapGenerator(image_size=[5, 5])
        hg.add_perlin_noise_layer()
        ref_image = hg.heightmap_image
        self.assertIsNotNone(ref_image)
        self.assertEqual(list(ref_image.shape), hg.image_size)
        self.assertGreater(np.nonzero(ref_image != 0)[0].size, 0)

        model = hg.as_model()

        model_sdf = model.to_sdf()
        self.assertIsNotNone(model_sdf)
        heightmap = model_sdf.links[0].visuals[0].geometry.heightmap
        self.assertIsNotNone(heightmap.uri)
        self.assertIn('.pcg', heightmap.uri.value)
        self.assertIn('file://', heightmap.uri.value)
        self.assertIn('/materials/textures/', heightmap.uri.value)

        image_path = Path(heightmap.uri.value)
        self.assertTrue(os.path.isfile(image_path.absolute_uri))

        # Clean up files
        os.remove(image_path.absolute_uri)

    def test_store_simplex_noise_map(self):
        if sys.version_info[0] < 3:
            return
        hg = HeightmapGenerator(image_size=[5, 5])
        hg.add_simplex_noise_layer()
        ref_image = hg.heightmap_image
        self.assertIsNotNone(ref_image)
        self.assertEqual(list(ref_image.shape), hg.image_size)
        self.assertGreater(np.nonzero(ref_image != 0)[0].size, 0)

        model = hg.as_model()

        model_sdf = model.to_sdf()
        self.assertIsNotNone(model_sdf)
        heightmap = model_sdf.links[0].visuals[0].geometry.heightmap
        self.assertIsNotNone(heightmap.uri)
        self.assertIn('.pcg', heightmap.uri.value)
        self.assertIn('file://', heightmap.uri.value)
        self.assertIn('/materials/textures/', heightmap.uri.value)

        image_path = Path(heightmap.uri.value)
        self.assertTrue(os.path.isfile(image_path.absolute_uri))

        # Clean up files
        os.remove(image_path.absolute_uri)

    def test_store_random_noise_map(self):
        if sys.version_info[0] < 3:
            return
        hg = HeightmapGenerator(image_size=[5, 5])
        hg.add_random_layer()
        ref_image = hg.heightmap_image
        self.assertIsNotNone(ref_image)
        self.assertEqual(list(ref_image.shape), hg.image_size)
        self.assertGreater(np.nonzero(ref_image != 0)[0].size, 0)

        model = hg.as_model()

        model_sdf = model.to_sdf()
        self.assertIsNotNone(model_sdf)
        heightmap = model_sdf.links[0].visuals[0].geometry.heightmap
        self.assertIsNotNone(heightmap.uri)
        self.assertIn('.pcg', heightmap.uri.value)
        self.assertIn('file://', heightmap.uri.value)
        self.assertIn('/materials/textures/', heightmap.uri.value)

        image_path = Path(heightmap.uri.value)
        self.assertTrue(os.path.isfile(image_path.absolute_uri))

        # Clean up files
        os.remove(image_path.absolute_uri)

    def test_biome(self):
        biome = Biome(n_moisture_zones=4, n_elevation_zones=3)
        self.assertTrue(biome.n_elevation_zones, 3)
        self.assertTrue(biome.n_moisture_zones, 4)

    def test_parse_biome_from_heightmap(self):
        add_custom_gazebo_resource_path(
            os.path.join(CUR_DIR, '..', 'examples', 'models'))

        # Parse the SDF file
        sdf = parse_sdf(
            os.path.join(
                CUR_DIR, '..', 'examples', 'models',
                'pcg_winding_valley_heightmap', 'model.sdf'))
        self.assertIsNotNone(sdf)

        sdf_heightmap = sdf.models[0].links[0].visuals[0].geometry.heightmap
        biome = Biome.from_sdf(sdf_heightmap)

        self.assertIsNotNone(biome)
        self.assertEqual(biome.n_elevation_zones, 3)
        self.assertEqual(biome.n_moisture_zones, 1)

        for i in range(biome.n_elevation_zones):
            tag = biome.get_biome(i, 0)
            self.assertIsNotNone(tag)
            self.assertIsNotNone(biome.get_diffuse(tag))
            self.assertIsNotNone(biome.get_normal(tag))

    def test_whittaker_biome(self):
        biome = WhittakerBiome()
        self.assertTrue(biome.n_elevation_zones, 4)
        self.assertTrue(biome.n_moisture_zones, 6)

        for e in range(biome.n_elevation_zones):
            for m in range(biome.n_moisture_zones):
                paths = biome.save_image(
                    moisture_zone=m, elevation_zone=e)
                self.assertTrue(os.path.isfile(paths['diffuse']))
                self.assertTrue(os.path.isfile(paths['normal']))

                diffuse = os.path.basename(paths['diffuse'])
                self.assertTrue(diffuse.endswith('.png'))
                os.remove(paths['diffuse'])

                normal = os.path.basename(paths['normal'])
                self.assertTrue(normal.endswith('.png'))
                os.remove(paths['normal'])

if __name__ == '__main__':
    unittest.main()


text:
# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import absolute_import, division, print_function, unicode_literals
import pytest
import numpy as np
from numpy.testing import assert_allclose
from astropy.io import fits
from ..geom import MapAxis, MapCoord
from ..hpx import HpxGeom, get_pix_size_from_nside, nside_to_order
from ..hpx import make_hpx_to_wcs_mapping, unravel_hpx_index, ravel_hpx_index
from ..hpx import get_hpxregion_dir, get_hpxregion_size, get_subpixels, get_superpixels

pytest.importorskip("scipy")
pytest.importorskip("healpy")
pytest.importorskip("numpy", "1.13.0")

hpx_allsky_test_geoms = [
    # 2D All-sky
    (8, False, "GAL", None, None),
    # 3D All-sky
    (8, False, "GAL", None, [MapAxis(np.logspace(0., 3., 4))]),
    # 3D All-sky w/ variable pixel size
    ([2, 4, 8], False, "GAL", None, [MapAxis(np.logspace(0., 3., 4))]),
    # 4D All-sky
    (
        8,
        False,
        "GAL",
        None,
        [
            MapAxis(np.logspace(0., 3., 3), name="axis0"),
            MapAxis(np.logspace(0., 2., 4), name="axis1"),
        ],
    ),
]

hpx_partialsky_test_geoms = [
    # 2D Partial-sky
    (8, False, "GAL", "DISK(110.,75.,10.)", None),
    # 3D Partial-sky
    (8, False, "GAL", "DISK(110.,75.,10.)", [MapAxis(np.logspace(0., 3., 4))]),
    # 3D Partial-sky w/ variable pixel size
    (
        [8, 16, 32],
        False,
        "GAL",
        "DISK(110.,75.,10.)",
        [MapAxis(np.logspace(0., 3., 4))],
    ),
    # 4D Partial-sky w/ variable pixel size
    (
        [[8, 16, 32], [8, 8, 16]],
        False,
        "GAL",
        "DISK(110.,75.,10.)",
        [
            MapAxis(np.logspace(0., 3., 3), name="axis0"),
            MapAxis(np.logspace(0., 2., 4), name="axis1"),
        ],
    ),
]

hpx_test_geoms = hpx_allsky_test_geoms + hpx_partialsky_test_geoms

def make_test_coords(geom, lon, lat):
    coords = [lon, lat] + [ax.center for ax in geom.axes]
    coords = np.meshgrid(*coords)
    coords = tuple([np.ravel(t) for t in coords])
    return MapCoord.create(coords)

def test_unravel_hpx_index():
    npix = np.array([2, 7])
    assert_allclose(
        unravel_hpx_index(np.array([0, 4]), npix), (np.array([0, 2]), np.array([0, 1]))
    )
    npix = np.array([[2, 7], [3, 1]])
    assert_allclose(
        unravel_hpx_index(np.array([0, 3, 10]), npix),
        (np.array([0, 1, 1]), np.array([0, 0, 1]), np.array([0, 1, 0])),
    )

def test_ravel_hpx_index():
    npix = np.array([2, 7])
    idx = (np.array([0, 2]), np.array([0, 1]))
    assert_allclose(ravel_hpx_index(idx, npix), np.array([0, 4]))
    npix = np.array([[2, 7], [3, 1]])
    idx = (np.array([0, 1, 1]), np.array([0, 0, 1]), np.array([0, 1, 0]))
    assert_allclose(ravel_hpx_index(idx, npix), np.array([0, 3, 10]))

def make_test_nside(nside, nside0, nside1):
    npix = 12 * nside ** 2
    nside_test = np.concatenate(
        (nside0 * np.ones(npix // 2, dtype=int), nside1 * np.ones(npix // 2, dtype=int))
    )
    return nside_test

@pytest.mark.parametrize(
    ("nside_subpix", "nside_superpix", "nest"),
    [
        (4, 2, True),
        (8, 2, True),
        (8, make_test_nside(8, 4, 2), True),
        (4, 2, False),
        (8, 2, False),
        (8, make_test_nside(8, 4, 2), False),
    ],
)
def test_get_superpixels(nside_subpix, nside_superpix, nest):
    import healpy as hp

    npix = 12 * nside_subpix ** 2
    subpix = np.arange(npix)
    ang_subpix = hp.pix2ang(nside_subpix, subpix, nest=nest)
    superpix = get_superpixels(subpix, nside_subpix, nside_superpix, nest=nest)
    pix1 = hp.ang2pix(nside_superpix, *ang_subpix, nest=nest)
    assert_allclose(superpix, pix1)

    subpix = subpix.reshape((12, -1))
    if not np.isscalar(nside_subpix):
        nside_subpix = nside_subpix.reshape((12, -1))
    if not np.isscalar(nside_superpix):
        nside_superpix = nside_superpix.reshape((12, -1))

    ang_subpix = hp.pix2ang(nside_subpix, subpix, nest=nest)
    superpix = get_superpixels(subpix, nside_subpix, nside_superpix, nest=nest)
    pix1 = hp.ang2pix(nside_superpix, *ang_subpix, nest=nest)
    assert_allclose(superpix, pix1)

@pytest.mark.parametrize(
    ("nside_superpix", "nside_subpix", "nest"),
    [(2, 4, True), (2, 8, True), (2, 4, False), (2, 8, False)],
)
def test_get_subpixels(nside_superpix, nside_subpix, nest):
    import healpy as hp

    npix = 12 * nside_superpix ** 2
    superpix = np.arange(npix)
    subpix = get_subpixels(superpix, nside_superpix, nside_subpix, nest=nest)
    ang1 = hp.pix2ang(nside_subpix, subpix, nest=nest)
    pix1 = hp.ang2pix(nside_superpix, *ang1, nest=nest)
    assert np.all(superpix[..., None] == pix1)

    superpix = superpix.reshape((12, -1))
    subpix = get_subpixels(superpix, nside_superpix, nside_subpix, nest=nest)
    ang1 = hp.pix2ang(nside_subpix, subpix, nest=nest)
    pix1 = hp.ang2pix(nside_superpix, *ang1, nest=nest)
    assert np.all(superpix[..., None] == pix1)

    pix1 = get_superpixels(subpix, nside_subpix, nside_superpix, nest=nest)
    assert np.all(superpix[..., None] == pix1)

def test_hpx_global_to_local():
    ax0 = np.linspace(0., 1., 3)
    ax1 = np.linspace(0., 1., 3)

    # 2D All-sky
    hpx = HpxGeom(16, False, "GAL")
    assert_allclose(hpx[0], np.array([0]))
    assert_allclose(hpx[633], np.array([633]))
    assert_allclose(hpx[0, 633], np.array([0, 633]))
    assert_allclose(hpx[np.array([0, 633])], np.array([0, 633]))

    # 3D All-sky
    hpx = HpxGeom(16, False, "GAL", axes=[ax0])
    assert_allclose(
        hpx[(np.array([177, 177]), np.array([0, 1]))], np.array([177, 177 + 3072])
    )

    # 2D Partial-sky
    hpx = HpxGeom(64, False, "GAL", region="DISK(110.,75.,2.)")
    assert_allclose(hpx[0, 633, 706], np.array([-1, 0, 2]))

    # 3D Partial-sky
    hpx = HpxGeom(64, False, "GAL", region="DISK(110.,75.,2.)", axes=[ax0])
    assert_allclose(hpx[633], np.array([0]))
    assert_allclose(hpx[49859], np.array([19]))
    assert_allclose(hpx[0, 633, 706, 49859, 49935], np.array([-1, 0, 2, 19, 21]))
    assert_allclose(
        hpx[np.array([0, 633, 706, 49859, 49935])], np.array([-1, 0, 2, 19, 21])
    )
    assert_allclose(
        hpx[(np.array([0, 633, 706, 707, 783]), np.array([0, 0, 0, 1, 1]))],
        np.array([-1, 0, 2, 19, 21]),
    )

    # 3D Partial-sky w/ variable bin size
    hpx = HpxGeom([32, 64], False, "GAL", region="DISK(110.,75.,2.)", axes=[ax0])

    assert_allclose(hpx[191], np.array([0]))
    assert_allclose(hpx[12995], np.array([6]))
    assert_allclose(hpx[0, 191, 233, 12995], np.array([-1, 0, 2, 6]))
    assert_allclose(
        hpx[(np.array([0, 191, 233, 707]), np.array([0, 0, 0, 1]))],
        np.array([-1, 0, 2, 6]),
    )

    # 4D Partial-sky w/ variable bin size
    hpx = HpxGeom(
        [[16, 32], [32, 64]], False, "GAL", region="DISK(110.,75.,2.)", axes=[ax0, ax1]
    )
    assert_allclose(hpx[3263], np.array([1]))
    assert_allclose(hpx[28356], np.array([11]))
    assert_allclose(hpx[(np.array([46]), np.array([0]), np.array([0]))], np.array([0]))

@pytest.mark.parametrize(
    ("nside", "nested", "coordsys", "region", "axes"), hpx_allsky_test_geoms
)
def test_hpxgeom_init_with_pix(nside, nested, coordsys, region, axes):
    geom = HpxGeom(nside, nested, coordsys, region=region, axes=axes)

    idx0 = geom.get_idx(flat=True)
    idx1 = tuple([t[::10] for t in idx0])
    geom = HpxGeom(nside, nested, coordsys, region=idx0, axes=axes)
    assert_allclose(idx0, geom.get_idx(flat=True))
    assert_allclose(len(idx0[0]), np.sum(geom.npix))
    geom = HpxGeom(nside, nested, coordsys, region=idx1, axes=axes)
    assert_allclose(idx1, geom.get_idx(flat=True))
    assert_allclose(len(idx1[0]), np.sum(geom.npix))

@pytest.mark.parametrize(
    ("nside", "nested", "coordsys", "region", "axes"), hpx_test_geoms
)
def test_hpxgeom_to_slice(nside, nested, coordsys, region, axes):
    geom = HpxGeom(nside, nested, coordsys, region=region, axes=axes)
    slices = tuple([slice(1, 2) for i in range(2, geom.ndim)])
    geom_slice = geom.to_slice(slices)
    assert_allclose(geom_slice.ndim, 2)
    assert_allclose(geom_slice.npix, np.squeeze(geom.npix[slices]))

    idx = geom.get_idx(flat=True)
    idx_slice = geom_slice.get_idx(flat=True)
    if geom.ndim > 2:
        m = np.all([np.in1d(t, [1]) for t in idx[1:]], axis=0)
        assert_allclose(idx_slice, (idx[0][m],))
    else:
        assert_allclose(idx_slice, idx)

    # Test slicing with explicit geometry
    geom = HpxGeom(
        nside, nested, coordsys, region=tuple([t[::3] for t in idx]), axes=axes
    )
    geom_slice = geom.to_slice(slices)
    assert_allclose(geom_slice.ndim, 2)
    assert_allclose(geom_slice.npix, np.squeeze(geom.npix[slices]))

    idx = geom.get_idx()
    idx_slice = geom_slice.get_idx()
    if geom.ndim > 2:
        m = np.all([np.in1d(t, [1]) for t in idx[1:]], axis=0)
        assert_allclose(idx_slice, (idx[0].flat[m],))
    else:
        assert_allclose(idx_slice, idx)

@pytest.mark.parametrize(
    ("nside", "nested", "coordsys", "region", "axes"), hpx_test_geoms
)
def test_hpxgeom_get_pix(nside, nested, coordsys, region, axes):
    geom = HpxGeom(nside, nested, coordsys, region=region, axes=axes)
    idx = geom.get_idx(local=False, flat=True)
    idx_local = geom.get_idx(local=True, flat=True)
    assert_allclose(idx, geom.local_to_global(idx_local))

    if axes is not None:
        idx_img = geom.get_idx(local=False, idx=tuple([1] * len(axes)), flat=True)
        idx_img_local = geom.get_idx(local=True, idx=tuple([1] * len(axes)), flat=True)
        assert_allclose(idx_img, geom.local_to_global(idx_img_local))

@pytest.mark.parametrize(
    ("nside", "nested", "coordsys", "region", "axes"), hpx_test_geoms
)
def test_hpxgeom_coord_to_idx(nside, nested, coordsys, region, axes):
    import healpy as hp

    geom = HpxGeom(nside, nested, coordsys, region=region, axes=axes)
    lon = np.array([112.5, 135., 105.])
    lat = np.array([75.3, 75.3, 74.6])
    coords = make_test_coords(geom, lon, lat)
    zidx = tuple([ax.coord_to_idx(t) for t, ax in zip(coords[2:], geom.axes)])

    if geom.nside.size > 1:
        nside = geom.nside[zidx]
    else:
        nside = geom.nside

    phi, theta = coords.phi, coords.theta
    idx = geom.coord_to_idx(coords)
    assert_allclose(hp.ang2pix(nside, theta, phi), idx[0])
    for i, z in enumerate(zidx):
        assert_allclose(z, idx[i + 1])

    # Test w/ coords outside the geometry
    lon = np.array([0.0, 5.0, 10.0])
    lat = np.array([75.3, 75.3, 74.6])
    coords = make_test_coords(geom, lon, lat)
    zidx = [ax.coord_to_idx(t) for t, ax in zip(coords[2:], geom.axes)]

    idx = geom.coord_to_idx(coords)
    if geom.region is not None:
        assert_allclose(np.full_like(coords[0], -1, dtype=int), idx[0])

    idx = geom.coord_to_idx(coords, clip=True)
    assert np.all(np.not_equal(np.full_like(coords[0], -1, dtype=int), idx[0]))

def test_hpxgeom_coord_to_pix():
    lon = np.array([110.25, 114., 105.])
    lat = np.array([75.3, 75.3, 74.6])
    z0 = np.array([0.5, 1.5, 2.5])
    z1 = np.array([3.5, 4.5, 5.5])
    ax0 = np.linspace(0., 3., 4)
    ax1 = np.linspace(3., 6., 4)

    pix64 = np.array([784, 785, 864])

    # 2D all-sky
    coords = (lon, lat)
    hpx = HpxGeom(64, False, "GAL")
    assert_allclose(hpx.coord_to_pix(coords)[0], pix64)

    # 2D partial-sky
    coords = (lon, lat)
    hpx = HpxGeom(64, False, "GAL", region="DISK(110.,75.,2.)")
    assert_allclose(hpx.coord_to_pix(coords)[0], pix64)

    # 3D partial-sky
    coords = (lon, lat, z0)
    hpx = HpxGeom(64, False, "GAL", region="DISK(110.,75.,2.)", axes=[ax0])
    assert_allclose(hpx.coord_to_pix(coords), (pix64, np.array([0, 1, 2])))

    # 3D partial-sky w/ variable bin size
    coords = (lon, lat, z0)
    nside = [16, 32, 64]
    hpx_bins = [HpxGeom(n, False, "GAL", region="DISK(110.,75.,2.)") for n in nside]
    hpx = HpxGeom(nside, False, "GAL", region="DISK(110.,75.,2.)", axes=[ax0])
    for i, (x, y, z) in enumerate(np.vstack(coords).T):
        pix0 = hpx.coord_to_pix((np.array([x]), np.array([y]), np.array([z])))
        pix1 = hpx_bins[i].coord_to_pix((np.array([x]), np.array([y])))
        assert_allclose(pix0[0], pix1[0])

    # 4D partial-sky
    coords = (lon, lat, z0, z1)
    hpx = HpxGeom(64, False, "GAL", region="DISK(110.,75.,2.)", axes=[ax0, ax1])
    assert_allclose(
        hpx.coord_to_pix(coords), (pix64, np.array([0, 1, 2]), np.array([0, 1, 2]))
    )

def test_hpx_nside_to_order():
    assert_allclose(nside_to_order(64), np.array([6]))
    assert_allclose(
        nside_to_order(np.array([10, 32, 42, 64, 128, 256])),
        np.array([-1, 5, -1, 6, 7, 8]),
    )

    order = np.linspace(1, 10, 10).astype(int)
    nside = 2 ** order
    assert_allclose(nside_to_order(nside), order)
    assert_allclose(nside_to_order(nside).reshape((2, 5)), order.reshape((2, 5)))

def test_hpx_get_pix_size_from_nside():
    assert_allclose(
        get_pix_size_from_nside(np.array([1, 2, 4])), np.array([32.0, 16.0, 8.0])
    )

def test_hpx_get_hpxregion_size():
    assert_allclose(get_hpxregion_size("DISK(110.,75.,2.)"), 2.0)

def test_hpxgeom_get_hpxregion_dir():
    refdir = get_hpxregion_dir("DISK(110.,75.,2.)", "GAL")
    assert_allclose(refdir.l.deg, 110.)
    assert_allclose(refdir.b.deg, 75.)

    refdir = get_hpxregion_dir(None, "GAL")
    assert_allclose(refdir.l.deg, 0.)
    assert_allclose(refdir.b.deg, 0.)

def test_hpxgeom_make_wcs():
    ax0 = np.linspace(0., 3., 4)

    hpx = HpxGeom(64, False, "GAL", region="DISK(110.,75.,2.)")
    wcs = hpx.make_wcs()
    assert_allclose(wcs.wcs.wcs.crval, np.array([110., 75.]))

    hpx = HpxGeom(64, False, "GAL", region="DISK(110.,75.,2.)", axes=[ax0])
    wcs = hpx.make_wcs()
    assert_allclose(wcs.wcs.wcs.crval, np.array([110., 75.]))

def test_hpxgeom_get_coord():
    ax0 = np.linspace(0., 3., 4)

    # 2D all-sky
    hpx = HpxGeom(16, False, "GAL")
    c = hpx.get_coord()
    assert_allclose(c[0][:3], np.array([45., 135., 225.]))
    assert_allclose(c[1][:3], np.array([87.075819, 87.075819, 87.075819]))

    # 3D all-sky
    hpx = HpxGeom(16, False, "GAL", axes=[ax0])
    c = hpx.get_coord()
    assert_allclose(c[0][0, :3], np.array([45., 135., 225.]))
    assert_allclose(c[1][0, :3], np.array([87.075819, 87.075819, 87.075819]))
    assert_allclose(c[2][0, :3], np.array([0.5, 0.5, 0.5]))

    # 2D partial-sky
    hpx = HpxGeom(64, False, "GAL", region="DISK(110.,75.,2.)")
    c = hpx.get_coord()
    assert_allclose(c[0][:3], np.array([107.5, 112.5, 106.57894737]))
    assert_allclose(c[1][:3], np.array([76.813533, 76.813533, 76.07742]))

    # 3D partial-sky
    hpx = HpxGeom(64, False, "GAL", region="DISK(110.,75.,2.)", axes=[ax0])
    c = hpx.get_coord()
    assert_allclose(c[0][0, :3], np.array([107.5, 112.5, 106.57894737]))
    assert_allclose(c[1][0, :3], np.array([76.813533, 76.813533, 76.07742]))
    assert_allclose(c[2][0, :3], np.array([0.5, 0.5, 0.5]))

    # 3D partial-sky w/ variable bin size
    hpx = HpxGeom([16, 32, 64], False, "GAL", region="DISK(110.,75.,2.)", axes=[ax0])
    c = hpx.get_coord(flat=True)
    assert_allclose(c[0][:3], np.array([117., 103.5, 112.5]))
    assert_allclose(c[1][:3], np.array([75.340734, 75.340734, 75.340734]))
    assert_allclose(c[2][:3], np.array([0.5, 1.5, 1.5]))

@pytest.mark.parametrize(
    ("nside", "nested", "coordsys", "region", "axes"), hpx_test_geoms
)
def test_hpxgeom_contains(nside, nested, coordsys, region, axes):
    geom = HpxGeom(nside, nested, coordsys, region=region, axes=axes)
    coords = geom.get_coord(flat=True)
    assert_allclose(geom.contains(coords), np.ones_like(coords[0], dtype=bool))

    if axes is not None:
        coords = [c[0] for c in coords[:2]] + [ax.edges[-1] + 1.0 for ax in axes]
        assert_allclose(geom.contains(coords), np.zeros((1,), dtype=bool))

    if geom.region is not None:
        coords = [0.0, 0.0] + [ax.center[0] for ax in geom.axes]
        assert_allclose(geom.contains(coords), np.zeros((1,), dtype=bool))

def test_make_hpx_to_wcs_mapping():
    ax0 = np.linspace(0., 1., 3)
    hpx = HpxGeom(16, False, "GAL", region="DISK(110.,75.,2.)")
    # FIXME construct explicit WCS projection here
    wcs = hpx.make_wcs()
    hpx2wcs = make_hpx_to_wcs_mapping(hpx, wcs)
    assert_allclose(
        hpx2wcs[0],
        np.array(
            [
                67,
                46,
                46,
                46,
                46,
                29,
                67,
                67,
                46,
                46,
                46,
                46,
                67,
                67,
                67,
                46,
                46,
                46,
                67,
                67,
                67,
                28,
                28,
                28,
                45,
                45,
                45,
                45,
                28,
                28,
                66,
                45,
                45,
                45,
                45,
                28,
            ]
        ),
    )
    assert_allclose(
        hpx2wcs[1],
        np.array(
            [
                0.11111111,
                0.09090909,
                0.09090909,
                0.09090909,
                0.09090909,
                1.,
                0.11111111,
                0.11111111,
                0.09090909,
                0.09090909,
                0.09090909,
                0.09090909,
                0.11111111,
                0.11111111,
                0.11111111,
                0.09090909,
                0.09090909,
                0.09090909,
                0.11111111,
                0.11111111,
                0.11111111,
                0.16666667,
                0.16666667,
                0.16666667,
                0.125,
                0.125,
                0.125,
                0.125,
                0.16666667,
                0.16666667,
                1.,
                0.125,
                0.125,
                0.125,
                0.125,
                0.16666667,
            ]
        ),
    )

    hpx = HpxGeom([8, 16], False, "GAL", region="DISK(110.,75.,2.)", axes=[ax0])
    hpx2wcs = make_hpx_to_wcs_mapping(hpx, wcs)
    assert_allclose(
        hpx2wcs[0],
        np.array(
            [
                [
                    15,
                    6,
                    6,
                    6,
                    6,
                    6,
                    15,
                    15,
                    6,
                    6,
                    6,
                    6,
                    15,
                    15,
                    15,
                    6,
                    6,
                    6,
                    15,
                    15,
                    15,
                    6,
                    6,
                    6,
                    15,
                    15,
                    15,
                    15,
                    6,
                    6,
                    15,
                    15,
                    15,
                    15,
                    15,
                    6,
                ],
                [
                    67,
                    46,
                    46,
                    46,
                    46,
                    29,
                    67,
                    67,
                    46,
                    46,
                    46,
                    46,
                    67,
                    67,
                    67,
                    46,
                    46,
                    46,
                    67,
                    67,
                    67,
                    28,
                    28,
                    28,
                    45,
                    45,
                    45,
                    45,
                    28,
                    28,
                    66,
                    45,
                    45,
                    45,
                    45,
                    28,
                ],
            ]
        ),
    )

def test_hpxgeom_from_header():
    pars = {
        "HPX_REG": "DISK(110.,75.,2.)",
        "EXTNAME": "SKYMAP",
        "NSIDE": 2 ** 6,
        "ORDER": 6,
        "PIXTYPE": "HEALPIX",
        "ORDERING": "RING",
        "COORDSYS": "CEL",
        "TTYPE1": "PIX",
        "TFORM1": "K",
        "TTYPE2": "CHANNEL1",
        "TFORM2": "D",
        "INDXSCHM": "EXPLICIT",
    }
    header = fits.Header()
    header.update(pars)
    hpx = HpxGeom.from_header(header)

    assert hpx.coordsys == pars["COORDSYS"]
    assert hpx.nest is False
    assert_allclose(hpx.nside, np.array([64]))

@pytest.mark.parametrize(
    ("nside", "nested", "coordsys", "region", "axes"), hpx_test_geoms
)
def test_hpxgeom_read_write(tmpdir, nside, nested, coordsys, region, axes):
    geom0 = HpxGeom(nside, nested, coordsys, region=region, axes=axes)
    hdu_bands = geom0.make_bands_hdu(hdu="BANDS")
    hdu_prim = fits.PrimaryHDU()
    hdu_prim.header.update(geom0.make_header())

    filename = str(tmpdir / "hpxgeom.fits")
    hdulist = fits.HDUList([hdu_prim, hdu_bands])
    hdulist.writeto(filename, overwrite=True)

    hdulist = fits.open(filename)
    geom1 = HpxGeom.from_header(hdulist[0].header, hdulist["BANDS"])

    assert_allclose(geom0.nside, geom1.nside)
    assert_allclose(geom0.npix, geom1.npix)
    assert_allclose(geom0.nest, geom1.nest)
    assert geom0.coordsys == geom1.coordsys

@pytest.mark.parametrize(
    ("nside", "nested", "coordsys", "region", "axes"), hpx_test_geoms
)
def test_hpxgeom_upsample(nside, nested, coordsys, region, axes):
    # NESTED
    geom = HpxGeom(nside, True, coordsys, region=region, axes=axes)
    geom_up = geom.upsample(2)
    assert_allclose(2 * geom.nside, geom_up.nside)
    assert_allclose(4 * geom.npix, geom_up.npix)
    coords = geom_up.get_coord(flat=True)
    assert np.all(geom.contains(coords))

    # RING
    geom = HpxGeom(nside, False, coordsys, region=region, axes=axes)
    geom_up = geom.upsample(2)
    assert_allclose(2 * geom.nside, geom_up.nside)
    assert_allclose(4 * geom.npix, geom_up.npix)
    coords = geom_up.get_coord(flat=True)
    assert np.all(geom.contains(coords))

@pytest.mark.parametrize(
    ("nside", "nested", "coordsys", "region", "axes"), hpx_test_geoms
)
def test_hpxgeom_downsample(nside, nested, coordsys, region, axes):
    # NESTED
    geom = HpxGeom(nside, True, coordsys, region=region, axes=axes)
    geom_down = geom.downsample(2)
    assert_allclose(geom.nside, 2 * geom_down.nside)
    coords = geom.get_coord(flat=True)
    assert np.all(geom_down.contains(coords))

    # RING
    geom = HpxGeom(nside, False, coordsys, region=region, axes=axes)
    geom_down = geom.downsample(2)
    assert_allclose(geom.nside, 2 * geom_down.nside)
    coords = geom.get_coord(flat=True)
    assert np.all(geom_down.contains(coords))

def test_hpxgeom_solid_angle():
    geom = HpxGeom.create(nside=8, coordsys="GAL", axes=[MapAxis.from_edges([0, 2, 3])])

    solid_angle = geom.solid_angle()

    assert solid_angle.shape == (1,)
    assert_allclose(solid_angle.value, 0.016362461737446838)

def test_geom_repr():
    geom = HpxGeom(nside=8)
    assert geom.__class__.__name__ in repr(geom)
    assert "nside" in repr(geom)

max_stars_repo_name: contrera/gammapy

max_issues_repo_name: contrera/gammapy

max_forks_repo_name: contrera/gammapy

reason: import numpy,from numpy,from astropy


text:
"""
Created on March 20, 2018

@author: Alejandro Molina
"""
import numpy as np
import collections
from collections import deque, OrderedDict
import logging

logger = logging.getLogger(__name__)

class Node(object):
    def __init__(self):
        self.id = 0
        self.scope = []

    @property
    def name(self):
        return "%sNode_%s" % (self.__class__.__name__, self.id)

    @property
    def parameters(self):
        raise Exception("Not Implemented")

    def __repr__(self):
        return self.name

    def __rmul__(self, weight):
        assert type(weight) == int or type(weight) == float
        self._tmp_weight = weight
        return self

    def __mul__(self, node):
        assert isinstance(node, Node)
        assert len(node.scope) > 0, "right node has no scope"
        assert len(self.scope) > 0, "left node has no scope"
        assert len(set(node.scope).intersection(set(self.scope))) == 0, "children's scope is not disjoint"
        result = Product()
        result.children.append(self)
        result.children.append(node)
        result.scope.extend(self.scope)
        result.scope.extend(node.scope)
        assign_ids(result)
        return result

    def __add__(self, node):
        assert isinstance(node, Node)
        assert hasattr(node, "_tmp_weight"), "right node has no weight"
        assert hasattr(self, "_tmp_weight"), "left node has no weight"
        assert len(node.scope) > 0, "right node has no scope"
        assert len(self.scope) > 0, "left node has no scope"
        assert set(node.scope) == (set(self.scope)), "children's scope are not the same"

        from numpy import isclose

        assert isclose(
            1.0, self._tmp_weight + node._tmp_weight
        ), "unnormalized weights, maybe trying to add many nodes at the same time?"

        result = Sum()
        result.children.append(self)
        result.weights.append(self._tmp_weight)
        result.children.append(node)
        result.weights.append(node._tmp_weight)
        result.scope.extend(self.scope)
        result._tmp_weight = self._tmp_weight + node._tmp_weight
        assign_ids(result)
        return result

class Sum(Node):
    def __init__(self, weights=None, children=None):
        Node.__init__(self)
        if weights is None:
            weights = []
        self.weights = weights

        if children is None:
            children = []
        self.children = children

    @property
    def parameters(self):
        sorted_children = sorted(self.children, key=lambda c: c.id)
        params = [(n.id, self.weights[i]) for i, n in enumerate(sorted_children)]
        return tuple(params)

class Product(Node):
    def __init__(self, children=None):
        Node.__init__(self)
        if children is None:
            children = []
        self.children = children

    @property
    def parameters(self):
        return tuple(map(lambda n: n.id, sorted(self.children, key=lambda c: c.id)))

class Leaf(Node):
    def __init__(self, scope=None):
        Node.__init__(self)
        if scope is not None:
            if type(scope) == int:
                self.scope.append(scope)
            elif type(scope) == list:
                self.scope.extend(scope)
            else:
                raise Exception("invalid scope type %s " % (type(scope)))

class Max(Node):
    def __init__(self, dec_values=None, children=None, feature_name = None):
        Node.__init__(self)
        if dec_values is None:
            dec_values = []
        self.dec_values = dec_values

        if  feature_name is None:
            feature_name = "Decision_Node"
        self.feature_name = feature_name

        if children is None:
            children = []
        self.children = children

class Context:
    def __init__(self, meta_types=None, domains=None, parametric_types=None, scope=None, feature_names=None):
        self.meta_types = meta_types
        self.domains = domains
        self.parametric_types = parametric_types
        self.feature_names = feature_names

        if scope is None and parametric_types is not None:
            self.scope = range(0,len(parametric_types))
        elif scope is None and meta_types is not None:
            self.scope = range(0, len(meta_types))
        else:
            self.scope = scope

        if  feature_names is not None:
            self.feature_names = dict(zip(self.scope, feature_names ))

        if meta_types is None and parametric_types is not None:
            self.meta_types = []
            for p in parametric_types:
                self.meta_types.append(p.type.meta_type)
            self.parametric_types = dict(zip(self.scope, self.parametric_types))

        self.meta_types = dict(zip(self.scope, self.meta_types ))

    def get_meta_types_by_scope(self, scopes):
        return [self.meta_types[s] for s in scopes]

    def get_domains_by_scope(self, scopes):
        return [self.domains[s] for s in scopes]

    def get_parametric_types_by_scope(self, scopes):
        return [self.parametric_types[s] for s in scopes]

    def add_domains(self, data):
        assert len(data.shape) == 2, "data is not 2D?"
        assert data.shape[1] == len(self.meta_types), "Data columns and metatype size doesn't match"

        from spn.structure.StatisticalTypes import MetaType

        domain = []

        for col, meta_type in enumerate(self.meta_types.values()):

            feature_meta_type = meta_type
            min_val = np.nanmin(data[:, col])
            max_val = np.nanmax(data[:, col])
            domain_values = [min_val, max_val]

            if feature_meta_type == MetaType.REAL or feature_meta_type == MetaType.BINARY:
                domain.append(domain_values)
            elif feature_meta_type == MetaType.DISCRETE:
                domain.append(np.arange(domain_values[0], domain_values[1] + 1, 1))
            else:
                raise Exception("Unkown MetaType " + str(feature_meta_type))

        self.domains = np.asanyarray(domain)
        self.domains = dict(zip(self.scope, self.domains))

        return self

def get_number_of_edges(node):
    return sum([len(c.children) for c in get_nodes_by_type(node, (Sum, Product))])

def get_number_of_nodes(spn, node_type=Node):
    return len(get_nodes_by_type(spn, node_type))

def get_parents(node, includ_pos=True):
    parents = OrderedDict({node: []})
    for n in get_nodes_by_type(node):
        if not isinstance(n, Leaf):
            for i, c in enumerate(n.children):
                parent_list = parents.get(c, None)
                if parent_list is None:
                    parents[c] = parent_list = []
                if includ_pos:
                    parent_list.append((n, i))
                else:
                    parent_list.append(n)
    return parents

def get_depth(node):
    node_depth = {}

    def count_layers(node):
        ndepth = node_depth.setdefault(node, 1)

        if hasattr(node, "children"):
            for c in node.children:
                node_depth.setdefault(c, ndepth + 1)

    bfs(node, count_layers)

    return max(node_depth.values())

def rebuild_scopes_bottom_up(node):
    # this function is not safe (updates in place)

    for n in get_topological_order(node):
        if isinstance(n, Leaf):
            continue

        new_scope = set()
        for c in n.children:
            new_scope.update(c.scope)
        n.scope = list(new_scope)

    return node

def bfs(root, func):
    seen, queue = set([root]), collections.deque([root])
    while queue:
        node = queue.popleft()
        func(node)
        if not isinstance(node, Leaf):
            for c in node.children:
                if c not in seen:
                    seen.add(c)
                    queue.append(c)

def get_topological_order(node):
    nodes = get_nodes_by_type(node)

    parents = OrderedDict({node: []})
    in_degree = OrderedDict()
    for n in nodes:
        in_degree[n] = in_degree.get(n, 0)
        if not isinstance(n, Leaf):
            for c in n.children:
                parent_list = parents.get(c, None)
                if parent_list is None:
                    parents[c] = parent_list = []
                parent_list.append(n)
                in_degree[n] += 1

    S = deque()  # Set of all nodes with no incoming edge
    for u in in_degree:
        if in_degree[u] == 0:
            S.appendleft(u)

    L = []  # Empty list that will contain the sorted elements

    while S:
        n = S.pop()  # remove a node n from S
        L.append(n)  # add n to tail of L

        for m in parents[n]:  # for each node m with an edge e from n to m do
            in_degree_m = in_degree[m] - 1  # remove edge e from the graph
            in_degree[m] = in_degree_m
            if in_degree_m == 0:  # if m has no other incoming edges then
                S.appendleft(m)  # insert m into S

    assert len(L) == len(nodes), "Graph is not DAG, it has at least one cycle"
    return L

def get_topological_order_layers(node):
    nodes = get_nodes_by_type(node)

    parents = OrderedDict({node: []})
    in_degree = OrderedDict()
    for n in nodes:
        in_degree[n] = in_degree.get(n, 0)
        if not isinstance(n, Leaf):
            for c in n.children:
                parent_list = parents.get(c, None)
                if parent_list is None:
                    parents[c] = parent_list = []
                parent_list.append(n)
                in_degree[n] += 1

    layer = []  # Set of all nodes with no incoming edge
    for u in in_degree:
        if in_degree[u] == 0:
            layer.append(u)

    L = [layer]  # add first layer

    added_nodes = len(layer)
    while True:
        layer = []

        for n in L[-1]:
            for m in parents[n]:  # for each node m with an edge e from n to m do
                in_degree_m = in_degree[m] - 1  # remove edge e from the graph
                in_degree[m] = in_degree_m
                if in_degree_m == 0:  # if m has no other incoming edges then
                    layer.append(m)  # insert m into layer

        if len(layer) == 0:
            break

        added_nodes += len(layer)
        L.append(layer)

    assert added_nodes == len(nodes), "Graph is not DAG, it has at least one cycle"
    return L

def get_nodes_by_type(node, ntype=Node):
    assert node is not None

    result = []

    def add_node(node):
        if isinstance(node, ntype):
            result.append(node)

    bfs(node, add_node)

    return result

def get_node_types(node, ntype=Node):
    assert node is not None

    result = set()

    def add_node(node):
        if isinstance(node, ntype):
            result.add(type(node))

    bfs(node, add_node)

    return result

def assign_ids(node, ids=None):
    if ids is None:
        ids = {}

    def assign_id(node):
        if node not in ids:
            ids[node] = len(ids)

        node.id = ids[node]

    bfs(node, assign_id)
    return node

def eval_spn_bottom_up(node, eval_functions, all_results=None, debug=False, **args):
    """
    Evaluates the spn bottom up

    :param node: spn root
    :param eval_functions: is a dictionary that contains k:Class of the node, v:lambda function that receives as parameters (node, args**) for leave nodes and (node, [children results], args**)
    :param all_results: is a dictionary that contains k:Class of the node, v:result of the evaluation of the lambda function for that node. It is used to store intermediate results so that non-tree graphs can be computed in O(n) size of the network
    :param debug: whether to present progress information on the evaluation
    :param args: free parameters that will be fed to the lambda functions.
    :return: the result of computing and propagating all the values throught the network
    """

    nodes = get_topological_order(node)

    if debug:
        from tqdm import tqdm

        nodes = tqdm(list(nodes))

    if all_results is None:
        all_results = {}
    else:
        all_results.clear()

    for node_type, func in eval_functions.items():
        if "_eval_func" not in node_type.__dict__:
            node_type._eval_func = []
        node_type._eval_func.append(func)
        node_type._is_leaf = issubclass(node_type, Leaf)
    leaf_func = eval_functions.get(Leaf, None)

    tmp_children_list = []
    len_tmp_children_list = 0
    for n in nodes:

        try:
            func = n.__class__._eval_func[-1]
            n_is_leaf = n.__class__._is_leaf
        except:
            if isinstance(n, Leaf) and leaf_func is not None:
                func = leaf_func
                n_is_leaf = True
            else:
                raise AssertionError("No lambda function associated with type: %s" % (n.__class__.__name__))

        if n_is_leaf:
            result = func(n, **args)
        else:
            len_children = len(n.children)
            if len_tmp_children_list < len_children:
                tmp_children_list.extend([None] * len_children)
                len_tmp_children_list = len(tmp_children_list)
            for i in range(len_children):
                ci = n.children[i]
                tmp_children_list[i] = all_results[ci]
            result = func(n, tmp_children_list[0:len_children], **args)

        all_results[n] = result

    for node_type, func in eval_functions.items():
        del node_type._eval_func[-1]
        if len(node_type._eval_func) == 0:
            delattr(node_type, "_eval_func")

    return all_results[node]

def eval_spn_top_down(root, eval_functions, all_results=None, parent_result=None, **args):
    """
    evaluates an spn top to down

    :param root: spnt root
    :param eval_functions: is a dictionary that contains k:Class of the node, v:lambda function that receives as parameters (node, [parent_results], args**) and returns {child : intermediate_result}. This intermediate_result will be passed to child as parent_result. If intermediate_result is None, no further propagation occurs
    :param all_results: is a dictionary that contains k:Class of the node, v:result of the evaluation of the lambda function for that node.
    :param parent_result: initial input to the root node
    :param args: free parameters that will be fed to the lambda functions.
    :return: the result of computing and propagating all the values throught the network
    """
    if all_results is None:
        all_results = {}
    else:
        all_results.clear()

    for node_type, func in eval_functions.items():
        if "_eval_func" not in node_type.__dict__:
            node_type._eval_func = []
        node_type._eval_func.append(func)

    all_results[root] = [parent_result]

    for layer in reversed(get_topological_order_layers(root)):
        for n in layer:
            func = n.__class__._eval_func[-1]

            param = all_results[n]
            result = func(n, param, **args)

            if result is not None and not isinstance(n, Leaf):
                assert isinstance(result, dict)

                for child, param in result.items():
                    if child not in all_results:
                        all_results[child] = []
                    all_results[child].append(param)

    for node_type, func in eval_functions.items():
        del node_type._eval_func[-1]
        if len(node_type._eval_func) == 0:
            delattr(node_type, "_eval_func")

    return all_results[root]

max_stars_repo_name: SwarajPawar/SPFlow

max_issues_repo_name: minimum-LaytonC/SPFlow

max_forks_repo_name: minimum-LaytonC/SPFlow

reason: import numpy,from numpy


text:
# ----------------------------------------------------------------------------
# Copyright (c) 2020, Franck Lejzerowicz.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------

import sys
import pandas as pd
import numpy as np
from os.path import splitext

from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import seaborn as sns

dat = '<DAT>'
cur_raref = '<RAREF>'
tab_pd = pd.read_table('<TAB_FP>', index_col=0)
meta_pd = pd.read_table('<META_FP>', dtype={'sample_name': str})
colors_sample = '<COL_SAMPLE>'
colors_feature = '<COL_FEATURE>'
stats_tax_dat = '<TAX_DAT>'
split_taxa_fp = '<TAXA_FP>'
level = '<LEVEL>'
cohort = '<COHORT>'
collapsed = '<COLL>'
graph_fp = '<GRAPH_FP>'
if tab_pd.shape[0] > 1000:
    print('No plotting all %s features...' % tab_pd.shape[0])
    with open('%s.txt' % splitext(graph_fp)[0], 'w') as o:
        pass
else:
    if colors_sample:
        meta_pd = meta_pd.rename(columns={meta_pd.columns[0]: 'SAMPLE_ID'})
        meta_pd = meta_pd.set_index('SAMPLE_ID')
        if set(colors_sample).issubset(meta_pd.columns):
            colors_in = list(set(colors_sample) & set(meta_pd.columns))

    tab_pd = np.log10(tab_pd + 1).stack().reset_index().rename(
        columns={'level_1': 'SAMPLE_ID', 0: 'log10_reads'})
    tab_pd = tab_pd.rename(columns={tab_pd.columns[0]: 'OBSERVATION_ID'})

    tax_cols = []
    if colors_feature and split_taxa_fp:
        tax_pd = pd.read_table(split_taxa_fp, index_col=0)
        if level != 'feature':
            if level not in collapsed:
                sys.exit(0)
            tax_pd = tax_pd.iloc[:, :collapsed[level]].drop_duplicates()
            tax_pd['OBSERVATION_ID'] = tax_pd.apply(
                func=lambda x: ';'.join(
                    [y for y in x if str(y) != 'nan']), axis=1)
        else:
            tax_pd = tax_pd.reset_index()
            tax_pd = tax_pd.rename(
                columns={tax_pd.columns[0]: 'OBSERVATION_ID'})

        tax_pd = tax_pd.set_index('OBSERVATION_ID')
        tax_cols = [tax_pd.columns.tolist()[(x - 1)] if isinstance(x, int)
                    and x not in tax_pd.columns and tax_pd.shape[1] >= x
                    else x for x in colors_feature]
        tax_cols = [x for x in tax_cols if x in tax_pd.columns]
        if tax_cols:
            tax_pd = tax_pd[tax_cols]

    graphs = pd.read_csv(graph_fp, header=0, sep=',', dtype={'SAMPLE_ID': str})
    samples_order = [y for x, y in sorted(
        dict(graphs[['SAMPLE_RANK', 'SAMPLE_ID']].values).items())][::-1]
    features_order = [y for x, y in sorted(
        dict(graphs[['OBSERVATION_RANK', 'OBSERVATION_ID']].values).items()
    )][::-1]
    graphs = graphs.merge(
        tab_pd, on=['SAMPLE_ID', 'OBSERVATION_ID'], how='left')

    matrix = graphs[
        ['OBSERVATION_ID', 'SAMPLE_ID', 'log10_reads']].pivot_table(
        values='log10_reads', index='OBSERVATION_ID', columns='SAMPLE_ID')
    matrix = matrix.loc[features_order, samples_order]
    cur_meta_pd = meta_pd.loc[
        matrix.columns.astype(str).tolist(), colors_in].copy()
    cur_meta_leg_hex = cur_meta_pd.apply(func=lambda x: dict(
        zip(x.unique(), sns.color_palette(
            palette='Set1', n_colors=x.unique().size).as_hex())
    )).to_dict()
    cur_meta_leg_rgb = cur_meta_pd.apply(func=lambda x: dict(
        zip(x.unique(), sns.color_palette(
            palette='Set1', n_colors=x.unique().size))
    )).to_dict()
    cur_meta_pd_hex = cur_meta_pd.apply(func=lambda x: x.map(dict(
        zip(x.unique(), sns.color_palette(
            palette='Set1', n_colors=x.unique().size).as_hex())
    )))
    cur_meta_pd_rgb = cur_meta_pd.apply(func=lambda x: x.map(dict(
        zip(x.unique(), sns.color_palette(
            palette='Set1', n_colors=x.unique().size))
    )))

    if tax_cols:
        cur_tax_pd = tax_pd.loc[
            matrix.index.astype(str).tolist(), :].copy()
        cur_tax_leg_hex = cur_tax_pd.apply(func=lambda x: dict(
            zip(x.unique(), sns.color_palette(
                palette='Paired', n_colors=x.unique().size).as_hex())
        )).to_dict()
        cur_tax_leg_rgb = cur_tax_pd.apply(func=lambda x: dict(
            zip(x.unique(), sns.color_palette(
                palette='Paired', n_colors=x.unique().size))
        )).to_dict()
        cur_tax_pd_hex = cur_tax_pd.apply(func=lambda x: x.map(dict(
            zip(x.unique(), sns.color_palette(
                palette='Paired', n_colors=x.unique().size).as_hex())
        )))
        cur_tax_pd_rgb = cur_tax_pd.apply(func=lambda x: x.map(dict(
            zip(x.unique(), sns.color_palette(
                palette='Paired', n_colors=x.unique().size))
        )))

    X, Y = matrix.shape
    graphs_pdf = '%s.pdf' % splitext(graph_fp)[0]
    graphs_pdf_complex = '%s_complex.pdf' % splitext(graphs_pdf)[0]
    with PdfPages(graphs_pdf_complex) as pdf:
        fig, ax = plt.subplots(figsize=(20, (20 * Y/X)))
        if tax_cols:
            g = sns.clustermap(
                matrix, col_cluster=False, row_cluster=False,
                linewidths=0.1, cmap='coolwarm',
                row_colors=cur_tax_pd_hex, col_colors=cur_meta_pd_hex,
                yticklabels=False, xticklabels=False
            )
            simples = [('feature', cur_tax_pd_rgb, cur_tax_leg_rgb),
                       ('sample', cur_meta_pd_rgb, cur_meta_leg_rgb)]
        else:
            g = sns.clustermap(
                matrix, col_cluster=False, row_cluster=False,
                linewidths=0.1, cmap='coolwarm',
                col_colors=cur_meta_pd_hex,
                yticklabels=False, xticklabels=False
            )
            simples = [('sample', cur_meta_pd_rgb, cur_meta_leg_rgb)]

        n = 0
        N = 1 / cur_meta_pd_hex.columns.size
        for cdx, col in enumerate(cur_meta_pd_hex.columns):
            n += N
            if not cdx:
                first_leg = g.ax_col_dendrogram.legend(handles=[
                    mpatches.Patch(label=x, color=y)
                    for x, y in cur_meta_leg_hex[col].items()
                ], loc="upper left")
                g.ax_col_dendrogram.add_artist(first_leg)
            else:
                g.ax_col_dendrogram.legend(handles=[
                    mpatches.Patch(label=x, color=y)
                    for x, y in cur_meta_leg_hex[col].items()
                ], loc="upper left", bbox_to_anchor=(n, 1))

        if tax_cols:
            n = 0
            N = 1 / cur_tax_pd_hex.columns.size
            for cdx, col in enumerate(cur_tax_pd_hex.columns):
                n += N
                if not cdx:
                    first_leg = g.ax_row_dendrogram.legend(handles=[
                        mpatches.Patch(label=x, color=y)
                        for x, y in cur_tax_leg_hex[col].items()
                    ], loc="lower left")
                    g.ax_row_dendrogram.add_artist(first_leg)
                else:
                    g.ax_row_dendrogram.legend(handles=[
                        mpatches.Patch(label=x, color=y)
                        for x, y in cur_tax_leg_hex[col].items()
                    ], loc="lower left", bbox_to_anchor=(0, n))

        g.ax_heatmap.set_xlabel('Samples (sorted by richness)')
        g.ax_heatmap.set_ylabel('Taxon (sorted by prevalence)')

        suptitle = '%s%s' % (dat, cur_raref)
        if cohort != 'ALL':
            suptitle = suptitle + '\nsamples subset: %s' % cohort
        plt.suptitle(suptitle, fontsize=20)
        plt.subplots_adjust(top=.95, hspace=0.3)
        pdf.savefig(bbox_inches='tight', dpi=300)
        plt.close()

    print('[Nestedness] Written:', graphs_pdf_complex)

    graphs_pdf_simples = '%s_simples.pdf' % splitext(graphs_pdf)[0]
    with PdfPages(graphs_pdf_simples) as pdf:
        for (typ, tab, leg) in simples:
            num = tab.columns.size
            fig, axes = plt.subplots(
                num, 1, figsize=(9, (tab.columns.size * 4)))
            for cdx, col in enumerate(tab.columns):
                leg_col2rgb = dict(
                    (x, leg[col][x]) for idx, x in enumerate(leg[col]))
                cols_d = tab[col].to_dict()
                if typ == 'sample':
                    mat = [[[np.nan, np.nan, np.nan, 0.]
                            if str(x[1]) == 'nan'
                            else ([y for y in cols_d[x[0]]] + [1.])
                            for x in row.reset_index().values]
                           for r, row in matrix.iterrows()]
                if typ == 'feature':
                    mat = [[[np.nan, np.nan, np.nan, 0.]
                            if str(x) == 'nan'
                            else ([y for y in cols_d[r]] + [1.])
                            for x in row] for r, row in matrix.iterrows()]
                mat = np.array(mat)
                hands = [mpatches.Patch(
                    label=x, color=y) for x, y in leg_col2rgb.items()]
                if num == 1:
                    axes.imshow(mat, aspect="auto")
                    axes.legend(handles=hands, bbox_to_anchor=(1.05, 1),
                                loc=2, borderaxespad=0.)
                    axes.set_title(col)
                    axes.set_xlabel('Samples (sorted by richness)')
                    axes.set_ylabel('%s (sorted by prevalence)' % level)
                else:
                    axes[cdx].imshow(mat, aspect="auto")
                    axes[cdx].legend(handles=hands,
                                     bbox_to_anchor=(1.05, 1),
                                     loc=2, borderaxespad=0.)
                    axes[cdx].set_title(col)
                    axes[cdx].set_xlabel('Samples (sorted by richness)')
                    axes[cdx].set_ylabel(
                        '%s (sorted by prevalence)' % level)
            suptitle = '%s%s (%s coloring)' % (dat, cur_raref, typ)
            top = .93
            if cohort != 'ALL':
                top -= 0.03
                suptitle = suptitle + '\nsamples subset: %s' % cohort
            plt.suptitle(suptitle, fontsize=15)
            plt.subplots_adjust(top=top, hspace=0.3)
            pdf.savefig(bbox_inches='tight', dpi=300)
            plt.close()

    print('[Nestedness] Written:', graphs_pdf_simples)


text:
### Source code / logs

import numpy as np
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile

from distutils.version import StrictVersion
from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image

sys.path.append("..")
from object_detection.utils import ops as utils_ops

if StrictVersion(tf.__version__) < StrictVersion('1.12.0'):
  raise ImportError('Please upgrade your TensorFlow installation to v1.12.*.')

import tensorflow.compat.v1 as tf
from object_detection.utils import label_map_util

from object_detection.utils import visualization_utils as vis_util

# # Model preparation

# ## Variables
#
# Any model exported using the `export_inference_graph.py` tool can be loaded here simply by changing `PATH_TO_FROZEN_GRAPH` to point to a new .pb file.

# What model to download.
MODEL_NAME='output/t9'

# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_CKPT = MODEL_NAME + '//frozen_inference_graph.pb'

# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS=os.path.join('F://Trashnet//Output', '//label_map.pbtxt')

detection_graph = tf.Graph()
with detection_graph.as_default():
  od_graph_def = tf.GraphDef()
  with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
    serialized_graph = fid.read()
    od_graph_def.ParseFromString(serialized_graph)
    tf.import_graph_def(od_graph_def, name='')

# ## Loading label map
# Label maps map indices to category names, so that when our convolution network predicts `5`, we know that this corresponds to `airplane`.  Here we use internal utility functions, but anything that returns a dictionary mapping integers to appropriate string labels would be fine

#category_index = {1: {'id': 1, 'name': 'plastic'}}
category_index = label_map_util.create_category_index_from_labelmap('F:/Trashnet/Output/label_map.pbtxt', use_display_name=True)

# ## Helper code

def load_image_into_numpy_array(image):
  (im_width, im_height) = image.size
  return np.array(image.getdata()).reshape(
      (im_height, im_width, 3)).astype(np.uint8)

# # Detection
PATH_TO_TEST_IMAGES_DIR = 'F://Plastic//val'
#TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, '1.jpg') for i in range(3, 4) ]
TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, '{0}.jpg'.format(i)) for i in range(1, 25) ]

# Size, in inches, of the output images.
IMAGE_SIZE = (12, 8)

def run_inference_for_single_image(image, graph):
  with graph.as_default():
    with tf.Session() as sess:
      # Get handles to input and output tensors
      ops = tf.get_default_graph().get_operations()
      all_tensor_names = {output.name for op in ops for output in op.outputs}
      tensor_dict = {}
      for key in [
          'num_detections', 'detection_boxes', 'detection_scores',
          'detection_classes', 'detection_masks'
      ]:
        tensor_name = key + ':0'
        if tensor_name in all_tensor_names:
          tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(
              tensor_name)
      if 'detection_masks' in tensor_dict:
        # The following processing is only for single image
        detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
        detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
        # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
        real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
        detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
        detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
        detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
            detection_masks, detection_boxes, image.shape[1], image.shape[2])
        detection_masks_reframed = tf.cast(
            tf.greater(detection_masks_reframed, 0.5), tf.uint8)
        # Follow the convention by adding back the batch dimension
        tensor_dict['detection_masks'] = tf.expand_dims(
            detection_masks_reframed, 0)
      image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')

      # Run inference
      output_dict = sess.run(tensor_dict,
                             feed_dict={image_tensor: image})

      # all outputs are float32 numpy arrays, so convert types as appropriate
      output_dict['num_detections'] = int(output_dict['num_detections'][0])
      output_dict['detection_classes'] = output_dict[
          'detection_classes'][0].astype(np.uint8)
      output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
      output_dict['detection_scores'] = output_dict['detection_scores'][0]
      if 'detection_masks' in output_dict:
        output_dict['detection_masks'] = output_dict['detection_masks'][0]
  return output_dict

for image_path in TEST_IMAGE_PATHS:
  image = Image.open(image_path)
  print(image_path)
  # the array based representation of the image will be used later in order to prepare the
  # result image with boxes and labels on it.
  image_np = load_image_into_numpy_array(image)
  # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
  image_np_expanded = np.expand_dims(image_np, axis=0)
  # Actual detection.
  output_dict = run_inference_for_single_image(image_np_expanded, detection_graph)
  # Visualization of the results of a detection.
  vis_util.visualize_boxes_and_labels_on_image_array(
      image_np,
      output_dict['detection_boxes'],
      output_dict['detection_classes'],
      output_dict['detection_scores'],
      category_index,
      instance_masks=output_dict.get('detection_masks'),
      use_normalized_coordinates=True,
      line_thickness=8)
  plt.figure(figsize=IMAGE_SIZE)
  plt.imshow(image_np)
  img_folder = MODEL_NAME+'/image'

  if not os.path.exists(img_folder):
    os.makedirs(img_folder)

  plt.savefig(img_folder + '/' + image_path.split('\\')[1])

#plt.show()

#requires matplotlib==3.0.1