Add gaussian_mixture_model_karyotype_assignment function to assign sex karyotype using Gaussian mixture models

[jkgoodrich]

Modify get_ploidy_cutoffs to accept an expression for grouping samples into 'XX' and 'XY' groupings.
Modify annotate_sex to allow for sex karyotype imputation to use a gaussian mixture model to determine grouping of samples into 'XX' and 'XY' for use in get_ploidy_cutoffs.
Modify annotate_sex to make computing f-stat and inferring sex karyotype optional.

[jkgoodrich]

If helpful for testing:

import hail as hl
from gnomad_qc.v4.resources.basics import (
    calling_intervals,
    get_gnomad_v4_vds,
)
from bokeh.plotting import output_notebook, show

hl.init(default_reference='GRCh38')
output_notebook()

vds = get_gnomad_v4_vds(
    remove_hard_filtered_samples=False,
    remove_hard_filtered_samples_no_sex=True,
    test=True,
)
mt = hl.vds.to_merged_sparse_mt(vds)
calling_intervals=calling_intervals(interval_name='intersection', calling_interval_padding=50).ht()
freq_ht = hl.read_table("gs://gnomad/v4.0/sample_qc/exomes/gnomad.exomes.v4.0.f_stat_sites.ht")
sex_ht = annotate_sex(vds, included_intervals=calling_intervals, sites_ht=freq_ht, f_stat_cutoff=-1, aaf_expr="AF", gt_expr="LGT", use_gaussian_mixture_model=True)
sex_ht2 = annotate_sex(mt, included_intervals=calling_intervals, sites_ht=freq_ht, f_stat_cutoff=-1, aaf_expr="AF", gt_expr="LGT", use_gaussian_mixture_model=True)
sex_ht3 = annotate_sex(vds, included_intervals=calling_intervals, sites_ht=freq_ht, f_stat_cutoff=-1, aaf_expr="AF", gt_expr="LGT", use_gaussian_mixture_model=True, variants_only_x_ploidy=True)
sex_ht4 = annotate_sex(mt, included_intervals=calling_intervals, sites_ht=freq_ht, f_stat_cutoff=-1, aaf_expr="AF", gt_expr="LGT", use_gaussian_mixture_model=True, variants_only_x_ploidy=True)

show(hl.plot.scatter(sex_ht.chrX_ploidy, sex_ht.chrY_ploidy, label=sex_ht.sex_karyotype))
show(hl.plot.scatter(sex_ht2.chrX_ploidy, sex_ht2.chrY_ploidy, label=sex_ht2.sex_karyotype))
show(hl.plot.scatter(sex_ht3.chrX_ploidy, sex_ht3.chrY_ploidy, label=sex_ht3.sex_karyotype))
show(hl.plot.scatter(sex_ht4.chrX_ploidy, sex_ht4.chrY_ploidy, label=sex_ht4.sex_karyotype))

[mike-w-wilson]

Overall this looks good to me. Mostly just small suggestions on docs and a couple errors to raise. However, I hit a bug in the old code when testing your example 2 and 4 above, that I cannot comment on in the PR but hit when testing your example 2 and 4 above -- using a MT in the annotate_sex function. The call to impute_sex_ploidy returns a {contig}_mean_dp field and not an autosomal_mean_dp

gnomad_methods/gnomad/utils/sparse_mt.py

Line 672 in 7a64393

f"{chrom}_mean_dp": hl.agg.filter(

and

gnomad_methods/gnomad/utils/sparse_mt.py

Line 682 in 7a64393

f"{chrom}_mean_dp": hl.agg.sum(

. In L455 the function attempts to rename autosomal_mean_dp if variants only is true but that field doesnt exist, and if it is false, it attempts to rename to {contig}_mean_dp which already exists. Unless there is a subsequent PR that will update this and it was accidentally split apart from this PR, well need to fix it. Do you want me to make a ticket for this or should this PR cover it?

[jkgoodrich]

Yes, the bug you found does get fixed in the next PR, but I might as well fix it here too

[mike-w-wilson]

LGTM!