Plot atlas on brain surface

plot_atlas_on_surf.py

from typing import Literal, List

import surfplot
import matplotlib as plt
import nibabel as nib
import numpy as np
import pandas as pd


def plot_atlas_on_surf(atlas_path: str,
                       data: pd.DataFrame,
                       cmap: str = 'viridis',
                       vlim: Optional[Tuple[float, float]] = None,
                       surf_type: Literal['white', 'pial', 'inflated', 'sphere', 'medial', 'sulc', 'vaavg'] = 'white',
                       views: List = ['lateral', 'medial', 'ventral']):
    import neuromaps.datasets
    import neuromaps.transforms
    import neuromaps.images
    
    # path to surface mesh
    surfaces = neuromaps.datasets.fetch_fsaverage(density='41k')
    lh, rh = surfaces[surf_type]
    
    # path to atlas
    volumetric_atlas = nib.load(atlas_path)
    atlas_label = volumetric_atlas.get_fdata()
    
    atlas_data = np.zeros(atlas_label.shape)
    for _, row in data.iterrows():
        atlas_data[atlas_label == row['idx']] = row['value']
    
    new_atlas_img = nib.Nifti1Image(atlas_data, affine=volumetric_atlas.affine, header=volumetric_atlas.header)
    
    # transform volumetric atlas to surface
    plot_lh_data, plot_rh_data = neuromaps.transforms.mni152_to_fsaverage(
        neuromaps.images.load_nifti(new_atlas_img),
        fsavg_density='41k',
        method='nearest',
    )
    
    p = surfplot.Plot(
        surf_lh=lh,
        surf_rh=rh,
        views=views,
        layout='grid',
        zoom=1.2,
        size=np.array([1000]) * (4, 6),
        brightness=0.7,
    )
    p.add_layer({'left': plot_lh_data, 'right': plot_rh_data}, cmap=cmap, color_range=vlim)
    fig = p.build(figsize=(10, 10), cbar_kws={'fontsize': 18})
    
    return fig

Example usage

p = argparse.ArgumentParser()
p.add_argument('lat_type', type=str, choices=['nonlat', 'bilat'])
args = p.parse_args()

lat_type = args.lat_type

atlas_path = '/d/gmi/1/simeon/electrode_localizer/seeglocgh/seegloc/atlases/AAL3v1_1mm.nii.gz'
atlas_label = '/d/gmi/1/simeon/electrode_localizer/seeglocgh/seegloc/atlases/AAL3v1_1mm.nii.txt'
power_data = f'/d/gmi/1/simeon/clas/output/p3_retro_intracranialpower/regions_{lat_type}_fits.csv'

df_label = pd.read_csv(atlas_label, sep=' ', header=None)
df_label = df_label.iloc[:, [0, 1]]
df_label.columns = ['idx', 'label']

df_label_right = df_label[df_label['label'].str.endswith('R')].copy()
df_label_right['label'] = df_label['label'].str.slice(None, -2)

df_label_left = df_label[df_label['label'].str.endswith('L')].copy()
df_label_left['label'] = df_label['label'].str.slice(None, -2)

df_label.set_index('label', inplace=True)
df_label_right.set_index('label', inplace=True)
df_label_left.set_index('label', inplace=True)

df_power = pd.read_csv(power_data)
df_power = df_power[df_power['pval'] < 0.05]
df_power.dropna(inplace=True)

if lat_type == 'bilat':
    df_power['idx_aal'] = df_power.iloc[:, 0].apply(lambda x: df_label.loc[x, 'idx'])
    df_power = df_power[['idx_aal', 'estimate']].rename(columns={'idx_aal': 'idx', 'estimate': 'value'})

else:
    df_power['idx_aalR'] = df_power.iloc[:, 0].apply(lambda x: df_label_right.loc[x, 'idx'])
    df_power['idx_aalL'] = df_power.iloc[:, 0].apply(lambda x: df_label_left.loc[x, 'idx'])
    
    df_power = pd.concat((df_power[['idx_aalR', 'estimate']].rename(columns={
        'idx_aalR': 'idx',
        'estimate': 'value'
    }), df_power[['idx_aalL', 'estimate']].rename(columns={
        'idx_aalL': 'idx',
        'estimate': 'value'
    })),
                         ignore_index=True)

fig = plot_atlas_on_surf(
    atlas_path,
    df_power,
    cmap='RdBu_r',
    surf_type='white',
    views=['lateral', 'medial', 'ventral'],
    vlim=tuple(np.array([-1, 1]) * np.max(np.abs(df_power['value']))),
)
fig.savefig(f'/d/gmi/1/simeon/clas/output/p3_retro_intracranialpower/sig_regions_{lat_type}.png', dpi=1200)
plt.show(block=True)