求torch_geometric.nn.radius_graph()的TensorFlow等效实现方法

直接用TensorFlow原生张量运算就能实现简洁的等价功能，不需要额外工具类，核心逻辑是计算节点间距离并筛选符合阈值的边对，代码如下：

import tensorflow as tf

def radius_graph(coords, r, loop=False):
    # coords: 节点坐标张量，形状 [节点数N, 维度D]
    # r: 距离阈值，仅保留距离≤r的节点间边
    # loop: 是否保留节点到自身的自环边
    
    # 计算所有节点对的坐标差（广播实现）
    diff = tf.expand_dims(coords, 0) - tf.expand_dims(coords, 1)
    # 计算欧氏距离的平方（避免开根号，提升效率并减少浮点误差）
    dist_sq = tf.reduce_sum(tf.square(diff), axis=-1)
    
    # 生成符合距离条件的掩码
    mask = dist_sq <= tf.square(r)
    if not loop:
        # 移除自环边（对角线位置的节点对）
        mask = tf.logical_and(mask, tf.logical_not(tf.eye(tf.shape(coords)[0], dtype=tf.bool)))
    
    # 提取满足条件的边索引，整理为[2, 边数E]的格式
    row, col = tf.where(mask)
    edge_index = tf.stack([row, col], axis=0)
    
    return edge_index

使用示例 #

# 生成10个3维随机节点坐标
node_coords = tf.random.normal((10, 3))
# 生成距离≤1.0的节点边，不保留自环
edge_index = radius_graph(node_coords, r=1.0, loop=False)
print("边索引形状:", edge_index.shape)  # 输出格式为[2, E]，E为符合条件的边数量

优化说明 #

• 用距离平方代替欧氏距离直接计算，既提升运算效率，又减少浮点精度损失
• 中小规模节点场景下，全量广播计算的方式足够简洁高效；若处理超大规模节点集，可结合tf.math.top_k或分块计算优化内存占用

内容的提问来源于stack exchange，提问作者playerJX1