def num_nodes(self):
"""
Number of nodes
:return: Number of nodes
"""
return union_len(self.x)
def num_edges(self):
"""
Number of edges
:return: Number of edges
"""
return union_len(self.edge_index[0])
def topk_pool(source_index, score, k=None, ratio=None):
"""
:param sorted_source_index: index of source node (of edge) or source graph (of node)
:param sorted_score: 1-D Array
:param k:
:param ratio:
:return: sampled_edge_index, sampled_edge_score, sample_index
"""
if k is None and ratio is None:
raise Exception("you should provide either k or ratio for topk_pool")
elif k is not None and ratio is not None:
raise Exception(
"you should provide either k or ratio for topk_pool, not both of them"
)
# currently, we consider the source_index is not sorted
# the option is preserved for future performance optimization
source_index_sorted = False
if source_index_sorted:
sorted_source_index = source_index
# sort score by source_index
sorted_score = score
else:
source_index_perm = tf.argsort(source_index)
sorted_source_index = tf.gather(source_index, source_index_perm)
sorted_score = tf.gather(score, source_index_perm)
sorted_score = tf.reshape(sorted_score, [-1])
num_targets = union_len(sorted_source_index)
target_ones = tf.ones([num_targets], dtype=tf.int32)
num_targets_for_sources = tf.math.segment_sum(target_ones,
sorted_source_index)
# number of columns for score matrix
num_cols = tf.reduce_max(num_targets_for_sources)
# max index of source + 1
num_seen_sources = tf.shape(num_targets_for_sources)[0]
min_score = tf.reduce_min(sorted_score)
num_targets_before = tf.concat([
tf.zeros([1], dtype=tf.int32),
tf.math.cumsum(num_targets_for_sources)[:-1]
],
axis=0)
target_index_for_source = tf.range(0, num_targets) - tf.gather(
num_targets_before, sorted_source_index)
score_matrix = tf.cast(tf.fill([num_seen_sources, num_cols],
min_score - 1.0),
dtype=tf.float32)
score_index = tf.stack([sorted_source_index, target_index_for_source],
axis=1)
score_matrix = tf.tensor_scatter_nd_update(score_matrix, score_index,
sorted_score)
sort_index = tf.argsort(score_matrix, axis=-1, direction="DESCENDING")
if k is not None:
node_k = tf.math.minimum(
tf.cast(tf.fill([num_seen_sources], k), dtype=tf.int32),
num_targets_for_sources)
else:
node_k = tf.cast(tf.math.ceil(
tf.cast(num_targets_for_sources, dtype=tf.float32) *
tf.cast(ratio, dtype=tf.float32)),
dtype=tf.int32)
left_k_index = tf.TensorArray(tf.int32,
size=0,
dynamic_size=True,
element_shape=[2])
num_rows = tf.shape(node_k)[0]
current_size = 0
for row_index in range(num_rows):
num_cols = node_k[row_index]
for col_index in range(num_cols):
left_k_index = left_k_index.write(current_size,
[row_index, col_index])
current_size += 1
left_k_index = left_k_index.stack()
# left_k_index = [[row_index, col_index]
# for row_index, num_cols in enumerate(node_k)
# for col_index in range(num_cols)]
# left_k_index = tf.convert_to_tensor(left_k_index, dtype=tf.int32)
sample_col_index = tf.gather_nd(sort_index, left_k_index)
sample_row_index = left_k_index[:, 0]
topk_index = tf.gather(num_targets_before,
sample_row_index) + sample_col_index
if source_index_sorted:
return topk_index
else:
return tf.gather(source_index_perm, topk_index)
return topk_index
def topk_pool(source_index, score, k=None, ratio=None):
"""
:param source_index: index of source node (of edge) or source graph (of node)
:param score: 1-D Array
:param k:
:param ratio:
:return: sampled_edge_index, sampled_edge_score, sample_index
"""
if k is None and ratio is None:
raise Exception("you should provide either k or ratio for topk_pool")
elif k is not None and ratio is not None:
raise Exception(
"you should provide either k or ratio for topk_pool, not both of them"
)
num_targets = union_len(source_index)
target_ones = tf.ones([num_targets], dtype=tf.int32)
num_targets_for_sources = tf.math.segment_sum(target_ones, source_index)
# number of columns for score matrix
num_cols = tf.reduce_max(num_targets_for_sources)
# max index of source + 1
num_seen_sources = num_targets_for_sources.shape[0]
min_score = tf.reduce_min(score)
num_targets_before = tf.concat([
tf.zeros([1], dtype=tf.int32),
tf.math.cumsum(num_targets_for_sources)[:-1]
],
axis=0)
target_index_for_source = tf.range(0, num_targets) - tf.gather(
num_targets_before, source_index[0])
score_matrix = tf.cast(tf.fill([num_seen_sources, num_cols],
min_score - 1.0),
dtype=tf.float32)
score_index = tf.stack([source_index, target_index_for_source], axis=1)
score_matrix = tf.tensor_scatter_nd_update(score_matrix, score_index,
score)
sort_index = tf.argsort(score_matrix, axis=-1, direction="DESCENDING")
if k is not None:
node_k = tf.math.minimum(
tf.cast(tf.fill([num_seen_sources], k), dtype=tf.int32),
num_targets_for_sources)
else:
node_k = tf.cast(tf.math.ceil(
tf.cast(num_targets_for_sources, dtype=tf.float32) *
tf.cast(ratio, dtype=tf.float32)),
dtype=tf.int32)
left_k_index = [[row_index, col_index]
for row_index, num_cols in enumerate(node_k.numpy())
for col_index in range(num_cols)]
left_k_index = tf.convert_to_tensor(left_k_index, dtype=tf.int32)
sample_col_index = tf.gather_nd(sort_index, left_k_index)
sample_row_index = left_k_index[:, 0]
topk_index = tf.gather(num_targets_before,
sample_row_index) + sample_col_index
return topk_index
def num_edges(self):
return union_len(self.edge_index[0])
def num_nodes(self):
return union_len(self.x)
