## mobicen / util / UnitDiskGraph.py @ b6158841

History | View | Annotate | Download (1.73 KB)

1 | e1cf8bea | LoreBz | ```
# https://stackoverflow.com/questions/32424604/find-all-nearest-neighbors-within-a-specific-distance
``` |
---|---|---|---|

2 | |||

3 | import networkx as nx |
||

4 | 1ef4948a | LoreBz | from scipy.spatial import KDTree |

5 | from scipy.spatial.distance import cdist |
||

6 | import code # code.interact(local=dict(globals(), **locals())) |
||

7 | import graph_tool as gt |
||

8 | e1cf8bea | LoreBz | |

9 | class UnitDiskGraph: |
||

10 | |||

11 | def __init__(self, points, radius): |
||

12 | 2dbcade9 | LoreBz | ```
self.nxG = nx.Graph()
``` |

13 | self.nxG.add_nodes_from(range(len(points))) |
||

14 | 1ef4948a | LoreBz | |

15 | 2dbcade9 | LoreBz | self.gtG = gt.Graph(directed=False) |

16 | 23c7ab1e | LoreBz | self.gtG.add_vertex(len(points)) |

17 | 2dbcade9 | LoreBz | self.edge_weights = self.gtG.new_edge_property('float') |

18 | self.gtG.edge_properties['weight'] = self.edge_weights |
||

19 | |||

20 | 1ef4948a | LoreBz | ```
self.generateGraphCDIST(points, radius)
``` |

21 | |||

22 | |||

23 | e1cf8bea | LoreBz | def genereateGraphFromKDtree(self, points, radius): |

24 | 1ef4948a | LoreBz | tree = KDTree(points) |

25 | e1cf8bea | LoreBz | edges = tree.query_pairs(r=radius) |

26 | 1ef4948a | LoreBz | ```
#edges = [e+(1.0,) for e in edges]
``` |

27 | ```
#pos = {k:points[k] for k in range(0,len(points))}
``` |
||

28 | ```
#self.G.add_weighted_edges_from(edges, weight='weight')
``` |
||

29 | for e in edges: |
||

30 | e = self.G.add_edge(e[0],e[1]) |
||

31 | self.edge_weights[e] = 1.0 |
||

32 | |||

33 | |||

34 | def generateGraphCDIST(self, points, radius): |
||

35 | ```
distM = cdist(points, points, 'euclidean')
``` |
||

36 | edges = [] |
||

37 | for r in range(len(points)): |
||

38 | for c in range(len(points)): |
||

39 | ```
if r==c:
``` |
||

40 | ```
continue
``` |
||

41 | ```
if distM[r][c] <= radius:
``` |
||

42 | 27b1b922 | LoreBz | ```
edges.append([r,c, 1.0])
``` |

43 | 2dbcade9 | LoreBz | self.nxG.add_weighted_edges_from(edges, weight='weight') |

44 | |||

45 | self.gtG.ep["weight"] = self.gtG.new_edge_property("float") |
||

46 | eprops = [self.gtG.ep["weight"]] |
||

47 | 23c7ab1e | LoreBz | ```
self.gtG.add_edge_list(edges, eprops=eprops)
``` |

48 | e1cf8bea | LoreBz | |

49 | def getGraph(self): |
||

50 | 2dbcade9 | LoreBz | return self.gtG, self.nxG |