# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import numpy as np
import nibabel as nib
from dipy.io import streamline
import os
from glob import glob

import matplotlib.pyplot as plt
import seaborn as sns
sns.set()

# %matplotlib inline
# %config InlineBackend.figure_format = 'retina'
# -

# define region labels and order:

labels = ['LH_CN', 'LH_SOC', 'LH_IC', 'LH_MGB',
          'RH_CN', 'RH_SOC', 'RH_IC', 'RH_MGB']
labels_atlas = ['LH_CN', 'RH_CN', 'LH_SOC', 'RH_SOC', 
          'LH_IC', 'RH_IC', 'LH_MGB', 'RH_MGB']
#label_order = list(range(1,9)) # if regions are already in same order as labels
label_order = [1, 3, 5, 7, 2, 4, 6, 8]

# define directories:

project_dir = os.path.abspath('/om2/user/ksitek/exvivo/')

# ## Full resolution, dilated ROI masks

processing_dir = os.path.join(project_dir, 'analysis/dipy/csd/',
                              'alow-0p001_angthr-75_minangle-10_fathresh-50_20190517_0.2mm_' \
                              'conj_kevin_v2_faruk_v1_dil-500um')
sl_dir = os.path.join(processing_dir, 'target_streamlines')

# load streamlines and check the number of streamlines:

sl_file = os.path.join(sl_dir, 'all_atlas_streamlines.trk')
streamlines, hdr = streamline.load_trk(sl_file)
total_sl_count = hdr['nb_streamlines']
print('total streamlines = %d'%(total_sl_count))

# +
# create a connectivity matrix from the seed-target streamline files
matrix_sl = np.zeros((len(labels), len(labels)))
for i, seed in enumerate(label_order):
    seed_label = labels_atlas[seed-1]
    print('seed: %s'%seed_label)
    for j, target in enumerate(label_order):
        target_label = labels_atlas[target-1]
        print('     target: %s'%target_label)
        try: 
            sl_file = os.path.join(sl_dir, 'target_streamlines_seed-%d_target-%d.trk'%(seed, target))
            streamlines, hdr = streamline.load_trk(sl_file)
            sl_count = hdr['nb_streamlines']
            print('             %d'%sl_count)
        except: sl_count = 0
        matrix_sl[i,j] = sl_count
matrix_sl = matrix_sl.astype(np.int) # convert to integers

# normalize the matrix by the total number of streamlines
norm_matrix_sl = matrix_sl/total_sl_count*100

# save matrix
txt_filename = os.path.join(sl_dir, 'sl_conn_mat.txt')
np.savetxt(txt_filename, matrix_sl)
# -

# create a custom color palette so that 0 is white (not light red)
new_cp = sns.color_palette("Reds",n_colors=10000)
new_cp[0] = [1,1,1]
new_cp[0:10]

# make plot labels
norm_mat_str = []
for rx, row in enumerate(norm_matrix_sl):
    norm_mat_str_row = [str(i) for i in row]
    str_row = [j[:4] for j in norm_mat_str_row]
    norm_mat_str.append(str_row)
norm_mat_str_arr = np.array(norm_mat_str)

# +
# make normalized streamline count plot
plt.clf()
plt.figure(figsize=(5,5))

mask = np.zeros_like(norm_matrix_sl)
mask[np.triu_indices_from(mask)] = True
with sns.axes_style("white"):
    ax = sns.heatmap(norm_matrix_sl, annot=norm_mat_str_arr, fmt="s", linewidths=.5,# ax=ax,
    #ax = sns.heatmap(norm_matrix_sl, annot=True, fmt=".02f", linewidths=.5,# ax=ax,
                xticklabels=labels, yticklabels=labels, square=True,#mask=mask, 
                cmap=new_cp, cbar=False);
ax.set(title='percent of total streamlines (n = %d)'%(total_sl_count));
plt.tight_layout()

# save plot
out_file = os.path.join(sl_dir, 'streamline_norm_heatplot.png')
print(out_file)
ax.figure.savefig(out_file, box_inches='tight', dpi=200)

# +
# make streamline count plot
plt.clf()
plt.figure(figsize=(5,5))

mask = np.zeros_like(matrix_sl)
mask[np.triu_indices_from(mask)] = True
with sns.axes_style("white"):
    ax = sns.heatmap(matrix_sl, annot=True, fmt="d", linewidths=.5,# ax=ax,
                xticklabels=labels, yticklabels=labels, square=True,#mask=mask, 
                cmap=new_cp, cbar=False);
ax.set(title='number of streamlines (n = %d)'%total_sl_count);
plt.tight_layout()

# save plot
out_file = os.path.join(sl_dir, 'streamline_heatplot.png')
print(out_file)
ax.figure.savefig(out_file, box_inches='tight', dpi=200)

# -

# ## Full resolution, original ROI masks

processing_dir = os.path.join(project_dir, 'analysis/dipy/csd/',
                              'alow-0p001_angthr-75_minangle-10_fathresh-50_20190517_0.2mm_' \
                              'conj_kevin_v2_faruk_v1')
sl_dir = os.path.join(processing_dir, 'target_streamlines')

# load streamlines and check the number of streamlines:

sl_file = os.path.join(sl_dir, 'all_atlas_streamlines.trk')
streamlines, hdr = streamline.load_trk(sl_file)
total_sl_count = hdr['nb_streamlines']
print('total streamlines = %d'%(total_sl_count))

# +
# create a connectivity matrix from the seed-target streamline files
matrix_sl = np.zeros((len(labels), len(labels)))
for i, seed in enumerate(label_order):
    seed_label = labels_atlas[seed-1]
    print('seed: %s'%seed_label)
    for j, target in enumerate(label_order):
        target_label = labels_atlas[target-1]
        print('     target: %s'%target_label)
        try: 
            sl_file = os.path.join(sl_dir, 'target_streamlines_seed-%d_target-%d.trk'%(seed, target))
            streamlines, hdr = streamline.load_trk(sl_file)
            sl_count = hdr['nb_streamlines']
            print('             %d'%sl_count)
        except: sl_count = 0
        matrix_sl[i,j] = sl_count
matrix_sl = matrix_sl.astype(np.int) # convert to integers

# normalize the matrix by the total number of streamlines
norm_matrix_sl = matrix_sl/total_sl_count*100

# save matrix
txt_filename = os.path.join(sl_dir, 'sl_conn_mat.txt')
np.savetxt(txt_filename, matrix_sl)
# -

# create a custom color palette so that 0 is white (not light red)
new_cp = sns.color_palette("Reds",n_colors=10000)
new_cp[0] = [1,1,1]
new_cp[0:10]

# make plot labels
norm_mat_str = []
for rx, row in enumerate(norm_matrix_sl):
    norm_mat_str_row = [str(i) for i in row]
    str_row = [j[:4] for j in norm_mat_str_row]
    norm_mat_str.append(str_row)
norm_mat_str_arr = np.array(norm_mat_str)

# +
# make normalized streamline count plot
plt.clf()
plt.figure(figsize=(5,5))

mask = np.zeros_like(norm_matrix_sl)
mask[np.triu_indices_from(mask)] = True
with sns.axes_style("white"):
    ax = sns.heatmap(norm_matrix_sl, annot=norm_mat_str_arr, fmt="s", linewidths=.5,# ax=ax,
    #ax = sns.heatmap(norm_matrix_sl, annot=True, fmt=".02f", linewidths=.5,# ax=ax,
                xticklabels=labels, yticklabels=labels, square=True,#mask=mask, 
                cmap=new_cp, cbar=False);
ax.set(title='percent of total streamlines (n = %d)'%(total_sl_count));
plt.tight_layout()

# save plot
out_file = os.path.join(sl_dir, 'streamline_norm_heatplot.png')
print(out_file)
ax.figure.savefig(out_file, box_inches='tight', dpi=200)

# +
# make streamline count plot
plt.clf()
plt.figure(figsize=(5,5))

mask = np.zeros_like(matrix_sl)
mask[np.triu_indices_from(mask)] = True
with sns.axes_style("white"):
    ax = sns.heatmap(matrix_sl, annot=True, fmt="d", linewidths=.5,# ax=ax,
                xticklabels=labels, yticklabels=labels, square=True,#mask=mask, 
                cmap=new_cp, cbar=False);
ax.set(title='number of streamlines (n = %d)'%total_sl_count);
plt.tight_layout()

# save plot
out_file = os.path.join(sl_dir, 'streamline_heatplot.png')
print(out_file)
ax.figure.savefig(out_file, box_inches='tight', dpi=200)

# -

# ## Downsampled resolution, original ROI masks

processing_dir = os.path.join(project_dir, 'analysis/dipy/csd/',
                              'alow-0p001_angthr-75_minangle-10_fathresh-50_20190523_1.05mm_' \
                              'conj_kevin_v2_faruk_v1')
sl_dir = os.path.join(processing_dir, 'target_streamlines')

# load streamlines and check the number of streamlines:

sl_file = os.path.join(sl_dir, 'all_atlas_streamlines.trk')
streamlines, hdr = streamline.load_trk(sl_file)
total_sl_count = hdr['nb_streamlines']
print('total streamlines = %d'%(total_sl_count))

# +
# create a connectivity matrix from the seed-target streamline files
matrix_sl = np.zeros((len(labels), len(labels)))
for i, seed in enumerate(label_order):
    seed_label = labels_atlas[seed-1]
    print('seed: %s'%seed_label)
    for j, target in enumerate(label_order):
        target_label = labels_atlas[target-1]
        print('     target: %s'%target_label)
        try: 
            sl_file = os.path.join(sl_dir, 'target_streamlines_seed-%d_target-%d.trk'%(seed, target))
            streamlines, hdr = streamline.load_trk(sl_file)
            sl_count = hdr['nb_streamlines']
            print('             %d'%sl_count)
        except: sl_count = 0
        matrix_sl[i,j] = sl_count
matrix_sl = matrix_sl.astype(np.int) # convert to integers

# normalize the matrix by the total number of streamlines
norm_matrix_sl = matrix_sl/total_sl_count*100

# save matrix
txt_filename = os.path.join(sl_dir, 'sl_conn_mat.txt')
np.savetxt(txt_filename, matrix_sl)
# -

# create a custom color palette so that 0 is white (not light red)
new_cp = sns.color_palette("Reds",n_colors=10000)
new_cp[0] = [1,1,1]
new_cp[0:10]

# make plot labels
norm_mat_str = []
for rx, row in enumerate(norm_matrix_sl):
    norm_mat_str_row = [str(i) for i in row]
    str_row = [j[:4] for j in norm_mat_str_row]
    norm_mat_str.append(str_row)
norm_mat_str_arr = np.array(norm_mat_str)

# +
# make normalized streamline count plot
plt.clf()
plt.figure(figsize=(5,5))

mask = np.zeros_like(norm_matrix_sl)
mask[np.triu_indices_from(mask)] = True
with sns.axes_style("white"):
    ax = sns.heatmap(norm_matrix_sl, annot=norm_mat_str_arr, fmt="s", linewidths=.5,# ax=ax,
    #ax = sns.heatmap(norm_matrix_sl, annot=True, fmt=".02f", linewidths=.5,# ax=ax,
                xticklabels=labels, yticklabels=labels, square=True,#mask=mask, 
                cmap=new_cp, cbar=False);
ax.set(title='percent of total streamlines (n = %d)'%(total_sl_count));
plt.tight_layout()

# save plot
out_file = os.path.join(sl_dir, 'streamline_norm_heatplot.png')
print(out_file)
ax.figure.savefig(out_file, box_inches='tight', dpi=200)

# +
# make streamline count plot
plt.clf()
plt.figure(figsize=(5,5))

mask = np.zeros_like(matrix_sl)
mask[np.triu_indices_from(mask)] = True
with sns.axes_style("white"):
    ax = sns.heatmap(matrix_sl, annot=True, fmt="d", linewidths=.5,# ax=ax,
                xticklabels=labels, yticklabels=labels, square=True,#mask=mask, 
                cmap=new_cp, cbar=False);
ax.set(title='number of streamlines (n = %d)'%total_sl_count);
plt.tight_layout()

# save plot
out_file = os.path.join(sl_dir, 'streamline_heatplot.png')
print(out_file)
ax.figure.savefig(out_file, box_inches='tight', dpi=200)

# -





StarcoderdataJupyterScriptsDedupFiltered