<commit_before>import numpy as np
import time
import glob
import sys
from os import makedirs, remove
from os.path import exists, join as pjoin
import nibabel
import cPickle
import commit.models
import commit.solvers
import amico.scheme
import amico.lut
import pyximport
pyximport.install( reload_support=True )


def setup( lmax = 12 ) :
    """General setup/initialization of the COMMIT framework."""
    amico.lut.precompute_rotation_matrices( lmax )


class Evaluation :
    """Class to hold all the information (data and parameters) when performing an
    evaluation with the COMMIT framework.
    """

    def __init__( self, study_path, subject ) :
        """Setup the data structures with default values.

        Parameters
        ----------
        study_path : string
            The path to the folder containing all the subjects from one study
        subject : string
            The path (relative to previous folder) to the subject folder
        """
        self.niiDWI     = None # set by "load_data" method
        self.scheme     = None # set by "load_data" method
        self.model      = None # set by "set_model" method
        self.KERNELS    = None # set by "load_kernels" method
        self.DICTIONARY = None # set by "load_dictionary" method
        self.THREADS    = None # set by "set_threads" method
        self.A          = None # set by "build_operator" method
        self.x          = None # set by "fit" method

        # store all the parameters of an evaluation with COMMIT
        self.CONFIG = {}
        self.set_config('study_path', study_path)
        self.set_config('subject', subject)
        self.set_config('DATA_path', pjoin( study_path, subject ))

        self.set_config('doNormalizeSignal', True)
        self.set_config('doMergeB0', True)
        self.set_config('doNormalizeKernels', True)
        self.set_config('doDemean', False)


    def set_config( self, key, value ) :
        self.CONFIG[ key ] = value

    def get_config( self, key ) :
        return self.CONFIG.get( key )


    def load_data( self, dwi_filename = 'DWI.nii', scheme_filename = 'DWI.scheme', b0_thr = 0 ) :
        """Load the diffusion signal and its corresponding acquisition scheme.

        Parameters
        ----------
        dwi_filename : string
            The file name of the DWI data, relative to the subject folder (default : 'DWI.nii')
        scheme_filename : string
            The file name of the corresponding acquisition scheme (default : 'DWI.scheme')
        b0_thr : float
            The threshold below which a b-value is considered a b0 (default : 0)
        """

        # Loading data and acquisition scheme
        tic = time.time()
        print '\n-> Loading data:'

        print '\t* DWI signal...'
        self.set_config('dwi_filename', dwi_filename)
        self.niiDWI  = nibabel.load( pjoin( self.get_config('DATA_path'), dwi_filename) )
        self.niiDWI_img = self.niiDWI.get_data().astype(np.float32)
        hdr = self.niiDWI.header if nibabel.__version__ >= '2.0.0' else self.niiDWI.get_header()
        self.set_config('dim', self.niiDWI_img.shape[0:3])
        self.set_config('pixdim', tuple( hdr.get_zooms()[:3] ))
        print '\t\t- dim    = %d x %d x %d x %d' % self.niiDWI_img.shape
        print '\t\t- pixdim = %.3f x %.3f x %.3f' % self.get_config('pixdim')

        print '\t* Acquisition scheme...'
        self.set_config('scheme_filename', scheme_filename)
        self.set_config('b0_thr', b0_thr)
        self.scheme = amico.scheme.Scheme( pjoin( self.get_config('DATA_path'), scheme_filename), b0_thr )
        print '\t\t- %d samples, %d shells' % ( self.scheme.nS, len(self.scheme.shells) )
        print '\t\t- %d @ b=0' % ( self.scheme.b0_count ),
        for i in xrange(len(self.scheme.shells)) :
            print ', %d @ b=%.1f' % ( len(self.scheme.shells[i]['idx']), self.scheme.shells[i]['b'] ),
        print

        if self.scheme.nS != self.niiDWI_img.shape[3] :
            raise ValueError( 'Scheme does not match with DWI data' )

        print '   [ %.1f seconds ]' % ( time.time() - tic )

        # Preprocessing
        tic = time.time()
        print '\n-> Preprocessing:'

        if self.get_config('doNormalizeSignal') :
            print '\t* Normalizing to b0...',
            sys.stdout.flush()
            mean = np.mean( self.niiDWI_img[:,:,:,self.scheme.b0_idx], axis=3 )
            idx = mean < 1
            mean[ idx ] = 1
            mean = 1 / mean
            mean[ idx ] = 0
            for i in xrange(self.scheme.nS) :
                self.niiDWI_img[:,:,:,i] *= mean
            print '[ min=%.2f,  mean=%.2f, max=%.2f ]' % ( self.niiDWI_img.min(), self.niiDWI_img.mean(), self.niiDWI_img.max() )

        if self.get_config('doMergeB0') :
            print '\t* Merging multiple b0 volume(s)...',
            mean = np.expand_dims( np.mean( self.niiDWI_img[:,:,:,self.scheme.b0_idx], axis=3 ), axis=3 )
            self.niiDWI_img = np.concatenate( (self.niiDWI_img[:,:,:,self.scheme.dwi_idx], mean), axis=3 )
        else :
            print '\t* Keeping all b0 volume(s)...',
        print '[ %d x %d x %d x %d ]' % self.niiDWI_img.shape

        if self.get_config('doDemean') :
            print '\t* Demeaning signal...',
            sys.stdout.flush()
            mean = np.repeat( np.expand_dims(np.mean(self.niiDWI_img,axis=3),axis=3), self.niiDWI_img.shape[3], axis=3 )
            self.niiDWI_img = self.niiDWI_img - mean
            print '[ min=%.2f,  mean=%.2f, max=%.2f ]' % ( self.niiDWI_img.min(), self.niiDWI_img.mean(), self.niiDWI_img.max() )

        print '   [ %.1f seconds ]' % ( time.time() - tic )


    def set_model( self, model_name ) :
        """Set the model to use to describe the signal contributions in each voxel.
<commit_msg>Fix normalization in case of synthetic data<commit_after>import numpy as np
import time
import glob
import sys
from os import makedirs, remove
from os.path import exists, join as pjoin
import nibabel
import cPickle
import commit.models
import commit.solvers
import amico.scheme
import amico.lut
import pyximport
pyximport.install( reload_support=True )


def setup( lmax = 12 ) :
    """General setup/initialization of the COMMIT framework."""
    amico.lut.precompute_rotation_matrices( lmax )


class Evaluation :
    """Class to hold all the information (data and parameters) when performing an
    evaluation with the COMMIT framework.
    """

    def __init__( self, study_path, subject ) :
        """Setup the data structures with default values.

        Parameters
        ----------
        study_path : string
            The path to the folder containing all the subjects from one study
        subject : string
            The path (relative to previous folder) to the subject folder
        """
        self.niiDWI     = None # set by "load_data" method
        self.scheme     = None # set by "load_data" method
        self.model      = None # set by "set_model" method
        self.KERNELS    = None # set by "load_kernels" method
        self.DICTIONARY = None # set by "load_dictionary" method
        self.THREADS    = None # set by "set_threads" method
        self.A          = None # set by "build_operator" method
        self.x          = None # set by "fit" method

        # store all the parameters of an evaluation with COMMIT
        self.CONFIG = {}
        self.set_config('study_path', study_path)
        self.set_config('subject', subject)
        self.set_config('DATA_path', pjoin( study_path, subject ))

        self.set_config('doNormalizeSignal', True)
        self.set_config('doMergeB0', True)
        self.set_config('doNormalizeKernels', True)
        self.set_config('doDemean', False)


    def set_config( self, key, value ) :
        self.CONFIG[ key ] = value

    def get_config( self, key ) :
        return self.CONFIG.get( key )


    def load_data( self, dwi_filename = 'DWI.nii', scheme_filename = 'DWI.scheme', b0_thr = 0 ) :
        """Load the diffusion signal and its corresponding acquisition scheme.

        Parameters
        ----------
        dwi_filename : string
            The file name of the DWI data, relative to the subject folder (default : 'DWI.nii')
        scheme_filename : string
            The file name of the corresponding acquisition scheme (default : 'DWI.scheme')
        b0_thr : float
            The threshold below which a b-value is considered a b0 (default : 0)
        """

        # Loading data and acquisition scheme
        tic = time.time()
        print '\n-> Loading data:'

        print '\t* DWI signal...'
        self.set_config('dwi_filename', dwi_filename)
        self.niiDWI  = nibabel.load( pjoin( self.get_config('DATA_path'), dwi_filename) )
        self.niiDWI_img = self.niiDWI.get_data().astype(np.float32)
        hdr = self.niiDWI.header if nibabel.__version__ >= '2.0.0' else self.niiDWI.get_header()
        self.set_config('dim', self.niiDWI_img.shape[0:3])
        self.set_config('pixdim', tuple( hdr.get_zooms()[:3] ))
        print '\t\t- dim    = %d x %d x %d x %d' % self.niiDWI_img.shape
        print '\t\t- pixdim = %.3f x %.3f x %.3f' % self.get_config('pixdim')

        print '\t* Acquisition scheme...'
        self.set_config('scheme_filename', scheme_filename)
        self.set_config('b0_thr', b0_thr)
        self.scheme = amico.scheme.Scheme( pjoin( self.get_config('DATA_path'), scheme_filename), b0_thr )
        print '\t\t- %d samples, %d shells' % ( self.scheme.nS, len(self.scheme.shells) )
        print '\t\t- %d @ b=0' % ( self.scheme.b0_count ),
        for i in xrange(len(self.scheme.shells)) :
            print ', %d @ b=%.1f' % ( len(self.scheme.shells[i]['idx']), self.scheme.shells[i]['b'] ),
        print

        if self.scheme.nS != self.niiDWI_img.shape[3] :
            raise ValueError( 'Scheme does not match with DWI data' )

        print '   [ %.1f seconds ]' % ( time.time() - tic )

        # Preprocessing
        tic = time.time()
        print '\n-> Preprocessing:'

        if self.get_config('doNormalizeSignal') :
            print '\t* Normalizing to b0...',
            sys.stdout.flush()
            mean = np.mean( self.niiDWI_img[:,:,:,self.scheme.b0_idx], axis=3 )
            idx = mean <= 0
            mean[ idx ] = 1
            mean = 1 / mean
            mean[ idx ] = 0
            for i in xrange(self.scheme.nS) :
                self.niiDWI_img[:,:,:,i] *= mean
            print '[ min=%.2f,  mean=%.2f, max=%.2f ]' % ( self.niiDWI_img.min(), self.niiDWI_img.mean(), self.niiDWI_img.max() )

        if self.get_config('doMergeB0') :
            print '\t* Merging multiple b0 volume(s)...',
            mean = np.expand_dims( np.mean( self.niiDWI_img[:,:,:,self.scheme.b0_idx], axis=3 ), axis=3 )
            self.niiDWI_img = np.concatenate( (self.niiDWI_img[:,:,:,self.scheme.dwi_idx], mean), axis=3 )
        else :
            print '\t* Keeping all b0 volume(s)...',
        print '[ %d x %d x %d x %d ]' % self.niiDWI_img.shape

        if self.get_config('doDemean') :
            print '\t* Demeaning signal...',
            sys.stdout.flush()
            mean = np.repeat( np.expand_dims(np.mean(self.niiDWI_img,axis=3),axis=3), self.niiDWI_img.shape[3], axis=3 )
            self.niiDWI_img = self.niiDWI_img - mean
            print '[ min=%.2f,  mean=%.2f, max=%.2f ]' % ( self.niiDWI_img.min(), self.niiDWI_img.mean(), self.niiDWI_img.max() )

        print '   [ %.1f seconds ]' % ( time.time() - tic )


    def set_model( self, model_name ) :
        """Set the model to use to describe the signal contributions in each voxel.

{'fraction_non_alphanumeric': 0.12962803882827936, 'fraction_numerical': 0.008933940383128598, 'mean_word_length': 2.263807120829829, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1722331', 'n_tokens_mistral': 3842, 'n_tokens_neox': 3740, 'n_words': 986}