from __future__ import annotations import os import sqlite3 import math import logging import inspect from typing import Any, Callable, Dict, Optional, Union, List, Type from collections import namedtuple from time import time as current_time from copy import deepcopy import networkx as nx from mesa import Model, Agent from . import agents as agentmod, datacollection, utils, time, network, events # TODO: maybe add metaclass to read attributes of a model class BaseEnvironment(Model): """ The environment is key in a simulation. It controls how agents interact, and what information is available to them. This is an opinionated version of `mesa.Model` class, which adds many convenience methods and abstractions. The environment parameters and the state of every agent can be accessed both by using the environment as a dictionary and with the environment's :meth:`soil.environment.Environment.get` method. """ collector_class = datacollection.SoilCollector schedule_class = time.TimedActivation def __new__(cls, *args: Any, seed="default", dir_path=None, collector_class: type = None, agent_reporters: Optional[Any] = None, model_reporters: Optional[Any] = None, tables: Optional[Any] = None, **kwargs: Any) -> Any: """Create a new model with a default seed value""" self = super().__new__(cls, *args, seed=seed, **kwargs) self.dir_path = dir_path or os.getcwd() collector_class = collector_class or cls.collector_class self.datacollector = collector_class( model_reporters=model_reporters, agent_reporters=agent_reporters, tables=tables, ) for k in dir(cls): v = getattr(cls, k) if isinstance(v, property): v = v.fget if getattr(v, "add_to_report", False): self.add_model_reporter(k, k) return self def __init__( self, *, id="unnamed_env", seed="default", dir_path=None, schedule=None, schedule_class=None, logger = None, agents: Optional[Dict] = None, collector_class: type = datacollection.SoilCollector, agent_reporters: Optional[Any] = None, model_reporters: Optional[Any] = None, tables: Optional[Any] = None, init: bool = True, **env_params, ): super().__init__() self.current_id = -1 self.id = id if logger: self.logger = logger else: self.logger = utils.logger.getChild(self.id) self.schedule = schedule if schedule is None: if schedule_class is None: schedule_class = self.schedule_class self.schedule = schedule_class(self) for (k, v) in env_params.items(): self[k] = v if agents: self.add_agents(**agents) if init: self.init() self.datacollector.collect(self) def init(self): pass @property def agents(self): return agentmod.AgentView(self.schedule._agents) def agent(self, *args, **kwargs): return agentmod.AgentView(self.schedule._agents).one(*args, **kwargs) def count_agents(self, *args, **kwargs): return sum(1 for i in self.agents(*args, **kwargs)) def agent_df(self, steps=False): df = self.datacollector.get_agent_vars_dataframe() if steps: df.index.rename(["step", "agent_id"], inplace=True) return df model_df = self.datacollector.get_model_vars_dataframe() df.index = df.index.set_levels(model_df.time, level=0).rename(["time", "agent_id"]) return df def model_df(self, steps=False): df = self.datacollector.get_model_vars_dataframe() if steps: return df df.index.rename("step", inplace=True) return df.reset_index().set_index("time") @property def now(self): if self.schedule: return self.schedule.time raise Exception( "The environment has not been scheduled, so it has no sense of time" ) def init_agents(self): pass def add_agent(self, agent_class, unique_id=None, **agent): if unique_id is None: unique_id = self.next_id() agent["unique_id"] = unique_id agent = dict(**agent) unique_id = agent.pop("unique_id", None) if unique_id is None: unique_id = self.next_id() a = agent_class(unique_id=unique_id, model=self, **agent) self.schedule.add(a) return a def add_agents(self, agent_classes: List[type], k, weights: Optional[List[float]] = None, **kwargs): if isinstance(agent_classes, type): agent_classes = [agent_classes] if weights is None: weights = [1] * len(agent_classes) for cls in self.random.choices(agent_classes, weights=weights, k=k): self.add_agent(agent_class=cls, **kwargs) def log(self, message, *args, level=logging.INFO, **kwargs): if not self.logger.isEnabledFor(level): return message = message + " ".join(str(i) for i in args) message = " @{:>3}: {}".format(self.now, message) for k, v in kwargs: message += " {k}={v} ".format(k, v) extra = {} extra["now"] = self.now extra["id"] = self.id return self.logger.log(level, message, extra=extra) def step(self): """ Advance one step in the simulation, and update the data collection and scheduler appropriately """ super().step() self.schedule.step() self.datacollector.collect(self) if self.logger.isEnabledFor(logging.DEBUG): msg = "Model data:\n" max_width = max(len(k) for k in self.datacollector.model_vars.keys()) for (k, v) in self.datacollector.model_vars.items(): msg += f"\t{k:<{max_width}}: {v[-1]:>6}\n" self.logger.debug(f"--- Steps: {self.schedule.steps:^5} - Time: {self.now:^5} --- " + msg) def add_model_reporter(self, name, func=None): if not func: func = name self.datacollector._new_model_reporter(name, func) def add_agent_reporter(self, name, reporter=None, agent_type=None): if not agent_type and not reporter: reporter = name elif agent_type: reporter = lambda a: reporter(a) if isinstance(a, agent_type) else None self.datacollector._new_agent_reporter(name, reporter) @classmethod def run(cls, *, name=None, iterations=1, num_processes=1, **kwargs): from .simulation import Simulation return Simulation(name=name or cls.__name__, model=cls, iterations=iterations, num_processes=num_processes, **kwargs).run() def __getitem__(self, key): try: return getattr(self, key) except AttributeError: raise KeyError(f"key {key} not found in environment") def __delitem__(self, key): return delattr(self, key) def __contains__(self, key): return hasattr(self, key) def __setitem__(self, key, value): setattr(self, key, value) def __str__(self): return str(dict(self)) def __len__(self): return sum(1 for n in self.keys()) def __iter__(self): return iter(self.agents()) def get(self, key, default=None): return self[key] if key in self else default def keys(self): return (k for k in self.__dict__ if k[0] != "_") class NetworkEnvironment(BaseEnvironment): """ The NetworkEnvironment is an environment that includes one or more networkx.Graph intances and methods to associate agents to nodes and vice versa. """ def __init__(self, *args, topology: Optional[Union[nx.Graph, str]] = None, agent_class: Optional[Type[agentmod.Agent]] = None, network_generator: Optional[Callable] = None, network_params: Optional[Dict] = {}, init=True, **kwargs): self.topology = topology self.network_generator = network_generator self.network_params = network_params if topology or network_params or network_generator: self.create_network(topology, generator=network_generator, **network_params) else: self.G = nx.Graph() super().__init__(*args, **kwargs, init=False) self.agent_class = agent_class if self.agent_class: self.populate_network(self.agent_class) self._check_agent_nodes() if init: self.init() self.datacollector.collect(self) def add_agent(self, agent_class, *args, node_id=None, topology=None, **kwargs): if node_id is None and topology is None: return super().add_agent(agent_class, *args, **kwargs) try: a = super().add_agent(agent_class, *args, node_id=node_id, **kwargs) except TypeError: self.logger.warning(f"Agent constructor for {agent_class} does not have a node_id attribute. Might be a bug.") a = super().add_agent(agent_class, *args, **kwargs) self.G.nodes[node_id]["agent"] = a return a def add_agents(self, *args, k=None, **kwargs): if not k and not self.G: raise ValueError("Cannot add agents to an empty network") super().add_agents(*args, k=k or len(self.G), **kwargs) def create_network(self, topology=None, generator=None, path=None, **network_params): if topology is not None: topology = network.from_topology(topology, dir_path=self.dir_path) elif path is not None: topology = network.from_topology(path, dir_path=self.dir_path) elif generator is not None: params = dict(generator=generator, dir_path=self.dir_path, seed=self.random, **network_params) try: topology = network.from_params(**params) except TypeError: del params["seed"] topology = network.from_params(**params) else: raise ValueError("topology must be a networkx.Graph or a string, or network_generator must be provided") self.G = topology def init_agents(self, *args, **kwargs): """Initialize the agents from a""" super().init_agents(*args, **kwargs) @property def network_agents(self): """Return agents still alive and assigned to a node in the network.""" for (id, data) in self.G.nodes(data=True): if "agent" in data: agent = data["agent"] if getattr(agent, "alive", True): yield agent def add_node(self, agent_class, unique_id=None, node_id=None, **kwargs): if unique_id is None: unique_id = self.next_id() if node_id is None: node_id = network.find_unassigned( G=self.G, shuffle=True, random=self.random ) if node_id is None: node_id = f"node_for_{unique_id}" if node_id not in self.G.nodes: self.G.add_node(node_id) assert "agent" not in self.G.nodes[node_id] a = self.add_agent( unique_id=unique_id, agent_class=agent_class, topology=self.G, node_id=node_id, **kwargs, ) a["visible"] = True return a def _check_agent_nodes(self): """ Detect nodes that have agents assigned to them. """ for (id, data) in self.G.nodes(data=True): if "agent_id" in data: agent = self.agents(data["agent_id"]) self.G.nodes[id]["agent"] = agent assert not getattr(agent, "node_id", None) or agent.node_id == id agent.node_id = id for agent in self.agents(): if hasattr(agent, "node_id"): node_id = agent["node_id"] if node_id not in self.G.nodes: raise ValueError(f"Agent {agent} is assigned to node {agent.node_id} which is not in the network") node = self.G.nodes[node_id] if node.get("agent") is not None and node["agent"] != agent: raise ValueError(f"Node {node_id} already has a different agent assigned to it") self.G.nodes[node_id]["agent"] = agent def add_agents(self, agent_classes: List[type], k=None, weights: Optional[List[float]] = None, **kwargs): if k is None: k = len(self.G) if not k: raise ValueError("Cannot add agents to an empty network") super().add_agents(agent_classes, k=k, weights=weights, **kwargs) def agent_for_node_id(self, node_id): return self.G.nodes[node_id].get("agent") def populate_network(self, agent_class: List[Model], weights: List[float] = None, **agent_params): if isinstance(agent_class, type): agent_class = [agent_class] else: agent_class = list(agent_class) if not weights: weights = [1] * len(agent_class) assert len(self.G) classes = self.random.choices(agent_class, weights, k=len(self.G)) toadd = [] for (cls, (node_id, node)) in zip(classes, self.G.nodes(data=True)): if "agent" in node: continue node["agent"] = None # Reserve toadd.append(dict(node_id=node_id, topology=self.G, agent_class=cls, **agent_params)) for d in toadd: a = self.add_agent(**d) self.G.nodes[d["node_id"]]["agent"] = a assert all("agent" in node for (_, node) in self.G.nodes(data=True)) assert len(list(self.network_agents)) class EventedEnvironment(BaseEnvironment): def broadcast(self, msg, sender=None, expiration=None, ttl=None, **kwargs): for agent in self.agents(**kwargs): if agent == sender: continue self.logger.debug(f"Telling {repr(agent)}: {msg} ttl={ttl}") try: inbox = agent._inbox except AttributeError: self.logger.info( f"Agent {agent.unique_id} cannot receive events because it does not have an inbox" ) continue # Allow for AttributeError exceptions in this part of the code inbox.append( events.Tell( payload=msg, sender=sender, expiration=expiration if ttl is None else self.now + ttl, ) ) class Environment(NetworkEnvironment, EventedEnvironment): """Default environment class, has both network and event capabilities"""