genotype_stats.py

import sys

import matplotlib.pyplot as plt

if __name__ == '__main__':
    # genotpye_file lines look like:
    # line_nr_of_region_file gt1 gt2 ...
    genotpye_file = sys.argv[1]
    with_quality = len(sys.argv) > 2 and sys.argv[2] == "quality"

    gt_counter = {}
    gt_quality = {"0/0": [], "0/1": [], "1/1": []}
    with open(genotpye_file, "r") as f:
        for line in f.read().splitlines():
            sp = line.split(" ")
            gts = sp[1:]
            if with_quality:
                gts, quality = gts[::2], gts[1::2]
                for gt, q in zip(gts, quality):
                    gt_quality[gt].append(int(q))

            gts = sorted(gts) # avoids having entries for "0/1 1/1" AND "1/1 0/1"
            if (gt_comb := " ".join(gts)) in gt_counter:
                gt_counter[gt_comb] += 1
            else:
                gt_counter[gt_comb] = 1

    total = sum(list(gt_counter.values()))
    for gt_comb, count in gt_counter.items():
        print(f"{count:3d} ({count/total*100:05.2f}%) cases found for genotype combinations {gt_comb}")

    if with_quality:
        labels, data = [*zip(*gt_quality.items())]
        plt.boxplot(data)
        plt.xticks(range(1, len(labels)+1), labels)
        plt.xlabel("Genotypes")
        plt.ylabel("Quality Scores")
        plt.savefig("genotype_quality_distribution.png")