1
0
mirror of https://github.com/gsi-upm/soil synced 2024-09-21 07:41:43 +00:00
soil/examples/custom_generator/generator_sim.py

39 lines
1004 B
Python
Raw Normal View History

from networkx import Graph
2022-10-13 20:43:16 +00:00
import random
import networkx as nx
from soil import Simulation, Environment, CounterModel, parameters
2022-10-17 18:23:57 +00:00
def mygenerator(n=5, n_edges=5):
2022-10-17 18:23:57 +00:00
"""
Just a simple generator that creates a network with n nodes and
n_edges edges. Edges are assigned randomly, only avoiding self loops.
2022-10-17 18:23:57 +00:00
"""
G = nx.Graph()
for i in range(n):
G.add_node(i)
2022-10-17 18:23:57 +00:00
for i in range(n_edges):
nodes = list(G.nodes)
2022-10-13 20:43:16 +00:00
n_in = random.choice(nodes)
nodes.remove(n_in) # Avoid loops
2022-10-13 20:43:16 +00:00
n_out = random.choice(nodes)
G.add_edge(n_in, n_out)
return G
class GeneratorEnv(Environment):
"""Using a custom generator for the network"""
generator: parameters.function = staticmethod(mygenerator)
def init(self):
self.create_network(generator=self.generator, n=10, n_edges=5)
self.add_agents(CounterModel)
sim = Simulation(model=GeneratorEnv, max_steps=10, interval=1)
if __name__ == '__main__':
sim.run(dump=False)