fcp2021 / fcp.py
fcp.py
Raw
#!/usr/bin/env python 

import argparse 
import matplotlib.pyplot as plt
from simulation import Simulation
from Project_Classes import Simulation, Animation, evolution_simulation

def main(*args):
    
    # ArgumentParser is a class and we are creating an object from it
    parser = argparse.ArgumentParser(description='Setting the parameters of the epidemic sample') # This line is required for any argparse lines of code

# Note: The order in which our parameters are called is not important 

    parser.add_argument('--size', metavar='N', type=int, default=50,
                        help='Defines the size of the matrix')
    parser.add_argument('--days', metavar='T', type=int, default=40,
                        help='Simulation running for T days')
    parser.add_argument('--asymp_cases', metavar='A', type=int, default=2,
                        help='Number of asymptomatic infected people introduced in the matrix')
    parser.add_argument('--symp_cases', metavar='S', type=int, default=2,
			help='Number of symptomatic infected people introduced in the matrix')
    parser.add_argument('--range_asymp', metavar='RA', type=int, default=2,
			help='Maximum range asymptomatic cases can infect people at')
    parser.add_argument('--range_symp', metavar='RS', type=int, default=1,
			help='Maximum range symptomatic cases can infect people at')
    parser.add_argument('--lockdown', metavar='LD', type=bool, default=False,
			help='Instauring lockdown for population in matrix')
    parser.add_argument('--inter_asymp', metavar='Interactions A', type=float, default=0.7,
			help='Average interactions asymptomatic cases with others')
    parser.add_argument('--inter_symp', metavar='Interactions S', type=float, default=0.6,
			help='Average interactions symptomatic infected cases with others')
    parser.add_argument('--file', metavar='F', type=str, default=None,
			help='Filename to save animation to')
    parser.add_argument('--plot', action='store_true',
			help='Generate a figure displaying the virus at different days in the simulation')
    parser.add_argument('--width', metavar='W', type=int, default=5,
			help='Number of columns in simulation plot')
    parser.add_argument('--height', metavar='H', type=int, default=3,
			help='Number of rows in simulation plot')

    args = parser.parse_args()


    # Setting up the simulation with the arguments from the parser (if no arguments, the default values are used)
    simulation = Simulation(args.size, args.asymp_cases, args.symp_cases, args.inter_asymp, args.range_asymp, args.inter_symp, args.range_symp, args.days, args.lockdown)
    simulation.init_array()

    if args.plot:
        evolution = evolution_simulation(simulation, args.width, args.height, args.days)
        
        if args.file is None:
            plt.show() # With : $ python fcp.py --plot
        else:
            evolution.savefig(args.file) # With : $ python fcp.py --plot --file=plot.pdf
    else:
        animation = Animation(args.days, args.size**2, simulation)
        
        if args.file == None:
            animation.show()
        else:
            animation.save(args.file)
	    # NOTE: ffmpeg required


if __name__ == "__main__":
    #
    # CLI entry point. The main() function can also be imported and called
    # with string arguments.
    #
    import sys
    main(*sys.argv[1:])