| 211 | |
| 212 | |
| 213 | class Distance(nn.Module): |
| 214 | def __init__(self, cutoff, max_num_neighbors=32, loop=True): |
| 215 | super(Distance, self).__init__() |
| 216 | self.cutoff = cutoff |
| 217 | self.max_num_neighbors = max_num_neighbors |
| 218 | self.loop = loop |
| 219 | |
| 220 | def forward(self, pos, batch): |
| 221 | edge_index = radius_graph(pos, r=self.cutoff, batch=batch, loop=self.loop, max_num_neighbors=self.max_num_neighbors) |
| 222 | edge_vec = pos[edge_index[0]] - pos[edge_index[1]] |
| 223 | |
| 224 | if self.loop: |
| 225 | mask = edge_index[0] != edge_index[1] |
| 226 | edge_weight = torch.zeros(edge_vec.size(0), device=edge_vec.device) |
| 227 | edge_weight[mask] = torch.norm(edge_vec[mask], dim=-1) |
| 228 | else: |
| 229 | edge_weight = torch.norm(edge_vec, dim=-1) |
| 230 | |
| 231 | return edge_index, edge_weight, edge_vec |
| 232 | |
| 233 | |
| 234 | class NeighborEmbedding(MessagePassing): |