Full Code of mp2893/med2vec for AI

Repository: mp2893/med2vec
Branch: master
Commit: ee7066dcb3c3
Files: 5
Total size: 28.6 KB

Directory structure:
gitextract_7k7owaz_/

├── .gitignore
├── LICENSE
├── README.md
├── med2vec.py
└── process_mimic.py

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitignore
================================================
.swp


================================================
FILE: LICENSE
================================================
Copyright (c) 2016, mp2893
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

* Neither the name of Med2Vec nor the names of its
  contributors may be used to endorse or promote products derived from
  this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


================================================
FILE: README.md
================================================
Med2Vec
=========================================

Med2Vec is a multi-layer representation learning tool for learning code representations and visit representations from EHR datasets.

[![Med2Vec Coordinate-wise Interpretation Demo](http://www.cc.gatech.edu/~echoi48/images/med2vec_interpret.png)](https://youtu.be/UR_f2rmMJkk?t=2m34s "Med2Vec Coordinate-wise Interpretation Demo - Click to Watch!")
Med2Vec embeddings not only help improve predictive performance of healthcare applications, but also enable the interpretation of the learned code representations in a coodinate-wise manner. You can see that these six coordinates (chosen by their strong correlation with patient severity level) of the code representation space demonstrate medically coherent groups of symptoms (diagnoses, medications, and procedures). 

#### Relevant Publications

Med2Vec implements an algorithm introduced in the following [paper](http://www.kdd.org/kdd2016/subtopic/view/multi-layer-representation-learning-for-medical-concepts):

    Multi-layer Representation Learning for Medical Concepts
	Edward Choi, Mohammad Taha Bahadori, Elizabeth Searles, Catherine Coffey, 
	Michael Thompson, James Bost, Javier Tejedor-Sojo, Jimeng Sun
	KDD 2016, pp.1495-1504

#### Running Med2Vec

**STEP 1: Installation**  

1. Install [python](https://www.python.org/), [Theano](http://deeplearning.net/software/theano/index.html). We use Python 2.7, Theano 0.7. Theano can be easily installed in Ubuntu as suggested [here](http://deeplearning.net/software/theano/install_ubuntu.html#install-ubuntu)

2. If you plan to use GPU computation, install [CUDA](https://developer.nvidia.com/cuda-downloads)

3. Download/clone the Med2Vec code  

**STEP 2: Fast way to test Med2Vec with MIMIC-III**

This step describes how to run, with minimum number of steps, Med2Vec using MIMIC-III. 

0. You will first need to request access for [MIMIC-III](https://mimic.physionet.org/gettingstarted/access/), a publicly avaiable electronic health records collected from ICU patients over 11 years. 

1. You can use "process_mimic.py" to process MIMIC-III dataset and generate a suitable training dataset for Med2Vec.
Place the script to the same location where the MIMIC-III CSV files are located, and run the script. 
The execution command is `python process_mimic.py ADMISSIONS.csv DIAGNOSES_ICD.csv <output file>`.
Instructions are described inside the script. 

2. Run Med2Vec using the ".seqs" file generated by process_mimic.py, using the following command.
`python med2vec.py <seqs file> 4894 <output path>`
where 4894 is the number of unique ICD9 diagnosis codes in the dataset.
As described in the paper, however, it is a good idea to use the grouped codes for training the Softmax component of Med2Vec. Therefore we recommend using the following command instead.
`python med2vec.py <seqs file> 4894 <output path> --label_file <3digitICD9.seqs file> --n_output_codes 942`
where 942 is the number of unique 3-digit ICD9 diagnosis codes in the dataset.
You can also use ".3digitICD9.seqs" to begin with, if you interested in learning the representation of 3-digit ICD9 codes only, using the following command.
`python med2vec.py <3digitICD9.seqs file> 942 <output path>`

3. As suggested in STEP 4, you might want to adjust the hyper-parameters. 
I recommend decreasing the `--batch_size` to 100 or so, since the default value 1,000 is too big considering the small number of patients in MIMIC-III datasets. 
There are only 7,500 patients who made more than a single visit, and most of them have only two visits.

**STEP 3: Preparing training data**  

1. Med2Vec training data need to be a Python Pickled list of list of medical codes (e.g. diagnosis codes, medication codes, or procedure codes). 
First, medical codes need to be converted to an integer. Then a single visit can be converted as a list of integers. 
For example, [5,8,15] means the patient was assigned with code 5, 8, and 15 at a certain visit. 
If a patient made two visits [1,2,3] and [4,5,6,7], it can be converted to a list of list [[1,2,3], [4,5,6,7]]. 
If there are multiple patients, each patient must be delimited by a list [-1]. 
For example, [[1,2,3], [4,5,6,7], [-1], [2,4], [8,3,1], [3]] means there are two patients where the first patient made two visits and the second patient made three visits. 
This list of list needs to be pickled using cPickle. We will refer to this file as the "visit file".

2. The total number of unique medical codes is required to run Med2Vec. 
For example, if the dataset is using 14,000 diagnosis codes and 11,000 procedure codes, the total number is 25,000. 
Note that using a huge number of codes could lead to memory problems, depending on your RAM/VRAM (thanks for the tip [tRosenflanz](https://github.com/tRosenflanz))

3. For a faster training, you can provide an additional dataset, which is simply the same dataset in step 1, but with grouped medical codes. 
For example, ICD9 diagnosis codes can be grouped into 283 categories by using [CCS](https://www.hcup-us.ahrq.gov/toolssoftware/ccs/ccs.jsp) groupers. 
You will still be able to learn the code representations for the original un-grouped codes. 
The grouped dataset is used only for speeding up the training speed. (Refer to section 4.4 of the paper) 
The grouped dataset should be prepared in the same way as the dataset in step 1. We will refer to this grouped dataset as the "label file".

4. Same as step 2, you will need to remember the total number of unique grouped codes if you plan to use this grouped dataset.

5. If you wish to use patient demographic information (e.g. age, weight, gender) you need to create a demographics vector for each visit the patient made. 
For example, if you are using age (real-valued) and ethnicity(categorical, assume 6 categories), you can create a vector such as [45.0, 0, 0, 0, 0, 1, 0]. 
Similar to the [-1] vector in step 1, each patient is delimited with an all-zero vector. 
Therefore the demographic information will be a pickled matrix where column size is the size of the demographics vector and row size is the number of total visits of all patients plus the delimiters. 
We will refer to this file as the "demo file".

6. Similar to step 2, you will need to remeber the size of the demographics vector if you plan to use the demo file. 
In the example of step 5, the size of the demographics vector is 7.

**STEP 4: Running Med2Vec**  

1. The minimum input you need to run Med2Vec is the visit file, the number of unique medical codes and the output path
`python med2vec <path/to/visit_file> <the number of unique medical codes> <path/to/output>`  

2. Specifying `--verbose` option will print training process after each 10 mini-batches.

3. Additional options can be specified such as the size of the code representation, the size of the visit representation and the number of epochs. Detailed information can be accessed by `python med2vec --help`

**STEP 5: Looking at your results**  

Med2Vec produces a model file after each epoch. The model file is generated by [numpy.savez_compressed](http://docs.scipy.org/doc/numpy-1.10.1/reference/generated/numpy.savez_compressed.html).

The 2D scatterplot of the learned code representations would look similar to [this](http://mp2893.com/scatterplot/nnsg_h200e49_category10.html).
(This is the scatterplot of the code representations trained with Non-negative Skip-gram, which is essentially Med2Vec minus the visit-level training)


================================================
FILE: med2vec.py
================================================
#################################################################
# Code written by Edward Choi (mp2893@gatech.edu)
# For bug report, please contact author using the email address
#################################################################

import sys, random
import numpy as np
import cPickle as pickle
from collections import OrderedDict
import argparse

import theano
import theano.tensor as T
from theano import config

def numpy_floatX(data):
	return np.asarray(data, dtype=config.floatX)

def unzip(zipped):
	new_params = OrderedDict()
	for k, v in zipped.iteritems():
		new_params[k] = v.get_value()
	return new_params

def init_params(options):
	params = OrderedDict()

	numXcodes = options['numXcodes']
	numYcodes = options['numYcodes']
	embDimSize= options['embDimSize']
	demoSize = options['demoSize']
	hiddenDimSize = options['hiddenDimSize']

	params['W_emb'] = np.random.uniform(-0.01, 0.01, (numXcodes, embDimSize)).astype(config.floatX) #emb matrix needs an extra dimension for the time
	params['b_emb'] = np.zeros(embDimSize).astype(config.floatX)
	params['W_hidden'] = np.random.uniform(-0.01, 0.01, (embDimSize+demoSize, hiddenDimSize)).astype(config.floatX) #emb matrix needs an extra dimension for the time
	params['b_hidden'] = np.zeros(hiddenDimSize).astype(config.floatX)
	if numYcodes > 0:
		params['W_output'] = np.random.uniform(-0.01, 0.01, (hiddenDimSize, numYcodes)).astype(config.floatX) #emb matrix needs an extra dimension for the time
		params['b_output'] = np.zeros(numYcodes).astype(config.floatX)
	else:
		params['W_output'] = np.random.uniform(-0.01, 0.01, (hiddenDimSize, numXcodes)).astype(config.floatX) #emb matrix needs an extra dimension for the time
		params['b_output'] = np.zeros(numXcodes).astype(config.floatX)

	return params

def load_params(options):
	params = np.load(options['modelFile'])
	return params

def init_tparams(params):
	tparams = OrderedDict()
	for k, v in params.iteritems():
		tparams[k] = theano.shared(v, name=k)
	return tparams

def build_model(tparams, options):
	x = T.matrix('x', dtype=config.floatX)
	d = T.matrix('d', dtype=config.floatX)
	y = T.matrix('y', dtype=config.floatX)
	mask = T.vector('mask', dtype=config.floatX)

	logEps = options['logEps']

	emb = T.maximum(T.dot(x, tparams['W_emb']) + tparams['b_emb'],0)
	if options['demoSize'] > 0: emb = T.concatenate((emb, d), axis=1)
	visit = T.maximum(T.dot(emb, tparams['W_hidden']) + tparams['b_hidden'],0)
	results = T.nnet.softmax(T.dot(visit, tparams['W_output']) + tparams['b_output'])
	
	mask1 = (mask[:-1] * mask[1:])[:,None]
	mask2 = (mask[:-2] * mask[1:-1] * mask[2:])[:,None]
	mask3 = (mask[:-3] * mask[1:-2] * mask[2:-1] * mask[3:])[:,None]
	mask4 = (mask[:-4] * mask[1:-3] * mask[2:-2] * mask[3:-1] * mask[4:])[:,None]
	mask5 = (mask[:-5] * mask[1:-4] * mask[2:-3] * mask[3:-2] * mask[4:-1] * mask[5:])[:,None]

	t = None
	if options['numYcodes'] > 0: t = y
	else: t = x

	forward_results =  results[:-1] * mask1
	forward_cross_entropy = -(t[1:] * T.log(forward_results + logEps) + (1. - t[1:]) * T.log(1. - forward_results + logEps))

	forward_results2 =  results[:-2] * mask2
	forward_cross_entropy2 = -(t[2:] * T.log(forward_results2 + logEps) + (1. - t[2:]) * T.log(1. - forward_results2 + logEps))

	forward_results3 =  results[:-3] * mask3
	forward_cross_entropy3 = -(t[3:] * T.log(forward_results3 + logEps) + (1. - t[3:]) * T.log(1. - forward_results3 + logEps))

	forward_results4 =  results[:-4] * mask4
	forward_cross_entropy4 = -(t[4:] * T.log(forward_results4 + logEps) + (1. - t[4:]) * T.log(1. - forward_results4 + logEps))

	forward_results5 =  results[:-5] * mask5
	forward_cross_entropy5 = -(t[5:] * T.log(forward_results5 + logEps) + (1. - t[5:]) * T.log(1. - forward_results5 + logEps))

	backward_results =  results[1:] * mask1
	backward_cross_entropy = -(t[:-1] * T.log(backward_results + logEps) + (1. - t[:-1]) * T.log(1. - backward_results + logEps))

	backward_results2 =  results[2:] * mask2
	backward_cross_entropy2 = -(t[:-2] * T.log(backward_results2 + logEps) + (1. - t[:-2]) * T.log(1. - backward_results2 + logEps))

	backward_results3 =  results[3:] * mask3
	backward_cross_entropy3 = -(t[:-3] * T.log(backward_results3 + logEps) + (1. - t[:-3]) * T.log(1. - backward_results3 + logEps))

	backward_results4 =  results[4:] * mask4
	backward_cross_entropy4 = -(t[:-4] * T.log(backward_results4 + logEps) + (1. - t[:-4]) * T.log(1. - backward_results4 + logEps))

	backward_results5 =  results[5:] * mask5
	backward_cross_entropy5 = -(t[:-5] * T.log(backward_results5 + logEps) + (1. - t[:-5]) * T.log(1. - backward_results5 + logEps))

	visit_cost1 = (forward_cross_entropy.sum(axis=1).sum(axis=0) + backward_cross_entropy.sum(axis=1).sum(axis=0)) / (mask1.sum() + logEps)
	visit_cost2 = (forward_cross_entropy2.sum(axis=1).sum(axis=0) + backward_cross_entropy2.sum(axis=1).sum(axis=0)) / (mask2.sum() + logEps)
	visit_cost3 = (forward_cross_entropy3.sum(axis=1).sum(axis=0) + backward_cross_entropy3.sum(axis=1).sum(axis=0)) / (mask3.sum() + logEps)
	visit_cost4 = (forward_cross_entropy4.sum(axis=1).sum(axis=0) + backward_cross_entropy4.sum(axis=1).sum(axis=0)) / (mask4.sum() + logEps)
	visit_cost5 = (forward_cross_entropy5.sum(axis=1).sum(axis=0) + backward_cross_entropy5.sum(axis=1).sum(axis=0)) / (mask5.sum() + logEps)

	windowSize = options['windowSize']
	visit_cost = visit_cost1
	if windowSize == 2:
		visit_cost = visit_cost1 + visit_cost2
	elif windowSize == 3:
		visit_cost = visit_cost1 + visit_cost2 + visit_cost3
	elif windowSize == 4:
		visit_cost = visit_cost1 + visit_cost2 + visit_cost3 + visit_cost4
	elif windowSize == 5:
		visit_cost = visit_cost1 + visit_cost2 + visit_cost3 + visit_cost4 + visit_cost5

	iVector = T.vector('iVector', dtype='int32')
	jVector = T.vector('jVector', dtype='int32')
	preVec = T.maximum(tparams['W_emb'],0)
	norms = (T.exp(T.dot(preVec, preVec.T))).sum(axis=1)
	emb_cost = -T.log((T.exp((preVec[iVector] * preVec[jVector]).sum(axis=1)) / norms[iVector]) + logEps)

	total_cost = visit_cost + T.mean(emb_cost) + options['L2_reg'] * (tparams['W_emb'] ** 2).sum()

	if options['demoSize'] > 0 and options['numYcodes'] > 0: return x, d, y, mask, iVector, jVector, total_cost
	elif options['demoSize'] == 0 and options['numYcodes'] > 0: return x, y, mask, iVector, jVector, total_cost
	elif options['demoSize'] > 0 and options['numYcodes'] == 0: return x, d, mask, iVector, jVector, total_cost
	else: return x, mask, iVector, jVector, total_cost

def adadelta(tparams, grads, x, mask, iVector, jVector, cost, options, d=None, y=None):
	zipped_grads = [theano.shared(p.get_value() * numpy_floatX(0.), name='%s_grad' % k) for k, p in tparams.iteritems()]
	running_up2 = [theano.shared(p.get_value() * numpy_floatX(0.), name='%s_rup2' % k) for k, p in tparams.iteritems()]
	running_grads2 = [theano.shared(p.get_value() * numpy_floatX(0.), name='%s_rgrad2' % k) for k, p in tparams.iteritems()]

	zgup = [(zg, g) for zg, g in zip(zipped_grads, grads)]
	rg2up = [(rg2, 0.95 * rg2 + 0.05 * (g ** 2)) for rg2, g in zip(running_grads2, grads)]

	if options['demoSize'] > 0 and options['numYcodes'] > 0:
		f_grad_shared = theano.function([x, d, y, mask, iVector, jVector], cost, updates=zgup + rg2up, name='adadelta_f_grad_shared')
	elif options['demoSize'] == 0 and options['numYcodes'] > 0:
		f_grad_shared = theano.function([x, y, mask, iVector, jVector], cost, updates=zgup + rg2up, name='adadelta_f_grad_shared')
	elif options['demoSize'] > 0 and options['numYcodes'] == 0:
		f_grad_shared = theano.function([x, d, mask, iVector, jVector], cost, updates=zgup + rg2up, name='adadelta_f_grad_shared')
	else:
		f_grad_shared = theano.function([x, mask, iVector, jVector], cost, updates=zgup + rg2up, name='adadelta_f_grad_shared')

	updir = [-T.sqrt(ru2 + 1e-6) / T.sqrt(rg2 + 1e-6) * zg for zg, ru2, rg2 in zip(zipped_grads, running_up2, running_grads2)]
	ru2up = [(ru2, 0.95 * ru2 + 0.05 * (ud ** 2)) for ru2, ud in zip(running_up2, updir)]
	param_up = [(p, p + ud) for p, ud in zip(tparams.values(), updir)]

	f_update = theano.function([], [], updates=ru2up + param_up, on_unused_input='ignore', name='adadelta_f_update')

	return f_grad_shared, f_update

def load_data(xFile, dFile, yFile):
	seqX = np.array(pickle.load(open(xFile, 'rb')))
	seqD = []
	if len(dFile) > 0: seqD = np.asarray(pickle.load(open(dFile, 'rb')), dtype=config.floatX)
	seqY = []
	if len(yFile) > 0: seqY = np.array(pickle.load(open(yFile, 'rb')))
	return seqX, seqD, seqY

def pickTwo(codes, iVector, jVector):
	for first in codes:
		for second in codes:
			if first == second: continue
			iVector.append(first)
			jVector.append(second)
	
def padMatrix(seqs, labels, options):
	n_samples = len(seqs)
	iVector = []
	jVector = []
	numXcodes = options['numXcodes']
	numYcodes = options['numYcodes']

	if numYcodes > 0:
		x = np.zeros((n_samples, numXcodes)).astype(config.floatX)
		y = np.zeros((n_samples, numYcodes)).astype(config.floatX)
		mask = np.zeros((n_samples,)).astype(config.floatX)
		for idx, (seq, label) in enumerate(zip(seqs, labels)):
			if not seq[0] == -1:
				x[idx][seq] = 1.
				y[idx][label] = 1.
				pickTwo(seq, iVector, jVector)
				mask[idx] = 1.
		return x, y, mask, iVector, jVector
	else:
		x = np.zeros((n_samples, numXcodes)).astype(config.floatX)
		mask = np.zeros((n_samples,)).astype(config.floatX)
		for idx, seq in enumerate(seqs):
			if not seq[0] == -1:
				x[idx][seq] = 1.
				pickTwo(seq, iVector, jVector)
				mask[idx] = 1.
		return x, mask, iVector, jVector

def train_med2vec(seqFile='seqFile.txt', 
				demoFile='demoFile.txt',
				labelFile='labelFile.txt',
				outFile='outFile.txt',
				modelFile='modelFile.txt',
				L2_reg=0.001,
				numXcodes=20000, 
				numYcodes=20000, 
				embDimSize=1000,
				hiddenDimSize=2000,
				batchSize=100,
				demoSize=2,
				logEps=1e-8,
				windowSize=1,
				verbose=False,
				maxEpochs=1000):

	options = locals().copy()
	print 'initializing parameters'
	params = init_params(options)
	#params = load_params(options)
	tparams = init_tparams(params)

	print 'building models'
	f_grad_shared = None
	f_update = None
	if demoSize > 0 and numYcodes > 0:
		x, d, y, mask, iVector, jVector, cost = build_model(tparams, options)
		grads = T.grad(cost, wrt=tparams.values())
		f_grad_shared, f_update = adadelta(tparams, grads, x, mask, iVector, jVector, cost, options, d=d, y=y)
	elif demoSize == 0 and numYcodes > 0:
		x, y, mask, iVector, jVector, cost = build_model(tparams, options)
		grads = T.grad(cost, wrt=tparams.values())
		f_grad_shared, f_update = adadelta(tparams, grads, x, mask, iVector, jVector, cost, options, y=y)
	elif demoSize > 0 and numYcodes == 0:
		x, d, mask, iVector, jVector, cost = build_model(tparams, options)
		grads = T.grad(cost, wrt=tparams.values())
		f_grad_shared, f_update = adadelta(tparams, grads, x, mask, iVector, jVector, cost, options, d=d)
	else:
		x, mask, iVector, jVector, cost = build_model(tparams, options)
		grads = T.grad(cost, wrt=tparams.values())
		f_grad_shared, f_update = adadelta(tparams, grads, x, mask, iVector, jVector, cost, options)

	print 'loading data'
	seqs, demos, labels = load_data(seqFile, demoFile, labelFile)
	n_batches = int(np.ceil(float(len(seqs)) / float(batchSize)))

	print 'training start'
	for epoch in xrange(maxEpochs):
		iteration = 0
		costVector = []
		for index in random.sample(range(n_batches), n_batches):
			batchX = seqs[batchSize*index:batchSize*(index+1)]
			batchY = []
			batchD = []
			if demoSize > 0 and numYcodes > 0:
				batchY = labels[batchSize*index:batchSize*(index+1)]
				x, y, mask, iVector, jVector = padMatrix(batchX, batchY, options)
				batchD = demos[batchSize*index:batchSize*(index+1)]
				cost = f_grad_shared(x, batchD, y, mask, iVector, jVector)
			elif demoSize == 0 and numYcodes > 0:
				batchY = labels[batchSize*index:batchSize*(index+1)]
				x, y, mask, iVector, jVector = padMatrix(batchX, batchY, options)
				cost = f_grad_shared(x, y, mask, iVector, jVector)
			elif demoSize > 0 and numYcodes == 0:
				x, mask, iVector, jVector = padMatrix(batchX, batchY, options)
				batchD = demos[batchSize*index:batchSize*(index+1)]
				cost = f_grad_shared(x, batchD, mask, iVector, jVector)
			else:
				x, mask, iVector, jVector = padMatrix(batchX, batchY, options)
				cost = f_grad_shared(x, mask, iVector, jVector)
			costVector.append(cost)
			f_update()
			if (iteration % 10 == 0) and verbose: print 'epoch:%d, iteration:%d/%d, cost:%f' % (epoch, iteration, n_batches, cost)
			iteration += 1
		print 'epoch:%d, mean_cost:%f' % (epoch, np.mean(costVector))
		tempParams = unzip(tparams)
		np.savez_compressed(outFile + '.' + str(epoch), **tempParams)

def parse_arguments(parser):
	parser.add_argument('seq_file', type=str, metavar='<visit_file>', help='The path to the Pickled file containing visit information of patients')
	parser.add_argument('n_input_codes', type=int, metavar='<n_input_codes>', help='The number of unique input medical codes')
	parser.add_argument('out_file', type=str, metavar='<out_file>', help='The path to the output models. The models will be saved after every epoch')
	parser.add_argument('--label_file', type=str, default='', help='The path to the Pickled file containing grouped visit information of patients. If you are not using a grouped output, do not use this option')
	parser.add_argument('--n_output_codes', type=int, default=0, help='The number of unique output medical codes (the number of unique grouped codes). If you are not using a grouped output, do not use this option')
	parser.add_argument('--demo_file', type=str, default='', help='The path to the Pickled file containing demographic information of patients. If you are not using patient demographic information, do not use this option')
	parser.add_argument('--demo_size', type=int, default=0, help='The size of the demographic information vector. If you are not using patient demographic information, do not use this option')
	parser.add_argument('--cr_size', type=int, default=200, help='The size of the code representation (default value: 200)')
	parser.add_argument('--vr_size', type=int, default=200, help='The size of the visit representation (default value: 200)')
	parser.add_argument('--batch_size', type=int, default=1000, help='The size of a single mini-batch (default value: 1000)')
	parser.add_argument('--n_epoch', type=int, default=10, help='The number of training epochs (default value: 10)')
	parser.add_argument('--L2_reg', type=float, default=0.001, help='L2 regularization for the code representation matrix W_c (default value: 0.001)')
	parser.add_argument('--window_size', type=int, default=1, choices=[1,2,3,4,5], help='The size of the visit context window (range: 1,2,3,4,5), (default value: 1)')
	parser.add_argument('--log_eps', type=float, default=1e-8, help='A small value to prevent log(0) (default value: 1e-8)')
	parser.add_argument('--verbose', action='store_true', help='Print output after every 10 mini-batches')
	args = parser.parse_args()
	return args

if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	args = parse_arguments(parser)

	train_med2vec(seqFile=args.seq_file, demoFile=args.demo_file, labelFile=args.label_file, outFile=args.out_file, numXcodes=args.n_input_codes, numYcodes=args.n_output_codes, embDimSize=args.cr_size, hiddenDimSize=args.vr_size, batchSize=args.batch_size, maxEpochs=args.n_epoch, L2_reg=args.L2_reg, demoSize=args.demo_size, windowSize=args.window_size, logEps=args.log_eps, verbose=args.verbose)


================================================
FILE: process_mimic.py
================================================
# This script processes MIMIC-III dataset and builds longitudinal diagnosis records for patients with at least two visits.
# The output data are cPickled, and suitable for training Doctor AI or RETAIN
# Written by Edward Choi (mp2893@gatech.edu)
# Usage: Put this script to the foler where MIMIC-III CSV files are located. Then execute the below command.
# python process_mimic.py ADMISSIONS.csv DIAGNOSES_ICD.csv <output file> 

# Output files
# <output file>.seqs: Dataset that follows the format described in the README.md.
# <output file>.types: Python dictionary that maps string diagnosis codes to integer diagnosis codes.
# <output file>.3digitICD9.seqs: Dataset that follows the format described in the README.md. This uses only the first 3 digits of the ICD9 diagnosis code.
# <output file>.3digitICD9.types: Python dictionary that maps 3-digit string diagnosis codes to integer diagnosis codes.

import sys
import cPickle as pickle
from datetime import datetime

def convert_to_icd9(dxStr):
	if dxStr.startswith('E'):
		if len(dxStr) > 4: return dxStr[:4] + '.' + dxStr[4:]
		else: return dxStr
	else:
		if len(dxStr) > 3: return dxStr[:3] + '.' + dxStr[3:]
		else: return dxStr
	
def convert_to_3digit_icd9(dxStr):
	if dxStr.startswith('E'):
		if len(dxStr) > 4: return dxStr[:4]
		else: return dxStr
	else:
		if len(dxStr) > 3: return dxStr[:3]
		else: return dxStr

if __name__ == '__main__':
	admissionFile = sys.argv[1]
	diagnosisFile = sys.argv[2]
	outFile = sys.argv[3]

	print 'Building pid-admission mapping, admission-date mapping'
	pidAdmMap = {}
	admDateMap = {}
	infd = open(admissionFile, 'r')
	infd.readline()
	for line in infd:
		tokens = line.strip().split(',')
		pid = int(tokens[1])
		admId = int(tokens[2])
		admTime = datetime.strptime(tokens[3], '%Y-%m-%d %H:%M:%S')
		admDateMap[admId] = admTime
		if pid in pidAdmMap: pidAdmMap[pid].append(admId)
		else: pidAdmMap[pid] = [admId]
	infd.close()

	print 'Building admission-dxList mapping'
	admDxMap = {}
	admDxMap_3digit = {}
	infd = open(diagnosisFile, 'r')
	infd.readline()
	for line in infd:
		tokens = line.strip().split(',')
		admId = int(tokens[2])
		dxStr = 'D_' + convert_to_icd9(tokens[4][1:-1]) ############## Uncomment this line and comment the line below, if you want to use the entire ICD9 digits.
		dxStr_3digit = 'D_' + convert_to_3digit_icd9(tokens[4][1:-1])

		if admId in admDxMap: 
			admDxMap[admId].append(dxStr)
		else: 
			admDxMap[admId] = [dxStr]

		if admId in admDxMap_3digit: 
			admDxMap_3digit[admId].append(dxStr_3digit)
		else: 
			admDxMap_3digit[admId] = [dxStr_3digit]
	infd.close()

	print 'Building pid-sortedVisits mapping'
	pidSeqMap = {}
	pidSeqMap_3digit = {}
	for pid, admIdList in pidAdmMap.iteritems():
		if len(admIdList) < 2: continue

		sortedList = sorted([(admDateMap[admId], admDxMap[admId]) for admId in admIdList])
		pidSeqMap[pid] = sortedList

		sortedList_3digit = sorted([(admDateMap[admId], admDxMap_3digit[admId]) for admId in admIdList])
		pidSeqMap_3digit[pid] = sortedList_3digit
	
	print 'Building pids, dates, strSeqs'
	pids = []
	dates = []
	seqs = []
	for pid, visits in pidSeqMap.iteritems():
		pids.append(pid)
		seq = []
		date = []
		for visit in visits:
			date.append(visit[0])
			seq.append(visit[1])
		dates.append(date)
		seqs.append(seq)
	
	print 'Building pids, dates, strSeqs for 3digit ICD9 code'
	seqs_3digit = []
	for pid, visits in pidSeqMap_3digit.iteritems():
		seq = []
		for visit in visits:
			seq.append(visit[1])
		seqs_3digit.append(seq)
	
	print 'Converting strSeqs to intSeqs, and making types'
	types = {}
	newSeqs = []
	for patient in seqs:
		newPatient = []
		for visit in patient:
			newVisit = []
			for code in visit:
				if code in types:
					newVisit.append(types[code])
				else:
					types[code] = len(types)
					newVisit.append(types[code])
			newPatient.append(newVisit)
		newSeqs.append(newPatient)
	
	print 'Converting strSeqs to intSeqs, and making types for 3digit ICD9 code'
	types_3digit = {}
	newSeqs_3digit = []
	for patient in seqs_3digit:
		newPatient = []
		for visit in patient:
			newVisit = []
			for code in set(visit):
				if code in types_3digit:
					newVisit.append(types_3digit[code])
				else:
					types_3digit[code] = len(types_3digit)
					newVisit.append(types_3digit[code])
			newPatient.append(newVisit)
		newSeqs_3digit.append(newPatient)
	
	print 'Re-formatting to Med2Vec dataset'
	seqs = []
	for patient in newSeqs:
		seqs.extend(patient)
		seqs.append([-1])
	seqs = seqs[:-1]

	seqs_3digit = []
	for patient in newSeqs_3digit:
		seqs_3digit.extend(patient)
		seqs_3digit.append([-1])
	seqs_3digit = seqs_3digit[:-1]

	pickle.dump(seqs, open(outFile+'.seqs', 'wb'), -1)
	pickle.dump(types, open(outFile+'.types', 'wb'), -1)
	pickle.dump(seqs_3digit, open(outFile+'.3digitICD9.seqs', 'wb'), -1)
	pickle.dump(types_3digit, open(outFile+'.3digitICD9.types', 'wb'), -1)