Python SIR model

SIR process class

class SIRProcess(ExplicitStateProcess):
    """
    S I R process,

    Attributes
    ----------
    beta - Infection rate.
    gamma - Recovery rate. 
    """
    def __init__(self, beta, gamma):
        
        super(SIRProcess, self).__init__(['S', 'I', 'R'],
                                         [],
                                         runNodeUpdate = True,
                                         runEdgeUpdate = False,
                                         runNetworkUpdate = False,
                                         constantTopology = True)
        self.beta = float(beta)
        self.gamma = float(gamma)

    
    def nodeUpdateRule(self, node, srcNetwork, dt):
        # Read original node state.
        srcState = node[1][self.STATE_ATTR_NAME]
        # By default we have not changed states, so set
        # the destination state to be the same as the source state.
        dstState = srcState

        # Start out with a dictionary of zero neighbors in each state.
        nNSt = dict(zip(self.nodeStateIds,[0]*len(self.nodeStateIds)))
        # Calculate the actual numbers and update dictionary.
        nNSt.update(attributeCount(neighbors_data_iter(srcNetwork, node[0]),
                                   self.STATE_ATTR_NAME))

        # Pick a random number.
        eventp = numpy.random.random_sample()
        # Go through each state name, and chose an action.
        if srcState == 'S':
            if eventp < self.beta*nNSt['I']*dt:
                dstState = 'I'
        elif srcState == 'I':
            if eventp < self.gamma*dt:
                dstState = 'R'

        node[1][self.STATE_ATTR_NAME] = dstState

        return node

SIR simulation configuration

# SIR with an own module.

# This is the simulation section.
[Simulation]
# Run the simulation this many iterations.
iterations = 500

# The time step taken each iteration.
dt = .1

# This is the python module containing the process we wish to use.
process_class_module = extended_SIR

# This is the name of the process object.
process_class = SIRProcess

# This is the name of the network generation function.
network_func = BA_networkx

# We need to add another module path
module_paths = ../modules/


# Network settings.
[NetworkParameters]
# Number of nodes.
n = 1000
# Number of edges to add in each iteration.
m = 2

# Defining the process parameter values.
# The contents of this section is dependent on
# the parameters of the process class as specified by
# the option process_class in the Simulation section.
[ProcessParameters]
# Infection rate.
beta = .9e-2
# Recovery rate. 
gamma = 0.04 

# The fraction of nodes, alternatively the number of nodes, that will be  assigned to each state initially.
# The state names must match those specified by the network process class.
[NodeStateDistribution]
# 95% S
S = 0.95
# 5% I
I = 0.05
# Zero recovered to start with.
R = 0

# Result output settings.
[Output]

# Output directory:
output_dir = ../output/

# This is the base name of all files generated by the run.
base_name = test_SIR_2

# If unique is defined as true, yes, 1, or on, unique file names will be created (time stamp added)
unique = yes

# If this is true, yes, 1, or on, a copy of the full program config, plus an Info
# section will be saved.
save_config = yes

# If this is true/yes/on, the network node states will be counted and saved as a csv file.
# Default value True.
# Note only valid if the current process support updates. If not nothing will be saved.
save_state_count = yes
 
# Count nodes every ... iterations. Value should be integer >= 1.
# Default value 1.
save_state_count_interval = 1

# If this is true, yes, 1, or on, a copy of the network will be saved.
# Save interval may be set using the save_network_interval key.
save_network = yes

# This control how often the network will be saved.
# A value <= 0 means only the initial network will be saved. A positive value
# n> 0, results in the initial network being saved plus every n:th iteration
# thereafter, as well as the last network.
# Default value 0.
save_network_interval = 0