mirror of https://github.com/gsi-upm/soil
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
21 lines
731 B
Python
21 lines
731 B
Python
3 years ago
|
from scipy.spatial import cKDTree as KDTree
|
||
3 years ago
|
import networkx as nx
|
||
2 years ago
|
from . import NetworkAgent
|
||
3 years ago
|
|
||
2 years ago
|
|
||
3 years ago
|
class Geo(NetworkAgent):
|
||
2 years ago
|
"""In this type of network, nodes have a "pos" attribute."""
|
||
3 years ago
|
|
||
1 year ago
|
def geo_search(self, radius, center=False, **kwargs):
|
||
2 years ago
|
"""Get a list of nodes whose coordinates are closer than *radius* to *node*."""
|
||
1 year ago
|
node = self.node_id
|
||
3 years ago
|
|
||
|
G = self.subgraph(**kwargs)
|
||
|
|
||
2 years ago
|
pos = nx.get_node_attributes(G, "pos")
|
||
3 years ago
|
if not pos:
|
||
|
return []
|
||
|
nodes, coords = list(zip(*pos.items()))
|
||
|
kdtree = KDTree(coords) # Cannot provide generator.
|
||
|
indices = kdtree.query_ball_point(pos[node], radius)
|
||
2 years ago
|
return [nodes[i] for i in indices if center or (nodes[i] != node)]
|