<commit_before>#!/bin/env python

import sys,pysam,argparse

# TODO: add readgroup mongering

# fixes unmapped reads that are marked as 'reverse'
def cleanup(read):
    if read.is_unmapped and read.is_reverse:
        read.is_reverse = False
    return read

def main(args):
    targetbam = pysam.Samfile(args.targetbam, 'rb')
    donorbam  = pysam.Samfile(args.donorbam, 'rb')
    outputbam = pysam.Samfile(args.outputbam, 'wb', template=targetbam)

    namechange = None

    if args.namechange:
        namechange = args.namechange

    # load reads from donorbam into dict 
    sys.stderr.write("loading donor reads into dictionary...\n")
    nr = 0
    rdict = {}
    for read in donorbam.fetch(until_eof=True):
        pairname = 'F' # read is first in pair
        if read.is_read2:
            pairname = 'S' # read is second in pair
        if not read.is_paired:
            pairname = 'U' # read is unpaired

        extqname = ','.join((read.qname,pairname))
        rdict[extqname] = read
        nr += 1

    sys.stderr.write("loaded " + str(nr) + " reads, reading " + args.targetbam + "...\n")


    used = {}
    for read in targetbam.fetch(until_eof=True):
        pairname = 'F' # read is first in pair
        if read.is_read2:
            pairname = 'S' # read is second in pair
        if not read.is_paired:
            pairname = 'U' # read is unpaired

        extqname = ','.join((read.qname,pairname))
        if extqname in rdict: 
            if args.keepqual:
                if namechange: # must set name _before_ setting quality (see pysam docs)
                    rdict[extqname].qname = args.namechange + rdict[extqname].qname    
                rdict[extqname].qual = read.qual
            rdict[extqname] = cleanup(rdict[extqname])
            outputbam.write(rdict[extqname])  # write read from donor .bam
            used[extqname] = True
        else:
            if namechange:
                qual = read.qual # temp
                read.qname = args.namechange + read.qname
                read.qual = qual
            read = cleanup(read)
            outputbam.write(read) # write read from target .bam

    nadded = 0
    # dump the unused reads from the donor if requested with --all
    if args.all:
        for extqname in rdict.keys():
            if extqname not in used:
                if namechange:
                    qual = rdict[extqname].qual # temp
                    rdict[extqname].qname = args.namechange + rdict[extqname].qname
                    rdict[extqname].qual = qual
                rdict[extqname] = cleanup(rdict[extqname])
                outputbam.write(rdict[extqname])
                nadded += 1
        sys.stderr.write("added " + str(nadded) + " reads due to --all\n")

    targetbam.close()
    donorbam.close()
    outputbam.close()<commit_msg>read group fill-in<commit_after>#!/bin/env python

import sys,pysam,argparse
from random import randint

# TODO: add readgroup mongering

def cleanup(read,RG):
    '''
    fixes unmapped reads that are marked as 'reverse'
    fill in read group at random from existing RGs if 
    RG tags are present in .bam header 
    '''
    if read.is_unmapped and read.is_reverse:
        read.is_reverse = False

    if RG:
        hasRG = False
        for tag in read.tags:
            if tag[0] == 'RG':
                hasRG = True

        if not hasRG:
            # add random read group from list in header
            newRG = RG[randint(0,len(RG)-1)]
            read.tags = read.tags + [("RG",newRG)]
    return read

def getRGs(bam):
    '''return list of RG IDs'''
    RG = []
    if 'RG' in bam.header:
        for headRG in bam.header['RG']:
            RG.append(headRG['ID'])
    return RG

def main(args):
    targetbam = pysam.Samfile(args.targetbam, 'rb')
    donorbam  = pysam.Samfile(args.donorbam, 'rb')
    outputbam = pysam.Samfile(args.outputbam, 'wb', template=targetbam)

    RG = getRGs(targetbam) # read groups

    namechange = None

    if args.namechange:
        namechange = args.namechange

    # load reads from donorbam into dict 
    sys.stderr.write("loading donor reads into dictionary...\n")
    nr = 0
    rdict = {}
    for read in donorbam.fetch(until_eof=True):
        pairname = 'F' # read is first in pair
        if read.is_read2:
            pairname = 'S' # read is second in pair
        if not read.is_paired:
            pairname = 'U' # read is unpaired

        extqname = ','.join((read.qname,pairname))
        rdict[extqname] = read
        nr += 1

    sys.stderr.write("loaded " + str(nr) + " reads, reading " + args.targetbam + "...\n")


    used = {}
    for read in targetbam.fetch(until_eof=True):
        pairname = 'F' # read is first in pair
        if read.is_read2:
            pairname = 'S' # read is second in pair
        if not read.is_paired:
            pairname = 'U' # read is unpaired

        extqname = ','.join((read.qname,pairname))
        if extqname in rdict: 
            if args.keepqual:
                if namechange: # must set name _before_ setting quality (see pysam docs)
                    rdict[extqname].qname = args.namechange + rdict[extqname].qname    
                rdict[extqname].qual = read.qual
            rdict[extqname] = cleanup(rdict[extqname],RG)
            outputbam.write(rdict[extqname])  # write read from donor .bam
            used[extqname] = True
        else:
            if namechange:
                qual = read.qual # temp
                read.qname = args.namechange + read.qname
                read.qual = qual
            read = cleanup(read,RG)
            outputbam.write(read) # write read from target .bam

    nadded = 0
    # dump the unused reads from the donor if requested with --all
    if args.all:
        for extqname in rdict.keys():
            if extqname not in used:
                if namechange:
                    qual = rdict[extqname].qual # temp
                    rdict[extqname].qname = args.namechange + rdict[extqname].qname
                    rdict[extqname].qual = qual
                rdict[extqname] = cleanup(rdict[extqname],RG)
                outputbam.write(rdict[extqname])
                nadded += 1
        sys.stderr.write("added " + str(nadded) + " reads due to --all\n")

    targetbam.close()
    donorbam.close()
    outputbam.close()
{'fraction_non_alphanumeric': 0.10418310970797158, 'fraction_numerical': 0.0023677979479084454, 'mean_word_length': 1.9077558513079393, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5747717', 'n_tokens_mistral': 2067, 'n_tokens_neox': 1967, 'n_words': 530}