<commit_before>#!/usr/bin/env python

import pysam
import sys
import os
import argparse
import subprocess

from random import randint, seed
from collections import defaultdict as dd
from collections import OrderedDict as od
from shutil import move
from uuid import uuid4


def getRG(tags):
    ''' fetch RG, tags is from a pysam.AlignedRead.tags, returns RG name '''
    for tag, val in tags:
        if tag == 'RG':
            return val
    return None


def makerandom(bamfn, reffn, outdir, samples=1000000):
    assert os.path.exists(reffn + '.fai'), "reference not indexed: " + reffn
    ref = pysam.Fastafile(reffn)
    bam = pysam.Samfile(bamfn, 'rb')

    # connect chromosomes end-to-end and track offsets, otherwise sampling will be biased
    chrlen = dd(int)
    for line in open(reffn + '.fai', 'r'):
        chrom, size = line.strip().split()[:2]
        size = int(size)
        chrlen[chrom] = size

    # calculate offsets
    offset = 0
    chroffset = dd(int)
    for chrom in sorted(chrlen.iterkeys()):
        offset += int(chrlen[chrom])
        chroffset[chrom] = offset

    genomelen = offset

    with open(outdir + '/' + 'matched.random.txt', 'w') as mrout:
        n = 1
        for read in bam.fetch(until_eof=True):
            rndoffset = randint(0,genomelen)
            rndchrom = None
            lastoffset = 0
            for chrom in sorted(chrlen.iterkeys()):
                offset = int(chroffset[chrom])
                assert lastoffset < offset
                if rndoffset >= lastoffset and rndoffset < offset:
                    rndchrom = chrom
                    rndloc = rndoffset - lastoffset
                lastoffset = offset

            rndstart, rndend = rndloc, rndloc+read.rlen

            # don't extend random intervals past the end of the chromosome
            if rndloc+read.rlen > chrlen[rndchrom]:
                rndstart, rndend = rndloc-read.rlen, rndloc

            assert rndstart < rndend
            mrout.write('\t'.join((rndchrom, str(rndstart), str(rndend))) + '\n')

            if n == samples:
                break

            if n % 1000000 == 0:
                print n,"segments sampled..."

            n += 1


    tmpfn = outdir + '/' + str(uuid4()) + '.tmp'

    with open(tmpfn, 'w') as tmp:
        p = subprocess.Popen(['sort', '-k1,1', '-k2,2n', outdir + '/' + 'matched.random.txt'], stdout=subprocess.PIPE)
        for line in p.stdout:
            tmp.write(line)

    move(tmpfn, outdir + '/' + 'matched.random.txt')

    subprocess.call(['bgzip', '-f', outdir + '/' + 'matched.random.txt'])
    assert os.path.exists(outdir + '/' + 'matched.random.txt.gz'), "bgzip failed: " + outdir + '/' + 'matched.random.txt'

    subprocess.call(['tabix', '-s', '1', '-b', '2', '-e', '3', '-f', outdir + '/' + 'matched.random.txt.gz'])
    assert os.path.exists(outdir + '/' + 'matched.random.txt.gz.tbi'), "tabix failed: " + outdir + '/' + 'matched.random.txt'

    return outdir + '/' + 'matched.random.txt.gz', n


def overlap(bamfn, bedfn, outdir, tabix=False):
    bam = None

    if tabix:
        bam = pysam.Tabixfile(bamfn)
    else:
        bam = pysam.Samfile(bamfn, 'rb')

    overlaps = {}

    with open(bedfn, 'r') as bed:
        for line in bed:
            chrom, start, end, strand, name, eltstart, eltend = line.strip().split()
            elt5p, elt3p, eltstart, eltend = map(int, (start, end, eltstart, eltend))

            if strand == '-':
                elt5p, elt3p = elt3p, elt5p

            # initialize counters
            if tabix:
                libs = ['sampled']
            else:
                assert 'RG' in bam.header
                libs = [rg['ID'] for rg in bam.header['RG']]

            for lib in libs:
                if lib not in overlaps:
                    overlaps[lib] = {}

                overlaps[lib][line.strip()] = {}
                overlaps[lib][line.strip()]['5pNonDup'] = 0
                overlaps[lib][line.strip()]['3pNonDup'] = 0
                overlaps[lib][line.strip()]['5pAll'] = 0
                overlaps[lib][line.strip()]['3pAll'] = 0


            # count 5p junction overlaps
            for read in bam.fetch(chrom, elt5p, elt5p+1):
                lib = 'sampled'

                if not tabix:
                    lib = getRG(read.tags)

                if tabix or not read.is_duplicate:
                    overlaps[lib][line.strip()]['5pNonDup'] += 1

                overlaps[lib][line.strip()]['5pAll'] += 1

            # count 3p junction overlaps
            for read in bam.fetch(chrom, elt3p, elt3p+1):
                lib='sampled'

                if not tabix:
                    lib = getRG(read.tags)

                if tabix or not read.is_duplicate:
                    overlaps[lib][line.strip()]['3pNonDup'] += 1

                overlaps[lib][line.strip()]['3pAll'] += 1

    return overlaps


def readcount(bamfn):
    bam = pysam.Samfile(bamfn, 'rb')
    counts = dd(int)
    for read in bam.fetch(until_eof=True):
        rg = getRG(read.tags)
        counts[rg] += 1

    return counts


def main(args):
    assert args.bam.endswith('.bam'), "not a BAM file: " + args.bam
    seed(int(args.seed))

    if not os.path.exists(args.outdir):
        os.mkdir(args.outdir)
    assert os.path.exists(args.outdir), 'could not create output directory: ' + args.outdir

    rndfn, rndcount = makerandom(args.bam, args.ref, args.outdir, samples=int(args.samples))
    real_overlaps = overlap(args.bam, args.elts, args.outdir)
    sim_overlaps  = overlap(rndfn, args.elts, args.outdir, tabix=True)

    normalised = od()
    enrichment = od()<commit_msg>count coverage using only mapped reads<commit_after>#!/usr/bin/env python

import pysam
import sys
import os
import argparse
import subprocess

from random import randint, seed
from collections import defaultdict as dd
from collections import OrderedDict as od
from shutil import move
from uuid import uuid4


def getRG(tags):
    ''' fetch RG, tags is from a pysam.AlignedRead.tags, returns RG name '''
    for tag, val in tags:
        if tag == 'RG':
            return val
    return None


def makerandom(bamfn, reffn, outdir, samples=1000000):
    assert os.path.exists(reffn + '.fai'), "reference not indexed: " + reffn
    ref = pysam.Fastafile(reffn)
    bam = pysam.Samfile(bamfn, 'rb')

    # connect chromosomes end-to-end and track offsets, otherwise sampling will be biased
    chrlen = dd(int)
    for line in open(reffn + '.fai', 'r'):
        chrom, size = line.strip().split()[:2]
        size = int(size)
        chrlen[chrom] = size

    # calculate offsets
    offset = 0
    chroffset = dd(int)
    for chrom in sorted(chrlen.iterkeys()):
        offset += int(chrlen[chrom])
        chroffset[chrom] = offset

    genomelen = offset

    with open(outdir + '/' + 'matched.random.txt', 'w') as mrout:
        n = 1
        for read in bam.fetch(until_eof=True):
            rndoffset = randint(0,genomelen)
            rndchrom = None
            lastoffset = 0
            for chrom in sorted(chrlen.iterkeys()):
                offset = int(chroffset[chrom])
                assert lastoffset < offset
                if rndoffset >= lastoffset and rndoffset < offset:
                    rndchrom = chrom
                    rndloc = rndoffset - lastoffset
                lastoffset = offset

            rndstart, rndend = rndloc, rndloc+read.rlen

            # don't extend random intervals past the end of the chromosome
            if rndloc+read.rlen > chrlen[rndchrom]:
                rndstart, rndend = rndloc-read.rlen, rndloc

            assert rndstart < rndend
            mrout.write('\t'.join((rndchrom, str(rndstart), str(rndend))) + '\n')

            if n == samples:
                break

            if n % 1000000 == 0:
                print n,"segments sampled..."

            n += 1


    tmpfn = outdir + '/' + str(uuid4()) + '.tmp'

    with open(tmpfn, 'w') as tmp:
        p = subprocess.Popen(['sort', '-k1,1', '-k2,2n', outdir + '/' + 'matched.random.txt'], stdout=subprocess.PIPE)
        for line in p.stdout:
            tmp.write(line)

    move(tmpfn, outdir + '/' + 'matched.random.txt')

    subprocess.call(['bgzip', '-f', outdir + '/' + 'matched.random.txt'])
    assert os.path.exists(outdir + '/' + 'matched.random.txt.gz'), "bgzip failed: " + outdir + '/' + 'matched.random.txt'

    subprocess.call(['tabix', '-s', '1', '-b', '2', '-e', '3', '-f', outdir + '/' + 'matched.random.txt.gz'])
    assert os.path.exists(outdir + '/' + 'matched.random.txt.gz.tbi'), "tabix failed: " + outdir + '/' + 'matched.random.txt'

    return outdir + '/' + 'matched.random.txt.gz', n


def overlap(bamfn, bedfn, outdir, tabix=False):
    bam = None

    if tabix:
        bam = pysam.Tabixfile(bamfn)
    else:
        bam = pysam.Samfile(bamfn, 'rb')

    overlaps = {}

    with open(bedfn, 'r') as bed:
        for line in bed:
            chrom, start, end, strand, name, eltstart, eltend = line.strip().split()
            elt5p, elt3p, eltstart, eltend = map(int, (start, end, eltstart, eltend))

            if strand == '-':
                elt5p, elt3p = elt3p, elt5p

            # initialize counters
            if tabix:
                libs = ['sampled']
            else:
                assert 'RG' in bam.header
                libs = [rg['ID'] for rg in bam.header['RG']]

            for lib in libs:
                if lib not in overlaps:
                    overlaps[lib] = {}

                overlaps[lib][line.strip()] = {}
                overlaps[lib][line.strip()]['5pNonDup'] = 0
                overlaps[lib][line.strip()]['3pNonDup'] = 0
                overlaps[lib][line.strip()]['5pAll'] = 0
                overlaps[lib][line.strip()]['3pAll'] = 0


            # count 5p junction overlaps
            for read in bam.fetch(chrom, elt5p, elt5p+1):
                lib = 'sampled'

                if not tabix:
                    lib = getRG(read.tags)

                if tabix or not read.is_duplicate:
                    overlaps[lib][line.strip()]['5pNonDup'] += 1

                overlaps[lib][line.strip()]['5pAll'] += 1

            # count 3p junction overlaps
            for read in bam.fetch(chrom, elt3p, elt3p+1):
                lib='sampled'

                if not tabix:
                    lib = getRG(read.tags)

                if tabix or not read.is_duplicate:
                    overlaps[lib][line.strip()]['3pNonDup'] += 1

                overlaps[lib][line.strip()]['3pAll'] += 1

    return overlaps


def readcount(bamfn):
    bam = pysam.Samfile(bamfn, 'rb')
    counts = dd(int)
    for read in bam.fetch():
        rg = getRG(read.tags)
        counts[rg] += 1

    return counts


def main(args):
    assert args.bam.endswith('.bam'), "not a BAM file: " + args.bam
    seed(int(args.seed))

    if not os.path.exists(args.outdir):
        os.mkdir(args.outdir)
    assert os.path.exists(args.outdir), 'could not create output directory: ' + args.outdir

    rndfn, rndcount = makerandom(args.bam, args.ref, args.outdir, samples=int(args.samples))
    real_overlaps = overlap(args.bam, args.elts, args.outdir)
    sim_overlaps  = overlap(rndfn, args.elts, args.outdir, tabix=True)

    normalised = od()
    enrichment = od()
{'fraction_non_alphanumeric': 0.13503940494111397, 'fraction_numerical': 0.010803152395289117, 'mean_word_length': 1.9752370916754478, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, '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': '5247969', 'n_tokens_mistral': 3927, 'n_tokens_neox': 3658, 'n_words': 908}