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.
256 lines
9.3 KiB
Python
256 lines
9.3 KiB
Python
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()
|