formatSumStatsforFgwas.py

#! /usr/bin/env python3

import pathlib
import gzip
import sys
import argparse
import textwrap

parser = argparse.ArgumentParser(
    prog="formatSumstatsForFgwas",
    usage="Generate input files for fGWAS using summary statistics",
    description=textwrap.dedent(
        """
        Use 1-based counting to specify column indexes.

        The options must specify the following columns:
            1. SNP ID
            2. Chromosome
            3. Position
            4. Allele Frequency
            5. Sample Size
            6. Z-score

        NOTE:
            (1) If Standard Error (SE) is available, it will be used over
        sample size and allele frequency.

            (2) If LNBF is available, it will override the Z-score,
            sample size, and allele frequency columns.
        """
    ),
    formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument("file", type=pathlib.Path, help="Path to the summary stats file")
parser.add_argument("outfile", type=pathlib.Path, help="Path to the output file")
parser.add_argument("--snpid", "-SNPID", type=int, help="Index for variant id (SNPID) column")
parser.add_argument("--chr", "-CHR", type=int, help="Index for chromosome (CHR) column")
parser.add_argument("--pos", "-POS", type=int, help="Index for position (POS) column")
parser.add_argument(
    "--af", "-F", type=int, help="Index for allele frequency (F) column"
)
parser.add_argument("--zscore", "-Z", type=int, help="Index for Z score (Z) column")
parser.add_argument("--size", "-N", type=int, help="Index for sample size (N) column")
parser.add_argument(
    "--se", "-SE", type=int, help="Index for standard error (SE)  size column"
)
parser.add_argument(
    "--lnbf", type=int, help="Index for natural log bayes factor (LNBF) column"
)
# parser.add_argument(
#     "--extra-cols",
#     nargs='+',
#     type=int,
#     help="Index for natural log bayes factor (LNBF) column"
# )
parser.add_argument(
    "--sep", type=str, default=None, help="Column separator (default: 'auto')"
)
parser.add_argument(
    "-pi",
    "--print-index",
    action="store_true",
    help="Print a (header -> column-index) for quickly selecting columns",
)
parser.add_argument(
    "--no-header", action="store_false", help="Specify if the file has no header"
)

if __name__ == "__main__":
    args = parser.parse_args()
    _open_fn, _mode = open, "r"

    if args.file.suffix == ".gz":
        _open_fn, _mode = gzip.open, "rt"

    if not args.file.exists():
        print("fatal: file not found", file=sys.stderr)
        sys.exit(1)

    if args.outfile.is_dir():
        print("fatal: outfile is not a file", file=sys.stderr)
        sys.exit(1)

    if not all([args.chr, args.pos, args.snpid]) and args.print_index == False:
        print(
            "error: must specify all necessary columns (--chr, --pos, --snpid, ...) OR --print-index"
        )
        parser.print_help(sys.stderr)
        sys.exit(1)

    if not (
        (args.lnbf)
        or all([args.zscore, args.size, args.af])
        or all([args.zscore, args.se])
        or (args.print_index == True)
    ):
        print(
            "error: must specify all necessary columns (--chr, --pos, --snpid, ...) OR --print-index"
        )
        parser.print_help(sys.stderr)
        sys.exit(1)

    if args.lnbf:
        print(
            "info: LNBF will override Z-score, Sample Size, and Allele Frequency columns.",
            file=sys.stderr,
        )

    if args.se:
        print(
            "info: SE will override Sample Size, and Allele Frequency columns.",
            file=sys.stderr,
        )

    col_index = {}

    header_vals = ["SNPID", "CHR", "POS", "F", "Z", "N", "SE", "LNBF"]
    all_vals = [
        args.snpid,
        args.chr,
        args.pos,
        args.af,
        args.zscore,
        args.size,
        args.se,
        args.lnbf,
    ]

    if not args.print_index:
        print("info: working..", file=sys.stderr)
    try:
        outfile = open(args.outfile, "w")
        header_line = " ".join(
            [header_vals[i] for i, _ in enumerate(all_vals) if _ is not None]
        )
        outfile.write("%s\n" % header_line)

        with _open_fn(args.file, _mode) as f:
            for i, line in enumerate(f):
                if i == 0 and args.no_header:
                    header = line.split(args.sep)
                    continue
                if args.print_index:
                    next_line = next(f).split(args.sep)
                    for i, val in enumerate(zip(header, next_line)):
                        print(f"{i + 1}. {val[0]} --> {val[1]}")
                    sys.exit(0)

                split_line = line.split(args.sep)
                vals = [split_line[x - 1] for x in all_vals if x is not None]

                outfile.write("%s\n" % " ".join(vals))
    except Exception as e:
        print("fatal: encountered an error -- %s" % e)
        sys.exit(1)
    finally:
        outfile.close()