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


/.gitignore:
--------------------------------------------------------------------------------
1 | .swp
2 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | Copyright (c) 2016, mp2893
 2 | All rights reserved.
 3 | 
 4 | Redistribution and use in source and binary forms, with or without
 5 | modification, are permitted provided that the following conditions are met:
 6 | 
 7 | * Redistributions of source code must retain the above copyright notice, this
 8 |   list of conditions and the following disclaimer.
 9 | 
10 | * Redistributions in binary form must reproduce the above copyright notice,
11 |   this list of conditions and the following disclaimer in the documentation
12 |   and/or other materials provided with the distribution.
13 | 
14 | * Neither the name of Med2Vec nor the names of its
15 |   contributors may be used to endorse or promote products derived from
16 |   this software without specific prior written permission.
17 | 
18 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
19 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
20 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
21 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
22 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
23 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
24 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
25 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
26 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
27 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
28 | 


--------------------------------------------------------------------------------
/process_mimic.py:
--------------------------------------------------------------------------------
  1 | # This script processes MIMIC-III dataset and builds longitudinal diagnosis records for patients with at least two visits.
  2 | # The output data are cPickled, and suitable for training Doctor AI or RETAIN
  3 | # Written by Edward Choi (mp2893@gatech.edu)
  4 | # Usage: Put this script to the foler where MIMIC-III CSV files are located. Then execute the below command.
  5 | # python process_mimic.py ADMISSIONS.csv DIAGNOSES_ICD.csv <output file> 
  6 | 
  7 | # Output files
  8 | # <output file>.seqs: Dataset that follows the format described in the README.md.
  9 | # <output file>.types: Python dictionary that maps string diagnosis codes to integer diagnosis codes.
 10 | # <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.
 11 | # <output file>.3digitICD9.types: Python dictionary that maps 3-digit string diagnosis codes to integer diagnosis codes.
 12 | 
 13 | import sys
 14 | import cPickle as pickle
 15 | from datetime import datetime
 16 | 
 17 | def convert_to_icd9(dxStr):
 18 | 	if dxStr.startswith('E'):
 19 | 		if len(dxStr) > 4: return dxStr[:4] + '.' + dxStr[4:]
 20 | 		else: return dxStr
 21 | 	else:
 22 | 		if len(dxStr) > 3: return dxStr[:3] + '.' + dxStr[3:]
 23 | 		else: return dxStr
 24 | 	
 25 | def convert_to_3digit_icd9(dxStr):
 26 | 	if dxStr.startswith('E'):
 27 | 		if len(dxStr) > 4: return dxStr[:4]
 28 | 		else: return dxStr
 29 | 	else:
 30 | 		if len(dxStr) > 3: return dxStr[:3]
 31 | 		else: return dxStr
 32 | 
 33 | if __name__ == '__main__':
 34 | 	admissionFile = sys.argv[1]
 35 | 	diagnosisFile = sys.argv[2]
 36 | 	outFile = sys.argv[3]
 37 | 
 38 | 	print 'Building pid-admission mapping, admission-date mapping'
 39 | 	pidAdmMap = {}
 40 | 	admDateMap = {}
 41 | 	infd = open(admissionFile, 'r')
 42 | 	infd.readline()
 43 | 	for line in infd:
 44 | 		tokens = line.strip().split(',')
 45 | 		pid = int(tokens[1])
 46 | 		admId = int(tokens[2])
 47 | 		admTime = datetime.strptime(tokens[3], '%Y-%m-%d %H:%M:%S')
 48 | 		admDateMap[admId] = admTime
 49 | 		if pid in pidAdmMap: pidAdmMap[pid].append(admId)
 50 | 		else: pidAdmMap[pid] = [admId]
 51 | 	infd.close()
 52 | 
 53 | 	print 'Building admission-dxList mapping'
 54 | 	admDxMap = {}
 55 | 	admDxMap_3digit = {}
 56 | 	infd = open(diagnosisFile, 'r')
 57 | 	infd.readline()
 58 | 	for line in infd:
 59 | 		tokens = line.strip().split(',')
 60 | 		admId = int(tokens[2])
 61 | 		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.
 62 | 		dxStr_3digit = 'D_' + convert_to_3digit_icd9(tokens[4][1:-1])
 63 | 
 64 | 		if admId in admDxMap: 
 65 | 			admDxMap[admId].append(dxStr)
 66 | 		else: 
 67 | 			admDxMap[admId] = [dxStr]
 68 | 
 69 | 		if admId in admDxMap_3digit: 
 70 | 			admDxMap_3digit[admId].append(dxStr_3digit)
 71 | 		else: 
 72 | 			admDxMap_3digit[admId] = [dxStr_3digit]
 73 | 	infd.close()
 74 | 
 75 | 	print 'Building pid-sortedVisits mapping'
 76 | 	pidSeqMap = {}
 77 | 	pidSeqMap_3digit = {}
 78 | 	for pid, admIdList in pidAdmMap.iteritems():
 79 | 		if len(admIdList) < 2: continue
 80 | 
 81 | 		sortedList = sorted([(admDateMap[admId], admDxMap[admId]) for admId in admIdList])
 82 | 		pidSeqMap[pid] = sortedList
 83 | 
 84 | 		sortedList_3digit = sorted([(admDateMap[admId], admDxMap_3digit[admId]) for admId in admIdList])
 85 | 		pidSeqMap_3digit[pid] = sortedList_3digit
 86 | 	
 87 | 	print 'Building pids, dates, strSeqs'
 88 | 	pids = []
 89 | 	dates = []
 90 | 	seqs = []
 91 | 	for pid, visits in pidSeqMap.iteritems():
 92 | 		pids.append(pid)
 93 | 		seq = []
 94 | 		date = []
 95 | 		for visit in visits:
 96 | 			date.append(visit[0])
 97 | 			seq.append(visit[1])
 98 | 		dates.append(date)
 99 | 		seqs.append(seq)
100 | 	
101 | 	print 'Building pids, dates, strSeqs for 3digit ICD9 code'
102 | 	seqs_3digit = []
103 | 	for pid, visits in pidSeqMap_3digit.iteritems():
104 | 		seq = []
105 | 		for visit in visits:
106 | 			seq.append(visit[1])
107 | 		seqs_3digit.append(seq)
108 | 	
109 | 	print 'Converting strSeqs to intSeqs, and making types'
110 | 	types = {}
111 | 	newSeqs = []
112 | 	for patient in seqs:
113 | 		newPatient = []
114 | 		for visit in patient:
115 | 			newVisit = []
116 | 			for code in visit:
117 | 				if code in types:
118 | 					newVisit.append(types[code])
119 | 				else:
120 | 					types[code] = len(types)
121 | 					newVisit.append(types[code])
122 | 			newPatient.append(newVisit)
123 | 		newSeqs.append(newPatient)
124 | 	
125 | 	print 'Converting strSeqs to intSeqs, and making types for 3digit ICD9 code'
126 | 	types_3digit = {}
127 | 	newSeqs_3digit = []
128 | 	for patient in seqs_3digit:
129 | 		newPatient = []
130 | 		for visit in patient:
131 | 			newVisit = []
132 | 			for code in set(visit):
133 | 				if code in types_3digit:
134 | 					newVisit.append(types_3digit[code])
135 | 				else:
136 | 					types_3digit[code] = len(types_3digit)
137 | 					newVisit.append(types_3digit[code])
138 | 			newPatient.append(newVisit)
139 | 		newSeqs_3digit.append(newPatient)
140 | 	
141 | 	print 'Re-formatting to Med2Vec dataset'
142 | 	seqs = []
143 | 	for patient in newSeqs:
144 | 		seqs.extend(patient)
145 | 		seqs.append([-1])
146 | 	seqs = seqs[:-1]
147 | 
148 | 	seqs_3digit = []
149 | 	for patient in newSeqs_3digit:
150 | 		seqs_3digit.extend(patient)
151 | 		seqs_3digit.append([-1])
152 | 	seqs_3digit = seqs_3digit[:-1]
153 | 
154 | 	pickle.dump(seqs, open(outFile+'.seqs', 'wb'), -1)
155 | 	pickle.dump(types, open(outFile+'.types', 'wb'), -1)
156 | 	pickle.dump(seqs_3digit, open(outFile+'.3digitICD9.seqs', 'wb'), -1)
157 | 	pickle.dump(types_3digit, open(outFile+'.3digitICD9.types', 'wb'), -1)
158 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | Med2Vec
 2 | =========================================
 3 | 
 4 | Med2Vec is a multi-layer representation learning tool for learning code representations and visit representations from EHR datasets.
 5 | 
 6 | [![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!")
 7 | 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). 
 8 | 
 9 | #### Relevant Publications
10 | 
11 | Med2Vec implements an algorithm introduced in the following [paper](http://www.kdd.org/kdd2016/subtopic/view/multi-layer-representation-learning-for-medical-concepts):
12 | 
13 |     Multi-layer Representation Learning for Medical Concepts
14 | 	Edward Choi, Mohammad Taha Bahadori, Elizabeth Searles, Catherine Coffey, 
15 | 	Michael Thompson, James Bost, Javier Tejedor-Sojo, Jimeng Sun
16 | 	KDD 2016, pp.1495-1504
17 | 
18 | #### Running Med2Vec
19 | 
20 | **STEP 1: Installation**  
21 | 
22 | 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)
23 | 
24 | 2. If you plan to use GPU computation, install [CUDA](https://developer.nvidia.com/cuda-downloads)
25 | 
26 | 3. Download/clone the Med2Vec code  
27 | 
28 | **STEP 2: Fast way to test Med2Vec with MIMIC-III**
29 | 
30 | This step describes how to run, with minimum number of steps, Med2Vec using MIMIC-III. 
31 | 
32 | 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. 
33 | 
34 | 1. You can use "process_mimic.py" to process MIMIC-III dataset and generate a suitable training dataset for Med2Vec.
35 | Place the script to the same location where the MIMIC-III CSV files are located, and run the script. 
36 | The execution command is `python process_mimic.py ADMISSIONS.csv DIAGNOSES_ICD.csv <output file>`.
37 | Instructions are described inside the script. 
38 | 
39 | 2. Run Med2Vec using the ".seqs" file generated by process_mimic.py, using the following command.
40 | `python med2vec.py <seqs file> 4894 <output path>`
41 | where 4894 is the number of unique ICD9 diagnosis codes in the dataset.
42 | 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.
43 | `python med2vec.py <seqs file> 4894 <output path> --label_file <3digitICD9.seqs file> --n_output_codes 942`
44 | where 942 is the number of unique 3-digit ICD9 diagnosis codes in the dataset.
45 | 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.
46 | `python med2vec.py <3digitICD9.seqs file> 942 <output path>`
47 | 
48 | 3. As suggested in STEP 4, you might want to adjust the hyper-parameters. 
49 | 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. 
50 | There are only 7,500 patients who made more than a single visit, and most of them have only two visits.
51 | 
52 | **STEP 3: Preparing training data**  
53 | 
54 | 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). 
55 | First, medical codes need to be converted to an integer. Then a single visit can be converted as a list of integers. 
56 | For example, [5,8,15] means the patient was assigned with code 5, 8, and 15 at a certain visit. 
57 | 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]]. 
58 | If there are multiple patients, each patient must be delimited by a list [-1]. 
59 | 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. 
60 | This list of list needs to be pickled using cPickle. We will refer to this file as the "visit file".
61 | 
62 | 2. The total number of unique medical codes is required to run Med2Vec. 
63 | For example, if the dataset is using 14,000 diagnosis codes and 11,000 procedure codes, the total number is 25,000. 
64 | 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))
65 | 
66 | 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. 
67 | 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. 
68 | You will still be able to learn the code representations for the original un-grouped codes. 
69 | The grouped dataset is used only for speeding up the training speed. (Refer to section 4.4 of the paper) 
70 | 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".
71 | 
72 | 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.
73 | 
74 | 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. 
75 | 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]. 
76 | Similar to the [-1] vector in step 1, each patient is delimited with an all-zero vector. 
77 | 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. 
78 | We will refer to this file as the "demo file".
79 | 
80 | 6. Similar to step 2, you will need to remeber the size of the demographics vector if you plan to use the demo file. 
81 | In the example of step 5, the size of the demographics vector is 7.
82 | 
83 | **STEP 4: Running Med2Vec**  
84 | 
85 | 1. The minimum input you need to run Med2Vec is the visit file, the number of unique medical codes and the output path
86 | `python med2vec <path/to/visit_file> <the number of unique medical codes> <path/to/output>`  
87 | 
88 | 2. Specifying `--verbose` option will print training process after each 10 mini-batches.
89 | 
90 | 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`
91 | 
92 | **STEP 5: Looking at your results**  
93 | 
94 | 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).
95 | 
96 | The 2D scatterplot of the learned code representations would look similar to [this](http://mp2893.com/scatterplot/nnsg_h200e49_category10.html).
97 | (This is the scatterplot of the code representations trained with Non-negative Skip-gram, which is essentially Med2Vec minus the visit-level training)
98 | 


--------------------------------------------------------------------------------
/med2vec.py:
--------------------------------------------------------------------------------
  1 | #################################################################
  2 | # Code written by Edward Choi (mp2893@gatech.edu)
  3 | # For bug report, please contact author using the email address
  4 | #################################################################
  5 | 
  6 | import sys, random
  7 | import numpy as np
  8 | import cPickle as pickle
  9 | from collections import OrderedDict
 10 | import argparse
 11 | 
 12 | import theano
 13 | import theano.tensor as T
 14 | from theano import config
 15 | 
 16 | def numpy_floatX(data):
 17 | 	return np.asarray(data, dtype=config.floatX)
 18 | 
 19 | def unzip(zipped):
 20 | 	new_params = OrderedDict()
 21 | 	for k, v in zipped.iteritems():
 22 | 		new_params[k] = v.get_value()
 23 | 	return new_params
 24 | 
 25 | def init_params(options):
 26 | 	params = OrderedDict()
 27 | 
 28 | 	numXcodes = options['numXcodes']
 29 | 	numYcodes = options['numYcodes']
 30 | 	embDimSize= options['embDimSize']
 31 | 	demoSize = options['demoSize']
 32 | 	hiddenDimSize = options['hiddenDimSize']
 33 | 
 34 | 	params['W_emb'] = np.random.uniform(-0.01, 0.01, (numXcodes, embDimSize)).astype(config.floatX) #emb matrix needs an extra dimension for the time
 35 | 	params['b_emb'] = np.zeros(embDimSize).astype(config.floatX)
 36 | 	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
 37 | 	params['b_hidden'] = np.zeros(hiddenDimSize).astype(config.floatX)
 38 | 	if numYcodes > 0:
 39 | 		params['W_output'] = np.random.uniform(-0.01, 0.01, (hiddenDimSize, numYcodes)).astype(config.floatX) #emb matrix needs an extra dimension for the time
 40 | 		params['b_output'] = np.zeros(numYcodes).astype(config.floatX)
 41 | 	else:
 42 | 		params['W_output'] = np.random.uniform(-0.01, 0.01, (hiddenDimSize, numXcodes)).astype(config.floatX) #emb matrix needs an extra dimension for the time
 43 | 		params['b_output'] = np.zeros(numXcodes).astype(config.floatX)
 44 | 
 45 | 	return params
 46 | 
 47 | def load_params(options):
 48 | 	params = np.load(options['modelFile'])
 49 | 	return params
 50 | 
 51 | def init_tparams(params):
 52 | 	tparams = OrderedDict()
 53 | 	for k, v in params.iteritems():
 54 | 		tparams[k] = theano.shared(v, name=k)
 55 | 	return tparams
 56 | 
 57 | def build_model(tparams, options):
 58 | 	x = T.matrix('x', dtype=config.floatX)
 59 | 	d = T.matrix('d', dtype=config.floatX)
 60 | 	y = T.matrix('y', dtype=config.floatX)
 61 | 	mask = T.vector('mask', dtype=config.floatX)
 62 | 
 63 | 	logEps = options['logEps']
 64 | 
 65 | 	emb = T.maximum(T.dot(x, tparams['W_emb']) + tparams['b_emb'],0)
 66 | 	if options['demoSize'] > 0: emb = T.concatenate((emb, d), axis=1)
 67 | 	visit = T.maximum(T.dot(emb, tparams['W_hidden']) + tparams['b_hidden'],0)
 68 | 	results = T.nnet.softmax(T.dot(visit, tparams['W_output']) + tparams['b_output'])
 69 | 	
 70 | 	mask1 = (mask[:-1] * mask[1:])[:,None]
 71 | 	mask2 = (mask[:-2] * mask[1:-1] * mask[2:])[:,None]
 72 | 	mask3 = (mask[:-3] * mask[1:-2] * mask[2:-1] * mask[3:])[:,None]
 73 | 	mask4 = (mask[:-4] * mask[1:-3] * mask[2:-2] * mask[3:-1] * mask[4:])[:,None]
 74 | 	mask5 = (mask[:-5] * mask[1:-4] * mask[2:-3] * mask[3:-2] * mask[4:-1] * mask[5:])[:,None]
 75 | 
 76 | 	t = None
 77 | 	if options['numYcodes'] > 0: t = y
 78 | 	else: t = x
 79 | 
 80 | 	forward_results =  results[:-1] * mask1
 81 | 	forward_cross_entropy = -(t[1:] * T.log(forward_results + logEps) + (1. - t[1:]) * T.log(1. - forward_results + logEps))
 82 | 
 83 | 	forward_results2 =  results[:-2] * mask2
 84 | 	forward_cross_entropy2 = -(t[2:] * T.log(forward_results2 + logEps) + (1. - t[2:]) * T.log(1. - forward_results2 + logEps))
 85 | 
 86 | 	forward_results3 =  results[:-3] * mask3
 87 | 	forward_cross_entropy3 = -(t[3:] * T.log(forward_results3 + logEps) + (1. - t[3:]) * T.log(1. - forward_results3 + logEps))
 88 | 
 89 | 	forward_results4 =  results[:-4] * mask4
 90 | 	forward_cross_entropy4 = -(t[4:] * T.log(forward_results4 + logEps) + (1. - t[4:]) * T.log(1. - forward_results4 + logEps))
 91 | 
 92 | 	forward_results5 =  results[:-5] * mask5
 93 | 	forward_cross_entropy5 = -(t[5:] * T.log(forward_results5 + logEps) + (1. - t[5:]) * T.log(1. - forward_results5 + logEps))
 94 | 
 95 | 	backward_results =  results[1:] * mask1
 96 | 	backward_cross_entropy = -(t[:-1] * T.log(backward_results + logEps) + (1. - t[:-1]) * T.log(1. - backward_results + logEps))
 97 | 
 98 | 	backward_results2 =  results[2:] * mask2
 99 | 	backward_cross_entropy2 = -(t[:-2] * T.log(backward_results2 + logEps) + (1. - t[:-2]) * T.log(1. - backward_results2 + logEps))
100 | 
101 | 	backward_results3 =  results[3:] * mask3
102 | 	backward_cross_entropy3 = -(t[:-3] * T.log(backward_results3 + logEps) + (1. - t[:-3]) * T.log(1. - backward_results3 + logEps))
103 | 
104 | 	backward_results4 =  results[4:] * mask4
105 | 	backward_cross_entropy4 = -(t[:-4] * T.log(backward_results4 + logEps) + (1. - t[:-4]) * T.log(1. - backward_results4 + logEps))
106 | 
107 | 	backward_results5 =  results[5:] * mask5
108 | 	backward_cross_entropy5 = -(t[:-5] * T.log(backward_results5 + logEps) + (1. - t[:-5]) * T.log(1. - backward_results5 + logEps))
109 | 
110 | 	visit_cost1 = (forward_cross_entropy.sum(axis=1).sum(axis=0) + backward_cross_entropy.sum(axis=1).sum(axis=0)) / (mask1.sum() + logEps)
111 | 	visit_cost2 = (forward_cross_entropy2.sum(axis=1).sum(axis=0) + backward_cross_entropy2.sum(axis=1).sum(axis=0)) / (mask2.sum() + logEps)
112 | 	visit_cost3 = (forward_cross_entropy3.sum(axis=1).sum(axis=0) + backward_cross_entropy3.sum(axis=1).sum(axis=0)) / (mask3.sum() + logEps)
113 | 	visit_cost4 = (forward_cross_entropy4.sum(axis=1).sum(axis=0) + backward_cross_entropy4.sum(axis=1).sum(axis=0)) / (mask4.sum() + logEps)
114 | 	visit_cost5 = (forward_cross_entropy5.sum(axis=1).sum(axis=0) + backward_cross_entropy5.sum(axis=1).sum(axis=0)) / (mask5.sum() + logEps)
115 | 
116 | 	windowSize = options['windowSize']
117 | 	visit_cost = visit_cost1
118 | 	if windowSize == 2:
119 | 		visit_cost = visit_cost1 + visit_cost2
120 | 	elif windowSize == 3:
121 | 		visit_cost = visit_cost1 + visit_cost2 + visit_cost3
122 | 	elif windowSize == 4:
123 | 		visit_cost = visit_cost1 + visit_cost2 + visit_cost3 + visit_cost4
124 | 	elif windowSize == 5:
125 | 		visit_cost = visit_cost1 + visit_cost2 + visit_cost3 + visit_cost4 + visit_cost5
126 | 
127 | 	iVector = T.vector('iVector', dtype='int32')
128 | 	jVector = T.vector('jVector', dtype='int32')
129 | 	preVec = T.maximum(tparams['W_emb'],0)
130 | 	norms = (T.exp(T.dot(preVec, preVec.T))).sum(axis=1)
131 | 	emb_cost = -T.log((T.exp((preVec[iVector] * preVec[jVector]).sum(axis=1)) / norms[iVector]) + logEps)
132 | 
133 | 	total_cost = visit_cost + T.mean(emb_cost) + options['L2_reg'] * (tparams['W_emb'] ** 2).sum()
134 | 
135 | 	if options['demoSize'] > 0 and options['numYcodes'] > 0: return x, d, y, mask, iVector, jVector, total_cost
136 | 	elif options['demoSize'] == 0 and options['numYcodes'] > 0: return x, y, mask, iVector, jVector, total_cost
137 | 	elif options['demoSize'] > 0 and options['numYcodes'] == 0: return x, d, mask, iVector, jVector, total_cost
138 | 	else: return x, mask, iVector, jVector, total_cost
139 | 
140 | def adadelta(tparams, grads, x, mask, iVector, jVector, cost, options, d=None, y=None):
141 | 	zipped_grads = [theano.shared(p.get_value() * numpy_floatX(0.), name='%s_grad' % k) for k, p in tparams.iteritems()]
142 | 	running_up2 = [theano.shared(p.get_value() * numpy_floatX(0.), name='%s_rup2' % k) for k, p in tparams.iteritems()]
143 | 	running_grads2 = [theano.shared(p.get_value() * numpy_floatX(0.), name='%s_rgrad2' % k) for k, p in tparams.iteritems()]
144 | 
145 | 	zgup = [(zg, g) for zg, g in zip(zipped_grads, grads)]
146 | 	rg2up = [(rg2, 0.95 * rg2 + 0.05 * (g ** 2)) for rg2, g in zip(running_grads2, grads)]
147 | 
148 | 	if options['demoSize'] > 0 and options['numYcodes'] > 0:
149 | 		f_grad_shared = theano.function([x, d, y, mask, iVector, jVector], cost, updates=zgup + rg2up, name='adadelta_f_grad_shared')
150 | 	elif options['demoSize'] == 0 and options['numYcodes'] > 0:
151 | 		f_grad_shared = theano.function([x, y, mask, iVector, jVector], cost, updates=zgup + rg2up, name='adadelta_f_grad_shared')
152 | 	elif options['demoSize'] > 0 and options['numYcodes'] == 0:
153 | 		f_grad_shared = theano.function([x, d, mask, iVector, jVector], cost, updates=zgup + rg2up, name='adadelta_f_grad_shared')
154 | 	else:
155 | 		f_grad_shared = theano.function([x, mask, iVector, jVector], cost, updates=zgup + rg2up, name='adadelta_f_grad_shared')
156 | 
157 | 	updir = [-T.sqrt(ru2 + 1e-6) / T.sqrt(rg2 + 1e-6) * zg for zg, ru2, rg2 in zip(zipped_grads, running_up2, running_grads2)]
158 | 	ru2up = [(ru2, 0.95 * ru2 + 0.05 * (ud ** 2)) for ru2, ud in zip(running_up2, updir)]
159 | 	param_up = [(p, p + ud) for p, ud in zip(tparams.values(), updir)]
160 | 
161 | 	f_update = theano.function([], [], updates=ru2up + param_up, on_unused_input='ignore', name='adadelta_f_update')
162 | 
163 | 	return f_grad_shared, f_update
164 | 
165 | def load_data(xFile, dFile, yFile):
166 | 	seqX = np.array(pickle.load(open(xFile, 'rb')))
167 | 	seqD = []
168 | 	if len(dFile) > 0: seqD = np.asarray(pickle.load(open(dFile, 'rb')), dtype=config.floatX)
169 | 	seqY = []
170 | 	if len(yFile) > 0: seqY = np.array(pickle.load(open(yFile, 'rb')))
171 | 	return seqX, seqD, seqY
172 | 
173 | def pickTwo(codes, iVector, jVector):
174 | 	for first in codes:
175 | 		for second in codes:
176 | 			if first == second: continue
177 | 			iVector.append(first)
178 | 			jVector.append(second)
179 | 	
180 | def padMatrix(seqs, labels, options):
181 | 	n_samples = len(seqs)
182 | 	iVector = []
183 | 	jVector = []
184 | 	numXcodes = options['numXcodes']
185 | 	numYcodes = options['numYcodes']
186 | 
187 | 	if numYcodes > 0:
188 | 		x = np.zeros((n_samples, numXcodes)).astype(config.floatX)
189 | 		y = np.zeros((n_samples, numYcodes)).astype(config.floatX)
190 | 		mask = np.zeros((n_samples,)).astype(config.floatX)
191 | 		for idx, (seq, label) in enumerate(zip(seqs, labels)):
192 | 			if not seq[0] == -1:
193 | 				x[idx][seq] = 1.
194 | 				y[idx][label] = 1.
195 | 				pickTwo(seq, iVector, jVector)
196 | 				mask[idx] = 1.
197 | 		return x, y, mask, iVector, jVector
198 | 	else:
199 | 		x = np.zeros((n_samples, numXcodes)).astype(config.floatX)
200 | 		mask = np.zeros((n_samples,)).astype(config.floatX)
201 | 		for idx, seq in enumerate(seqs):
202 | 			if not seq[0] == -1:
203 | 				x[idx][seq] = 1.
204 | 				pickTwo(seq, iVector, jVector)
205 | 				mask[idx] = 1.
206 | 		return x, mask, iVector, jVector
207 | 
208 | def train_med2vec(seqFile='seqFile.txt', 
209 | 				demoFile='demoFile.txt',
210 | 				labelFile='labelFile.txt',
211 | 				outFile='outFile.txt',
212 | 				modelFile='modelFile.txt',
213 | 				L2_reg=0.001,
214 | 				numXcodes=20000, 
215 | 				numYcodes=20000, 
216 | 				embDimSize=1000,
217 | 				hiddenDimSize=2000,
218 | 				batchSize=100,
219 | 				demoSize=2,
220 | 				logEps=1e-8,
221 | 				windowSize=1,
222 | 				verbose=False,
223 | 				maxEpochs=1000):
224 | 
225 | 	options = locals().copy()
226 | 	print 'initializing parameters'
227 | 	params = init_params(options)
228 | 	#params = load_params(options)
229 | 	tparams = init_tparams(params)
230 | 
231 | 	print 'building models'
232 | 	f_grad_shared = None
233 | 	f_update = None
234 | 	if demoSize > 0 and numYcodes > 0:
235 | 		x, d, y, mask, iVector, jVector, cost = build_model(tparams, options)
236 | 		grads = T.grad(cost, wrt=tparams.values())
237 | 		f_grad_shared, f_update = adadelta(tparams, grads, x, mask, iVector, jVector, cost, options, d=d, y=y)
238 | 	elif demoSize == 0 and numYcodes > 0:
239 | 		x, y, mask, iVector, jVector, cost = build_model(tparams, options)
240 | 		grads = T.grad(cost, wrt=tparams.values())
241 | 		f_grad_shared, f_update = adadelta(tparams, grads, x, mask, iVector, jVector, cost, options, y=y)
242 | 	elif demoSize > 0 and numYcodes == 0:
243 | 		x, d, mask, iVector, jVector, cost = build_model(tparams, options)
244 | 		grads = T.grad(cost, wrt=tparams.values())
245 | 		f_grad_shared, f_update = adadelta(tparams, grads, x, mask, iVector, jVector, cost, options, d=d)
246 | 	else:
247 | 		x, mask, iVector, jVector, cost = build_model(tparams, options)
248 | 		grads = T.grad(cost, wrt=tparams.values())
249 | 		f_grad_shared, f_update = adadelta(tparams, grads, x, mask, iVector, jVector, cost, options)
250 | 
251 | 	print 'loading data'
252 | 	seqs, demos, labels = load_data(seqFile, demoFile, labelFile)
253 | 	n_batches = int(np.ceil(float(len(seqs)) / float(batchSize)))
254 | 
255 | 	print 'training start'
256 | 	for epoch in xrange(maxEpochs):
257 | 		iteration = 0
258 | 		costVector = []
259 | 		for index in random.sample(range(n_batches), n_batches):
260 | 			batchX = seqs[batchSize*index:batchSize*(index+1)]
261 | 			batchY = []
262 | 			batchD = []
263 | 			if demoSize > 0 and numYcodes > 0:
264 | 				batchY = labels[batchSize*index:batchSize*(index+1)]
265 | 				x, y, mask, iVector, jVector = padMatrix(batchX, batchY, options)
266 | 				batchD = demos[batchSize*index:batchSize*(index+1)]
267 | 				cost = f_grad_shared(x, batchD, y, mask, iVector, jVector)
268 | 			elif demoSize == 0 and numYcodes > 0:
269 | 				batchY = labels[batchSize*index:batchSize*(index+1)]
270 | 				x, y, mask, iVector, jVector = padMatrix(batchX, batchY, options)
271 | 				cost = f_grad_shared(x, y, mask, iVector, jVector)
272 | 			elif demoSize > 0 and numYcodes == 0:
273 | 				x, mask, iVector, jVector = padMatrix(batchX, batchY, options)
274 | 				batchD = demos[batchSize*index:batchSize*(index+1)]
275 | 				cost = f_grad_shared(x, batchD, mask, iVector, jVector)
276 | 			else:
277 | 				x, mask, iVector, jVector = padMatrix(batchX, batchY, options)
278 | 				cost = f_grad_shared(x, mask, iVector, jVector)
279 | 			costVector.append(cost)
280 | 			f_update()
281 | 			if (iteration % 10 == 0) and verbose: print 'epoch:%d, iteration:%d/%d, cost:%f' % (epoch, iteration, n_batches, cost)
282 | 			iteration += 1
283 | 		print 'epoch:%d, mean_cost:%f' % (epoch, np.mean(costVector))
284 | 		tempParams = unzip(tparams)
285 | 		np.savez_compressed(outFile + '.' + str(epoch), **tempParams)
286 | 
287 | def parse_arguments(parser):
288 | 	parser.add_argument('seq_file', type=str, metavar='<visit_file>', help='The path to the Pickled file containing visit information of patients')
289 | 	parser.add_argument('n_input_codes', type=int, metavar='<n_input_codes>', help='The number of unique input medical codes')
290 | 	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')
291 | 	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')
292 | 	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')
293 | 	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')
294 | 	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')
295 | 	parser.add_argument('--cr_size', type=int, default=200, help='The size of the code representation (default value: 200)')
296 | 	parser.add_argument('--vr_size', type=int, default=200, help='The size of the visit representation (default value: 200)')
297 | 	parser.add_argument('--batch_size', type=int, default=1000, help='The size of a single mini-batch (default value: 1000)')
298 | 	parser.add_argument('--n_epoch', type=int, default=10, help='The number of training epochs (default value: 10)')
299 | 	parser.add_argument('--L2_reg', type=float, default=0.001, help='L2 regularization for the code representation matrix W_c (default value: 0.001)')
300 | 	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)')
301 | 	parser.add_argument('--log_eps', type=float, default=1e-8, help='A small value to prevent log(0) (default value: 1e-8)')
302 | 	parser.add_argument('--verbose', action='store_true', help='Print output after every 10 mini-batches')
303 | 	args = parser.parse_args()
304 | 	return args
305 | 
306 | if __name__ == '__main__':
307 | 	parser = argparse.ArgumentParser()
308 | 	args = parse_arguments(parser)
309 | 
310 | 	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)
311 | 


--------------------------------------------------------------------------------