sample_qc.py

"""

Compute stratified sample QC metrics on filtered variants. Apply stratified
sample filters based on computed QC metrics. Export result as separated
HailTables. Optionally, export results as GBZ-compresed TSV files.

usage: sample_qc.py [-h] [--exome_cohort EXOME_COHORT] [--write_to_file]
                    [--overwrite] [--default_ref_genome DEFAULT_REF_GENOME]

optional arguments:
  -h, --help            show this help message and exit
  --exome_cohort EXOME_COHORT
                        One of <chd_ukbb> or <chd_ddd>
  --write_to_file       Write output to BGZ-compressed file
  --overwrite           Overwrite pre-existing data
  --default_ref_genome DEFAULT_REF_GENOME
                        Default reference genome to start Hail

"""

import argparse
import logging
from typing import Dict, List, Optional

import hail as hl
from gnomad.utils.gen_stats import merge_stats_counters_expr
from hail.utils.misc import divide_null

from utils.expressions import bi_allelic_expr
from utils.filter import (filter_low_conf_regions,
                          filter_to_autosomes,
                          remove_telomeres_centromes)

from utils.data_utils import (get_sample_pop_qc,
                              get_sample_qc_ht_path,
                              get_mt_data)

from utils.config import NFS_DIR

logging.basicConfig(format="%(levelname)s (%(name)s %(lineno)s): %(message)s")
logger = logging.getLogger("Sample QC")
logger.setLevel(logging.INFO)

# hdfs_dir = HDFS_DIR  # set via WES_HDFS_DIR env var
nfs_dir = NFS_DIR


def merge_sample_qc_expr(
        sample_qc_exprs: List[hl.expr.StructExpression],
) -> hl.expr.StructExpression:
    """
    Creates an expression that merges results from non-overlapping strata of hail.sample_qc

    E.g.:

    - Compute autosomes and sex chromosomes metrics separately, then merge results
    - Compute bi-allelic and multi-allelic metrics separately, then merge results

    Note regarding the merging of ``dp_stats`` and ``gq_stats``:
    Because ``n`` is needed to aggregate ``stdev``, ``n_called`` is used for this purpose.
    This should work very well on a standard GATK VCF and it essentially assumes that:

    - samples that are called have `DP` and `GQ` fields
    - samples that are not called do not have `DP` and `GQ` fields

    Even if these assumptions are broken for some genotypes, it shouldn't matter too much.

    :param sample_qc_exprs: List of sample QC struct expressions for each stratification
    :return: Combined sample QC results
    """

    # List of metrics that can be aggregated by summing
    additive_metrics = [
        "n_called",
        "n_not_called",
        "n_hom_ref",
        "n_het",
        "n_hom_var",
        "n_snp",
        "n_insertion",
        "n_deletion",
        "n_singleton",
        "n_transition",
        "n_transversion",
        "n_star",
    ]

    # List of metrics that are ratio of summed metrics (name, nominator, denominator)
    ratio_metrics = [
        ("call_rate", "n_called", "n_not_called"),
        ("r_ti_tv", "n_transition", "n_transversion"),
        ("r_het_hom_var", "n_het", "n_hom_var"),
        ("r_insertion_deletion", "n_insertion", "n_deletion"),
    ]

    # List of metrics that are struct generated by a stats counter
    stats_metrics = ["gq_stats", "dp_stats"]

    # Gather metrics present in sample qc fields
    sample_qc_fields = set(sample_qc_exprs[0])
    for sample_qc_expr in sample_qc_exprs[1:]:
        sample_qc_fields = sample_qc_fields.union(set(sample_qc_expr))

    # Merge additive metrics in sample qc fields
    merged_exprs = {
        metric: hl.sum([sample_qc_expr[metric] for sample_qc_expr in sample_qc_exprs])
        for metric in additive_metrics
        if metric in sample_qc_fields
    }

    # Merge ratio metrics in sample qc fields
    merged_exprs.update(
        {
            metric: hl.float64(divide_null(merged_exprs[nom], merged_exprs[denom]))
            for metric, nom, denom in ratio_metrics
            if nom in sample_qc_fields and denom in sample_qc_fields
        }
    )

    # Merge stats counter metrics in sample qc fields
    # Use n_called as n for DP and GQ stats
    if "n_called" in sample_qc_fields:
        merged_exprs.update(
            {
                metric: merge_stats_counters_expr(
                    [
                        sample_qc_expr[metric].annotate(n=sample_qc_expr.n_called)
                        for sample_qc_expr in sample_qc_exprs
                    ]
                ).drop("n")
                for metric in stats_metrics
            }
        )

    return hl.struct(**merged_exprs)


def compute_stratified_sample_qc(
        mt: hl.MatrixTable,
        strata: Dict[str, hl.expr.BooleanExpression],
        tmp_ht_prefix: Optional[str],
        gt_expr: Optional[hl.expr.CallExpression],
) -> hl.Table:
    """
    Runs hl.sample_qc on different strata and then also merge the results into a single expression.
    Note that strata should be non-overlapping,  e.g. SNV vs indels or bi-allelic vs multi-allelic

    :param mt: Input MT
    :param strata: Strata names and filtering expressions
    :param tmp_ht_prefix: Optional path prefix to write the intermediate strata results to (recommended for larger datasets)
    :param gt_expr: Optional entry field storing the genotype (if not specified, then it is assumed that it is stored in mt.GT)
    :return: Sample QC table, including strat-specific numbers
    """
    mt = mt.select_rows(**strata)

    # if gt_expr is not None:
    #    mt = mt.select_entries(GT=gt_expr)
    # else:
    #    mt = mt.select_entries("GT")

    strat_hts = {}
    for strat in strata:
        strat_sample_qc_ht = hl.sample_qc(mt.filter_rows(mt[strat])).cols()
        if tmp_ht_prefix is not None:
            strat_sample_qc_ht = strat_sample_qc_ht.checkpoint(
                tmp_ht_prefix + f"_{strat}.ht", overwrite=True
            )
        else:
            strat_sample_qc_ht = strat_sample_qc_ht.persist()
        strat_hts[strat] = strat_sample_qc_ht

    sample_qc_ht = strat_hts.pop(list(strata)[0])
    sample_qc_ht = sample_qc_ht.select(
        **{f"{list(strata)[0]}_sample_qc": sample_qc_ht.sample_qc},
        **{
            f"{strat}_sample_qc": strat_hts[strat][sample_qc_ht.key].sample_qc
            for strat in list(strata)[1:]
        },
    )

    sample_qc_ht = sample_qc_ht.annotate(
        sample_qc=merge_sample_qc_expr(list(sample_qc_ht.row_value.values()))
    )

    return sample_qc_ht


def compute_sample_qc(mt: hl.MatrixTable) -> hl.Table:
    """
    Perform sample QC on the raw split matrix table using `compute_stratified_sample_qc`.
    :return: Table containing sample QC metrics
    :rtype: hl.Table
    """
    logger.info("Computing sample QC")

    # mt = mt.select_entries("GT")

    # Remove centromeres and telomeres incase they were included
    mt = filter_low_conf_regions(
        mt,
        filter_lcr=True,  # TODO: include also decoy and low coverage exome regions
        filter_segdup=True
    )

    # filter to autosomes
    mt = filter_to_autosomes(mt)

    # filter telomeres/centromes
    mt = remove_telomeres_centromes(mt)

    # filter coding variants
    # mt = filter_cds_regions(mt)

    sample_qc_ht = compute_stratified_sample_qc(
        mt,
        strata={
            "bi_allelic": bi_allelic_expr(mt),
            "multi_allelic": ~bi_allelic_expr(mt),
        }, tmp_ht_prefix=None,
        gt_expr=None
    )

    # Remove annotations that cannot be computed from the sparse format
    # sample_qc_ht = sample_qc_ht.annotate(
    #    **{
    #        x: sample_qc_ht[x].drop(
    #            "n_called", "n_not_called", "n_filtered", "call_rate"
    #        )
    #        for x in sample_qc_ht.row_value
    #   }
    # )

    return sample_qc_ht.repartition(100)


def compute_stratified_metrics_filter(ht: hl.Table, qc_metrics: List[str], strata: List[str] = None) -> hl.Table:
    """
    Compute median, MAD, and upper and lower thresholds for each metric used in pop- and platform-specific outlier filtering
    :param MatrixTable ht: HT containing relevant sample QC metric annotations
    :param list qc_metrics: list of metrics for which to compute the critical values for filtering outliers
    :param list of str strata: List of annotations used for stratification. These metrics should be discrete types!
    :return: Table grouped by pop and platform, with upper and lower threshold values computed for each sample QC metric
    :rtype: Table
    """

    def make_pop_filters_expr(ht: hl.Table, qc_metrics: List[str]) -> hl.expr.SetExpression:
        return hl.set(hl.filter(lambda x: hl.is_defined(x),
                                [hl.or_missing(ht[f'fail_{metric}'], metric) for metric in qc_metrics]))

    ht = ht.select(*strata, **ht.sample_qc.select(*qc_metrics)).key_by('s').persist()

    def get_metric_expr(ht, metric):
        metric_values = hl.agg.collect(ht[metric])
        metric_median = hl.median(metric_values)
        metric_mad = 1.4826 * hl.median(hl.abs(metric_values - metric_median))
        return hl.struct(
            median=metric_median,
            mad=metric_mad,
            upper=metric_median + 4 * metric_mad if metric != 'callrate' else 1,
            lower=metric_median - 4 * metric_mad if metric != 'callrate' else 0.99
        )

    agg_expr = hl.struct(**{metric: get_metric_expr(ht, metric) for metric in qc_metrics})
    if strata:
        ht = ht.annotate_globals(
            metrics_stats=ht.aggregate(hl.agg.group_by(hl.tuple([ht[x] for x in strata]), agg_expr)))
    else:
        ht = ht.annotate_globals(metrics_stats={(): ht.aggregate(agg_expr)})

    strata_exp = hl.tuple([ht[x] for x in strata]) if strata else hl.tuple([])

    fail_exprs = {
        f'fail_{metric}':
            (ht[metric] >= ht.metrics_stats[strata_exp][metric].upper) |
            (ht[metric] <= ht.metrics_stats[strata_exp][metric].lower)
        for metric in qc_metrics}
    ht = ht.transmute(**fail_exprs)
    pop_platform_filters = make_pop_filters_expr(ht, qc_metrics)
    return ht.annotate(pop_platform_filters=pop_platform_filters)


def checkpoint_sample_qc(
        ht: hl.Table,
        dataset: str,
        overwrite: bool,
) -> hl.Table:
    """Checkpoint the sample QC table with metrics to disk."""
    ht = ht.checkpoint(
        get_sample_qc_ht_path(dataset=dataset, part='high_conf_autosomes'),
        overwrite=overwrite,
        _read_if_exists=not overwrite
    )

    return ht


def export_sample_qc_tsv(ht: hl.Table, dataset: str, write_to_file: bool) -> None:
    """Optionally export the sample QC table as a flattened BGZ-compressed TSV."""
    # Export HT to file
    if write_to_file:
        (ht.flatten().export(
            f"{get_sample_qc_ht_path(dataset=dataset, part='high_conf_autosomes')}.tsv.bgz")
         )


def annotate_pop_platform(ht: hl.Table, dataset: str) -> hl.Table:
    """Annotate the sample QC table with predicted population and QC platform labels."""
    # annotate sample population and platform qc info
    pop_qc = hl.read_table(
        get_sample_qc_ht_path(dataset=dataset, part='population_qc')
    )
    platform_qc = hl.read_table(
        get_sample_qc_ht_path(dataset=dataset, part='platform_pca')
    )

    ann_expr = {'qc_pop': pop_qc[ht.s].predicted_pop,
                'qc_platform': platform_qc[ht.s].qc_platform
                }

    return ht.annotate(**ann_expr)


def compute_and_checkpoint_stratified_filter(
        ht: hl.Table,
        dataset: str,
        overwrite: bool,
        write_to_file: bool,
) -> hl.Table:
    """Compute stratified metrics filter, checkpoint the result, and optionally export a TSV."""
    # Apply stratified sample filters based on defined QC metrics
    exome_qc_metrics = ['n_snp',
                        'r_ti_tv',
                        'r_insertion_deletion',
                        'n_insertion',
                        'n_deletion',
                        'r_het_hom_var']

    print('Computing stratified metrics filters...')
    exome_pop_platform_filter_ht = compute_stratified_metrics_filter(ht,
                                                                     exome_qc_metrics,
                                                                     ['qc_pop', 'qc_platform'])

    exome_pop_platform_filter_ht = exome_pop_platform_filter_ht.checkpoint(
        get_sample_qc_ht_path(dataset=dataset, part='stratified_metrics_filter'),
        overwrite=overwrite,
        _read_if_exists=not overwrite
    )

    # Export HT to file
    if write_to_file:
        (exome_pop_platform_filter_ht.export(
            f"{get_sample_qc_ht_path(dataset=dataset, part='stratified_metrics_filter')}.tsv.bgz")
         )

    return exome_pop_platform_filter_ht


def main(args):
    # Start Hail
    hl.init(default_reference=args.default_ref_genome)

    # Import unfiltered split MT
    mt = get_mt_data(dataset=args.exome_cohort, part='raw')

    # Compute stratified sample_qc (biallelic and multi-allelic sites)
    sample_qc_ht = compute_sample_qc(mt)

    # Write HT with sample QC metrics
    sample_qc_ht = checkpoint_sample_qc(
        sample_qc_ht, args.exome_cohort, args.overwrite
    )

    sample_qc_ht = annotate_pop_platform(sample_qc_ht, dataset=args.exome_cohort)

    export_sample_qc_tsv(sample_qc_ht, args.exome_cohort, args.write_to_file)

    compute_and_checkpoint_stratified_filter(
        sample_qc_ht, args.exome_cohort, args.overwrite, args.write_to_file
    )

    # Stop Hail
    hl.stop()

    print("Finished!")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--exome_cohort', help="One of <chd_ukbb> or <chd_ddd>",
                        type=str, default=None)
    parser.add_argument('--write_to_file', help='Write output to BGZ-compressed file',
                        action='store_true')
    parser.add_argument('--overwrite', help='Overwrite pre-existing data',
                        action='store_true')
    parser.add_argument('--default_ref_genome', help='Default reference genome to start Hail',
                        type=str, default='GRCh38')

    args = parser.parse_args()

    main(args)