# Buggy code
def restrict_forward_to_label(fwd, labels):
    """Restrict forward operator to labels.

    Parameters
    ----------
    fwd : Forward
        Forward operator.
    labels : label object | list
        Label object or list of label objects.

    Returns
    -------
    fwd_out : dict
        Restricted forward operator.

    See Also
    --------
    restrict_forward_to_stc
    """
    message = "labels must be instance of Label or a list of Label."
    vertices = [np.array([], int), np.array([], int)]

    if not isinstance(labels, list):
        labels = [labels]

    # Get vertices separately of each hemisphere from all label
    for label in labels:
        if not isinstance(label, Label):
            raise TypeError(message + " Instead received %s" % type(label))
        i = 0 if label.hemi == "lh" else 1
        vertices[i] = np.append(vertices[i], label.vertices)
    # Remove duplicates and sort
    vertices = [np.unique(vert_hemi) for vert_hemi in vertices]

    fwd_out = deepcopy(fwd)
    fwd_out["source_rr"] = np.zeros((0, 3))
    fwd_out["nsource"] = 0
    fwd_out["source_nn"] = np.zeros((0, 3))
    fwd_out["sol"]["data"] = np.zeros((fwd["sol"]["data"].shape[0], 0))
    fwd_out["sol"]["ncol"] = 0
    nuse_lh = fwd["src"][0]["nuse"]

    for i in range(2):
        fwd_out["src"][i]["vertno"] = np.array([], int)
        fwd_out["src"][i]["nuse"] = 0
        fwd_out["src"][i]["inuse"] = fwd["src"][i]["inuse"].copy()
        fwd_out["src"][i]["inuse"].fill(0)
        fwd_out["src"][i]["use_tris"] = np.array([], int)
        fwd_out["src"][i]["nuse_tri"] = np.array([0])

        # src_sel is idx to cols in fwd that are in any label per hemi
        src_sel = np.intersect1d(fwd["src"][i]["vertno"], vertices[i])
        src_sel = np.searchsorted(fwd["src"][i]["vertno"], src_sel)

        # Reconstruct each src
        vertno = fwd["src"][i]["vertno"][src_sel]
        fwd_out["src"][i]["inuse"][vertno] = 1
        fwd_out["src"][i]["nuse"] += len(vertno)
        fwd_out["src"][i]["vertno"] = np.where(fwd_out["src"][i]["inuse"])[0]

        # Reconstruct part of fwd that is not sol data
        src_sel += i * nuse_lh  # Add column shift to right hemi
        fwd_out["source_rr"] = np.vstack(
            [fwd_out["source_rr"], fwd["source_rr"][src_sel]]
        )
        fwd_out["nsource"] += len(src_sel)

        if is_fixed_orient(fwd):
            idx = src_sel
        else:
            idx = (3 * src_sel[:, None] + np.arange(3)).ravel()

        fwd_out["source_nn"] = np.vstack([fwd_out["source_nn"], fwd["source_nn"][idx]])
        fwd_out["sol"]["data"] = np.hstack(
            [fwd_out["sol"]["data"], fwd["sol"]["data"][:, idx]]
        )
        fwd_out["sol"]["ncol"] += len(idx)

    return fwd_out


# Traceback (IndexError)
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-300-d1587ecf7da4> in <module>()
4 stc = mne.read_source_estimate(path + '/MEG/sample/sample_audvis-meg')
5 fwd_restricted = mne.forward.restrict_forward_to_stc(fwd, stc)
----> 6 mne.write_forward_solution('restricted-fwd.fif', fwd_r, overwrite=True)

/data/marijn/mne-python/mne/forward/forward.pyc in write_forward_solution(fname, fwd, overwrite, verbose)

/data/marijn/mne-python/mne/utils.pyc in verbose(function, *args, **kwargs)
724         with use_log_level(verbose_level):
725             return function(*args, **kwargs)
--> 726     return function(*args, **kwargs)
727
728

/data/marijn/mne-python/mne/forward/forward.pyc in write_forward_solution(fname, fwd, overwrite, verbose)
732     # Write the source spaces (again)
733     #
--> 734     _write_source_spaces_to_fid(fid, src)
735     n_vert = sum([ss['nuse'] for ss in src])
736     n_col = fwd['sol']['data'].shape[1]

/data/marijn/mne-python/mne/source_space.pyc in _write_source_spaces_to_fid(fid, src, verbose)

/data/marijn/mne-python/mne/utils.pyc in verbose(function, *args, **kwargs)
724         with use_log_level(verbose_level):
725             return function(*args, **kwargs)
--> 726     return function(*args, **kwargs)
727
728

/data/marijn/mne-python/mne/source_space.pyc in _write_source_spaces_to_fid(fid, src, verbose)
926         logger.info('    Write a source space...')
927         start_block(fid, FIFF.FIFFB_MNE_SOURCE_SPACE)
--> 928         _write_one_source_space(fid, s, verbose)
929         end_block(fid, FIFF.FIFFB_MNE_SOURCE_SPACE)
930         logger.info('    [done]')

/data/marijn/mne-python/mne/source_space.pyc in _write_one_source_space(fid, this, verbose)
1016         write_int(fid, FIFF.FIFF_MNE_SOURCE_SPACE_NUSE_TRI, this['nuse_tri'])
1017         write_int_matrix(fid, FIFF.FIFF_MNE_SOURCE_SPACE_USE_TRIANGLES,
-> 1018                          this['use_tris'] + 1)
1019
1020     if this['type'] == 'vol':

/data/marijn/mne-python/mne/io/write.pyc in write_int_matrix(fid, kind, mat)
172
173     dims = np.empty(3, dtype=np.int32)
--> 174     dims[0] = mat.shape[1]
175     dims[1] = mat.shape[0]
176     dims[2] = 2

IndexError: tuple index out of range

# Fixed code
def restrict_forward_to_label(fwd, labels):
    """Restrict forward operator to labels.

    Parameters
    ----------
    fwd : Forward
        Forward operator.
    labels : label object | list
        Label object or list of label objects.

    Returns
    -------
    fwd_out : dict
        Restricted forward operator.

    See Also
    --------
    restrict_forward_to_stc
    """
    message = "labels must be instance of Label or a list of Label."
    vertices = [np.array([], int), np.array([], int)]

    if not isinstance(labels, list):
        labels = [labels]

    # Get vertices separately of each hemisphere from all label
    for label in labels:
        if not isinstance(label, Label):
            raise TypeError(message + " Instead received %s" % type(label))
        i = 0 if label.hemi == "lh" else 1
        vertices[i] = np.append(vertices[i], label.vertices)
    # Remove duplicates and sort
    vertices = [np.unique(vert_hemi) for vert_hemi in vertices]

    fwd_out = deepcopy(fwd)
    fwd_out["source_rr"] = np.zeros((0, 3))
    fwd_out["nsource"] = 0
    fwd_out["source_nn"] = np.zeros((0, 3))
    fwd_out["sol"]["data"] = np.zeros((fwd["sol"]["data"].shape[0], 0))
    fwd_out["sol"]["ncol"] = 0
    nuse_lh = fwd["src"][0]["nuse"]

    for i in range(2):
        fwd_out["src"][i]["vertno"] = np.array([], int)
        fwd_out["src"][i]["nuse"] = 0
        fwd_out["src"][i]["inuse"] = fwd["src"][i]["inuse"].copy()
        fwd_out["src"][i]["inuse"].fill(0)
        fwd_out["src"][i]["use_tris"] = np.array([[]], int)
        fwd_out["src"][i]["nuse_tri"] = np.array([0])

        # src_sel is idx to cols in fwd that are in any label per hemi
        src_sel = np.intersect1d(fwd["src"][i]["vertno"], vertices[i])
        src_sel = np.searchsorted(fwd["src"][i]["vertno"], src_sel)

        # Reconstruct each src
        vertno = fwd["src"][i]["vertno"][src_sel]
        fwd_out["src"][i]["inuse"][vertno] = 1
        fwd_out["src"][i]["nuse"] += len(vertno)
        fwd_out["src"][i]["vertno"] = np.where(fwd_out["src"][i]["inuse"])[0]

        # Reconstruct part of fwd that is not sol data
        src_sel += i * nuse_lh  # Add column shift to right hemi
        fwd_out["source_rr"] = np.vstack(
            [fwd_out["source_rr"], fwd["source_rr"][src_sel]]
        )
        fwd_out["nsource"] += len(src_sel)

        if is_fixed_orient(fwd):
            idx = src_sel
        else:
            idx = (3 * src_sel[:, None] + np.arange(3)).ravel()

        fwd_out["source_nn"] = np.vstack([fwd_out["source_nn"], fwd["source_nn"][idx]])
        fwd_out["sol"]["data"] = np.hstack(
            [fwd_out["sol"]["data"], fwd["sol"]["data"][:, idx]]
        )
        fwd_out["sol"]["ncol"] += len(idx)

    return fwd_out