Basic disease modelling in python

disease.py

from random import random, sample
import networkx as NX
import numpy as NP
from matplotlib import pyplot as MPL

class Graph:
    def __init__(self, vertices, edges):
        self.vertices = vertices
        self.edges = edges


def random_population(n, alpha):
    """Give a random graph for a population of size n, with familiarity alpha.

    Idea: a uniformly random graph does not represent populations accurately,
    since cliques should be common. Ideally, tune alpha so that the average
    person knows 200 others."""

    # Distribute the population over a plane
    pop = [(random(), random()) for _ in range(n)]

    # Adjacency is dependent on distance
    adj = NP.zeros((n, n))
    for i in range(n):
        for j in range(i + 1, n):
            itoj = tuple(pop[j][k] - pop[i][k] for k in (0, 1))
            d2 = itoj[0] ** 2 + itoj[1] ** 2
            dn = d2 / (2 ** 0.5)
            if random() < (1 - dn) ** alpha:
                adj[i][j] = adj[j][i] = 1
    return adj

def neighbours(adj, x):
    for i in range(len(adj)):
        if adj[x][i] > 0:
            yield i

def run_infection(adj, start_size, recovery_time, mortality_per_day, infectiousness):
    n = len(adj)
    infections = {x: recovery_time for x in sample(range(n), start_size)}
    num_infected = start_size
    dead = set()
    while num_infected > 0:
        print("Infected: {}/{}".format(num_infected, n - len(dead)))
        new_infections = []
        for x in infections:
            infections[x] -= 1
            if random() < mortality_per_day:
                dead.add(x)
            for y in neighbours(adj, x):
                # Can't get infected once recovered
                if y not in dead and y not in infections and random() < infectiousness:
                    new_infections.append(y)
        for x in new_infections:
            infections[x] = recovery_time
        num_infected = sum(1 for x in infections if infections[x] > 0)
    print("{}/{} people were never infected".format(n - len(infections), n - len(dead)))

run_infection(random_population(1000, 50), 1, 3, 0.05, 0.05)




#G = NX.Graph(random_population(1000, 50))
#NX.draw(G)
#MPL.show()