""" DisMod III Covariate Simulation Study

For more realistic simulated data, use the age, year, and country data

from a specified disease model. For example:

$ python2.5 generate_covariate_data.py 175

"""

import sys

import os

import optparse

import csv

import random

import pylab as pl

import numpy as np

import scipy.linalg

import pymc as mc

import simplejson as json

import dismod3

import dismod3.utils

from dismod3.utils import clean, trim, NEARLY_ZERO

from dismod3.disease_json import DiseaseJson

import dismod3.gbd_disease_model as model

# simulation parameters

usage = 'usage: %prog [options] disease_model_id'

parser = optparse.OptionParser(usage)

parser.add_option('-n', '--studysize', dest='study_size', default='10000',

help='number of subjects in each age-range')

parser.add_option('-d', '--dispersion', dest='dispersion', default='10000',

help='dispersion of study-level beta-binomial in data generation process')

(options, args) = parser.parse_args()

# check that args are correct

if len(args) == 1:

try:

id = int(args[0])

except ValueError:

parser.error('disease_model_id must be an integer')

exit()

else:

parser.error('incorrect number of arguments')

exit()

# fetch requested model

dm = dismod3.get_disease_model(id)

# define ground truth

age_len = dismod3.MAX_AGE

ages = np.arange(age_len, dtype='float')

print 'defining model transition parameters'

truth = {}

# all-cause mortality-rate

m = np.array(

0.2467051 ])

for region in dismod3.gbd_regions:

for year in dismod3.gbd_years:

for sex in dismod3.gbd_sexes:

key = dismod3.gbd_key_for('%s', region, year, sex)

if clean(region) == 'north_america_high_income':

regional_offset = 0.

else:

regional_offset = -.5

time_offset = (int(year)-1997)/10.

if clean(sex) == 'male':

sex_offset = .1

else:

sex_offset = 0.

# incidence rate

i = mc.invlogit(mc.logit(.012 * mc.invlogit((ages - 44) / 3)) + regional_offset + time_offset + sex_offset)

truth[key % 'incidence'] = i

# remission rate

r = 0. * ages

truth[key % 'remission'] = r

# excess-mortality rate

f = .085 * (ages / 100) ** 2.5

truth[key % 'excess-mortality'] = f

## compartmental model (bins S, C, D, M)

SCDM = np.zeros([4, age_len])

SCDM[0,0] = 1.

for a in range(age_len - 1):

A = [[-i[a]-m[a], r[a] , 0., 0.],

[ i[a] , -r[a]-m[a]-f[a], 0., 0.],

[ m[a], m[a] , 0., 0.],

[ 0., f[a], 0., 0.]]

SCDM[:,a+1] = trim(np.dot(scipy.linalg.expm2(A), SCDM[:,a]), 0, 1)

S = SCDM[0,:]

C = SCDM[1,:]

# prevalence = # with condition / (# with condition + # without)

p = C / (S + C + NEARLY_ZERO)

truth[key % 'prevalence'] = p

truth[key % 'relative-risk'] = (m + f) / m

# duration = E[time in bin C]

pr_exit = 1 - r - m - f

X = np.empty(len(pr_exit))

t = 1.

for a in xrange(len(X) - 1, -1, -1):

X[a] = t * pr_exit[a]

t = 1 + X[a]

truth[key % 'duration'] = X

# generate synthetic data from fictitious ground truth, using data age

# structure from selected disease model

print '\nsimulating noisy realizations'

dispersion = float(options.dispersion)

n = float(options.study_size)

def generate_synthetic_data(truth, key, d):

return d

data = []

for d in dm.data:

t = d['data_type']

r = d['gbd_region']

y = d['year_start']

s = d['sex']

key = dismod3.gbd_key_for(t.replace(' data',''), r, y, s)

if key not in truth.keys():

continue

data += [generate_synthetic_data(truth, key, d)]

def data_dict_for_csv(d):

return c

f_file = open('simulated_data.tsv', 'w')

csv_f = csv.writer(f_file, dialect=csv.excel_tab)

col_names = sorted(data_dict_for_csv(data[0]).keys())

csv_f.writerow(col_names)

for d in data:

dd = data_dict_for_csv(d)

csv_f.writerow([dd[c] for c in col_names])

f_file.close()

# upload a new disease model which knows ground truth (but needs to

# have the data from the csv loaded separately)

dm.data = []

dm.params.pop('id')

dm.id = -1

for key in truth:

dm.set_truth(key, truth[key])

url = dismod3.post_disease_model(dm)

print url