-
Notifications
You must be signed in to change notification settings - Fork 0
/
graph_helpers.py
399 lines (297 loc) · 14.9 KB
/
graph_helpers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
import torch
from torch_geometric.nn import radius_graph, knn_graph
from torch_geometric.utils import subgraph
import utils
from math import sqrt
def getImPos(rows,cols,start_row=0,start_col=0):
row_space = torch.arange(start_row,rows+start_row)
col_space = torch.arange(start_col,cols+start_col)
col_image,row_image = torch.meshgrid(col_space,row_space,indexing='xy')
im_pos = torch.reshape(torch.stack((row_image,col_image),dim=-1),(rows*cols,2))
return im_pos
def convertImPos(im_pos,flip_y=True):
# Cast to float for clustering based methods
pos2D = im_pos.float()
# Switch rows,cols to x,y
pos2D[:,[1,0]] = pos2D[:,[0,1]]
if flip_y:
# Flip to y-axis to match mathematical definition and edges2Selections settings
pos2D[:,1] = torch.amax(pos2D[:,1]) - pos2D[:,1]
return pos2D
def grid2Edges(locs):
# Assume locs are already spaced at a distance of 1 structure
edge_index = radius_graph(locs,1.44,loop=True)
return edge_index
def radius2Edges(locs,r=1.0):
edge_index = radius_graph(locs,r,loop=True)
return edge_index
def knn2Edges(locs,knn=9):
edge_index = knn_graph(locs,knn,loop=True)
return edge_index
def surface2Edges(pos3D,normals,up_vector=None,k_neighbors=9):
if up_vector is None:
up_vector = torch.tensor([[0.0,1.0,0.0]]).to(pos3D.device)
# K Nearest Neighbors graph
edge_index = knn_graph(pos3D,k_neighbors,loop=True)
# Cull neighbors based on normals (dot them together)
culling = torch.sum(torch.multiply(normals[edge_index[1]],normals[edge_index[0]]),dim=1)
edge_index = edge_index[:,torch.where(culling>0)[0]]
# For each node, rotate based on Grahm-Schmidt Orthognalization
norms = normals[edge_index[0]]
z_dir = norms
z_dir = z_dir/torch.linalg.norm(z_dir,dim=1,keepdims=True) # Make sure it is a unit vector
#x_dir = torch.cross(up_vector,norms,dim=1)
x_dir = utils.cross(up_vector,norms) # torch.cross doesn't broadcast properly in some versions of torch
x_dir = x_dir/torch.linalg.norm(x_dir,dim=1,keepdims=True)
#y_dir = torch.cross(norms,x_dir,dim=1)
y_dir = utils.cross(norms,x_dir)
y_dir = y_dir/torch.linalg.norm(y_dir,dim=1,keepdims=True)
directions = (pos3D[edge_index[1]] - pos3D[edge_index[0]])
# Perform rotation by multiplying out rotation matrix
temp = torch.clone(directions) # Buffer
directions[:,0] = temp[:,0] * x_dir[:,0] + temp[:,1] * x_dir[:,1] + temp[:,2] * x_dir[:,2]
directions[:,1] = temp[:,0] * y_dir[:,0] + temp[:,1] * y_dir[:,1] + temp[:,2] * y_dir[:,2]
#directions[:,2] = temp[:,0] * z_dir[:,0] + temp[:,1] * z_dir[:,1] + temp[:,2] * z_dir[:,2]
# Drop z coordinate
directions = directions[:,:2]
return edge_index, directions
def edges2Selections(edge_index,directions,interpolated=True,bary_d=None,y_down=False):
# Current Ordering
# 4 3 2
# 5 0 1
# 6 7 8
if y_down:
vectorList = torch.tensor([[1,0],[sqrt(2)/2,-sqrt(2)/2],[0,-1],[-sqrt(2)/2,-sqrt(2)/2],[-1,0],[-sqrt(2)/2,sqrt(2)/2],[0,1],[sqrt(2)/2,sqrt(2)/2]],dtype=torch.float).transpose(1,0)
else:
vectorList = torch.tensor([[1,0],[sqrt(2)/2,sqrt(2)/2],[0,1],[-sqrt(2)/2,sqrt(2)/2],[-1,0],[-sqrt(2)/2,-sqrt(2)/2],[0,-1],[sqrt(2)/2,-sqrt(2)/2]],dtype=torch.float).transpose(1,0)
if interpolated:
if bary_d is None:
edge_index,selections,interps = interpolateSelections(edge_index,directions,vectorList)
else:
edge_index,selections,interps = interpolateSelections_barycentric(edge_index,directions,bary_d,vectorList)
interps = normalizeEdges(edge_index,selections,interps)
return edge_index,selections,interps
else:
selections = torch.argmax(torch.matmul(directions,vectorList),dim=1) + 1
selections[torch.where(torch.sum(torch.abs(directions),axis=1) == 0)] = 0 # Same cell selection
return selections
def makeEdges(prev_sources,prev_targets,prev_selections,sources,targets,selection,reverse=True):
sources = sources.flatten()
targets = targets.flatten()
prev_sources += sources.tolist()
prev_targets += targets.tolist()
prev_selections += len(sources)*[selection]
if reverse:
prev_sources += targets
prev_targets += sources
prev_selections += len(sources)*[utils.reverse_selection(selection)]
return prev_sources,prev_targets,prev_selections
def maskNodes(mask,x):
node_mask = torch.where(mask)
x = x[node_mask]
return x
def maskPoints(mask,x,y):
mask = torch.squeeze(mask)
x0 = torch.floor(x).long()
x1 = x0 + 1
y0 = torch.floor(y).long()
y1 = y0 + 1
x0 = torch.clip(x0, 0, mask.shape[1]-1);
x1 = torch.clip(x1, 0, mask.shape[1]-1);
y0 = torch.clip(y0, 0, mask.shape[0]-1);
y1 = torch.clip(y1, 0, mask.shape[0]-1);
Ma = mask[ y0, x0 ]
Mb = mask[ y1, x0 ]
Mc = mask[ y0, x1 ]
Md = mask[ y1, x1 ]
node_mask = torch.where(torch.logical_and(torch.logical_and(torch.logical_and(Ma,Mb),Mc),Md))[0]
return node_mask
def maskGraph(mask,edge_index,selections,interps=None):
edge_index,_,edge_mask = subgraph(mask,edge_index,relabel_nodes=True,return_edge_mask=True)
selections = selections[edge_mask]
if interps:
interps = interps[edge_mask]
return edge_index, selections, interps
else:
return edge_index, selections
def interpolateSelections(edge_index,directions,vectorList=None):
if vectorList is None:
# Current Ordering
# 4 3 2
# 5 0 1
# 6 7 8
vectorList = torch.tensor([[1,0],[sqrt(2)/2,sqrt(2)/2],[0,1],[-sqrt(2)/2,sqrt(2)/2],[-1,0],[-sqrt(2)/2,-sqrt(2)/2],[0,-1],[sqrt(2)/2,-sqrt(2)/2]],dtype=torch.float).transpose(1,0)
# Normalize directions for simplicity of calculations
dir_norm = torch.linalg.norm(directions,dim=1,keepdims=True)
directions = directions/dir_norm
#locs = torch.where(dir_norm > 1)[0]
#directions[locs] = directions[locs]/dir_norm[locs]
values = torch.matmul(directions,vectorList)
best = torch.unsqueeze(torch.argmax(values,dim=1),1)
best_val = torch.take_along_dim(values,best,dim=1)
# Look at both neighbors to see who is closer
lower_val = torch.take_along_dim(values,(best-1) % 8,dim=1)
upper_val = torch.take_along_dim(values,(best+1) % 8,dim=1)
comp_vals = torch.cat((lower_val,upper_val),dim=1)
second_best_vals = torch.amax(comp_vals,dim=1)
second_best = torch.argmax(comp_vals,dim=1)
# Find the interpolation value (in terms of angles)
best_val = torch.minimum(best_val[:,0],torch.tensor(1,device=directions.device)) # Prep for arccos function
angle_best = torch.arccos(best_val)
angle_second_best = torch.arccos(second_best_vals)
angle_vals = angle_best/(angle_second_best + angle_best)
# Use negative values for clockwise selections
clockwise = torch.where(second_best == 0)[0]
angle_vals[clockwise] = -angle_vals[clockwise]
# Handle computation problems at the poles
angle_vals = torch.nan_to_num(angle_vals)
# Make Selections
selections = best[:,0] + 1
# Same cell selection
same_locs = torch.where(edge_index[0] == edge_index[1])
selections[same_locs] = 0
angle_vals[same_locs] = 0
# Make starting interp_values
interps = torch.ones_like(angle_vals)
interps -= torch.abs(angle_vals)
# Add new edges
pos_interp_locs = torch.where(angle_vals > 1e-2)[0]
pos_interps = angle_vals[pos_interp_locs]
pos_edges = edge_index[:,pos_interp_locs]
pos_selections = selections[pos_interp_locs] + 1
pos_selections[torch.where(pos_selections>8)] = 1 # Account for wrap around
neg_interp_locs = torch.where(angle_vals < -1e-2)[0]
neg_interps = torch.abs(angle_vals[neg_interp_locs])
neg_edges = edge_index[:,neg_interp_locs]
neg_selections = selections[neg_interp_locs] - 1
neg_selections[torch.where(neg_selections<1)] = 8 # Account for wrap around
edge_index = torch.cat((edge_index,pos_edges,neg_edges),dim=1)
selections = torch.cat((selections,pos_selections,neg_selections),dim=0)
interps = torch.cat((interps,pos_interps,neg_interps),dim=0)
return edge_index,selections,interps
def interpolateSelections_barycentric(edge_index,directions,d,vectorList=None):
if vectorList is None:
# Current Ordering
# 4 3 2
# 5 0 1
# 6 7 8
vectorList = torch.tensor([[1,0],[sqrt(2)/2,-sqrt(2)/2],[0,-1],[-sqrt(2)/2,-sqrt(2)/2],[-1,0],[-sqrt(2)/2,sqrt(2)/2],[0,1],[sqrt(2)/2,sqrt(2)/2]],dtype=torch.float).transpose(1,0).to(directions.device)
# Preprune central selections and reappend them at the end
same_locs = torch.where(edge_index[0] == edge_index[1])
same_edges = edge_index[:,same_locs[0]]
different_locs = torch.where(edge_index[0] != edge_index[1])
edge_index = edge_index[:,different_locs[0]]
directions = directions[different_locs[0]]
# Normalize directions for simplicity of calculations
dir_norm = torch.linalg.norm(directions,dim=1,keepdims=True)
unit_directions = directions/dir_norm
#locs = torch.where(dir_norm > 1)[0]
#directions[locs] = directions[locs]/dir_norm[locs]
values = torch.matmul(unit_directions,vectorList)
best = torch.unsqueeze(torch.argmax(values,dim=1),1)
#best_val = torch.take_along_dim(values,best,dim=1)
# Look at both neighbors to see who is closer
lower_val = torch.take_along_dim(values,(best-1) % 8,dim=1)
upper_val = torch.take_along_dim(values,(best+1) % 8,dim=1)
comp_vals = torch.cat((lower_val,upper_val),dim=1)
second_best = torch.argmax(comp_vals,dim=1)
#second_best_vals = torch.amax(comp_vals,dim=1)
# Convert into uv cooridnates for barycentric interpolation calculation
# /|
# / |v
# /__|
# u
scaled_directions = torch.abs(directions/d)
u = torch.amax(scaled_directions,dim=1)
v = torch.amin(scaled_directions,dim=1)
# Force coordinates to be within the triangle
boundary_check = torch.where(u > d)
v[boundary_check] /= u[boundary_check]
u[boundary_check] = 1.0
# Precalculated barycentric values from linear matrix solve
I0 = 1 - u
I1 = u - v
I2 = v
# Make first selections and proper interps
selections = best[:,0] + 1
interps = I1
even_sels = torch.where(selections % 2 == 0)
interps[even_sels] = I2[even_sels] # Corners get different weights
# Make new edges for the central selections
central_edges = torch.clone(edge_index).to(edge_index.device)
central_selections = torch.zeros_like(selections)
central_interps = I0
# Make new edges for the last selection
pos_locs = torch.where(second_best==1)[0]
pos_edges = edge_index[:,pos_locs]
pos_selections = selections[pos_locs] + 1
pos_selections[torch.where(pos_selections>8)] = 1 #Account for wrap around
pos_interps = I1[pos_locs]
even_sels = torch.where(pos_selections % 2 == 0)
pos_interps[even_sels] = I2[pos_locs][even_sels]
neg_locs = torch.where(second_best==0)[0]
neg_edges = edge_index[:,neg_locs]
neg_selections = selections[neg_locs] - 1
neg_selections[torch.where(neg_selections<1)] = 8 # Account for wrap around
neg_interps = I1[neg_locs]
even_sels = torch.where(neg_selections % 2 == 0)
neg_interps[even_sels] = I2[neg_locs][even_sels]
# Account for the previously pruned same node edges
same_selections = torch.zeros(same_edges.shape[1],dtype=torch.long)
same_interps = torch.ones(same_edges.shape[1],dtype=torch.float)
# Combine
edge_index = torch.cat((edge_index,central_edges,pos_edges,neg_edges,same_edges),dim=1)
selections = torch.cat((selections,central_selections,pos_selections,neg_selections,same_selections),dim=0)
interps = torch.cat((interps,central_interps,pos_interps,neg_interps,same_interps),dim=0)
#edge_index = torch.cat((edge_index,central_edges,pos_edges,neg_edges),dim=1)
#selections = torch.cat((selections,central_selections,pos_selections,neg_selections),dim=0)
#interps = torch.cat((interps,central_interps,pos_interps,neg_interps),dim=0)
# Account for edges to the same node
#same_locs = torch.where(edge_index[0] == edge_index[1])
#selections[same_locs] = 0
#interps[same_locs] = 1
return edge_index,selections,interps
def normalizeEdges(edge_index,selections,interps=None,kernel_norm=False):
'''Given an edge_index and selections, normalize the edges for each node so that
aggregation of edges with interps = 1. If interps is given, use a weighted average.
if kernel_norm = True, account for missing selections by increasing weight on other selections.'''
N = torch.max(edge_index) + 1
S = torch.max(selections) + 1
total_weight = torch.zeros((N,S),dtype=torch.float).to(edge_index.device)
if interps is None:
interps = torch.ones(len(selections),dtype=torch.float).to(edge_index.device)
# Aggregate all edges to determine normalizations per selection
nodes = edge_index[0]
#total_weight[nodes,selections] += interps
total_weight.index_put_((nodes,selections),interps,accumulate=True)
# Reassign interps accordingly
if kernel_norm:
row_totals = torch.sum(total_weight,dim=1)
interps = interps * S/row_totals[nodes]
else:
norms = total_weight[nodes,selections]
norms[torch.where(norms < 1e-6)] = 1e-6 # Avoid divide by zero error
interps = interps/norms
return interps
def simplifyGraph(edge_index,selections,edge_lengths):
# Take the shortest edge for the set of the same selections on a given node
num_edges = edge_index.shape[1]
# Keep track of which nodes have been visited
keep_edges = torch.zeros(num_edges,dtype=torch.bool).to(edge_index.device)
previous_best_distance = 100000*torch.ones((torch.amax(edge_index)+1,torch.amax(selections)+1),dtype=torch.long).to(edge_index.device)
previous_best_edge = -1*torch.ones((torch.amax(edge_index)+1,torch.amax(selections)+1),dtype=torch.long).to(edge_index.device)
for i in range(num_edges):
start_node = edge_index[0,i]
#end_node = edge_index[1,i]
selection = selections[i]
distance = edge_lengths[i]
if distance < previous_best_distance[start_node,selection]:
previous_best_distance[start_node,selection] = distance
keep_edges[i] = True
prev = previous_best_edge[start_node,selection]
if prev != -1:
keep_edges[prev] = False
previous_best_edge[start_node,selection] = i
edge_index = edge_index[:,torch.where(keep_edges)[0]]
selections = selections[torch.where(keep_edges)]
return edge_index, selections