Binomial Test for Identifying Regions of Enriched Differentiation

binomial_test_window_analysis.py

"""
Copyright 2017 Ronald J. Nowling

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import argparse

import numpy as np
from scipy import stats

def stream_snp_positions(flname):
    with open(flname) as fl:
        for ln in fl:
            cols = ln.split()
            yield int(cols[1])
            
def parseargs():
    args = argparse.ArgumentParser()

    args.add_argument("--snps-fl",
                      type=str,
                      required=True,
                      help="Significant SNPs")

    args.add_argument("--window-size",
                      type=int,
                      required=True)

    args.add_argument("--output-windows",
                      type=str,
                      required=True,
                      help="Output window data")

    args.add_argument("--alpha",
                      type=float,
                      default=0.01,
                      help="Significance threshold")

    return args.parse_args()

if __name__ == "__main__":
    args = parseargs()

    sig_positions = []
    max_pos = 0
    for pos in stream_snp_positions(args.snps_fl):
        max_pos = max(pos, max_pos)
        sig_positions.append(int(pos))

    bins = range(0, max_pos, args.window_size)
    if max_pos % args.window_size != 0:
        bins.append(max_pos)

    hist, _ = np.histogram(sig_positions,
                           bins=bins)

    n_bins = float(len(hist))
    n_snps = len(sig_positions)

    windows_stats = []
    total_prob = 0.0
    for i, counts in enumerate(hist):
        # assume uniform prob, adjusted by bin size
        bin_size = bins[i+1] - bins[i]
        expected_prob = float(bin_size) / max_pos
        total_prob += expected_prob
        
        pvalue = stats.binom_test(counts, n_snps, expected_prob, alternative="greater")
        is_sig = pvalue <= args.alpha
        windows_stats.append((bins[i], bins[i+1], counts, is_sig, pvalue))

    if args.output_windows:
        with open(args.output_windows, "w") as fl:
            for start, end, count, is_sig, pvalue in windows_stats:
                fl.write("%s\t%s\t%s\t%s\t%s\n" % (start, end, count, is_sig, pvalue))