import networkx as nx
import random

from ndlib.models.CompositeModel import CompositeModel
from ndlib.models.ContinuousModel import ContinuousModel
from ndlib.models.compartments.NodeStochastic import NodeStochastic
from ndlib.models.compartments.NodeNumericalAttribute import NodeNumericalAttribute

import ndlib.models.ModelConfig as mc

from bokeh.io import show
from ndlib.viz.bokeh.DiffusionTrend import DiffusionTrend

def intial_status(node, graph):
    addiction = 0
    return addiction

def craving_model(node, graph, current_val, variable):
    craving = nx.get_node_attributes(graph, 'craving')[node]
    self_control = nx.get_node_attributes(graph, 'self_control')[node]
    return min(current_val + craving - self_control, 1)

# Network definition
g = nx.erdos_renyi_graph(n=10, p=0.1)

# Extra network setup
attr = {n: {'craving': random.random(), 'self_control': random.random()} for n in g.nodes()}
nx.set_node_attributes(g, attr)

print('initial nodes:')
for n in g.nodes(data=True):
    print(n)
print()

# Model definition
addiction_model = ContinuousModel(g)
addiction_model.add_status('addiction')

# Compartments
condition = NodeNumericalAttribute('self_control', value='craving', op='<')

# Rules
addiction_model.add_rule('addiction', craving_model, condition)

# Configuration
config = mc.Configuration()
config.add_model_parameter('fraction_infected', 0.1)
addiction_model.set_initial_status(intial_status, config)

# Simulation
iterations = addiction_model.iteration_bunch(10, node_status=True)
print()
print(iterations)
print()
@@ -1,39 +1,191 @@
-import networkx as nx%0Aimport random
+# TODO Dynamic graph, plots, attribute updating, multiple status implementation (?), optimize speed%0A%0Aimport networkx as nx%0Aimport random%0Aimport matplotlib.pyplot as plt%0Aimport numpy as np
 %0A%0Afr
@@ -937,16 +937,17 @@
 aph(n=10
+0
 , p=0.1)
@@ -1647,17 +1647,18 @@
 n_bunch(
-1
+20
 0, node_
@@ -1670,16 +1670,17 @@
 s=True)%0A
+%0A
 print()%0A
@@ -1696,16 +1696,334 @@
 rations)%0Aprint()
+%0A%0A### Plots / data manipulation%0A# Mean status delta per iterations%0Ameans = %5B%5D%0Afor it in iterations:%0A    deltas = list(it%5B'status_delta'%5D.values())%0A    if len(deltas) %3E 0:%0A        means.append(sum(deltas) / len(deltas))%0A    else:%0A        means.append(0)%0A%0Ax = np.arange(0, len(iterations))%0A%0Aplt.plot(x, means)%0Aplt.show()

Implement basic mean of change per iteration plot