rapid-cortical-mapping/train.py at main · HuthLab/rapid-cortical-mapping · 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
import sys, os
import numpy as np
from em.ridge import ridge, bootstrap_ridge
import consts, paths
from utils_cm import read_stimulus, read_mri, make_valinds, save_npz

def train_reference(key, references, feature, stimuli, em_alpha = None, cm_alpha = None, nboots = consts.NBOOTS):

    # optimize encoding model ridge parameter
    if em_alpha is None:
        alphas = consts.ALPHAS
        single_alpha = False
    else:
        alphas = [em_alpha]
        single_alpha = True

    # fit bootstrap encoding models for parameter optimization and voxel selection
    em_corrs = {}
    for reference in references:
        stim = read_stimulus(feature, stimuli)
        resp = read_mri(reference, stimuli)
        _, _, em_corrs[reference] = bootstrap_ridge(stim, resp, valinds = make_valinds(len(resp)), use_corr = False, nboots = nboots, alphas = alphas, single_alpha = single_alpha)

    # choose selective voxels
    if em_alpha is None:
        all_corrs = np.hstack([em_corrs[subject].mean(2) for subject in references])
        em_alpha = consts.ALPHAS[np.argmax(all_corrs.mean(1))]
        em_corrs = {k : v[np.argmax(all_corrs.mean(1))].mean(1) for k, v in em_corrs.items()}
    else:
        em_corrs = {k : v[0].mean(1) for k, v in em_corrs.items()}
    vox = {k : np.argsort(v)[-10000:] for k, v in em_corrs.items()}

    # fit encoding models
    wt = {}
    for reference in references:
        stim = read_stimulus(feature, stimuli)
        resp = read_mri(reference, stimuli, vox = vox[reference])
        wt[reference] = ridge(stim, resp, em_alpha)

    # optimize converter ridge parameter
    if cm_alpha is None:
        all_corrs = []
        for r1 in references:
            resp1 = read_mri(r1, stimuli, vox = vox[r1])
            for r2 in references:
                if r1 == r2: continue
                resp2 = read_mri(r2, stimuli, vox = vox[r2])
                _, _, cm_corrs = bootstrap_ridge(resp1, resp2, valinds = make_valinds(len(resp1)), use_corr = False, nboots = nboots, alphas = consts.ALPHAS)
                all_corrs.append(cm_corrs.mean(2))
        cm_alpha = consts.ALPHAS[np.hstack(all_corrs).mean(1).argmax()]

    vox = {f'vox_{k}' : v for k, v in vox.items()}
    wt = {f'wt_{k}' : v for k, v in wt.items()}
    save_npz(paths.CONFIG % key, references = references, stimuli = stimuli, feature = feature, em_alpha = em_alpha, cm_alpha = cm_alpha, **vox, **wt)

def train_cross(key, goal, stimuli):

    # load data
    config = np.load(paths.CONFIG % key)
    rwt = {}
    rresp = {}
    for reference in config['references']:
        rvox = config[f'vox_{reference}']
        rwt[reference] = config[f'wt_{reference}']
        rresp[reference] = read_mri(reference, stimuli, vox = rvox)
    rwt = np.hstack([rwt[reference] for reference in config['references']])
    rresp = np.hstack([rresp[reference] for reference in config['references']])
    gresp = read_mri(goal, stimuli)

    # fit model
    converter = ridge(rresp, gresp, config['cm_alpha'].item())
    gwt = rwt.dot(converter)

    return gwt

def train_within(key, goal, feature, stimuli):

    # load data
    config = np.load(paths.CONFIG % key)
    stim = read_stimulus(feature, stimuli)
    resp = read_mri(goal, stimuli)

    # fit model
    gwt = ridge(stim, resp, config['em_alpha'].item())

    return gwt