In [0]:
import sys
import _pickle as pickle
import numpy as np
from datetime import datetime

In [0]:
from google.colab import drive
drive.mount('/content/gdrive', force_remount=True)


Mounted at /content/gdrive

In [0]:
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

In [0]:
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

In [0]:
# input arguments
binary_count = 'binary'

In [0]:
root_dir = "/content/gdrive/My Drive/"

if binary_count == 'count':
  base_dir = root_dir + 'GOSH/Synthetic Data/medgan/count/'
else:
  base_dir = root_dir + 'GOSH/Synthetic Data/medgan/binary/'

raw_data_dir = root_dir + 'GOSH/Synthetic Data/medgan/mimic/'
processed_data_dir = base_dir + 'processed_mimic/'
model_dir = base_dir + 'models/'
gen_data_dir = base_dir + 'generated_data/'

In [0]:
admissionFile = raw_data_dir + 'ADMISSIONS.csv'
diagnosisFile = raw_data_dir + 'DIAGNOSES_ICD.csv'
outFile = processed_data_dir + 'processed_mimic'

In [0]:
if binary_count != 'binary' and binary_count != 'count':
    print('You must choose either binary or count.')

In [0]:
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 = {}
infd = open(diagnosisFile, 'r')
infd.readline()
for line in infd:
    tokens = line.strip().split(',')
    admId = int(tokens[2])
    # Uncomment this line and comment the line below, if you want to use the entire ICD9 digits.
    dxStr = 'D_' + convert_to_icd9(tokens[4][1:-1])
    #dxStr = 'D_' + convert_to_3digit_icd9(tokens[4][1:-1])
    if admId in admDxMap: admDxMap[admId].append(dxStr)
    else: admDxMap[admId] = [dxStr]
infd.close()

print('Building pid-sortedVisits mapping')
pidSeqMap = {}
for pid, admIdList in pidAdmMap.items():
    #if len(admIdList) < 2: continue
    sortedList = sorted([(admDateMap[admId], admDxMap[admId]) for admId in admIdList])
    pidSeqMap[pid] = sortedList

print('Building pids, dates, strSeqs')
pids = []
dates = []
seqs = []
for pid, visits in pidSeqMap.items():
    pids.append(pid)
    seq = []
    date = []
    for visit in visits:
        date.append(visit[0])
        seq.append(visit[1])
    dates.append(date)
    seqs.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('Constructing the matrix')
numPatients = len(newSeqs)
numCodes = len(types)
matrix = np.zeros((numPatients, numCodes)).astype('float32')
for i, patient in enumerate(newSeqs):
    for visit in patient:
        for code in visit:
            if binary_count == 'binary':
                matrix[i][code] = 1.
            else:
                matrix[i][code] += 1.

pickle.dump(pids, open(outFile+'.pids', 'wb'), -1)
pickle.dump(matrix, open(outFile+'.matrix', 'wb'), -1)
pickle.dump(types, open(outFile+'.types', 'wb'), -1)


Building pid-admission mapping, admission-date mapping
Building admission-dxList mapping
Building pid-sortedVisits mapping
Building pids, dates, strSeqs
Converting strSeqs to intSeqs, and making types
Constructing the matrix

In [0]: