tsne_helpers.py

# Assign a color to each ICD code for use in plots
icd_codes = icd_df.icd_code.unique()
icd_codes.sort()
colors = cm.rainbow(np.linspace(0,1,len(icd_codes)))
icd_colors = dict(zip(icd_codes, colors))
row_colors = list()
for idx, row in icd_df.iterrows():
    row_icd = row['icd_code']
    row_color = icd_colors[row_icd]
    row_colors.append(row_color)
row_colors = np.array(row_colors)


def plot_tsne(embedding):
    fig = plt.figure(figsize=(10,10))
    plt.scatter([p[0] for p in embedding], [p[1] for p in embedding], color=row_colors)
    plt.title('2D t-SNE Plot Colored by ICD Code')

    
def run_plot_tsne(options=dict()):
    tsne_model = manifold.TSNE(**options)
    tsne_output = tsne_model.fit_transform(scaled_data)
    plot_tsne(tsne_output)