fsl_mrs

#!/usr/bin/env python

# fsl_mrs - wrapper script for MRS fitting
#
# Author: Saad Jbabdi <saad@fmrib.ox.ac.uk>
#         William Carke <william.clarke@ndcn.ox.ac.uk>
#
# Copyright (C) 2019 University of Oxford 
# SHBASECOPYRIGHT

# Quick imports
#import argparse
import configargparse

from fsl_mrs import __version__
from fsl_mrs.utils.splash import splash

# NOTE!!!! THERE ARE MORE IMPORTS IN THE CODE BELOW (AFTER ARGPARSING)



def main():
    # Parse command-line arguments
    p = configargparse.ArgParser(add_config_file_help=False,description="FSL Magnetic Resonance Spectroscopy Wrapper Script")


    # utility for hiding certain arguments
    def hide_args(arglist):
        for action in arglist:
            action.help=p.SUPPRESS

        
    #p = argparse.ArgumentParser(description='FSL Magnetic Resonance Spectroscopy Tool')
    p.add_argument('-v','--version', action='version', version=__version__)
    
    required     = p.add_argument_group('required arguments')
    fitting_args = p.add_argument_group('fitting options')
    optional     = p.add_argument_group('additional options')

    # REQUIRED ARGUMENTS
    required.add_argument('--data',
                          required=True,type=str,metavar='<str>',
                          help='input FID file')
    required.add_argument('--basis',
                          required=True,type=str,metavar='<str>',
                          help='.BASIS file or folder containing basis spectra (will read all files within)')
    required.add_argument('--output',
                          required=True,type=str,metavar='<str>',
                          help='output folder')
    
    # FITTING ARGUMENTS
    fitting_args.add_argument('--algo',default='Newton',type=str,
                   help='algorithm [Newton (fast, default) or MH (slow)]')
    fitting_args.add_argument('--ignore',type=str,nargs='+',metavar='METAB',
                   help='ignore certain metabolites [repeatable]')
    fitting_args.add_argument('--keep',type=str,nargs='+',metavar='METAB',
                   help='only keep these metabolites')
    fitting_args.add_argument('--combine',type=str,nargs='+',action='append',metavar='METAB',
                   help='combine certain metabolites [repeatable]')
    fitting_args.add_argument('--ppmlim',default=(.2,4.2),type=float,nargs=2,metavar=('LOW','HIGH'),
                   help='limit the fit to a freq range (default=(.2,4.2))')
    fitting_args.add_argument('--h2o',default=None,type=str,metavar='H2O',
                   help='input .H2O file for quantification')
    fitting_args.add_argument('--baseline_order',default=2,type=int,metavar=('ORDER'),
                   help='order of baseline polynomial (default=2,-1 disables)')
    fitting_args.add_argument('--metab_groups',default=0,nargs='+',type=str_or_int_arg,
                              help="metabolite groups: list of groups or list of names for indept groups.")
    fitting_args.add_argument('--add_MM',type=bool,
                              help="include default macromolecule peaks")

    
    # ADDITONAL OPTIONAL ARGUMENTS
    optional.add_argument('--t1',type=str,default=None,metavar='IMAGE',
                          help='structural image (for report)')
    optional.add_argument('--TE',type=float,default=None,metavar='TE',
                          help='Echo time for relaxation correction (ms)')
    optional.add_argument('--tissue_fractions',type=float,nargs=3,default=[0.5,0.5,0.0],metavar='GM WM CSF',
                          help='Fractional tissue volumes for WM, GM, CSF. Defaults to 0.5, 0.5, 0.0.')
    optional.add_argument('--central_frequency',default=None,type=float,
                           help='central frequency in Hz')
    optional.add_argument('--dwell_time',default=None,type=float,
                           help='dwell time in seconds')
    optional.add_argument('--report',action="store_true",
                          help='output html report')
    optional.add_argument('--verbose',action="store_true",
                          help='spit out verbose info')
    optional.add_argument('--phase_correct',action="store_true",
                          help='do phase correction')
    optional.add_argument('--overwrite',action="store_true",
                          help='overwrite existing output folder')
    optional.add_argument('--conj_fid',dest='conjfid',action="store_true",help='Force conjugation of FID')
    optional.add_argument('--no_conj_fid',dest='conjfid',action="store_false",help='Forbid automatic conjugation of FID')
    optional.add_argument('--conj_basis',dest='conjbasis',action="store_true",help='Force conjugation of basis')
    optional.add_argument('--no_conj_basis',dest='conjbasis',action="store_false",help='Forbid automatic conjugation of basis')
    optional.set_defaults(conjfid=None,conjbasis=None)
    optional.add_argument('--no_rescale',action="store_false",help='Forbid rescaling of FID/basis/H2O.')
    optional.add('--config', required=False, is_config_file=True, help='configuration file')


    #hide_args([h1,h2])
    
    # Parse command-line arguments
    args = p.parse_args()

    # Output kickass splash screen
    if args.verbose:
        splash(logo='mrs')


    # ######################################################
    # DO THE IMPORTS AFTER PARSING TO SPEED UP HELP DISPLAY
    import time, os, sys, shutil, warnings
    import numpy as np
    from fsl_mrs.core import MRS
    from fsl_mrs.utils import mrs_io
    from fsl_mrs.utils import report
    from fsl_mrs.utils import fitting
    from fsl_mrs.utils import plotting
    from fsl_mrs.utils import misc
    import datetime
    import plotly
    # ######################################################
    if not args.verbose:
        warnings.filterwarnings("ignore")

    
    # Check if output folder exists
    overwrite = args.overwrite
    if os.path.exists(args.output):
        if not overwrite:
            print("Folder '{}' exists. Are you sure you want to delete it? [Y,N]".format(args.output))
            response = input()
            overwrite = response.upper() == "Y"
        if not overwrite:
            print('Early stopping...')
            exit()
        else:
            shutil.rmtree(args.output)
            os.makedirs(args.output,exist_ok=True)
    else:
        os.makedirs(args.output,exist_ok=True)


    # Save chosen arguments
    with open(os.path.join(args.output,"options.txt"),"w") as f:
        f.write(str(args))
        f.write("\n--------\n")
        f.write(p.format_values())
        
        
    #######  Do the work #######

    # Read data/h2o/basis
    if args.verbose:        
        print('--->> Read input data and basis\n')
        print('  {}'.format(args.data))
        print('  {}\n'.format(args.basis))


    FID,dataheader            = mrs_io.read_FID(args.data)
    basis, names, basisheader = mrs_io.read_basis(args.basis)
        
    if args.h2o is not None:
        H2O,_                 = mrs_io.read_FID(args.h2o)
    else:
        H2O = None
      
    # Collect useful info
    if args.central_frequency is not None:
        cf = args.central_frequency
    elif dataheader['centralFrequency'] is not None:
        cf = dataheader['centralFrequency']
        if args.verbose:
            print(f'      Detected central frequency in header info cf = {cf:0.6f} MHz')
    else:
        raise(Exception('Cannot determine central frequency. Please either set it or include it in data header'))

    if args.dwell_time is not None:
        bw = 1/args.dwell_time
    elif dataheader['bandwidth'] is not None:
        bw = dataheader['bandwidth']
        if args.verbose:
            print(f'      Detected bandwidth in header info bw = {bw:0.1f} Hz')
    else:
        raise(Exception('Cannot determine bandwidth. Please either set it or include it in data header'))

    # Fix case where basis file contains no header info (e.g. .RAW)
    if basisheader is None:
        basisheader = {'bandwidth':bw,'dwelltime':1/bw,'centralFrequency':cf}
    else:
        basisheader = basisheader[0]
        
    # Instantiate MRS object
    MRSargs = {'FID':FID,'basis':basis,'basis_hdr':basisheader,'names':names,'H2O':H2O,'cf':cf,'bw':bw}
    mrs     = MRS(**MRSargs)
    
    # Check the FID and basis / conjugate
    if args.conjfid is not None:
        if args.conjfid:
            mrs.conj_FID()
    else:
        conjugated = mrs.check_FID(repair=True)
        if args.verbose:
            if conjugated == 1:
                warnings.warn('FID has been checked and conjugated. Please check!',UserWarning)

    if args.conjbasis is not None:
        if args.conjbasis:
            mrs.conj_Basis()
    else:
        conjugated = mrs.check_Basis(repair=True)
        if args.verbose:
            if conjugated == 1:
                warnings.warn('Basis has been checked and conjugated. Please check!',UserWarning)
    
    # Rescale FID, H2O and basis to have nice range
    if not args.no_rescale:
        mrs.rescaleForFitting()

    # Do phase correction
    if args.phase_correct:
        if args.verbose:
            print('--->>  Phase correction\n')
        mrs.FID  = misc.phase_correct(mrs,mrs.FID)
        mrs.Spec = misc.FIDToSpec(mrs.FID)
              
        
    # Keep/Ignore metabolites
    mrs.keep(args.keep)
    mrs.ignore(args.ignore)    

    # Do the fitting here
    if args.verbose:
        print('--->> Start fitting\n\n')
        print('    Algorithm = [{}]\n'.format(args.algo))
        start = time.time()


    ppmlim=args.ppmlim
    if ppmlim is not None:
        ppmlim=(ppmlim[0],ppmlim[1])
 


    # Parse metabolite groups
    metab_groups = args.metab_groups
    if isinstance(metab_groups,list):
        if isinstance(metab_groups[0],str):
            tmp = [0]*mrs.numBasis
            grpcounter = 0
            for n in metab_groups:
                grpcounter += 1
                tmp[mrs.names.index(n)] = grpcounter
            metab_groups = tmp
        elif isinstance(metab_groups[0],int):
            if metab_groups == [0]:
                metab_groups = [0]*mrs.numBasis
            elif len(metab_groups) != mrs.numBasis:
                raise(Exception('Found {} metab_groups but there are {} basis functions'.format(len(metab_groups),mrs.numBasis)))
    elif metab_groups == 0:
        metab_groups = [0]*mrs.numBasis

    # Include Macromolecules? These should have their own metab groups
    if args.add_MM is not None:
        if not args.verbose:
            print('Adding macromolecules')
        nMM = mrs.add_MM_peaks()
        G   = [i+max(metab_groups)+1 for i in range(nMM)]
        metab_groups += G


    Fitargs = {'ppmlim':ppmlim,
               'method':args.algo,'baseline_order':args.baseline_order,
               'metab_groups':metab_groups}

    if args.verbose:
        print(mrs)
        print('Fitting args:')
        print(Fitargs)

    res = fitting.fit_FSLModel(mrs,**Fitargs)

    # Quantification
    # Echo time
    if args.TE is not None:
        echotime = args.TE*1E-3
    elif 'meta' in basisheader:
        if 'TE' in basisheader['meta']:
            echotime = basisheader['meta']['TE']
            if echotime > 1.0: # Assume in ms.
                echotime *= 1E-3
            else:
                echotime = None
    else:
        echotime = None
    # Internal and Water quantification if requested
    if (mrs.H2O is None) or (echotime is None):
        res.calculateConcScaling(mrs,referenceMetab=['Cr','PCr'])
    else:        
        tfrac = {'WM':args.tissue_fractions[0],'GM':args.tissue_fractions[1],'CSF':args.tissue_fractions[2]}
        res.calculateConcScaling(mrs,
                                referenceMetab=['Cr','PCr'],
                                waterRefFID=mrs.H2O,
                                tissueFractions=tfrac,
                                TE=echotime,
                                verbose=args.verbose)
    # Combine metabolites.
    if args.combine is not None:
        res.combine(args.combine)
    stop = time.time()

    
    # Report on the fitting
    if args.verbose:
        print('    Fitting lasted          : {:.3f} secs.\n'.format(stop-start))
    # Save output files
    if args.verbose:
        print('--->> Saving output files to {}\n'.format(args.output))


    res.to_file(filename=os.path.join(args.output,'results_table.csv'),what='concentrations')
    res.to_file(filename=os.path.join(args.output,'qc.csv'),what='qc')
    res.to_file(filename=os.path.join(args.output,'all_parameters.csv'),what='parameters')

    # Save image of MRS voxel
    if args.t1 is not None:
        datatype = mrs_io.check_datatype(args.data)
        if datatype == 'NIFTI':            
            fig = plotting.plot_world_orient(args.t1,args.data)
            fig.savefig(os.path.join(args.output,'voxel_location.png'))

    # Create short HTML report
    #fig = plotting.plotly_fit(mrs,res,ppmlim=ppmlim,proj='abs')
    #plotly.io.write_html(fig, file=os.path.join(args.output,'short_report.html'))
    

    # Save quick summary figure
    report.fitting_summary_fig(mrs,res,
                               filename=os.path.join(args.output,'fit_summary.png'))
    
    
    # Create interactive HTML report
    if args.report:
        report.create_report(mrs,res,
                             filename=os.path.join(args.output,'report.html'),
                             fidfile=args.data,
                             basisfile=args.basis,
                             h2ofile=args.h2o,
                             outdir=args.output,
                             date=datetime.datetime.now().strftime("%Y-%m-%d %H:%M"))
                         

     
    if args.verbose:
        print('\n\n\nDone.')

def str_or_int_arg(x):    
    try:                           
        return int(x)
    except:
        return x
        
if __name__ == '__main__':
    main()