ProbabilisticParcellation/plot.py at main · DiedrichsenLab/ProbabilisticParcellation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import sys
import ProbabilisticParcellation.evaluate as ev
import ProbabilisticParcellation.util as ut
import ProbabilisticParcellation.export_atlas as ea
import ProbabilisticParcellation.scripts.parcel_hierarchy as ph
import ProbabilisticParcellation.scripts.evaluate_atlas as eva
from Functional_Fusion.dataset import *
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sb
import pandas as pd
import torch as pt
import nitools as nt
import surfAnalysisPy as surf

atlas_dir = (
    "/Volumes/diedrichsen_data$/data/Cerebellum/ProbabilisticParcellationModel/Atlases/"
)
conn_dir = "/Volumes/diedrichsen_data$/data/Cerebellum/connectivity/maps/"
surf_dir = surf.plot._surf_dir


def plot_parcel_summary(parcel="D1", atlas="NettekovenSym32", space="MNISymC2"):
    patlas = nb.load(f"{atlas_dir}{atlas}_space-{space}_probseg.nii")
    _, cmap, labels = nt.read_lut(f"{atlas_dir}{atlas}.lut")
    pseg = suit.flatmap.vol_to_surf(patlas, stats="nanmean", space="MNISymC")

    labels = labels[1:]
    idx = np.array([l.startswith(parcel) for l in labels])
    iidx = np.where(idx)[0]
    p = pseg[:, idx].sum(axis=1)

    fig = plt.figure(figsize=(22, 16))
    spec = fig.add_gridspec(2, 6)

    axC = fig.add_subplot(spec[0:2, 2:4])
    ax = suit.flatmap.plot(
        p,
        cscale=[0, 0.4],
        render="matplotlib",
        cmap="Reds",
        new_figure=False,
        overlay_type="func",
        colorbar=False,
    )

    # Now do connectivity maps
    conn_map = nb.load(conn_dir + "Fusion_L2_06.pscalar.nii")
    weights = nt.cifti.surf_from_cifti(conn_map)
    sc = conn_map.header.get_axis(0).name
    cidx = np.empty((2,), dtype=int)
    for c in range(2):
        cidx[c] = np.where(sc == labels[iidx[c]])[0][0]

    flat = []
    # Use the mirrored flatmap for the left hemisphere
    flat.append(nb.load(surf_dir + "/fs_L/fs_LR.32k.Lm.flat.surf.gii"))
    flat.append(nb.load(surf_dir + "/fs_R/fs_LR.32k.R.flat.surf.gii"))
    border = []
    border.append(surf_dir + "/fs_L/fs_LR.32k.L.border")
    border.append(surf_dir + "/fs_R/fs_LR.32k.R.border")

    axH = np.empty((2, 2), dtype=object)
    axH[0, 0] = fig.add_subplot(spec[0, 0:2])
    axH[1, 0] = fig.add_subplot(spec[1, 0:2])
    axH[0, 1] = fig.add_subplot(spec[0, 4:])
    axH[1, 1] = fig.add_subplot(spec[1, 4:])

    for h in range(2):
        for c in range(2):
            plt.axes(axH[h, c])
            surf.plot.plotmap(
                weights[h][cidx[c], :],
                flat[h],
                underlay=None,
                overlay_type="func",
                cmap="bwr",
                cscale=[-0.002, 0.002],
                borders=border[h],
            )

    # fig.suptitle(parcel)

def plot_parcel_prob(parcel="D1", atlas="NettekovenSym32",
                     space="MNISymC2",
                     backgroundcolor='w',
                     bordercolor='k'):
    patlas = nb.load(f"{atlas_dir}{atlas}_space-{space}_probseg.nii")
    _, cmap, labels = nt.read_lut(f"{atlas_dir}{atlas}.lut")
    pseg = suit.flatmap.vol_to_surf(patlas, stats="nanmean", space="MNISymC")

    labels = labels[1:]
    idx = np.array([l.startswith(parcel) for l in labels])
    iidx = np.where(idx)[0]
    p = pseg[:, idx].sum(axis=1)

    suit.flatmap.plot(
        p,
        cscale=[0, 0.4],
        render="matplotlib",
        cmap="Reds",
        new_figure=False,
        overlay_type="func",
        colorbar=False,
        bordersize=4,
        backgroundcolor=backgroundcolor,
        bordercolor=bordercolor
    )


if __name__ == "__main__":
    plot_parcel_summary(parcel="M3")
    pass