Show the equivalence of performing source localization in source space and source localizing regression coefficients

regression.py

from pathlib import Path

import mne
import numpy as np
from sklearn.linear_model import LinearRegression

data_path = mne.datasets.kiloword.data_path()
epochs = mne.read_epochs(data_path / "kword_metadata-epo.fif")
print(epochs.metadata.head())

# Set electrode positions and apply an average reference.
# We set the reference as a projector, because that's what the source estimation wants
# later on.
montage = mne.channels.make_standard_montage("standard_1005")
epochs.set_montage(montage)
epochs.set_eeg_reference("average", projection=True)
epochs.apply_proj()  # need to call this to apply the projector

# Trim and resample to speed things up.
epochs.resample(100)

# Perform a linear regression between the concreteness rating of the word stimulus and
# the EEG signal. First on the sensor level.
X = epochs.metadata[["Concreteness"]]
Y = epochs.get_data().reshape(len(epochs), -1)
model = LinearRegression().fit(X, Y)
coef = model.coef_.reshape(len(epochs.ch_names), len(epochs.times))
coef_sensor = mne.EvokedArray(coef, epochs.info, tmin=epochs.times[0])
coef_sensor.plot_joint(times=0.4)  # we should see an effect in the n400 time window

# Perform source estimation on the EEG.
# For this example, we don't have any MRI nor sensor locations, we we'll be using a
# template: https://mne.tools/stable/auto_tutorials/forward/35_eeg_no_mri.html
subjects_dir = Path("./subjects")
fs_dir = mne.datasets.fetch_fsaverage(subjects_dir)

subject = "fsaverage"
trans = "fsaverage"  # MNE has a built-in fsaverage transformation
bem = fs_dir / "bem" / "fsaverage-5120-5120-5120-bem-sol.fif"

# For this example, don't use many source points so our computations are faster.
src = mne.setup_source_space(
    "fsaverage", spacing="oct5", subjects_dir=subjects_dir, n_jobs=-1
)


# Build the head model
fwd = mne.make_forward_solution(
    epochs.info, trans=trans, src=src, bem=bem, eeg=True, mindist=5.0
)

# Show the head model
mne.viz.plot_alignment(
    epochs.info,
    src=src,
    eeg=["original", "projected"],
    trans=trans,
    show_axes=True,
    mri_fiducials=True,
    dig="fiducials",
    subjects_dir=subjects_dir,
)

# Perform source estimation.
inv = mne.minimum_norm.make_inverse_operator(
    epochs.info, fwd, mne.make_ad_hoc_cov(epochs.info)
)
stc_epochs = mne.minimum_norm.apply_inverse_epochs(
    epochs,
    inv,
    lambda2=1 / 9,
    pick_ori="vector",  # keep source orientation!
    method="sLORETA",
)

# Perform the regression in source space.
Y = np.array([e.data for e in stc_epochs]).reshape(len(stc_epochs), -1)
model = LinearRegression().fit(X, Y)
coef = model.coef_.reshape(stc_epochs[0].data.shape)
coef_source = mne.VectorSourceEstimate(
    coef,
    stc_epochs[0].vertices,
    stc_epochs[0].tmin,
    stc_epochs[0].tstep,
    stc_epochs[0].subject,
)

# Source localize the sensor-level coefs.
coef_source2 = mne.minimum_norm.apply_inverse(
    coef_sensor, inv, lambda2=1 / 9, pick_ori="vector", method="sLORETA"
)

# Plot the result of the two approaches side-by-side
brain1 = coef_source.plot(hemi="both", initial_time=0.4)
brain2 = coef_source2.plot(hemi="both", initial_time=0.4)
mne.viz.link_brains([brain1, brain2])

# Are these two approaches equivalent?
assert np.allclose(coef_source.data, coef_source2.data)