<commit_before>import csv
import numpy as np
import os
import pandas as pd
import sys

from pandas import DataFrame


def read_patients_table(mimic3_path):
    pats = DataFrame.from_csv(os.path.join(mimic3_path, 'PATIENTS_DATA_TABLE.csv'))
    pats = pats[['SUBJECT_ID', 'GENDER', 'DOB', 'DOD']]
    pats.DOB = pd.to_datetime(pats.DOB)
    pats.DOD = pd.to_datetime(pats.DOD)
    return pats

def read_admissions_table(mimic3_path):
    admits = DataFrame.from_csv(os.path.join(mimic3_path, 'ADMISSIONS_DATA_TABLE.csv'))
    admits = admits[['SUBJECT_ID', 'HADM_ID', 'DEATHTIME', 'ETHNICITY', 'DIAGNOSIS']]
    admits.DEATHTIME = pd.to_datetime(admits.DEATHTIME)
    return admits

def read_icustays_table(mimic3_path):
    stays = DataFrame.from_csv(os.path.join(mimic3_path, 'ICUSTAYS_DATA_TABLE.csv'))
    stays.INTIME = pd.to_datetime(stays.INTIME)
    stays.OUTTIME = pd.to_datetime(stays.OUTTIME)
    return stays

def read_icd_diagnoses_table(mimic3_path):
    codes = DataFrame.from_csv(os.path.join(mimic3_path, 'D_ICD_DIAGNOSES_DATA_TABLE.csv'))
    codes = codes[['ICD9_CODE','SHORT_TITLE','LONG_TITLE']]
    diagnoses = DataFrame.from_csv(os.path.join(mimic3_path, 'DIAGNOSES_ICD_DATA_TABLE.csv'))
    diagnoses = diagnoses.merge(codes, how='inner', left_on='ICD9_CODE', right_on='ICD9_CODE')
    diagnoses[['SUBJECT_ID','HADM_ID','SEQ_NUM']] = diagnoses[['SUBJECT_ID','HADM_ID','SEQ_NUM']].astype(int)
    return diagnoses

def read_events_table_by_row(mimic3_path, table):
    nb_rows = { 'chartevents': 263201376, 'labevents': 27872576, 'outputevents': 4349340 }
    reader = csv.DictReader(open(os.path.join(mimic3_path, table.upper() + '_DATA_TABLE.csv'), 'r'))
    for i,row in enumerate(reader):
        if 'ICUSTAY_ID' not in row:
            row['ICUSTAY_ID'] = ''
        yield row, i, nb_rows[table.lower()]

def count_icd_codes(diagnoses, output_path=None):
    codes = diagnoses[['ICD9_CODE','SHORT_TITLE','LONG_TITLE']].drop_duplicates().set_index('ICD9_CODE')
    codes['COUNT'] = diagnoses.groupby('ICD9_CODE')['ICUSTAY_ID'].count()
    codes.COUNT = codes.COUNT.fillna(0).astype(int)
    codes = codes.ix[codes.COUNT>0]
    if output_path:
        codes.to_csv(output_path, index_label='ICD9_CODE')
    return codes.sort_values('COUNT', ascending=False).reset_index()

def remove_icustays_with_transfers(stays):
    stays = stays.ix[(stays.FIRST_WARDID == stays.LAST_WARDID) & (stays.FIRST_CAREUNIT == stays.LAST_CAREUNIT)]
    return stays[['SUBJECT_ID', 'HADM_ID', 'ICUSTAY_ID', 'DBSOURCE', 'INTIME', 'OUTTIME', 'LOS']]

def merge_on_subject(table1, table2):
    return table1.merge(table2, how='inner', left_on=['SUBJECT_ID'], right_on=['SUBJECT_ID'])

def merge_on_subject_admission(table1, table2):
    return table1.merge(table2, how='inner', left_on=['SUBJECT_ID', 'HADM_ID'], right_on=['SUBJECT_ID', 'HADM_ID'])

def add_age_to_icustays(stays):
    stays['AGE'] = (stays.INTIME - stays.DOB).apply(lambda s: s / np.timedelta64(1, 's')) / 60./60/24/365
    stays.AGE.ix[stays.AGE<0] = 90
    return stays

def add_mortality_to_icustays(stays):
    mortality = stays.DOD.notnull() & ((stays.INTIME <= stays.DOD) & (stays.OUTTIME >= stays.DOD))
    mortality = mortality | (stays.DEATHTIME.notnull() & ((stays.INTIME <= stays.DEATHTIME) & (stays.OUTTIME >= stays.DEATHTIME)))
    stays['MORTALITY'] = mortality.astype(int)
    return stays

def filter_admissions_on_nb_icustays(stays, min_nb_stays=1, max_nb_stays=1):
    to_keep = stays.groupby('HADM_ID').count()[['ICUSTAY_ID']].reset_index()
    to_keep = to_keep.ix[(to_keep.ICUSTAY_ID>=min_nb_stays)&(to_keep.ICUSTAY_ID<=max_nb_stays)][['HADM_ID']]
    stays = stays.merge(to_keep, how='inner', left_on='HADM_ID', right_on='HADM_ID')
    return stays

def filter_icustays_on_age(stays, min_age=18, max_age=np.inf):
    stays = stays.ix[(stays.AGE>=min_age)&(stays.AGE<=max_age)]
    return stays

def filter_diagnoses_on_stays(diagnoses, stays):
    return diagnoses.merge(stays[['SUBJECT_ID', 'HADM_ID', 'ICUSTAY_ID']].drop_duplicates(), how='inner',
                           left_on=['SUBJECT_ID', 'HADM_ID'], right_on=['SUBJECT_ID', 'HADM_ID'])

def break_up_stays_by_subject(stays, output_path, subjects=None, verbose=1):
    subjects = stays.SUBJECT_ID.unique() if subjects is None else subjects
    nb_subjects = subjects.shape[0]
    for i, subject_id in enumerate(subjects):
        if verbose:
            sys.stdout.write('\rSUBJECT {0} of {1}...'.format(i+1, nb_subjects))
        dn = os.path.join(output_path, str(subject_id))
        try:
            os.makedirs(dn)
        except:
            pass

        stays.ix[stays.SUBJECT_ID == subject_id].sort_values(by='INTIME').to_csv(os.path.join(dn, 'stays.csv'), index=False)
    if verbose:
        sys.stdout.write('DONE!\n')

def break_up_diagnoses_by_subject(diagnoses, output_path, subjects=None, verbose=1):
    subjects = stays.SUBJECT_ID.unique() if subjects is None else subjects
    nb_subjects = subjects.shape[0]
    for i, subject_id in enumerate(subjects):
        if verbose:
            sys.stdout.write('\rSUBJECT {0} of {1}...'.format(i+1, nb_subjects))
        dn = os.path.join(output_path, str(subject_id))
        try:
            os.makedirs(dn)
        except:
            pass

        diagnoses.ix[diagnoses.SUBJECT_ID == subject_id].sort_values(by=['ICUSTAY_ID','SEQ_NUM']).to_csv(os.path.join(dn, 'diagnoses.csv'), index=False)
    if verbose:
        sys.stdout.write('DONE!\n')

def read_events_table_and_break_up_by_subject(mimic3_path, table, output_path, items_to_keep=None, subjects_to_keep=None, verbose=1):
    obs_header = [ 'SUBJECT_ID', 'HADM_ID', 'ICUSTAY_ID', 'CHARTTIME', 'ITEMID', 'VALUE', 'VALUEUOM' ]
    if items_to_keep is not None:
        items_to_keep = set([ str(s) for s in items_to_keep ])
    if subjects_to_keep is not None:
        subjects_to_keep = set([ str(s) for s in subjects_to_keep ])

    class nonlocal: pass
    nonlocal.curr_subject_id = ''
    nonlocal.last_write_no = 0
    nonlocal.last_write_nb_rows = 0
    nonlocal.last_write_subject_id = ''
    nonlocal.curr_obs = []
    def write_current_observations():
        nonlocal.last_write_no += 1
        nonlocal.last_write_nb_rows = len(nonlocal.curr_obs)
        nonlocal.last_write_subject_id = nonlocal.curr_subject_id
        dn = os.path.join(output_path, str(nonlocal.curr_subject_id))
        try:
            os.makedirs(dn)
        except:
            pass
        fn = os.path.join(dn, 'events.csv')
        if not os.path.exists(fn) or not os.path.isfile(fn):
            f = open(fn, 'w')
            f.write(','.join(obs_header) + '\n')
            f.close()
        w = csv.DictWriter(open(fn, 'a'), fieldnames=obs_header, quoting=csv.QUOTE_MINIMAL)
        w.writerows(nonlocal.curr_obs)
        nonlocal.curr_obs = []
    
    for row, row_no, nb_rows in read_events_table_by_row(mimic3_path, table):
        if verbose:
            if nonlocal.last_write_no != '':
                sys.stdout.write('\rprocessing {0}: ROW {1} of {2}...last write '
                                 '({3}) {4} rows for subject {5}'.format(table, row_no, nb_rows,
                                                                         nonlocal.last_write_no,
                                                                         nonlocal.last_write_nb_rows,
                                                                         nonlocal.last_write_subject_id))
            else:
                sys.stdout.write('\rprocessing {0}: ROW {1} of {2}...'.format(table, row_no, nb_rows))
        
        if (subjects_to_keep is not None and row['SUBJECT_ID'] not in subjects_to_keep):
            continue
        if (items_to_keep is not None and row['ITEMID'] not in items_to_keep):
            continue
        
        row_out = { 'SUBJECT_ID': row['SUBJECT_ID'],
                    'HADM_ID': row['HADM_ID'],
                    'ICUSTAY_ID': '' if 'ICUSTAY_ID' not in row else row['ICUSTAY_ID'],
                    'CHARTTIME': row['CHARTTIME'],
                    'ITEMID': row['ITEMID'],
                    'VALUE': row['VALUE'],
                    'VALUEUOM': row['VALUEUOM'] }
        if nonlocal.curr_subject_id != '' and nonlocal.curr_subject_id != row['SUBJECT_ID']:
            write_current_observations()
        nonlocal.curr_obs.append(row_out)
        nonlocal.curr_subject_id = row['SUBJECT_ID']
        
    if nonlocal.curr_subject_id != '':
        write_current_observations()

    if verbose:
        sys.stdout.write('\rprocessing {0}: ROW {1} of {2}...last write '
                         '({3}) {4} rows for subject {5}...DONE!\n'.format(table, row_no, nb_rows,
                                                                 nonlocal.last_write_no,
                                                                 nonlocal.last_write_nb_rows,
                                                                 nonlocal.last_write_subject_id))
<commit_msg>Remove _DATA_TABLE from CSV names<commit_after>import csv
import numpy as np
import os
import pandas as pd
import sys

from pandas import DataFrame


def read_patients_table(mimic3_path):
    pats = DataFrame.from_csv(os.path.join(mimic3_path, 'PATIENTS.csv'))
    pats = pats[['SUBJECT_ID', 'GENDER', 'DOB', 'DOD']]
    pats.DOB = pd.to_datetime(pats.DOB)
    pats.DOD = pd.to_datetime(pats.DOD)
    return pats

def read_admissions_table(mimic3_path):
    admits = DataFrame.from_csv(os.path.join(mimic3_path, 'ADMISSIONS.csv'))
    admits = admits[['SUBJECT_ID', 'HADM_ID', 'DEATHTIME', 'ETHNICITY', 'DIAGNOSIS']]
    admits.DEATHTIME = pd.to_datetime(admits.DEATHTIME)
    return admits

def read_icustays_table(mimic3_path):
    stays = DataFrame.from_csv(os.path.join(mimic3_path, 'ICUSTAYS.csv'))
    stays.INTIME = pd.to_datetime(stays.INTIME)
    stays.OUTTIME = pd.to_datetime(stays.OUTTIME)
    return stays

def read_icd_diagnoses_table(mimic3_path):
    codes = DataFrame.from_csv(os.path.join(mimic3_path, 'D_ICD_DIAGNOSES.csv'))
    codes = codes[['ICD9_CODE','SHORT_TITLE','LONG_TITLE']]
    diagnoses = DataFrame.from_csv(os.path.join(mimic3_path, 'DIAGNOSES_ICD.csv'))
    diagnoses = diagnoses.merge(codes, how='inner', left_on='ICD9_CODE', right_on='ICD9_CODE')
    diagnoses[['SUBJECT_ID','HADM_ID','SEQ_NUM']] = diagnoses[['SUBJECT_ID','HADM_ID','SEQ_NUM']].astype(int)
    return diagnoses

def read_events_table_by_row(mimic3_path, table):
    nb_rows = { 'chartevents': 263201376, 'labevents': 27872576, 'outputevents': 4349340 }
    reader = csv.DictReader(open(os.path.join(mimic3_path, table.upper() + '.csv'), 'r'))
    for i,row in enumerate(reader):
        if 'ICUSTAY_ID' not in row:
            row['ICUSTAY_ID'] = ''
        yield row, i, nb_rows[table.lower()]

def count_icd_codes(diagnoses, output_path=None):
    codes = diagnoses[['ICD9_CODE','SHORT_TITLE','LONG_TITLE']].drop_duplicates().set_index('ICD9_CODE')
    codes['COUNT'] = diagnoses.groupby('ICD9_CODE')['ICUSTAY_ID'].count()
    codes.COUNT = codes.COUNT.fillna(0).astype(int)
    codes = codes.ix[codes.COUNT>0]
    if output_path:
        codes.to_csv(output_path, index_label='ICD9_CODE')
    return codes.sort_values('COUNT', ascending=False).reset_index()

def remove_icustays_with_transfers(stays):
    stays = stays.ix[(stays.FIRST_WARDID == stays.LAST_WARDID) & (stays.FIRST_CAREUNIT == stays.LAST_CAREUNIT)]
    return stays[['SUBJECT_ID', 'HADM_ID', 'ICUSTAY_ID', 'DBSOURCE', 'INTIME', 'OUTTIME', 'LOS']]

def merge_on_subject(table1, table2):
    return table1.merge(table2, how='inner', left_on=['SUBJECT_ID'], right_on=['SUBJECT_ID'])

def merge_on_subject_admission(table1, table2):
    return table1.merge(table2, how='inner', left_on=['SUBJECT_ID', 'HADM_ID'], right_on=['SUBJECT_ID', 'HADM_ID'])

def add_age_to_icustays(stays):
    stays['AGE'] = (stays.INTIME - stays.DOB).apply(lambda s: s / np.timedelta64(1, 's')) / 60./60/24/365
    stays.AGE.ix[stays.AGE<0] = 90
    return stays

def add_mortality_to_icustays(stays):
    mortality = stays.DOD.notnull() & ((stays.INTIME <= stays.DOD) & (stays.OUTTIME >= stays.DOD))
    mortality = mortality | (stays.DEATHTIME.notnull() & ((stays.INTIME <= stays.DEATHTIME) & (stays.OUTTIME >= stays.DEATHTIME)))
    stays['MORTALITY'] = mortality.astype(int)
    return stays

def filter_admissions_on_nb_icustays(stays, min_nb_stays=1, max_nb_stays=1):
    to_keep = stays.groupby('HADM_ID').count()[['ICUSTAY_ID']].reset_index()
    to_keep = to_keep.ix[(to_keep.ICUSTAY_ID>=min_nb_stays)&(to_keep.ICUSTAY_ID<=max_nb_stays)][['HADM_ID']]
    stays = stays.merge(to_keep, how='inner', left_on='HADM_ID', right_on='HADM_ID')
    return stays

def filter_icustays_on_age(stays, min_age=18, max_age=np.inf):
    stays = stays.ix[(stays.AGE>=min_age)&(stays.AGE<=max_age)]
    return stays

def filter_diagnoses_on_stays(diagnoses, stays):
    return diagnoses.merge(stays[['SUBJECT_ID', 'HADM_ID', 'ICUSTAY_ID']].drop_duplicates(), how='inner',
                           left_on=['SUBJECT_ID', 'HADM_ID'], right_on=['SUBJECT_ID', 'HADM_ID'])

def break_up_stays_by_subject(stays, output_path, subjects=None, verbose=1):
    subjects = stays.SUBJECT_ID.unique() if subjects is None else subjects
    nb_subjects = subjects.shape[0]
    for i, subject_id in enumerate(subjects):
        if verbose:
            sys.stdout.write('\rSUBJECT {0} of {1}...'.format(i+1, nb_subjects))
        dn = os.path.join(output_path, str(subject_id))
        try:
            os.makedirs(dn)
        except:
            pass

        stays.ix[stays.SUBJECT_ID == subject_id].sort_values(by='INTIME').to_csv(os.path.join(dn, 'stays.csv'), index=False)
    if verbose:
        sys.stdout.write('DONE!\n')

def break_up_diagnoses_by_subject(diagnoses, output_path, subjects=None, verbose=1):
    subjects = stays.SUBJECT_ID.unique() if subjects is None else subjects
    nb_subjects = subjects.shape[0]
    for i, subject_id in enumerate(subjects):
        if verbose:
            sys.stdout.write('\rSUBJECT {0} of {1}...'.format(i+1, nb_subjects))
        dn = os.path.join(output_path, str(subject_id))
        try:
            os.makedirs(dn)
        except:
            pass

        diagnoses.ix[diagnoses.SUBJECT_ID == subject_id].sort_values(by=['ICUSTAY_ID','SEQ_NUM']).to_csv(os.path.join(dn, 'diagnoses.csv'), index=False)
    if verbose:
        sys.stdout.write('DONE!\n')

def read_events_table_and_break_up_by_subject(mimic3_path, table, output_path, items_to_keep=None, subjects_to_keep=None, verbose=1):
    obs_header = [ 'SUBJECT_ID', 'HADM_ID', 'ICUSTAY_ID', 'CHARTTIME', 'ITEMID', 'VALUE', 'VALUEUOM' ]
    if items_to_keep is not None:
        items_to_keep = set([ str(s) for s in items_to_keep ])
    if subjects_to_keep is not None:
        subjects_to_keep = set([ str(s) for s in subjects_to_keep ])

    class nonlocal: pass
    nonlocal.curr_subject_id = ''
    nonlocal.last_write_no = 0
    nonlocal.last_write_nb_rows = 0
    nonlocal.last_write_subject_id = ''
    nonlocal.curr_obs = []
    def write_current_observations():
        nonlocal.last_write_no += 1
        nonlocal.last_write_nb_rows = len(nonlocal.curr_obs)
        nonlocal.last_write_subject_id = nonlocal.curr_subject_id
        dn = os.path.join(output_path, str(nonlocal.curr_subject_id))
        try:
            os.makedirs(dn)
        except:
            pass
        fn = os.path.join(dn, 'events.csv')
        if not os.path.exists(fn) or not os.path.isfile(fn):
            f = open(fn, 'w')
            f.write(','.join(obs_header) + '\n')
            f.close()
        w = csv.DictWriter(open(fn, 'a'), fieldnames=obs_header, quoting=csv.QUOTE_MINIMAL)
        w.writerows(nonlocal.curr_obs)
        nonlocal.curr_obs = []
    
    for row, row_no, nb_rows in read_events_table_by_row(mimic3_path, table):
        if verbose:
            if nonlocal.last_write_no != '':
                sys.stdout.write('\rprocessing {0}: ROW {1} of {2}...last write '
                                 '({3}) {4} rows for subject {5}'.format(table, row_no, nb_rows,
                                                                         nonlocal.last_write_no,
                                                                         nonlocal.last_write_nb_rows,
                                                                         nonlocal.last_write_subject_id))
            else:
                sys.stdout.write('\rprocessing {0}: ROW {1} of {2}...'.format(table, row_no, nb_rows))
        
        if (subjects_to_keep is not None and row['SUBJECT_ID'] not in subjects_to_keep):
            continue
        if (items_to_keep is not None and row['ITEMID'] not in items_to_keep):
            continue
        
        row_out = { 'SUBJECT_ID': row['SUBJECT_ID'],
                    'HADM_ID': row['HADM_ID'],
                    'ICUSTAY_ID': '' if 'ICUSTAY_ID' not in row else row['ICUSTAY_ID'],
                    'CHARTTIME': row['CHARTTIME'],
                    'ITEMID': row['ITEMID'],
                    'VALUE': row['VALUE'],
                    'VALUEUOM': row['VALUEUOM'] }
        if nonlocal.curr_subject_id != '' and nonlocal.curr_subject_id != row['SUBJECT_ID']:
            write_current_observations()
        nonlocal.curr_obs.append(row_out)
        nonlocal.curr_subject_id = row['SUBJECT_ID']
        
    if nonlocal.curr_subject_id != '':
        write_current_observations()

    if verbose:
        sys.stdout.write('\rprocessing {0}: ROW {1} of {2}...last write '
                         '({3}) {4} rows for subject {5}...DONE!\n'.format(table, row_no, nb_rows,
                                                                 nonlocal.last_write_no,
                                                                 nonlocal.last_write_nb_rows,
                                                                 nonlocal.last_write_subject_id))

{'fraction_non_alphanumeric': 0.13587564303287855, 'fraction_numerical': 0.011406844106463879, 'mean_word_length': 2.888888888888889, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 32, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3489128', 'n_tokens_mistral': 6481, 'n_tokens_neox': 6299, 'n_words': 1020}