<jupyter_start><jupyter_text>Set up a Function for Cost Evaluations<jupyter_code>#####################################
##                                 ##
##   Generic cost fn Evaluation    ##
##                                 ##
#####################################

#Define the particular cost function

def evaluate_cost(t,u_t,S_t,I_t,Ic_t,cost_params):
    #unpack cost params
    Omega = cost_params
    
    #evaluate costs from infections/death
    nt = len(t)
    pos = 2
    cost_health = np.sum(Ic_t[pos,:,nt-1])
  
    #evaluate costs from implementing controls:
    cost_NPI = Omega*np.trapz(np.sum(((1-u_t)**2)*S_t,0),t)
    return (cost_health + cost_NPI)


#Set a function to compute outcomes and evaluate cost function

def get_cost(x, gradx, IC, cost_params, evaluate_cost,print_figs = False):
    
    
    #set up the time-dependent contact matrix based on control scheme
    nu = int((len(x) + 2)/(M + 1))
    t_con = np.append(np.append(0, x[:(nu - 2)]),Tf)
    u_con = np.reshape(x[(nu - 2):],(M,nu))
    def contactMatrix(t):
        u = np.zeros(M)
        if t == 0:
            return contactMatrix_0(t) #this is needed because rescalings are done inside pyrosstsi based on this
        else:
            for i in range(M):
                u[i] = np.interp(t,t_con,u_con[i,:])
            return np.outer(u,u)*contactMatrix_0(t)

    parameters['contactMatrix'] = contactMatrix
    
    #let's not messs around with poorly sorted lists
    for i in range(nu - 1):
        if t_con[i + 1] < t_con[i]:
            return float('inf')
    
    #set up and solve
    model = pyrosstsi.deterministic.Simulator(parameters)
    data  = model.simulate(IC)
    #unpack and rescale simulation output
    t = data['t']; S_t = data['S_t']; I_t = data['I_t']; Ic_t = data['Ic_t'] 
    
    #compute cost function
    nt = len(t)
    u_t = np.zeros((M,nt))
    for j in range(M):
        u_t[j,:] = np.interp(t,t_con,u_con[j,:])
        
    #Standard SIR plot
    if print_figs:
        plt.figure()    
        plt.plot(t,np.sum(S_t,0),'r', label = 'Susceptible')           #all susceptible
        plt.plot(t,np.sum(I_t,0),'b', label = 'Infected')              #all Infected
        plt.plot(t,np.sum(Ic_t[0,:,:],0),'g', label = subclasses[0])   #all recovered
        plt.legend()
        plt.xlabel('time (days)');       plt.xlim(0,Tf)
        plt.ylabel('number of people');  plt.ylim(0,1)
    
    return evaluate_cost(t,u_t,S_t,I_t,Ic_t,cost_params)


###################################
##                               ##
##       Optimize Controls       ##
##                               ##
###################################

###################################
##     Herd Immunity Example     ##
###################################


#Define input parameters to the cost function:
Omega = 10**-3   #input parameter to trial cost function.  Omega scales cost of lockdown to statistical value of life.
                 #low values represent cases where lockdowns are cheap (or disease is deadly)
                 #Generalizations should include age-structuring.
    
#bundle all input parameters together
cost_params = Omega
    
nu = 5           #how many interpolating points (including endpoints)

#set an initial guess (evenly spaced points, no intervention)
x0 = np.append(np.linspace(Tf/(nu - 1), Tf*(nu-2)/(nu-1),nu-2), np.ones(nu*M))
nx = len(x0)

#We use nlopt for optimization.
opt = nlopt.opt(nlopt.LN_COBYLA, nx)
opt.set_lower_bounds(np.zeros(nx))
opt.set_upper_bounds(np.append(Tf*np.ones(nu - 2),np.ones(nx - (nu - 2))))
opt.set_min_objective(lambda x, gradx: get_cost(x, gradx, IC, cost_params,evaluate_cost))
opt.set_xtol_rel(1e-3)#-4
x_opt = opt.optimize(x0)
minf = opt.last_optimum_value()

#group results and plot
t_con = np.append(np.append(0, x_opt[:(nu - 2)]),Tf)
u_con = np.reshape(x_opt[(nu - 2):],(M,nu))

plt.figure()
for i in range(M):
    lbl = ['age group ', str(i)]
    lbl = ''.join(lbl)
    plt.plot(t_con,u_con[i,:], label = lbl)
plt.legend()
plt.xlabel('time (days)');       plt.xlim(0,Tf)
plt.ylabel('fraction of normal contacts');  plt.ylim(0,1)

print(get_cost(x_opt, [], IC, cost_params,evaluate_cost,True))
###################################
##                               ##
##       Optimize Controls       ##
##                               ##
###################################

###################################
##        Lockdown Example       ##
###################################


#Define input parameters to the cost function:
Omega = 10**-3*2/T   #input parameter to trial cost function.  Omega scales cost of lockdown to statistical value of life.
                 #low values represent cases where lockdowns are cheap (or disease is deadly)
                 #Generalizations should include age-structuring.
    
#bundle all input parameters together
cost_params = Omega
    
nu = 5           #how many interpolating points (including endpoints)

#set an initial guess (evenly spaced points, no intervention)
x0 = np.append(np.linspace(Tf/(nu - 1), Tf*(nu-2)/(nu-1),nu-2), np.ones(nu*M))
nx = len(x0)

#We use nlopt for optimization.
opt = nlopt.opt(nlopt.LN_COBYLA, nx)
opt.set_lower_bounds(np.zeros(nx))
opt.set_upper_bounds(np.append(Tf*np.ones(nu - 2),np.ones(nx - (nu - 2))))
opt.set_min_objective(lambda x, gradx: get_cost(x, gradx, IC, cost_params,evaluate_cost))
opt.set_xtol_rel(1e-3)#-4
x_opt = opt.optimize(x0)
minf = opt.last_optimum_value()

#group results and plot
t_con = np.append(np.append(0, x_opt[:(nu - 2)]),Tf)
u_con = np.reshape(x_opt[(nu - 2):],(M,nu))

plt.figure()
for i in range(M):
    lbl = ['age group ', str(i)]
    lbl = ''.join(lbl)
    plt.plot(t_con,u_con[i,:], label = lbl)
plt.legend()
plt.xlabel('time (days)');       plt.xlim(0,Tf)
plt.ylabel('fraction of normal contacts');  plt.ylim(0,1)

print(get_cost(x_opt, [], IC, cost_params,evaluate_cost,True))
Omega*T<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1931069437404967, 'fraction_numerical': 0.011488427099172157, 'mean_word_length': 3.008124576844956, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 5, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '64798', 'license': 'MIT', 'n_tokens_mistral': 2125, 'n_tokens_neox': 2033, 'n_words': 475, 'path': 'examples/tsi/ex02_control.ipynb', 'repo_name': 'hidekb/pyross', 'size': '5919'}