embedding_truth.py

import torch
import numpy as np
import awkward as ak

def make_embedding_truth(truth_label_by_hits, device='cpu'):
    truth_ordering = np.argsort(truth_label_by_hits)
    uniques, counts = np.unique(truth_label_by_hits, return_counts=True)
    truth_indices = ak.JaggedArray.fromcounts(counts, truth_ordering) # delete noise?   
    pairwise_truth = truth_indices.choose(2)
    other_samples = np.zeros(pairwise_truth.counts.sum() + (pairwise_truth.counts==0).sum(), dtype=np.int64)
    edge_index = np.zeros((2, pairwise_truth.counts.sum() + other_samples.size), dtype=np.int64)
    truth_label = np.zeros((pairwise_truth.counts.sum() + other_samples.size, ), dtype=np.int64)
    current_offset = 0
    for i, pairs in enumerate(pairwise_truth):
        mask = np.ones((truth_indices.size), dtype=np.bool)
        mask[i] = False
        if pairs.size > 0:
            in_class = np.random.choice(truth_indices[i], pairs.size)            
            not_class = np.random.choice(truth_indices[mask].content, pairs.size)
            edge_index[0][current_offset:current_offset+2*pairs.size] = np.concatenate((pairs.i0, in_class))            
            edge_index[1][current_offset:current_offset+2*pairs.size] = np.concatenate((pairs.i1, not_class))
            truth_label[current_offset:current_offset+pairs.size] = 1
            truth_label[current_offset+pairs.size:current_offset+2*pairs.size] = -1
            current_offset += 2*pairs.size
        else:
            in_class = truth_indices[i]
            not_class = np.random.choice(truth_indices[mask].content, 1)
            edge_index[0][current_offset:current_offset+1] = in_class
            edge_index[1][current_offset:current_offset+1] = not_class
            truth_label[current_offset:current_offset+1] = -1
            current_offset += 1
    return torch.from_numpy(edge_index).to(device), torch.from_numpy(truth_label).to(device)