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.
89 lines
2.7 KiB
Python
89 lines
2.7 KiB
Python
from . import BaseAgent
|
|
|
|
|
|
class NetworkAgent(BaseAgent):
|
|
def __init__(self, *args, topology, node_id, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
assert topology is not None
|
|
assert node_id is not None
|
|
self.G = topology
|
|
assert self.G
|
|
self.node_id = node_id
|
|
|
|
def count_neighbors(self, state_id=None, **kwargs):
|
|
return len(self.get_neighbors(state_id=state_id, **kwargs))
|
|
|
|
def iter_neighbors(self, **kwargs):
|
|
return self.iter_agents(limit_neighbors=True, **kwargs)
|
|
|
|
def get_neighbors(self, **kwargs):
|
|
return list(self.iter_neighbors())
|
|
|
|
@property
|
|
def node(self):
|
|
return self.G.nodes[self.node_id]
|
|
|
|
def iter_agents(self, unique_id=None, *, limit_neighbors=False, **kwargs):
|
|
unique_ids = None
|
|
if isinstance(unique_id, list):
|
|
unique_ids = set(unique_id)
|
|
elif unique_id is not None:
|
|
unique_ids = set(
|
|
[
|
|
unique_id,
|
|
]
|
|
)
|
|
|
|
if limit_neighbors:
|
|
neighbor_ids = set()
|
|
for node_id in self.G.neighbors(self.node_id):
|
|
if self.G.nodes[node_id].get("agent") is not None:
|
|
neighbor_ids.add(node_id)
|
|
if unique_ids:
|
|
unique_ids = unique_ids & neighbor_ids
|
|
else:
|
|
unique_ids = neighbor_ids
|
|
if not unique_ids:
|
|
return
|
|
unique_ids = list(unique_ids)
|
|
yield from super().iter_agents(unique_id=unique_ids, **kwargs)
|
|
|
|
def subgraph(self, center=True, **kwargs):
|
|
include = [self] if center else []
|
|
G = self.G.subgraph(
|
|
n.node_id for n in list(self.get_agents(**kwargs) + include)
|
|
)
|
|
return G
|
|
|
|
def remove_node(self):
|
|
self.debug(f"Removing node for {self.unique_id}: {self.node_id}")
|
|
self.G.remove_node(self.node_id)
|
|
self.node_id = None
|
|
|
|
def add_edge(self, other, edge_attr_dict=None, *edge_attrs):
|
|
if self.node_id not in self.G.nodes(data=False):
|
|
raise ValueError(
|
|
"{} not in list of existing agents in the network".format(
|
|
self.unique_id
|
|
)
|
|
)
|
|
if other.node_id not in self.G.nodes(data=False):
|
|
raise ValueError(
|
|
"{} not in list of existing agents in the network".format(other)
|
|
)
|
|
|
|
self.G.add_edge(
|
|
self.node_id, other.node_id, edge_attr_dict=edge_attr_dict, *edge_attrs
|
|
)
|
|
|
|
def die(self, remove=True):
|
|
if not self.alive:
|
|
return None
|
|
if remove:
|
|
self.remove_node()
|
|
return super().die()
|
|
|
|
|
|
NetAgent = NetworkAgent
|