def main(args):
hl.init(log='/frequency_data_generation.log', default_reference='GRCh38')
logger.info("Reading sparse MT and metadata table...")
mt = get_gnomad_v3_mt(key_by_locus_and_alleles=True)
meta_ht = meta.ht().select('pop', 'sex', 'project_id', 'release', 'sample_filters')
if args.test:
logger.info("Filtering to chr20:1-1000000")
mt = hl.filter_intervals(mt, [hl.parse_locus_interval('chr20:1-1000000')])
mt = hl.experimental.sparse_split_multi(mt, filter_changed_loci=True)
logger.info("Annotating sparse MT with metadata...")
mt = mt.annotate_cols(meta=meta_ht[mt.s])
mt = mt.filter_cols(mt.meta.release)
samples = mt.count_cols()
logger.info(f"Running frequency table prep and generation pipeline on {samples} samples")
logger.info("Computing adj and sex adjusted genotypes.")
mt = mt.annotate_entries(
GT=adjusted_sex_ploidy_expr(mt.locus, mt.GT, mt.meta.sex),
adj=get_adj_expr(mt.GT, mt.GQ, mt.DP, mt.AD)
)
logger.info("Densify-ing...")
mt = hl.experimental.densify(mt)
mt = mt.filter_rows(hl.len(mt.alleles) > 1)
logger.info("Generating frequency data...")
mt = annotate_freq(
mt,
sex_expr=mt.meta.sex,
pop_expr=mt.meta.pop
)
# Select freq, FAF and popmax
faf, faf_meta = faf_expr(mt.freq, mt.freq_meta, mt.locus, POPS_TO_REMOVE_FOR_POPMAX)
mt = mt.select_rows(
'freq',
faf=faf,
popmax=pop_max_expr(mt.freq, mt.freq_meta, POPS_TO_REMOVE_FOR_POPMAX)
)
mt = mt.annotate_globals(faf_meta=faf_meta)
# Annotate quality metrics histograms, as these also require densifying
mt = mt.annotate_rows(
**qual_hist_expr(mt.GT, mt.GQ, mt.DP, mt.AD)
)
logger.info("Writing out frequency data...")
if args.test:
mt.rows().write("gs://gnomad-tmp/gnomad_freq/chr20_1_1000000_freq.ht", overwrite=True)
else:
mt.rows().write(freq.path, overwrite=args.overwrite)
def main(args):
hl.init(log='/frequency_data_generation.log', default_reference='GRCh38')
logger.info("Reading sparse MT and metadata table...")
mt = get_gnomad_v3_mt(key_by_locus_and_alleles=True)
meta_ht = meta.ht().select('pop', 'sex', 'project_id', 'release', 'sample_filters')
if args.test:
logger.info("Filtering to chr20:1-1000000")
mt = hl.filter_intervals(mt, [hl.parse_locus_interval('chr20:1-1000000')])
mt = hl.experimental.sparse_split_multi(mt, filter_changed_loci=True)
logger.info("Annotating sparse MT with metadata...")
mt = mt.annotate_cols(meta=meta_ht[mt.s])
mt = mt.filter_cols(mt.meta.release)
samples = mt.count_cols()
logger.info(f"Running frequency table prep and generation pipeline on {samples} samples")
logger.info("Computing adj and sex adjusted genotypes.")
mt = mt.annotate_entries(
GT=adjusted_sex_ploidy_expr(mt.locus, mt.GT, mt.meta.sex),
adj=get_adj_expr(mt.GT, mt.GQ, mt.DP, mt.AD)
)
logger.info("Densify-ing...")
mt = hl.experimental.densify(mt)
mt = mt.filter_rows(hl.len(mt.alleles) > 1)
logger.info("Setting het genotypes at sites with >1% AF (using v3.0 frequencies) and > 0.9 AB to homalt...")
# hotfix for depletion of homozygous alternate genotypes
# Using v3.0 AF to avoid an extra frequency calculation
# TODO: Using previous callset AF works for small incremental changes to a callset, but we need to revisit for large increments
freq_ht = freq.versions["3"].ht()
freq_ht = freq_ht.select(AF=freq_ht.freq[0].AF)
mt = mt.annotate_entries(
GT=hl.cond(
(freq_ht[mt.row_key].AF > 0.01)
& mt.GT.is_het()
& (mt.AD[1] / mt.DP > 0.9),
hl.call(1, 1),
mt.GT,
)
)
logger.info("Calculating InbreedingCoefficient...")
# NOTE: This is not the ideal location to calculate this, but added here to avoid another densify
mt = mt.annotate_rows(InbreedingCoeff=bi_allelic_site_inbreeding_expr(mt.GT))
logger.info("Generating frequency data...")
mt = annotate_freq(
mt,
sex_expr=mt.meta.sex,
pop_expr=mt.meta.pop
)
# Select freq, FAF and popmax
faf, faf_meta = faf_expr(mt.freq, mt.freq_meta, mt.locus, POPS_TO_REMOVE_FOR_POPMAX)
mt = mt.select_rows(
'InbreedingCoeff',
'freq',
faf=faf,
popmax=pop_max_expr(mt.freq, mt.freq_meta, POPS_TO_REMOVE_FOR_POPMAX)
)
mt = mt.annotate_globals(faf_meta=faf_meta)
# Annotate quality metrics histograms, as these also require densifying
mt = mt.annotate_rows(
**qual_hist_expr(mt.GT, mt.GQ, mt.DP, mt.AD)
)
logger.info("Writing out frequency data...")
if args.test:
mt.rows().write("gs://gnomad-tmp/gnomad_freq/chr20_1_1000000_freq.ht", overwrite=True)
else:
mt.rows().write(freq.path, overwrite=args.overwrite)
def main(args):
subsets = args.subsets
hl.init(
log=
f"/generate_frequency_data{'.' + '_'.join(subsets) if subsets else ''}.log",
default_reference="GRCh38",
)
invalid_subsets = []
n_subsets_use_subpops = 0
for s in subsets:
if s not in SUBSETS:
invalid_subsets.append(s)
if s in COHORTS_WITH_POP_STORED_AS_SUBPOP:
n_subsets_use_subpops += 1
if invalid_subsets:
raise ValueError(
f"{', '.join(invalid_subsets)} subset(s) are not one of the following official subsets: {SUBSETS}"
)
if n_subsets_use_subpops & (n_subsets_use_subpops != len(subsets)):
raise ValueError(
f"All or none of the supplied subset(s) should be in the list of cohorts that need to use subpops instead "
f"of pops in frequency calculations: {COHORTS_WITH_POP_STORED_AS_SUBPOP}"
)
try:
logger.info("Reading full sparse MT and metadata table...")
mt = get_gnomad_v3_mt(
key_by_locus_and_alleles=True,
release_only=not args.include_non_release,
samples_meta=True,
)
if args.test:
logger.info("Filtering to two partitions on chr20")
mt = hl.filter_intervals(
mt, [hl.parse_locus_interval("chr20:1-1000000")])
mt = mt._filter_partitions(range(2))
mt = hl.experimental.sparse_split_multi(mt, filter_changed_loci=True)
if args.include_non_release:
logger.info("Filtering MT columns to high quality samples")
total_sample_count = mt.count_cols()
mt = mt.filter_cols(mt.meta.high_quality)
high_quality_sample_count = mt.count_cols()
logger.info(
f"Filtered {total_sample_count - high_quality_sample_count} from the full set of {total_sample_count} "
f"samples...")
if subsets:
mt = mt.filter_cols(hl.any([mt.meta.subsets[s] for s in subsets]))
logger.info(
f"Running frequency generation pipeline on {mt.count_cols()} samples in {', '.join(subsets)} subset(s)..."
)
else:
logger.info(
f"Running frequency generation pipeline on {mt.count_cols()} samples..."
)
logger.info("Computing adj and sex adjusted genotypes...")
mt = mt.annotate_entries(
GT=adjusted_sex_ploidy_expr(mt.locus, mt.GT,
mt.meta.sex_imputation.sex_karyotype),
adj=get_adj_expr(mt.GT, mt.GQ, mt.DP, mt.AD),
)
logger.info("Densify-ing...")
mt = hl.experimental.densify(mt)
mt = mt.filter_rows(hl.len(mt.alleles) > 1)
# Temporary hotfix for depletion of homozygous alternate genotypes
logger.info(
"Setting het genotypes at sites with >1% AF (using v3.0 frequencies) and > 0.9 AB to homalt..."
)
# Load v3.0 allele frequencies to avoid an extra frequency calculation
# NOTE: Using previous callset AF works for small incremental changes to a callset, but we will need to revisit for large increments
freq_ht = get_freq(version="3").ht()
freq_ht = freq_ht.select(AF=freq_ht.freq[0].AF)
mt = mt.annotate_entries(GT=hl.cond(
(freq_ht[mt.row_key].AF > 0.01)
& mt.GT.is_het()
& (mt.AD[1] / mt.DP > 0.9),
hl.call(1, 1),
mt.GT,
))
logger.info("Generating frequency data...")
if subsets:
mt = annotate_freq(
mt,
sex_expr=mt.meta.sex_imputation.sex_karyotype,
pop_expr=mt.meta.population_inference.pop
if not n_subsets_use_subpops else
mt.meta.project_meta.project_subpop,
# NOTE: TGP and HGDP labeled populations are highly specific and are stored in the project_subpop meta field
)
# NOTE: no FAFs or popmax needed for subsets
mt = mt.select_rows("freq")
logger.info(
f"Writing out frequency data for {', '.join(subsets)} subset(s)..."
)
if args.test:
mt.rows().write(
get_checkpoint_path(
f"chr20_test_freq.{'_'.join(subsets)}"),
overwrite=True,
)
else:
mt.rows().write(get_freq(subset="_".join(subsets)).path,
overwrite=args.overwrite)
else:
logger.info("Computing age histograms for each variant...")
mt = mt.annotate_cols(age=hl.if_else(
hl.is_defined(mt.meta.project_meta.age),
mt.meta.project_meta.age,
mt.meta.project_meta.age_alt,
# NOTE: most age data is stored as integers in 'age' annotation, but for a select number of samples, age is stored as a bin range and 'age_alt' corresponds to an integer in the middle of the bin
))
mt = mt.annotate_rows(**age_hists_expr(mt.adj, mt.GT, mt.age))
# Compute callset-wide age histogram global
mt = mt.annotate_globals(age_distribution=mt.aggregate_cols(
hl.agg.hist(mt.age, 30, 80, 10)))
mt = annotate_freq(
mt,
sex_expr=mt.meta.sex_imputation.sex_karyotype,
pop_expr=mt.meta.population_inference.pop,
downsamplings=DOWNSAMPLINGS,
)
# Remove all loci with raw AC=0
mt = mt.filter_rows(mt.freq[1].AC > 0)
logger.info("Calculating InbreedingCoeff...")
# NOTE: This is not the ideal location to calculate this, but added here to avoid another densify
mt = mt.annotate_rows(
InbreedingCoeff=bi_allelic_site_inbreeding_expr(mt.GT))
logger.info("Computing filtering allele frequencies and popmax...")
faf, faf_meta = faf_expr(mt.freq, mt.freq_meta, mt.locus,
POPS_TO_REMOVE_FOR_POPMAX)
mt = mt.select_rows(
"InbreedingCoeff",
"freq",
faf=faf,
popmax=pop_max_expr(mt.freq, mt.freq_meta,
POPS_TO_REMOVE_FOR_POPMAX),
)
mt = mt.annotate_globals(
faf_meta=faf_meta,
faf_index_dict=make_faf_index_dict(faf_meta))
mt = mt.annotate_rows(popmax=mt.popmax.annotate(
faf95=mt.faf[mt.faf_meta.index(
lambda x: x.values() == ["adj", mt.popmax.pop])].faf95))
logger.info("Annotating quality metrics histograms...")
# NOTE: these are performed here as the quality metrics histograms also require densifying
mt = mt.annotate_rows(
qual_hists=qual_hist_expr(mt.GT, mt.GQ, mt.DP, mt.AD, mt.adj))
ht = mt.rows()
ht = ht.annotate(
qual_hists=hl.Struct(
**{
i.replace("_adj", ""): ht.qual_hists[i]
for i in ht.qual_hists if "_adj" in i
}),
raw_qual_hists=hl.Struct(**{
i: ht.qual_hists[i]
for i in ht.qual_hists if "_adj" not in i
}),
)
logger.info("Writing out frequency data...")
if args.test:
ht.write(get_checkpoint_path("chr20_test_freq"),
overwrite=True)
else:
ht.write(get_freq().path, overwrite=args.overwrite)
finally:
logger.info("Copying hail log to logging bucket...")
hl.copy_log(f"{qc_temp_prefix()}logs/")
def generate_frequency_data(
mt: hl.MatrixTable,
calculate_downsampling: bool = False,
calculate_by_platform: bool = False) -> Tuple[hl.Table, hl.Table]:
"""
:param MatrixTable mt: Input MatrixTable
:param bool calculate_downsampling: Calculate frequencies for downsampled data
:param bool calculate_by_platform: Calculate frequencies for PCR-free data
"""
if calculate_downsampling:
mt, downsamplings = generate_downsamplings_cumulative(mt)
print(f'Got {len(downsamplings)} downsamplings: {downsamplings}')
cut_dict = {
'pop':
hl.agg.filter(hl.is_defined(mt.meta.pop), hl.agg.counter(mt.meta.pop)),
'sex':
hl.agg.filter(hl.is_defined(mt.meta.sex),
hl.agg.collect_as_set(mt.meta.sex)),
'subpop':
hl.agg.filter(
hl.is_defined(mt.meta.subpop) & hl.is_defined(mt.meta.pop),
hl.agg.collect_as_set(
hl.struct(subpop=mt.meta.subpop, pop=mt.meta.pop)))
}
if calculate_by_platform:
cut_dict['platform'] = hl.agg.filter(
hl.is_defined(mt.meta.qc_platform),
hl.agg.collect_as_set(mt.meta.qc_platform))
cut_data = mt.aggregate_cols(hl.struct(**cut_dict))
sample_group_filters = [({}, True)]
sample_group_filters.extend([({
'pop': pop
}, mt.meta.pop == pop) for pop in cut_data.pop] + [({
'sex': sex
}, mt.meta.sex == sex) for sex in cut_data.sex] +
[({
'pop': pop,
'sex': sex
}, (mt.meta.sex == sex) & (mt.meta.pop == pop))
for sex in cut_data.sex
for pop in cut_data.pop] +
[({
'subpop': subpop.subpop,
'pop': subpop.pop
}, mt.meta.subpop == subpop.subpop)
for subpop in cut_data.subpop])
if calculate_by_platform:
sample_group_filters.extend([({
'platform': str(platform)
}, mt.meta.qc_platform == platform) for platform in cut_data.platform])
if calculate_downsampling:
sample_group_filters.extend([({
'downsampling': str(ds),
'pop': 'global'
}, mt.downsampling.global_idx < ds) for ds in downsamplings])
sample_group_filters.extend([
({
'downsampling': str(ds),
'pop': pop
}, (mt.downsampling.pop_idx < ds) & (mt.meta.pop == pop))
for ds in downsamplings for pop, pop_count in cut_data.pop.items()
if ds <= pop_count
])
mt = mt.select_cols(group_membership=tuple(x[1]
for x in sample_group_filters),
project_id=mt.meta.project_id,
age=mt.meta.age)
mt = mt.select_rows()
frequency_expression = []
meta_expressions = []
for i in range(len(sample_group_filters)):
subgroup_dict = sample_group_filters[i][0]
subgroup_dict['group'] = 'adj'
call_stats = hl.agg.filter(mt.group_membership[i] & mt.adj,
hl.agg.call_stats(mt.GT, mt.alleles))
call_stats_bind = hl.bind(
lambda cs: cs.annotate(AC=cs.AC[1],
AF=cs.AF[1],
homozygote_count=cs.homozygote_count[1]),
call_stats)
frequency_expression.append(call_stats_bind)
meta_expressions.append(subgroup_dict)
raw_stats = hl.agg.call_stats(mt.GT, mt.alleles)
raw_stats_bind = hl.bind(
lambda cs: cs.annotate(
AC=cs.AC[1], AF=cs.AF[1], homozygote_count=cs.homozygote_count[1]),
raw_stats)
frequency_expression.insert(1, raw_stats_bind)
meta_expressions.insert(1, {'group': 'raw'})
print(f'Calculating {len(frequency_expression)} aggregators...')
global_expression = {'freq_meta': meta_expressions}
mt = mt.annotate_rows(
freq=frequency_expression,
age_hist_het=hl.agg.filter(mt.adj & mt.GT.is_het(),
hl.agg.hist(mt.age, 30, 80, 10)),
age_hist_hom=hl.agg.filter(mt.adj & mt.GT.is_hom_var(),
hl.agg.hist(mt.age, 30, 80, 10)))
if calculate_downsampling:
global_expression['downsamplings'] = downsamplings
mt = mt.annotate_globals(**global_expression)
sample_data = mt.cols()
pops = set(cut_data.pop.keys())
[pops.discard(x) for x in POPS_TO_REMOVE_FOR_POPMAX]
mt = mt.annotate_rows(
popmax=pop_max_expr(mt.freq,
mt.freq_meta,
pops_to_exclude=set(POPS_TO_REMOVE_FOR_POPMAX)),
faf=add_faf_expr(mt.freq, mt.freq_meta, mt.locus, populations=pops),
project_max=project_max_expr(mt.project_id, mt.GT, mt.alleles))
return mt.rows(), sample_data