In [1]:
import networkx as nx

number_of_nodes = 10
G = nx.complete_graph(number_of_nodes)

In [2]:
import random
from nxsim import BaseNetworkAgent

class ZombieOutbreak(BaseNetworkAgent):
    def __init__(self, environment=None, agent_id=0, state=()):
        super().__init__(environment=environment, agent_id=agent_id, state=state)
        self.bite_prob = 0.05

    def run(self):
        while True:
            if self.state['id'] == 1:
                self.zombify()
                yield self.env.timeout(1)
            else:
                yield self.env.event()

    def zombify(self):
        normal_neighbors = self.get_neighboring_agents(state_id=0)
        for neighbor in normal_neighbors:
            if random.random() < self.bite_prob:
                neighbor.state['id'] = 1 # zombie
                print(self.env.now, self.id, neighbor.id, sep='\t')
                break

In [3]:
from nxsim import NetworkSimulation

# Initialize agent states. Let's assume everyone is normal.
init_states = [{'id': 0, } for _ in range(number_of_nodes)]  # add keys as as necessary, but "id" must always refer to that state category

# Seed a zombie
init_states[5] = {'id': 1}
sim = NetworkSimulation(topology=G, states=init_states, agent_type=ZombieOutbreak, 
                        max_time=30, num_trials=1, logging_interval=1.0)

In [4]:
sim.run_simulation()


Starting simulations...
---Trial 0---
Setting up agents...
2	5	8
3	5	1
4	5	3
6	5	9
8	5	7
13	5	0
15	5	2
17	5	6
22	5	4
Written 30 items to pickled binary file: sim_01/log.0.state.pickled
Simulation completed.

In [5]:
%matplotlib inline
nx.draw(G)


/usr/local/lib/python3.4/site-packages/matplotlib/font_manager.py:1279: UserWarning: findfont: Font family ['Helvetica'] not found. Falling back to Bitstream Vera Sans
  (prop.get_family(), self.defaultFamily[fontext]))

In [6]:
from nxsim import BaseLoggingAgent
trial = BaseLoggingAgent.open_trial_state_history(dir_path='sim_01', trial_id=0)

In [7]:
import numpy as np
from matplotlib import pyplot as plt
%matplotlib inline
zombie_census = [sum([1 for node_id, state in g.items() if state['id'] == 1]) for t,g in trial.items()]
plt.plot(zombie_census)


Out[7]:
[<matplotlib.lines.Line2D at 0x10f242cf8>]

In [7]: