<commit_before>import os
import gzip
import numpy as np
import nibabel as nib

from ..algorithm.graph_tool import node_attr2array
from .io import GiftiReader, CiftiReader


def read_mgh_mgz(filepath):

    ext = os.path.splitext(filepath)[1]
    if ext == ".mgz":
        openfile = gzip.open
    elif ext == ".mgh":
        openfile = open
    else:
        raise ValueError("The data must be a mgh or mgz file!")

    fobj = openfile(filepath, "rb")
    # We have to use np.fromstring here as gzip fileobjects don't work
    # with np.fromfile; same goes for try/finally instead of with statement
    try:
        v = np.fromstring(fobj.read(4), ">i4")[0]
        if v != 1:
            # I don't actually know what versions this code will read, so to be
            # on the safe side, let's only let version 1 in for now.
            # Scalar data might also be in curv format (e.g. lh.thickness)
            # in which case the first item in the file is a magic number.
            raise NotImplementedError("Scalar data file version not supported")
        ndim1 = np.fromstring(fobj.read(4), ">i4")[0]
        ndim2 = np.fromstring(fobj.read(4), ">i4")[0]
        ndim3 = np.fromstring(fobj.read(4), ">i4")[0]
        nframes = np.fromstring(fobj.read(4), ">i4")[0]
        datatype = np.fromstring(fobj.read(4), ">i4")[0]
        # Set the number of bytes per voxel and numpy data type according to
        # FS codes
        databytes, typecode = {0: (1, ">i1"), 1: (4, ">i4"), 3: (4, ">f4"),
                               4: (2, ">h")}[datatype]
        # Ignore the rest of the header here, just seek to the data
        fobj.seek(284)
        nbytes = ndim1 * ndim2 * ndim3 * nframes * databytes
        # Read in all the data, keep it in flat representation
        # (is this ever a problem?)
        _data = np.fromstring(fobj.read(nbytes), typecode)
    finally:
        fobj.close()

    data = []
    if _data.ndim == 4:
        for idx in range(_data.shape[3]):
            data.append(np.ravel(_data[..., idx], order='F'))
    else:
        data.append(np.ravel(_data, order="F"))
    data = np.array(data).T

    return data


def read_scalar_data(fpath, n_vtx=None, brain_structure=None):

    fname = os.path.basename(fpath)
    suffix0 = fname.split('.')[-1]
    suffix1 = fname.split('.')[-2]
    islabel = False
    if suffix0 in ('curv', 'thickness', 'sulc', 'area'):
        data = nib.freesurfer.read_morph_data(fpath)
        data = data.astype(np.float64)  # necessary for visualization
        data = np.atleast_2d(data).T

    elif suffix0 == 'label':
        islabel = True
        _data = nib.freesurfer.read_label(fpath)
        if n_vtx is None:
            raise RuntimeError("Reading label as scalar data need specify the number of vertices.")
        else:
            if np.max(_data) < n_vtx:
                label_array = np.zeros(n_vtx, np.int)
                label_array[_data] = 1
                data = np.array([label_array]).T
            else:
                raise RuntimeError('vertices number mismatch!')

    elif suffix0 == 'nii':
        if suffix1 in ('dscalar', 'dtseries', 'dlabel'):
            if suffix1 == 'dlabel':
                islabel = True
            reader = CiftiReader(fpath)
            data = reader.get_data(brain_structure, True).T
        else:
            Warning('The data will be regarded as a nifti file.')
            _data = nib.load(fpath).get_data()
            data = []
            if _data.ndim == 4:
                for idx in range(_data.shape[3]):
                    data.append(np.ravel(_data[..., idx], order='F'))
            else:
                data.append(np.ravel(_data, order="F"))
            data = np.array(data).T

    elif suffix0 == 'gz' and suffix1 == 'nii':
        _data = nib.load(fpath).get_data()
        data = []
        if _data.ndim == 4:
            for idx in range(_data.shape[3]):
                data.append(np.ravel(_data[..., idx], order='F'))
        else:
            data.append(np.ravel(_data, order="F"))
        data = np.array(data).T

    elif suffix0 in ('mgh', 'mgz'):
        data = read_mgh_mgz(fpath)
        data = data.astype(np.float64)

    elif suffix0 == 'gii':
        if suffix1 == 'label':
            islabel = True
        data = nib.load(fpath).darrays[0].data
        data = np.atleast_2d(data).T

    else:
        raise RuntimeError('Unsupported data type.')

    if n_vtx is not None:
        if data.shape[0] != n_vtx:
            raise RuntimeError('vertices number mismatch!')<commit_msg>Refine the NIFTI file reading function<commit_after>import os
import gzip
import numpy as np
import nibabel as nib

from ..algorithm.graph_tool import node_attr2array
from .io import GiftiReader, CiftiReader


def read_mgh_mgz(filepath):

    ext = os.path.splitext(filepath)[1]
    if ext == ".mgz":
        openfile = gzip.open
    elif ext == ".mgh":
        openfile = open
    else:
        raise ValueError("The data must be a mgh or mgz file!")

    fobj = openfile(filepath, "rb")
    # We have to use np.fromstring here as gzip fileobjects don't work
    # with np.fromfile; same goes for try/finally instead of with statement
    try:
        v = np.fromstring(fobj.read(4), ">i4")[0]
        if v != 1:
            # I don't actually know what versions this code will read, so to be
            # on the safe side, let's only let version 1 in for now.
            # Scalar data might also be in curv format (e.g. lh.thickness)
            # in which case the first item in the file is a magic number.
            raise NotImplementedError("Scalar data file version not supported")
        ndim1 = np.fromstring(fobj.read(4), ">i4")[0]
        ndim2 = np.fromstring(fobj.read(4), ">i4")[0]
        ndim3 = np.fromstring(fobj.read(4), ">i4")[0]
        nframes = np.fromstring(fobj.read(4), ">i4")[0]
        datatype = np.fromstring(fobj.read(4), ">i4")[0]
        # Set the number of bytes per voxel and numpy data type according to
        # FS codes
        databytes, typecode = {0: (1, ">i1"), 1: (4, ">i4"), 3: (4, ">f4"),
                               4: (2, ">h")}[datatype]
        # Ignore the rest of the header here, just seek to the data
        fobj.seek(284)
        nbytes = ndim1 * ndim2 * ndim3 * nframes * databytes
        # Read in all the data, keep it in flat representation
        # (is this ever a problem?)
        _data = np.fromstring(fobj.read(nbytes), typecode)
    finally:
        fobj.close()

    data = []
    if _data.ndim == 4:
        for idx in range(_data.shape[3]):
            data.append(np.ravel(_data[..., idx], order='F'))
    else:
        data.append(np.ravel(_data, order="F"))
    data = np.array(data).T

    return data


def read_scalar_data(fpath, n_vtx=None, brain_structure=None):

    fname = os.path.basename(fpath)
    suffix0 = fname.split('.')[-1]
    suffix1 = fname.split('.')[-2]
    islabel = False
    if suffix0 in ('curv', 'thickness', 'sulc', 'area'):
        data = nib.freesurfer.read_morph_data(fpath)
        data = data.astype(np.float64)  # necessary for visualization
        data = np.atleast_2d(data).T

    elif suffix0 == 'label':
        islabel = True
        _data = nib.freesurfer.read_label(fpath)
        if n_vtx is None:
            raise RuntimeError("Reading label as scalar data need specify the number of vertices.")
        else:
            if np.max(_data) < n_vtx:
                label_array = np.zeros(n_vtx, np.int)
                label_array[_data] = 1
                data = np.array([label_array]).T
            else:
                raise RuntimeError('vertices number mismatch!')

    elif suffix0 == 'nii':
        if suffix1 in ('dscalar', 'dtseries', 'dlabel'):
            if suffix1 == 'dlabel':
                islabel = True
            reader = CiftiReader(fpath)
            data = reader.get_data(brain_structure, True).T
        else:
            Warning('The data will be regarded as a nifti file.\n'
                    'If the data in the file is a 1D array, it will be regard as one surface map.\n'
                    'If the data in the file is a 2D array, each row of it will be regard as a surface map.\n'
                    'If the number of dimension is larger than 2, an error will be thrown.')
            data = nib.load(fpath).get_data()
            if data.ndim > 2:
                raise ValueError("The number of dimension of data array is larger than 2.")
            else:
                data = np.atleast_2d(data).T

    elif suffix0 == 'gz' and suffix1 == 'nii':
        Warning('The data will be regarded as a nifti file.\n'
                'If the data in the file is a 1D array, it will be regard as one surface map.\n'
                'If the data in the file is a 2D array, each row of it will be regard as a surface map.\n'
                'If the number of dimension is larger than 2, an error will be thrown.')
        data = nib.load(fpath).get_data()
        if data.ndim > 2:
            raise ValueError("The number of dimension of data array is larger than 2.")
        else:
            data = np.atleast_2d(data).T

    elif suffix0 in ('mgh', 'mgz'):
        data = read_mgh_mgz(fpath)
        data = data.astype(np.float64)

    elif suffix0 == 'gii':
        if suffix1 == 'label':
            islabel = True
        data = nib.load(fpath).darrays[0].data
        data = np.atleast_2d(data).T

    else:
        raise RuntimeError('Unsupported data type.')

    if n_vtx is not None:
        if data.shape[0] != n_vtx:
            raise RuntimeError('vertices number mismatch!')
{'fraction_non_alphanumeric': 0.11829077644606566, 'fraction_numerical': 0.015633142261594582, 'mean_word_length': 2.0444162436548226, 'pattern_counts': {'":': 4, '<': 5, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4508966', 'n_tokens_mistral': 3175, 'n_tokens_neox': 3017, 'n_words': 941}