1
0
mirror of https://github.com/gsi-upm/soil synced 2025-07-12 07:52:22 +00:00
soil/soil/agents/Geo.py
J. Fernando Sánchez 2869b1e1e6 Clean-up
* Removed old/unnecessary models
* Added a `simulation.{iter_}from_py` method to load simulations from python
files
* Changed tests of examples to run programmatic simulations
* Fixed programmatic examples
2022-11-13 20:31:05 +01:00

21 lines
741 B
Python

from scipy.spatial import cKDTree as KDTree
import networkx as nx
from . import NetworkAgent
class Geo(NetworkAgent):
"""In this type of network, nodes have a "pos" attribute."""
def geo_search(self, radius, agent=None, center=False, **kwargs):
"""Get a list of nodes whose coordinates are closer than *radius* to *node*."""
node = agent.node
G = self.subgraph(**kwargs)
pos = nx.get_node_attributes(G, "pos")
if not pos:
return []
nodes, coords = list(zip(*pos.items()))
kdtree = KDTree(coords) # Cannot provide generator.
indices = kdtree.query_ball_point(pos[node], radius)
return [nodes[i] for i in indices if center or (nodes[i] != node)]