def test_restrict_forward_to_label():
    """Test restriction of source space to label
    """
    fwd = read_forward_solution(fname_meeg, force_fixed=True)
    fwd = pick_types_forward(fwd, meg=True)

    label_path = op.join(data_path, 'MEG', 'sample', 'labels')
    labels = ['Aud-lh', 'Vis-rh']
    label_lh = read_label(op.join(label_path, labels[0] + '.label'))
    label_rh = read_label(op.join(label_path, labels[1] + '.label'))

    fwd_out = restrict_forward_to_label(fwd, [label_lh, label_rh])

    src_sel_lh = np.intersect1d(fwd['src'][0]['vertno'], label_lh.vertices)
    src_sel_lh = np.searchsorted(fwd['src'][0]['vertno'], src_sel_lh)
    vertno_lh = fwd['src'][0]['vertno'][src_sel_lh]

    nuse_lh = fwd['src'][0]['nuse']
    src_sel_rh = np.intersect1d(fwd['src'][1]['vertno'], label_rh.vertices)
    src_sel_rh = np.searchsorted(fwd['src'][1]['vertno'], src_sel_rh)
    vertno_rh = fwd['src'][1]['vertno'][src_sel_rh]
    src_sel_rh += nuse_lh

    assert_equal(fwd_out['sol']['ncol'], len(src_sel_lh) + len(src_sel_rh))
    assert_equal(fwd_out['src'][0]['nuse'], len(src_sel_lh))
    assert_equal(fwd_out['src'][1]['nuse'], len(src_sel_rh))
    assert_equal(fwd_out['src'][0]['vertno'], vertno_lh)
    assert_equal(fwd_out['src'][1]['vertno'], vertno_rh)

    fwd = read_forward_solution(fname_meeg, force_fixed=False)
    fwd = pick_types_forward(fwd, meg=True)

    label_path = op.join(data_path, 'MEG', 'sample', 'labels')
    labels = ['Aud-lh', 'Vis-rh']
    label_lh = read_label(op.join(label_path, labels[0] + '.label'))
    label_rh = read_label(op.join(label_path, labels[1] + '.label'))

    fwd_out = restrict_forward_to_label(fwd, [label_lh, label_rh])

    src_sel_lh = np.intersect1d(fwd['src'][0]['vertno'], label_lh.vertices)
    src_sel_lh = np.searchsorted(fwd['src'][0]['vertno'], src_sel_lh)
    vertno_lh = fwd['src'][0]['vertno'][src_sel_lh]

    nuse_lh = fwd['src'][0]['nuse']
    src_sel_rh = np.intersect1d(fwd['src'][1]['vertno'], label_rh.vertices)
    src_sel_rh = np.searchsorted(fwd['src'][1]['vertno'], src_sel_rh)
    vertno_rh = fwd['src'][1]['vertno'][src_sel_rh]
    src_sel_rh += nuse_lh

    assert_equal(fwd_out['sol']['ncol'],
                 3 * (len(src_sel_lh) + len(src_sel_rh)))
    assert_equal(fwd_out['src'][0]['nuse'], len(src_sel_lh))
    assert_equal(fwd_out['src'][1]['nuse'], len(src_sel_rh))
    assert_equal(fwd_out['src'][0]['vertno'], vertno_lh)
    assert_equal(fwd_out['src'][1]['vertno'], vertno_rh)

def test_apply_mne_inverse_epochs():
    """Test MNE with precomputed inverse operator on Epochs
    """
    inverse_operator = read_inverse_operator(fname_inv)
    label_lh = read_label(fname_label % 'Aud-lh')
    label_rh = read_label(fname_label % 'Aud-rh')
    event_id, tmin, tmax = 1, -0.2, 0.5
    raw = fiff.Raw(fname_raw)

    picks = fiff.pick_types(raw.info, meg=True, eeg=False, stim=True,
                            ecg=True, eog=True, include=['STI 014'],
                            exclude='bads')
    reject = dict(grad=4000e-13, mag=4e-12, eog=150e-6)
    flat = dict(grad=1e-15, mag=1e-15)

    events = read_events(fname_event)[:15]
    epochs = Epochs(raw, events, event_id, tmin, tmax, picks=picks,
                    baseline=(None, 0), reject=reject, flat=flat)
    stcs = apply_inverse_epochs(epochs, inverse_operator, lambda2, "dSPM",
                                label=label_lh, pick_ori="normal")

    assert_true(len(stcs) == 4)
    assert_true(3 < stcs[0].data.max() < 10)
    assert_true(stcs[0].subject == 'sample')

    data = sum(stc.data for stc in stcs) / len(stcs)
    flip = label_sign_flip(label_lh, inverse_operator['src'])

    label_mean = np.mean(data, axis=0)
    label_mean_flip = np.mean(flip[:, np.newaxis] * data, axis=0)

    assert_true(label_mean.max() < label_mean_flip.max())

    # test extracting a BiHemiLabel
    stcs_rh = apply_inverse_epochs(epochs, inverse_operator, lambda2, "dSPM",
                                   label=label_rh, pick_ori="normal")
    stcs_bh = apply_inverse_epochs(epochs, inverse_operator, lambda2, "dSPM",
                                   label=label_lh + label_rh,
                                   pick_ori="normal")

    n_lh = len(stcs[0].data)
    assert_array_almost_equal(stcs[0].data, stcs_bh[0].data[:n_lh])
    assert_array_almost_equal(stcs_rh[0].data, stcs_bh[0].data[n_lh:])

    # test without using a label (so delayed computation is used)
    stcs = apply_inverse_epochs(epochs, inverse_operator, lambda2, "dSPM",
                                pick_ori="normal")
    assert_true(stcs[0].subject == 'sample')
    label_stc = stcs[0].in_label(label_rh)
    assert_true(label_stc.subject == 'sample')
    assert_array_almost_equal(stcs_rh[0].data, label_stc.data)

def test_apply_mne_inverse_fixed_raw():
    """Test MNE with fixed-orientation inverse operator on Raw
    """
    raw = fiff.Raw(fname_raw)
    start = 3
    stop = 10
    _, times = raw[0, start:stop]
    label_lh = read_label(fname_label % 'Aud-lh')

    # create a fixed-orientation inverse operator
    fwd = read_forward_solution(fname_fwd, force_fixed=False, surf_ori=True)
    noise_cov = read_cov(fname_cov)
    inv_op = make_inverse_operator(raw.info, fwd, noise_cov,
                                   loose=None, depth=0.8, fixed=True)

    stc = apply_inverse_raw(raw, inv_op, lambda2, "dSPM",
                            label=label_lh, start=start, stop=stop, nave=1,
                            pick_ori=None, buffer_size=None)

    stc2 = apply_inverse_raw(raw, inv_op, lambda2, "dSPM",
                             label=label_lh, start=start, stop=stop, nave=1,
                             pick_ori=None, buffer_size=3)

    assert_true(stc.subject == 'sample')
    assert_true(stc2.subject == 'sample')
    assert_array_almost_equal(stc.times, times)
    assert_array_almost_equal(stc2.times, times)
    assert_array_almost_equal(stc.data, stc2.data)

def test_apply_mne_inverse_raw():
    """Test MNE with precomputed inverse operator on Raw."""
    start = 3
    stop = 10
    raw = read_raw_fif(fname_raw)
    label_lh = read_label(fname_label % 'Aud-lh')
    _, times = raw[0, start:stop]
    inverse_operator = read_inverse_operator(fname_full)
    inverse_operator = prepare_inverse_operator(inverse_operator, nave=1,
                                                lambda2=lambda2, method="dSPM")
    for pick_ori in [None, "normal", "vector"]:
        stc = apply_inverse_raw(raw, inverse_operator, lambda2, "dSPM",
                                label=label_lh, start=start, stop=stop, nave=1,
                                pick_ori=pick_ori, buffer_size=None,
                                prepared=True)

        stc2 = apply_inverse_raw(raw, inverse_operator, lambda2, "dSPM",
                                 label=label_lh, start=start, stop=stop,
                                 nave=1, pick_ori=pick_ori,
                                 buffer_size=3, prepared=True)

        if pick_ori is None:
            assert_true(np.all(stc.data > 0))
            assert_true(np.all(stc2.data > 0))

        assert_true(stc.subject == 'sample')
        assert_true(stc2.subject == 'sample')
        assert_array_almost_equal(stc.times, times)
        assert_array_almost_equal(stc2.times, times)
        assert_array_almost_equal(stc.data, stc2.data)

def test_tfr_with_inverse_operator():
    """Test time freq with MNE inverse computation"""

    tmin, tmax, event_id = -0.2, 0.5, 1

    # Setup for reading the raw data
    raw = io.Raw(fname_data)
    events = find_events(raw, stim_channel='STI 014')
    inverse_operator = read_inverse_operator(fname_inv)
    inv = prepare_inverse_operator(inverse_operator, nave=1,
                                   lambda2=1. / 9., method="dSPM")

    raw.info['bads'] += ['MEG 2443', 'EEG 053']  # bads + 2 more

    # picks MEG gradiometers
    picks = pick_types(raw.info, meg=True, eeg=False, eog=True,
                       stim=False, exclude='bads')

    # Load condition 1
    event_id = 1
    events3 = events[:3]  # take 3 events to keep the computation time low
    epochs = Epochs(raw, events3, event_id, tmin, tmax, picks=picks,
                    baseline=(None, 0), reject=dict(grad=4000e-13, eog=150e-6),
                    preload=True)

    # Compute a source estimate per frequency band
    bands = dict(alpha=[10, 10])
    label = read_label(fname_label)

    stcs = source_band_induced_power(epochs, inv, bands,
                                     n_cycles=2, use_fft=False, pca=True,
                                     label=label, prepared=True)

    stc = stcs['alpha']
    assert_true(len(stcs) == len(list(bands.keys())))
    assert_true(np.all(stc.data > 0))
    assert_array_almost_equal(stc.times, epochs.times)

    stcs_no_pca = source_band_induced_power(epochs, inv, bands,
                                            n_cycles=2, use_fft=False,
                                            pca=False, label=label,
                                            prepared=True)

    assert_array_almost_equal(stcs['alpha'].data, stcs_no_pca['alpha'].data)

    # Compute a source estimate per frequency band
    epochs = Epochs(raw, events[:10], event_id, tmin, tmax, picks=picks,
                    baseline=(None, 0), reject=dict(grad=4000e-13, eog=150e-6),
                    preload=True)

    frequencies = np.arange(7, 30, 2)  # define frequencies of interest
    power, phase_lock = source_induced_power(epochs, inv,
                                             frequencies, label,
                                             baseline=(-0.1, 0),
                                             baseline_mode='percent',
                                             n_cycles=2, n_jobs=1,
                                             prepared=True)
    assert_true(np.all(phase_lock > 0))
    assert_true(np.all(phase_lock <= 1))
    assert_true(np.max(power) > 10)

def test_restrict_forward_to_label():
    """Test restriction of source space to label
    """
    fwd = read_forward_solution(fname_meeg, force_fixed=True)
    fwd = pick_types_forward(fwd, meg=True)

    label_path = op.join(data_path, "MEG", "sample", "labels")
    labels = ["Aud-lh", "Vis-rh"]
    label_lh = read_label(op.join(label_path, labels[0] + ".label"))
    label_rh = read_label(op.join(label_path, labels[1] + ".label"))

    fwd_out = restrict_forward_to_label(fwd, [label_lh, label_rh])

    src_sel_lh = np.intersect1d(fwd["src"][0]["vertno"], label_lh.vertices)
    src_sel_lh = np.searchsorted(fwd["src"][0]["vertno"], src_sel_lh)

    src_sel_rh = np.intersect1d(fwd["src"][1]["vertno"], label_rh.vertices)
    src_sel_rh = np.searchsorted(fwd["src"][1]["vertno"], src_sel_rh) + len(fwd["src"][0]["vertno"])

    assert_equal(fwd_out["sol"]["ncol"], len(src_sel_lh) + len(src_sel_rh))
    assert_equal(fwd_out["src"][0]["nuse"], len(src_sel_lh))
    assert_equal(fwd_out["src"][1]["nuse"], len(src_sel_rh))
    assert_equal(fwd_out["src"][0]["vertno"], src_sel_lh)
    assert_equal(fwd_out["src"][1]["vertno"], src_sel_rh)

    fwd = read_forward_solution(fname_meeg, force_fixed=False)
    fwd = pick_types_forward(fwd, meg=True)

    label_path = op.join(data_path, "MEG", "sample", "labels")
    labels = ["Aud-lh", "Vis-rh"]
    label_lh = read_label(op.join(label_path, labels[0] + ".label"))
    label_rh = read_label(op.join(label_path, labels[1] + ".label"))

    fwd_out = restrict_forward_to_label(fwd, [label_lh, label_rh])

    src_sel_lh = np.intersect1d(fwd["src"][0]["vertno"], label_lh.vertices)
    src_sel_lh = np.searchsorted(fwd["src"][0]["vertno"], src_sel_lh)

    src_sel_rh = np.intersect1d(fwd["src"][1]["vertno"], label_rh.vertices)
    src_sel_rh = np.searchsorted(fwd["src"][1]["vertno"], src_sel_rh) + len(fwd["src"][0]["vertno"])

    assert_equal(fwd_out["sol"]["ncol"], 3 * (len(src_sel_lh) + len(src_sel_rh)))
    assert_equal(fwd_out["src"][0]["nuse"], len(src_sel_lh))
    assert_equal(fwd_out["src"][1]["nuse"], len(src_sel_rh))
    assert_equal(fwd_out["src"][0]["vertno"], src_sel_lh)
    assert_equal(fwd_out["src"][1]["vertno"], src_sel_rh)

def test_source_psd_epochs():
    """Test multi-taper source PSD computation in label from epochs"""

    raw = fiff.Raw(fname_data)
    inverse_operator = read_inverse_operator(fname_inv)
    label = read_label(fname_label)

    event_id, tmin, tmax = 1, -0.2, 0.5
    lambda2, method = 1. / 9., 'dSPM'
    bandwidth = 8.
    fmin, fmax = 0, 100

    picks = fiff.pick_types(raw.info, meg=True, eeg=False, stim=True,
                            ecg=True, eog=True, include=['STI 014'],
                            exclude='bads')
    reject = dict(grad=4000e-13, mag=4e-12, eog=150e-6)

    events = find_events(raw)
    epochs = Epochs(raw, events, event_id, tmin, tmax, picks=picks,
                    baseline=(None, 0), reject=reject)

    # only look at one epoch
    epochs.drop_bad_epochs()
    one_epochs = epochs[:1]

    # return list
    stc_psd = compute_source_psd_epochs(one_epochs, inverse_operator,
                                        lambda2=lambda2, method=method,
                                        pick_normal=True, label=label,
                                        bandwidth=bandwidth,
                                        fmin=fmin, fmax=fmax)[0]

    # return generator
    stcs = compute_source_psd_epochs(one_epochs, inverse_operator,
                                     lambda2=lambda2, method=method,
                                     pick_normal=True, label=label,
                                     bandwidth=bandwidth,
                                     fmin=fmin, fmax=fmax,
                                     return_generator=True)

    for stc in stcs:
        stc_psd_gen = stc

    assert_array_almost_equal(stc_psd.data, stc_psd_gen.data)

    # compare with direct computation
    stc = apply_inverse_epochs(one_epochs, inverse_operator,
                               lambda2=lambda2, method=method,
                               pick_normal=True, label=label)[0]

    sfreq = epochs.info['sfreq']
    psd, freqs = multitaper_psd(stc.data, sfreq=sfreq, bandwidth=bandwidth,
                                fmin=fmin, fmax=fmax)

    assert_array_almost_equal(psd, stc_psd.data)
    assert_array_almost_equal(freqs, stc_psd.times)

def test_apply_inverse_operator():
    """Test MNE inverse application
    """
    inverse_operator = read_inverse_operator(fname_full)
    evoked = _get_evoked()

    # Inverse has 306 channels - 4 proj = 302
    assert_true(compute_rank_inverse(inverse_operator) == 302)

    # Inverse has 306 channels - 4 proj = 302
    assert_true(compute_rank_inverse(inverse_operator) == 302)

    stc = apply_inverse(evoked, inverse_operator, lambda2, "MNE")
    assert_true(stc.subject == 'sample')
    assert_true(stc.data.min() > 0)
    assert_true(stc.data.max() < 10e-9)
    assert_true(stc.data.mean() > 1e-11)

    # test if using prepared and not prepared inverse operator give the same
    # result
    inv_op = prepare_inverse_operator(inverse_operator, nave=evoked.nave,
                                      lambda2=lambda2, method="MNE")
    stc2 = apply_inverse(evoked, inv_op, lambda2, "MNE")
    assert_array_almost_equal(stc.data, stc2.data)
    assert_array_almost_equal(stc.times, stc2.times)

    stc = apply_inverse(evoked, inverse_operator, lambda2, "sLORETA")
    assert_true(stc.subject == 'sample')
    assert_true(stc.data.min() > 0)
    assert_true(stc.data.max() < 10.0)
    assert_true(stc.data.mean() > 0.1)

    stc = apply_inverse(evoked, inverse_operator, lambda2, "dSPM")
    assert_true(stc.subject == 'sample')
    assert_true(stc.data.min() > 0)
    assert_true(stc.data.max() < 35)
    assert_true(stc.data.mean() > 0.1)

    # test without using a label (so delayed computation is used)
    label = read_label(fname_label % 'Aud-lh')
    stc = apply_inverse(evoked, inv_op, lambda2, "MNE")
    stc_label = apply_inverse(evoked, inv_op, lambda2, "MNE",
                              label=label)
    assert_equal(stc_label.subject, 'sample')
    label_stc = stc.in_label(label)
    assert_true(label_stc.subject == 'sample')
    assert_array_almost_equal(stc_label.data, label_stc.data)

    # Test we get errors when using custom ref or no average proj is present
    evoked.info['custom_ref_applied'] = True
    assert_raises(ValueError, apply_inverse, evoked, inv_op, lambda2, "MNE")
    evoked.info['custom_ref_applied'] = False
    evoked.info['projs'] = []  # remove EEG proj
    assert_raises(ValueError, apply_inverse, evoked, inv_op, lambda2, "MNE")

def test_source_psd():
    """Test source PSD computation in label."""
    raw = read_raw_fif(fname_data)
    inverse_operator = read_inverse_operator(fname_inv)
    label = read_label(fname_label)
    tmin, tmax = 0, 20  # seconds
    fmin, fmax = 55, 65  # Hz
    n_fft = 2048
    stc = compute_source_psd(raw, inverse_operator, lambda2=1. / 9.,
                             method="dSPM", tmin=tmin, tmax=tmax,
                             fmin=fmin, fmax=fmax, pick_ori="normal",
                             n_fft=n_fft, label=label, overlap=0.1)
    assert_true(stc.times[0] >= fmin * 1e-3)
    assert_true(stc.times[-1] <= fmax * 1e-3)
    # Time max at line frequency (60 Hz in US)
    assert_true(59e-3 <= stc.times[np.argmax(np.sum(stc.data, axis=0))] <=
                61e-3)

def test_apply_mne_inverse_fixed_raw():
    """Test MNE with fixed-orientation inverse operator on Raw."""
    raw = read_raw_fif(fname_raw)
    start = 3
    stop = 10
    _, times = raw[0, start:stop]
    label_lh = read_label(fname_label % 'Aud-lh')

    # create a fixed-orientation inverse operator
    fwd = read_forward_solution_meg(fname_fwd, force_fixed=False,
                                    surf_ori=True)
    noise_cov = read_cov(fname_cov)
    assert_raises(ValueError, make_inverse_operator,
                  raw.info, fwd, noise_cov, loose=1., fixed=True)
    inv_op = make_inverse_operator(raw.info, fwd, noise_cov,
                                   fixed=True, use_cps=True)

    inv_op2 = prepare_inverse_operator(inv_op, nave=1,
                                       lambda2=lambda2, method="dSPM")
    stc = apply_inverse_raw(raw, inv_op2, lambda2, "dSPM",
                            label=label_lh, start=start, stop=stop, nave=1,
                            pick_ori=None, buffer_size=None, prepared=True)

    stc2 = apply_inverse_raw(raw, inv_op2, lambda2, "dSPM",
                             label=label_lh, start=start, stop=stop, nave=1,
                             pick_ori=None, buffer_size=3, prepared=True)

    stc3 = apply_inverse_raw(raw, inv_op, lambda2, "dSPM",
                             label=label_lh, start=start, stop=stop, nave=1,
                             pick_ori=None, buffer_size=None)

    assert_true(stc.subject == 'sample')
    assert_true(stc2.subject == 'sample')
    assert_array_almost_equal(stc.times, times)
    assert_array_almost_equal(stc2.times, times)
    assert_array_almost_equal(stc3.times, times)
    assert_array_almost_equal(stc.data, stc2.data)
    assert_array_almost_equal(stc.data, stc3.data)

def test_apply_inverse_operator():
    """Test MNE inverse application."""
    inverse_operator = read_inverse_operator(fname_full)
    evoked = _get_evoked()

    # Inverse has 306 channels - 4 proj = 302
    assert_true(compute_rank_inverse(inverse_operator) == 302)

    # Inverse has 306 channels - 4 proj = 302
    assert_true(compute_rank_inverse(inverse_operator) == 302)

    stc = apply_inverse(evoked, inverse_operator, lambda2, "MNE")
    assert_true(stc.subject == 'sample')
    assert_true(stc.data.min() > 0)
    assert_true(stc.data.max() < 10e-9)
    assert_true(stc.data.mean() > 1e-11)

    # test if using prepared and not prepared inverse operator give the same
    # result
    inv_op = prepare_inverse_operator(inverse_operator, nave=evoked.nave,
                                      lambda2=lambda2, method="MNE")
    stc2 = apply_inverse(evoked, inv_op, lambda2, "MNE")
    assert_array_almost_equal(stc.data, stc2.data)
    assert_array_almost_equal(stc.times, stc2.times)

    stc = apply_inverse(evoked, inverse_operator, lambda2, "sLORETA")
    assert_true(stc.subject == 'sample')
    assert_true(stc.data.min() > 0)
    assert_true(stc.data.max() < 10.0)
    assert_true(stc.data.mean() > 0.1)

    stc = apply_inverse(evoked, inverse_operator, lambda2, "dSPM")
    assert_true(stc.subject == 'sample')
    assert_true(stc.data.min() > 0)
    assert_true(stc.data.max() < 35)
    assert_true(stc.data.mean() > 0.1)

    # test without using a label (so delayed computation is used)
    label = read_label(fname_label % 'Aud-lh')
    stc = apply_inverse(evoked, inv_op, lambda2, "MNE")
    stc_label = apply_inverse(evoked, inv_op, lambda2, "MNE",
                              label=label)
    assert_equal(stc_label.subject, 'sample')
    label_stc = stc.in_label(label)
    assert_true(label_stc.subject == 'sample')
    assert_allclose(stc_label.data, label_stc.data)

    # Test that no errors are raised with loose inverse ops and picking normals
    noise_cov = read_cov(fname_cov)
    fwd = read_forward_solution_meg(fname_fwd)
    inv_op2 = make_inverse_operator(evoked.info, fwd, noise_cov, loose=1,
                                    fixed='auto', depth=None)
    apply_inverse(evoked, inv_op2, 1 / 9., method='MNE',
                  pick_ori='normal')

    # Test we get errors when using custom ref or no average proj is present
    evoked.info['custom_ref_applied'] = True
    assert_raises(ValueError, apply_inverse, evoked, inv_op, lambda2, "MNE")
    evoked.info['custom_ref_applied'] = False
    evoked.info['projs'] = []  # remove EEG proj
    assert_raises(ValueError, apply_inverse, evoked, inv_op, lambda2, "MNE")

def test_source_psd_epochs():
    """Test multi-taper source PSD computation in label from epochs."""
    raw = read_raw_fif(fname_data)
    inverse_operator = read_inverse_operator(fname_inv)
    label = read_label(fname_label)

    event_id, tmin, tmax = 1, -0.2, 0.5
    lambda2, method = 1. / 9., 'dSPM'
    bandwidth = 8.
    fmin, fmax = 0, 100

    picks = pick_types(raw.info, meg=True, eeg=False, stim=True,
                       ecg=True, eog=True, include=['STI 014'],
                       exclude='bads')
    reject = dict(grad=4000e-13, mag=4e-12, eog=150e-6)

    events = find_events(raw, stim_channel='STI 014')
    epochs = Epochs(raw, events, event_id, tmin, tmax, picks=picks,
                    baseline=(None, 0), reject=reject)

    # only look at one epoch
    epochs.drop_bad()
    one_epochs = epochs[:1]

    inv = prepare_inverse_operator(inverse_operator, nave=1,
                                   lambda2=1. / 9., method="dSPM")
    # return list
    stc_psd = compute_source_psd_epochs(one_epochs, inv,
                                        lambda2=lambda2, method=method,
                                        pick_ori="normal", label=label,
                                        bandwidth=bandwidth,
                                        fmin=fmin, fmax=fmax,
                                        prepared=True)[0]

    # return generator
    stcs = compute_source_psd_epochs(one_epochs, inv,
                                     lambda2=lambda2, method=method,
                                     pick_ori="normal", label=label,
                                     bandwidth=bandwidth,
                                     fmin=fmin, fmax=fmax,
                                     return_generator=True,
                                     prepared=True)

    for stc in stcs:
        stc_psd_gen = stc

    assert_array_almost_equal(stc_psd.data, stc_psd_gen.data)

    # compare with direct computation
    stc = apply_inverse_epochs(one_epochs, inv,
                               lambda2=lambda2, method=method,
                               pick_ori="normal", label=label,
                               prepared=True)[0]

    sfreq = epochs.info['sfreq']
    psd, freqs = _psd_multitaper(stc.data, sfreq=sfreq, bandwidth=bandwidth,
                                 fmin=fmin, fmax=fmax)

    assert_array_almost_equal(psd, stc_psd.data)
    assert_array_almost_equal(freqs, stc_psd.times)

    # Check corner cases caused by tiny bandwidth
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter('always')
        compute_source_psd_epochs(one_epochs, inv,
                                  lambda2=lambda2, method=method,
                                  pick_ori="normal", label=label,
                                  bandwidth=0.01, low_bias=True,
                                  fmin=fmin, fmax=fmax,
                                  return_generator=False,
                                  prepared=True)
        compute_source_psd_epochs(one_epochs, inv,
                                  lambda2=lambda2, method=method,
                                  pick_ori="normal", label=label,
                                  bandwidth=0.01, low_bias=False,
                                  fmin=fmin, fmax=fmax,
                                  return_generator=False,
                                  prepared=True)
    assert_true(len(w) >= 2)
    assert_true(any('not properly use' in str(ww.message) for ww in w))
    assert_true(any('Bandwidth too small' in str(ww.message) for ww in w))

def test_restrict_forward_to_label():
    """Test restriction of source space to label
    """
    fwd = read_forward_solution(fname_meeg)
    fwd = convert_forward_solution(fwd, surf_ori=True, force_fixed=True,
                                   use_cps=True)
    fwd = pick_types_forward(fwd, meg=True)

    label_path = op.join(data_path, 'MEG', 'sample', 'labels')
    labels = ['Aud-lh', 'Vis-rh']
    label_lh = read_label(op.join(label_path, labels[0] + '.label'))
    label_rh = read_label(op.join(label_path, labels[1] + '.label'))

    fwd_out = restrict_forward_to_label(fwd, [label_lh, label_rh])

    src_sel_lh = np.intersect1d(fwd['src'][0]['vertno'], label_lh.vertices)
    src_sel_lh = np.searchsorted(fwd['src'][0]['vertno'], src_sel_lh)
    vertno_lh = fwd['src'][0]['vertno'][src_sel_lh]

    nuse_lh = fwd['src'][0]['nuse']
    src_sel_rh = np.intersect1d(fwd['src'][1]['vertno'], label_rh.vertices)
    src_sel_rh = np.searchsorted(fwd['src'][1]['vertno'], src_sel_rh)
    vertno_rh = fwd['src'][1]['vertno'][src_sel_rh]
    src_sel_rh += nuse_lh

    assert_equal(fwd_out['sol']['ncol'], len(src_sel_lh) + len(src_sel_rh))
    assert_equal(fwd_out['src'][0]['nuse'], len(src_sel_lh))
    assert_equal(fwd_out['src'][1]['nuse'], len(src_sel_rh))
    assert_equal(fwd_out['src'][0]['vertno'], vertno_lh)
    assert_equal(fwd_out['src'][1]['vertno'], vertno_rh)

    fwd = read_forward_solution(fname_meeg)
    fwd = pick_types_forward(fwd, meg=True)

    label_path = op.join(data_path, 'MEG', 'sample', 'labels')
    labels = ['Aud-lh', 'Vis-rh']
    label_lh = read_label(op.join(label_path, labels[0] + '.label'))
    label_rh = read_label(op.join(label_path, labels[1] + '.label'))

    fwd_out = restrict_forward_to_label(fwd, [label_lh, label_rh])

    src_sel_lh = np.intersect1d(fwd['src'][0]['vertno'], label_lh.vertices)
    src_sel_lh = np.searchsorted(fwd['src'][0]['vertno'], src_sel_lh)
    vertno_lh = fwd['src'][0]['vertno'][src_sel_lh]

    nuse_lh = fwd['src'][0]['nuse']
    src_sel_rh = np.intersect1d(fwd['src'][1]['vertno'], label_rh.vertices)
    src_sel_rh = np.searchsorted(fwd['src'][1]['vertno'], src_sel_rh)
    vertno_rh = fwd['src'][1]['vertno'][src_sel_rh]
    src_sel_rh += nuse_lh

    assert_equal(fwd_out['sol']['ncol'],
                 3 * (len(src_sel_lh) + len(src_sel_rh)))
    assert_equal(fwd_out['src'][0]['nuse'], len(src_sel_lh))
    assert_equal(fwd_out['src'][1]['nuse'], len(src_sel_rh))
    assert_equal(fwd_out['src'][0]['vertno'], vertno_lh)
    assert_equal(fwd_out['src'][1]['vertno'], vertno_rh)

    # Test saving the restricted forward object. This only works if all fields
    # are properly accounted for.
    temp_dir = _TempDir()
    fname_copy = op.join(temp_dir, 'copy-fwd.fif')
    write_forward_solution(fname_copy, fwd_out, overwrite=True)
    fwd_out_read = read_forward_solution(fname_copy)
    compare_forwards(fwd_out, fwd_out_read)

def test_mxne_inverse():
    """Test (TF-)MxNE inverse computation."""
    # Read noise covariance matrix
    cov = read_cov(fname_cov)

    # Handling average file
    loose = 0.0
    depth = 0.9

    evoked = read_evokeds(fname_data, condition=0, baseline=(None, 0))
    evoked.crop(tmin=-0.05, tmax=0.2)

    evoked_l21 = evoked.copy()
    evoked_l21.crop(tmin=0.081, tmax=0.1)
    label = read_label(fname_label)

    forward = read_forward_solution(fname_fwd)
    forward = convert_forward_solution(forward, surf_ori=True)

    # Reduce source space to make test computation faster
    inverse_operator = make_inverse_operator(evoked_l21.info, forward, cov,
                                             loose=loose, depth=depth,
                                             fixed=True, use_cps=True)
    stc_dspm = apply_inverse(evoked_l21, inverse_operator, lambda2=1. / 9.,
                             method='dSPM')
    stc_dspm.data[np.abs(stc_dspm.data) < 12] = 0.0
    stc_dspm.data[np.abs(stc_dspm.data) >= 12] = 1.
    weights_min = 0.5

    # MxNE tests
    alpha = 70  # spatial regularization parameter

    stc_prox = mixed_norm(evoked_l21, forward, cov, alpha, loose=loose,
                          depth=depth, maxit=300, tol=1e-8,
                          active_set_size=10, weights=stc_dspm,
                          weights_min=weights_min, solver='prox')
    with pytest.warns(None):  # CD
        stc_cd = mixed_norm(evoked_l21, forward, cov, alpha, loose=loose,
                            depth=depth, maxit=300, tol=1e-8,
                            active_set_size=10, weights=stc_dspm,
                            weights_min=weights_min, solver='cd',
                            pca=False)  # pca=False deprecated, doesn't matter
    stc_bcd = mixed_norm(evoked_l21, forward, cov, alpha, loose=loose,
                         depth=depth, maxit=300, tol=1e-8, active_set_size=10,
                         weights=stc_dspm, weights_min=weights_min,
                         solver='bcd')
    assert_array_almost_equal(stc_prox.times, evoked_l21.times, 5)
    assert_array_almost_equal(stc_cd.times, evoked_l21.times, 5)
    assert_array_almost_equal(stc_bcd.times, evoked_l21.times, 5)
    assert_allclose(stc_prox.data, stc_cd.data, rtol=1e-3, atol=0.0)
    assert_allclose(stc_prox.data, stc_bcd.data, rtol=1e-3, atol=0.0)
    assert_allclose(stc_cd.data, stc_bcd.data, rtol=1e-3, atol=0.0)
    assert stc_prox.vertices[1][0] in label.vertices
    assert stc_cd.vertices[1][0] in label.vertices
    assert stc_bcd.vertices[1][0] in label.vertices

    with pytest.warns(None):  # CD
        dips = mixed_norm(evoked_l21, forward, cov, alpha, loose=loose,
                          depth=depth, maxit=300, tol=1e-8, active_set_size=10,
                          weights=stc_dspm, weights_min=weights_min,
                          solver='cd', return_as_dipoles=True)
    stc_dip = make_stc_from_dipoles(dips, forward['src'])
    assert isinstance(dips[0], Dipole)
    assert stc_dip.subject == "sample"
    _check_stcs(stc_cd, stc_dip)

    with pytest.warns(None):  # CD
        stc, _ = mixed_norm(evoked_l21, forward, cov, alpha, loose=loose,
                            depth=depth, maxit=300, tol=1e-8,
                            active_set_size=10, return_residual=True,
                            solver='cd')
    assert_array_almost_equal(stc.times, evoked_l21.times, 5)
    assert stc.vertices[1][0] in label.vertices

    # irMxNE tests
    with pytest.warns(None):  # CD
        stc = mixed_norm(evoked_l21, forward, cov, alpha,
                         n_mxne_iter=5, loose=loose, depth=depth,
                         maxit=300, tol=1e-8, active_set_size=10,
                         solver='cd')
    assert_array_almost_equal(stc.times, evoked_l21.times, 5)
    assert stc.vertices[1][0] in label.vertices
    assert stc.vertices == [[63152], [79017]]

    # Do with TF-MxNE for test memory savings
    alpha = 60.  # overall regularization parameter
    l1_ratio = 0.01  # temporal regularization proportion

    stc, _ = tf_mixed_norm(evoked, forward, cov,
                           loose=loose, depth=depth, maxit=100, tol=1e-4,
                           tstep=4, wsize=16, window=0.1, weights=stc_dspm,
                           weights_min=weights_min, return_residual=True,
                           alpha=alpha, l1_ratio=l1_ratio)
    assert_array_almost_equal(stc.times, evoked.times, 5)
    assert stc.vertices[1][0] in label.vertices

    pytest.raises(ValueError, tf_mixed_norm, evoked, forward, cov,
                  alpha=101, l1_ratio=0.03)
    pytest.raises(ValueError, tf_mixed_norm, evoked, forward, cov,
                  alpha=50., l1_ratio=1.01)

def test_mxne_inverse():
    """Test (TF-)MxNE inverse computation"""
    # Handling forward solution
    evoked = fiff.Evoked(fname_data, setno=1, baseline=(None, 0))

    # Read noise covariance matrix
    cov = read_cov(fname_cov)

    # Handling average file
    setno = 0
    loose = None
    depth = 0.9

    evoked = fiff.read_evoked(fname_data, setno=setno, baseline=(None, 0))
    evoked.crop(tmin=-0.1, tmax=0.4)

    evoked_l21 = copy.deepcopy(evoked)
    evoked_l21.crop(tmin=0.08, tmax=0.1)
    label = read_label(fname_label)
    weights_min = 0.5
    forward = read_forward_solution(fname_fwd, force_fixed=False,
                                    surf_ori=True)

    # Reduce source space to make test computation faster
    inverse_operator = make_inverse_operator(evoked.info, forward, cov,
                                             loose=loose, depth=depth,
                                             fixed=True)
    stc_dspm = apply_inverse(evoked_l21, inverse_operator, lambda2=1. / 9.,
                             method='dSPM')
    stc_dspm.data[np.abs(stc_dspm.data) < 12] = 0.0
    stc_dspm.data[np.abs(stc_dspm.data) >= 12] = 1.

    # MxNE tests
    alpha = 60  # spatial regularization parameter

    stc_prox = mixed_norm(evoked_l21, forward, cov, alpha, loose=None,
                          depth=0.9, maxit=1000, tol=1e-8, active_set_size=10,
                          solver='prox')
    stc_cd = mixed_norm(evoked_l21, forward, cov, alpha, loose=None,
                        depth=0.9, maxit=1000, tol=1e-8, active_set_size=10,
                        solver='cd')
    assert_array_almost_equal(stc_prox.times, evoked_l21.times, 5)
    assert_array_almost_equal(stc_cd.times, evoked_l21.times, 5)
    assert_array_almost_equal(stc_prox.data, stc_cd.data, 5)
    assert_true(stc_prox.vertno[1][0] in label.vertices)
    assert_true(stc_cd.vertno[1][0] in label.vertices)

    stc, _ = mixed_norm(evoked_l21, forward, cov, alpha, loose=None,
                        depth=depth, maxit=500, tol=1e-4, active_set_size=10,
                        weights=stc_dspm, weights_min=weights_min,
                        return_residual=True)

    assert_array_almost_equal(stc.times, evoked_l21.times, 5)
    assert_true(stc.vertno[1][0] in label.vertices)

    # Do with TF-MxNE for test memory savings
    alpha_space = 60.  # spatial regularization parameter
    alpha_time = 1.  # temporal regularization parameter

    stc, _ = tf_mixed_norm(evoked, forward, cov, alpha_space, alpha_time,
                           loose=loose, depth=depth, maxit=100, tol=1e-4,
                           tstep=4, wsize=16, window=0.1, weights=stc_dspm,
                           weights_min=weights_min, return_residual=True)

    assert_array_almost_equal(stc.times, evoked.times, 5)
    assert_true(stc.vertno[1][0] in label.vertices)

fname_cov = op.join(data_path, 'MEG', 'sample',
                            'sample_audvis-cov.fif')
fname_fwd = op.join(data_path, 'MEG', 'sample',
                            'sample_audvis-meg-oct-6-fwd.fif')
                            # 'sample_audvis-meg-eeg-oct-6-fwd.fif')
fname_raw = op.join(data_path, 'MEG', 'sample',
                            'sample_audvis_filt-0-40_raw.fif')
fname_event = op.join(data_path, 'MEG', 'sample',
                            'sample_audvis_filt-0-40_raw-eve.fif')
label = 'Aud-lh'
fname_label = op.join(data_path, 'MEG', 'sample', 'labels', '%s.label' % label)

inverse_operator = read_inverse_operator(fname_inv)
inverse_operator_fixed = read_inverse_operator(fname_inv_fixed)
inverse_operator_vol = read_inverse_operator(fname_vol_inv)
label = read_label(fname_label)
noise_cov = read_cov(fname_cov)
raw = fiff.Raw(fname_raw)
snr = 3.0
lambda2 = 1.0 / snr ** 2


def _compare(a, b):
    if isinstance(a, dict):
        assert_true(isinstance(b, dict))
        for k, v in a.iteritems():
            if not k in b:
                raise ValueError('%s not in %s' % (k, b))
            _compare(v, b[k])
    elif isinstance(a, list):
        assert_true(len(a) == len(b))

def test_apply_inverse_operator():
    """Test MNE inverse application."""
    # use fname_inv as it will be faster than fname_full (fewer verts and chs)
    inverse_operator = read_inverse_operator(fname_inv)
    evoked = _get_evoked()

    # Inverse has 306 channels - 4 proj = 302
    assert (compute_rank_inverse(inverse_operator) == 302)

    # Inverse has 306 channels - 4 proj = 302
    assert (compute_rank_inverse(inverse_operator) == 302)

    stc = apply_inverse(evoked, inverse_operator, lambda2, "MNE")
    assert stc.subject == 'sample'
    assert stc.data.min() > 0
    assert stc.data.max() < 13e-9
    assert stc.data.mean() > 1e-11

    # test if using prepared and not prepared inverse operator give the same
    # result
    inv_op = prepare_inverse_operator(inverse_operator, nave=evoked.nave,
                                      lambda2=lambda2, method="MNE")
    stc2 = apply_inverse(evoked, inv_op, lambda2, "MNE")
    assert_array_almost_equal(stc.data, stc2.data)
    assert_array_almost_equal(stc.times, stc2.times)

    # This is little more than a smoke test...
    stc = apply_inverse(evoked, inverse_operator, lambda2, "sLORETA")
    assert stc.subject == 'sample'
    assert stc.data.min() > 0
    assert stc.data.max() < 10.0
    assert stc.data.mean() > 0.1

    stc = apply_inverse(evoked, inverse_operator, lambda2, "eLORETA")
    assert stc.subject == 'sample'
    assert stc.data.min() > 0
    assert stc.data.max() < 3.0
    assert stc.data.mean() > 0.1

    stc = apply_inverse(evoked, inverse_operator, lambda2, "dSPM")
    assert stc.subject == 'sample'