<commit_before>#!/usr/bin/env python3
import pandas as pd
import os, numpy, argparse
from tuba_seq.fastq import fastqDF, MasterRead
from tuba_seq.shared import logPrint
from rpy2.robjects.packages import importr
from rpy2.robjects import pandas2ri
pandas2ri.activate()
from collections import defaultdict
import warnings

fastq_ext = '.fastq'

default_master_read = ''.join((
    'GCGCACGTCTGCCGCGCTGTTCTCCTCTTCCTCATCTCCGGGACCCGGA',# forward flank
    '........',                                         # sgID
    'AA.....TT.....AA.....',                            # random barcode
    'ATGCCCAAGAAGAAGAGGAAGGTGTCCAATTTACTGACCGTACACCAAAATTTGCCTGCATTACCGGTCGATGCAACGAGTGATGAGGTTCGCAAGAACCT')) # aft flank

############### Input Parameters ########################

parser = argparse.ArgumentParser(description="Prepare FASTQ files for DADA training & clustering.",
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument('input_dir', type=str, help='Input directory with fastq files.') 
parser.add_argument('--master_read', type=str, default=default_master_read, 
    help="Outline of the amplicon sequence, degenerate bases can be either 'N' or '.'")

parser.add_argument('-t', '--training_dir', default='training/', help='Directory to save files for error training.') 
parser.add_argument('-o', '--output_dir', default='preprocessed/', help='Directory to save files for barcode clustering.') 
parser.add_argument("-v", "--verbose", help='Output more Info', action="store_true")
parser.add_argument('-p', '--parallel', action='store_true', help='Multi-threaded operation')
parser.add_argument('-s', '--search_blast', action='store_true', help='Use NCBI BLAST algorithm to identify contaminations in samples')
parser.add_argument('-l', '--local_blast', action='store_true', 
    help='Use local NCBI BLAST+ algorithm to accelerate searching (if present), see tuba_seq/blast.py.')
parser.add_argument('-d', '--derep', action='store_true', help='De-replicate output fastQ files for DADA2 to minimize file sizes.')
parser.add_argument('-f', '--fraction', type=float, default=0.05, help='Minimum fraction of total reads to elicit a BLAST-search of an unknown sequence.')
parser.add_argument('-k',  '--skip', action='store_true', help='Skip files that already exist in output directories.')

parser.add_argument('-a', '--allowable_deviation', type=int, default=4, help="Length of Indel to tolerate before discarding reads.")
parser.add_argument('--alignment_flank', type=int, default=22, help='# of bases flanking the degenerate region to be used to score the quality of the read.')
parser.add_argument('--training_flank', type=int, default=18, help='# of bases flanking the degenerate region to be used to develop the DADA2 error model.')
parser.add_argument('-M', '--min_align_score', type=float, default=0.6, help='Minimum alignment score needed to keep read, Range [0, 1).')
parser.add_argument('--compression', default='bz2', choices=['bz2', 'gzip', 'none'], help='Compression algorithm for saved file.')
###############################################################################
args = parser.parse_args()

master_read = MasterRead(args.master_read, args)
output_dirs = args.training_dir, args.output_dir
dada2 = importr("dada2")
R_base = importr('base')

if args.derep or args.compression == 'none':
    compression_ext = ''
elif args.compression == 'gzip':
    from gzip import open       
    compression_ext = '.gz'
elif args.compression == 'bz2':
    from bz2 import open
    compression_ext = '.bz2'

if args.parallel:
    from tuba_seq.pmap import low_memory_pmap as map 

for Dir in output_dirs:
    os.makedirs(Dir, exist_ok=True)

files = [os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir) if fastq_ext in f]
max_file_size = 1e9 if args.parallel else 2e9

Log = logPrint(args)
Log('Processing {:} samples found in {:}.'.format(len(files), args.input_dir), print_line=True)
def process_sample(filename):
    sample = os.path.basename(filename.partition(fastq_ext)[0])
    filenames = [os.path.join(Dir, sample)+fastq_ext+compression_ext for Dir in output_dirs]
    if args.skip and all(map(os.path.isfile, filenames)):
        Log("Skipping "+sample+" (already exists).")
        return None
    if os.path.getsize(filename) > max_file_size:
         warnings.warn("""
{:} is {:.1f} GB. preprocess.py must load this entire fastq file into memory, 
and gzip decompression will inflate this size ~10x. You'll need about twice this
amount of memory for successful execution.""".format(filename, os.path.getsize(filename)/1e9), RuntimeWarning)
         if args.parallel:
            warnings.warn("Running preprocess.py in non-multi-threaded form saves memory (multi-threading processes several fastq files simultaneously).", RuntimeWarning)

    reads = fastqDF.from_file(filename, fake_header=True, use_Illumina_filter=True)
    init_reads = len(reads)
    gb = reads.set_index("DNA")['QC'].groupby(level='DNA') 

    with open(filenames[0], 'wb') as training_file, open(filenames[1], 'wb') as cluster_file:
        reads_tuples = gb.agg(master_read.process_read_set, reads.construct_read_set, training_file, cluster_file)
    counts = gb.count()
    output = pd.DataFrame(reads_tuples.tolist(), columns=['outcome', 'score'], index=counts.index)
    output['Reads'] = counts.values
    unaligned = output.query('outcome == "Unaligned"')['Reads']
    PhiX = pandas2ri.ri2py(dada2.isPhiX(pandas2ri.py2ri(unaligned.index))) == 1
    contaminants = unaligned.loc[(~PhiX)&(unaligned>int(args.fraction*init_reads))]
    tallies = output.groupby('outcome')['Reads'].sum()
    tallies['PhiX'] = unaligned.loc[PhiX].sum()
    tallies['Unaligned'] = tallies['Unaligned'] - tallies['PhiX'] if 'Unaligned' in tallies else 0
    percents = defaultdict(float, tallies/init_reads)
    Log('Sample {:}: {:.2f}M Reads, {Unaligned:.1%} unaligned, {PhiX:.1%} PhiX, {Wrong Barcode Length:.1%} inappropriate barcode length, {Clustered:.1%} Cluster-able.'.format(sample, init_reads*1e-6, **percents))
    tallies['Initial'] = init_reads
    return dict(tallies=tallies,
                info=reads.info, 
                unknown_DNAs=contaminants/init_reads,
                scores=output.groupby('score').sum())

_reports = list(map(process_sample, files))

dfs = {r['info'].loc['Sample']:r for r in _reports if r is not None}
meta_datas = ['tallies', 'unknown_DNAs', 'info', 'scores']
output = {meta_data:pd.concat({d['info'].loc['Sample']:d[meta_data] for d in _reports if d is not None}, names=['Sample']) for meta_data in meta_datas}

def derep(filename):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore") 
        try:
            derep = dada2.derepFastq(filename, verbose=args.verbose)
            R_base.saveRDS(derep, file=filename.replace(fastq_ext, '.rds'))
        except Exception: 
            print("Could not derep "+filename)
        else:
            os.remove(filename)

if args.derep:
    Log("De-replicating outputs for DADA2...")
    list(map(derep, [os.path.join(Dir, filename) for Dir in output_dirs for filename in os.listdir(Dir) if fastq_ext in filename]))

if args.search_blast and len(output['unknown_DNAs']) > 0: 
    from tuba_seq.blast import sleuth_DNAs
    gb = output['unknown_DNAs'].groupby(level='DNA')
    DNAs = list(gb.keys())
    Log("BLAST-searching", len(DNAs), "common unknown sequence(s)...", True)
    searches = sleuth_DNAs(DNAs, local_blast=args.local_blast)<commit_msg>Update preprocess.py<commit_after>#!/usr/bin/env python3
import pandas as pd
import os, numpy, argparse
from tuba_seq.fastq import fastqDF, MasterRead
from tuba_seq.shared import logPrint
from rpy2.robjects.packages import importr
from rpy2.robjects import pandas2ri
pandas2ri.activate()
from collections import defaultdict
import warnings

fastq_ext = '.fastq'

default_master_read = ''.join((
    'GCGCACGTCTGCCGCGCTGTTCTCCTCTTCCTCATCTCCGGGACCCGGA',# forward flank
    '........',                                         # sgID
    'AA.....TT.....AA.....',                            # random barcode
    'ATGCCCAAGAAGAAGAGGAAGGTGTCCAATTTACTGACCGTACACCAAAATTTGCCTGCATTACCGGTCGATGCAACGAGTGATGAGGTTCGCAAGAACCT')) # aft flank

############### Input Parameters ########################

parser = argparse.ArgumentParser(description="Prepare FASTQ files for DADA training & clustering.",
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument('input_dir', type=str, help='Input directory with fastq files.') 
parser.add_argument('--master_read', type=str, default=default_master_read, 
    help="Outline of the amplicon sequence, degenerate bases can be either 'N' or '.'")

parser.add_argument('-t', '--training_dir', default='training/', help='Directory to save files for error training.') 
parser.add_argument('-o', '--output_dir', default='preprocessed/', help='Directory to save files for barcode clustering.') 
parser.add_argument("-v", "--verbose", help='Output more Info', action="store_true")
parser.add_argument('-p', '--parallel', action='store_true', help='Multi-threaded operation')
parser.add_argument('-s', '--search_blast', action='store_true', help='Use NCBI BLAST algorithm to identify contaminations in samples')
parser.add_argument('-l', '--local_blast', action='store_true', 
    help='Use local NCBI BLAST+ algorithm to accelerate searching (if present), see tuba_seq/blast.py.')
parser.add_argument('-d', '--derep', action='store_true', help='De-replicate output fastQ files for DADA2 to minimize file sizes.')
parser.add_argument('-f', '--fraction', type=float, default=0.05, help='Minimum fraction of total reads to elicit a BLAST-search of an unknown sequence.')
parser.add_argument('-k',  '--skip', action='store_true', help='Skip files that already exist in output directories.')

parser.add_argument('-a', '--allowable_deviation', type=int, default=4, help="Length of Indel to tolerate before discarding reads.")
parser.add_argument('--alignment_flank', type=int, default=22, help='# of bases flanking the degenerate region to be used to score the quality of the read.')
parser.add_argument('--training_flank', type=int, default=18, help='# of bases flanking the degenerate region to be used to develop the DADA2 error model.')
parser.add_argument('-M', '--min_align_score', type=float, default=0.6, help='Minimum alignment score needed to keep read, Range [0, 1).')
parser.add_argument('--compression', default='bz2', choices=['bz2', 'gzip', 'none'], help='Compression algorithm for saved file.')
###############################################################################
args = parser.parse_args()

master_read = MasterRead(args.master_read, args)
output_dirs = args.training_dir, args.output_dir
dada2 = importr("dada2")
R_base = importr('base')

if args.derep or args.compression == 'none':
    compression_ext = ''
elif args.compression == 'gzip':
    from gzip import open       
    compression_ext = '.gz'
elif args.compression == 'bz2':
    from bz2 import open
    compression_ext = '.bz2'

if args.parallel:
    from tuba_seq.pmap import pmap 

for Dir in output_dirs:
    os.makedirs(Dir, exist_ok=True)

files = [os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir) if fastq_ext in f]
max_file_size = 1e9 if args.parallel else 2e9

Log = logPrint(args)
Log('Processing {:} samples found in {:}.'.format(len(files), args.input_dir), print_line=True)

tallies = {}
infos = {}
unknown_DNAs = {}
scores = {}

for filename in files:
    sample = os.path.basename(filename.partition(fastq_ext)[0])
    filenames = [os.path.join(Dir, sample)+fastq_ext+compression_ext for Dir in output_dirs]
    if args.skip and all(map(os.path.isfile, filenames)):
        Log("Skipping "+sample+" (already exists).")
        return None
    if os.path.getsize(filename) > max_file_size:
         warnings.warn("""
{:} is {:.1f} GB. preprocess.py must load this entire fastq file into memory, 
and gzip decompression will inflate this size ~10x. You'll need about twice this
amount of memory for successful execution.""".format(filename, os.path.getsize(filename)/1e9), RuntimeWarning)
         if args.parallel:
            warnings.warn("Running preprocess.py in non-multi-threaded form saves memory (multi-threading processes several fastq files simultaneously).", RuntimeWarning)

    reads = fastqDF.from_file(filename, fake_header=True, use_Illumina_filter=True)
    init_reads = len(reads)
    gb = reads.set_index("DNA")['QC'].groupby(level='DNA') 

    with open(filenames[0], 'wb') as training_file, open(filenames[1], 'wb') as cluster_file:
        reads_tuples = gb.agg(master_read.process_read_set, reads.construct_read_set, training_file, cluster_file)
    counts = gb.count()
    output = pd.DataFrame(reads_tuples.tolist(), columns=['outcome', 'score'], index=counts.index)
    output['Reads'] = counts.values
    unaligned = output.query('outcome == "Unaligned"')['Reads']
    PhiX = pandas2ri.ri2py(dada2.isPhiX(pandas2ri.py2ri(unaligned.index))) == 1
    contaminants = unaligned.loc[(~PhiX)&(unaligned>int(args.fraction*init_reads))]
    tallies = output.groupby('outcome')['Reads'].sum()
    tallies['PhiX'] = unaligned.loc[PhiX].sum()
    tallies['Unaligned'] = tallies['Unaligned'] - tallies['PhiX'] if 'Unaligned' in tallies else 0
    percents = defaultdict(float, tallies/init_reads)
    Log('Sample {:}: {:.2f}M Reads, {Unaligned:.1%} unaligned, {PhiX:.1%} PhiX, {Wrong Barcode Length:.1%} inappropriate barcode length, {Clustered:.1%} Cluster-able.'.format(sample, init_reads*1e-6, **percents))
    tallies['Initial'] = init_reads
    dict(tallies=tallies,
                info=reads.info, 
                unknown_DNAs=contaminants/init_reads,
                scores=output.groupby('score').sum())


dfs = {r['info'].loc['Sample']:r for r in _reports if r is not None}
meta_datas = ['tallies', 'unknown_DNAs', 'info', 'scores']
output = {meta_data:pd.concat({d['info'].loc['Sample']:d[meta_data] for d in _reports if d is not None}, names=['Sample']) for meta_data in meta_datas}

def derep(filename):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore") 
        try:
            derep = dada2.derepFastq(filename, verbose=args.verbose)
            R_base.saveRDS(derep, file=filename.replace(fastq_ext, '.rds'))
        except Exception: 
            print("Could not derep "+filename)
        else:
            os.remove(filename)

if args.derep:
    Log("De-replicating outputs for DADA2...")
    list(map(derep, [os.path.join(Dir, filename) for Dir in output_dirs for filename in os.listdir(Dir) if fastq_ext in filename]))

if args.search_blast and len(output['unknown_DNAs']) > 0: 
    from tuba_seq.blast import sleuth_DNAs
    gb = output['unknown_DNAs'].groupby(level='DNA')
    DNAs = list(gb.keys())
    Log("BLAST-searching", len(DNAs), "common unknown sequence(s)...", True)
    searches = sleuth_DNAs(DNAs, local_blast=args.local_blast)
{'fraction_non_alphanumeric': 0.1524361794307805, 'fraction_numerical': 0.007331866959941345, 'mean_word_length': 4.47991234477721, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6110793', 'n_tokens_mistral': 4994, 'n_tokens_neox': 4751, 'n_words': 1166}