reassign_matrix.py

#!/bin/env python

import sys
import getopt
from aggregate_matrix import matrix_iterator


def main(h_matrix=sys.stdin, h_a=sys.stdin, h_b=sys.stdin, h_out=sys.stdout):
    """The main program loop."""
    iter_matrix = matrix_iterator(h_matrix)
    iter_a = matrix_iterator(h_a)
    iter_b = matrix_iterator(h_b)

    for fields in iter_matrix:
        replace_a = iter_a.next()
        replace_b = iter_b.next()

        # don't consider non-overlapping fragments
        if replace_a[3] != "." and replace_b[3] != ".":
            data = []
            data.extend(replace_a[3:6])
            data.extend(replace_b[3:6])
            data.append("")
            data.append(fields[7])
            h_out.write("%s\n" % "\t".join(data))


if __name__ == "__main__":

    def usage(message="", error=None):
        """Print the usage information."""
        sys.stdout.write("""
Message %(message)s

Called as %(commandline)s

Usage:
    python %(tool)s -m [matrix file] -a [bed a] -b [bed b] -o [output file] 

Options:
    -m/--matrix   [file]  the matrix file
    -a/--bed_a    [file]  the BED A file
    -b/--bed_b    [file]  the BED B file
    -o/--output   [file]  the output matrix file, default stdout
    -h/--help     []  print this message

""" % {
    "message": message,
    "tool": sys.argv[0],
    "commandline": " ".join(sys.argv)
})
        if error is not None:
            sys.exit(error)
        # end of usage

    # main variables
    fstream_matrix = None
    fstream_a = None
    fstream_b = None
    fstream_output = sys.stdout

    # option parsing
    shortopt = "m:a:b:o:h"
    longopt = ["matrix=", "bed_a=", "bed_b=", "output=", "help"]

    try:
        opts, args = getopt.getopt(sys.argv[1:], shortopt, longopt)
        for o, a in opts:
            if o in ("-m", "--matrix"):
                fstream_matrix = open(a, "r")
            elif o in ("-a", "--bed_a"):
                fstream_a = open(a, "r")
            elif o in ("-b", "--bed_b"):
                fstream_b = open(a, "r")
            elif o in ("-o", "--output"):
                fstream_output = open(a, "w")
            elif o in ('-h', '--help'):
                usage("Help was asked", error=0)
    except getopt.GetoptError as err:
        usage(str(err), error=2)

    if fstream_matrix is None:
        usage("Matrix file not specified")
    if fstream_a is None:
        usage("BED a file not specified")
    if fstream_b is None:
        usage("BED b file not specified")

    # run the program loop
    main(fstream_matrix, fstream_a, fstream_b, fstream_output)

    if not fstream_matrix.closed and fstream_matrix is not sys.stdin:
        fstream_matrix.close()
    if not fstream_a.closed and fstream_a is not sys.stdin:
        fstream_a.close()
    if not fstream_b.closed and fstream_b is not sys.stdin:
        fstream_b.close()
    if not fstream_output.closed and fstream_output is not sys.stdout:
        fstream_output.close()