<commit_before>#!/usr/bin/env python
# encoding: utf-8
"""
Tools to handle reads sequenced with unique molecular identifiers (UMIs).
"""
import os
import sys
import doctest
from re import findall
from pysam import Samfile
from toolshed import nopen
from itertools import islice, izip
from collections import Counter, defaultdict
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

from ._version import __version__

os.environ['TERM'] = 'linux'

IUPAC = {"A":"A","T":"T","C":"C","G":"G","R":"GA","Y":"TC",
         "M":"AC","K":"GT","S":"GC","W":"AT","H":"ACT",
         "B":"GTC","V":"GCA","D":"GAT","N":"GATC"}


class Fastq(object):
    def __init__(self, args):
        self.name = args[0][1:]
        self.seq = args[1]
        self.qual = args[3]
        assert len(self.seq) == len(self.qual)

    def __repr__(self):
        return "Fastq({name})".format(name=self.name)

    def __str__(self):
        return "@{name}\n{seq}\n+\n{qual}".format(name=self.name,
                seq=self.seq, qual=self.qual)


def umi_from_name(name):
    """
    extract the UMI sequence from the read name.

    >>> umi_from_name("cluster_1017333:UMI_GCCGCA")
    'GCCGCA'
    """
    return findall(r'UMI_([\w]*)', name)[0].strip()


def get_chromosomes(sam_header):
    """
    parse sam header to return SN values within SQ lines.
    """
    chromosomes = []

    for line in sam_header.split("\n"):
        if not line.startswith("@SQ"): continue
        line = line.strip().split("\t")

        for token in line:
            if not token.startswith("SN"): continue
            chromosome = token.split(":", 1)[1]
            chromosomes.append(chromosome)

    assert len(chromosomes) > 0
    return chromosomes


def process_bam(args):
    """
    removes duplicate reads characterized by their UMI at any given start
    location.

    abam    input bam with potential duplicate UMIs
    bbam    output bam after removing duplicate UMIs
    """
    with Samfile(args.abam, 'rb') as in_bam, \
            Samfile(args.bbam, 'wb', template=in_bam) as out_bam:
        chromosomes = get_chromosomes(in_bam.text)

        for chrom in chromosomes:
            print >>sys.stderr, "processing chromosome", chrom

            umi_idx = defaultdict(set)
            read_counts = Counter()

            for read in in_bam.fetch(chrom):
                if read.is_unmapped: continue

                # get the iupac umi sequence
                umi = umi_from_name(read.qname)

                # get actual read start3
                # read.pos accounts for 5' soft clipping
                if read.is_reverse:
                    # read.alen alignment length accounting for 3' soft clipping
                    # UMIs are then compared to reads with the same start
                    read_start = read.pos + read.alen
                else:
                    read_start = read.pos

                # add count for this start
                read_counts[read_start] += 1

                # if umi is not present, write it out
                if read_start not in umi_idx:
                    out_bam.write(read)

                elif umi not in umi_idx[read_start]:
                    out_bam.write(read)

                # keep track of unique UMIs - set eliminates duplicates
                umi_idx[read_start].add(umi)

            # process before and after counts over chrom
            for start, before_count in sorted(read_counts.items()):
                print "{chrom}\t{start}\t{stop}\t{before}\t{after}".format(
                            chrom=chrom,
                            start=start,
                            stop=start+1,
                            before=before_count,
                            after=len(umi_idx[start]))


def readfq(fq):
    with nopen(fq) as fh:
        fqclean = (x.strip("\r\n") for x in fh if x.strip())
        while True:
            rd = [x for x in islice(fqclean, 4)]
            if not rd: raise StopIteration
            assert all(rd) and len(rd) == 4
            yield Fastq(rd)


def valid_umi(iupac, umi):
    """
    parse UMI sequence to validate against IUPAC sequence.

    >>> valid_umi("NNNV", "ACGT")<commit_msg>cleaning up the UMI seen logic<commit_after>#!/usr/bin/env python
# encoding: utf-8
"""
Tools to handle reads sequenced with unique molecular identifiers (UMIs).
"""
import os
import sys
import doctest
from re import findall
from pysam import Samfile
from toolshed import nopen
from itertools import islice, izip
from collections import Counter, defaultdict
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

from ._version import __version__

os.environ['TERM'] = 'linux'

IUPAC = {"A":"A","T":"T","C":"C","G":"G","R":"GA","Y":"TC",
         "M":"AC","K":"GT","S":"GC","W":"AT","H":"ACT",
         "B":"GTC","V":"GCA","D":"GAT","N":"GATC"}


class Fastq(object):
    def __init__(self, args):
        self.name = args[0][1:]
        self.seq = args[1]
        self.qual = args[3]
        assert len(self.seq) == len(self.qual)

    def __repr__(self):
        return "Fastq({name})".format(name=self.name)

    def __str__(self):
        return "@{name}\n{seq}\n+\n{qual}".format(name=self.name,
                seq=self.seq, qual=self.qual)


def umi_from_name(name):
    """
    extract the UMI sequence from the read name.

    >>> umi_from_name("cluster_1017333:UMI_GCCGCA")
    'GCCGCA'
    """
    return findall(r'UMI_([\w]*)', name)[0].strip()


def get_chromosomes(sam_header):
    """
    parse sam header to return SN values within SQ lines.
    """
    chromosomes = []

    for line in sam_header.split("\n"):
        if not line.startswith("@SQ"): continue
        line = line.strip().split("\t")

        for token in line:
            if not token.startswith("SN"): continue
            chromosome = token.split(":", 1)[1]
            chromosomes.append(chromosome)

    assert len(chromosomes) > 0
    return chromosomes


def process_bam(args):
    """
    removes duplicate reads characterized by their UMI at any given start
    location.

    abam    input bam with potential duplicate UMIs
    bbam    output bam after removing duplicate UMIs
    """
    with Samfile(args.abam, 'rb') as in_bam, \
            Samfile(args.bbam, 'wb', template=in_bam) as out_bam:
        chromosomes = get_chromosomes(in_bam.text)

        for chrom in chromosomes:
            print >>sys.stderr, "processing chromosome", chrom

            umi_idx = defaultdict(set)
            read_counts = Counter()

            for read in in_bam.fetch(chrom):
                if read.is_unmapped: continue

                # get the iupac umi sequence
                umi = umi_from_name(read.qname)

                # get actual read start
                # read.pos accounts for 5' soft clipping
                if read.is_reverse:
                    # read.alen alignment length accounting for 3' soft clipping
                    # UMIs are then compared to reads with the same start
                    read_start = read.pos + read.alen
                else:
                    read_start = read.pos

                # add count for this start; counts all reads
                read_counts[read_start] += 1

                # check if UMI seen
                if umi in umi_idx[read_start]: continue

                # keep track of unique UMIs - set eliminates duplicates
                umi_idx[read_start].add(umi)

                out_bam.write(read)

            # process before and after counts over chrom
            for start, before_count in sorted(read_counts.items()):
                print "{chrom}\t{start}\t{stop}\t{before}\t{after}".format(
                            chrom=chrom,
                            start=start,
                            stop=start+1,
                            before=before_count,
                            after=len(umi_idx[start]))


def readfq(fq):
    with nopen(fq) as fh:
        fqclean = (x.strip("\r\n") for x in fh if x.strip())
        while True:
            rd = [x for x in islice(fqclean, 4)]
            if not rd: raise StopIteration
            assert all(rd) and len(rd) == 4
            yield Fastq(rd)


def valid_umi(iupac, umi):
    """
    parse UMI sequence to validate against IUPAC sequence.

    >>> valid_umi("NNNV", "ACGT")
{'fraction_non_alphanumeric': 0.11534302535143579, 'fraction_numerical': 0.005406704313348552, 'mean_word_length': 2.0115774240231548, 'pattern_counts': {'":': 32, '<': 3, '<?xml version=': 0, '>': 21, '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': '5246819', 'n_tokens_mistral': 2636, 'n_tokens_neox': 2437, 'n_words': 681}