Statistics
| Branch: | Revision:

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