# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## Load contact matrices from Wuhan and tailor it to the camp of interest

import pandas as pd
import numpy as np

wuhan_matrix=pd.read_csv('Contact_matrix_wuhan.csv',index_col=0)
moria_params=pd.read_csv('moria_params.csv',index_col=0)

moria_params_pop=moria_params[moria_params['Variable'] == 'Population_structure']

moria_params_pop_value=moria_params_pop.Value.values

contact_matrix=wuhan_matrix.values

contact_matrix.shape

# +
# we have 1-16 in 5 year bands

age_limits=np.array([0,10,20,30,40,50,60,70,80])
n_categories=len(age_limits)-1
ind_limits=(age_limits/5).astype(int)

p=[]
for cc in range(n_categories):
    for i in range(ind_limits[cc],ind_limits[cc+1]):
        p.append(moria_params_pop_value[cc]/(ind_limits[cc+1]-ind_limits[cc]))

M=np.zeros((n_categories,n_categories))

for rr in range(n_categories):
    for cc in range(n_categories):
        V2=0
        sump=sum(p[(ind_limits[cc]+1):ind_limits[cc+1] ])*sum(p[(ind_limits[rr]+1):ind_limits[rr+1]])
        for i in range(ind_limits[cc]+1,ind_limits[cc+1]):
            for j in range(ind_limits[rr]+1,ind_limits[rr+1]):
                    V2+=contact_matrix[j-1,i-1]*p[i-1]*p[j-1]/sump
        M[rr,cc] = V2


# -

col_names=[]
for i in range(M.shape[0]):
    col_names.append('Age '+str(age_limits[i])+'-'+str(age_limits[i+1]))
moria_contact_dict={}
for col,index in zip(col_names,list(range(n_categories))):
    moria_contact_dict[col]=M[:,index]

moria_contact_matrix=pd.DataFrame.from_dict(moria_contact_dict)

moria_contact_matrix.to_csv('moria_contact_matrix.csv')



StarcoderdataJupyterScriptsDedupFiltered