<commit_before>#!/usr/bin/env python
from __future__ import print_function, division
import argparse
import subprocess
import os
import re
import shutil
import multiprocessing
from datetime import datetime

import logging

def make_arg_parser():
    parser = argparse.ArgumentParser(description='This is the commandline interface for NINJA-SHI7',
                                     usage='python ninja_shi7.py -i <input> -o <output> -t_trim <threads>...')

    # parser.add_argument('-adaptor', '--adaptor_type', help='Set the type of the adaptor (default: None)', choices=[None, 'Nextera', 'TruSeq3', 'TruSeq2'], default=None)
    parser.add_argument('-SE', help='Run in Single End mode (default: Disabled)', dest='single_end', action='store_true')
    parser.add_argument('--axe_adaptors', help='Path to the adaptor file.', default=None)
    parser.add_argument('--flash', help='Disable FLASH stiching (default: Enabled)', dest='flash', action='store_false')
    parser.add_argument('--no_trim', help='Disable the TRIMMER (default: Enabled)', dest='trim', action='store_false')
    parser.add_argument('--no_allow_outies', help='Disable "outie" orientation (default: Enabled)', dest='allow_outies', action='store_false')
    parser.add_argument('--no_convert_fasta', help='Disable convert FASTQS to FASTA (default: Enabled)', dest='convert_fasta', action='store_false')
    parser.add_argument('--no_combine_fasta', help='Disable the FASTA append mode (default: Enabled)', dest='combine_fasta', action='store_false')
    parser.add_argument('--shell', help='Use shell in Python system calls, NOT RECOMMENDED (default: Disabled)', dest='shell', action='store_true')
    parser.add_argument('-i', '--input', help='Set the directory path of the fastq directory', required=True)
    parser.add_argument('-o', '--output', help='Set the directory path of the output (default: cwd)', default=os.getcwd())
    parser.add_argument('-t', '--threads', help='Set the number of threads (default: %(default)s)',
                        default=multiprocessing.cpu_count())
    parser.add_argument('-m', '--min_overlap',
                        help='Set the minimum overlap length between two reads (default: %(default)s)', default=20, type=int)
    parser.add_argument('-M', '--max_overlap',
                        help='Set the maximum overlap length between two reads (default: %(default)s)', default=700, type=int)
    parser.add_argument('-trim_l', '--trim_length', help='Set the trim length (default: %(default)s)', default=150, type=int)
    parser.add_argument('-trim_q', '--trim_qual', help='Set the trim qual (default: %(default)s)', default=20, type=int)
    parser.set_defaults(flash=True, trim=True, allow_outies=True, convert_fasta=True, combine_fasta=True, shell=False, single_end=False)
    return parser

def run_command(cmd, shell=False):
    """
    Run prepared behave command in shell and return its output.
    :param cmd: Well-formed behave command to run.
    :param shell: Force subprocess to use shell, not recommended
    :return: Command output as string.
    """

    try:
        cmd = [str(i) for i in cmd]
        logging.debug(' '.join(cmd))
        output = subprocess.check_output(
            cmd,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            shell=shell,
            cwd=os.getcwd(),
        )

    except subprocess.CalledProcessError as e:
        output = e.output

    return output

def format_basename(filename):
    return '.'.join(re.sub('[^0-9a-zA-Z]+', '.', re.sub('_L001', '', re.sub('_001', '', os.path.basename(filename)))).split('.')[:-1])

def whitelist(dir, whitelist):
    for root, subdirs, files in os.walk(dir):
            for file in files:
                if os.path.join(root, file) not in whitelist:
                    os.remove(os.path.join(root, file))

def read_fastq(fh):
    line = next(fh)
    while line:
        title = line[1:].strip()
        data = ''
        qualities = ''
        flag = True
        line = next(fh)
        while line and (flag or len(data) != len(qualities)):
            if line[0] == '+':
                flag = False
            elif flag:
                data += line.strip()
            else:
                qualities += line.strip()
            line = next(fh)
        yield title, data, qualities

def split_fwd_rev(paths):
    path_R1_fastqs = [f for f in paths if 'R1' in os.path.basename(f)]
    path_R2_fastqs = [f for f in paths if 'R2' in os.path.basename(f)]
    if len(path_R1_fastqs) != len(path_R2_fastqs) or len(path_R1_fastqs) < 1:
        raise ValueError('Error: The input directory %s must contain at least one pair of R1 & R2 fastq file!' % os.path.dirname(paths[0]))
    return path_R1_fastqs, path_R2_fastqs

def axe_adaptors_single_end(input_fastqs, output_path, adapters, threads=1, shell=False):
    output_filenames = []
    for fastq in input_fastqs:
        output_fastq = os.path.join(output_path, format_basename(fastq) + '.fastq')
        trim_cmd = ['trimmomatic', 'SE', fastq, output_fastq, 'ILLUMINACLIP:%s:2:30:10:2:true' % adapters, '-threads', threads]
        logging.info(run_command(trim_cmd, shell=shell))
        output_filenames.append(output_fastq)
    return output_filenames

def axe_adaptors_paired_end(input_fastqs, output_path, adapters, threads=1, shell=False):
    path_R1_fastqs, path_R2_fastqs = split_fwd_rev(input_fastqs)
    output_filenames = []
    for input_path_R1, input_path_R2 in zip(path_R1_fastqs, path_R2_fastqs):
        output_path_R1 = os.path.join(output_path, format_basename(input_path_R1) + '.fastq')
        unpaired_R1 = os.path.join(output_path, 'unpaired.%s' % os.path.basename(input_path_R1))
        output_path_R2 = os.path.join(output_path, format_basename(input_path_R2) + '.fastq')
        unpaired_R2 = os.path.join(output_path, 'unpaired.%s' % os.path.basename(input_path_R1))
        trim_cmd = ['trimmomatic', 'PE', input_path_R1, input_path_R2, output_path_R1, unpaired_R1,  output_path_R2, unpaired_R2, 'ILLUMINACLIP:%s:2:30:10:2:true' % adapters, '-threads', threads]
        logging.info(run_command(trim_cmd, shell=shell))
        output_filenames.append(output_path_R1)
        output_filenames.append(output_path_R2)
    return output_filenames

def flash(input_fastqs, output_path, max_overlap, min_overlap, allow_outies, threads=1, shell=False):
    path_R1_fastqs, path_R2_fastqs = split_fwd_rev(input_fastqs)
    for input_path_R1, input_path_R2 in zip(path_R1_fastqs, path_R2_fastqs):
        flash_cmd = ['flash', input_path_R1, input_path_R2, '-o', format_basename(re.sub('R1', '', input_path_R1)), '-d', output_path, '-M', max_overlap, '-m', min_overlap, '-t', threads]
        if allow_outies:
            flash_cmd.append('-O')
        logging.info(run_command(flash_cmd, shell=shell))
    return [os.path.join(output_path, f) for f in os.listdir(output_path) if f.endswith('extendedFrags.fastq')]

def trimmer(input_fastqs, output_path, trim_length, trim_qual, threads=1, shell=False):
    [logging.info(filename) for filename in input_fastqs]
    output_filenames = []
    for path_input_fastq in input_fastqs:
        path_output_fastq = os.path.join(output_path, re.sub('.extendedFrags', '', format_basename(path_input_fastq)) + '.trimmed.fastq')
        ninja_shi7_cmd = ['ninja_shi7_trimmer', path_input_fastq, path_output_fastq, trim_length, trim_qual, 'FLOOR', 5, 'ASS_QUALITY', 30]
        logging.info(run_command(ninja_shi7_cmd, shell=shell))
        output_filenames.append(path_output_fastq)
    return output_filenames

def convert_fastqs(input_fastqs, output_path):
    output_filenames = []
    for path_input_fastq in input_fastqs:
        with open(path_input_fastq) as inf_fastq:
            gen_fastq = read_fastq(inf_fastq)
            output_filename = os.path.join(output_path, re.sub('.trimmed', '', format_basename(path_input_fastq)) + '.fna')
            with open(output_filename, 'w') as outf_fasta:
                for title, seq, quals in gen_fastq:
                    outf_fasta.write('>%s\n%s\n' % (title, seq))
        output_filenames.append(output_filename)
    return output_filenames

def convert_combine_fastqs(input_fastqs, output_path):
    output_filename = os.path.join(output_path, 'combined_seqs.fna')
    with open(output_filename, 'w') as outf_fasta:
        for path_input_fastq in input_fastqs:
                basename = re.sub('.trimmed', '', format_basename(path_input_fastq))
                with open(path_input_fastq) as inf_fastq:
                    gen_fastq = read_fastq(inf_fastq)
                    for i, (title, seq, quals) in enumerate(gen_fastq):
                        outf_fasta.write('>%s_%i %s\n%s\n' % (basename, i, title, seq))
    return [output_filename]

def main():
    # gcc -m64 -O3 ninja_shi7.c -o ninja_shi7
    # gcc -m64 -O3 -msse4.1 ninja_shi7.c -o ninja_shi7
    start_time = datetime.now()

    parser = make_arg_parser()
    args = parser.parse_args()

    logging.basicConfig(filename=os.path.join(args.output, 'ninja_shi7.log'), filemode='w', level=logging.DEBUG, format='%(asctime)s %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p')

    # FIRST CHECK IF THE INPUT AND OUTPUT PATH EXIST. IF DO NOT, RAISE EXCEPTION AND EXIT
    if not os.path.exists(args.input):
        raise ValueError('Error: Input directory %s doesn\'t exist!' % args.input)

    if not os.path.exists(args.output):
        os.makedirs(args.output)

    if os.path.exists(os.path.join(args.output, 'temp')):
        shutil.rmtree(os.path.join(args.output, 'temp'))
        logging.info('Existing temp directory deleted.')
        os.makedirs(os.path.join(args.output, 'temp'))
    else:
        os.makedirs(os.path.join(args.output, 'temp'))

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # AXE_ADAPTORS

    path_fastqs = [os.path.join(args.input, f) for f in os.listdir(args.input) if f.endswith('fastq')]
    logging.debug(path_fastqs)

    if args.axe_adaptors:
        if args.single_end:
            path_fastqs = axe_adaptors_single_end(path_fastqs, os.path.join(args.output, 'temp'), args.axe_adaptors, threads=args.threads, shell=args.shell)
        else:
            path_fastqs = axe_adaptors_paired_end(path_fastqs, os.path.join(args.output, 'temp'), args.axe_adaptors, threads=args.threads, shell=args.shell)
        whitelist(os.path.join(args.output, 'temp'), path_fastqs)
        logging.info('AXE_ADAPTORS done!')

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # FLASH
    if args.flash:
        if not args.single_end:
            path_fastqs = flash(path_fastqs, os.path.join(args.output, 'temp'), args.max_overlap, args.min_overlap, args.allow_outies, threads=args.threads, shell=args.shell)
            whitelist(os.path.join(args.output, 'temp'), path_fastqs)
            logging.info('FLASH done!')
        else:
            logging.warning('Single End mode enabled with FLASH. Skipping this step.')

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # CREATE_TRIMMER_GENERAL

    if args.trim:
        path_fastqs = trimmer(path_fastqs, os.path.join(args.output, 'temp'), args.trim_length, args.trim_qual, threads=args.threads, shell=args.shell)
        whitelist(os.path.join(args.output, 'temp'), path_fastqs)
        logging.info('CREATE_TRIMMER_GENERAL done!')

    # ">[SAMPLENAME]_[INDX starting at 0] HEADER"
    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # CONVERT FASTA TO FASTQ
    if args.convert_fasta:
        if args.combine_fasta:
            path_fastqs = convert_combine_fastqs(path_fastqs, os.path.join(args.output, 'temp'))
            logging.info('Convert FASTQs to FASTAs done!')
            logging.info('Combine FASTAs done!')
        else:
            path_fastqs = convert_fastqs(path_fastqs, os.path.join(args.output, 'temp'))
            logging.info('Convert FASTQs to FASTAs done!')

    for file in path_fastqs:
        dest = os.path.join(args.output, os.path.basename(file))
        if os.path.exists(dest):
            os.remove(dest)
        shutil.move(file, args.output)
    shutil.rmtree(os.path.join(args.output, 'temp'))
    logging.info('Execution time: %s' % (datetime.now() - start_time))

if __name__ == '__main__':
    main()
<commit_msg>Ever step creates its own folder.<commit_after>#!/usr/bin/env python
from __future__ import print_function, division
import argparse
import subprocess
import os
import re
import shutil
import multiprocessing
from datetime import datetime

import logging

def make_arg_parser():
    parser = argparse.ArgumentParser(description='This is the commandline interface for NINJA-SHI7',
                                     usage='python ninja_shi7.py -i <input> -o <output> -t_trim <threads>...')

    # parser.add_argument('-adaptor', '--adaptor_type', help='Set the type of the adaptor (default: None)', choices=[None, 'Nextera', 'TruSeq3', 'TruSeq2'], default=None)
    parser.add_argument('-SE', help='Run in Single End mode (default: Disabled)', dest='single_end', action='store_true')
    parser.add_argument('--axe_adaptors', help='Path to the adaptor file.', default=None)
    parser.add_argument('--flash', help='Disable FLASH stiching (default: Enabled)', dest='flash', action='store_false')
    parser.add_argument('--no_trim', help='Disable the TRIMMER (default: Enabled)', dest='trim', action='store_false')
    parser.add_argument('--no_allow_outies', help='Disable "outie" orientation (default: Enabled)', dest='allow_outies', action='store_false')
    parser.add_argument('--no_convert_fasta', help='Disable convert FASTQS to FASTA (default: Enabled)', dest='convert_fasta', action='store_false')
    parser.add_argument('--no_combine_fasta', help='Disable the FASTA append mode (default: Enabled)', dest='combine_fasta', action='store_false')
    parser.add_argument('--shell', help='Use shell in Python system calls, NOT RECOMMENDED (default: Disabled)', dest='shell', action='store_true')
    parser.add_argument('-i', '--input', help='Set the directory path of the fastq directory', required=True)
    parser.add_argument('-o', '--output', help='Set the directory path of the output (default: cwd)', default=os.getcwd())
    parser.add_argument('-t', '--threads', help='Set the number of threads (default: %(default)s)',
                        default=multiprocessing.cpu_count())
    parser.add_argument('-m', '--min_overlap',
                        help='Set the minimum overlap length between two reads (default: %(default)s)', default=20, type=int)
    parser.add_argument('-M', '--max_overlap',
                        help='Set the maximum overlap length between two reads (default: %(default)s)', default=700, type=int)
    parser.add_argument('-trim_l', '--trim_length', help='Set the trim length (default: %(default)s)', default=150, type=int)
    parser.add_argument('-trim_q', '--trim_qual', help='Set the trim qual (default: %(default)s)', default=20, type=int)
    parser.set_defaults(flash=True, trim=True, allow_outies=True, convert_fasta=True, combine_fasta=True, shell=False, single_end=False)
    return parser

def run_command(cmd, shell=False):
    """
    Run prepared behave command in shell and return its output.
    :param cmd: Well-formed behave command to run.
    :param shell: Force subprocess to use shell, not recommended
    :return: Command output as string.
    """

    try:
        cmd = [str(i) for i in cmd]
        logging.debug(' '.join(cmd))
        output = subprocess.check_output(
            cmd,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            shell=shell,
            cwd=os.getcwd(),
        )

    except subprocess.CalledProcessError as e:
        output = e.output

    return output

def format_basename(filename):
    return '.'.join(re.sub('[^0-9a-zA-Z]+', '.', re.sub('_L001', '', re.sub('_001', '', os.path.basename(filename)))).split('.')[:-1])

def whitelist(dir, whitelist):
    for root, subdirs, files in os.walk(dir):
            for file in files:
                if os.path.join(root, file) not in whitelist:
                    os.remove(os.path.join(root, file))

def read_fastq(fh):
    line = next(fh)
    while line:
        title = line[1:].strip()
        data = ''
        qualities = ''
        flag = True
        line = next(fh)
        while line and (flag or len(data) != len(qualities)):
            if line[0] == '+':
                flag = False
            elif flag:
                data += line.strip()
            else:
                qualities += line.strip()
            line = next(fh)
        yield title, data, qualities

def split_fwd_rev(paths):
    path_R1_fastqs = [f for f in paths if 'R1' in os.path.basename(f)]
    path_R2_fastqs = [f for f in paths if 'R2' in os.path.basename(f)]
    if len(path_R1_fastqs) != len(path_R2_fastqs) or len(path_R1_fastqs) < 1:
        raise ValueError('Error: The input directory %s must contain at least one pair of R1 & R2 fastq file!' % os.path.dirname(paths[0]))
    return path_R1_fastqs, path_R2_fastqs

def axe_adaptors_single_end(input_fastqs, output_path, adapters, threads=1, shell=False):
    output_filenames = []
    for fastq in input_fastqs:
        output_fastq = os.path.join(output_path, format_basename(fastq) + '.fastq')
        trim_cmd = ['trimmomatic', 'SE', fastq, output_fastq, 'ILLUMINACLIP:%s:2:30:10:2:true' % adapters, '-threads', threads]
        logging.info(run_command(trim_cmd, shell=shell))
        output_filenames.append(output_fastq)
    return output_filenames

def axe_adaptors_paired_end(input_fastqs, output_path, adapters, threads=1, shell=False):
    path_R1_fastqs, path_R2_fastqs = split_fwd_rev(input_fastqs)
    output_filenames = []
    for input_path_R1, input_path_R2 in zip(path_R1_fastqs, path_R2_fastqs):
        output_path_R1 = os.path.join(output_path, format_basename(input_path_R1) + '.fastq')
        unpaired_R1 = os.path.join(output_path, 'unpaired.%s' % os.path.basename(input_path_R1))
        output_path_R2 = os.path.join(output_path, format_basename(input_path_R2) + '.fastq')
        unpaired_R2 = os.path.join(output_path, 'unpaired.%s' % os.path.basename(input_path_R1))
        trim_cmd = ['trimmomatic', 'PE', input_path_R1, input_path_R2, output_path_R1, unpaired_R1,  output_path_R2, unpaired_R2, 'ILLUMINACLIP:%s:2:30:10:2:true' % adapters, '-threads', threads]
        logging.info(run_command(trim_cmd, shell=shell))
        output_filenames.append(output_path_R1)
        output_filenames.append(output_path_R2)
    return output_filenames

def flash(input_fastqs, output_path, max_overlap, min_overlap, allow_outies, threads=1, shell=False):
    path_R1_fastqs, path_R2_fastqs = split_fwd_rev(input_fastqs)
    for input_path_R1, input_path_R2 in zip(path_R1_fastqs, path_R2_fastqs):
        flash_cmd = ['flash', input_path_R1, input_path_R2, '-o', format_basename(re.sub('R1', '', input_path_R1)), '-d', output_path, '-M', max_overlap, '-m', min_overlap, '-t', threads]
        if allow_outies:
            flash_cmd.append('-O')
        logging.info(run_command(flash_cmd, shell=shell))
    return [re.sub('.extendedFrags', '', sos.path.join(output_path, f)) for f in os.listdir(output_path) if f.endswith('extendedFrags.fastq')]

def trimmer(input_fastqs, output_path, trim_length, trim_qual, threads=1, shell=False):
    [logging.info(filename) for filename in input_fastqs]
    output_filenames = []
    for path_input_fastq in input_fastqs:
        path_output_fastq = os.path.join(output_path, format_basename(path_input_fastq) + '.fastq')
        ninja_shi7_cmd = ['ninja_shi7_trimmer', path_input_fastq, path_output_fastq, trim_length, trim_qual, 'FLOOR', 5, 'ASS_QUALITY', 30]
        logging.info(run_command(ninja_shi7_cmd, shell=shell))
        output_filenames.append(path_output_fastq)
    return output_filenames

def convert_fastqs(input_fastqs, output_path):
    output_filenames = []
    for path_input_fastq in input_fastqs:
        with open(path_input_fastq) as inf_fastq:
            gen_fastq = read_fastq(inf_fastq)
            output_filename = os.path.join(output_path, re.sub('.trimmed', '', format_basename(path_input_fastq)) + '.fna')
            with open(output_filename, 'w') as outf_fasta:
                for title, seq, quals in gen_fastq:
                    outf_fasta.write('>%s\n%s\n' % (title, seq))
        output_filenames.append(output_filename)
    return output_filenames

def convert_combine_fastqs(input_fastqs, output_path):
    output_filename = os.path.join(output_path, 'combined_seqs.fna')
    with open(output_filename, 'w') as outf_fasta:
        for path_input_fastq in input_fastqs:
                basename = format_basename(path_input_fastq)
                with open(path_input_fastq) as inf_fastq:
                    gen_fastq = read_fastq(inf_fastq)
                    for i, (title, seq, quals) in enumerate(gen_fastq):
                        outf_fasta.write('>%s_%i %s\n%s\n' % (basename, i, title, seq))
    return [output_filename]

def main():
    # gcc -m64 -O3 ninja_shi7.c -o ninja_shi7
    # gcc -m64 -O3 -msse4.1 ninja_shi7.c -o ninja_shi7
    start_time = datetime.now()

    parser = make_arg_parser()
    args = parser.parse_args()

    logging.basicConfig(filename=os.path.join(args.output, 'ninja_shi7.log'), filemode='w', level=logging.DEBUG, format='%(asctime)s %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p')

    # FIRST CHECK IF THE INPUT AND OUTPUT PATH EXIST. IF DO NOT, RAISE EXCEPTION AND EXIT
    if not os.path.exists(args.input):
        raise ValueError('Error: Input directory %s doesn\'t exist!' % args.input)

    if not os.path.exists(args.output):
        os.makedirs(args.output)

    if os.path.exists(os.path.join(args.output, 'temp')):
        shutil.rmtree(os.path.join(args.output, 'temp'))
        logging.info('Existing temp directory deleted.')
        os.makedirs(os.path.join(args.output, 'temp'))
    else:
        os.makedirs(os.path.join(args.output, 'temp'))

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # AXE_ADAPTORS

    path_fastqs = [os.path.join(args.input, f) for f in os.listdir(args.input) if f.endswith('fastq')]
    logging.debug(path_fastqs)

    if args.axe_adaptors:
        axe_output = os.path.join(args.output, 'temp', 'axe')
        os.makedirs(axe_output)
        if args.single_end:
            path_fastqs = axe_adaptors_single_end(path_fastqs, axe_output, args.axe_adaptors, threads=args.threads, shell=args.shell)
        else:
            path_fastqs = axe_adaptors_paired_end(path_fastqs, axe_output, args.axe_adaptors, threads=args.threads, shell=args.shell)
        whitelist(os.path.join(args.output, 'temp'), path_fastqs)
        logging.info('AXE_ADAPTORS done!')

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # FLASH
    if args.flash:
        if not args.single_end:
            flash_output = os.path.join(args.output, 'temp', 'flash')
            os.makedirs(flash_output)
            path_fastqs = flash(path_fastqs, flash_output, args.max_overlap, args.min_overlap, args.allow_outies, threads=args.threads, shell=args.shell)
            whitelist(os.path.join(args.output, 'temp'), path_fastqs)
            logging.info('FLASH done!')
        else:
            logging.warning('Single End mode enabled with FLASH. Skipping this step.')

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # CREATE_TRIMMER_GENERAL

    if args.trim:
        trimmer_output = os.path.join(args.output, 'temp', 'trimmer')
        os.makedirs(trimmer_output)
        path_fastqs = trimmer(path_fastqs, trimmer_output, args.trim_length, args.trim_qual, threads=args.threads, shell=args.shell)
        whitelist(os.path.join(args.output, 'temp'), path_fastqs)
        logging.info('CREATE_TRIMMER_GENERAL done!')

    # ">[SAMPLENAME]_[INDX starting at 0] HEADER"
    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # CONVERT FASTA TO FASTQ
    if args.convert_fasta:
        convert_output = os.path.join(args.output, 'temp', 'convert')
        os.makedirs(convert_output)
        if args.combine_fasta:
            path_fastqs = convert_combine_fastqs(path_fastqs, convert_output)
            logging.info('Convert FASTQs to FASTAs done!')
            logging.info('Combine FASTAs done!')
        else:
            path_fastqs = convert_fastqs(path_fastqs, convert_output)
            logging.info('Convert FASTQs to FASTAs done!')

    for file in path_fastqs:
        dest = os.path.join(args.output, os.path.basename(file))
        if os.path.exists(dest):
            os.remove(dest)
        shutil.move(file, args.output)
    shutil.rmtree(os.path.join(args.output, 'temp'))
    logging.info('Execution time: %s' % (datetime.now() - start_time))

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import print_function, division
import argparse
import subprocess
import os
import re
import shutil
import multiprocessing
from datetime import datetime

import logging

def make_arg_parser():
    parser = argparse.ArgumentParser(description='This is the commandline interface for NINJA-SHI7',
                                     usage='python ninja_shi7.py -i <input> -o <output> -t_trim <threads>...')

    # parser.add_argument('-adaptor', '--adaptor_type', help='Set the type of the adaptor (default: None)', choices=[None, 'Nextera', 'TruSeq3', 'TruSeq2'], default=None)
    parser.add_argument('-SE', help='Run in Single End mode (default: Disabled)', dest='single_end', action='store_true')
    parser.add_argument('--axe_adaptors', help='Path to the adaptor file.', default=None)
    parser.add_argument('--flash', help='Disable FLASH stiching (default: Enabled)', dest='flash', action='store_false')
    parser.add_argument('--no_trim', help='Disable the TRIMMER (default: Enabled)', dest='trim', action='store_false')
    parser.add_argument('--no_allow_outies', help='Disable "outie" orientation (default: Enabled)', dest='allow_outies', action='store_false')
    parser.add_argument('--no_convert_fasta', help='Disable convert FASTQS to FASTA (default: Enabled)', dest='convert_fasta', action='store_false')
    parser.add_argument('--no_combine_fasta', help='Disable the FASTA append mode (default: Enabled)', dest='combine_fasta', action='store_false')
    parser.add_argument('--shell', help='Use shell in Python system calls, NOT RECOMMENDED (default: Disabled)', dest='shell', action='store_true')
    parser.add_argument('-i', '--input', help='Set the directory path of the fastq directory', required=True)
    parser.add_argument('-o', '--output', help='Set the directory path of the output (default: cwd)', default=os.getcwd())
    parser.add_argument('-t', '--threads', help='Set the number of threads (default: %(default)s)',
                        default=multiprocessing.cpu_count())
    parser.add_argument('-m', '--min_overlap',
                        help='Set the minimum overlap length between two reads (default: %(default)s)', default=20, type=int)
    parser.add_argument('-M', '--max_overlap',
                        help='Set the maximum overlap length between two reads (default: %(default)s)', default=700, type=int)
    parser.add_argument('-trim_l', '--trim_length', help='Set the trim length (default: %(default)s)', default=150, type=int)
    parser.add_argument('-trim_q', '--trim_qual', help='Set the trim qual (default: %(default)s)', default=20, type=int)
    parser.set_defaults(flash=True, trim=True, allow_outies=True, convert_fasta=True, combine_fasta=True, shell=False, single_end=False)
    return parser

def run_command(cmd, shell=False):
    """
    Run prepared behave command in shell and return its output.
    :param cmd: Well-formed behave command to run.
    :param shell: Force subprocess to use shell, not recommended
    :return: Command output as string.
    """

    try:
        cmd = [str(i) for i in cmd]
        logging.debug(' '.join(cmd))
        output = subprocess.check_output(
            cmd,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            shell=shell,
            cwd=os.getcwd(),
        )

    except subprocess.CalledProcessError as e:
        output = e.output

    return output

def format_basename(filename):
    return '.'.join(re.sub('[^0-9a-zA-Z]+', '.', re.sub('_L001', '', re.sub('_001', '', os.path.basename(filename)))).split('.')[:-1])

def whitelist(dir, whitelist):
    for root, subdirs, files in os.walk(dir):
            for file in files:
                if os.path.join(root, file) not in whitelist:
                    os.remove(os.path.join(root, file))

def read_fastq(fh):
    line = next(fh)
    while line:
        title = line[1:].strip()
        data = ''
        qualities = ''
        flag = True
        line = next(fh)
        while line and (flag or len(data) != len(qualities)):
            if line[0] == '+':
                flag = False
            elif flag:
                data += line.strip()
            else:
                qualities += line.strip()
            line = next(fh)
        yield title, data, qualities

def split_fwd_rev(paths):
    path_R1_fastqs = [f for f in paths if 'R1' in os.path.basename(f)]
    path_R2_fastqs = [f for f in paths if 'R2' in os.path.basename(f)]
    if len(path_R1_fastqs) != len(path_R2_fastqs) or len(path_R1_fastqs) < 1:
        raise ValueError('Error: The input directory %s must contain at least one pair of R1 & R2 fastq file!' % os.path.dirname(paths[0]))
    return path_R1_fastqs, path_R2_fastqs

def axe_adaptors_single_end(input_fastqs, output_path, adapters, threads=1, shell=False):
    output_filenames = []
    for fastq in input_fastqs:
        output_fastq = os.path.join(output_path, format_basename(fastq) + '.fastq')
        trim_cmd = ['trimmomatic', 'SE', fastq, output_fastq, 'ILLUMINACLIP:%s:2:30:10:2:true' % adapters, '-threads', threads]
        logging.info(run_command(trim_cmd, shell=shell))
        output_filenames.append(output_path)
    return output_filenames

def axe_adaptors_paired_end(input_fastqs, output_path, adapters, threads=1, shell=False):
    path_R1_fastqs, path_R2_fastqs = split_fwd_rev(input_fastqs)
    output_filenames = []
    for input_path_R1, input_path_R2 in zip(path_R1_fastqs, path_R2_fastqs):
        output_path_R1 = os.path.join(output_path, format_basename(input_path_R1) + '.fastq')
        unpaired_R1 = os.path.join(output_path, 'unpaired.%s' % os.path.basename(input_path_R1))
        output_path_R2 = os.path.join(output_path, format_basename(input_path_R2) + '.fastq')
        unpaired_R2 = os.path.join(output_path, 'unpaired.%s' % os.path.basename(input_path_R1))
        trim_cmd = ['trimmomatic', 'PE', input_path_R1, input_path_R2, output_path_R1, unpaired_R1,  output_path_R2, unpaired_R2, 'ILLUMINACLIP:%s:2:30:10:2:true' % adapters, '-threads', threads]
        logging.info(run_command(trim_cmd, shell=shell))
        output_filenames.append(output_path_R1)
        output_filenames.append(output_path_R2)
    return output_filenames

def flash(input_fastqs, output_path, max_overlap, min_overlap, allow_outies, threads=1, shell=False):
    path_R1_fastqs, path_R2_fastqs = split_fwd_rev(input_fastqs)
    for input_path_R1, input_path_R2 in zip(path_R1_fastqs, path_R2_fastqs):
        flash_cmd = ['flash', input_path_R1, input_path_R2, '-o', format_basename(re.sub('R1', '', input_path_R1)), '-d', output_path, '-M', max_overlap, '-m', min_overlap, '-t', threads]
        if allow_outies:
            flash_cmd.append('-O')
        logging.info(run_command(flash_cmd, shell=shell))
    return [os.path.join(output_path, f) for f in os.listdir(output_path) if f.endswith('extendedFrags.fastq')]

def trimmer(input_fastqs, output_path, trim_length, trim_qual, threads=1, shell=False):
    [logging.info(filename) for filename in input_fastqs]
    output_filenames = []
    for path_input_fastq in input_fastqs:
        path_output_fastq = os.path.join(output_path, re.sub('.extendedFrags', '', format_basename(path_input_fastq)) + '.fastq')
        ninja_shi7_cmd = ['ninja_shi7_trimmer', path_input_fastq, path_output_fastq, trim_length, trim_qual, 'FLOOR', 5, 'ASS_QUALITY', 30]
        logging.info(run_command(ninja_shi7_cmd, shell=shell))
        output_filenames.append(path_output_fastq)
    return output_filenames

def convert_fastqs(input_fastqs, output_path):
    output_filenames = []
    for path_input_fastq in input_fastqs:
        with open(path_input_fastq) as inf_fastq:
            gen_fastq = read_fastq(inf_fastq)
            output_filename = os.path.join(output_path, format_basename(path_input_fastq) + '.fna')
            with open(output_filename, 'w') as outf_fasta:
                for title, seq, quals in gen_fastq:
                    outf_fasta.write('>%s\n%s\n' % (title, seq))
        output_filenames.append(output_filename)
    return output_filenames

def convert_combine_fastqs(input_fastqs, output_path, basenames):
    output_filename = os.path.join(output_path, 'combined_seqs.fna')
    with open(output_filename, 'w') as outf_fasta:
        for path_input_fastq, basename in zip(input_fastqs, basenames):
                with open(path_input_fastq) as inf_fastq:
                    gen_fastq = read_fastq(inf_fastq)
                    for i, (title, seq, quals) in enumerate(gen_fastq):
                        outf_fasta.write('>%s_%i %s\n%s\n' % (basename, i, title, seq))
    return [output_filename]

def main():
    # gcc -m64 -O3 ninja_shi7.c -o ninja_shi7
    # gcc -m64 -O3 -msse4.1 ninja_shi7.c -o ninja_shi7
    start_time = datetime.now()

    parser = make_arg_parser()
    args = parser.parse_args()

    logging.basicConfig(filename=os.path.join(args.output, 'ninja_shi7.log'), filemode='w', level=logging.DEBUG, format='%(asctime)s %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p')

    # FIRST CHECK IF THE INPUT AND OUTPUT PATH EXIST. IF DO NOT, RAISE EXCEPTION AND EXIT
    if not os.path.exists(args.input):
        raise ValueError('Error: Input directory %s doesn\'t exist!' % args.input)

    if not os.path.exists(args.output):
        os.makedirs(args.output)

    if os.path.exists(os.path.join(args.output, 'temp')):
        shutil.rmtree(os.path.join(args.output, 'temp'))
        logging.info('Existing temp directory deleted.')
        os.makedirs(os.path.join(args.output, 'temp'))
    else:
        os.makedirs(os.path.join(args.output, 'temp'))

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # AXE_ADAPTORS

    path_fastqs = [os.path.join(args.input, f) for f in os.listdir(args.input) if f.endswith('fastq')]
    logging.debug(path_fastqs)

    if args.axe_adaptors:
        if args.single_end:
            path_fastqs = axe_adaptors_single_end(path_fastqs, os.path.join(args.output, 'temp'), args.axe_adaptors, threads=args.threads, shell=args.shell)
        else:
            path_fastqs = axe_adaptors_paired_end(path_fastqs, os.path.join(args.output, 'temp'), args.axe_adaptors, threads=args.threads, shell=args.shell)
        whitelist(os.path.join(args.output, 'temp'), path_fastqs)
        logging.info('AXE_ADAPTORS done!')

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # FLASH
    if args.flash:
        if not args.single_end:
            path_fastqs = flash(path_fastqs, os.path.join(args.output, 'temp'), args.max_overlap, args.min_overlap, args.allow_outies, threads=args.threads, shell=args.shell)
            whitelist(os.path.join(args.output, 'temp'), path_fastqs)
            logging.info('FLASH done!')
        else:
            logging.warning('Single End mode enabled with FLASH. Skipping this step.')

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # CREATE_TRIMMER_GENERAL

    if args.trim:
        path_fastqs = trimmer(path_fastqs, os.path.join(args.output, 'temp'), args.trim_length, args.trim_qual, threads=args.threads, shell=args.shell)
        whitelist(os.path.join(args.output, 'temp'), path_fastqs)
        logging.info('CREATE_TRIMMER_GENERAL done!')

    # ">[SAMPLENAME]_[INDX starting at 0] HEADER"
    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # CONVERT FASTA TO FASTQ
    if args.convert_fasta:
        if args.combine_fasta:
            basenames = [format_basename(filename) for filename in path_fastqs]
            path_fastqs = convert_combine_fastqs(path_fastqs, os.path.join(args.output, 'temp'), basenames=basenames)
            logging.info('Convert FASTQs to FASTAs done!')
            logging.info('Combine FASTAs done!')
        else:
            path_fastqs = convert_fastqs(path_fastqs, os.path.join(args.output, 'temp'))
            logging.info('Convert FASTQs to FASTAs done!')

    for file in path_fastqs:
        dest = os.path.join(args.output, os.path.basename(file))
        if os.path.exists(dest):
            os.remove(dest)
        shutil.move(file, args.output)
    shutil.rmtree(os.path.join(args.output, 'temp'))
    logging.info('Execution time: %s' % (datetime.now() - start_time))

if __name__ == '__main__':
    main()
<commit_msg>Updated the filename.<commit_after>#!/usr/bin/env python
from __future__ import print_function, division
import argparse
import subprocess
import os
import re
import shutil
import multiprocessing
from datetime import datetime

import logging

def make_arg_parser():
    parser = argparse.ArgumentParser(description='This is the commandline interface for NINJA-SHI7',
                                     usage='python ninja_shi7.py -i <input> -o <output> -t_trim <threads>...')

    # parser.add_argument('-adaptor', '--adaptor_type', help='Set the type of the adaptor (default: None)', choices=[None, 'Nextera', 'TruSeq3', 'TruSeq2'], default=None)
    parser.add_argument('-SE', help='Run in Single End mode (default: Disabled)', dest='single_end', action='store_true')
    parser.add_argument('--axe_adaptors', help='Path to the adaptor file.', default=None)
    parser.add_argument('--flash', help='Disable FLASH stiching (default: Enabled)', dest='flash', action='store_false')
    parser.add_argument('--no_trim', help='Disable the TRIMMER (default: Enabled)', dest='trim', action='store_false')
    parser.add_argument('--no_allow_outies', help='Disable "outie" orientation (default: Enabled)', dest='allow_outies', action='store_false')
    parser.add_argument('--no_convert_fasta', help='Disable convert FASTQS to FASTA (default: Enabled)', dest='convert_fasta', action='store_false')
    parser.add_argument('--no_combine_fasta', help='Disable the FASTA append mode (default: Enabled)', dest='combine_fasta', action='store_false')
    parser.add_argument('--shell', help='Use shell in Python system calls, NOT RECOMMENDED (default: Disabled)', dest='shell', action='store_true')
    parser.add_argument('-i', '--input', help='Set the directory path of the fastq directory', required=True)
    parser.add_argument('-o', '--output', help='Set the directory path of the output (default: cwd)', default=os.getcwd())
    parser.add_argument('-t', '--threads', help='Set the number of threads (default: %(default)s)',
                        default=multiprocessing.cpu_count())
    parser.add_argument('-m', '--min_overlap',
                        help='Set the minimum overlap length between two reads (default: %(default)s)', default=20, type=int)
    parser.add_argument('-M', '--max_overlap',
                        help='Set the maximum overlap length between two reads (default: %(default)s)', default=700, type=int)
    parser.add_argument('-trim_l', '--trim_length', help='Set the trim length (default: %(default)s)', default=150, type=int)
    parser.add_argument('-trim_q', '--trim_qual', help='Set the trim qual (default: %(default)s)', default=20, type=int)
    parser.set_defaults(flash=True, trim=True, allow_outies=True, convert_fasta=True, combine_fasta=True, shell=False, single_end=False)
    return parser

def run_command(cmd, shell=False):
    """
    Run prepared behave command in shell and return its output.
    :param cmd: Well-formed behave command to run.
    :param shell: Force subprocess to use shell, not recommended
    :return: Command output as string.
    """

    try:
        cmd = [str(i) for i in cmd]
        logging.debug(' '.join(cmd))
        output = subprocess.check_output(
            cmd,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            shell=shell,
            cwd=os.getcwd(),
        )

    except subprocess.CalledProcessError as e:
        output = e.output

    return output

def format_basename(filename):
    return '.'.join(re.sub('[^0-9a-zA-Z]+', '.', re.sub('_L001', '', re.sub('_001', '', os.path.basename(filename)))).split('.')[:-1])

def whitelist(dir, whitelist):
    for root, subdirs, files in os.walk(dir):
            for file in files:
                if os.path.join(root, file) not in whitelist:
                    os.remove(os.path.join(root, file))

def read_fastq(fh):
    line = next(fh)
    while line:
        title = line[1:].strip()
        data = ''
        qualities = ''
        flag = True
        line = next(fh)
        while line and (flag or len(data) != len(qualities)):
            if line[0] == '+':
                flag = False
            elif flag:
                data += line.strip()
            else:
                qualities += line.strip()
            line = next(fh)
        yield title, data, qualities

def split_fwd_rev(paths):
    path_R1_fastqs = [f for f in paths if 'R1' in os.path.basename(f)]
    path_R2_fastqs = [f for f in paths if 'R2' in os.path.basename(f)]
    if len(path_R1_fastqs) != len(path_R2_fastqs) or len(path_R1_fastqs) < 1:
        raise ValueError('Error: The input directory %s must contain at least one pair of R1 & R2 fastq file!' % os.path.dirname(paths[0]))
    return path_R1_fastqs, path_R2_fastqs

def axe_adaptors_single_end(input_fastqs, output_path, adapters, threads=1, shell=False):
    output_filenames = []
    for fastq in input_fastqs:
        output_fastq = os.path.join(output_path, format_basename(fastq) + '.fastq')
        trim_cmd = ['trimmomatic', 'SE', fastq, output_fastq, 'ILLUMINACLIP:%s:2:30:10:2:true' % adapters, '-threads', threads]
        logging.info(run_command(trim_cmd, shell=shell))
        output_filenames.append(output_fastq)
    return output_filenames

def axe_adaptors_paired_end(input_fastqs, output_path, adapters, threads=1, shell=False):
    path_R1_fastqs, path_R2_fastqs = split_fwd_rev(input_fastqs)
    output_filenames = []
    for input_path_R1, input_path_R2 in zip(path_R1_fastqs, path_R2_fastqs):
        output_path_R1 = os.path.join(output_path, format_basename(input_path_R1) + '.fastq')
        unpaired_R1 = os.path.join(output_path, 'unpaired.%s' % os.path.basename(input_path_R1))
        output_path_R2 = os.path.join(output_path, format_basename(input_path_R2) + '.fastq')
        unpaired_R2 = os.path.join(output_path, 'unpaired.%s' % os.path.basename(input_path_R1))
        trim_cmd = ['trimmomatic', 'PE', input_path_R1, input_path_R2, output_path_R1, unpaired_R1,  output_path_R2, unpaired_R2, 'ILLUMINACLIP:%s:2:30:10:2:true' % adapters, '-threads', threads]
        logging.info(run_command(trim_cmd, shell=shell))
        output_filenames.append(output_path_R1)
        output_filenames.append(output_path_R2)
    return output_filenames

def flash(input_fastqs, output_path, max_overlap, min_overlap, allow_outies, threads=1, shell=False):
    path_R1_fastqs, path_R2_fastqs = split_fwd_rev(input_fastqs)
    for input_path_R1, input_path_R2 in zip(path_R1_fastqs, path_R2_fastqs):
        flash_cmd = ['flash', input_path_R1, input_path_R2, '-o', format_basename(re.sub('R1', '', input_path_R1)), '-d', output_path, '-M', max_overlap, '-m', min_overlap, '-t', threads]
        if allow_outies:
            flash_cmd.append('-O')
        logging.info(run_command(flash_cmd, shell=shell))
    return [os.path.join(output_path, f) for f in os.listdir(output_path) if f.endswith('extendedFrags.fastq')]

def trimmer(input_fastqs, output_path, trim_length, trim_qual, threads=1, shell=False):
    [logging.info(filename) for filename in input_fastqs]
    output_filenames = []
    for path_input_fastq in input_fastqs:
        path_output_fastq = os.path.join(output_path, re.sub('.extendedFrags', '', format_basename(path_input_fastq)) + '.fastq')
        ninja_shi7_cmd = ['ninja_shi7_trimmer', path_input_fastq, path_output_fastq, trim_length, trim_qual, 'FLOOR', 5, 'ASS_QUALITY', 30]
        logging.info(run_command(ninja_shi7_cmd, shell=shell))
        output_filenames.append(path_output_fastq)
    return output_filenames

def convert_fastqs(input_fastqs, output_path):
    output_filenames = []
    for path_input_fastq in input_fastqs:
        with open(path_input_fastq) as inf_fastq:
            gen_fastq = read_fastq(inf_fastq)
            output_filename = os.path.join(output_path, format_basename(path_input_fastq) + '.fna')
            with open(output_filename, 'w') as outf_fasta:
                for title, seq, quals in gen_fastq:
                    outf_fasta.write('>%s\n%s\n' % (title, seq))
        output_filenames.append(output_filename)
    return output_filenames

def convert_combine_fastqs(input_fastqs, output_path, basenames):
    output_filename = os.path.join(output_path, 'combined_seqs.fna')
    with open(output_filename, 'w') as outf_fasta:
        for path_input_fastq, basename in zip(input_fastqs, basenames):
                with open(path_input_fastq) as inf_fastq:
                    gen_fastq = read_fastq(inf_fastq)
                    for i, (title, seq, quals) in enumerate(gen_fastq):
                        outf_fasta.write('>%s_%i %s\n%s\n' % (basename, i, title, seq))
    return [output_filename]

def main():
    # gcc -m64 -O3 ninja_shi7.c -o ninja_shi7
    # gcc -m64 -O3 -msse4.1 ninja_shi7.c -o ninja_shi7
    start_time = datetime.now()

    parser = make_arg_parser()
    args = parser.parse_args()

    logging.basicConfig(filename=os.path.join(args.output, 'ninja_shi7.log'), filemode='w', level=logging.DEBUG, format='%(asctime)s %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p')

    # FIRST CHECK IF THE INPUT AND OUTPUT PATH EXIST. IF DO NOT, RAISE EXCEPTION AND EXIT
    if not os.path.exists(args.input):
        raise ValueError('Error: Input directory %s doesn\'t exist!' % args.input)

    if not os.path.exists(args.output):
        os.makedirs(args.output)

    if os.path.exists(os.path.join(args.output, 'temp')):
        shutil.rmtree(os.path.join(args.output, 'temp'))
        logging.info('Existing temp directory deleted.')
        os.makedirs(os.path.join(args.output, 'temp'))
    else:
        os.makedirs(os.path.join(args.output, 'temp'))

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # AXE_ADAPTORS

    path_fastqs = [os.path.join(args.input, f) for f in os.listdir(args.input) if f.endswith('fastq')]
    logging.debug(path_fastqs)

    if args.axe_adaptors:
        if args.single_end:
            path_fastqs = axe_adaptors_single_end(path_fastqs, os.path.join(args.output, 'temp'), args.axe_adaptors, threads=args.threads, shell=args.shell)
        else:
            path_fastqs = axe_adaptors_paired_end(path_fastqs, os.path.join(args.output, 'temp'), args.axe_adaptors, threads=args.threads, shell=args.shell)
        whitelist(os.path.join(args.output, 'temp'), path_fastqs)
        logging.info('AXE_ADAPTORS done!')

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # FLASH
    if args.flash:
        if not args.single_end:
            path_fastqs = flash(path_fastqs, os.path.join(args.output, 'temp'), args.max_overlap, args.min_overlap, args.allow_outies, threads=args.threads, shell=args.shell)
            whitelist(os.path.join(args.output, 'temp'), path_fastqs)
            logging.info('FLASH done!')
        else:
            logging.warning('Single End mode enabled with FLASH. Skipping this step.')

    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # CREATE_TRIMMER_GENERAL

    if args.trim:
        path_fastqs = trimmer(path_fastqs, os.path.join(args.output, 'temp'), args.trim_length, args.trim_qual, threads=args.threads, shell=args.shell)
        whitelist(os.path.join(args.output, 'temp'), path_fastqs)
        logging.info('CREATE_TRIMMER_GENERAL done!')

    # ">[SAMPLENAME]_[INDX starting at 0] HEADER"
    # -----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
    # CONVERT FASTA TO FASTQ
    if args.convert_fasta:
        if args.combine_fasta:
            basenames = [format_basename(filename) for filename in path_fastqs]
            path_fastqs = convert_combine_fastqs(path_fastqs, os.path.join(args.output, 'temp'), basenames=basenames)
            logging.info('Convert FASTQs to FASTAs done!')
            logging.info('Combine FASTAs done!')
        else:
            path_fastqs = convert_fastqs(path_fastqs, os.path.join(args.output, 'temp'))
            logging.info('Convert FASTQs to FASTAs done!')

    for file in path_fastqs:
        dest = os.path.join(args.output, os.path.basename(file))
        if os.path.exists(dest):
            os.remove(dest)
        shutil.move(file, args.output)
    shutil.rmtree(os.path.join(args.output, 'temp'))
    logging.info('Execution time: %s' % (datetime.now() - start_time))

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>from itertools import groupby

from .numbers import dec, inc

def partition_by(f, coll):
    for _, group in groupby(coll, key=f):
        yield list(group)

def is_arithmetic(items):
    return len(items) == 2 \
        or len(items) >= 3 and items[-1] - items[-2] == items[-2] - items[-3]

def is_geometric(items):
    return len(items) >= 3 and items[-1] / items[-2] == items[-2] / items[-3]

def guess_succ(items, end):
    if is_geometric(items):
        koef = items[-1] / items[-2]
        return lambda x: x * koef
    elif is_arithmetic(items):
        step = items[-1] - items[-2]
        return lambda x: x + step
    else:
        return inc if end >= items[-1] else dec

def parse_key(key):
    is_ellipsis = lambda x: x is Ellipsis
    items, _, end = partition_by(is_ellipsis, key)
    assert len(items) >= 1
    assert len(end) == 1
    end = end[0]

    if callable(items[-1]) and len(items) > 1 and not callable(items[-2]):
        succ = items.pop()
    else:
        succ = guess_succ(items, end)

    assert len(items) >= 1
    return items, succ, end

def get_arity(succ, items):
    try:
        return int(succ.__code__.co_argcount)
    except (AttributeError, ValueError, TypeError):
        for n in range(len(items)):
            try:
                succ(*items[n:])
                return len(items) - n
            except TypeError:
                pass
        else:
            succ(*items) # reraise

class Seq(object):
    def __getitem__(self, key):
        if isinstance(key, int):
            return range(1, key + 1)
        else:
            items, succ, end = parse_key(key)

            arity = get_arity(succ, items)
            while True:
                x = succ(*items[-arity:])
                if items[-1] < end < x or items[-1] > end > x:
                    break
                items.append(x)
                if x == end:
                    break

            return items

seq = Seq()

from operator import __add__
from whatever import _

def test_seq():
    assert seq[1, ..., 3] == [1,2,3]
    assert seq[1, 3, ..., 9] == [1, 3, 5, 7, 9]
    assert seq[1, 2, 4, ..., 10] == [1, 2, 4, 8]
    assert seq[1, _ + 3, ..., 10] == [1, 4, 7, 10]
    assert seq[1, 1, _ + _, ..., 10] == [1, 1, 2, 3, 5, 8]
    assert seq[1, 1, __add__, ..., 10] == [1, 1, 2, 3, 5, 8]

    assert seq[3, ..., 1] == [3, 2, 1]
    assert seq[1, -2, 4, ..., 10] == [1, -2, 4, -8]
    assert seq[-1, 2, -4, ..., 10] == [-1, 2, -4, 8, -16]

    # seq[1, 10]
    # seq[10]
<commit_msg>Renamed Seq to SeqMaker<commit_after>from itertools import groupby

from .numbers import dec, inc

def partition_by(f, coll):
    for _, group in groupby(coll, key=f):
        yield list(group)

def is_arithmetic(items):
    return len(items) == 2 \
        or len(items) >= 3 and items[-1] - items[-2] == items[-2] - items[-3]

def is_geometric(items):
    return len(items) >= 3 and items[-1] / items[-2] == items[-2] / items[-3]

def guess_succ(items, end):
    if is_geometric(items):
        koef = items[-1] / items[-2]
        return lambda x: x * koef
    elif is_arithmetic(items):
        step = items[-1] - items[-2]
        return lambda x: x + step
    else:
        return inc if end >= items[-1] else dec

def parse_key(key):
    is_ellipsis = lambda x: x is Ellipsis
    items, _, end = partition_by(is_ellipsis, key)
    assert len(items) >= 1
    assert len(end) == 1
    end = end[0]

    if callable(items[-1]) and len(items) > 1 and not callable(items[-2]):
        succ = items.pop()
    else:
        succ = guess_succ(items, end)

    assert len(items) >= 1
    return items, succ, end

def get_arity(succ, items):
    try:
        return int(succ.__code__.co_argcount)
    except (AttributeError, ValueError, TypeError):
        for n in range(len(items)):
            try:
                succ(*items[n:])
                return len(items) - n
            except TypeError:
                pass
        else:
            succ(*items) # reraise

class SeqMaker(object):
    def __getitem__(self, key):
        if isinstance(key, int):
            return range(1, key + 1)
        else:
            items, succ, end = parse_key(key)

            arity = get_arity(succ, items)
            while True:
                x = succ(*items[-arity:])
                if items[-1] < end < x or items[-1] > end > x:
                    break
                items.append(x)
                if x == end:
                    break

            return items

seq = SeqMaker()

from operator import __add__
from whatever import _

def test_seq():
    assert seq[1, ..., 3] == [1,2,3]
    assert seq[1, 3, ..., 9] == [1, 3, 5, 7, 9]
    assert seq[1, 2, 4, ..., 10] == [1, 2, 4, 8]
    assert seq[1, _ + 3, ..., 10] == [1, 4, 7, 10]
    assert seq[1, 1, _ + _, ..., 10] == [1, 1, 2, 3, 5, 8]
    assert seq[1, 1, __add__, ..., 10] == [1, 1, 2, 3, 5, 8]

    assert seq[3, ..., 1] == [3, 2, 1]
    assert seq[1, -2, 4, ..., 10] == [1, -2, 4, -8]
    assert seq[-1, 2, -4, ..., 10] == [-1, 2, -4, 8, -16]

    # seq[1, 10]
    # seq[10]
<|endoftext|>


<commit_before>#
# furl - URL manipulation made simple.
#
# Arthur Grunseid
# grunseid.com
# grunseid@gmail.com
#
# License: Build Amazing Things (Unlicense)

__title__ = 'furl'
__version__ = '0.4.3'
__license__ = 'Unlicense'
__author__ = 'Arthur Grunseid'
__contact__ = 'grunseid@gmail.com'
__url__ = 'https://github.com/gruns/furl'

from .furl import *
<commit_msg>Furl v0.4.4.<commit_after>#
# furl - URL manipulation made simple.
#
# Arthur Grunseid
# grunseid.com
# grunseid@gmail.com
#
# License: Build Amazing Things (Unlicense)

__title__ = 'furl'
__version__ = '0.4.4'
__license__ = 'Unlicense'
__author__ = 'Arthur Grunseid'
__contact__ = 'grunseid@gmail.com'
__url__ = 'https://github.com/gruns/furl'

from .furl import *
<|endoftext|>


<commit_before># emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
from __future__ import with_statement
"""
This module implements fMRI Design Matrix creation.

The DesignMatrix object is just a container that represents the design matrix.
Computations of the different parts of the design matrix are confined
to the make_dmtx() function, that instantiates the DesignMatrix object.
All the remainder are just ancillary functions.

Design matrices contain three different types of regressors:

1. Task-related regressors, that result from the convolution
   of the experimental paradigm regressors with hemodynamic models
   A hemodynamic model is one of:
   'spm' : linear filter used in the SPM software
   'glover' : linear filter estimated by G.Glover
   'fir' : finite impulse response model, generic linear filter

2. User-specified regressors, that represent information available on
   the data, e.g. motion parameters, physiological data resampled at
   the acquisition rate, or sinusoidal regressors that model the
   signal at a frequency of interest.

3. Drift regressors, that represent low_frequency phenomena of no
   interest in the data; they need to be included to reduce variance
   estimates.

Author: Bertrand Thirion, 2009-2015
"""
import numpy as np
from warnings import warn

from .hemodynamic_models import compute_regressor, _orthogonalize
from .experimental_paradigm import check_paradigm
from pandas import DataFrame

######################################################################
# Ancillary functions
######################################################################

def _poly_drift(order, frame_times):
    """Create a polynomial drift matrix

    Parameters
    ----------
    order : int,
        number of polynomials in the drift model

    frame_times : array of shape(n_scans),
        time stamps used to sample polynomials

    Returns
    -------
    pol : ndarray, shape(n_scans, order + 1)
         estimated polynomial drifts plus a constant regressor
    """
    order = int(order)
    pol = np.zeros((np.size(frame_times), order + 1))
    tmax = float(frame_times.max())
    for k in range(order + 1):
        pol[:, k] = (frame_times / tmax) ** k
    pol = _orthogonalize(pol)
    pol = np.hstack((pol[:, 1:], pol[:, :1]))
    return pol

def _cosine_drift(period_cut, frame_times):
    """Create a cosine drift matrix with periods greater or equal to
    period_cut.

    Parameters
    ----------
    period_cut : float
         Cut period of the low-pass filter (in sec)

    frame_times : array of shape(n_scans)
         The sampling times (in sec)

    Returns
    -------
    cosine_drift:  array of shape(n_scans, n_drifts)
             cosine drifts plus a constant regressor at cosine_drift[:, -1]

    Ref: http://en.wikipedia.org/wiki/Discrete_cosine_transform DCT-II
    """
    n_frames = len(frame_times)
    n_times = np.arange(n_frames)
    hfcut = 1. / period_cut  # input parameter is the period
    dt = frame_times[1] - frame_times[0]
    order = int(np.floor(2 * n_frames * hfcut * dt))
    # s.t. hfcut = 1 / (2 * dt) yields n_frames
    cosine_drift = np.zeros((n_frames, order))
    normalizer = np.sqrt(2.0 / n_frames)

    for k in range(1, order):
        cosine_drift[:, k - 1] = normalizer * np.cos(
            (np.pi / n_frames) * (n_times + .5) * k)

    cosine_drift[:, order - 1] = 1.
    return cosine_drift

def _blank_drift(frame_times):
    """ Create the blank drift matrix

    Returns
    -------
    np.ones_like(frame_times)
    """
    return np.reshape(np.ones_like(frame_times), (np.size(frame_times), 1))

def _make_drift(drift_model, frame_times, order=1, period_cut=128.):
    """Create the drift matrix

    Parameters
    ----------
    drift_model : {'polynomial', 'cosine', 'blank'},
        string that specifies the desired drift model

    frame_times : array of shape(n_scans),
        list of values representing the desired TRs

    order : int, optional,
           order of the drift model (in case it is polynomial)

    period_cut : float, optional (defaults to 128),
           period cut in case of a cosine model (in seconds)

    Returns
    -------
    drift: array of shape(n_scans, n_drifts),
        the drift matrix

    names: list of length(n_drifts),
        the associated names
    """
    drift_model = drift_model.lower()   # for robust comparisons
    if drift_model == 'polynomial':
        drift = _poly_drift(order, frame_times)
    elif drift_model == 'cosine':
        drift = _cosine_drift(period_cut, frame_times)
    elif drift_model == 'blank':
        drift = _blank_drift(frame_times)
    else:
        raise NotImplementedError("Unknown drift model %r" % (drift_model))
    names = []
    for k in range(1, drift.shape[1]):
        names.append('drift_%d' % k)
    names.append('constant')
    return drift, names

def _convolve_regressors(paradigm, hrf_model, frame_times, fir_delays=[0],
                         min_onset=-24):
    """ Creation of  a matrix that comprises
    the convolution of the conditions onset with a certain hrf model

    Parameters
    ----------
    paradigm : DataFrame instance,
        Descriptor of an experimental paradigm
        see .experimental_paradigm to check the specification
        for these to be valid paradigm descriptors

    hrf_model : {'canonical', 'canonical with derivative', 'fir'}
        string that specifies the hemodynamic response function

    frame_times: array of shape(n_scans)
        the targeted timing for the design matrix

    fir_delays : array-like of shape(nb_onsets), optional,
        in case of FIR design, yields the array of delays
        used in the FIR model

    min_onset: float, optional (default: -24),
        minimal onset relative to frame_times[0] (in seconds)
        events that start before frame_times[0] + min_onset are not considered

    Returns
    -------
    regressor_matrix: array of shape(n_scans, n_regressors),
        contains the convolved regressors
        associated with the experimental conditions

    regressor_names: list of strings,
        the regressor names, that depend on the hrf model used
        if 'canonical' then this is identical to the input names
        if 'canonical with derivative', then two names are produced for
        input name 'name': 'name' and 'name_derivative'
    """
    regressor_names = []
    regressor_matrix = None
    if hrf_model == 'fir':
        oversampling = 1
    else:
        oversampling = 16

    name, onset, duration, modulation = check_paradigm(paradigm)
    for condition in np.unique(paradigm.name):
        condition_mask = (name == condition)
        exp_condition = (onset[condition_mask],
                         duration[condition_mask],
                         modulation[condition_mask])
        reg, names = compute_regressor(
            exp_condition, hrf_model, frame_times, con_id=condition,
            fir_delays=fir_delays, oversampling=oversampling,
            min_onset=min_onset)
        regressor_names += names
        if regressor_matrix is None:
            regressor_matrix = reg
        else:
            regressor_matrix = np.hstack((regressor_matrix, reg))
    return regressor_matrix, regressor_names

def _full_rank(X, cmax=1e15):
    """ Computes the condition number of X and if it is larger than cmax,
    returns a matrix with a condition number smaller than cmax.

    Parameters
    ----------
    X : array of shape(nrows, ncols)
        input array

    cmax : float, optional (default:1.e-15),
        tolerance for condition number

    Returns
    -------
    X : array of shape(nrows, ncols)
        output array

    cond : float,
        actual condition number
    """
    U, s, V = np.linalg.svd(X, 0)
    smax, smin = s.max(), s.min()
    cond = smax / smin
    if cond < cmax:
        return X, cond

    warn('Matrix is singular at working precision, regularizing...')
    lda = (smax - cmax * smin) / (cmax - 1)
    X = np.dot(U, np.dot(np.diag(s + lda), V))
    return X, cmax

######################################################################
# Design matrix creation
######################################################################

def make_design_matrix(
    frame_times, paradigm=None, hrf_model='canonical',
    drift_model='cosine', period_cut=128, drift_order=1, fir_delays=[0],
    add_regs=None, add_reg_names=None, min_onset=-24):
    """ Generate a design matrix from the input parameters

    Parameters
    ----------
    frame_times : array of shape(n_frames),
        the timing of the scans

    paradigm : DataFrame instance, optional
        description of the experimental paradigm

    hrf_model : string, optional,
        that specifies the hemodynamic response function
        it can be 'canonical', 'canonical with derivative' or 'fir'

    drift_model : string, optional
        specifies the desired drift model,
        to be chosen among 'polynomial', 'cosine', 'blank'

    period_cut : float, optional
        cut period of the low-pass filter in seconds

    drift_order : int, optional
        order of the drift model (in case it is polynomial)

    fir_delays : array of shape(n_onsets) or list, optional,
        in case of FIR design, yields the array of delays used in the FIR model

    add_regs : array of shape(n_frames, n_add_reg), optional
        additional user-supplied regressors

    add_reg_names : list of (n_add_reg) strings, optional
        if None, while n_add_reg > 0, these will be termed
        'reg_%i', i = 0..n_add_reg - 1

    min_onset : float, optional
        minimal onset relative to frame_times[0] (in seconds)
        events that start before frame_times[0] + min_onset are not considered

    Returns
    -------
    design_matrix: DataFrame instance,
        describing the design matrix
    """
    # check arguments
    # check that additional regressor specification is correct
    n_add_regs = 0
    if add_regs is not None:
        if add_regs.shape[0] == np.size(add_regs):
            add_regs = np.reshape(add_regs, (np.size(add_regs), 1))
        n_add_regs = add_regs.shape[1]
        assert add_regs.shape[0] == np.size(frame_times), ValueError(
            'incorrect specification of additional regressors: '
            'length of regressors provided: %s, number of ' +
            'time-frames: %s' % (add_regs.shape[0], np.size(frame_times)))

    # check that additional regressor names are well specified
    if add_reg_names is None:
        add_reg_names = ['reg%d' % k for k in range(n_add_regs)]
    elif len(add_reg_names) != n_add_regs:
        raise ValueError(
            'Incorrect number of additional regressor names was provided'
            '(%s provided, %s expected) % (len(add_reg_names),'
            'n_add_regs)')

    # computation of the matrix
    names = []
    matrix = np.zeros((frame_times.size, 0))

    # step 1: paradigm-related regressors
    if paradigm is not None:
        # create the condition-related regressors
        matrix, names = _convolve_regressors(
            paradigm, hrf_model.lower(), frame_times, fir_delays, min_onset)

    # step 2: additional regressors
    if add_regs is not None:
        # add user-supplied regressors and corresponding names
        matrix = np.hstack((matrix, add_regs))
        names += add_reg_names

    # step 3: drifts
    drift, dnames = _make_drift(drift_model.lower(), frame_times, drift_order,
                                period_cut)
    matrix = np.hstack((matrix, drift))
    names += dnames

    # step 4: Force the design matrix to be full rank at working precision
    matrix, _ = _full_rank(matrix)

    design_matrix = DataFrame(
        np.hstack((frame_times[:, np.newaxis], matrix)),
        columns=['frame_times'] + names)
    return design_matrix

def check_design_matrix(design_matrix):
    """ Check that the provided DataFrame is indeed a valid design matrix
    descriptor, and returns a triplet of fields

    Parameters
    ----------
    design matrix : pandas DataFrame,
        describes a design matrix

    Returns
    -------
    frame_times : array of shape (n_frames),
        sampling times of the design matrix

    matrix : array of shape (n_frames, n_regressors), dtype='f'
        numerical values for the design matrix

    names : array of shape (n_events), dtype='f'
           per-event onset time (in seconds)
    """
    names = design_matrix.keys()
    if 'frame_times' not in names:
        raise ValueError('The provided DataFrame does not contain the'
                         'mandatory frame_times field')
    frame_times = design_matrix['frame_times']
    names = [key for key in names if key != 'frame_times']
    matrix = np.array([design_matrix[key] for key in names]).T
    return frame_times, matrix, names

def plot_design_matrix(design_matrix, rescale=True, ax=None):
    """ plot a design matrix provided as a DataFrame

    Parameters
    ----------
    design matrix : pandas DataFrame,
        describes a design matrix

    rescale: bool, optional
        rescale columns magnitude for visualization or not

    ax: axis handle, optional
        Handle to axis onto which we will draw design matrix

    Returns
    -------
    ax: axis handle
    """
    # We import _set_mpl_backend because just the fact that we are
    # importing it sets the backend
    from nilearn.plotting import _set_mpl_backend
    # avoid unhappy pyflackes
    _set_mpl_backend
    import matplotlib.pyplot as plt

    # normalize the values per column for better visualization
    _, X, names = check_design_matrix(design_matrix)
    if rescale:
        X = X / np.maximum(1.e-12, np.sqrt(np.sum(X ** 2, 0)))
    if ax is None:
        plt.figure()
        ax = plt.subplot(1, 1, 1)

    ax.imshow(X, interpolation='Nearest', aspect='auto')
    ax.set_label('conditions')
    ax.set_ylabel('scan number')

    ax.set_xticks(range(len(names)))
    ax.set_xticklabels(names, rotation=60, ha='right')
    return ax

def design_matrix_light(
    frame_times, paradigm=None, hrf_model='canonical',
    drift_model='cosine', period_cut=128, drift_order=1, fir_delays=[0],
    add_regs=None, add_reg_names=None, min_onset=-24, path=None):
    """ Idem make_design_matrix, but only returns the computed matrix
    and associated names """
    dmtx = make_design_matrix(frame_times, paradigm, hrf_model,
    drift_model, period_cut, drift_order, fir_delays,
    add_regs, add_reg_names, min_onset)
    _, matrix, names = check_design_matrix(dmtx)
    return matrix, names
<commit_msg>oops. one late commit<commit_after># emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
from __future__ import with_statement
"""
This module implements fMRI Design Matrix creation.

The DesignMatrix object is just a container that represents the design matrix.
Computations of the different parts of the design matrix are confined
to the make_dmtx() function, that instantiates the DesignMatrix object.
All the remainder are just ancillary functions.

Design matrices contain three different types of regressors:

1. Task-related regressors, that result from the convolution
   of the experimental paradigm regressors with hemodynamic models
   A hemodynamic model is one of:
   'spm' : linear filter used in the SPM software
   'glover' : linear filter estimated by G.Glover
   'fir' : finite impulse response model, generic linear filter

2. User-specified regressors, that represent information available on
   the data, e.g. motion parameters, physiological data resampled at
   the acquisition rate, or sinusoidal regressors that model the
   signal at a frequency of interest.

3. Drift regressors, that represent low_frequency phenomena of no
   interest in the data; they need to be included to reduce variance
   estimates.

Author: Bertrand Thirion, 2009-2015
"""
import numpy as np
from warnings import warn

from .hemodynamic_models import compute_regressor, _orthogonalize
from .experimental_paradigm import check_paradigm
from pandas import DataFrame

######################################################################
# Ancillary functions
######################################################################

def _poly_drift(order, frame_times):
    """Create a polynomial drift matrix

    Parameters
    ----------
    order : int,
        number of polynomials in the drift model

    frame_times : array of shape(n_scans),
        time stamps used to sample polynomials

    Returns
    -------
    pol : ndarray, shape(n_scans, order + 1)
         estimated polynomial drifts plus a constant regressor
    """
    order = int(order)
    pol = np.zeros((np.size(frame_times), order + 1))
    tmax = float(frame_times.max())
    for k in range(order + 1):
        pol[:, k] = (frame_times / tmax) ** k
    pol = _orthogonalize(pol)
    pol = np.hstack((pol[:, 1:], pol[:, :1]))
    return pol

def _cosine_drift(period_cut, frame_times):
    """Create a cosine drift matrix with periods greater or equal to
    period_cut.

    Parameters
    ----------
    period_cut : float
         Cut period of the low-pass filter (in sec)

    frame_times : array of shape(n_scans)
         The sampling times (in sec)

    Returns
    -------
    cosine_drift:  array of shape(n_scans, n_drifts)
             cosine drifts plus a constant regressor at cosine_drift[:, -1]

    Ref: http://en.wikipedia.org/wiki/Discrete_cosine_transform DCT-II
    """
    n_frames = len(frame_times)
    n_times = np.arange(n_frames)
    hfcut = 1. / period_cut  # input parameter is the period
    dt = frame_times[1] - frame_times[0]
    order = int(np.floor(2 * n_frames * hfcut * dt))
    # s.t. hfcut = 1 / (2 * dt) yields n_frames
    cosine_drift = np.zeros((n_frames, order))
    normalizer = np.sqrt(2.0 / n_frames)

    for k in range(1, order):
        cosine_drift[:, k - 1] = normalizer * np.cos(
            (np.pi / n_frames) * (n_times + .5) * k)

    cosine_drift[:, order - 1] = 1.
    return cosine_drift

def _blank_drift(frame_times):
    """ Create the blank drift matrix

    Returns
    -------
    np.ones_like(frame_times)
    """
    return np.reshape(np.ones_like(frame_times), (np.size(frame_times), 1))

def _make_drift(drift_model, frame_times, order=1, period_cut=128.):
    """Create the drift matrix

    Parameters
    ----------
    drift_model : {'polynomial', 'cosine', 'blank'},
        string that specifies the desired drift model

    frame_times : array of shape(n_scans),
        list of values representing the desired TRs

    order : int, optional,
           order of the drift model (in case it is polynomial)

    period_cut : float, optional (defaults to 128),
           period cut in case of a cosine model (in seconds)

    Returns
    -------
    drift: array of shape(n_scans, n_drifts),
        the drift matrix

    names: list of length(n_drifts),
        the associated names
    """
    drift_model = drift_model.lower()   # for robust comparisons
    if drift_model == 'polynomial':
        drift = _poly_drift(order, frame_times)
    elif drift_model == 'cosine':
        drift = _cosine_drift(period_cut, frame_times)
    elif drift_model == 'blank':
        drift = _blank_drift(frame_times)
    else:
        raise NotImplementedError("Unknown drift model %r" % (drift_model))
    names = []
    for k in range(1, drift.shape[1]):
        names.append('drift_%d' % k)
    names.append('constant')
    return drift, names

def _convolve_regressors(paradigm, hrf_model, frame_times, fir_delays=[0],
                         min_onset=-24):
    """ Creation of  a matrix that comprises
    the convolution of the conditions onset with a certain hrf model

    Parameters
    ----------
    paradigm : DataFrame instance,
        Descriptor of an experimental paradigm
        see .experimental_paradigm to check the specification
        for these to be valid paradigm descriptors

    hrf_model : {'canonical', 'canonical with derivative', 'fir'}
        string that specifies the hemodynamic response function

    frame_times: array of shape(n_scans)
        the targeted timing for the design matrix

    fir_delays : array-like of shape(nb_onsets), optional,
        in case of FIR design, yields the array of delays
        used in the FIR model

    min_onset: float, optional (default: -24),
        minimal onset relative to frame_times[0] (in seconds)
        events that start before frame_times[0] + min_onset are not considered

    Returns
    -------
    regressor_matrix: array of shape(n_scans, n_regressors),
        contains the convolved regressors
        associated with the experimental conditions

    regressor_names: list of strings,
        the regressor names, that depend on the hrf model used
        if 'canonical' then this is identical to the input names
        if 'canonical with derivative', then two names are produced for
        input name 'name': 'name' and 'name_derivative'
    """
    regressor_names = []
    regressor_matrix = None
    if hrf_model == 'fir':
        oversampling = 1
    else:
        oversampling = 16

    name, onset, duration, modulation = check_paradigm(paradigm)
    for condition in np.unique(paradigm.name):
        condition_mask = (name == condition)
        exp_condition = (onset[condition_mask],
                         duration[condition_mask],
                         modulation[condition_mask])
        reg, names = compute_regressor(
            exp_condition, hrf_model, frame_times, con_id=condition,
            fir_delays=fir_delays, oversampling=oversampling,
            min_onset=min_onset)
        regressor_names += names
        if regressor_matrix is None:
            regressor_matrix = reg
        else:
            regressor_matrix = np.hstack((regressor_matrix, reg))
    return regressor_matrix, regressor_names

def _full_rank(X, cmax=1e15):
    """ Computes the condition number of X and if it is larger than cmax,
    returns a matrix with a condition number smaller than cmax.

    Parameters
    ----------
    X : array of shape(nrows, ncols)
        input array

    cmax : float, optional (default:1.e-15),
        tolerance for condition number

    Returns
    -------
    X : array of shape(nrows, ncols)
        output array

    cond : float,
        actual condition number
    """
    U, s, V = np.linalg.svd(X, 0)
    smax, smin = s.max(), s.min()
    cond = smax / smin
    if cond < cmax:
        return X, cond

    warn('Matrix is singular at working precision, regularizing...')
    lda = (smax - cmax * smin) / (cmax - 1)
    X = np.dot(U, np.dot(np.diag(s + lda), V))
    return X, cmax

######################################################################
# Design matrix creation
######################################################################

def make_design_matrix(
    frame_times, paradigm=None, hrf_model='canonical',
    drift_model='cosine', period_cut=128, drift_order=1, fir_delays=[0],
    add_regs=None, add_reg_names=None, min_onset=-24):
    """ Generate a design matrix from the input parameters

    Parameters
    ----------
    frame_times : array of shape(n_frames),
        the timing of the scans

    paradigm : DataFrame instance, optional
        description of the experimental paradigm

    hrf_model : string, optional,
        that specifies the hemodynamic response function
        it can be 'canonical', 'canonical with derivative' or 'fir'

    drift_model : string, optional
        specifies the desired drift model,
        to be chosen among 'polynomial', 'cosine', 'blank'

    period_cut : float, optional
        cut period of the low-pass filter in seconds

    drift_order : int, optional
        order of the drift model (in case it is polynomial)

    fir_delays : array of shape(n_onsets) or list, optional,
        in case of FIR design, yields the array of delays used in the FIR model

    add_regs : array of shape(n_frames, n_add_reg), optional
        additional user-supplied regressors

    add_reg_names : list of (n_add_reg) strings, optional
        if None, while n_add_reg > 0, these will be termed
        'reg_%i', i = 0..n_add_reg - 1

    min_onset : float, optional
        minimal onset relative to frame_times[0] (in seconds)
        events that start before frame_times[0] + min_onset are not considered

    Returns
    -------
    design_matrix: DataFrame instance,
        describing the design matrix
    """
    # check arguments
    # check that additional regressor specification is correct
    n_add_regs = 0
    if add_regs is not None:
        if add_regs.shape[0] == np.size(add_regs):
            add_regs = np.reshape(add_regs, (np.size(add_regs), 1))
        n_add_regs = add_regs.shape[1]
        assert add_regs.shape[0] == np.size(frame_times), ValueError(
            'incorrect specification of additional regressors: '
            'length of regressors provided: %s, number of ' +
            'time-frames: %s' % (add_regs.shape[0], np.size(frame_times)))

    # check that additional regressor names are well specified
    if add_reg_names is None:
        add_reg_names = ['reg%d' % k for k in range(n_add_regs)]
    elif len(add_reg_names) != n_add_regs:
        raise ValueError(
            'Incorrect number of additional regressor names was provided'
            '(%s provided, %s expected) % (len(add_reg_names),'
            'n_add_regs)')

    # computation of the matrix
    names = []
    matrix = np.zeros((frame_times.size, 0))

    # step 1: paradigm-related regressors
    if paradigm is not None:
        # create the condition-related regressors
        matrix, names = _convolve_regressors(
            paradigm, hrf_model.lower(), frame_times, fir_delays, min_onset)

    # step 2: additional regressors
    if add_regs is not None:
        # add user-supplied regressors and corresponding names
        matrix = np.hstack((matrix, add_regs))
        names += add_reg_names

    # step 3: drifts
    drift, dnames = _make_drift(drift_model.lower(), frame_times, drift_order,
                                period_cut)
    matrix = np.hstack((matrix, drift))
    names += dnames

    # step 4: Force the design matrix to be full rank at working precision
    matrix, _ = _full_rank(matrix)

    design_matrix = DataFrame(
        np.hstack((frame_times[:, np.newaxis], matrix)),
        columns=['frame_times'] + names)
    return design_matrix

def check_design_matrix(design_matrix):
    """ Check that the provided DataFrame is indeed a valid design matrix
    descriptor, and returns a triplet of fields

    Parameters
    ----------
    design matrix : pandas DataFrame,
        describes a design matrix

    Returns
    -------
    frame_times : array of shape (n_frames),
        sampling times of the design matrix

    matrix : array of shape (n_frames, n_regressors), dtype='f'
        numerical values for the design matrix

    names : array of shape (n_events), dtype='f'
           per-event onset time (in seconds)
    """
    names = design_matrix.keys()
    if 'frame_times' not in names:
        raise ValueError('The provided DataFrame does not contain the'
                         'mandatory frame_times field')
    frame_times = design_matrix['frame_times']
    names = [key for key in names if key != 'frame_times']
    matrix = np.array([design_matrix[key] for key in names]).T
    return frame_times, matrix, names

def plot_design_matrix(design_matrix, rescale=True, ax=None):
    """ plot a design matrix provided as a DataFrame

    Parameters
    ----------
    design matrix : pandas DataFrame,
        describes a design matrix

    rescale: bool, optional
        rescale columns magnitude for visualization or not

    ax: axis handle, optional
        Handle to axis onto which we will draw design matrix

    Returns
    -------
    ax: axis handle
    """
    # We import _set_mpl_backend because just the fact that we are
    # importing it sets the backend
    from nilearn.plotting import _set_mpl_backend
    # avoid unhappy pyflackes
    _set_mpl_backend
    import matplotlib.pyplot as plt

    # normalize the values per column for better visualization
    _, X, names = check_design_matrix(design_matrix)
    if rescale:
        X = X / np.maximum(1.e-12, np.sqrt(np.sum(X ** 2, 0)))
    if ax is None:
        plt.figure()
        ax = plt.subplot(1, 1, 1)

    ax.imshow(X, interpolation='Nearest', aspect='auto')
    ax.set_label('conditions')
    ax.set_ylabel('scan number')

    ax.set_xticks(range(len(names)))
    ax.set_xticklabels(names, rotation=60, ha='right')
    return ax
<|endoftext|>