69 lines
2.1 KiB
Python
Executable File
69 lines
2.1 KiB
Python
Executable File
import math, random, os, time, sys, io
|
|
import json
|
|
import numpy as np
|
|
import networkx as nx
|
|
import scipy
|
|
import time
|
|
from networkx.drawing.nx_pydot import write_dot
|
|
|
|
def clean_shuffle_graph(G):
|
|
random_seed_state = int(random.random()*100000) # quick hack to go back to random afterwards
|
|
random.seed(42)
|
|
node_mapping = dict(zip(sorted(G.nodes()), sorted(G.nodes(), key=lambda _: random.random()))) # maybe sorted not really deterministic
|
|
G = nx.relabel_nodes(G, node_mapping)
|
|
G = nx.convert_node_labels_to_integers(G)
|
|
if not nx.is_connected(G):
|
|
print('Graph is not connected, try a differnt one.')
|
|
assert(nx.is_connected(G))
|
|
random.seed(random_seed_state)
|
|
return G
|
|
|
|
def gen_sis(G, steps = 1000, inf_rate=1.0, rec_rate=2.0, noise=0.1, statesList = None):
|
|
S = [1., 0.]
|
|
I = [0., 1.]
|
|
states = [random.choice([S, I]) for i in range(G.number_of_nodes())]
|
|
for _ in range(steps):
|
|
rates = np.zeros(G.number_of_nodes())
|
|
for n in range(G.number_of_nodes()):
|
|
rates[n] = noise
|
|
if states[n] == I:
|
|
rates[n] += rec_rate
|
|
if states[n] == S:
|
|
rates[n] += inf_rate * len([n_j for n_j in G.neighbors(n) if states[n_j] == I])
|
|
rates[n] = 1.0/rates[n] # numpy uses mean as rate param
|
|
jump_time = np.random.exponential(rates)
|
|
change_n = np.argmin(jump_time)
|
|
states[change_n] = S if states[change_n] == I else I
|
|
statesList.append(states)
|
|
return states
|
|
|
|
statesList = []
|
|
steps = 1000 + random.choice(range(1000))
|
|
G_grid10x10 = nx.grid_2d_graph(10,10)
|
|
G = clean_shuffle_graph(G_grid10x10)
|
|
TS_data = gen_sis(G, steps=steps,statesList=statesList)
|
|
# print(G.number_of_nodes())
|
|
# f = open("graph.dot", "rw")
|
|
# nx.drawing.nx_pydot.write_dot(G,f)
|
|
# f.write("\n")
|
|
# f.close()
|
|
# Writing the json file to stdout
|
|
# f = StringIO("")
|
|
dotGraph = ""
|
|
|
|
#Write dot file to string
|
|
with io.StringIO() as f:
|
|
write_dot(G, f)
|
|
f.seek(0)
|
|
dotGraph = f.read()
|
|
|
|
output = {
|
|
"name" : "testGraph",
|
|
"dotGraph" : dotGraph,
|
|
"steps" : steps,
|
|
"states" : statesList,
|
|
}
|
|
|
|
jsonOutput = json.dumps(output)
|
|
print(jsonOutput)
|