# from pynndescent import NNDescent from utils import lovasz_losses as L from utils import utils import torch from torch import nn from torch import optim import numpy as np import networkx as nx import random # def crossings(pos, G, k2i, sampleSize, sampleOn='edges', reg_coef=1, niter=30): # crossing_segs_sample = utils.sample_crossings(pos, G, k2i, sampleSize, sampleOn) # # if len(crossing_segs_sample) < sampleSize*0.5: # if sampleOn=='edges' and len(crossing_segs_sample) == 0: # crossing_segs_sample = utils.sample_crossings(pos, G, k2i, sampleSize, sampleOn='crossings') # if len(crossing_segs_sample) > 0: # pos_segs = pos[crossing_segs_sample.flatten()].view(-1,4,2) # w = (torch.rand(pos_segs.shape[0], 2, 1)-0.5).requires_grad_(True) # b = (torch.rand(pos_segs.shape[0], 1, 1)-0.5).requires_grad_(True) # relu = nn.ReLU() # o = optim.SGD([w,b], lr=0.01, momentum=0.5, nesterov=True) # for _ in range(niter): # pred = pos_segs.detach() @ w + b # ## assume labels of nodes in the first edges are -1 # ## now flip the pred of those nodes so that now we want every pred to be +1 # pred[:,:2,:] = -pred[:,:2,:] # loss_svm = relu(1-pred).sum() + reg_coef * w.pow(2).sum() # o.zero_grad() # loss_svm.backward() # o.step() # pred = pos_segs @ w.detach() + b.detach() # pred[:,:2,:] = -pred[:,:2,:] # loss_crossing = relu(1-pred).sum() # return loss_crossing # else: # ##return dummy loss # return (pos[0,0]*0).sum() # def angular_resolution(pos, G, k2i, sampleSize=2): # samples = utils.sample_nodes(G, sampleSize) # neighbors = [list(G.neighbors(s)) for s in samples] # sampleIndices = [k2i[s] for s in samples] # neighborIndices = [[k2i[n] for n in nei] for nei in neighbors] # samples = pos[sampleIndices] # neighbors = [pos[nei] for nei in neighborIndices] # angles = [utils.get_angles(rs) for rs in rays] # if len(angles) > 0: # loss = sum([torch.exp(-a*len(a)).sum() for a in angles]) # # loss = sum([(a - np.pi*20/len(a)).pow(2).sum() for a in angles]) # else: # loss = pos[0,0]*0##dummy output # return loss def angular_resolution(pos, G, k2i, sampleSize=2, sample=None): relu = nn.ReLU() if sample is None: sample = utils.sample_nodes(G, sampleSize) sample = [s for s in sample if len(list(G.neighbors(s)))>=2] if len(sample)==0: loss = pos[0,0]*0##dummy output return loss degrees = torch.tensor([G.degree(s) for s in sample]) neighbors = [random.choices(list(G.neighbors(s)),k=2) for s in sample] sampleIndices = [k2i[s] for s in sample] neighborIndices = [[k2i[n] for n in nei] for nei in neighbors] sample = pos[sampleIndices] neighbors = torch.stack([pos[nei] for nei in neighborIndices]) rays = neighbors - sample.unsqueeze(1) sim = cos_sim(rays[:,0,:], rays[:,1,:]) else: degrees, a,b,c,d = sample a,b,c,d = pos[a], pos[b], pos[c], pos[d] sim = cos_sim(b-a, d-c) angles = torch.acos(sim.clamp_(-0.99,0.99)) if sim.shape[0] > 0: # loss = ((cos+1)**2 / 4).sum() # loss = torch.exp(-angles).sum() # loss = relu(-angles + 2*np.pi/degrees).pow(2).mean() optimal = 2*np.pi/degrees loss = bce(relu((-angles + optimal)/optimal), torch.zeros_like(angles)) # loss = torch.exp(-angles/optimal).mean() else: loss = pos[0,0]*0##dummy output return loss def gabriel(pos, G, k2i, sample=None, sampleSize=64): if sample is not None: nodes, edges = sample edges = torch.stack(edges, dim=-1) else: edges = utils.sample_edges(G, sampleSize) nodes = utils.sample_nodes(G, sampleSize) edges = np.array([(k2i[e0], k2i[e1]) for e0,e1 in edges]) nodes = np.array([k2i[n] for n in nodes]) m,n = len(nodes), len(edges) node_pos = pos[nodes] edge_pos = pos[edges.flatten()].reshape([-1,2,2]) centers = edge_pos.mean(1) radii = (edge_pos[:,0,:] - edge_pos[:,1,:]).norm(dim=1, keepdim=True)/2 # centers = centers.repeat(1,m).view(-1, 2) # radii = radii.repeat(1,m).view(-1, 1) # node_pos = node_pos.repeat(n,1) relu = nn.ReLU() # print((node_pos-centers).norm(dim=1)) loss = relu(radii+0.01 - (node_pos-centers).norm(dim=1)).pow(2) # loss = relu(radii - (node_pos-centers).norm(dim=1)).pow(2) loss = loss.mean() return loss bce = torch.nn.BCELoss() cos_sim = torch.nn.CosineSimilarity() def crossing_angle_maximization(pos, G, k2i, i2k, sample=None, sample_labels=None, sampleSize=16, sampleOn='edges'): # edge_list = list(G.edges) # if sampleOn == 'edges': # sample_indices = np.random.choice(len(edge_list), sampleSize, replace=False) # edge_samples = [edge_list[i] for i in sample_indices] # crossing_segs_sample = utils.find_crossings(pos, edge_samples, k2i) # elif sampleOn == 'crossings': # crossing_segs = utils.find_crossings(pos, edge_list, k2i) # crossing_count = crossing_segs.shape[0] # sample_indices = np.random.choice(crossing_count, min(sampleSize,crossing_count), replace=False) # crossing_segs_sample = crossing_segs[sample_indices] # if len(crossing_segs_sample) > 0: # pos_segs = pos[crossing_segs_sample.flatten()].view(-1,4,2) if sample is None: crossing_segs_sample = utils.sample_crossings(pos, G, k2i, sampleSize, sampleOn) if sampleOn=='edges' and len(crossing_segs_sample) == 0: crossing_segs_sample = utils.sample_crossings(pos, G, k2i, sampleSize, sampleOn='crossings') sample = crossing_segs_sample pos_segs = pos[sample.flatten()].view(-1,4,2) if sample_labels is None: sample_labels = utils.are_edge_pairs_crossed(pos_segs.view(-1,8)) pos_segs = pos[sample.flatten()].view(-1,4,2) v1 = pos_segs[:,1] - pos_segs[:,0] v2 = pos_segs[:,3] - pos_segs[:,2] sim = cos_sim(v1, v2) # angles = torch.acos(sim.clamp_(-0.99,0.99)) # return (sample_labels*(angles-np.pi/2)).pow(2).mean() # return (sample_labels * sim**2).mean() return (sample_labels * sim**2 / (1-sim**2+1e-6)).mean() # return bce( # sample_labels * (sim**2).clamp_(1e-4, 1-1e-4), # torch.zeros_like(sim) # ) # else: # return pos[0,0]*0##dummy loss def aspect_ratio(pos, sampleSize=None, sample=None, # rotation_angles=torch.arange(7,dtype=torch.float)/7*(np.pi/2), target=[1,1]): if target[0] < target[1]: target = target[::-1] if sample is not None: sample = pos[sample,:] elif sampleSize is None or sampleSize=='full': sample = pos else: n = pos.shape[0] i = np.random.choice(n, min(n,sampleSize), replace=False) sample = pos[i,:] bce = nn.BCELoss(reduction='sum') singular_values = torch.svd(sample-sample.mean(dim=0)).S return bce(singular_values[1]/singular_values[0], torch.tensor(target[1]/target[0])) # mean = sample.mean(dim=0, keepdim=True) # sample -= mean # scale = sample.max().detach() # sample = sample/scale * 1 # cos = torch.cos(rotation_angles) # sin = torch.sin(rotation_angles) # rot = torch.stack([cos, sin, -sin, cos], 1).view(len(rotation_angles), 2, 2) # sample = sample.matmul(rot) # softmax = nn.Softmax(dim=1) # # print(softmax(sample)) # max_hat = (softmax(sample) * sample).sum(1) # min_hat = (softmax(-sample) * sample).sum(1) # w = max_hat[:,0] - min_hat[:,0] # h = max_hat[:,1] - min_hat[:,1] # estimate = torch.stack([w,h], 1) # estimate /= estimate.sum(1, keepdim=True) # # print(estimate) # target = torch.tensor(target_width_to_height, dtype=torch.float) # target /= target.sum() # target = target.repeat(len(rotation_angles), 1) # bce = nn.BCELoss(reduction='mean') # mse = nn.MSELoss(reduction='mean') # return mse(estimate, target) def vertex_resolution(pos, sampleSize=None, sample=None, target=0.1, prev_target_dist=1, prev_weight=1): pairwiseDistance = nn.PairwiseDistance() relu = nn.ReLU() # softmax = nn.Softmax(dim=0) # softmin = nn.Softmin(dim=0) n = pos.shape[0] if sample is None: if sampleSize is None or sampleSize=='full': sample = pos else: i = np.random.choice(n, min(n,sampleSize), replace=False) sample = pos[i,:] m = sample.shape[0] a = sample.repeat([1,m]).view(-1,2) b = sample.repeat([m,1]) pdist = pairwiseDistance(a, b) else: a = pos[sample[:,0],:] b = pos[sample[:,1],:] pdist = pairwiseDistance(a,b) # dmax = (softmax(pdist)*pdist).sum().detach() dmax = pdist.max().detach() target_dist = target*dmax ## exponentially smoothed target_dist smoothness = 0.1 weight = prev_weight*smoothness + 1 target_dist = ( max(target_dist, prev_target_dist) +min(target_dist, prev_target_dist)*smoothness )/weight # loss = len(pdist)*(softmin(pdist).detach() * relu((target_dist - pdist)/target_dist)).sum() # loss = relu((target_dist - pdist)/targetDist).sum() # loss = relu(target_dist - pdist).sum() loss = relu(1 - pdist/target_dist).pow(2).mean() return loss, target_dist, weight def neighborhood_preseration(pos, G, adj, k2i, i2k, degrees, max_degree, sample=None, n_roots=2, depth_limit=1, neg_sample_rate=0.5, device='cpu'): if sample is not None: pos_samples = [] for root_index in sample: root = i2k[root_index] G_sub = nx.bfs_tree(G, root, depth_limit=depth_limit) pos_samples += [k2i[n] for n in G_sub.nodes] pos_samples = sorted(set(pos_samples)) else: pos_samples = [] for _ in range(n_roots): root = i2k[random.randint(0, len(G)-1)] G_sub = nx.bfs_tree(G, root, depth_limit=depth_limit) pos_samples += [k2i[n] for n in G_sub.nodes] pos_samples = sorted(set(pos_samples)) n_neg_samples = int(neg_sample_rate * len(pos_samples)) neg_samples = [random.randint(0, len(G)-1) for _ in range(n_neg_samples)] ##negative samples samples = sorted(set( pos_samples + neg_samples )) ##remove duplicates pos = pos[samples,:] adj = adj[samples,:][:, samples] n,m = pos.shape x = pos ## pdist x0 = x.repeat(1, n).view(-1,m) x1 = x.repeat(n, 1) pdist = nn.PairwiseDistance()(x0, x1).view(n, n) ## k_dist # Option 1, mean of k-th and (k+1)th-NN as threshold # degrees = degrees[samples] # max_degree = degrees.max() # n_neighbors = max(2, min(max_degree+1, n)) # n_trees = min(64, 5 + int(round((n) ** 0.5 / 20.0))) # n_iters = max(5, int(round(np.log2(n)))) # knn_search_index = NNDescent( # x.detach().numpy(), # n_neighbors=n_neighbors, # n_trees=n_trees, # n_iters=n_iters, # max_candidates=60, # ) # knn_indices, knn_dists = knn_search_index.neighbor_graph # kmax = knn_dists.shape[1]-1 # k_dist = np.array([ # ( knn_dists[i,min(kmax, k)] + knn_dists[i,min(kmax, k+1)] ) / 2 # for i,k in enumerate(degrees) # ]) # pred = torch.from_numpy(k_dist.astype(np.float32)).view(-1,1) - pdist ## Option 2, fraction of diameter of drawing as threshold ## TODO adaptive dist (e.g. normalize by diameter of the drawing, 0.1*diameter) # diameter = pdist.max().item() # k_dist = 0.1 * max(diameter, 1.0) # pred = -pdist + k_dist ## Option 3, constant threshold k_dist = 1.5 pred = -pdist + k_dist # pred = torch.sign(-pdist + k_dist) * (-pdist + k_dist).pow(2) ## Option 4 (TODO), threshold by furthest neighboring node target = adj + torch.eye(adj.shape[0], device=device) loss = L.lovasz_hinge(pred, target) return loss def ideal_edge_length(pos, G, k2i, targetLengths=None, sampleSize=None, sample=None, reduce='mean'): if targetLengths is None: targetLengths = {e:1 for e in G.edges} n,m = pos.shape[0], pos.shape[1] if sample is not None: edges = sample elif sampleSize is not None: edges = random.sample(G.edges, sampleSize) else: edges = G.edges sourceIndices, targetIndices = zip(*[ [k2i[e0], k2i[e1]] for e0,e1 in edges]) source = pos[sourceIndices,:] target = pos[targetIndices,:] edgeLengths = (source-target).norm(dim=1) targetLengths = torch.tensor([targetLengths[e] for e in edges]) eu = ((edgeLengths-targetLengths)/targetLengths).pow(2) if reduce == 'sum': return eu.sum() elif reduce == 'mean': return eu.mean() def stress(pos, D, W, sampleSize=None, sample=None, reduce='mean'): if sample is None: n,m = pos.shape[0], pos.shape[1] if sampleSize is not None: i0 = np.random.choice(n, sampleSize) i1 = np.random.choice(n, sampleSize) x0 = pos[i0,:] x1 = pos[i1,:] D = torch.tensor([D[i,j] for i, j in zip(i0, i1)]) W = torch.tensor([W[i,j] for i, j in zip(i0, i1)]) else: x0 = pos.repeat(1, n).view(-1,m) x1 = pos.repeat(n, 1) D = D.view(-1) W = W.view(-1) else: x0 = pos[sample[:,0],:] x1 = pos[sample[:,1],:] D = torch.tensor([D[i,j] for i, j in sample]) W = torch.tensor([W[i,j] for i, j in sample]) pdist = nn.PairwiseDistance()(x0, x1) # wbound = (1/4 * diff.abs().min()).item() # W.clamp_(0, wbound) res = W*(pdist-D)**2 if reduce == 'sum': return res.sum() elif reduce == 'mean': return res.mean()