#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jun 29 03:08:28 2020

@author: zahrashah
"""


import networkx as nx
import EoN
import matplotlib.pyplot as plt

G = nx.grid_2d_graph(300,300) 

#can change where the infection starts by changing the this list

initial_infections = [(x,y) for (x,y) in G if 35<x<45 and 30<y<40]
pos = {node:node for node in G}
sim_kwargs = {'pos': pos}
#you can change tau and gamma here
sim = EoN.fast_SIR(G, 2.0, 1.0, initial_infecteds = initial_infections,
               tmax = 100, return_full_data=True, sim_kwargs = sim_kwargs)

animation=sim.animate(ts_plots=[["S"], ["I"], ["R"]], node_size = 4)
animation.save('SIR2D.mp4', fps=5)