import os import time import imp import sys import yaml import networkx as nx from networkx.readwrite import json_graph from copy import deepcopy import pickle from nxsim import NetworkSimulation from . import agents, utils, environment, basestring from .utils import logger class SoilSimulation(NetworkSimulation): """ Subclass of nsim.NetworkSimulation with three main differences: 1) agent type can be specified by name or by class. 2) instead of just one type, an network_agents can be used. The distribution specifies the weight (or probability) of each agent type in the topology. This is an example distribution: :: [ {'agent_type': 'agent_type_1', 'weight': 0.2, 'state': { 'id': 0 } }, {'agent_type': 'agent_type_2', 'weight': 0.8, 'state': { 'id': 1 } } ] In this example, 20% of the nodes will be marked as type 'agent_type_1'. 3) if no initial state is given, each node's state will be set to `{'id': 0}`. """ def __init__(self, name=None, topology=None, network_params=None, network_agents=None, agent_type=None, states=None, default_state=None, interval=1, dump=False, dir_path=None, num_trials=1, max_time=100, agent_module=None, load_module=None, seed=None, environment_agents=None, environment_params=None): if topology is None: topology = utils.load_network(network_params, dir_path=dir_path) elif isinstance(topology, basestring) or isinstance(topology, dict): topology = json_graph.node_link_graph(topology) self.load_module = load_module self.topology = nx.Graph(topology) self.network_params = network_params self.name = name or 'UnnamedSimulation' self.num_trials = num_trials self.max_time = max_time self.default_state = default_state or {} self.dir_path = dir_path or os.getcwd() self.interval = interval self.seed = str(seed) or str(time.time()) self.dump = dump self.environment_params = environment_params or {} if load_module: path = sys.path + [self.dir_path] f, fp, desc = imp.find_module(load_module, path) imp.load_module('soil.agents.custom', f, fp, desc) environment_agents = environment_agents or [] self.environment_agents = self._convert_agent_types(environment_agents) distro = self.calculate_distribution(network_agents, agent_type) self.network_agents = self._convert_agent_types(distro) self.states = self.validate_states(states, self.topology) def calculate_distribution(self, network_agents=None, agent_type=None): if network_agents: network_agents = deepcopy(network_agents) elif agent_type: network_agents = [{'agent_type': agent_type}] else: return [] # Calculate the thresholds total = sum(x.get('weight', 1) for x in network_agents) acc = 0 for v in network_agents: upper = acc + (v.get('weight', 1)/total) v['threshold'] = [acc, upper] acc = upper return network_agents def serialize_distribution(self): d = self._convert_agent_types(self.network_agents, to_string=True) for v in d: if 'threshold' in v: del v['threshold'] return d def _convert_agent_types(self, ind, to_string=False): d = deepcopy(ind) for v in d: agent_type = v['agent_type'] if to_string and not isinstance(agent_type, str): v['agent_type'] = str(agent_type.__name__) elif not to_string and isinstance(agent_type, str): v['agent_type'] = agents.agent_types[agent_type] return d def validate_states(self, states, topology): states = states or [] # Validate states to avoid ignoring states during # initialization if isinstance(states, dict): for x in states: assert x in self.topology.node else: assert len(states) <= len(self.topology) return states def run_simulation(self): return self.run() def run(self): return list(self.run_simulation_gen()) def run_simulation_gen(self, *args, **kwargs): with utils.timer('simulation'): for i in range(self.num_trials): res = self.run_trial(i) if self.dump: res.dump_gexf(self.dir_path) res.dump_csv(self.dir_path) yield res if self.dump: logger.info('Dumping results to {}'.format(self.dir_path)) self.dump_pickle(self.dir_path) self.dump_yaml(self.dir_path) else: logger.info('NOT dumping results') def run_trial(self, trial_id=0, dump=False, dir_path=None): """Run a single trial of the simulation Parameters ---------- trial_id : int """ # Set-up trial environment and graph logger.info('Trial: {}'.format(trial_id)) env_name = '{}_trial_{}'.format(self.name, trial_id) env = environment.SoilEnvironment(name=env_name, topology=self.topology.copy(), seed=self.seed+env_name, initial_time=0, dump=self.dump, interval=self.interval, network_agents=self.network_agents, states=self.states, default_state=self.default_state, environment_agents=self.environment_agents, simulation=self, **self.environment_params) # Set up agents on nodes logger.info('\tRunning') with utils.timer('trial'): env.run(until=self.max_time) return env def to_dict(self): return self.__getstate__() def to_yaml(self): return yaml.dump(self.to_dict()) def dump_yaml(self, dir_path=None, file_name=None): dir_path = dir_path or self.dir_path if not os.path.exists(dir_path): os.makedirs(dir_path) if not file_name: file_name = os.path.join(dir_path, '{}.dumped.yml'.format(self.name)) with open(file_name, 'w') as f: f.write(self.to_yaml()) def dump_pickle(self, dir_path=None, pickle_name=None): dir_path = dir_path or self.dir_path if not os.path.exists(dir_path): os.makedirs(dir_path) if not pickle_name: pickle_name = os.path.join(dir_path, '{}.simulation.pickle'.format(self.name)) with open(pickle_name, 'wb') as f: pickle.dump(self, f) def __getstate__(self): state = self.__dict__.copy() state['topology'] = json_graph.node_link_data(self.topology) state['network_agents'] = self.serialize_distribution() state['environment_agents'] = self._convert_agent_types(self.environment_agents, to_string=True) return state def __setstate__(self, state): self.__dict__ = state self.topology = json_graph.node_link_graph(state['topology']) self.network_agents = self._convert_agent_types(self.network_agents) self.environment_agents = self._convert_agent_types(self.environment_agents) return state def from_config(config, G=None): config = list(utils.load_config(config)) if len(config) > 1: raise AttributeError('Provide only one configuration') config = config[0][0] sim = SoilSimulation(**config) return sim def run_from_config(*configs, dump=True, results_dir=None, timestamp=False): if not results_dir: results_dir = 'soil_output' for config_def in configs: for config, cpath in utils.load_config(config_def): name = config.get('name', 'unnamed') logger.info("Using config(s): {name}".format(name=name)) sim = SoilSimulation(**config) if timestamp: sim_folder = '{}_{}'.format(sim.name, time.strftime("%Y-%m-%d_%H:%M:%S")) else: sim_folder = sim.name sim.dir_path = os.path.join(results_dir, sim_folder) sim.dump = dump logger.info('Dumping results to {} : {}'.format(sim.dir_path, dump)) results = sim.run_simulation()