<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Merge peaks in a file by clustering instead of simple overlap merge.

===============================================================================

        AUTHOR: <NAME>, <EMAIL>
  ORGANIZATION: Stanford University
       LICENSE: MIT License, property of Stanford, use as you wish
       VERSION: 0.1
       CREATED: 2016-29-24 16:08
 Last modified: 2016-08-25 16:41

   DESCRIPTION: When calling ATACseq peaks with MACS in different populations,
                a single peak can be shifted slightly in such a way that
                multiple peaks are collapsed into a single peak when using
                bedtools merge.

                To avoid this issue, this script uses information about the
                prior peak as well as the next peak and tries to create
                clusters that way. The default overlap to be within a single
                cluster is 75%, this can be set using the -o flag.

                This script uses a file parsing iterator that can be easily
                modified to use any file format, the default format is bed.

          NOTE: File must contain all peaks and be sorted by coordinate first.

===============================================================================
"""
import sys
import bz2
import gzip
import argparse
from collections import OrderedDict

from tqdm import tqdm
import logme

logme.MIN_LEVEL = 'debug'


def peak_merge(peak_file, outfile=sys.stdout, overlap=.75, logfile=sys.stderr):
    """Merge peaks.

    :peak_file: A file handle or sequence file, currently bed only. File
                extension used for parsing, gzip or bzip compression OK, file
                handle OK.
    :outfile:   A bed file of merged peaks with the following fields:
                chr, start, end, count, mean_fold_change, mean_log10pval
                count is the number of members that went into the peak.
    :overlap:   The amount a peak can overlap a prior peak before being moved
                into a new cluster.
    :logfile:   A file to contain some summary stats.
    """
    # Make sure overlap is specified
    if not isinstance(overlap, float):
        overlap = .75
        logme.log('Overlap not specified or not float, using .75', 'info')
    # Specify an offset that is subtracted from peaks to avoid clustering of
    # minor overlaps
    offset         = 4
    # Load file generator
    if isinstance(peak_file, str):
        line_count = int(check_output(
            "wc -l data/all_peaks_sorted.bed | sed 's/ .*//'", shell=True)
            .decode().rstrip())
    else:
        line_count = None
    pfile          = peak_file_parser(peak_file)
    # Create objects to track stats
    lines          = 0
    clusters       = 0
    cluster_sizes  = {}
    # First peak
    prior_peak     = next(pfile)
    cluster        = Cluster(prior_peak)
    lines         += 1
    clusters      += 1
    # Open outfile and run algorithm
    with open_zipped(outfile, 'w') as fout:
        # Loop through peaks
        #  for peak in tqdm(pfile, total=line_count, unit='lines'):
        for peak in pfile:
            lines += 1
            overlap_prior = peak.start - offset < prior_peak.end
            logme.log('Prior overlap: {}'.format(overlap_prior), 'debug')
            try:
                next_peak    = next(pfile)
                overlap_next = peak.end - offset < next_peak.start
                logme.log('Next overlap: {}'.format(overlap_next), 'debug')
            # Last peak
            except StopIteration:
                next_peak    = None
                overlap_next = False
                logme.log('Last peak of file: {}_{}'.format(
                    peak.chrom, peak.start), 'debug')

            ######################
            #  Actual Algorithm  #
            ######################

            # Overlap both, decide by overlap amount.
            if overlap_prior and overlap_next:
                prior_overlap_amount = abs(float(prior_peak.end - peak.start) /
                                           float(cluster.len))
                logme.log('Prior overlap amount: {}'
                          .format(prior_overlap_amount), 'debug')
                #  next_overlap_amount  = (float(peak.end - next_peak.start) /
                #                          float(next_peak.end - next_peak.start))
                prior_peak = peak
                # Overlap prior more than 75%, add to prior cluster.
                if prior_overlap_amount > overlap:
                    cluster.add(peak)
                    continue
                # Overlap prior less than 75%, make new cluster.
                else:
                    # Stats
                    if cluster.count in cluster_sizes:
                        cluster_sizes[cluster.count] += 1
                    else:
                        cluster_sizes[cluster.count]  = 1
                    clusters    += 1
                    # Write cluster and make new one
                    cluster.write(fout)
                    cluster = Cluster(peak)
                    continue
            # Overlap only prior, add to cluster.
            elif overlap_prior:
                cluster.add(peak)
                prior_peak = peak
                continue
            # Overlap next or none, make new cluster.
            else:
                # Stats
                if cluster.count in cluster_sizes:
                    cluster_sizes[cluster.count] += 1
                else:
                    cluster_sizes[cluster.count]  = 1
                clusters    += 1
                # Write cluster and make new one
                cluster.write(fout)
                cluster      = Cluster(peak)
                prior_peak   = peak
                continue

    # Print stats
    logfile.write('\n')
    logme.log('Clustering complete,\nstats:', 'info')
    logfile.write('Total lines:\t{}\n'.format(lines) +
                  'Total clusters:\t{}\n'.format(clusters) +
                  'Total clustered:{}\n'.format(sum(cluster_sizes.values())) +
                  'Cluster sizes:\n')
    for k, v in OrderedDict(cluster_sizes).items():
        logfile.write('\t{}:\t{}\n'.format(k, v))


class Cluster(object):

    """A single merged peak."""

    def __init__(self, peak):
        """Create self.

        :peak: A Peak object.
        """
        self.name        = '{}_{}'.format(peak.chrom, peak.start)
        self.peaks       = [peak]
        self.chrom       = peak.chrom
        self.start       = peak.start
        self.end         = peak.end
        self.len         = self.start - self.end
        self.fold_change = peak.fold_change
        self.log10p      = peak.log10p
        self.count       = 1

    def add(self, peak):
        """Add a peak to this cluster.

        :peak: A Peak object.
        """
        self.peaks.append(peak)<commit_msg>Bug fix<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Merge peaks in a file by clustering instead of simple overlap merge.

===============================================================================

        AUTHOR: <NAME>, <EMAIL>
  ORGANIZATION: Stanford University
       LICENSE: MIT License, property of Stanford, use as you wish
       VERSION: 0.1
       CREATED: 2016-29-24 16:08
 Last modified: 2016-08-25 16:45

   DESCRIPTION: When calling ATACseq peaks with MACS in different populations,
                a single peak can be shifted slightly in such a way that
                multiple peaks are collapsed into a single peak when using
                bedtools merge.

                To avoid this issue, this script uses information about the
                prior peak as well as the next peak and tries to create
                clusters that way. The default overlap to be within a single
                cluster is 75%, this can be set using the -o flag.

                This script uses a file parsing iterator that can be easily
                modified to use any file format, the default format is bed.

          NOTE: File must contain all peaks and be sorted by coordinate first.

===============================================================================
"""
import sys
import bz2
import gzip
import argparse
from subprocess import check_output
from collections import OrderedDict

from tqdm import tqdm
import logme

logme.MIN_LEVEL = 'debug'


def peak_merge(peak_file, outfile=sys.stdout, overlap=.75, logfile=sys.stderr):
    """Merge peaks.

    :peak_file: A file handle or sequence file, currently bed only. File
                extension used for parsing, gzip or bzip compression OK, file
                handle OK.
    :outfile:   A bed file of merged peaks with the following fields:
                chr, start, end, count, mean_fold_change, mean_log10pval
                count is the number of members that went into the peak.
    :overlap:   The amount a peak can overlap a prior peak before being moved
                into a new cluster.
    :logfile:   A file to contain some summary stats.
    """
    # Make sure overlap is specified
    if not isinstance(overlap, float):
        overlap = .75
        logme.log('Overlap not specified or not float, using .75', 'info')
    # Specify an offset that is subtracted from peaks to avoid clustering of
    # minor overlaps
    offset         = 4
    # Load file generator
    if isinstance(peak_file, str):
        line_count = int(check_output(
            "wc -l data/all_peaks_sorted.bed | sed 's/ .*//'", shell=True)
            .decode().rstrip())
    else:
        line_count = None
    pfile          = peak_file_parser(peak_file)
    # Create objects to track stats
    lines          = 0
    clusters       = 0
    cluster_sizes  = {}
    # First peak
    prior_peak     = next(pfile)
    cluster        = Cluster(prior_peak)
    lines         += 1
    clusters      += 1
    # Open outfile and run algorithm
    with open_zipped(outfile, 'w') as fout:
        # Loop through peaks
        #  for peak in tqdm(pfile, total=line_count, unit='lines'):
        for peak in pfile:
            lines += 1
            overlap_prior = peak.start - offset < prior_peak.end
            logme.log('Prior overlap: {}'.format(overlap_prior), 'debug')
            try:
                next_peak    = next(pfile)
                overlap_next = peak.end - offset < next_peak.start
                logme.log('Next overlap: {}'.format(overlap_next), 'debug')
            # Last peak
            except StopIteration:
                next_peak    = None
                overlap_next = False
                logme.log('Last peak of file: {}_{}'.format(
                    peak.chrom, peak.start), 'debug')

            ######################
            #  Actual Algorithm  #
            ######################

            # Overlap both, decide by overlap amount.
            if overlap_prior and overlap_next:
                prior_overlap_amount = abs(float(prior_peak.end - peak.start) /
                                           float(cluster.len))
                logme.log('Prior overlap amount: {}'
                          .format(prior_overlap_amount), 'debug')
                #  next_overlap_amount  = (float(peak.end - next_peak.start) /
                #                          float(next_peak.end - next_peak.start))
                prior_peak = peak
                # Overlap prior more than 75%, add to prior cluster.
                if prior_overlap_amount > overlap:
                    cluster.add(peak)
                    continue
                # Overlap prior less than 75%, make new cluster.
                else:
                    # Stats
                    if cluster.count in cluster_sizes:
                        cluster_sizes[cluster.count] += 1
                    else:
                        cluster_sizes[cluster.count]  = 1
                    clusters    += 1
                    # Write cluster and make new one
                    cluster.write(fout)
                    cluster = Cluster(peak)
                    continue
            # Overlap only prior, add to cluster.
            elif overlap_prior:
                cluster.add(peak)
                prior_peak = peak
                continue
            # Overlap next or none, make new cluster.
            else:
                # Stats
                if cluster.count in cluster_sizes:
                    cluster_sizes[cluster.count] += 1
                else:
                    cluster_sizes[cluster.count]  = 1
                clusters    += 1
                # Write cluster and make new one
                cluster.write(fout)
                cluster      = Cluster(peak)
                prior_peak   = peak
                continue

    # Print stats
    logfile.write('\n')
    logme.log('Clustering complete,\nstats:', 'info')
    logfile.write('Total lines:\t{}\n'.format(lines) +
                  'Total clusters:\t{}\n'.format(clusters) +
                  'Total clustered:{}\n'.format(
                      sum([k*v for k,v in cluster_sizes.items()])) +
                  'Cluster sizes:\n')
    for k, v in OrderedDict(cluster_sizes).items():
        logfile.write('\t{}:\t{}\n'.format(k, v))


class Cluster(object):

    """A single merged peak."""

    def __init__(self, peak):
        """Create self.

        :peak: A Peak object.
        """
        self.name        = '{}_{}'.format(peak.chrom, peak.start)
        self.peaks       = [peak]
        self.chrom       = peak.chrom
        self.start       = peak.start
        self.end         = peak.end
        self.len         = self.start - self.end
        self.fold_change = peak.fold_change
        self.log10p      = peak.log10p
        self.count       = 1

    def add(self, peak):
        """Add a peak to this cluster.

        :peak: A Peak object.
        """
        self.peaks.append(peak)
{'fraction_non_alphanumeric': 0.1081197501615335, 'fraction_numerical': 0.008471534209203819, 'mean_word_length': 1.6042250888016452, 'pattern_counts': {'":': 0, '<': 11, '<?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': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4072286', 'n_tokens_mistral': 3853, 'n_tokens_neox': 3475, 'n_words': 1167}