<commit_before>"""

preprocess.py

This directory contains functions and classes for preprocessing raw data into
matrices and vectors we can run our EM algorithm on. A big part of this is to
reduce the overall size of our input matrixes to make our algorithm run quickly
and without using too much memory.

<NAME> (<EMAIL>)
Wed Apr  6 10:48:21 PDT 2016

"""

import sys
import numpy
import collections

import phylotree


class HapVarBaseMatrix(object):
    """
    This class is designed to simplify a pretty large 3 dimensional matrix of
    indicator variables. Since our matrix is sparse and only need to look up
    single cells, we can store this information in a simplied way. The first
    part is that there are only two possible values that we store, a value if
    the haplogroup, variant site, observed base is expected (1) or unexpected
    (0). In practice, we will use fractional values (e.g. 0.991 and 0.003) to
    allow for sequencing error or de novo mutations.  Second, we keep a
    dictionary of haplogroup, variant site, base tuples that mark where a
    non-reference base occurs. If an entry is not present in the dictionary, it
    must be the case that the reference base is the "expected" base and all
    others are "unexpected". Finally, we will now make use of mutation counts
    available from Phylotree to adjust the probabilities based on how "mutable"
    a site appears to be.
    """
    def __init__(self, refseq, phylo, mut_wt=0.01, mut_max=0.5):
        """ Initialzes the HapVarBaseMatrix. """
        self.refseq = refseq
        self.phylo  = phylo
        self.mut_wt = mut_wt
        self.mut_max = mut_max
        self.markers = dict()

        self.add_hap_markers(phylo.hap_var)
        return

    def add_hap_markers(self, hap_var):
        """
        Populates the table with bases that are associated with haplogroups.
        """
        for hap in hap_var:
            self.markers[hap] = dict()
            for var in hap_var[hap]:
                pos = phylotree.pos_from_var(var)
                der = phylotree.der_allele(var)
                if der != self.refseq[pos]:
                    self.markers[hap][pos] = der
        return

    def _prob(self, hap_pos, pos, base):
        """
        Returns the "probability" of observing this base, at this position in
        this haplogroup, given the dictionary that describes the derived
        alleles of this haplogroup.
        """
        mut_prob = min(self.mut_wt * sum(self.phylo.variants[pos].values()),
                       self.mut_max)
        if pos in hap_pos:
            # Does this haplogroup carry a derived base?
            if hap_pos[pos] == base:
                # Is it the one we observed?
                return 1.0 - mut_prob
        else:
            # No derived base at this position, does our observed match ref.
            if self.refseq[pos] == base:
                return 1.0 - mut_prob
        return mut_prob / 3.0

    def prob_for_vars(self, hap, pos_obs):
        """
        Finds the probability of observing the read signature from the given
        haplotype by calculating the product of the probability of all
        hap, pos, base tuples it represents.
        """
        total = 1
        hap_pos = self.markers[hap]
        for pos, obs in pos_obs:
            total *= self._prob(hap_pos, pos, obs)
        return total


def process_reads(bamfile, var_pos, min_mq, min_bq):
    """
    Reads in a set of reads from a BAM file and makes observations for known
    variant positions. Each read is simplified into a signature of observed
    base per variant site. At the same time, build a table of observed bases
    for every position in the reference.
    """
    read_obs = collections.defaultdict(dict)
    base_obs = collections.defaultdict(collections.Counter)
    var_pos = set(var_pos)
    for aln in bamfile.fetch():
        if aln.mapping_quality >= min_mq:
            for qpos, rpos in aln.get_aligned_pairs(matches_only=True):
                qpos = int(qpos)
                rpos = int(rpos)
                if (aln.query_qualities is None
                  or aln.query_qualities[qpos] >= min_bq):
                    # Store the observation for the ref position.
                    obs = aln.query_sequence[qpos].upper()
                    base_obs[rpos][obs] += 1

                    if rpos in var_pos:
                        # Add this to the list, if this is a known var site.
                        if rpos in read_obs[aln.query_name]:
                            # check if this position has been observed before.
                            if read_obs[aln.query_name][rpos] != obs:
                                read_obs[aln.query_name][rpos] = "N"
                        else:
                            read_obs[aln.query_name][rpos] = obs
    # Finished, do one pass to remove Ns
    for aln_id in read_obs:
        read_obs[aln_id] = {pos:base for pos, base in
                            read_obs[aln_id].items() if base != 'N'}
    return read_obs, base_obs


def read_signature(obs_by_pos):
    """
    Returns a string made from the positions where bases were able to be
    observed from the read.
    """
    return ','.join(["%d:%s" % (pos, obs_by_pos[pos])
                     for pos in sorted(obs_by_pos)])


def pos_obs_from_sig(read_sig):
    """
    Returns a list of position, observation pairs described in the read
    signature string.
    """
    def pos_obs(var):
        """ Splits var into a int position and a string base. """
        pos, obs = var.split(':')
        return int(pos), obs
    return [pos_obs(var) for var in read_sig.split(',')]


def reduce_reads(read_obs):
    """
    Takes a dictionary of read IDs and base observations at variable positions
    and returns a dictionary observation signatures (a string describing the
    base observations made at variant sites) that maps to a list of read IDs
    that carry that signature.
    """
    read_sigs = collections.defaultdict(list)
    for read_id in read_obs:
        sig = read_signature(read_obs[read_id])
        read_sigs[sig].append(read_id)
    return read_sigs


def build_em_matrix(refseq, phylo, reads, haplogroups, args):
    """
    Returns the matrix that describes the probabiliy of each read
    originating in each haplotype.
    """
    hvb_mat = HapVarBaseMatrix(refseq, phylo)
    read_hap_mat = numpy.empty((len(reads), len(haplogroups)))

    if args.verbose:
        sys.stderr.write('Building EM input matrix...\n')

    for i in xrange(len(reads)):
        pos_obs = pos_obs_from_sig(reads[i])
        for j in xrange(len(haplogroups)):
            read_hap_mat[i, j] = hvb_mat.prob_for_vars(haplogroups[j], pos_obs)
        if args.verbose and (i + 1) % 500 == 0:
            sys.stderr.write('  processed %d fragments...\n' % (i + 1))

    if args.verbose:
        sys.stderr.write('Done.\n\n')

    return read_hap_mat


def build_em_input(bamfile, refseq, phylo, args):
    """
    Builds the matrix that describes the "probability" that a read originated
    from a specific haplogroup.
    """
    var_pos = phylo.get_variant_pos()
    read_obs, base_obs = process_reads(bamfile, var_pos,
                                       args.min_mq, args.min_bq)
    read_sigs = reduce_reads(read_obs)

    if args.verbose:
        sys.stderr.write('Using %d aligned fragments (MQ>=%d) '
                         '(%d distinct sub-haplotypes)\n\n'
                           % (len(read_obs), args.min_mq, len(read_sigs)))

    # This is now the order we will be using for the matrix.
    haplogroups = sorted(phylo.hap_var)
    reads = sorted(read_sigs)
    weights = numpy.array([len(read_sigs[r]) for r in reads])

    em_matrix = build_em_matrix(refseq, phylo, reads, haplogroups, args)

    # make a list mapping matrix indexes to read IDs from bam.
    read_ids = [read_sigs[reads[i]] for i in xrange(reads)]

    return em_matrix, weights, haplogroups, read_ids, base_obs


def main():
    """ Main function for simple testing. """
#   args = {'min_mq':30, 'min_bq':30, 'verbose':True}
#   if len(sys.argv) > 2:
#       phy_fn = sys.argv[1]
#       bam_fn = sys.argv[2]<commit_msg>Removed call to AlignmentFile.fetch() from process_reads().<commit_after>"""

preprocess.py

This directory contains functions and classes for preprocessing raw data into
matrices and vectors we can run our EM algorithm on. A big part of this is to
reduce the overall size of our input matrixes to make our algorithm run quickly
and without using too much memory.

<NAME> (<EMAIL>)
Wed Apr  6 10:48:21 PDT 2016

"""

import sys
import numpy
import collections

import phylotree


class HapVarBaseMatrix(object):
    """
    This class is designed to simplify a pretty large 3 dimensional matrix of
    indicator variables. Since our matrix is sparse and only need to look up
    single cells, we can store this information in a simplied way. The first
    part is that there are only two possible values that we store, a value if
    the haplogroup, variant site, observed base is expected (1) or unexpected
    (0). In practice, we will use fractional values (e.g. 0.991 and 0.003) to
    allow for sequencing error or de novo mutations.  Second, we keep a
    dictionary of haplogroup, variant site, base tuples that mark where a
    non-reference base occurs. If an entry is not present in the dictionary, it
    must be the case that the reference base is the "expected" base and all
    others are "unexpected". Finally, we will now make use of mutation counts
    available from Phylotree to adjust the probabilities based on how "mutable"
    a site appears to be.
    """
    def __init__(self, refseq, phylo, mut_wt=0.01, mut_max=0.5):
        """ Initialzes the HapVarBaseMatrix. """
        self.refseq = refseq
        self.phylo  = phylo
        self.mut_wt = mut_wt
        self.mut_max = mut_max
        self.markers = dict()

        self.add_hap_markers(phylo.hap_var)
        return

    def add_hap_markers(self, hap_var):
        """
        Populates the table with bases that are associated with haplogroups.
        """
        for hap in hap_var:
            self.markers[hap] = dict()
            for var in hap_var[hap]:
                pos = phylotree.pos_from_var(var)
                der = phylotree.der_allele(var)
                if der != self.refseq[pos]:
                    self.markers[hap][pos] = der
        return

    def _prob(self, hap_pos, pos, base):
        """
        Returns the "probability" of observing this base, at this position in
        this haplogroup, given the dictionary that describes the derived
        alleles of this haplogroup.
        """
        mut_prob = min(self.mut_wt * sum(self.phylo.variants[pos].values()),
                       self.mut_max)
        if pos in hap_pos:
            # Does this haplogroup carry a derived base?
            if hap_pos[pos] == base:
                # Is it the one we observed?
                return 1.0 - mut_prob
        else:
            # No derived base at this position, does our observed match ref.
            if self.refseq[pos] == base:
                return 1.0 - mut_prob
        return mut_prob / 3.0

    def prob_for_vars(self, hap, pos_obs):
        """
        Finds the probability of observing the read signature from the given
        haplotype by calculating the product of the probability of all
        hap, pos, base tuples it represents.
        """
        total = 1
        hap_pos = self.markers[hap]
        for pos, obs in pos_obs:
            total *= self._prob(hap_pos, pos, obs)
        return total


def process_reads(alns, var_pos, min_mq, min_bq):
    """
    Read alignments from an iterator/iterable and makes observations for known
    variant positions. Each read is simplified into a signature of observed
    base per variant site. At the same time, build a table of observed bases
    for every position in the reference.

    Args:
        alns: An iterator/iterable of pysam AlignedSegments
        var_pos: A sorted list of all known variant positions from Phylotree.
        min_mq: The minimum mapping quality score required for a read to be
                considered.
        min_bq: The minimum base quality score required for a base observation.

    Returns:
        read_obs: A dictionary that maps AlignedSegment query_names to
                  dictionaries mapping reference positions to base
                  observations.
        base_obs: A dictionary containing counts of base observations for each
                  position in the reference.
    """
    read_obs = collections.defaultdict(dict)
    base_obs = collections.defaultdict(collections.Counter)
    var_pos = set(var_pos)
    for aln in alns:
        if aln.mapping_quality >= min_mq:
            for qpos, rpos in aln.get_aligned_pairs(matches_only=True):
                qpos = int(qpos)
                rpos = int(rpos)
                if (aln.query_qualities is None
                  or aln.query_qualities[qpos] >= min_bq):
                    # Store the observation for the ref position.
                    obs = aln.query_sequence[qpos].upper()
                    base_obs[rpos][obs] += 1

                    if rpos in var_pos:
                        # Add this to the list, if this is a known var site.
                        if rpos in read_obs[aln.query_name]:
                            # check if this position has been observed before.
                            if read_obs[aln.query_name][rpos] != obs:
                                read_obs[aln.query_name][rpos] = "N"
                        else:
                            read_obs[aln.query_name][rpos] = obs
    # Finished, do one pass to remove Ns
    for aln_id in read_obs:
        read_obs[aln_id] = {pos:base for pos, base in
                            read_obs[aln_id].items() if base != 'N'}
    return read_obs, base_obs


def read_signature(obs_by_pos):
    """
    Returns a string made from the positions where bases were able to be
    observed from the read.
    """
    return ','.join(["%d:%s" % (pos, obs_by_pos[pos])
                     for pos in sorted(obs_by_pos)])


def pos_obs_from_sig(read_sig):
    """
    Returns a list of position, observation pairs described in the read
    signature string.
    """
    def pos_obs(var):
        """ Splits var into a int position and a string base. """
        pos, obs = var.split(':')
        return int(pos), obs
    return [pos_obs(var) for var in read_sig.split(',')]


def reduce_reads(read_obs):
    """
    Takes a dictionary of read IDs and base observations at variable positions
    and returns a dictionary observation signatures (a string describing the
    base observations made at variant sites) that maps to a list of read IDs
    that carry that signature.
    """
    read_sigs = collections.defaultdict(list)
    for read_id in read_obs:
        sig = read_signature(read_obs[read_id])
        read_sigs[sig].append(read_id)
    return read_sigs


def build_em_matrix(refseq, phylo, reads, haplogroups, args):
    """
    Returns the matrix that describes the probabiliy of each read
    originating in each haplotype.
    """
    hvb_mat = HapVarBaseMatrix(refseq, phylo)
    read_hap_mat = numpy.empty((len(reads), len(haplogroups)))

    if args.verbose:
        sys.stderr.write('Building EM input matrix...\n')

    for i in xrange(len(reads)):
        pos_obs = pos_obs_from_sig(reads[i])
        for j in xrange(len(haplogroups)):
            read_hap_mat[i, j] = hvb_mat.prob_for_vars(haplogroups[j], pos_obs)
        if args.verbose and (i + 1) % 500 == 0:
            sys.stderr.write('  processed %d fragments...\n' % (i + 1))

    if args.verbose:
        sys.stderr.write('Done.\n\n')

    return read_hap_mat


def build_em_input(bamfile, refseq, phylo, args):
    """
    Builds the matrix that describes the "probability" that a read originated
    from a specific haplogroup.
    """
    var_pos = phylo.get_variant_pos()

    input_alns = bamfile.fetch()
    read_obs, base_obs = process_reads(input_alns, var_pos,
                                       args.min_mq, args.min_bq)
    read_sigs = reduce_reads(read_obs)

    if args.verbose:
        sys.stderr.write('Using %d aligned fragments (MQ>=%d) '
                         '(%d distinct sub-haplotypes)\n\n'
                           % (len(read_obs), args.min_mq, len(read_sigs)))

    # This is now the order we will be using for the matrix.
    haplogroups = sorted(phylo.hap_var)
    reads = sorted(read_sigs)
    weights = numpy.array([len(read_sigs[r]) for r in reads])

    em_matrix = build_em_matrix(refseq, phylo, reads, haplogroups, args)

    # make a list mapping matrix indexes to read IDs from bam.
    read_ids = [read_sigs[reads[i]] for i in xrange(reads)]

    return em_matrix, weights, haplogroups, read_ids, base_obs


def main():
    """ Main function for simple testing. """
#   args = {'min_mq':30, 'min_bq':30, 'verbose':True}
#   if len(sys.argv) > 2:
#       phy_fn = sys.argv[1]
#       bam_fn = sys.argv[2]
{'fraction_non_alphanumeric': 0.08245736977341743, 'fraction_numerical': 0.005606166783461808, 'mean_word_length': 2.251993923281428, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1592744', 'n_tokens_mistral': 5275, 'n_tokens_neox': 4997, 'n_words': 1767}