(data, aggr: str = 'sink', k: int = 10)
| 159 | |
| 160 | |
| 161 | def get_pagerank_weights(data, aggr: str = 'sink', k: int = 10): |
| 162 | def _compute_pagerank(edge_index, damp: float = 0.85, k: int = 10): |
| 163 | num_nodes = edge_index.max().item() + 1 |
| 164 | deg_out = degree(edge_index[0]) |
| 165 | x = torch.ones((num_nodes,)).to(edge_index.device).to(torch.float32) |
| 166 | |
| 167 | for i in range(k): |
| 168 | edge_msg = x[edge_index[0]] / deg_out[edge_index[0]] |
| 169 | agg_msg = scatter(edge_msg, edge_index[1], reduce='sum') |
| 170 | |
| 171 | x = (1 - damp) * x + damp * agg_msg |
| 172 | |
| 173 | return x |
| 174 | |
| 175 | pv = _compute_pagerank(data.edge_index, k=k) |
| 176 | pv_row = pv[data.edge_index[0]].to(torch.float32) |
| 177 | pv_col = pv[data.edge_index[1]].to(torch.float32) |
| 178 | s_row = torch.log(pv_row) |
| 179 | s_col = torch.log(pv_col) |
| 180 | if aggr == 'sink': |
| 181 | s = s_col |
| 182 | elif aggr == 'source': |
| 183 | s = s_row |
| 184 | elif aggr == 'mean': |
| 185 | s = (s_col + s_row) * 0.5 |
| 186 | else: |
| 187 | s = s_col |
| 188 | |
| 189 | return normalize(s), pv |
| 190 | |
| 191 | |
| 192 | def drop_edge_by_weight(edge_index, weights, drop_prob: float, threshold: float = 0.7): |
nothing calls this directly
no test coverage detected