<jupyter_start><jupyter_text>Fitting an SIR epidemic modelHere we fit a stochastic epidemic model. The inputs are time and $R_0$ and the outputs are the number of infected.<jupyter_code>%pylab inline
import sys
sys.path.append('../')
import importlib
import pydra
import keras
import seaborn as sns
import plot_utils as pu
import tensorflowjs as tfjs
importlib.reload(pydra)
importlib.reload(pu)
import tensorflow as tf
from tqdm import tqdm
tf.__version__<jupyter_output><empty_output><jupyter_text>Define SIR model function<jupyter_code>def SIR(beta,gamma,N=100,T=10):
    Is = np.zeros(T)
    Ss = np.zeros(T)
    Is[0] = 1
    Ss[0] = N-Is[0]
    ts = range(T)
    for t in ts[1:]:
        infections = np.random.poisson(beta*Is[t-1]*Ss[t-1]/N)
        recoveries = np.random.poisson(gamma*Is[t-1])
        Is[t] = max(Is[t-1] + infections - recoveries,0)
        Ss[t] = max(Ss[t-1] - infections,0)
    Fs = 100.*(N - Ss[-1])/float(N)
    return ts,Ss,Is,Fs<jupyter_output><empty_output><jupyter_text>Test SIR model<jupyter_code>beta = 0.2
ts,_,Is,_ = SIR(beta,0.1,N=1000,T=100)
inputs = np.vstack((beta*np.ones(100),ts))
outputs = Is
plt.plot(ts,Is)
inputs = np.zeros(2,None)
outputs = np.zeros(None)
betas = np.random.uniform(high=1.0,low=0.0,size=3000)
inputs = None
outputs = None
for beta in betas:
    ts,_,Is,_ = SIR(beta,0.1,N=1000,T=100)
    Is = Is.astype(int)
    #Is[Is==0] = 1e-10 #if zero replace with small value.
    binput = np.vstack((beta*np.ones(100),ts))
    if inputs is None:
        inputs = binput.copy()
    else:
        inputs = np.hstack((binput,inputs))
    if outputs is None:
        outputs = Is
    else:
        outputs = np.hstack((Is,outputs))
print(inputs.shape)
print(outputs.shape)
outputs = np.atleast_2d(outputs)
print(outputs.shape)<jupyter_output>(2, 300000)
(300000,)
(1, 300000)<jupyter_text>Instantiate Pydra model<jupyter_code>model = pydra.Pydra(input_size=2,output_size=1,layers=5,learning_rate=0.01,output_distributions='Poisson',
                    cluster_size=5,activation='relu',dense_layer_size=32,params={'binomial_n':float(100)})
history = model.fit(inputs.T, outputs.T, batch_size=1000, epochs=10, verbose=1)
plt.figure(figsize=(12,8));
plt.plot(history.history['loss']); plt.xlabel('Epoch'); plt.ylabel('loss');
def plot_for_beta(beta):
    n = 1000
    t_test = np.float32(np.linspace(0.,100,num=n))

    x_test = np.vstack((beta*np.ones(n),t_test)).T

    output = model.generate_mdn_sample_from_ouput(x_test)
    plt.plot(t_test,output,'bo',alpha=0.5);
betas = [0.05,0.1,0.15,0.5]
plt.figure(figsize=(10,8));
for i,beta in enumerate(betas):
    plt.subplot(2,2,i+1);
    plot_for_beta(beta);
    plt.title(str(beta));
betas = np.linspace(0.,0.5,num=9)
plt.figure(figsize=(12,8));
n=100
for i,beta in enumerate(betas):
    plt.subplot(3,3,i+1);
    t_test = np.float32(np.linspace(0.0,100,num=n))
    x_test = np.vstack((beta*np.ones(n),t_test)).T
    output = model.predict(x_test)
    mu,sigma = pydra.get_stats(output,distribution='Poisson')
    plt.plot(mu)
    plt.plot(mu+np.sqrt(sigma),'k--');
    plt.plot(mu-np.sqrt(sigma),'k--');
    plt.title(str(beta));
    if i<6: plt.xticks([],[]);<jupyter_output><empty_output><jupyter_text>Perform KS-test<jupyter_code>from scipy.stats import ks_2samp



def gen_comparison_sample(beta,t=100,n=100):
    
    t_test = np.float32(np.linspace(0.0,100,num=n))
    x_test = np.vstack((beta*np.ones(n),t_test)).T
    output = model.predict(x_test)
    y_test = pydra.generate_mdn_sample_from_ouput(output, output.size,distribution='Poisson')
    ts,_,y_original,_ = SIR(beta,0.1,N=1000,T=100)
    return y_original, y_test


def compare_ks(y_test,y_original):
    res = ks_2samp(y_test,y_original)
    return res.statistic

def ks_test_over_linspace(beta,n=100,sample_size=50):
    ksvals = np.zeros(n)
    Ytest = np.zeros((n,sample_size))
    Yoriginal = np.zeros((n,sample_size))
    for i in range(sample_size):
        y_original,y_test = gen_comparison_sample(beta,t=100,n=100)
        Ytest[:,i] = y_test
        Yoriginal[:,i] = y_original
    for i in range(ksvals.size):
        ksvals[i] = compare_ks(Ytest[i,:].flatten(),Yoriginal[i,:].flatten())
    xinput = np.linspace(0.,100.,num=n)
    return xinput,ksvals

xinput,pvals = ks_test_over_linspace(0.25)
plt.plot(xinput,pvals)
def ks_evaluation_plot_statistics(beta):
    reps = 100
    n=100
    Xinput = np.zeros((reps,n))
    Pval = np.zeros((reps,n))
    for i in tqdm(range(reps)):

        xinput,pvals = ks_test_over_linspace(beta)
        Xinput[i,:] = xinput
        Pval[i,:] = pvals
    return Pval,xinput
Pval,xinput = ks_evaluation_plot_statistics(0.4)
plt.figure(figsize=(10,8));
pmean = np.mean(Pval,axis=0)
lc,uc = np.percentile(Pval,[2.5,97.5],axis=0)
plt.plot(xinput,pmean);
plt.fill_between(xinput,lc,uc,alpha=0.1);
plt.xlabel('time (t)',fontsize=18);
plt.ylabel('KS-statistic',fontsize=18);
cval = np.sqrt(0.02)*1.36
plt.plot([0,100.],[cval,cval],'k--');
plt.title(r'$\beta=0.4$');
#plt.xlim(0.,5.);
Pval,xinput = ks_evaluation_plot_statistics(0.05)
plt.figure(figsize=(10,8));
pmean = np.mean(Pval,axis=0)
lc,uc = np.percentile(Pval,[2.5,97.5],axis=0)
plt.plot(xinput,pmean);
plt.fill_between(xinput,lc,uc,alpha=0.1);
plt.xlabel('time (t)',fontsize=18);
plt.ylabel('KS-statistic',fontsize=18);
cval = np.sqrt(0.02)*1.36
plt.plot([0,100.],[cval,cval],'k--');
plt.title(r'$\beta=0.05$');
Pval,xinput = ks_evaluation_plot_statistics(0.25)
plt.figure(figsize=(10,8));
pmean = np.mean(Pval,axis=0)
lc,uc = np.percentile(Pval,[2.5,97.5],axis=0)
plt.plot(xinput,pmean);
plt.fill_between(xinput,lc,uc,alpha=0.1);
plt.xlabel('time (t)',fontsize=18);
plt.ylabel('KS-statistic',fontsize=18);
cval = np.sqrt(0.02)*1.36
plt.plot([0,100.],[cval,cval],'k--');
plt.title(r'$\beta=0.25$');<jupyter_output><empty_output><jupyter_text>Fit simpler model  Create simple test model data<jupyter_code>betas = np.random.uniform(high=1.0,low=0.0,size=1000)
inputs = None
outputs = None
for beta in betas:
    ts,_,Is,_ = SIR(beta,0.1,N=1000,T=100)
    binput = np.vstack((beta*np.ones(100),ts))
    
    output = model.predict(binput.T)
    mu,sigma = pydra.get_stats(output,distribution='Gamma')
    boutput = np.vstack((mu,sigma))
    if inputs is None:
        inputs = binput.copy()
    else:
        inputs = np.hstack((binput,inputs))
    
    if outputs is None:
        outputs = boutput.copy()
    else:
        outputs = np.hstack((boutput,outputs))
print(outputs.shape)
print(inputs.shape)
import keras
simple_model = keras.models.Sequential()
simple_model.add(keras.layers.Dense(units=128, activation='relu', input_dim=2))
simple_model.add(keras.layers.Dense(units=128, activation='relu',))
simple_model.add(keras.layers.Dense(units=128, activation='relu',))
simple_model.add(keras.layers.Dense(units=2, activation=None))
simple_model.compile(loss='mean_squared_error',
              optimizer='adam',
              metrics=['accuracy'])
simple_model.fit(x=inputs.T,y=outputs.T,batch_size=500,epochs=100)
betas = np.linspace(0.,0.5,num=9)
plt.figure(figsize=(12,8));
for i,beta in enumerate(betas):
    plt.subplot(3,3,i+1);
    
    t_test = np.float32(np.linspace(0.,100,num=n))
    x_test = np.vstack((beta*np.ones(n),t_test)).T
    output = simple_model.predict(x_test)

    original_output = model.predict(x_test)
    mu,sigma = pydra.get_stats(original_output,distribution='Gamma')

    plt.plot(t_test,output[:,0],label='mean');
    plt.plot(t_test,mu,label='true mean');


    plt.plot(t_test,output[:,1],label='variance');
    plt.plot(t_test,sigma,label='true variance');
tfjs.converters.save_keras_model(simple_model, './SIR_model')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.17759206059023244, 'fraction_numerical': 0.04596500391747192, 'mean_word_length': 4.824334600760456, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '315902', 'license': 'MIT', 'n_tokens_mistral': 3381, 'n_tokens_neox': 3165, 'n_words': 432, 'path': 'notebooks/SIR_binomial_demo.ipynb', 'repo_name': 'QCaudron/pydra', 'size': '7658'}