-
Notifications
You must be signed in to change notification settings - Fork 0
/
graph.py
157 lines (126 loc) · 5.9 KB
/
graph.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
import numpy as np
import torch
class NeighborFinder:
def __init__(self, adj_list, uniform=False):
"""
Params
------
node_idx_l: List[int]
node_ts_l: List[int]
off_set_l: List[int], such that node_idx_l[off_set_l[i]:off_set_l[i + 1]] = adjacent_list[i]
"""
node_idx_l, node_ts_l, edge_idx_l, off_set_l = self.init_off_set(adj_list)
self.node_idx_l = node_idx_l
self.node_ts_l = node_ts_l
self.edge_idx_l = edge_idx_l
self.off_set_l = off_set_l
self.uniform = uniform
def init_off_set(self, adj_list):
"""
Params
------
adj_list: List[List[int]]
"""
n_idx_l = []
n_ts_l = []
e_idx_l = []
off_set_l = [0]
for i in range(len(adj_list)):
curr = adj_list[i]
curr = sorted(curr, key=lambda x: x[1])
n_idx_l.extend([x[0] for x in curr])
e_idx_l.extend([x[1] for x in curr])
n_ts_l.extend([x[2] for x in curr])
off_set_l.append(len(n_idx_l))
n_idx_l = np.array(n_idx_l)
n_ts_l = np.array(n_ts_l)
e_idx_l = np.array(e_idx_l)
off_set_l = np.array(off_set_l)
assert(len(n_idx_l) == len(n_ts_l))
assert(off_set_l[-1] == len(n_ts_l))
return n_idx_l, n_ts_l, e_idx_l, off_set_l
def find_before(self, src_idx, cut_time):
"""
Params
------
src_idx: int
cut_time: float
"""
node_idx_l = self.node_idx_l
node_ts_l = self.node_ts_l
edge_idx_l = self.edge_idx_l
off_set_l = self.off_set_l
neighbors_idx = node_idx_l[off_set_l[src_idx]:off_set_l[src_idx + 1]]
neighbors_ts = node_ts_l[off_set_l[src_idx]:off_set_l[src_idx + 1]]
neighbors_e_idx = edge_idx_l[off_set_l[src_idx]:off_set_l[src_idx + 1]]
if len(neighbors_idx) == 0 or len(neighbors_ts) == 0:
return neighbors_idx, neighbors_ts, neighbors_e_idx
left = 0
right = len(neighbors_idx) - 1
while left + 1 < right:
mid = (left + right) // 2
curr_t = neighbors_ts[mid]
if curr_t < cut_time:
left = mid
else:
right = mid
if neighbors_ts[right] < cut_time:
return neighbors_idx[:right], neighbors_e_idx[:right], neighbors_ts[:right]
else:
return neighbors_idx[:left], neighbors_e_idx[:left], neighbors_ts[:left]
def get_temporal_neighbor(self, src_idx_l, cut_time_l, num_neighbors=20):
"""
Params
------
src_idx_l: List[int]
cut_time_l: List[float],
num_neighbors: int
"""
assert(len(src_idx_l) == len(cut_time_l))
out_ngh_node_batch = np.zeros((len(src_idx_l), num_neighbors)).astype(np.int32)
out_ngh_t_batch = np.zeros((len(src_idx_l), num_neighbors)).astype(np.float32)
out_ngh_eidx_batch = np.zeros((len(src_idx_l), num_neighbors)).astype(np.int32)
for i, (src_idx, cut_time) in enumerate(zip(src_idx_l, cut_time_l)):
ngh_idx, ngh_eidx, ngh_ts = self.find_before(src_idx, cut_time)
if len(ngh_idx) > 0:
if self.uniform:
sampled_idx = np.random.randint(0, len(ngh_idx), num_neighbors)
out_ngh_node_batch[i, :] = ngh_idx[sampled_idx]
out_ngh_t_batch[i, :] = ngh_ts[sampled_idx]
out_ngh_eidx_batch[i, :] = ngh_eidx[sampled_idx]
# resort based on time
pos = out_ngh_t_batch[i, :].argsort()
out_ngh_node_batch[i, :] = out_ngh_node_batch[i, :][pos]
out_ngh_t_batch[i, :] = out_ngh_t_batch[i, :][pos]
out_ngh_eidx_batch[i, :] = out_ngh_eidx_batch[i, :][pos]
else:
ngh_ts = ngh_ts[:num_neighbors]
ngh_idx = ngh_idx[:num_neighbors]
ngh_eidx = ngh_eidx[:num_neighbors]
assert(len(ngh_idx) <= num_neighbors)
assert(len(ngh_ts) <= num_neighbors)
assert(len(ngh_eidx) <= num_neighbors)
out_ngh_node_batch[i, num_neighbors - len(ngh_idx):] = ngh_idx
out_ngh_t_batch[i, num_neighbors - len(ngh_ts):] = ngh_ts
out_ngh_eidx_batch[i, num_neighbors - len(ngh_eidx):] = ngh_eidx
return out_ngh_node_batch, out_ngh_eidx_batch, out_ngh_t_batch
def find_k_hop(self, k, src_idx_l, cut_time_l, num_neighbors=20):
"""Sampling the k-hop sub graph
"""
x, y, z = self.get_temporal_neighbor(src_idx_l, cut_time_l, num_neighbors)
node_records = [x]
eidx_records = [y]
t_records = [z]
for _ in range(k -1):
ngn_node_est, ngh_t_est = node_records[-1], t_records[-1] # [N, *([num_neighbors] * (k - 1))]
orig_shape = ngn_node_est.shape
ngn_node_est = ngn_node_est.flatten()
ngn_t_est = ngh_t_est.flatten()
out_ngh_node_batch, out_ngh_eidx_batch, out_ngh_t_batch = self.get_temporal_neighbor(ngn_node_est, ngn_t_est, num_neighbors)
out_ngh_node_batch = out_ngh_node_batch.reshape(*orig_shape, num_neighbors) # [N, *([num_neighbors] * k)]
out_ngh_eidx_batch = out_ngh_eidx_batch.reshape(*orig_shape, num_neighbors)
out_ngh_t_batch = out_ngh_t_batch.reshape(*orig_shape, num_neighbors)
node_records.append(out_ngh_node_batch)
eidx_records.append(out_ngh_eidx_batch)
t_records.append(out_ngh_t_batch)
return node_records, eidx_records, t_records