Represents a diffusion of an infection on a graph. a_true is the adjacency matrix. This version is slow. Lots of optimizations to be made.

diffusion.py

import numpy
import math
import random

class Diffusion(object):
    def __init__(self, a_true, prob_model):
        N = numpy.size(a_true, 0)
        assert a_true.shape == (N,N), a_true.shape

        seed_node_index = math.floor(random.random() * N)

        infection_times = -1 * numpy.ones((1,N))
        assert infection_times.shape == (1,N), infection_times.shape
        infection_times[0][seed_node_index] = 0

        susceptible = (a_true[:,seed_node_index:seed_node_index+1] > 0).transpose()
        assert susceptible.shape == (1,N), susceptible.shape
        susceptible[0][seed_node_index] = False # The seed node is not susceptible

        unpropagated = numpy.zeros((1,N))
        assert unpropagated.shape == (1,N), unpropagated.shape
        unpropagated[0][seed_node_index] = True

        assert len(susceptible[0]) == N
        assert len(unpropagated[0]) == N
        assert len(infection_times[0]) == N

        while(numpy.sum(unpropagated) > 0):
            unpropagated_indexes = numpy.flatnonzero(unpropagated)

            assert len(unpropagated_indexes) < N

            seed_node_index = numpy.random.choice(unpropagated_indexes)

            current_time = infection_times[0][seed_node_index]
            unpropagated[0][seed_node_index] = 0

            assert numpy.random.rand(1,N).shape == (1,N), numpy.random.rand(1,N).shape
            assert susceptible.shape == (1,N), susceptible.shape

            a_vector = a_true[:,seed_node_index:seed_node_index+1].transpose()
            assert a_vector.shape == (1,N), a_vector.shape

            new_infections = (a_vector > numpy.random.rand(1,N))
            assert new_infections.shape == (1,N), new_infections.shape

            new_infections &= susceptible

            num_new = numpy.sum(new_infections)
            if num_new > 0:
                unpropagated += new_infections
                susceptible ^= new_infections
                times = current_time + prob_model(1, num_new)
                print times.shape
                infection_times[new_infections] = times[0]

        self.times = infection_times