ps8b.py

# Problem Set 8: Simulating the Spread of Disease and Virus Population Dynamics 

import numpy
import random
import pylab

''' 
Begin helper code
'''

class NoChildException(Exception):
    """
    NoChildException is raised by the reproduce() method in the SimpleVirus
    and ResistantVirus classes to indicate that a virus particle does not
    reproduce. You can use NoChildException as is, you do not need to
    modify/add any code.
    """

'''
End helper code
'''

#
# PROBLEM 2
#
class SimpleVirus(object):

    """
    Representation of a simple virus (does not model drug effects/resistance).
    """
    def __init__(self, maxBirthProb, clearProb):
        """
        Initialize a SimpleVirus instance, saves all parameters as attributes
        of the instance.        
        maxBirthProb: Maximum reproduction probability (a float between 0-1)        
        clearProb: Maximum clearance probability (a float between 0-1).
        """
        self.maxBirthProb = maxBirthProb
        self.clearProb = clearProb
        # TODO

    def getMaxBirthProb(self):
        """
        Returns the max birth probability.
        """
        return self.maxBirthProb
        # TODO

    def getClearProb(self):
        """
        Returns the clear probability.
        """
        return self.clearProb
        # TODO

    def doesClear(self):
        """ Stochastically determines whether this virus particle is cleared from the
        patient's body at a time step. 
        returns: True with probability self.getClearProb and otherwise returns
        False.
        """
        if random.random() < self.getClearProb():
            return True
        else:
            return False
        # TODO

    
    def reproduce(self, popDensity):
        """
        Stochastically determines whether this virus particle reproduces at a
        time step. Called by the update() method in the Patient and
        TreatedPatient classes. The virus particle reproduces with probability
        self.getMaxBirthProb * (1 - popDensity).
        
        If this virus particle reproduces, then reproduce() creates and returns
        the instance of the offspring SimpleVirus (which has the same
        maxBirthProb and clearProb values as its parent).         

        popDensity: the population density (a float), defined as the current
        virus population divided by the maximum population.         
        
        returns: a new instance of the SimpleVirus class representing the
        offspring of this virus particle. The child should have the same
        maxBirthProb and clearProb values as this virus. Raises a
        NoChildException if this virus particle does not reproduce.               
        """
        if random.random() <= (self.getMaxBirthProb() * (1 - popDensity)):
            return SimpleVirus(self.getMaxBirthProb(), self.getClearProb())
        else:
            raise NoChildException
        # TODO



class Patient(object):
    """
    Representation of a simplified patient. The patient does not take any drugs
    and his/her virus populations have no drug resistance.
    """    

    def __init__(self, viruses, maxPop):
        """
        Initialization function, saves the viruses and maxPop parameters as
        attributes.

        viruses: the list representing the virus population (a list of
        SimpleVirus instances)

        maxPop: the maximum virus population for this patient (an integer)
        """
        self.viruses = viruses
        self.maxPop = maxPop
        # TODO

    def getViruses(self):
        """
        Returns the viruses in this Patient.
        """
        # TODO
        return self.viruses


    def getMaxPop(self):
        """
        Returns the max population.
        """
        # TODO
        return self.maxPop


    def getTotalPop(self):
        """
        Gets the size of the current total virus population. 
        returns: The total virus population (an integer)
        """
        return len(self.viruses)
        # TODO        


    def update(self):
        """
        Update the state of the virus population in this patient for a single
        time step. update() should execute the following steps in this order:
        
        - Determine whether each virus particle survives and updates the list
        of virus particles accordingly.   
        
        - The current population density is calculated. This population density
          value is used until the next call to update() 
        
        - Based on this value of population density, determine whether each 
          virus particle should reproduce and add offspring virus particles to 
          the list of viruses in this patient.                    

        returns: The total virus population at the end of the update (an
        integer)
        """

        # TODO
        #1
        survivors = []
        for virus in self.getViruses():
            if not virus.doesClear():
                survivors.append(virus)
        self.viruses = survivors
        #2
        popDensity = float(len(self.getViruses()))/self.getMaxPop()
        #3
        for virus in survivors:
            try:
                self.viruses.append(virus.reproduce(popDensity))
            except NoChildException:
                continue
        return self.getTotalPop()

'''
virus = SimpleVirus(1.0, 1.0)
patient = Patient([virus], 100)
virus.getClearProb()
virus.getMaxBirthProb()
#virus.reproduce(.5)
patient.getViruses()
patient.getMaxPop()
patient.getTotalPop()
for i in xrange(100):
    patient.update()
print patient.getViruses()
'''
#
# PROBLEM 3
#
def simulationWithoutDrug(numViruses=100, maxPop=1000, maxBirthProb=0.1, clearProb=0.05,
                          numTrials=30):
    """
    Run the simulation and plot the graph for problem 3 (no drugs are used,
    viruses do not have any drug resistance).    
    For each of numTrials trial, instantiates a patient, runs a simulation
    for 300 timesteps, and plots the average virus population size as a
    function of time.

    numViruses: number of SimpleVirus to create for patient (an integer)
    maxPop: maximum virus population for patient (an integer)
    maxBirthProb: Maximum reproduction probability (a float between 0-1)        
    clearProb: Maximum clearance probability (a float between 0-1)
    numTrials: number of simulation runs to execute (an integer)
    """
    for trial in range(numTrials):
        virusLevelResults = runTrial(300, numViruses, maxPop, maxBirthProb, clearProb)
        print trial

        if trial == 0:  # First trial
            # Save results of the first trial
            accumulatedVirusLevels = virusLevelResults
        else:
            # Accumulate second and subsequent results to same list
            for i in range(len(virusLevelResults)):
                accumulatedVirusLevels[i] += virusLevelResults[i]

    # At end of the trials, average the accumulated results over
    # the number of trials carried out (see HINT in problem spec.)
    for i in range(len(accumulatedVirusLevels)):
        accumulatedVirusLevels[i] /= float(numTrials)

    # PLOT THE RESULTS OF THE SIMULATION
    # Let X-axis scale default to length of Y-values list (301)
    pylab.plot(accumulatedVirusLevels, label = 'SampleVirus')
    pylab.title('Population of SimpleVirus with '+str(numTrials)+' trials')
    pylab.ylabel('Average number of viruses in patient')
    pylab.xlabel('Time steps (arbitrary units)')
    pylab.legend(loc = 'best')
    pylab.show()
    # TODO
def runTrial(numTimeSteps, numViruses, maxPop, maxBirthProb, clearProb):

    # Create numViruses virus instances with:
    #   maxBirthProb 
    #   clearProb
    # attribute values
    viruses = []
    for i in xrange(numViruses):
        viruses.append(SimpleVirus(maxBirthProb, clearProb))

    # Instantiate a patient with the list of numVirus virus instances
    # and a Maximum Sustainable Virus Population of maxPop
    # one line of code
    patient = Patient(viruses, maxPop)
    # The following pseudocode simulates, records and returns changes
    # to the virus population over elapsedTimeSteps steps.
    # Note that the patient has numViruses viruses at start.

    virusLevelsThisTrial = []
    # NOTE: You SHOULD NOT record the starting level of viruses in the
    #       patient as the virus population at t=0.
    # 
    # The Problem Statement/Specification for Problem 3 is quite clear. 
    # It says: "The population at time=0 is the population after the 
    # first call to update". So the next line of pseudocode should
    # not be imnplemented! If you are having problems with the grader
    # AND you had implemented the next line, take it out. Your code
    # should pass the grader.

    #REMOVE THIS LINE# virusLevelsThisTrial.append(...) # numViruses

    # Record virus population changes over numTimeSteps number
    # of updates. Do the update first then record the virus
    # level/population that results. 
    for i in xrange(numTimeSteps):
        patient.update()#... call update() method
        virusLevelsThisTrial.append(patient.getTotalPop())#... append to allVirusLevelsThisTrial # use the getter method
        # Note: in the above we can roll those two steps into one by
        #       appending directly the value returned by update().
        #       But see the "Discussion Section" later on for the
        #       reason I adopt the former approach.

    return virusLevelsThisTrial  # a list with numTimeStep values (300)


#
# PROBLEM 4
#
class ResistantVirus(SimpleVirus):
    """
    Representation of a virus which can have drug resistance.
    """   

    def __init__(self, maxBirthProb, clearProb, resistances, mutProb):
        """
        Initialize a ResistantVirus instance, saves all parameters as attributes
        of the instance.

        maxBirthProb: Maximum reproduction probability (a float between 0-1)       

        clearProb: Maximum clearance probability (a float between 0-1).

        resistances: A dictionary of drug names (strings) mapping to the state
        of this virus particle's resistance (either True or False) to each drug.
        e.g. {'guttagonol':False, 'srinol':False}, means that this virus
        particle is resistant to neither guttagonol nor srinol.

        mutProb: Mutation probability for this virus particle (a float). This is
        the probability of the offspring acquiring or losing resistance to a drug.
        """
        self.resistances = resistances
        self.mutProb = mutProb
        SimpleVirus.__init__(self, maxBirthProb, clearProb)
        # TODO


    def getResistances(self):
        """
        Returns the resistances for this virus.
        """
        # TODO
        return self.resistances

    def getMutProb(self):
        """
        Returns the mutation probability for this virus.
        """
        # TODO
        return self.mutProb

    def isResistantTo(self, drug):
        """
        Get the state of this virus particle's resistance to a drug. This method
        is called by getResistPop() in TreatedPatient to determine how many virus
        particles have resistance to a drug.       

        drug: The drug (a string)

        returns: True if this virus instance is resistant to the drug, False
        otherwise.
        """
        try:
            return self.getResistances()[drug]
        except KeyError:
            return False
        # TODO


    def reproduce(self, popDensity, activeDrugs):
        """
        Stochastically determines whether this virus particle reproduces at a
        time step. Called by the update() method in the TreatedPatient class.

        A virus particle will only reproduce if it is resistant to ALL the drugs
        in the activeDrugs list. For example, if there are 2 drugs in the
        activeDrugs list, and the virus particle is resistant to 1 or no drugs,
        then it will NOT reproduce.

        Hence, if the virus is resistant to all drugs
        in activeDrugs, then the virus reproduces with probability:      

        self.getMaxBirthProb * (1 - popDensity).                       

        If this virus particle reproduces, then reproduce() creates and returns
        the instance of the offspring ResistantVirus (which has the same
        maxBirthProb and clearProb values as its parent). The offspring virus
        will have the same maxBirthProb, clearProb, and mutProb as the parent.

        For each drug resistance trait of the virus (i.e. each key of
        self.resistances), the offspring has probability 1-mutProb of
        inheriting that resistance trait from the parent, and probability
        mutProb of switching that resistance trait in the offspring.       

        For example, if a virus particle is resistant to guttagonol but not
        srinol, and self.mutProb is 0.1, then there is a 10% chance that
        that the offspring will lose resistance to guttagonol and a 90%
        chance that the offspring will be resistant to guttagonol.
        There is also a 10% chance that the offspring will gain resistance to
        srinol and a 90% chance that the offspring will not be resistant to
        srinol.

        popDensity: the population density (a float), defined as the current
        virus population divided by the maximum population       

        activeDrugs: a list of the drug names acting on this virus particle
        (a list of strings).

        returns: a new instance of the ResistantVirus class representing the
        offspring of this virus particle. The child should have the same
        maxBirthProb and clearProb values as this virus. Raises a
        NoChildException if this virus particle does not reproduce.
        """
        isResistantToActiveDrugs = True
        for drug in activeDrugs:
            if not self.isResistantTo(drug):
                isResistantToActiveDrugs = False
                break
            else:
                continue
        if isResistantToActiveDrugs and (random.random() <= (self.getMaxBirthProb() * (1 - popDensity))) :
            offspringResistances = {}
            for drug in self.getResistances().keys():
                resistanceDict = self.getResistances()
                if resistanceDict[drug]:
                    if random.random()<(1-self.getMutProb()):
                        offspringResistances[drug] = True
                    else:
                        offspringResistances[drug] = False
                else:
                    if random.random()<self.getMutProb():
                        offspringResistances[drug] = True
                    else:
                        offspringResistances[drug] = False
            return ResistantVirus(self.getMaxBirthProb(), self.getClearProb(), offspringResistances, self.getMutProb())
        else:
            raise NoChildException
        # TODO

'''    
a = ResistantVirus(.99,.01,{'srinol':True}, .2)
j=0
for i in xrange(100):
    try:
        a.reproduce(.1, [])
        print j
        j+=1
    except NoChildException:
        continue
'''
class TreatedPatient(Patient):
    """
    Representation of a patient. The patient is able to take drugs and his/her
    virus population can acquire resistance to the drugs he/she takes.
    """

    def __init__(self, viruses, maxPop):
        """
        Initialization function, saves the viruses and maxPop parameters as
        attributes. Also initializes the list of drugs being administered
        (which should initially include no drugs).              

        viruses: The list representing the virus population (a list of
        virus instances)

        maxPop: The  maximum virus population for this patient (an integer)
        """
        self.prescriptions = []
        Patient.__init__(self, viruses, maxPop)
        # TODO


    def addPrescription(self, newDrug):
        """
        Administer a drug to this patient. After a prescription is added, the
        drug acts on the virus population for all subsequent time steps. If the
        newDrug is already prescribed to this patient, the method has no effect.

        newDrug: The name of the drug to administer to the patient (a string).

        postcondition: The list of drugs being administered to a patient is updated
        """
        if newDrug not in self.getPrescriptions():
            self.prescriptions.append(newDrug)
        # TODO


    def getPrescriptions(self):
        """
        Returns the drugs that are being administered to this patient.

        returns: The list of drug names (strings) being administered to this
        patient.
        """
        return self.prescriptions
        # TODO


    def getResistPop(self, drugResist):
        """
        Get the population of virus particles resistant to the drugs listed in
        drugResist.       

        drugResist: Which drug resistances to include in the population (a list
        of strings - e.g. ['guttagonol'] or ['guttagonol', 'srinol'])

        returns: The population of viruses (an integer) with resistances to all
        drugs in the drugResist list.
        """
        count = 0
        for virus in self.getViruses():
            for drug in drugResist:
                if virus.isResistantTo(drug):
                    count += 1
        return count
        # TODO


    def update(self):
        """
        Update the state of the virus population in this patient for a single
        time step. update() should execute these actions in order:

        - Determine whether each virus particle survives and update the list of
          virus particles accordingly

        - The current population density is calculated. This population density
          value is used until the next call to update().

        - Based on this value of population density, determine whether each 
          virus particle should reproduce and add offspring virus particles to 
          the list of viruses in this patient.
          The list of drugs being administered should be accounted for in the
          determination of whether each virus particle reproduces.

        returns: The total virus population at the end of the update (an
        integer)
        """
        #1
        survivors = []
        for virus in self.getViruses():
            if not virus.doesClear():
                survivors.append(virus)
        self.viruses = survivors
        #2
        popDensity = float(len(self.getViruses()))/self.getMaxPop()
        #3
        for virus in survivors:
            try:
                self.viruses.append(virus.reproduce(popDensity, self.getPrescriptions()))
            except NoChildException:
                continue
        return self.getTotalPop()
            
        # TODO



#
# PROBLEM 5
#
def simulationWithDrug(numViruses = 10, maxPop = 1000, maxBirthProb = .9, clearProb = .2, resistances = {'guttagonol':True, 'srinol':True},
                       mutProb = .9, numTrials = 10):
    """
    Runs simulations and plots graphs for problem 5.

    For each of numTrials trials, instantiates a patient, runs a simulation for
    150 timesteps, adds guttagonol, and runs the simulation for an additional
    150 timesteps.  At the end plots the average virus population size
    (for both the total virus population and the guttagonol-resistant virus
    population) as a function of time.

    numViruses: number of ResistantVirus to create for patient (an integer)
    maxPop: maximum virus population for patient (an integer)
    maxBirthProb: Maximum reproduction probability (a float between 0-1)        
    clearProb: maximum clearance probability (a float between 0-1)
    resistances: a dictionary of drugs that each ResistantVirus is resistant to
                 (e.g., {'guttagonol': False})
    mutProb: mutation probability for each ResistantVirus particle
             (a float between 0-1). 
    numTrials: number of simulation runs to execute (an integer)
    
    """

    # TODO
    for trial in range(numTrials):
        (totalVirusResult, resistantVirusResult) = runTrial2(150, 150,
              numViruses, maxPop,
              maxBirthProb, clearProb,
              resistances, mutProb)

        if trial == 0:  # First trial
            # Save results of the first trial...
            totalVirusPop = totalVirusResult
            resistantVirusPop = resistantVirusResult
        else:
            # ...then accumulate second and subsequent sets
            for i in range(len(totalVirusResult)):
                totalVirusPop[i] += totalVirusResult[i]
                resistantVirusPop[i] += resistantVirusResult[i]

    # End of the numTrial trials. Average the results accumulated during
    # the trials over the number of trials. (See HINT in problem spec.)
    for i in range(len(totalVirusPop)):
        totalVirusPop[i] /= float(numTrials)
        resistantVirusPop[i] /= float(numTrials)

    # PLOT THE RESULTS OF THE SIMULATION AS TWO PLOTS
    # Let X-axis scale default to length of Y-values lists (300)
    # Here is how we deal with plotting two sets of y-values
    pylab.plot(totalVirusPop, label = 'TotalVirusPop') 
    pylab.plot(resistantVirusPop, label = 'ResistantVirusPop')
    pylab.title('Population of SimpleVirus with '+str(numTrials)+' trials')
    pylab.ylabel('Average number of viruses in patient')
    pylab.xlabel('Time steps (arbitrary units)')
    pylab.legend(loc = 'best')
    pylab.show()
def runTrial2(numTimeSteps1, numTimeSteps2,
              numViruses, maxPop,
              maxBirthProb, clearProb,
              resistances, mutProb):
    # Create numViruses virus instances with specified:
    #   maxBirthProb 
    #   clearProb
    #   resistances
    #   mutProb
    # attribute values
    viruses = []
    for i in xrange(numViruses):
        viruses.append(ResistantVirus(maxBirthProb, clearProb, resistances, mutProb))
    # Instantiate a patient with the list of numVirus virus instances
    # and a Maximum Sustainable Virus Population of maxPop
    patient = TreatedPatient(viruses, maxPop)  # one line of code

    # Simulate changes to the virus population over elapsedTimeSteps
    # steps.
    # The patient has numViruses virus level at start.

    allVirusLevelsThisTrial = []
    resistantVirusLevelsThisTrial = []
    # Record starting level of viruses in patient...
    # NO! See the BEWARE! comments above..
    #REMOVE LINE# allVirusLevelsThisTrial.append(...)  # total virus...
    #REMOVE LINE# resistantVirusLevelsThisTrial.append(...) # resistant virus...

    # Record virus population changes over numTimeSteps1 
    # number of updates. Do the update first then record the
    # two virus populations of interest that results.
    for i in xrange(numTimeSteps1):
        patient.update()#... call update() method
        allVirusLevelsThisTrial.append(patient.getTotalPop())#... append to allVirusLevelsThisTrial # use the getter method
        # Note below, looking at virus population resistant to 'guttagonol'
        resistantVirusLevelsThisTrial.append(patient.getResistPop(['guttagonol'])) # use the getter method

    # Add the drug 'guttagonol' to the patient and run a further
    # numTimeSteps2 number of updates
    patient.addPrescription('guttagonol')
    for i in xrange(numTimeSteps2):
        patient.update()#... call update() method
        allVirusLevelsThisTrial.append(patient.getTotalPop())#... append to allVirusLevelsThisTrial # use the getter method
        # Note below, looking at virus population resistant to 'guttagonol'
        resistantVirusLevelsThisTrial.append(patient.getResistPop(['guttagonol'])) # use the getter method

    return (allVirusLevelsThisTrial, resistantVirusLevelsThisTrial)