example.py

import h5py
import scipy.sparse
import numpy as np
import matplotlib.pyplot as plt


def load_hdf5_array(file_name, key=None, slice=slice(0, None)):
    """Function to load data from an hdf file.

    Parameters
    ----------
    file_name: string
        hdf5 file name.
    key: string
        Key name to load. If not provided, all keys will be loaded.
    slice: slice, or tuple of slices
        Load only a slice of the hdf5 array. It will load `array[slice]`.
        Use a tuple of slices to get a slice in multiple dimensions.

    Returns
    -------
    result : array or dictionary
        Array, or dictionary of arrays (if `key` is None).
    """
    with h5py.File(file_name, mode='r') as hf:
        if key is None:
            data = dict()
            for k in hf.keys():
                data[k] = hf[k][slice]
            return data
        else:
            return hf[key][slice]


def load_hdf5_sparse_array(file_name, key):
    """Load a scipy sparse array from an hdf file

    Parameters
    ----------
    file_name : string
        File name containing array to be loaded.
    key : string
        Name of variable to be loaded.

    Notes
    -----
    This function relies on variables being stored with specific naming
    conventions, so cannot be used to load arbitrary sparse arrays.
    """
    with h5py.File(file_name, mode='r') as hf:
        data = (hf['%s_data' % key], hf['%s_indices' % key],
                hf['%s_indptr' % key])
        sparsemat = scipy.sparse.csr_matrix(data, shape=hf['%s_shape' % key])
    return sparsemat


def save_hdf5_dataset(file_name, dataset, mode='w'):
    """Save a dataset of arrays and sparse arrays.

    Parameters
    ----------
    file_name : str
        Full name of the file.
    dataset : dict of arrays
        Mappers to save.
    mode : str
        File opening model.
        Use 'w' to write from scratch, 'a' to add to existing file.
    """
    print("Saving... ", end="", flush=True)

    with h5py.File(file_name, mode=mode) as hf:
        for name, array in dataset.items():

            if scipy.sparse.issparse(array):  # sparse array
                array = array.tocsr()
                hf.create_dataset(name + '_indices', data=array.indices,
                                  compression='gzip')
                hf.create_dataset(name + '_data', data=array.data,
                                  compression='gzip')
                hf.create_dataset(name + '_indptr', data=array.indptr,
                                  compression='gzip')
                hf.create_dataset(name + '_shape', data=array.shape,
                                  compression='gzip')
            else:  # dense array
                hf.create_dataset(name, data=array, compression='gzip')

    print("Saved %s" % file_name)


def map_voxels_to_flatmap(voxels, mapper_file):
    """Generate flatmap image from voxel array using a mapper.

    This function maps an array of voxels into a flattened representation
    of an individual subject's brain.

    Parameters
    ----------
    voxels: array of shape (n_voxels, )
        Voxel values to be mapped.
    mapper_file: string
        File containing mapping arrays for a particular subject.

    Returns
    -------
    image : array of shape (width, height)
        Flatmap image.
    """
    voxel_to_flatmap = load_hdf5_sparse_array(mapper_file, 'voxel_to_flatmap')
    flatmap_mask = load_hdf5_array(mapper_file, 'flatmap_mask')

    badmask = np.array(voxel_to_flatmap.sum(1) > 0).ravel()
    img = (np.nan * np.ones(flatmap_mask.shape)).astype(voxels.dtype)
    mimg = (np.nan * np.ones(badmask.shape)).astype(voxels.dtype)
    mimg[badmask] = (voxel_to_flatmap * voxels.ravel())[badmask].astype(
        mimg.dtype)
    img[flatmap_mask] = mimg
    return img.T[::-1]


def plot_flatmap_from_mapper(voxels, mapper_file, ax=None, alpha=0.7,
                             cmap='inferno', vmin=None, vmax=None,
                             with_curvature=True, with_rois=True,
                             with_colorbar=True,
                             colorbar_location=(.4, .9, .2, .05)):
    """Plot a flatmap from a mapper file.

    Note that this function does not have the full capability of pycortex,
    (like cortex.quickshow) since it is based on flatmap mappers and not on the
    original brain surface of the subject.

    Parameters
    ----------
    voxels : array of shape (n_voxels, )
        Data to be plotted.
    mapper_file : str
        File name of the mapper.
    ax : matplotlib Axes or None.
        Axes where the figure will be plotted.
        If None, a new figure is created.
    alpha : float in [0, 1], or array of shape (n_voxels, )
        Transparency of the flatmap.
    cmap : str
        Name of the matplotlib colormap.
    vmin : float or None
        Minimum value of the colormap. If None, use the 1st percentile of the
        `voxels` array.
    vmax : float or None
        Minimum value of the colormap. If None, use the 99th percentile of the
        `voxels` array.
    with_curvature : bool
        If True, show the curvature below the data layer.
    with_rois : bool
        If True, show the ROIs labels above the data layer.
    colorbar_location : [left, bottom, width, height]
        Location of the colorbar. All quantities are in fractions of figure
        width and height.

    Returns
    -------
    ax : matplotlib Axes
        Axes where the figure has been plotted.
    """
    # create a figure
    if ax is None:
        flatmap_mask = load_hdf5_array(mapper_file, key='flatmap_mask')
        figsize = np.array(flatmap_mask.shape) / 100.
        fig = plt.figure(figsize=figsize)
        ax = fig.add_axes((0, 0, 1, 1))
        ax.axis('off')

    # process plotting parameters
    if vmin is None:
        vmin = np.percentile(voxels, 1)
    if vmax is None:
        vmax = np.percentile(voxels, 99)
    if isinstance(alpha, np.ndarray):
        alpha = map_voxels_to_flatmap(alpha, mapper_file)

    # plot the data
    image = map_voxels_to_flatmap(voxels, mapper_file)
    cimg = ax.imshow(image, aspect='equal', zorder=1, alpha=alpha, cmap=cmap,
                     vmin=vmin, vmax=vmax)

    if with_colorbar:
        try:
            cbar = ax.inset_axes(colorbar_location)
        except AttributeError:  # for matplotlib < 3.0
            cbar = ax.figure.add_axes(colorbar_location)
        ax.figure.colorbar(cimg, cax=cbar, orientation='horizontal')

    # plot additional layers if present
    with h5py.File(mapper_file, mode='r') as hf:
        if with_curvature and "flatmap_curvature" in hf.keys():
            curvature = load_hdf5_array(mapper_file, key='flatmap_curvature')
            background = np.swapaxes(curvature, 0, 1)[::-1]
        else:
            background = map_voxels_to_flatmap(np.ones_like(voxels),
                                               mapper_file)
        ax.imshow(background, aspect='equal', cmap='gray', vmin=0, vmax=1,
                  zorder=0)

        if with_rois and "flatmap_rois" in hf.keys():
            rois = load_hdf5_array(mapper_file, key='flatmap_rois')
            ax.imshow(
                np.swapaxes(rois, 0, 1)[::-1], aspect='equal',
                interpolation='bicubic', zorder=2)

    return ax


if __name__ == "__main__":
    """
    Example for how to load the fMRI test data,
    and display the explainable variance on one subject's cortical surface.

    More examples at https://github.com/gallantlab/voxelwise_tutorials
    """
    import os

    directory = os.path.abspath('.')
    subject = "S01"

    # Load fMRI responses on the test set
    file_name = os.path.join(directory, 'responses',
                             f'{subject}_responses.hdf')
    Y_test = load_hdf5_array(file_name, "Y_test")

    # compute the explainable variance per voxel, based on the test set repeats
    mean_var = np.mean(np.var(Y_test, axis=1), axis=0)
    var_mean = np.var(np.mean(Y_test, axis=0), axis=0)
    explainable_variance = var_mean / mean_var

    # Map to subject flatmap
    mapper_file = os.path.join(directory, 'mappers', f'{subject}_mappers.hdf')
    plot_flatmap_from_mapper(explainable_variance, mapper_file, vmin=0,
                             vmax=0.7)
    plt.show()