def epi_curve(event_db): """Display curve with the number of infected individuals at each time point""" num_infected = [0]*np.max(event_db.time) for t in range(0, np.max(event_db.time)): num_infected[t] = find_nonrecovered(event_db, t+1).shape[0] plt.plot(range(1, np.max(event_db.time)+1), num_infected, linewidth=2)
def sir_likelihood(pop, event_db, recoverytimes, alpha, beta, gamma): """Determine the likelihood of the data from an SIR simulation given alpha, beta, and gamma. """ daily_likelihood=[0]*(np.max(event_db.time)-1) t=1 nonrecovered=find_nonrecovered(event_db, t) susceptible=find_susceptible(pop, event_db, t) for t in range(1, np.max(event_db.time)): new_nonrecovered=find_nonrecovered(event_db, t+1) new_susceptible=find_susceptible(pop, event_db, t+1) infection_probs=infect_prob(pop, alpha, beta, nonrecovered, susceptible) def new_nonrecovered_func(x): return any(susceptible.ind_ID[x] == new_nonrecovered.ind_ID) recovery_index = map(new_nonrecovered_func, susceptible.index) daily_likelihood[t-1]=np.prod(np.subtract(1, infection_probs[np.where(np.invert(recovery_index))]))*np.prod(infection_probs[np.where(recovery_index)]) nonrecovered=new_nonrecovered susceptible=new_susceptible return np.prod(daily_likelihood) *geometric_likelihood(recoverytimes, (1./gamma))
def plot_sir(pop, event_db, time): """This will create a simple scatterplot of susceptible, infectious, and recovered individuals at a given time """ i=find_nonrecovered(event_db, time) s=find_susceptible(pop, event_db, time) r=find_recovered(event_db, time) status=pd.DataFrame({"status":np.append(np.append(np.repeat("i", i.shape[0]), np.repeat("s", s.shape[0])),np.repeat("r", r.shape[0])), "colour":np.append(np.append(np.repeat("#a6cee3", i.shape[0]), np.repeat("#1f78b4", s.shape[0])),np.repeat("#b2df8a", r.shape[0])), "ind_ID":np.append(np.append(i.ind_ID, s.ind_ID), r.ind_ID), "x":pop.iloc[np.append(np.append(i.ind_ID, s.ind_ID), r.ind_ID)].x, "y":pop.iloc[np.append(np.append(i.ind_ID, s.ind_ID), r.ind_ID)].y}) plt.scatter(status.x, status.y, c=status.colour, s=60, edgecolors='none')