Python Exercise: Vampire Infection Simulation

Python Code Exercise

This code exercise makes use of Python3 and programming concepts to achieve a fun goal – in this case to build a program that simulates the outcome of a population of vampires infecting a large human population and the effect of vampire hunters in the mix (culling the vampire population.)

Is this influenced by the TV show, The Strain?

YES. It is true, I used the lore specifics of the strain to come up with this simulation. In the strain it took 1 to 2 days to get a human bitten to become a vampire. Similarly I use the same time span to calculate this as well. I also used the population size of a large city in New York (as in the Strain) and started with only a few vampire hunters (as in the story.)

Code Concepts Covered:

  • Functions (creating, calling, calling from within another function)
  • Tuples (returning a tuple, referencing an index of the tuple)
  • Dictionaries (creating, reading from, adding to)
  • Flow Control (if, elif)
  • Loops (while, for)
  • Random module
  • Matplotlib module (generating real time plots in Python3)
  • Error Handling


A couple weeks ago, I was getting into bed and I had an idea for a Zombie Infection simulation (written in code.)  After creating the simulation last week, someone at my job suggested I try doing a vampire simulation.    The Vampire concept was more complex then the Zombies.  Here’s the assumptions/rules the logic will abide by:


  • Vampirism is an infection that can turn a human into a vampire ( -1 from humans and + 1 to vampires)
  • Vampires only hunt at night
  • The general human population is large, and does not actively hunt vampires (unaware of their presence or scared)
  • The Vampire Hunter population is small, but grows at a rate each 24 hours (and in proportion to the Vampire population.)
  • Each vampire within it’s population has a roll/attempt to kill a human
  • If a vampire fails in it’s attempt to kill a human, the human has a percentage chance to kill 1 to X vampires.
  • If a vampire hunter fails to kill a vampire, then there is a chance the vampire hunter will be turned into a vampire.

Population Control

  • Vampire Hunters grow at a rate of Z, every 24 hours
  • Humans who are bit by Vampires, do not contribute to the Vampire population for 2 days (the gestation time of the vampire infection to fully turn a human bit, into a vampire)
  • If the Vampire population increases above 300,000 a potential cure will be introduced that can reduce the population of vampires by X% and add to the human population the same X%
  • If the Vampire Hunter population increases above 30,000 a potential trap of the vampires will turn Y% of Vampire Hunters into Vampires


For starters, you can clone my repo at:

To run the code you will need Python3, as well as the marplotlib module installed (you can install it with pip3, like so: pip3 install matplotlib )

To run the code, you can load it in your IDE and run the script, or from the command line just run:

Code Walkthrough

Let’s dissect the code.  Keep in mind I’m still learning Python, so my choices here may not be best practices.

import random
import csv
import matplotlib.pyplot as plt

I use 3 modules here, random (to get random numbers, which are used to calculate percentage chances at winning a round), csv (to output the results to a csv file) and matplotlib to plot real time data as it streams through the simulation, in a graph.

gestation_period = {}  
filename = random.randint(1100,5500)

The dictionary, gestation_period, was something I added later.  What it does is keep track of how many previously bit humans become vampires after X days.  For example, if 100 humans were bit on day 1, and the gestation period of the disease (time it takes to turn a human into a full vampire) is 2 days, then on day 3 we would have 100 new vampires added to the vampire population.

Without it, the vampire population would explode too fast.  It was instant vampire creation.  I wanted to try and simulate the Strain storyline as best as possible, so I put some time in there between the infection and becoming a full vampire.

To track this, I made use of a dictionary.  When the main function (vampire_simulation) is run it kicks off a function called Vampire attack.  That function calculates how many humans were “killed.”  Then it calls a function called gestation that stores the amount of humans “killed” and notes the current day (day 1, 2, 3, etc) and adds 2 to that and stores that key:value into the gestation_period dictionary.

def vampire_simulation(human_pop, vampire_pop, v_hunter_pop):
    day = 1
    hours_of_day = range(1,25) 

This is the start of the main function. Here we take 3 params:

  • the starting human population
  • the starting vampire population
  • the starting vampire hunter population

The simulation needs to keep track of the hours of the day, but also each day.  If a simulation goes beyond a certain period of time (days), I want to introduce more complexity (like population control measures and random events.)  So I start a counter of day = 1, and each time I go 24 hours, I increment that by 1.

A variable called hours_of_day is a collection of a range of numbers from 1 to 24.  This is accomplished with range(1,25) which gives us an integer from 1 to 24.

    while human_pop > 0 and vampire_pop > 0:

This is the core loop. While the human pop and vampire pop are not dead (0), we will continue to run what’s below.

The plt.ion() method allows us to redraw the plot via matplotlib.

        for h in hours_of_day:
            if h == 24:
                day = day + 1
                r = reinforcements(day,v_hunter_pop, vampire_pop)
                v_hunter_pop = r
                print("Today is day " + str(day) + " since the infection started.")

            if human_pop < 1:
                print("The general human population is gone")
            if vampire_pop < 1:
                print("humans won, eliminating the vampire threat")

Within the while loop (while the human and vampire populations exist), I set a for loop. I lop through each hour of the day. However, if the hour of the day is 24, that’s special. I want to make sure on the 24th hour I:

  • increment the day counter by 1
  • Vampire Hunter reinforcements are added at the start of each new day

I also want to make sure that if, we start the loop and the human population or vampire population has hit 0, we break out of the loops and report the result of who won.

                if h <= 4 or h > 20:
                    print("it's day: " + str(day) + " hour:" + str(h))
                    print(str(vampire_pop) + " vampire's roam the night")
                    va = vampire_attack(human_pop, vampire_pop)
                    # print(va)
                    human_pop = va[0]
                    vh = v_hunter_attack(h,vampire_pop,v_hunter_pop, human_pop)
                    vampire_pop = vh[0]

Within this else clause is our vampire/human struggle!

If h (hours) is less than or equal to 4 (4am) and greater then 8pm, then this is the nighttime, and Vampires only hunt at this time.

So here we kick off the vampire_attack method and we assign it to a variable va (for vampire attack.) The reason I assign a variable to this, is that the return is going to be a tuple. I need to reference the items returned to me in the tuple.

For example, the vampire_attack method will return several data objects, like the human population and vampire population post attack.  It will come back like so: (700,20).   To grab the first value, we would do the va[0] and va[1] for the second variable.

                    if v_hunter_pop > 30000:
                        vd = vampire_deceipt(v_hunter_pop,vampire_pop)
                        v_hunter_pop = vd[0]
                        vampire_pop = vd[1]

There’s an embedded condition here used to pull back the vampire hunter population if it gets too large. If the population goes above 30,000 there is another method called that has a chance of converting vampire hunters into vampires.

                        if gestation_period[day]:
                            vampire_pop = vampire_pop + gestation_period[day]
                    print("humans now: " + str(human_pop))
                    print("vampire hunters: " + str(v_hunter_pop))
                    print("vampires now: " + str(vampire_pop))

Next we need to know if any new vampires are added to the vampire pool. After thinking about, this might be better to put this logic in the if h == 24 (at the start of each new day.) But, here’s where I originally put the logic, right here in the main method.

What I’m doing is checking the dictionary to see if the current day has any humans converting into vampires (remember there is a gestation period of 2 days.) So why the try/except?

What if there is no data in the dictionary for that day? Day 1 won’t have any data. So it would halt the program on an error that there is no value in the dictionary for day 1. To compensate for this I make use of error handling.

Error handling allows me to try to do something, and if an error is returned I handle the error (instead of stopping the application.) In this case I use try: to try to check the dictionary. If I get a value returned, GREAT! If not, the code rolls to the except case… and here I’m just passing. That means “ignore it and move on.”

That completes the NIGHTTIME simulation. Now for the daylight hours!

                elif h >= 5 or h >= 19:                     
                     print("it's day: " + str(day) + " hour:" + str(h)) 
                     vh = v_hunter_attack(h,vampire_pop,v_hunter_pop,human_pop)
                     vampire_pop = vh[0]
                     if vampire_pop > 300000:
                        sw = secret_weapon(human_pop, vampire_pop)
                        vampire_pop = sw[1]
                        human_pop = sw[0]
                    print("humans now: " + str(human_pop))
                    print("vampire hunters: " + str(v_hunter_pop))
                    print("vampires now: " + str(vampire_pop))

Similar to the nighttime condition, we check the hours for daylight. If the hours are between 5AM and 7PM then we allow the hunters to kill vampires, and vampires do not initiate any attacks in this time period.

We call a method vampire_hunter and assign the results to a variable vh. A tuple is returned and we pull out the data we need from the tuple with it’s index.

There is a population control mechanism here as well. If the vampires grow too large, we introduce randomness to convert some of them into humans … this is handled in a method called secret_weapon.

                with open(str(filename)+'.csv', 'a') as csvfile:
                    vampirewriter = csv.writer(csvfile)
                    vampirewriter.writerow([day, h, human_pop, vampire_pop, v_hunter_pop])

Outside the if conditions (but within the FOR LOOP) a method called plotting is invoked. Passed into it is the hour of the day, the human population and the vampire population. Below this method call is a call to make a csv file and add a row to it each time this is kicked off (each hour.)

Finally, outside all of this, I call plt (matplotlib) and it’s show method. This renders the graphing.

def plotting(h,human_pop,vampire_pop):
    # plt.scatter(h,vampire_pop,v_hunter_pop,alpha=0.5)
    plt.title('Vampire Infection')
    plt.ylabel('Vampire Population')
    plt.xlabel('time (hours)')
    plt.ylabel('Human Population')

The above is the plotting method. Just a few points here… plt.draw() needs to be added to draw each updated graph (updating it each time the for loop is run), and we need to add a plt.pause(.1) to allow the graph to keep up with the function.

Also, Python3 has a little bit different set up then Python2.  In Python 2, the could be placed at the top (before the loop is started) but I found in Python 3, it has to come after the loop.

Vampire Attack Function!

def vampire_attack(human_pop, vampire_pop):
     if human_pop < 0 or vampire_pop < 0: return v_range = range(0,vampire_pop) for v in v_range: v_roll = random.randint(1,100) + 20 h_roll = random.randint(1,100) if v_roll > h_roll:
            human_casualties = random.randint(1,2)
            vampire_pop += human_casualties
            human_pop -= human_casualties
        elif h_roll > v_roll:
            v_kill_chance = random.randint(1,100)
            if v_kill_chance > 70 < 80: vampire_pop -= 1 + random.randint(0,10) elif v_kill_chance >= 80:
                vampire_pop = int(vampire_pop - (vampire_pop * .1))
    return human_pop, vampire_pop

In the Vampire attack function I check that the human/vampire pop is greater then 0. In the case of a large population, the loop could have started with a positive value, but part way through the loop processing, we could have hit 0. This is a double check. If we got a 0, we return nothing back.

Each Vampire gets a chance!

I get a range from 0 to the amount of vampires. Then I iterate through that range, and each vampire gets an attack opportunity. They roll a random number with a bonus to convert a human (+20, which is a 20% bonus.)

If the vampire attack is a success (vampire roll is higher then the human roll), then we subtract the amount of human casualties from the human pool and add it to the vampire pop… But remember, those vampire converts are not instantly part of the vampire population. No, they will be returned but put into the gestation dictionary, to be added to the vampire pool 2 days from the day of the attack.

Humans have a chance!

If a vampire fails to kill a human, it’s possible the human could kill one or more vampires. So some logic is added here to account for that as the v_kill chance.

Finally what’s returned is the human and vampire populations, which comes back as a tuple.

Vampire Hunter Attack Function

def v_hunter_attack(hour, vampire_pop, v_hunter_pop, human_pop):
    hunter_range = range(0,v_hunter_pop)
    for h in hunter_range:
        if hour > 5 < 20: hunter_roll = random.randint(1,100) + 20 v_roll = random.randint(1,100) - 20 if hunter_roll > v_roll:
                vampire_pop -= 1

This function takes the hour of the day, the vampire, vampire hunter and human population counts as params. The reason we pass hour into this, is that depending on the time of day, there is a percentage bonus given to one group vs. another.

If the hour is daylight (6am – 7pm), we give a + 20% increase in effectiveness to the Vampire Hunters and lower the vampire effectiveness by 20%.

                hunter_fail_roll = random.randint(1,100)
                if hunter_fail_roll > 90:
                    v_hunter_pop -= 1

Like all good role playing games, if you fail to do something, there is a potential for disaster. In the code above, if a vampire hunter fails to kill a vampire, it is potentially possible that they will be killed. During daylight hours, it’s more rare, and only a 9% chance (greater then 90 of 100) that a hunter is killed.

            hunter_roll = random.randint(1,100)
            v_roll = random.randint(1,100) + 20
            if hunter_roll > v_roll:
                vampire_pop -= 1
                hunter_fail_roll = random.randint(1,100)
                if hunter_fail_roll > 50:
                    v_hunter_pop -= 1
    return vampire_pop, v_hunter_pop

Ah but in the evening hours things are different… during the evening, the vampires get a + 20% chance of killing the vampire hunters… and vampire hunters have no bonus (nor penalty) in the evening.

If a hunter is successful, we remove 1 vampire from the vampire population. However, if the vampire hunter fails to kill it’s human target, the potential danger now, is 50% chance that the vampire hunter will be killed!

def secret_weapon(human_pop, vampire_pop):
        print("SECRET WEAPON")
        vc_roll = random.randint(1,100)
        if vc_roll > 70:
            vampire_cure = int(vampire_pop - (vampire_pop * 0.3))
            vampire_pop = vampire_cure
            human_pop = human_pop + vampire_cure
        return human_pop, vampire_pop

The secret weapon method above is called from the main vampire_simulation for loop, ONLY if the vampire population is above 300,000.

If above 300,000 then we make a roll from 1 to 100 and if we are above 70 (a 29% chance), then the vampire population is cut by 30% and added back to the human population.

The reason for this method is to keep the simulation going longer and get more interesting results.

def vampire_deceipt(v_hunter_pop, vampire_pop):
    print("Trickery has led to many hunters turned to the Night")
    hunter_conversion = int(v_hunter_pop - (v_hunter_pop * 0.2))
    v_hunter_pop = hunter_conversion
    vampire_pop = vampire_pop + hunter_conversion
    return v_hunter_pop, vampire_pop

Similar to the secret weapon method, the vampire deceipt method can convert vampire hunters into vampires. This method is only invoked form the main vampire_simulation function ONLY if the vampire hunter population is greater then 30,000.

def reinforcements(day, v_hunter_pop, vampire_pop):
    days = range(0,day) 
    for d in days:
        if d < 5: v_hunter_pop += 1 * int((vampire_pop * 0.01)/2) else: random_reinforcements = random.randint(1,100) if random_reinforcements > 95:
                v_hunter_pop += random.randint(1,3) * int((vampire_pop * 0.01))
                v_hunter_pop += 1 * int((vampire_pop * 0.01)/2)
    return v_hunter_pop

This reinforcement method above is called each day (on the 24th hour.) It runs a check on what day it currently is… day 1, 2, 3, 4, 5, etc. It takes the day as a range from 0 to the current day and iterates over it to increment how many vampire hunters are added to the pool.

This increases quite a bit… day1, not so much, but on day 2, it gets 2 iterations and day 3, it gets 3 and day 4, 4…

Sample Kick Off of Simulation


The above takes a population of 8 million people, 100 vampires and 2 vampire hunters and starts the simulation with those values.


As I created each method, I tested it out with sample data… below are some unit tests:

# Unit Tests:
print(v_hunter_attack(19, 10, 9))
print(reinforcements(3, 5))



Code Walk Through:

Demo of the Real Time Plotting with Matplotlib:

Leave a Reply

Your email address will not be published. Required fields are marked *