From 65678ed1a4baf7253f83b9faf807a5a69824834f Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Mon, 10 Mar 2025 13:32:46 -0600 Subject: [PATCH 01/38] Large rearrangement of the constraint pipeline so it's more streamlined --- .../pipeline/constraint_pipeline.py | 948 ++++---- gnomad_constraint/resources/resource_utils.py | 177 +- gnomad_constraint/utils/constraint.py | 1946 +++++++++-------- 3 files changed, 1569 insertions(+), 1502 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 77163354..ac6cee23 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -23,29 +23,33 @@ import argparse import logging -from typing import List +from typing import List, Optional import hail as hl -from gnomad.resources.grch38.gnomad import DOWNSAMPLINGS, all_sites_an -from gnomad.utils.constraint import build_models, explode_downsamplings_oe -from gnomad.utils.filtering import filter_x_nonpar, filter_y_nonpar +from gnomad.resources.grch38.gnomad import all_sites_an +from gnomad.utils.constraint import ( + annotate_with_mu, + assemble_constraint_context_ht, + build_models, + explode_downsamplings_oe, +) from gnomad.utils.reference_genome import get_reference_genome +from gnomad.utils.vep import update_loftee_end_trunc_filter from gnomad_qc.resource_utils import ( PipelineResourceCollection, PipelineStepResourceCollection, ) -from hail.utils.misc import new_temp_file import gnomad_constraint.resources.resource_utils as constraint_res from gnomad_constraint.utils.constraint import ( - add_vep_context_annotations, - annotate_context_ht, - apply_models, + aggregate_per_variant_expected_ht, calculate_gerp_cutoffs, calculate_mu_by_downsampling, compute_constraint_metrics, - create_observed_and_possible_ht, + create_per_variant_expected_ht, + create_training_set, prepare_ht_for_constraint_calculations, + print_global_struct, ) logging.basicConfig( @@ -58,14 +62,12 @@ def filter_for_test( ht: hl.Table, - data_type: str, use_gene_list: bool = False, ) -> hl.Table: """ Filter `ht` to chr20, chrX, and chrY or a gene list for testing. :param ht: Table to filter. - :param data_type: Data type of `ht`. :param use_gene_list: Whether to use a gene list for testing instead of all of chr20, chrX, and chrY for testing. :return: Filtered Table for testing. @@ -94,161 +96,240 @@ def filter_for_test( hl.parse_locus_interval(c, reference_genome=rg) for c in [rg.contigs[19], rg.x_contigs[0], rg.y_contigs[0]] ] - logger.info( - "Filtering the %s HT to chr20, chrX, and chrY for testing...", - data_type, - ) + logger.info("Filtering the context HT to chr20, chrX, and chrY for testing...") + ht = hl.filter_intervals(ht, keep) return ht +def run_prepare_context( + resources: PipelineResourceCollection, + test: bool = False, + test_gene_list: bool = False, +) -> hl.Table: + """ + Annotate the context Table with coverage, AN, and frequency annotations. + + Uses `assemble_constraint_context_ht` to annotate the context Table with annotations + that are used in downstream steps of the constraint pipeline. + + :param resources: PipelineResourceCollection containing resources for the constraint + pipeline. + :param test: Whether to filter the context Table to only chr20, chrX, and chrY for + testing. + :param test_gene_list: Whether to filter the context Table to a gene list for + testing. + :return: Annotated context Table. + """ + # We use naive_coalesce on the context Table because it has a large number of + # partitions which caused some issues with Hail 0.2.133. 5000 partitions was a + # number that worked well for the context Table in the past. + ht = resources.context_ht.ht().naive_coalesce(5000) + + if test: + ht = filter_for_test(ht, use_gene_list=test_gene_list) + + def _build_ht_dict(ht_name: str, keep: List[str] = None): + dts = ["exomes", "genomes"] + hts = {d: getattr(resources, f"{d}_{ht_name}_ht").ht() for d in dts} + return {d: t.select(*keep) for d, t in hts.items()} if keep else hts + + # There was a bug in the GERP cutoffs used to filter transcripts with the + # "END_TRUNC" filter in the LOFTEE VEP plugin resulting in some transcripts + # being considered "HC" when they should have been "LC". We use the + # `update_loftee_end_trunc_filter` function to correct this issue. + ht = ht.annotate( + vep=ht.vep.annotate( + transcript_consequences=update_loftee_end_trunc_filter( + ht.vep.transcript_consequences + ) + ) + ) + ht = assemble_constraint_context_ht( + ht, + coverage_hts=_build_ht_dict("coverage"), + an_hts=_build_ht_dict("an"), + freq_hts=_build_ht_dict("sites", ["freq"]), + filter_hts=_build_ht_dict("sites", ["filters"]), + methylation_ht=resources.methylation_ht.ht(), + gerp_ht=constraint_res.get_gerp_ht(get_reference_genome(ht.locus).name), + transformation_funcs=None, + ) + + # Add annotation for exome coverage and genomic region (autosome/PAR, X non-PAR, + # Y non-PAR). + genomic_region_expr = ( + hl.case() + .when(ht.locus.in_autosome_or_par(), "autosome_or_par") + .when(ht.locus.in_x_nonpar(), "chrx_nonpar") + .when(ht.locus.in_y_nonpar(), "chry_nonpar") + .or_missing() + ) + + # TODO: Make these resources. + am_ht = hl.read_table( + "gs://gnomad/v4.1/constraint/resources/alpha_missense_filters.ht" + ) + adj_r_ht = hl.read_table( + "gs://gnomad/v4.1/constraint/resources/ncc_adj_r_per_base_WG.ht" + ).key_by("locus") + multisfs_ht = hl.read_table( + "gs://gnomad/v4.1/constraint/resources/julia/constraint/multisfs.dedup.ht" + ) + am_keyed = am_ht[ht.key] + ht = ht.annotate( + coverage=hl.struct( + exomes=ht.coverage.exomes.select("mean", "median_approx"), + genomes=ht.coverage.genomes.select("mean", "median_approx"), + ), + AN=hl.struct( + exomes=ht.AN.exomes[0], + genomes=ht.AN.genomes[0], + ), + genomic_region=genomic_region_expr, + alpha_missense=hl.struct( + pathogenicity=am_keyed.am_pathogenicity, + per_98=hl.or_else(am_keyed.am_per_98, False), + per_99=hl.or_else(am_keyed.am_per_99, False), + over_0_999=hl.or_else(am_keyed.am_0_999, False), + ), + adj_r=adj_r_ht[ht.locus].adj_r, + sfs_bin=multisfs_ht[ht.key].Freq_bin_9, + ) + + return ht + + def get_constraint_resources( version: str, - use_v2_release_mutation_ht: bool, - use_v2_release_context_ht: bool, custom_vep_annotation: str, overwrite: bool, test: bool, models: List[str] = ["plateau", "coverage"], + post_fix: Optional[str] = None, ) -> PipelineResourceCollection: """ Get PipelineResourceCollection for all resources needed in the constraint pipeline. :param version: Version of constraint resources to use. - :param use_v2_release_mutation_ht: Whether to use the v2 release mutation ht. - :param use_v2_release_context_ht: Whether to use the v2 release context ht. :param custom_vep_annotation: Custom VEP annotation to use for applying models resources. :param overwrite: Whether to overwrite existing resources. :param test: Whether to use test resources. :param models: List of models to use. Default is ["plateau", "coverage"]. + :param post_fix: Optional post-fix to append to resource paths. :return: PipelineResourceCollection containing resources for all steps of the constraint pipeline. """ - data_types = constraint_res.DATA_TYPES - regions = constraint_res.GENOMIC_REGIONS # Initialize constraint pipeline resource collection. constraint_pipeline = PipelineResourceCollection( pipeline_name="constraint", overwrite=overwrite, ) - # Make dictionary for allele number tables. - an_hts = {} - if int(version[0]) >= 4: - an_hts["exomes_an_ht"] = all_sites_an("exomes") - an_hts["genomes_an_ht"] = all_sites_an("genomes") - # Create resource collection for each step of the constraint pipeline. context_res = constraint_res.get_vep_context_ht(version) context_build = get_reference_genome(context_res.ht().locus).name + + # Make dictionary for prepare_context input Tables. + input_hts = { + "context_ht": context_res, + "methylation_ht": constraint_res.get_methylation_ht(context_build), + } + for d in ["exomes", "genomes"]: + input_hts[f"{d}_coverage_ht"] = constraint_res.get_coverage_ht(d, version) + input_hts[f"{d}_sites_ht"] = constraint_res.get_sites_resource(d, version) + input_hts[f"{d}_an_ht"] = all_sites_an(d) + prepare_context = PipelineStepResourceCollection( "--prepare-context-ht", output_resources={ "annotated_context_ht": constraint_res.get_annotated_context_ht( - version, use_v2_release_context_ht, test + version, test, post_fix ) }, - input_resources={ - "gnomAD resources": { - "context_ht": context_res, - "exomes_coverage_ht": constraint_res.get_coverage_ht("exomes", version), - "genomes_coverage_ht": constraint_res.get_coverage_ht( - "genomes", version - ), - "methylation_ht": constraint_res.get_methylation_ht(context_build), - **an_hts, - }, - }, + input_resources={"gnomAD resources": input_hts}, ) - # For genomes need a preprocessed ht for autosome_par. - # For exomes and context need a preprocessed ht for autosome_par, chrX, - # and chrY. preprocess_data = PipelineStepResourceCollection( - "--preprocess-data", + "preprocess data for downstream steps", output_resources={ - f"preprocessed_{r}_{d}_ht": constraint_res.get_preprocessed_ht( - d, version, r, test - ) - for r in regions - for d in data_types - if (r == "autosome_par") | (d != "genomes") + "temp_preprocess_data_ht": constraint_res.get_preprocessed_ht( + version, test, post_fix + ), }, pipeline_input_steps=[prepare_context], - add_input_resources={ - "gnomAD sites resources": { - f"{d}_sites_ht": constraint_res.get_sites_resource(d, version) - for d in data_types - if d != "context" - } - }, ) calculate_gerp_cutoffs = PipelineStepResourceCollection( "--calculate-gerp-cutoffs", output_resources={}, - pipeline_input_steps=[preprocess_data], + pipeline_input_steps=[prepare_context], ) calculate_mutation_rate = PipelineStepResourceCollection( "--calculate-mutation-rate", output_resources={ - "mutation_ht": constraint_res.get_mutation_ht( - version, test, use_v2_release_mutation_ht - ) + "mutation_ht": constraint_res.get_mutation_ht(version, test, post_fix) }, pipeline_input_steps=[preprocess_data], ) create_training_set = PipelineStepResourceCollection( "--create-training-set", output_resources={ - **{ - f"train_{r}_ht": constraint_res.get_training_dataset(version, r, test) - for r in regions - }, - **{ - f"train_{r}_tsv": constraint_res.get_training_tsv_path(version, r, test) - for r in regions - }, + f"train_ht": constraint_res.get_training_dataset(version, test, post_fix), + f"train_tsv": constraint_res.get_training_tsv_path(version, test, post_fix), }, pipeline_input_steps=[preprocess_data, calculate_mutation_rate], ) build_models = PipelineStepResourceCollection( "--build-models", output_resources={ - f"model_{r}_{m}": constraint_res.get_models(m, version, r, test) + f"model_{m}": constraint_res.get_models(m, version, test, post_fix) for m in models - for r in regions }, pipeline_input_steps=[create_training_set], ) - apply_models = PipelineStepResourceCollection( - "--apply-models", + apply_models_per_variant = PipelineStepResourceCollection( + "--apply-models-per-variant", output_resources={ - f"apply_{r}_ht": constraint_res.get_predicted_proportion_observed_dataset( - custom_vep_annotation, version, r, test + "per_variant_apply_ht": constraint_res.get_per_variant_expected_dataset( + custom_vep_annotation, version, test, post_fix ) - for r in regions }, pipeline_input_steps=[preprocess_data, calculate_mutation_rate, build_models], ) + aggregate_per_variant_expected = PipelineStepResourceCollection( + "--aggregate-per-variant-expected", + output_resources={ + f"apply_ht": constraint_res.get_apply_models( + custom_vep_annotation, version, test, post_fix + ) + }, + pipeline_input_steps=[ + apply_models_per_variant, + calculate_mutation_rate, + build_models, + ], + ) compute_constraint_metrics = PipelineStepResourceCollection( "--compute-constraint-metrics", output_resources={ "constraint_metrics_ht": constraint_res.get_constraint_metrics_dataset( - version, test + custom_vep_annotation, version, test, post_fix ) }, - pipeline_input_steps=[apply_models], + pipeline_input_steps=[aggregate_per_variant_expected], ) export_tsv = PipelineStepResourceCollection( "--export-tsv", output_resources={ "constraint_metrics_tsv": constraint_res.get_constraint_tsv_path( - version, test + version, test, post_fix ), "downsampling_constraint_metrics_tsv": ( - constraint_res.get_downsampling_constraint_tsv_path(version, test) + constraint_res.get_downsampling_constraint_tsv_path( + version, test, post_fix + ) ), }, pipeline_input_steps=[compute_constraint_metrics], @@ -263,7 +344,8 @@ def get_constraint_resources( "calculate_mutation_rate": calculate_mutation_rate, "create_training_set": create_training_set, "build_models": build_models, - "apply_models": apply_models, + "apply_models_per_variant": apply_models_per_variant, + "aggregate_per_variant_expected": aggregate_per_variant_expected, "compute_constraint_metrics": compute_constraint_metrics, "export_tsv": export_tsv, } @@ -278,72 +360,37 @@ def main(args): log="/constraint_pipeline.log", tmp_dir="gs://gnomad-tmp-4day", ) - regions = constraint_res.GENOMIC_REGIONS version = args.version test_gene_list = args.test_gene_list test = args.test or test_gene_list + post_fix = args.post_fix overwrite = args.overwrite - - max_af = args.max_af - pops = args.pops - use_v2_release_mutation_ht = args.use_v2_release_mutation_ht custom_vep_annotation = args.custom_vep_annotation - gerp_lower_cutoff = args.gerp_lower_cutoff - gerp_upper_cutoff = args.gerp_upper_cutoff - coverage_metric = args.coverage_metric - coverage_model_type = args.coverage_model_type + skip_coverage_model = args.skip_coverage_model + log10_coverage = args.use_logarithmic_coverage_model if version not in constraint_res.VERSIONS: raise ValueError("The requested version of resource Tables is not available.") - # If "global" is the only population specified for v4, use the pared-down - # downsampling list. - downsamplings = ( - DOWNSAMPLINGS["v4"] if ((pops == ["global"]) & (int(version[0]) == 4)) else None - ) - logger.info("The following downsamplings will be used: %s", downsamplings) - - # If pops not specified, set to empty Tuple - if not pops: - pops = () - - # Drop chromosome Y from version v4.0 (can add back in when obtain chrY - # methylation data). - if int(version[0]) >= 4: - # TODO: check why there is no Y-par in the context_ht. - regions.remove("chry_nonpar") - # TODO: Add chromosome X back in after complete evaluation for autosome_par. - regions.remove("chrx_nonpar") - # Define variable indicating whether or not the gnomAD version is greater - # than or equal to v4. - version_4_and_above = True - else: - version_4_and_above = False - - # Generate both "plateau" and "coverage" models unless specified to skip - # the coverage model. - models = ["plateau", "coverage"] if not args.skip_coverage_model else ["plateau"] - - # Check the version if 4.0 or later is using "exomes_AN_percent" as coverage_metric. - if coverage_metric == "exomes_AN_percent" and not version_4_and_above: + if version == "2.1.1": raise ValueError( - "Allele number tables are not available for versions prior to v4.0." + "Version 2.1.1 is no longer supported by this constraint pipeline script." + "Please refer to Commit 39928d1 for the last version of the script that" + "supports v2.1.1." ) - if coverage_model_type == "logarithmic": - log10_coverage = True - elif coverage_model_type == "linear": - log10_coverage = False + # Generate both "plateau" and "coverage" models unless specified to skip the + # coverage model. + models = ["plateau", "coverage"] if not skip_coverage_model else ["plateau"] # Construct resources with paths for intermediate Tables generated in the pipeline. resources = get_constraint_resources( version, - use_v2_release_mutation_ht, - args.use_v2_release_context_ht, custom_vep_annotation, overwrite, test, models, + post_fix, ) try: @@ -353,85 +400,10 @@ def main(args): ) res = resources.prepare_context res.check_resource_existence() - context_ht = res.context_ht.ht() - if test: - context_ht = filter_for_test( - context_ht, "raw context", use_gene_list=test_gene_list - ) - - coverage_hts = { - "exomes": res.exomes_coverage_ht.ht(), - "genomes": res.genomes_coverage_ht.ht(), - } - an_hts = ( - {"exomes": res.exomes_an_ht.ht(), "genomes": res.genomes_an_ht.ht()} - if version_4_and_above - else {} - ) - - annotate_context_ht( - context_ht, - coverage_hts, - an_hts, - res.methylation_ht.ht(), - constraint_res.get_gerp_ht(get_reference_genome(context_ht.locus).name), - ).write(res.annotated_context_ht.path, overwrite) - - if args.preprocess_data: - logger.info( - "Adding VEP context annotations and preparing tables for constraint" - " calculations..." - ) - res = resources.preprocess_data - res.check_resource_existence() - context_ht = res.annotated_context_ht.ht() + ht = run_prepare_context(res, test=test, test_gene_list=test_gene_list) + ht.write(res.annotated_context_ht.path, overwrite) - # Add annotations used in constraint calculations. - for data_type in constraint_res.DATA_TYPES: - if data_type != "context": - ht = getattr(res, f"{data_type}_sites_ht").ht() - else: - ht = context_ht - - if test: - ht = filter_for_test(ht, data_type, use_gene_list=test_gene_list) - - # Add annotations from VEP context Table to genome and exome Tables. - if data_type != "context": - ht = add_vep_context_annotations(ht, context_ht) - - # Filter input Table and add annotations used in constraint - # calculations. - ht = prepare_ht_for_constraint_calculations( - ht, - require_exome_coverage=(data_type == "exomes"), - coverage_metric=coverage_metric, - ) - # Filter to locus that is on an autosome. - # TODO: Add back in pseudoautosomal regions once have X/Y methylation - # data. - ht.filter(ht.locus.in_autosome()).write( - getattr(res, f"preprocessed_autosome_par_{data_type}_ht").path, - overwrite=overwrite, - ) - # Sex chromosomes are analyzed separately, since they are biologically - # different from the autosomes. - if data_type != "genomes": - if "chrx_nonpar" in regions: - filter_x_nonpar(ht).write( - getattr( - res, f"preprocessed_chrx_nonpar_{data_type}_ht" - ).path, - overwrite=overwrite, - ) - if "chry_nonpar" in regions: - filter_y_nonpar(ht).write( - getattr( - res, f"preprocessed_chry_nonpar_{data_type}_ht" - ).path, - overwrite=overwrite, - ) - logger.info("Done with preprocessing genome and exome Table.") + logger.info("Done annotating the VEP context Table.") if args.calculate_gerp_cutoffs: logger.warning( @@ -440,157 +412,154 @@ def main(args): ) res = resources.calculate_gerp_cutoffs res.check_resource_existence() + ht = res.annotated_context_ht.ht() gerp_lower_cutoff, gerp_upper_cutoff = calculate_gerp_cutoffs( - res.preprocessed_autosome_par_context_ht.ht() + ht.filter(ht.genomic_region == "autosome_par") ) logger.info( - "Calculated new GERP cutoffs: using a lower GERP cutoff of %f and an" - " upper GERP cutoff of %f.", + "Calculated new GERP cutoffs: using a lower GERP cutoff of %f " + "and an upper GERP cutoff of %f.", gerp_lower_cutoff, gerp_upper_cutoff, ) + if args.preprocess_data: + logger.info( + "Preprocessing the context Table for all downstream constraint steps..." + ) + res = resources.preprocess_data + res.check_resource_existence() + ht = res.annotated_context_ht.ht() + ht = prepare_ht_for_constraint_calculations( + (filter_for_test(ht, use_gene_list=test_gene_list) if test else ht), + exome_coverage_metric=args.exome_coverage_metric, + gen_ancs=args.genetic_ancestry_groups, + include_downsamplings=args.include_downsamplings, + calculate_mutation_rate_min_cov=args.calculate_mutation_rate_min_cov, + calculate_mutation_rate_max_cov=args.calculate_mutation_rate_max_cov, + calculate_mutation_rate_gerp_lower_cutoff=args.calculate_mutation_rate_gerp_lower_cutoff, + calculate_mutation_rate_gerp_upper_cutoff=args.calculate_mutation_rate_gerp_upper_cutoff, + max_af=args.max_af, + build_model_low_cov_cutoff=args.pipeline_low_coverage_filter, + build_model_high_cov_cutoff=args.build_model_high_cov_definition, + build_model_upper_cov_cutoff=args.build_model_upper_cov_cutoff, + apply_model_low_cov_cutoff=args.pipeline_low_coverage_filter, + apply_model_high_cov_cutoff=args.apply_model_high_cov_definition, + skip_coverage_model=skip_coverage_model, + ) + ht.write(res.temp_preprocess_data_ht.path, overwrite=overwrite) + + logger.info("Done preprocessing the context Table.") + if args.calculate_mutation_rate: logger.info("Calculating mutation rate...") res = resources.calculate_mutation_rate res.check_resource_existence() - # Calculate mutation rate using the downsampling with size 1000 genomes in - # genome site Table. - calculate_mu_by_downsampling( - res.preprocessed_autosome_par_genomes_ht.ht(), - res.preprocessed_autosome_par_context_ht.ht(), - recalculate_all_possible_summary=True, - pops=pops, - min_cov=args.min_cov, - max_cov=args.max_cov, - gerp_lower_cutoff=gerp_lower_cutoff, - gerp_upper_cutoff=gerp_upper_cutoff, - ).repartition(args.mutation_rate_partitions).write( - res.mutation_ht.path, overwrite=overwrite - ) + # Use new shuffle method to prevent shuffle errors. + hl._set_flags(use_new_shuffle="1") + + ht = calculate_mu_by_downsampling(res.temp_preprocess_data_ht.ht()) + ht = ht.repartition(args.mutation_rate_partitions) + ht.write(res.mutation_ht.path, overwrite=overwrite) + hl._set_flags(use_new_shuffle=None) + + logger.info("Done calculating mutation rate.") - # Create training datasets that include possible and observed variant counts - # for building models. if args.create_training_set: - logger.info("Counting possible and observed variant counts...") + logger.info( + "Computing the observed and possible counts of synonymous variants to" + "use as a training set for the plateau and coverage models..." + ) res = resources.create_training_set res.check_resource_existence() - # Create training datasets for sites on autosomes/pseudoautosomal regions, - # chromosome X, and chromosome Y. - for r in regions: - op_ht = create_observed_and_possible_ht( - getattr(res, f"preprocessed_{r}_exomes_ht").ht(), - getattr(res, f"preprocessed_{r}_context_ht").ht(), - res.mutation_ht.ht().select("mu_snp"), - max_af=max_af, - pops=pops, - grouping=(coverage_metric,), - coverage_metric=coverage_metric, - partition_hint=args.training_set_partition_hint, - low_coverage_filter=args.pipeline_low_coverage_filter, - transcript_for_synonymous_filter=( - "mane_select" if version_4_and_above else "canonical" - ), # Switch to using MANE Select transcripts rather than canonical for gnomAD v4 and later versions. - global_annotation="training_dataset_params", - ) - if use_v2_release_mutation_ht: - op_ht = op_ht.annotate_globals(use_v2_release_mutation_ht=True) - op_ht.write(getattr(res, f"train_{r}_ht").path, overwrite=overwrite) - op_ht.export(getattr(res, f"train_{r}_tsv")) + ht = create_training_set( + res.temp_preprocess_data_ht.ht(), + res.mutation_ht.ht(), + partition_hint=args.training_set_partition_hint, + ) + + # TODO: Remove repartition once partition_hint bugs are resolved. + ht = ht.repartition(args.training_set_partition_hint) + ht = ht.checkpoint(res.train_ht.path, overwrite=overwrite) + ht.export(res.train_tsv) logger.info("Done with creating training dataset.") if args.build_models: + logger.info("Building plateau and coverage models...") res = resources.build_models res.check_resource_existence() - - # Build plateau and coverage models for autosomes/pseudoautosomal regions, - # chromosome X, and chromosome Y. - for r in regions: - # TODO: Remove repartition once partition_hint bugs are resolved. - training_ht = getattr(res, f"train_{r}_ht").ht() - training_ht = training_ht.repartition(args.training_set_partition_hint) - - logger.info("Building %s plateau and coverage models...", r) - coverage_model, plateau_models = build_models( - coverage_ht=training_ht, - coverage_expr=training_ht[coverage_metric], - weighted=args.use_weights, - pops=pops, - high_cov_definition=args.high_cov_definition, - upper_cov_cutoff=args.upper_cov_cutoff, - skip_coverage_model=True if args.skip_coverage_model else False, - log10_coverage=log10_coverage, - ) + ht = res.train_ht.ht() + print_global_struct(ht.build_models_globals) + coverage_model, plateau_models = build_models( + ht, + ht.exomes_coverage, + model_group_expr=ht.build_model, + skip_coverage_model=skip_coverage_model, + log10_coverage=log10_coverage, + ) + hl.experimental.write_expression( + plateau_models, res.model_plateau.path, overwrite=overwrite + ) + if not args.skip_coverage_model: hl.experimental.write_expression( - plateau_models, - getattr(res, f"model_{r}_plateau").path, - overwrite=overwrite, + coverage_model, res.model_coverage.path, overwrite=overwrite ) - if not args.skip_coverage_model: - hl.experimental.write_expression( - coverage_model, - getattr(res, f"model_{r}_coverage").path, - overwrite=overwrite, - ) - logger.info("Done building %s models.", r) - - if args.apply_models: - res = resources.apply_models + + logger.info("Done building models.") + + if args.apply_models_per_variant: + logger.info( + "Applying plateau and coverage models (if specified) per variant to " + "compute the per-variant expected variant count..." + ) + res = resources.apply_models_per_variant res.check_resource_existence() - # TODO: Remove repartition once partition write bugs are resolved. - mutation_ht = res.mutation_ht.ht().select("mu_snp") - mutation_ht = mutation_ht = mutation_ht.repartition( - args.mutation_rate_partitions + # Use new shuffle method to prevent shuffle errors. + hl._set_flags(use_new_shuffle="1") + + ht = res.temp_preprocess_data_ht.ht() + print_global_struct(ht.apply_models_globals) + ht = create_per_variant_expected_ht( + annotate_with_mu(ht, res.mutation_ht.ht().select("mu_snp")), + res.model_plateau.he(), + coverage_model=None if skip_coverage_model else res.model_coverage.he(), + log10_coverage=log10_coverage, ) + ht.write(res.per_variant_apply_ht.path, overwrite=overwrite) + hl._set_flags(use_new_shuffle=None) - # Apply separate plateau models for sites on autosomes/pseudoautosomal - # regions, chromosome X, and chromosome Y. Use autosomes/pseudoautosomal - # coverage models for all contigs (Note: should test separate coverage models - # for XX/XY in the future). - for r in regions: - logger.info( - "Applying %s plateau and autosome coverage models (if specified)" - " and computing expected variant count and observed:expected" - " ratio...", - r, - ) - oe_ht = apply_models( - exome_ht=getattr(res, f"preprocessed_{r}_exomes_ht").ht(), - context_ht=getattr(res, f"preprocessed_{r}_context_ht").ht(), - mutation_ht=mutation_ht, - plateau_models=getattr(res, f"model_{r}_plateau").he(), - coverage_model=( - getattr(res, "model_autosome_par_coverage").he() - if not args.skip_coverage_model - else None - ), - log10_coverage=log10_coverage, - max_af=max_af, - pops=pops, - downsamplings=downsamplings, - obs_pos_count_partition_hint=args.apply_obs_pos_count_partition_hint, - expected_variant_partition_hint=args.apply_expected_variant_partition_hint, - custom_vep_annotation=custom_vep_annotation, - coverage_metric=coverage_metric, - high_cov_definition=args.high_cov_definition, - low_coverage_filter=args.pipeline_low_coverage_filter, - use_mane_select=( - True - if version_4_and_above - and custom_vep_annotation != "worst_csq_by_gene" - else False - ), # Group by MANE Select transcripts in addition canonical for gnomAD v4 and later versions. - ) - if use_v2_release_mutation_ht: - oe_ht = oe_ht.annotate_globals(use_v2_release_mutation_ht=True) - oe_ht.write(getattr(res, f"apply_{r}_ht").path, overwrite=overwrite) + logger.info("Done computing per-variant expected variant count.") + if args.aggregate_per_variant_expected: logger.info( - "Done computing expected variant count and observed:expected ratio." + "Aggregating per-variant expected variant count by transcript, " + "consequence annotations, and consequence modifier annotations..." + ) + res = resources.aggregate_per_variant_expected + res.check_resource_existence() + + # Use new shuffle method to prevent shuffle errors. + hl._set_flags(use_new_shuffle="1") + + ht = res.per_variant_apply_ht.ht() + ht = ht.annotate(**{f"am_{k}": v for k, v in ht.alpha_missense.items()}) + ht = aggregate_per_variant_expected_ht( + ht, + res.mutation_ht.ht().select("mu_snp"), + custom_vep_annotation=custom_vep_annotation, + additional_grouping=("am_per_98", "am_over_0_999", "am_per_99"), + use_mane_select=True, + ) + ht.write(res.apply_ht.path, overwrite=overwrite) + hl._set_flags(use_new_shuffle=None) + + logger.info( + "Done aggregating per-variant expected variant count by transcript, " + "consequence annotations, and consequence modifier annotations." ) if args.compute_constraint_metrics: @@ -601,24 +570,18 @@ def main(args): res = resources.compute_constraint_metrics res.check_resource_existence() - # Combine Tables of expected variant counts at autosomes/pseudoautosomal - # regions, chromosome X, and chromosome Y sites. - hts = [getattr(res, f"apply_{r}_ht").ht() for r in regions] - union_ht = hts[0].union(*hts[1:]) - union_ht = union_ht.repartition(args.compute_constraint_metrics_partitions) - union_ht = union_ht.checkpoint( - new_temp_file(prefix="constraint_apply_union", extension="ht") - ) + # Use new shuffle method to prevent shuffle errors. + hl._set_flags(use_new_shuffle="1") # Compute constraint metrics. + ht = res.apply_ht.ht() compute_constraint_metrics( - ht=union_ht, + ht=res.apply_ht.ht(), gencode_ht=constraint_res.get_gencode_ht(version), - pops=pops, keys=tuple( [ i - for i in list(union_ht.key) + for i in list(ht.key) if i in ["gene", "transcript", "canonical", "mane_select", "gene_id"] ] @@ -633,12 +596,10 @@ def main(args): raw_z_outlier_threshold_lower_missense=args.raw_z_outlier_threshold_lower_missense, raw_z_outlier_threshold_lower_syn=args.raw_z_outlier_threshold_lower_syn, raw_z_outlier_threshold_upper_syn=args.raw_z_outlier_threshold_upper_syn, - # OS (other splice) is not implemented for build 38. - include_os=not version_4_and_above, - use_mane_select_over_canonical=version_4_and_above, - ).select_globals("version", "apply_model_params", "sd_raw_z").write( - res.constraint_metrics_ht.path, overwrite=overwrite - ) + # ).select_globals( + # "version", "apply_model_params", "constraint_meta", "sd_raw_z" + ).write(res.constraint_metrics_ht.path, overwrite=overwrite) + hl._set_flags(use_new_shuffle=None) logger.info("Done with computing constraint metrics.") if args.export_tsv: @@ -649,7 +610,7 @@ def main(args): ht = res.constraint_metrics_ht.ht() # If downsamplings per genetic ancestry group are present, export # downsamplings to a separate tsv and drop from the main metrics tsv. - if pops: + if args.genetic_ancestry_groups: downsampling_ht = explode_downsamplings_oe( ht, downsampling_meta=hl.eval(ht.apply_model_params.downsampling_meta), @@ -687,6 +648,12 @@ def main(args): type=str, default=constraint_res.CURRENT_VERSION, ) + parser.add_argument( + "--post-fix", + help="Post-fix to append to the output file names.", + type=str, + default=None, + ) parser.add_argument( "--test", help=( @@ -704,11 +671,7 @@ def main(args): ), action="store_true", ) - - prepare_context_args = parser.add_argument_group( - "Prepare context Table args", "Arguments used for preparing the context Table." - ) - prepare_context_args.add_argument( + parser.add_argument( "--prepare-context-ht", help=( "Prepare the context Table by splitting multiallelic sites and adding " @@ -717,34 +680,6 @@ def main(args): ), action="store_true", ) - - preprocess_data_args = parser.add_argument_group( - "Preprocess data args", "Arguments used for preprocessing the data." - ) - - preprocess_data_args.add_argument( - "--preprocess-data", - help=( - "Whether to prepare the exome, genome, and context Table for constraint" - " calculations by adding necessary coverage, methylation level, and VEP" - " annotations." - ), - action="store_true", - ) - - preprocess_data_args.add_argument( - "--use-v2-release-context-ht", - help="Whether to use the annotated context Table for the v2 release.", - action="store_true", - ) - - preprocess_data_args.add_argument( - "--coverage-metric", - help="Name of metric to use to assess coverage, such as 'exome_coverage' or 'exomes_AN_percent'. Default is 'exome_coverage'.", - type=str, - default="exome_coverage", - ) - parser.add_argument( "--calculate-gerp-cutoffs", help=( @@ -757,34 +692,17 @@ def main(args): action="store_true", ) - parser.add_argument( - "--pipeline-low-coverage-filter", - help=( - "Lower median coverage cutoff to use throughout the pipeline. Sites with" - " coverage below this cutoff will be excluded when creating the training" - " set, building and applying models, and computing constraint metrics." - " Default is 30." - ), - type=int, - default=30, - ) - - mutation_rate_args = parser.add_argument_group( - "Calculate mutation rate args", - "Arguments used for calculating the mutation rate.", + preprocess_args = parser.add_argument_group( + "Preprocess data args", + "All arguments used for preprocessing data for downstream steps.", ) - - recalculate_mutation_rate = mutation_rate_args.add_argument( - "--calculate-mutation-rate", - help=( - "Calculate baseline mutation rate for each substitution and context using" - " downsampling data." - ), + preprocess_args.add_argument( + "--preprocess-data", + help="Preprocess the context Table for downstream constraint steps.", action="store_true", ) - - mutation_rate_args.add_argument( - "--min-cov", + preprocess_args.add_argument( + "--calculate-mutation-rate-min-cov", help=( "Minimum coverage required to keep a site when calculating the mutation" " rate. Default is 15." @@ -792,8 +710,8 @@ def main(args): type=int, default=15, ) - mutation_rate_args.add_argument( - "--max-cov", + preprocess_args.add_argument( + "--calculate-mutation-rate-max-cov", help=( "Maximum coverage required to keep a site when calculating the mutation" " rate. Default is 60." @@ -801,8 +719,8 @@ def main(args): type=int, default=60, ) - mutation_rate_args.add_argument( - "--gerp-lower-cutoff", + preprocess_args.add_argument( + "--calculate-mutation-rate-gerp-lower-cutoff", help=( "Minimum GERP score for variant to be included when calculating the" " mutation rate. Default is -3.9885 (precalculated on the GRCh37 context" @@ -811,8 +729,8 @@ def main(args): type=float, default=-3.9885, ) - mutation_rate_args.add_argument( - "--gerp-upper-cutoff", + preprocess_args.add_argument( + "--calculate-mutation-rate-gerp-upper-cutoff", help=( "Maximum GERP score for variant to be included when calculating the" " mutation rate. Default is 2.6607 (precalculated on the GRCh37 context" @@ -821,11 +739,103 @@ def main(args): type=float, default=2.6607, ) + preprocess_args.add_argument( + "--exome-coverage-metric", + help=( + "Name of metric to use to assess exome coverage, such as 'median', 'AN', or" + "'AN_percent'. Default is 'AN_percent'." + ), + type=str, + default="AN_percent", + ) + preprocess_args.add_argument( + "--pipeline-low-coverage-filter", + help=( + "Lower exome coverage cutoff to use throughout the pipeline. Sites with" + " coverage below this cutoff will be excluded when creating the training" + " set, building and applying models, and computing constraint metrics." + " Default is 0." + ), + type=int, + default=0, + ) + preprocess_args.add_argument( + "--max-af", + help=( + "Maximum variant allele frequency to use when filtering variants for " + "training and applying models." + ), + type=float, + default=0.001, + ) + preprocess_args.add_argument( + "--genetic-ancestry-groups", + nargs="+", + help=( + "Populations on which to build models, apply models, and or compute metrics " + "on. Default is None." + ), + choices=["afr", "amr", "eas", "nfe", "sas"], + default=None, + ) + preprocess_args.add_argument( + "--include-downsamplings", + help="Include downsamplings in the constraint pipeline.", + action="store_true", + ) + preprocess_args.add_argument( + "--skip-coverage-model", + help="Omit computing and applying the coverage model.", + action="store_true", + ) + preprocess_args.add_argument( + "--build-model-upper-cov-cutoff", + help=( + "Upper exome coverage cutoff. Sites with coverage above this cutoff are" + " excluded from the high coverage Table when building the models. Default" + " is None." + ), + type=int, + default=None, + ) + preprocess_args.add_argument( + "--build-model-high-cov-definition", + help=( + "Lower exome coverage cutoff to use to define high coverage sites when " + "building models. Sites with coverage below this cutoff are excluded from " + "the high coverage Table when building models. Default is 90." + ), + type=int, + default=90, + ) + preprocess_args.add_argument( + "--apply-model-high-cov-definition", + help=( + "Lower exome coverage cutoff to use to define high coverage sites when " + "applying models. Sites with coverage below this cutoff are excluded from " + "the high coverage Table when applying models. Default is 90." + ), + type=int, + default=90, + ) + + mutation_rate_args = parser.add_argument_group( + "Calculate mutation rate args", + "Arguments used for calculating the mutation rate.", + ) + mutation_rate_args.add_argument( + "--calculate-mutation-rate", + help=( + "Calculate baseline mutation rate for each substitution and context using" + " downsampling data." + ), + action="store_true", + ) mutation_rate_args.add_argument( "--mutation-rate-partitions", help=( - "Number of partitions to which the mutation rate Table should be" - " repartitioned." + "Number of partitions to which the mutation rate Table should be " + "repartitioned." ), type=int, default=1, @@ -834,67 +844,32 @@ def main(args): training_set_args = parser.add_argument_group( "Training set args", "Arguments used for creating the training set." ) - training_set_args.add_argument( "--create-training-set", help=( - "Count the observed variants and possible variants by exome coverage at" - " synonymous sites." + "Count the observed variants and possible variants by exome coverage at " + "synonymous sites." ), action="store_true", ) - training_set_args.add_argument( "--training-set-partition-hint", help=( - "Target number of partitions for aggregation when counting variants for" - " training datasets." + "Target number of partitions for aggregation when counting variants for " + "training datasets." ), type=int, default=100, ) - # `max-af` is an arg for both `--create-training-set` and `--apply-models` - maximum_af = training_set_args.add_argument( - "--max-af", - help="Maximum variant allele frequency to keep.", - type=float, - default=0.001, - ) - - # `populations` is an arg for `--create-training-set`, `--apply-models`, `--build-models`, and `compute_constraint_args` - populations = training_set_args.add_argument( - "--pops", - nargs="+", - help=( - "Populations on which to train models, build models, apply models, and or" - " compute metrics on. Downsamplings for the specified population will be" - " included." - ), - choices=["global", "afr", "amr", "eas", "nfe", "sas"], - default=None, - ) - - use_v2_release_mutation_rate = training_set_args.add_argument( - "--use-v2-release-mutation-ht", - help="Whether to use the mutatation rate computed for the v2 release.", - action="store_true", - ) - - mutation_rate_parser = parser.add_mutually_exclusive_group(required=False) - mutation_rate_parser._group_actions.append(use_v2_release_mutation_rate) - mutation_rate_parser._group_actions.append(recalculate_mutation_rate) - build_models_args = parser.add_argument_group( "Build models args", "Arguments used for building models." ) - build_models_args.add_argument( "--build-models", help="Build plateau and coverage models.", action="store_true", ) - build_models_args.add_argument( "--use-weights", help=( @@ -903,61 +878,39 @@ def main(args): ), action="store_true", ) - build_models_args.add_argument( - "--upper-cov-cutoff", - help=( - "Upper median coverage cutoff. Sites with coverage above this cutoff are" - " excluded from the high coverage Table when building the models. Default" - " is 100." - ), - type=int, - default=100, - ) - - build_models_args.add_argument( - "--high-cov-definition", + cov_model_type = build_models_args.add_argument( + "--use-logarithmic-coverage-model", help=( - "Lower median coverage cutoff to use to define high coverage sites. Sites" - " with coverage below this cutoff are excluded from the high coverage Table" - " when building and applying the models. Default is 30." + "Use a logarithmic model for low coverage sites when building and applying " + "the coverage model." ), - type=int, - default=30, - ) - - build_models_args.add_argument( - "--skip-coverage-model", - help="Omit computing and applying the coverage model.", action="store_true", ) - cov_model_type = build_models_args.add_argument( - "--coverage-model-type", + parser.add_argument( + "--apply-models-per-variant", help=( - "Type of model to use for low coverage sites when building and applying the coverage model, either 'linear' or 'logarithmic'. Default is 'logarithmic'." + "Apply plateau and coverage models to variants in exome sites Table and" + " context Table to compute expected variant counts per variant." ), - type=str, - choices=["linear", "logarithmic"], - default="logarithmic", + action="store_true", ) - build_models_args._group_actions.append(populations) - - apply_models_args = parser.add_argument_group( - "Apply models args", - "Arguments used for applying the plateau and coverage models.", + aggregate_per_variant_expected_args = parser.add_argument_group( + "Aggregate per variant expected args", + "Arguments used for applying aggregating the per variant expected values.", ) - - apply_models_args.add_argument( - "--apply-models", + aggregate_per_variant_expected_args.add_argument( + "--aggregate-per-variant-expected", help=( - "Apply plateau and coverage models to variants in exome sites Table and" - " context Table to compute expected variant counts." + "Aggregate the per-variant expected variant counts to get the expected " + "variant counts for each transcript by consequence annotation and " + "modifier." ), action="store_true", ) - apply_models_args.add_argument( + aggregate_per_variant_expected_args.add_argument( "--apply-obs-pos-count-partition-hint", help=( "Target number of partitions for aggregation when counting observed and" @@ -966,7 +919,7 @@ def main(args): type=int, default=2000, ) - apply_models_args.add_argument( + aggregate_per_variant_expected_args.add_argument( "--apply-expected-variant-partition-hint", help=( "Target number of partitions for sum aggregators after applying models to" @@ -975,7 +928,7 @@ def main(args): type=int, default=1000, ) - apply_models_args.add_argument( + aggregate_per_variant_expected_args.add_argument( "--custom-vep-annotation", help=( "Custom VEP annotation to be used to annotate transcript when" @@ -985,10 +938,7 @@ def main(args): default="transcript_consequences", choices=constraint_res.CUSTOM_VEP_ANNOTATIONS, ) - apply_models_args._group_actions.append(maximum_af) - apply_models_args._group_actions.append(populations) - apply_models_args._group_actions.append(use_v2_release_mutation_rate) - apply_models_args._group_actions.append(cov_model_type) + aggregate_per_variant_expected_args._group_actions.append(cov_model_type) compute_constraint_args = parser.add_argument_group( "Computate constraint metrics args", @@ -1005,14 +955,12 @@ def main(args): compute_constraint_args.add_argument( "--compute-constraint-metrics-partitions", help=( - "Number of partitions to which the unioned Table of expected variant counts" - " for autosomes/pseudoautosomal regions, chromosome X, and chromosome Y " - " should be reaprtitioned." + "Number of partitions to which the Table of expected variant counts should " + "be reaprtitioned." ), type=int, default=1000, ) - compute_constraint_args.add_argument( "--min-diff-convergence", help=( @@ -1024,7 +972,6 @@ def main(args): type=float, default=0.001, ) - compute_constraint_args.add_argument( "--expectation-null", help=( @@ -1063,7 +1010,6 @@ def main(args): type=float, default=-8.0, ) - # NOTE: gnomAD v2 used raw z thresholds of +/- 5. compute_constraint_args.add_argument( "--raw-z-outlier-threshold-lower-missense", help=( @@ -1094,13 +1040,11 @@ def main(args): type=float, default=8.0, ) - parser.add_argument( + compute_constraint_args.add_argument( "--export-tsv", help="Export constraint metrics to tsv file.", action="store_true", ) - compute_constraint_args._group_actions.append(populations) - args = parser.parse_args() main(args) diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 3f497786..94cc51e3 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -55,8 +55,29 @@ Low coverage sites require an extra calibration when computing the proportion of expected variation. """ +MU_GROUPING = ("context", "ref", "alt", "methylation_level") +""" +Annotations used to group variants for the mutation rate calculation. +""" + +CALIBRATION_GROUPING = ("genomic_region", "build_model", "cpg", "exomes_coverage") +""" +Annotations used to group variants for the mutation rate calibration. +""" + +AGGREGATE_SUM_FIELDS = ( + "mu_snp", + "observed_variants", + "possible_variants", + "predicted_proportion_observed", + "coverage_correction", + "expected_variants", +) +""" +Fields to sum (or array sum) when aggregating the expected counts Table. +""" + -# VEP context Table. def get_vep_context_ht(version: str) -> TableResource: """ Return VEP context Table corresponding to specified gnomAD version. @@ -74,7 +95,12 @@ def get_vep_context_ht(version: str) -> TableResource: raise ValueError("Not a valid gnomAD version -- must be either 2.1.1 or 4.x!") -def get_constraint_root(version: str = CURRENT_VERSION, test: bool = False) -> str: +def get_constraint_root( + version: str = CURRENT_VERSION, + test: bool = False, + post_fix=None, + temp: bool = False, +) -> str: """ Return path to constraint root folder. @@ -82,11 +108,16 @@ def get_constraint_root(version: str = CURRENT_VERSION, test: bool = False) -> s :param test: Whether to use a tmp path. :return: Root path to constraint resources. """ - return ( - f"gs://gnomad-tmp/gnomad_v{version}_testing/constraint" - if test - else f"gs://gnomad/v{version}/constraint" - ) + post_fix = post_fix or "" + if post_fix: + post_fix = f"_{post_fix}" + + if test: + return f"gs://gnomad-tmp/gnomad_v{version}_testing/constraint{post_fix}" + if temp: + return f"gs://gnomad-tmp/gnomad_v{version}/constraint{post_fix}" + + return f"gs://gnomad/v{version}/constraint{post_fix}" def get_sites_resource(data_type: str, version: str = CURRENT_VERSION) -> BaseResource: @@ -138,8 +169,8 @@ def get_methylation_ht(build: str) -> TableResource: methylation_chrx = ref_grch38.methylation_sites_chrx.ht() methylation_autosomes = ref_grch38.methylation_sites.ht() methylation_ht = methylation_autosomes.union(methylation_chrx) - tmp_path = get_constraint_root(version=build, test=True) - methylation_ht = methylation_ht.checkpoint(tmp_path, overwrite=True) + tmp_path = get_checkpoint_path(f"methylation_{build}") + methylation_ht.checkpoint(tmp_path, _read_if_exists=True) return TableResource(path=tmp_path) else: raise ValueError("Build must be one of 'GRCh37' or 'GRCh38'.") @@ -167,7 +198,7 @@ def get_coverage_ht( def get_mutation_ht( version: str = CURRENT_VERSION, test: bool = False, - use_v2_release_mutation_ht: bool = False, + post_fix: Optional[str] = None, ) -> TableResource: """ Return mutation Table that includes the baseline mutation rate for each substitution and context. @@ -175,53 +206,39 @@ def get_mutation_ht( :param version: The version of the Table. Default is CURRENT_VERSION. :param test: Whether the Table is for testing purposes and only contains sites in chr20, chrX, and chrY. Default is False. - :param use_v2_release_mutation_ht: Whether to use the precomputed gnomAD v2.1.1 released mutation rate table. :return: Mutation rate Table. """ - if use_v2_release_mutation_ht: - return TableResource( - path="gs://gcp-public-data--gnomad/papers/2019-flagship-lof/v1.0/model/mutation_rate_methylation_bins.ht", - ) - else: - check_param_scope(version) - return TableResource( - f"{get_constraint_root(version, test)}/mutation_rate/gnomad.v{version}.mutation_rate.ht" - ) + check_param_scope(version) + return TableResource( + f"{get_constraint_root(version, test, post_fix)}/mutation_rate/gnomad.v{version}.mutation_rate.ht" + ) def get_annotated_context_ht( version: str = CURRENT_VERSION, - use_v2_context_ht: bool = False, test: bool = False, + post_fix: Optional[str] = None, ) -> TableResource: """ Return TableResource of annotated context Table. :param version: One of the release versions (`VERSIONS`). Default is `CURRENT_VERSION`. - :param use_v2_context_ht: Whether to use annotated context Table that was produced - for gnomAD v2. Default is False. :param test: Whether the Table is for testing purposes and only contains sites in chr20, chrX, and chrY. Default is False. :return: TableResource of annotated context Table. """ - if use_v2_context_ht: - return TableResource( - "gs://gcp-public-data--gnomad/papers/2019-flagship-lof/v1.0/context/Homo_sapiens_assembly19.fasta.snps_only.vep_20181129.ht" - ) - check_param_scope(version) return TableResource( - f"{get_constraint_root(version, test)}/preprocessed_data/annotated_context.ht" + f"{get_constraint_root(version, test, post_fix)}/preprocessed_data/annotated_context.ht" ) def get_preprocessed_ht( - data_type: str, version: str = CURRENT_VERSION, - genomic_region: str = "autosome_par", test: bool = False, + post_fix: Optional[str] = None, ) -> TableResource: """ Return TableResource of preprocessed genome, exomes, and context Table. @@ -232,69 +249,62 @@ def get_preprocessed_ht( The context Table will have annotations added by `prepare_ht_for_constraint_calculations()`. - :param data_type: One of "exomes", "genomes" or "context. :param version: One of the release versions (`VERSIONS`). Default is `CURRENT_VERSION`. - :param genomic_region: The genomic region of the resource. One of "autosome_par", - "chrx_nonpar", "chry_nonpar". Default is "autosome_par". :param test: Whether the Table is for testing purposes and only contains sites in chr20, chrX, and chrY. Default is False. - :return: TableResource of processed genomes, exomes, or context Table. + :return: TableResource of processed context Table. """ - check_param_scope(version, genomic_region, data_type) + check_param_scope(version) return TableResource( - f"{get_constraint_root(version, test)}/preprocessed_data/gnomad.v{version}.{data_type}.preprocessed.{genomic_region}.ht" + f"{get_constraint_root(version, test, post_fix, temp=True)}/preprocessed_data/gnomad.v{version}.context.preprocessed.ht" ) def get_training_dataset( version: str = CURRENT_VERSION, - genomic_region: str = "autosome_par", test: bool = False, + post_fix: Optional[str] = None, ) -> TableResource: """ Return TableResource of training dataset with observed and possible variant count. :param version: One of the release versions (`VERSIONS`). Default is `CURRENT_VERSION`. - :param genomic_region: The genomic region of the resource. One of "autosome_par", - "chrx_nonpar", or "chry_nonpar". Default is "autosome_par". :param test: Whether the Table is for testing purpose and only contains sites in chr20, chrX, and chrY. Default is False. :return: TableResource of training dataset. """ - check_param_scope(version, genomic_region) + check_param_scope(version) return TableResource( - f"{get_constraint_root(version, test)}/training_data/gnomad.v{version}.constraint_training.{genomic_region}.ht" + f"{get_constraint_root(version, test, post_fix)}/training_data/gnomad.v{version}.constraint_training.ht" ) def get_training_tsv_path( version: str = CURRENT_VERSION, - genomic_region: str = "autosome_par", test: bool = False, + post_fix: Optional[str] = None, ) -> str: """ Return tsv of training dataset with observed and possible variant count. :param version: One of the release versions (`VERSIONS`). Default is `CURRENT_VERSION`. - :param genomic_region: The genomic region of the resource. One of "autosome_par", - "chrx_nonpar", or "chry_nonpar". Default is "autosome_par". :param test: Whether the Table is for testing purpose and only contains sites in chr20, chrX, and chrY. Default is False. :return: TSV path of training dataset. """ - check_param_scope(version, genomic_region) + check_param_scope(version) - return f"{get_constraint_root(version, test)}/training_data/gnomad.v{version}.constraint_training.{genomic_region}.tsv.bgz" + return f"{get_constraint_root(version, test, post_fix)}/training_data/gnomad.v{version}.constraint_training.tsv.bgz" def get_models( model_type: str, version: str = CURRENT_VERSION, - genomic_region: str = "autosome_par", test: bool = False, + post_fix: Optional[str] = None, ) -> ExpressionResource: """ Return path to a HailExpression that contains desired model type. @@ -302,25 +312,47 @@ def get_models( :param model_type: The type of model. One of "plateau", "coverage". Default is None. :param version: One of the release versions (`VERSIONS`). Default is `CURRENT_VERSION`. - :param genomic_region: The genomic region of the resource. One of "autosome_par", - "chrx_non_par", or "chry_non_par". Default is "autosome_par". :param test: Whether the Table is for testing purpose and only contains sites in chr20, chrX, and chrY. Default is False. :return: Path to the specified model. """ + check_param_scope(version=version, model_type=model_type) + return ExpressionResource( + f"{get_constraint_root(version, test, post_fix)}/models/gnomad.v{version}.{model_type}.he" + ) + + +def get_per_variant_expected_dataset( + custom_vep_annotation: str = "transcript_consequences", + version: str = CURRENT_VERSION, + test: bool = False, + post_fix: Optional[str] = None, +) -> TableResource: + """ + Return TableResource containing the expected variant counts and observed:expected ratio. + + :param custom_vep_annotation: The VEP annotation used to customize the constraint + model (one of "transcript_consequences" or "worst_csq_by_gene"). + :param version: One of the release versions (`VERSIONS`). Default is + `CURRENT_VERSION`. + :param test: Whether the Table is for testing purpose and only contains sites in + chr20, chrX, and chrY. Default is False. + :return: Path of the model. + """ check_param_scope( - version=version, genomic_region=genomic_region, model_type=model_type + version=version, + custom_vep_annotation=custom_vep_annotation, ) - return ExpressionResource( - f"{get_constraint_root(version, test)}/models/gnomad.v{version}.{model_type}.{genomic_region}.he" + return TableResource( + f"{get_constraint_root(version, test, post_fix)}/apply_models/{custom_vep_annotation}/gnomad.v{version}.per_variant_expected.ht" ) -def get_predicted_proportion_observed_dataset( +def get_apply_models( custom_vep_annotation: str = "transcript_consequences", version: str = CURRENT_VERSION, - genomic_region: str = "autosome_par", test: bool = False, + post_fix: Optional[str] = None, ) -> TableResource: """ Return TableResource containing the expected variant counts and observed:expected ratio. @@ -329,25 +361,24 @@ def get_predicted_proportion_observed_dataset( model (one of "transcript_consequences" or "worst_csq_by_gene"). :param version: One of the release versions (`VERSIONS`). Default is `CURRENT_VERSION`. - :param genomic_region: The genomic region of the resource. One of "autosome_par", - "chrx_non_par", or "chry_non_par". Default is "autosome_par". :param test: Whether the Table is for testing purpose and only contains sites in chr20, chrX, and chrY. Default is False. :return: Path of the model. """ check_param_scope( version=version, - genomic_region=genomic_region, custom_vep_annotation=custom_vep_annotation, ) return TableResource( - f"{get_constraint_root(version, test)}/predicted_proportion_observed/{custom_vep_annotation}/gnomad.v{version}.predicted_proportion_observed.{genomic_region}.ht" + f"{get_constraint_root(version, test, post_fix)}/apply_models/{custom_vep_annotation}/gnomad.v{version}.apply.per_variant_expected.ht" ) def get_constraint_metrics_dataset( + custom_vep_annotation: str = "transcript_consequences", version: str = CURRENT_VERSION, test: bool = False, + post_fix: Optional[str] = None, ) -> TableResource: """ Return TableResource of pLI scores, observed:expected ratio, 90% confidence interval around the observed:expected ratio, and z scores. @@ -361,13 +392,14 @@ def get_constraint_metrics_dataset( check_param_scope(version=version) return TableResource( - f"{get_constraint_root(version, test)}/metrics/gnomad.v{version}.constraint_metrics.ht" + f"{get_constraint_root(version, test, post_fix)}/metrics/{custom_vep_annotation}/gnomad.v{version}.constraint_metrics.ht" ) def get_constraint_tsv_path( version: str = CURRENT_VERSION, test: bool = False, + post_fix: Optional[str] = None, ) -> str: """ Return tsv path of pLI scores, observed:expected ratio, 90% confidence interval around the observed:expected ratio, and z scores. @@ -379,12 +411,13 @@ def get_constraint_tsv_path( """ check_param_scope(version=version) - return f"{get_constraint_root(version, test)}/metrics/tsv/gnomad.v{version}.constraint_metrics.tsv" + return f"{get_constraint_root(version, test, post_fix)}/metrics/tsv/gnomad.v{version}.constraint_metrics.tsv" def get_downsampling_constraint_tsv_path( version: str = CURRENT_VERSION, test: bool = False, + post_fix: Optional[str] = None, ) -> str: """ Return tsv path of downsampling observed and expected counts. @@ -396,12 +429,11 @@ def get_downsampling_constraint_tsv_path( """ check_param_scope(version=version) - return f"{get_constraint_root(version, test)}/metrics/tsv/gnomad.v{version}.downsampling_constraint_metrics.tsv.bgz" + return f"{get_constraint_root(version, test, post_fix)}/metrics/tsv/gnomad.v{version}.downsampling_constraint_metrics.tsv.bgz" def check_param_scope( version: Optional[str] = None, - genomic_region: Optional[str] = None, data_type: Optional[str] = None, model_type: Optional[str] = None, custom_vep_annotation: Optional[str] = None, @@ -412,16 +444,12 @@ def check_param_scope( If version is specified, return the genome build of the version as a string. :param version: One of the release versions (`VERSIONS`). Default is None. - :param genomic_region: The genomic region of the resource. One of "autosome_par", - "chrx_non_par", or "chry_non_par". Default is None. :param data_type: One of "exomes", "genomes" or "context". Default is None. :param model_type: One of "plateau", "coverage". Default is None. :param custom_vep_annotation: The VEP annotation used to customize the constraint model (one of "transcript_consequences" or "worst_csq_by_gene"). Default is None. :return: Genome build of version as a string or None. """ - if genomic_region and genomic_region not in GENOMIC_REGIONS: - raise ValueError(f"genomic_region must be one of: {GENOMIC_REGIONS}!") if data_type and data_type not in DATA_TYPES: raise ValueError(f"data_type must be one of: {DATA_TYPES}!") if model_type and model_type not in MODEL_TYPES: @@ -439,7 +467,9 @@ def check_param_scope( return "GRCh38" -def get_logging_path(name: str, version: str = CURRENT_VERSION) -> str: +def get_logging_path( + name: str, version: str = CURRENT_VERSION, post_fix: Optional[str] = None +) -> str: """ Create a path for Hail log files. @@ -448,11 +478,14 @@ def get_logging_path(name: str, version: str = CURRENT_VERSION) -> str: `CURRENT_VERSION`. :return: Output log path. """ - return f"{get_constraint_root(version, test=True)}/logging/{name}.log" + return f"{get_constraint_root(version, test=True, post_fix=post_fix)}/logging/{name}.log" def get_checkpoint_path( - name: str, version: str = CURRENT_VERSION, mt: bool = False + name: str, + version: str = CURRENT_VERSION, + mt: bool = False, + post_fix: Optional[str] = None, ) -> str: """ Create a checkpoint path for Table or MatrixTable. @@ -462,7 +495,7 @@ def get_checkpoint_path( :param bool mt: Whether path is for a MatrixTable. Default is False. :return: Output checkpoint path. """ - return f'{get_constraint_root(version, test=True)}/checkpoint_files/{name}.{"mt" if mt else "ht"}' + return f'{get_constraint_root(version, test=True, post_fix=post_fix)}/checkpoint_files/{name}.{"mt" if mt else "ht"}' def get_gencode_ht(version: str) -> hl.Table: @@ -475,6 +508,6 @@ def get_gencode_ht(version: str) -> hl.Table: if int(version[0]) == 2: return ref_grch37.gencode.ht() elif int(version[0]) == 4: - return ref_grch38.gencode.ht() + return ref_grch38.gencode.ht(read_args={"_n_partitions": 500}) else: raise ValueError("Version must be within gnomAD v2 or v4.") diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 07d710b1..469c3fe0 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -1,43 +1,43 @@ """Script containing utility functions used in the constraint pipeline.""" +import functools import logging -from typing import Dict, List, Optional, Tuple +import operator +from typing import Dict, List, Optional, Tuple, Union import hail as hl import numpy as np +from gnomad.assessment.summary_stats import generate_filter_combinations +from gnomad.resources.grch38.gnomad import DOWNSAMPLINGS from gnomad.utils.constraint import ( add_gencode_transcript_annotations, + aggregate_expected_variants_expr, annotate_exploded_vep_for_constraint_groupings, annotate_mutation_type, annotate_with_mu, + apply_models, + apply_plateau_models, calculate_raw_z_score, calculate_raw_z_score_sd, - collapse_strand, - compute_expected_variants, + calibration_model_group_expr, compute_pli, - count_variants_by_group, + count_observed_and_possible_by_group, + coverage_correction_expr, get_constraint_flags, - get_downsampling_freq_indices, - oe_aggregation_expr, oe_confidence_interval, - trimer_from_heptamer, + single_variant_count_expr, + single_variant_observed_and_possible_expr, + weighted_agg_sum_expr, ) -from gnomad.utils.filtering import ( - add_filters_expr, - filter_by_numeric_expr_range, - filter_for_mu, - filter_to_autosomes, -) -from gnomad.utils.reference_genome import get_reference_genome -from gnomad.utils.vep import ( - add_most_severe_csq_to_tc_within_vep_root, - filter_vep_transcript_csqs, -) -from hail.utils.misc import new_temp_file +from gnomad.utils.filtering import add_filters_expr +from gnomad.utils.vep import filter_vep_transcript_csqs_expr +from hail.utils.misc import divide_null, new_temp_file from gnomad_constraint.resources.resource_utils import ( + AGGREGATE_SUM_FIELDS, + CALIBRATION_GROUPING, COVERAGE_CUTOFF, - get_checkpoint_path, + MU_GROUPING, ) logging.basicConfig( @@ -48,441 +48,731 @@ logger.setLevel(logging.INFO) -def add_vep_context_annotations( - ht: hl.Table, annotated_context_ht: hl.Table -) -> hl.Table: +# TODO: For now I am leaving this here instead of moving to gnomad_methods because +# there is another PR in gnomad_methods that might change the way this function is +# implemented. +def filter_freq_for_constraint( + freq_expr: hl.ArrayExpression, + freq_meta_expr: List[Dict[str, str]], + gen_ancs: Optional[List[str]] = None, + downsamplings: Optional[List[int]] = None, + downsampling_gen_ancs: Optional[List[str]] = None, + gen_anc_label: str = "gen_anc", +) -> Tuple[hl.ArrayExpression, List[Dict[str, str]]]: """ - Add annotations from VEP context Table to gnomAD data. - - Function adds the following annotations: - - context - - methylation - - coverage - - gerp + Filter the frequency array for constraint calculations. + + The frequency array is filtered to include only adj frequencies for + the populations in `pops` and the downsamplings in `downsamplings`, for the + populations in `downsampling_pops`. + + No matter the input, the frequency array is always filtered to include the + "adj" frequency for the full dataset. + + If `downsamplings` is None, no downsamplings are included. If `downsamplings` is + provided, and `downsampling_pops` is None, only the "global" downsampling is + included. If `downsampling_pops` is provided, the downsamplings for the populations + in `downsampling_pops` are included as well as the "global" downsampling. + + :param freq_expr: Frequency array. + :param freq_meta_expr: Frequency metadata array. + :param gen_ancs: Optional list of genetic ancestries to include in the frequency + array. Default is None. + :param downsamplings: Optional list of downsamplings to include in the frequency + array. Default is None. + :param downsampling_gen_ancs: Optional list of genetic ancestries to include + downsamplings frequencies for. Default is None. + :param gen_anc_label: Label for the genetic ancestry field in the frequency + metadata. Default is "gen_anc". + :return: Filtered frequency array and metadata. + """ + freq_meta = hl.eval(freq_meta_expr) + meta_keep = [{"group": "adj"}] + + if gen_ancs is not None: + meta_keep += [{"group": "adj", gen_anc_label: pop} for pop in gen_ancs] + + if downsamplings is not None: + downsampling_pops = ["global"] + (downsampling_gen_ancs or []) + meta_keep += [ + {"group": "adj", gen_anc_label: pop, "downsampling": str(ds)} + for pop in downsampling_pops + for ds in downsamplings + ] - Function drops `a_index`, `was_split`, and`colocated_variants` annotations from - gnomAD data. + meta_keep = [m for m in meta_keep if m in freq_meta] + freq_expr = hl.array([freq_expr[freq_meta.index(m)] for m in meta_keep]) - .. note:: - Function expects that multiallelic variants in the VEP context Table have been - split. + return freq_expr, meta_keep - Function also adds 'an_strata_sample_count' to globals if present. - :param ht: gnomAD exomes or genomes public Hail Table. - :param annotated_context_ht: VEP context Table. - :return: Table with annotations. +def get_annotations_for_computing_mu( + locus_expr: hl.expr.LocusExpression, + genomes_filter_expr: hl.expr.StructExpression, + genomes_freq_expr: hl.expr.ArrayExpression, + genomes_freq_meta: List[Dict[str, str]], + genomes_coverage_expr: hl.expr.Int32Expression, + gerp_expr: hl.expr.Float64Expression, + most_severe_consequence_expr: hl.expr.StringExpression, + gen_ancs: Optional[List[str]] = None, + downsampling_level: int = 1000, + min_cov: int = 15, + max_cov: int = 60, + gerp_lower_cutoff: float = -3.9885, + gerp_upper_cutoff: float = 2.6607, + ac_cutoff: int = 5, +) -> Tuple[hl.expr.StructExpression, hl.expr.StructExpression]: """ - context_ht = annotated_context_ht.drop("a_index", "was_split") - context_ht = context_ht.annotate(vep=context_ht.vep.drop("colocated_variants")) - if "an_strata_sample_count" in context_ht.globals: - ht = ht.annotate_globals( - an_strata_sample_count=context_ht.index_globals().an_strata_sample_count - ) - ht = ht.annotate(**context_ht[ht.key]) - return ht + Get the annotations that are needed to compute the mutation rate. + + The function will return the following annotations: + + - genomes_freq: Frequency array for the genomes dataset, filtered to the + requested genetic ancestries and downsampling level. + - observed_variants: This annotation is an array, where each element + corresponds to whether the variant is observed in the genomes dataset for the + frequency group at the corresponding index in the `genomes_freq` array. + Must PASS genome filters, have AC <= 'ac_cutoff' at the specified + `downsampling_level`, and have a genome mean coverage >= `min_cov` + and <= `max_cov`. The boolean value is stored as an integer (0 or 1). + - possible_variants: Whether the variant is considered a possible variant in + the genomes dataset. This includes variants not in the genome dataset (genome + AF undefined), or also considered in the observed variant set. The boolean + value is stored as an integer (0 or 1). + + The observed and possible variant annotations are set to missing if the variant + does not meet the following criteria: + + - Is autosomal. + - Has a most severe transcript consequence of: "intron_variant" or + "intergenic_variant". + - Is at a site with GERP > `gerp_lower_cutoff` and < `gerp_upper_cutoff`. + + The function also returns a struct of the mutation rate globals: + + - freq_meta: Frequency metadata for the genomes dataset, filtered to the + requested genetic ancestries and downsampling level. + - ac_cutoff: Allele count cutoff used for the mutation rate calculation. + - min_cov: Minimum genome coverage used for the mutation rate calculation. + - max_cov: Maximum genome coverage used for the mutation rate calculation. + - gerp_lower_cutoff: Minimum GERP score used for the mutation rate calculation. + - gerp_upper_cutoff: Maximum GERP score used for the mutation rate calculation. + - downsampling_level: Downsampling level used for the mutation rate calculation. + - downsampling_idx: Index of the downsampling level in the frequency metadata. + - most_severe_consequence: List of most severe transcript consequences used for + the mutation rate calculation. + .. note:: -def prepare_ht_for_constraint_calculations( - ht: hl.Table, - require_exome_coverage: bool = True, - coverage_metric: str = "exome_coverage", -) -> hl.Table: + Values for `gerp_lower_cutoff` and `gerp_upper_cutoff` default to -3.9885 and + 2.6607, respectively. These values were precalculated on the GRCh37 context + table and define the 5th and 95th percentiles. + + :param locus_expr: Locus expression. + :param genomes_filter_expr: Filter expression for the genomes dataset. + :param genomes_freq_expr: Frequency array for the genomes dataset. + :param genomes_freq_meta: Frequency metadata for the genomes dataset. + :param genomes_coverage_expr: Mean genome coverage expression. + :param gerp_expr: GERP score expression. + :param most_severe_consequence_expr: Most severe consequence expression. + :param gen_ancs: List of genetic ancestries to filter the genome frequency array to. + Default is None, which includes only the full genome dataset. + :param downsampling_level: Downsampling level to use for the mutation rate + calculation. Default is 1000. + :param min_cov: Minimum genome coverage for variant to be included. Default is 15. + :param max_cov: Maximum genome coverage for variant to be included. Default is 60. + :param gerp_lower_cutoff: Minimum GERP score for variant to be included. Default + is -3.9885. + :param gerp_upper_cutoff: Maximum GERP score for variant to be included. Default + is 2.6607. + :param ac_cutoff: Allele count cutoff for variant to be included. Default is 5. + :return: Tuple containing the observed and possible variant annotations and the + globals. """ - Filter input Table and add annotations used in constraint calculations. - - Function filters to SNPs, removes rows with undefined contexts, collapses strands - to deduplicate trimer or heptamer contexts, and annotates the input Table. - - The following annotations are added to the output Table: - - ref - - alt - - methylation_level - - exome_coverage - - pass_filters - Whether the variant passed all variant filters - - annotations added by `annotate_mutation_type()`, `collapse_strand()`, and - `add_most_severe_csq_to_tc_within_vep_root()` - - :param ht: Input Table to be annotated. - :param require_exome_coverage: Filter to sites where exome coverage is defined. - Default is True. - :param coverage_metric: Name for metric to use for coverage. Default is "exome_coverage". - :return: Table with annotations. - """ - ht = trimer_from_heptamer(ht) - - if "filters" in ht.row_value.keys(): - ht = ht.annotate(pass_filters=hl.len(ht.filters) == 0) + genomes_freq_expr, genomes_freq_meta = filter_freq_for_constraint( + genomes_freq_expr, + genomes_freq_meta, + gen_ancs=None, + downsamplings=[downsampling_level], + downsampling_gen_ancs=gen_ancs, + gen_anc_label="pop", + ) + downsampling_idx = genomes_freq_meta.index( + {"group": "adj", "pop": "global", "downsampling": str(downsampling_level)} + ) - # Add annotations for 'ref' and 'alt'. - ht = ht.annotate(ref=ht.alleles[0], alt=ht.alleles[1]) + # Filter to autosomal sites (remove pseudoautosomal regions). + keep_expr = locus_expr.in_autosome() - # Filter to SNPs and context fields where the bases are either A, T, C, or G. - ht = ht.filter(hl.is_snp(ht.ref, ht.alt) & ht.context.matches(f"[ATCG]{{{3}}}")) + # Filter to sites with mean genome coverage between min_cov and max_cov. + keep_expr &= (genomes_coverage_expr >= min_cov) & (genomes_coverage_expr <= max_cov) - # Annotate mutation type (such as "CpG", "non-CpG transition", "transversion") and - # collapse strands to deduplicate the context. - ht = annotate_mutation_type(collapse_strand(ht)) + # Filter to sites where the GERP score is between 'gerp_lower_cutoff' and + # 'gerp_upper_cutoff' (ideally these values will define the 5th and 95th + # percentile of the genome-wide distribution). + keep_expr &= (gerp_expr > gerp_lower_cutoff) & (gerp_expr < gerp_upper_cutoff) - # Obtain field name for median exome coverage. - # TODO: Edit coverage field once decide what to use for v4. - exome_median_cov_field = ( - "median_approx" if "median_approx" in ht.coverage.exomes else "median" + # Filter so that the most severe annotation is 'intron_variant' or + # 'intergenic_variant' + keep_expr &= (most_severe_consequence_expr == "intron_variant") | ( + most_severe_consequence_expr == "intergenic_variant" ) - # Define methylation level cutoffs based on fields present in the 'methylation' - # annotation. - if "MEAN" in ht.methylation: - # The GRCh37 methylation resource provides a MEAN score ranging from 0-1. - methylation_expr = ht.methylation.MEAN - methylation_cutoffs = (0.6, 0.2) - elif "methylation_level" in ht.methylation: - # The GRCh38 methylation resource provides a score ranging from 0-15 for autosomes. The - # determination of this score is described in Chen et al: - # https://www.biorxiv.org/content/10.1101/2022.03.20.485034v2.full - # For chrX, methylation scores reange from 0-12, but these scores are not directly comparable - # to the autosome scores (chrX and autosomes were analyzed separately and levels are relative). - # Cutoffs to translate these scores to the 0-2 methylation level were determined by - # correlating these scores with the GRCh37 liftover scores. Proposed cutoffs are: - # 0, 1-5, 6+ for autosomes, and 0, 1-3, 4+ for chrX. - methylation_expr = ht.methylation.methylation_level - methylation_cutoffs = hl.if_else(ht.locus.contig != "chrX", (5, 0), (3, 0)) - else: - raise ValueError( - "No 'methylation_level' or 'MEAN' found in 'methylation' annotation." - ) - - # Add annotations for methylation level and median exome coverage. - ht = ht.annotate( - methylation_level=( - hl.case() - .when(ht.cpg & (methylation_expr > methylation_cutoffs[0]), 2) - .when(ht.cpg & (methylation_expr > methylation_cutoffs[1]), 1) - .default(0) + # Set up the criteria to keep high-quality sites, and sites found in less than or + # equal to 'ac_cutoff' copies in the downsampled set. + # Count possible variants in context Table, only keeping variants not in the genome + # dataset, or with AC <= 'ac_cutoff' and passing filters. + genomes_filter_freq_expr = genomes_freq_expr[downsampling_idx] + keep_expr &= hl.or_else( + (hl.len(genomes_filter_expr) == 0) & (genomes_filter_freq_expr.AC <= ac_cutoff), + True, + ) + obs_pos_expr = hl.struct( + genomes_freq=genomes_freq_expr, + **hl.or_missing( + keep_expr, single_variant_observed_and_possible_expr(genomes_freq_expr) ), - exome_coverage=ht.coverage.exomes[exome_median_cov_field], + ) + obs_pos_globals = hl.struct( + freq_meta=genomes_freq_meta, + ac_cutoff=ac_cutoff, + min_cov=min_cov, + max_cov=max_cov, + gerp_lower_cutoff=gerp_lower_cutoff, + gerp_upper_cutoff=gerp_upper_cutoff, + genetic_ancestry_groups=gen_ancs or hl.missing(hl.tarray(hl.tstr)), + downsampling_level=downsampling_level, + downsampling_idx=downsampling_idx, + most_severe_consequence=["intron_variant", "intergenic_variant"], ) - # Modify allele number annotations if present. - if "AN" in ht.row_value: - ht = ht.annotate( - exomes_AN=ht.AN.exomes[0], - exomes_AN_raw=ht.AN.exomes[1], - genomes_AN=ht.AN.genomes, - ) + return obs_pos_expr, obs_pos_globals + +def get_exome_coverage_expr( + ht: hl.Table, + exome_coverage_metric: str = "AN_percent", +) -> hl.expr.Int32Expression: + """ + Get the exome coverage expression based on the specified metric. + + The requested `exome_coverage_metric` is extracted from the exome coverage + annotations in the input `ht`: + + - "median": the expression returned is "median_approx" if it exists in + `ht.coverage.exomes`, otherwise "median". + - "AN": the expression returned is the exomes allele number (`ht.AN.exomes`). + - "AN_percent": the expression returned is the percent of samples with a + non-missing genotype, which is the exomes allele number (`ht.AN.exomes`) + divided by the total number of alleles in the exomes dataset (pulled from + `ht.an_globals.exomes.strata_sample_count` * 2) multiplied by 100. + + :param ht: Input Table with exome coverage information. + :param exome_coverage_metric: Metric to use for exome coverage. One of ["median", + "AN", "AN_percent"]. Default is "AN_percent". + :return: Exome coverage expression. + """ + if exome_coverage_metric == "median": + # Obtain field name for median exome coverage. + exome_coverage_metric = ( + "median_approx" if "median_approx" in ht.coverage.exomes else "median" + ) + cov_expr = ht.coverage.exomes[exome_coverage_metric] + elif exome_coverage_metric == "AN": + cov_expr = ht.AN.exomes + elif exome_coverage_metric == "AN_percent": # Calculate total allele number from strata_sample_count and annotate - # exomes_AN_percent (percent samples with AN) - ht = ht.annotate( - exomes_AN_percent=hl.int( - ht.exomes_AN / (ht.an_strata_sample_count.exomes[0] * 2) * 100 - ), - exomes_AN_percent_raw=hl.int( - ht.exomes_AN_raw / (ht.an_strata_sample_count.exomes[1] * 2) * 100 - ), + # exomes_AN_percent (percent samples with AN). + an_sample_count = ht.an_globals.exomes.strata_sample_count + an_meta = ht.an_globals.exomes.strata_meta + + # Get total AN count taking into account XX and XY samples for X and Y non-PAR. + xx_index = an_meta.index({"group": "adj", "sex": "XX"}) + xy_index = an_meta.index({"group": "adj", "sex": "XY"}) + xx_an_sample_count = an_sample_count[xx_index] + xy_an_sample_count = an_sample_count[xy_index] + an_count = ( + hl.case() + .when(ht.locus.in_x_nonpar(), (xx_an_sample_count * 2) + xy_an_sample_count) + .when(ht.locus.in_y_nonpar(), xy_an_sample_count) + .default(an_sample_count[0] * 2) ) - # Add most_severe_consequence annotation to 'transcript_consequences' within the - # vep root annotation. - ht = add_most_severe_csq_to_tc_within_vep_root(ht) + cov_expr = hl.int((ht.AN.exomes / an_count) * 100) + else: + raise ValueError( + f"Exome coverage metric must be one of ['median', 'AN', 'AN_percent'], not {exome_coverage_metric}" + ) - if require_exome_coverage: - # Filter out locus with undefined coverage_metric. - ht = ht.filter(hl.is_defined(ht[coverage_metric])) + logger.info("Setting 'exome_coverage' to %s", exome_coverage_metric) - return ht + return cov_expr -def create_observed_and_possible_ht( - exome_ht: hl.Table, - context_ht: hl.Table, - mutation_ht: hl.Table, +def get_exomes_observed_and_possible( + exomes_filter_expr: hl.expr.SetExpression, + exomes_freq_expr: hl.expr.ArrayExpression, + exomes_freq_meta: List[Dict[str, str]], + exomes_coverage_expr: hl.expr.Int32Expression, + gen_ancs: Optional[List[str]] = None, + include_downsamplings: bool = False, max_af: float = 0.001, - keep_annotations: Tuple[str] = ( - "context", - "ref", - "alt", - "methylation_level", - ), - pops: Tuple[str] = (), - downsamplings: Optional[List[int]] = None, - grouping: Tuple[str] = (), - coverage_metric: str = "exome_coverage", - partition_hint: int = 100, - filter_coverage_over_0: bool = False, - low_coverage_filter: int = None, - transcript_for_synonymous_filter: str = None, - global_annotation: Optional[str] = None, -) -> hl.Table: +) -> Tuple[hl.expr.StructExpression, hl.expr.StructExpression]: """ - Count the observed variants and possible variants by substitution, context, methylation level, and additional `grouping`. + Get the observed and possible variants for the exomes dataset. + + The function returns a struct indicating whether the variant should be included in + the observed and possible variant counts for the exomes dataset. The struct includes + the following fields: + + - observed_variants: This annotation is an array, where each element corresponds + to whether the variant is observed in the exomes dataset for the frequency + group at the corresponding index in the `exomes_freq` array and has an + AF <= 0.001. The boolean value is stored as an integer (0 or 1). + - possible_variants: Whether the variant is considered a possible variant in the + exomes dataset. This includes variants not in the exome dataset (exome AF + undefined), or also considered in the observed variant set. The boolean value + is stored as an integer (0 or 1). + + The observed and possible variant annotations are set to missing if the exome + coverage is undefined or the variant does not pass the exome filters. + + The function also returns a struct with the global parameters for the observed and + possible variant annotations: + + - exomes_freq_meta: Frequency metadata for the exomes dataset. + - genetic_ancestry_groups: List of genetic ancestry groups used for the + observed and possible variant annotations. + - downsamplings: List of downsamplings used for the observed and possible + variant annotations. + + :param exomes_filter_expr: Filter expression for the exomes dataset. + :param exomes_freq_expr: Frequency array for the exomes dataset. + :param exomes_freq_meta: Frequency metadata for the exomes dataset. + :param exomes_coverage_expr: Exome coverage expression. + :param gen_ancs: List of genetic ancestries to filter the exome frequency array to. + Default is None, which includes only the full exomes dataset. + :param include_downsamplings: Whether to include downsamplings in the observed and + possible variant annotations. Default is False. + :param max_af: Maximum allele frequency to consider a variant as observed. Default + is 0.001. + :return: Tuple containing the observed and possible variant annotations and the + globals. + """ + # If downsamplings are requested and 'genetic_ancestry_groups' is not specified, + # use the pared-down downsamplings list. + downsamplings = [m["downsampling"] for m in exomes_freq_meta if "downsampling" in m] + downsamplings = DOWNSAMPLINGS["v4"] if gen_ancs is None else downsamplings + downsamplings = downsamplings if include_downsamplings else None + logger.info("The following downsamplings will be used: %s", downsamplings) + + # Filter frequency array for computing the observed expression on all requested + # populations and downsamplings. + exomes_freq_expr, exomes_freq_meta = filter_freq_for_constraint( + exomes_freq_expr, + exomes_freq_meta, + gen_ancs=gen_ancs, + downsamplings=downsamplings, + downsampling_gen_ancs=gen_ancs if downsamplings is not None else None, + ) + + # If the exome coverage is undefined or the variant does not pass the exome filters, + # set the observed and possible variant annotations to missing. Otherwise, set the + # observed and possible variant annotations based on the frequency array. + obs_pos_expr = hl.struct( + **hl.or_missing( + hl.is_defined(exomes_coverage_expr) + & hl.or_else(hl.len(exomes_filter_expr) == 0, True), + single_variant_observed_and_possible_expr(exomes_freq_expr, max_af=max_af), + ) + ) + obs_pos_globals = hl.struct( + exomes_freq_meta=exomes_freq_meta, + genetic_ancestry_groups=gen_ancs or hl.missing(hl.tarray(hl.tstr)), + downsamplings=downsamplings or hl.missing(hl.tarray(hl.tstr)), + max_af=max_af, + ) - Prior to computing variant counts the following variants are removed: - - Variants not observed by any samples in the dataset: `(freq_expr.AC > 0)` - - Low-quality variants: `exome_ht.pass_filters` - - Variants with allele frequency above `max_af` cutoff: `(freq_expr.AF <= - max_af)` - - Variants that are not synonymous or in the canonical/MANE Select transcript if specified + return obs_pos_expr, obs_pos_globals - For each substitution, context, methylation level, and exome coverage, the rest of - variants in `exome_ht` are counted and annotated as `observed_variants`, and the - rest of variants in `context_ht` are counted and annotated as `possible_variants`. - The final Table is the outer-join of the filtered `exome_ht` and `context_ht` with - the `observed_variants` and `possible_variants` annotations. - The returned Table includes the following annotations: - - context - trinucleotide genomic context - - ref - the reference allele - - alt - the alternate base - - methylation_level - methylation_level - - observed_variants - observed variant counts in `exome_ht` - - possible_variants - possible variant counts in `context_ht` - - downsampling_counts_{pop} - variant counts in downsamplings for populations - in `pops` - - mu_snp - SNP mutation rate - - annotations added by `annotate_mutation_type` - - :param exome_ht: Preprocessed exome Table. - :param context_ht: Preprocessed context Table. - :param mutation_ht: Preprocessed mutation rate Table. - :param max_af: Maximum allele frequency for a variant to be included in returned - counts. Default is 0.001. - :param keep_annotations: Annotations to keep in the context Table. - :param pops: List of populations to use for downsampling counts. Default is (). - :param downsamplings: Optional List of integers specifying what downsampling - indices to obtain. Default is None, which will return all downsampling counts. - :param grouping: Annotations other than 'context', 'ref', 'alt', and - `methylation_level` to group by when counting variants. Default is - ('exome_coverage',). - :param partition_hint: Target number of partitions for aggregation. Default is 100. - :param filter_coverage_over_0: Whether to filter the exome Table and context Table - to variants with `coverage_metric` larger than 0. Default is False. - :param low_coverage_filter: Lower median coverage cutoff for coverage filter. Sites - with coverage below this cutoff will be removed from the `exome_ht` and - 'context_ht'. - :param transcript_for_synonymous_filter: Transcript to use when filtering to - synonymous variants. Choices: ["mane_select", "canonical", None]. If "canonical", will - filter to variants with a synonymous consequence in Ensembl canonical - transcripts. If "mane_select", will filter to variants with a synonymous consequence - in MANE Select transcripts. If None, no transcript/synonymous filter will be - applied. Default is None. - :param global_annotation: The annotation name to use as a global StructExpression - annotation containing input parameter values. If no value is supplied, this - global annotation will not be added. Default is None. - :param coverage_metric: Name for metric to use for coverage. Default is "exome_coverage". - :return: Table with observed variant and possible variant count. +def get_build_calibration_model_annotation( + exomes_coverage_expr: hl.expr.Int32Expression, + transcript_csq_expr: hl.expr.ArrayExpression, + cpg_expr: hl.expr.BooleanExpression, + genomic_region_expr: hl.expr.StringExpression, + synonymous_transcript_filter_field: str = "mane_select", + low_cov_cutoff: Optional[int] = None, + high_cov_cutoff: int = COVERAGE_CUTOFF, + upper_cov_cutoff: Optional[int] = None, + skip_coverage_model: bool = False, +) -> Tuple[hl.expr.StructExpression, hl.expr.StructExpression]: """ - logger.info("Setting coverage_metric to %s", coverage_metric) - - if low_coverage_filter is not None: - context_ht = context_ht.filter( - context_ht[coverage_metric] >= low_coverage_filter + Get the annotation and globals for building the calibration models. + + The build model grouping is set to missing if the variant is not a + "synonymous_variant" in a canonical or MANE Select transcript (depending on + `synonymous_transcript_filter_field`). Otherwise, it is a struct with the following + fields detailed in `calibration_model_group_expr`. + + :param exomes_coverage_expr: Exome coverage expression. + :param transcript_csq_expr: Transcript consequences expression. + :param cpg_expr: CpG expression. + :param genomic_region_expr: Genomic region expression. + :param synonymous_transcript_filter_field: Field used to filter to variants with a + transcript consequence of "synonymous_variant". Default is "mane_select". + :param low_cov_cutoff: Low coverage cutoff for the build models step. Default is + None. + :param high_cov_cutoff: High coverage cutoff for the build models step. Default is + COVERAGE_CUTOFF. + :param upper_cov_cutoff: Upper coverage cutoff for the build models step. Default is + None. + :param skip_coverage_model: Whether the coverage model should be skipped during the + build models step. Default is False. + :return: Tuple containing the build model expression and the global parameters. + """ + # Determine the canonical and mane_select parameters for + # 'filter_vep_transcript_csqs_expr' based on 'synonymous_transcript_filter_field'. + if synonymous_transcript_filter_field == "canonical": + canonical, mane_select = True, False + elif synonymous_transcript_filter_field == "mane_select": + canonical, mane_select = False, True + else: + raise ValueError( + "synonymous_transcript_filter_field must be either 'canonical' or " + "'mane_select'" ) - exome_ht = exome_ht.filter(exome_ht[coverage_metric] >= low_coverage_filter) - # Allele frequency information for high-quality genotypes (GQ >= 20; DP >= 10; and - # AB >= 0.2 for heterozygous calls) in all release samples in gnomAD. - freq_expr = exome_ht.freq[0] + # Filter the VEP transcript consequences to include only synonymous transcripts. + syn_csq_expr = filter_vep_transcript_csqs_expr( + transcript_csq_expr, + synonymous=True, + ensembl_only=True, + canonical=canonical, + mane_select=mane_select, + ) - # Set up the criteria to exclude variants not observed in the dataset, low-quality - # variants, variants with allele frequency above the `max_af` cutoff, and variants - # with exome coverage larger than 0 if requested. - keep_criteria = ( - (freq_expr.AC > 0) & exome_ht.pass_filters & (freq_expr.AF <= max_af) + # Define whether the variant should be included in the high or low coverage model. + build_expr = calibration_model_group_expr( + exomes_coverage_expr, + cpg_expr, + low_cov_cutoff=0 if low_cov_cutoff is None else low_cov_cutoff, + high_cov_cutoff=high_cov_cutoff, + upper_cov_cutoff=upper_cov_cutoff, + skip_coverage_model=skip_coverage_model, + additional_grouping_exprs={"genomic_region": genomic_region_expr}, ) - if filter_coverage_over_0: - keep_criteria &= exome_ht[coverage_metric] > 0 - keep_annotations += grouping + return hl.or_missing(syn_csq_expr.length() > 0, build_expr) - logger.info("Setting keep annotations to %s", keep_annotations) - # Keep variants that satisfy the criteria above. - filtered_exome_ht = exome_ht.filter(keep_criteria) +# TODO: We don't really need this, I just found int helpful to look over the +# chosen parameters. +def print_global_struct(t: Union[hl.Table, hl.Struct, hl.StructExpression]) -> None: + """ + Print the global struct. - # Filter context ht to sites with defined exome coverage. - context_ht = context_ht.filter(hl.is_defined(context_ht[coverage_metric])) + :param t: Table with globals or globals struct to print. + :return: None + """ + if isinstance(t, hl.Table): + t = t.globals + if isinstance(t, hl.StructExpression): + t = hl.eval(t) - # If requested keep only variants that are synonymous in either MANE Select or - # canonical transcripts. - if transcript_for_synonymous_filter is not None: - if transcript_for_synonymous_filter == "canonical": - canonical, mane_select = True, False - elif transcript_for_synonymous_filter == "mane_select": - canonical, mane_select = False, True - else: - raise ValueError( - "If transcript_for_synonymous_filter is not None, must be either" - " 'canonical' or 'mane_select'" - ) - filtered_exome_ht = filter_vep_transcript_csqs( - exome_ht.filter(keep_criteria), canonical=canonical, mane_select=mane_select - ) - context_ht = filter_vep_transcript_csqs( - context_ht, canonical=canonical, mane_select=mane_select - ) - # Count the observed variants in the entire Table and in each downsampling grouped - # by `grouping`, context, ref, alt, and methylation_level. - observed_ht = count_variants_by_group( - filtered_exome_ht.select(*list(keep_annotations) + ["freq"]), - additional_grouping=grouping, - partition_hint=partition_hint, - count_downsamplings=pops, - use_table_group_by=True, - max_af=max_af, + def _get_pretty_print_globals(global_struct: hl.Struct, level: int = 1) -> str: + output = "" + level_tab = "".join([" "] * level) + for k, v in global_struct.items(): + if isinstance(v, hl.Struct): + v = f"\n{_get_pretty_print_globals(v, level + 1)}" + + output += f"{level_tab}{k}: {v}\n" + + return output + + logger.info( + "\nThe following parameters were used: \n%s", _get_pretty_print_globals(t) ) - # TODO: Remove repartition once partition_hint bugs are resolved. - observed_ht = observed_ht.repartition(partition_hint) - observed_ht = observed_ht.transmute(observed_variants=observed_ht.variant_count) - # Filter the `exome_ht` to rows that don’t match the criteria above. - # Anti join the `context_ht` with filtered `exome_ht`, so that `context_ht` only - # has rows that match the criteria above in the `exome_ht` or are never in - # the `exome_ht`. - context_ht = context_ht.select(*keep_annotations).anti_join( - exome_ht.filter(keep_criteria, keep=False) +def prepare_ht_for_constraint_calculations( + ht: hl.Table, + exome_coverage_metric: str = "median", + gen_ancs: Optional[List[str]] = None, + include_downsamplings: bool = False, + mu_downsampling_level: int = 1000, + calculate_mutation_rate_min_cov: int = 15, + calculate_mutation_rate_max_cov: int = 60, + calculate_mutation_rate_gerp_lower_cutoff: float = -3.9885, + calculate_mutation_rate_gerp_upper_cutoff: float = 2.6607, + calculate_mutation_rate_ac_cutoff: int = 5, + max_af: float = 0.001, + build_model_low_cov_cutoff: Optional[int] = None, + build_model_high_cov_cutoff: int = COVERAGE_CUTOFF, + build_model_upper_cov_cutoff: Optional[int] = None, + apply_model_low_cov_cutoff: Optional[int] = None, + apply_model_high_cov_cutoff: int = COVERAGE_CUTOFF, + skip_coverage_model: bool = False, + synonymous_transcript_filter_field: str = "mane_select", +) -> hl.Table: + """ + Prepare Table for constraint calculations. + + This function is a wrapper around the functions that generate the annotations + required for the constraint calculations. Please see the following functions for + more information on the annotations generated: + + - `get_annotations_for_computing_mu` + - `get_exomes_observed_and_possible` + - `get_build_calibration_model_annotation` + - `get_apply_calibration_model_annotation` + + :param ht: Annotated context Table. + :param exome_coverage_metric: Metric to use for exome coverage. One of ["median", + "AN", "AN_percent"]. Default is "median". + :param gen_ancs: List of genetic ancestries to filter the frequency arrays to. + Default is None, which includes only the full dataset. + :param include_downsamplings: Whether to include downsamplings in the observed and + possible variant annotations. Default is False. + :param mu_downsampling_level: Downsampling level to use for the mutation rate + calculation. Default is 1000. + :param calculate_mutation_rate_min_cov: Minimum genome coverage for variant to be + included in the mutation rate calculation. Default is 15. + :param calculate_mutation_rate_max_cov: Maximum genome coverage for variant to be + included in the mutation rate calculation. Default is 60. + :param calculate_mutation_rate_gerp_lower_cutoff: Minimum GERP score for variant to + be included in the mutation rate calculation. Default is -3.9885. + :param calculate_mutation_rate_gerp_upper_cutoff: Maximum GERP score for variant to + be included in the mutation rate calculation. Default is 2.6607. + :param calculate_mutation_rate_ac_cutoff: Allele count cutoff for variant to be + included in the mutation rate calculation. Default is 5. + :param max_af: Maximum allele frequency to consider a variant as observed. Default + is 0.001. + :param build_model_low_cov_cutoff: Low coverage cutoff for the build models step. + Default is None. + :param build_model_high_cov_cutoff: High coverage cutoff for the build models step. + Default is COVERAGE_CUTOFF. + :param build_model_upper_cov_cutoff: Upper coverage cutoff for the build models + step. Default is None. + :param apply_model_low_cov_cutoff: Low coverage cutoff for the apply models step. + Default is COVERAGE_CUTOFF. + :param apply_model_high_cov_cutoff: High coverage cutoff for the apply models step. + Default is COVERAGE_CUTOFF. + :param skip_coverage_model: Whether the coverage model should be skipped during the + build and apply models steps. Default is False. + :param synonymous_transcript_filter_field: Field used to filter to variants with a + transcript consequence of "synonymous_variant". Default is "canonical". + :return: Table with the computed annotations. + """ + # Get the annotations relevant for computing the mutation rate. + compute_mu_expr, compute_mu_globals = get_annotations_for_computing_mu( + ht.locus, + ht.filters.genomes, + ht.freq.genomes, + ht.freq_globals.genomes.freq_meta, + ht.coverage.genomes.mean, + ht.gerp, + ht.vep.most_severe_consequence, + gen_ancs=gen_ancs, + downsampling_level=mu_downsampling_level, + min_cov=calculate_mutation_rate_min_cov, + max_cov=calculate_mutation_rate_max_cov, + gerp_lower_cutoff=calculate_mutation_rate_gerp_lower_cutoff, + gerp_upper_cutoff=calculate_mutation_rate_gerp_upper_cutoff, + ac_cutoff=calculate_mutation_rate_ac_cutoff, ) - # Count the possible variants in the context Table grouped by by - # `grouping`, context, ref, alt, and methylation_level. - possible_ht = count_variants_by_group( - context_ht, - additional_grouping=grouping, - partition_hint=partition_hint, - use_table_group_by=True, + # Get an observed and possible variant annotation for the exomes dataset. + exomes_coverage_expr = get_exome_coverage_expr(ht, exome_coverage_metric) + exomes_obs_pos_expr, exomes_obs_pos_globals = get_exomes_observed_and_possible( + ht.filters.exomes, + ht.freq.exomes, + hl.eval(ht.freq_globals.exomes.freq_meta), + exomes_coverage_expr, + gen_ancs=gen_ancs, + include_downsamplings=include_downsamplings, + max_af=max_af, ) - possible_ht = annotate_with_mu(possible_ht, mutation_ht) - possible_ht = possible_ht.transmute(possible_variants=possible_ht.variant_count) - # Outer join the Tables with possible variant counts and observed variant counts. - ht = observed_ht.join(possible_ht, "outer") - ht = ht.checkpoint(new_temp_file(prefix="constraint", extension="ht")) + # Get the annotations relevant for building the calibration models. + build_expr = get_build_calibration_model_annotation( + exomes_coverage_expr, + ht.vep.transcript_consequences, + ht.cpg, + ht.genomic_region, + synonymous_transcript_filter_field=synonymous_transcript_filter_field, + low_cov_cutoff=build_model_low_cov_cutoff, + high_cov_cutoff=build_model_high_cov_cutoff, + upper_cov_cutoff=build_model_upper_cov_cutoff, + skip_coverage_model=skip_coverage_model, + ) - # Annotate the Table with 'cpg' and 'mutation_type' (one of "CpG", "non-CpG - # transition", or "transversion"). - ht = annotate_mutation_type(ht) + # Get the annotations relevant for applying the calibration models. + apply_expr = calibration_model_group_expr( + exomes_coverage_expr, + ht.cpg, + low_cov_cutoff=apply_model_low_cov_cutoff, + high_cov_cutoff=apply_model_high_cov_cutoff, + skip_coverage_model=skip_coverage_model, + additional_grouping_exprs={"genomic_region": ht.genomic_region}, + ) - if global_annotation: - ht = ht.annotate_globals( - **{global_annotation: hl.struct(max_af=max_af, pops=pops)} - ) + # Annotate the Table with the computed annotations, and select only the relevant + # fields. + ht = ht.annotate( + exomes_coverage=exomes_coverage_expr, + compute_mu=compute_mu_expr, + calibrate_mu=hl.struct( + **exomes_obs_pos_expr, build_model=build_expr, apply_model=apply_expr + ), + ) + ht = ht.drop("freq") + + # Build a struct with the global parameters for building the calibration models. + mis_int = hl.missing(hl.tint) + handle_none = lambda x: x if x is not None else mis_int + ht = ht.select_globals( + calculate_mu_globals=compute_mu_globals, + build_models_globals=hl.struct( + synonymous_transcript_filter_field=synonymous_transcript_filter_field, + low_cov_cutoff=handle_none(build_model_low_cov_cutoff), + high_cov_cutoff=build_model_high_cov_cutoff, + upper_cov_cutoff=handle_none(build_model_upper_cov_cutoff), + skip_coverage_model=skip_coverage_model, + ), + apply_models_globals=hl.struct( + low_cov_cutoff=handle_none(apply_model_low_cov_cutoff), + high_cov_cutoff=apply_model_high_cov_cutoff, + skip_coverage_model=skip_coverage_model, + ), + **exomes_obs_pos_globals, + ) + + print_global_struct(ht) return ht -def apply_models( - exome_ht: hl.Table, - context_ht: hl.Table, +def create_training_set( + ht: hl.Table, mutation_ht: hl.Table, + partition_hint=100, +) -> hl.Table: + """ + Create the training set for the constraint model. + + The input `ht` should be prepared using `prepare_ht_for_constraint_calculations`. + The `ht` is filtered to include only the rows that have a build model annotation. + The observed and possible variants are counted by group and annotated with the + mutation rate. The Table is then checkpointed to avoid memory and shuffle issues. + + :param ht: Table prepared using `prepare_ht_for_constraint_calculations`. + :param mutation_ht: Mutation rate Table. + :param partition_hint: Partition hint for the Table. Default is 100. + :return: Training set Table. + """ + # Selecting the only fields that are needed for the training set and filtering out + # the rows that are not needed, then checkpointing the Table. This is added to + # help avoid memory and shuffle issues. + ht = ht.annotate(**ht.calibrate_mu) + ht = ht.filter(hl.is_defined(ht.build_model)) + ht = ht.select( + *MU_GROUPING, *CALIBRATION_GROUPING, "observed_variants", "possible_variants" + ) + ht = ht.checkpoint(new_temp_file("create_training_set", "ht")) + + # Aggregate and count the observed and possible variants by group. + ht = count_observed_and_possible_by_group( + ht, + ht.possible_variants, + ht.observed_variants, + additional_grouping=("methylation_level",) + CALIBRATION_GROUPING, + partition_hint=partition_hint, + ) + + # Annotate with mutation rate. + ht = annotate_with_mu(ht, mutation_ht) + + return ht + + +def create_per_variant_expected_ht( + ht: hl.Table, plateau_models: hl.StructExpression, - coverage_model: Optional[Tuple[float, float]] = None, + coverage_model: Tuple[float, float], log10_coverage: bool = True, - max_af: float = 0.001, - keep_annotations: Tuple[str] = ( - "context", - "ref", - "alt", - "methylation_level", - ), - pops: Tuple[str] = (), - downsamplings: Optional[List[int]] = None, - obs_pos_count_partition_hint: int = 2000, - expected_variant_partition_hint: int = 1000, - custom_vep_annotation: str = None, - coverage_metric: str = "exome_coverage", - high_cov_definition: int = COVERAGE_CUTOFF, - low_coverage_filter: int = None, - use_mane_select: bool = True, + filter_to_apply_variants: bool = True, ) -> hl.Table: """ - Compute the expected number of variants and observed:expected ratio using plateau models and coverage model. - - This function sums the number of possible variants times the mutation rate for all - variants, and applies the calibration model separately for CpG transitions and - other sites. For sites with coverage lower than the coverage cutoff, the value - obtained from the previous step is multiplied by the coverage correction factor. - These values are summed across the set of variants of interest to obtain the - expected number of variants. - - A brief view of how to get the expected number of variants: - mu_agg = the number of possible variants * the mutation rate (all variants) - predicted_proportion_observed = sum(plateau model slope * mu_agg + plateau model intercept) (separately for CpG transitions and other sites) - if 0 < coverage < coverage cutoff: - coverage_correction = coverage_model slope * log10(coverage) + coverage_model intercept - expected_variants = sum(predicted_proportion_observed * coverage_correction) - else: - expected_variants = sum(predicted_proportion_observed) - The expected_variants are summed across the set of variants of interest to - obtain the final expected number of variants. - - Function adds the following annotations all grouped by groupings (output of - `annotate_exploded_vep_for_constraint_groupings()`): - - observed_variants - observed variant counts annotated by `count_variants` - function - - predicted_proportion_observed (including those for each population) - the sum - of mutation rate adjusted by plateau models and possible variant counts - - possible_variants (including those for each population if `pops` is - specified) - the sum of possible variant counts derived from the context - Table - - expected_variants (including those for each population if `pops` is - specified) - the sum of expected variant counts - - mu - sum(mu_snp * possible_variant * coverage_correction) - - obs_exp - observed:expected ratio - - annotations annotated by `annotate_exploded_vep_for_constraint_groupings()` - - :param exome_ht: Exome sites Table (output of `prepare_ht_for_constraint_calculations - ()`) filtered to autosomes and pseudoautosomal regions. - :param context_ht: Context Table (output of `prepare_ht_for_constraint_calculations - ()`) filtered to autosomes and pseudoautosomal regions. - :param mutation_ht: Mutation rate Table with 'mu_snp' field. - :param plateau_models: Linear models (output of `build_models()` in - gnomad_methods`), with the values of the dictionary formatted as a - StrucExpression of intercept and slope, that calibrates mutation rate to - proportion observed for high coverage exome. It includes models for CpG sites, - non-CpG sites, and each population in `POPS`. - :param coverage_model: A linear model (output of `build_models()` in - gnomad_methods), formatted as a Tuple of intercept and slope, that calibrates a - given coverage level to observed:expected ratio. It's a correction factor for - low coverage sites. - :param log10_coverage: Whether to convert coverage sites with log10 when building the coverage model. Default is True. - :param max_af: Maximum allele frequency for a variant to be included in returned - counts. Default is 0.001. - :param keep_annotations: Annotations to keep in the context Table and exome Table. - :param pops: List of populations to use for downsampling counts. Default is (). - :param downsamplings: Optional List of integers specifying what downsampling - indices to obtain. Default is None, which will return all downsampling counts. - :param obs_pos_count_partition_hint: Target number of partitions for - aggregation when counting variants. Default is 2000. - :param expected_variant_partition_hint: Target number of partitions for sum - aggregators when computation is done. Default is 1000. - :param custom_vep_annotation: The customized model (one of - "transcript_consequences" or "worst_csq_by_gene"). Default is None. - :param coverage_metric: Name for metric to use for coverage. Default is "exome_coverage". - :param high_cov_definition: Median coverage cutoff. Sites with coverage above this cutoff - are considered well covered and was used to build plateau models. Sites - below this cutoff have low coverage and was used to build coverage models. - Default is `COVERAGE_CUTOFF`. - :param low_coverage_filter: Lower median coverage cutoff for coverage filter. - Sites with coverage below this cutoff will be removed from`exome_ht` and - 'context_ht'. - :param use_mane_select: Use MANE Select transcripts in grouping. - Only used when `custom_vep_annotation` is set to 'transcript_consequences'. - Default is True. - - :return: Table with `expected_variants` (expected variant counts) and `obs_exp` - (observed:expected ratio) annotations. + Create the per-variant expected Table. + + The input `ht` should be prepared using `prepare_ht_for_constraint_calculations`. + The `ht` is filtered to include only the rows that have an apply model annotation ( + if `filter_to_apply_variants` is True). The Table is then annotated with the + expected number of variants using `apply_models`. See the function `apply_models` + for more information on the expected annotations. + + :param ht: Table prepared using `prepare_ht_for_constraint_calculations`. + :param plateau_models: Plateau models for the constraint calculations. + :param coverage_model: Coverage model for the constraint calculations. + :param log10_coverage: Whether to use log10 coverage. Default is True. + :param filter_to_apply_variants: Whether to filter to only the rows with an apply + model annotation. Default is True. + :return: Per-variant expected Table """ - # Filter context ht to sites with defined exome coverage_metric. - context_ht = context_ht.filter(hl.is_defined(context_ht[coverage_metric])) + ht = ht.annotate(**ht.calibrate_mu) + + if filter_to_apply_variants: + ht = ht.filter(hl.is_defined(ht.apply_model)) - if low_coverage_filter is not None: - context_ht = context_ht.filter( - context_ht[coverage_metric] >= low_coverage_filter + ht = ht.annotate( + **apply_models( + ht.mu_snp, + plateau_models.get(ht.apply_model), + ht.possible_variants, + coverage_model=coverage_model, + coverage_expr=ht.exomes_coverage, + model_group_expr=ht.apply_model, + log10_coverage=log10_coverage, ) - exome_ht = exome_ht.filter(exome_ht[coverage_metric] >= low_coverage_filter) + ) + ht = ht.annotate_globals( + apply_models_globals=ht.apply_models_globals.annotate( + plateau_models=plateau_models, + coverage_model=coverage_model, + log10_coverage=log10_coverage, + ) + ) - # Add necessary constraint annotations for grouping. + return ht + + +def aggregate_per_variant_expected_ht( + ht, + mutation_ht: hl.Table, + additional_grouping: Tuple = (), + custom_vep_annotation: str = "transcript_consequences", + use_mane_select: bool = False, +): + """ + Aggregate the per-variant expected Table. + + The input `ht` should be the Table returned by `create_per_variant_expected_ht`. + The Table is exploded by the VEP annotation and aggregated by "genomic_region", + "context", "ref", "alt", "methylation_level", groupings returned by + `annotate_exploded_vep_for_constraint_groupings` and fields in + `additional_grouping` to get the observed and expected counts. + + :param ht: Table returned by `create_per_variant_expected_ht`. + :param mutation_ht: Mutation rate Table. + :param additional_grouping: Additional fields to group by. Default is (). + :param custom_vep_annotation: Custom VEP annotation to use. Default is + :param use_mane_select: Whether to include MANE Select as a group. Default is False. + :return: Table with the observed and expected counts. + """ + include_canonical_group = False + include_mane_select_group = False if custom_vep_annotation == "worst_csq_by_gene": vep_annotation = "worst_csq_by_gene" if use_mane_select: @@ -490,175 +780,43 @@ def apply_models( "'mane_select' cannot be set to True when custom_vep_annotation is set" " to 'worst_csq_by_gene'." ) - else: - vep_annotation = "transcript_consequences" + vep_annotation = custom_vep_annotation include_canonical_group = True include_mane_select_group = use_mane_select - context_ht, _ = annotate_exploded_vep_for_constraint_groupings( - ht=context_ht, - coverage_expr=context_ht[coverage_metric], - vep_annotation=vep_annotation, - include_canonical_group=include_canonical_group, - include_mane_select_group=include_mane_select_group, + ht = ht.select( + "genomic_region", + *MU_GROUPING, + *additional_grouping, + *AGGREGATE_SUM_FIELDS, + "vep", ) - exome_ht, grouping = annotate_exploded_vep_for_constraint_groupings( - ht=exome_ht, - coverage_expr=exome_ht[coverage_metric], + + ht, groupings = annotate_exploded_vep_for_constraint_groupings( + ht=ht, vep_annotation=vep_annotation, include_canonical_group=include_canonical_group, include_mane_select_group=include_mane_select_group, ) + ht = annotate_with_mu(ht, mutation_ht) + ht = ht.checkpoint(new_temp_file("annotate_exploded_vep", "ht")) - # Compute observed and possible variant counts. - ht = create_observed_and_possible_ht( - exome_ht=exome_ht, - context_ht=context_ht, - mutation_ht=mutation_ht, - max_af=max_af, - keep_annotations=keep_annotations, - pops=pops, - downsamplings=downsamplings, - grouping=grouping, - coverage_metric=coverage_metric, - partition_hint=obs_pos_count_partition_hint, - filter_coverage_over_0=True, - transcript_for_synonymous_filter=None, - ) - - # NOTE: In v2 ht.mu_snp was incorrectly multiplied here by possible_variants, but this multiplication has now been moved, - # so that it is applied after the regression within compute_expected_variants. - mu_expr = ht.mu_snp - poss_expr = ht.possible_variants - # Determine coverage correction to use based on coverage value. If no - # coverage model is provided, set to 1 as long as coverage > 0. - if log10_coverage: - logger.info("Converting coverage sites by log10.") - cov_value = hl.log10(ht.coverage) - else: - cov_value = ht.coverage - - cov_corr_expr = ( - hl.case() - .when(ht.coverage == 0, 0) - .when(ht.coverage >= high_cov_definition, 1) - .default( - (coverage_model[1] * cov_value + coverage_model[0]) - if coverage_model is not None - else 1 - ) - ) - - # Generate sum aggregators for 'mu' on the entire dataset. - agg_expr = {"mu": hl.agg.sum(mu_expr * cov_corr_expr)} - agg_expr.update( - compute_expected_variants( - ht=ht, - plateau_models_expr=plateau_models, - mu_expr=mu_expr, - cov_corr_expr=cov_corr_expr, - possible_variants_expr=poss_expr, - cpg_expr=ht.cpg, - ) - ) - downsampling_meta = {} - for pop in pops: - agg_expr.update( - compute_expected_variants( - ht=ht, - plateau_models_expr=plateau_models, - mu_expr=mu_expr, - cov_corr_expr=cov_corr_expr, - possible_variants_expr=poss_expr, - cpg_expr=ht.cpg, - pop=pop, - ) - ) + ht = ht.group_by( + "genomic_region", *MU_GROUPING, *groupings, *additional_grouping + ).aggregate(**aggregate_expected_variants_expr(ht)) - # Store which downsamplings are obtained for each pop in a - # downsampling_meta dictionary. - ds = hl.eval(get_downsampling_freq_indices(ht.freq_meta, pop=pop)) - key_names = {key for _, meta_dict in ds for key in meta_dict.keys()} - genetic_ancestry_label = "gen_anc" if "gen_anc" in key_names else "pop" - downsampling_meta[pop] = [ - x[1]["downsampling"] - for x in ds - if (x[1][genetic_ancestry_label] == pop) - & ( - int(x[1]["downsampling"]) in downsamplings - if downsamplings is not None - else True - ) - ] - - # Remove coverage from grouping. - grouping = list(grouping) - grouping.remove("coverage") - - # Aggregate the sum aggregators grouped by `grouping`. - ht = ( - ht.group_by(*grouping) - .partition_hint(expected_variant_partition_hint) - .aggregate(**agg_expr) - ) - - # TODO: Remove repartition once partition_hint bugs are resolved. - ht = ht.repartition(expected_variant_partition_hint) - - # Annotate global annotations. - coverage_model_global = coverage_model if coverage_model else "None" - ht = ht.annotate_globals( - apply_model_params=hl.struct( - max_af=max_af, - genetic_ancestry_groups=pops, - plateau_models=plateau_models, - coverage_model=coverage_model_global, - high_cov_definition=high_cov_definition, - coverage_metric=coverage_metric, - log10_coverage=log10_coverage, - downsampling_meta=downsampling_meta if downsampling_meta else "None", - ) - ) - # Compute the observed:expected ratio. - return ht.annotate(obs_exp=ht.observed_variants / ht.expected_variants) + return ht.naive_coalesce(1000) +# TODO: Move this up after review in this location. def calculate_mu_by_downsampling( - genome_ht: hl.Table, - context_ht: hl.Table, - recalculate_all_possible_summary: bool = True, - omit_methylation: bool = False, - count_singletons: bool = False, - keep_annotations: Tuple[str] = ( - "context", - "ref", - "alt", - "methylation_level", - ), - ac_cutoff: int = 5, - downsampling_level: int = 1000, + ht: hl.Table, + additional_grouping: Tuple[str] = ("methylation_level",), total_mu: float = 1.2e-08, - pops: Tuple[str] = (), - min_cov: int = 15, - max_cov: int = 60, - gerp_lower_cutoff: float = -3.9885, - gerp_upper_cutoff: float = 2.6607, ) -> hl.Table: """ - Calculate mutation rate using the downsampling with size specified by `downsampling_level` in genome sites Table. - - Prior to computing mutation rate, only the following variants are kept: - - variants with the mean coverage in the gnomAD genomes between `min_cov` and - `max_cov`. - - variants where the most severe consequence was 'intron_variant' or - 'intergenic_variant'. - - variants with the GERP score between `gerp_lower_cutoff` and - `gerp_upper_cutoff` (these default to -3.9885 and 2.6607, respectively - - these values were precalculated on the GRCh37 context Table and define the - 5th and 95th percentiles). - - high-quality variants: `genome_ht.pass_filters`. - - variants with allele count below `ac_cutoff`: `(freq_expr.AC <= ac_cutoff)`. + Calculate mutation rate. The returned Table includes the following annotations: - context - trinucleotide genomic context. @@ -670,98 +828,19 @@ def calculate_mu_by_downsampling( - mu_snp - SNP mutation rate. - annotations added by `annotate_mutation_type`. - :param genome_ht: Genome sites Table for autosome/pseudoautosomal regions. - :param context_ht: Context Table for autosome/pseudoautosomal regions. - :param recalculate_all_possible_summary: Whether to calculate possible - variants using context Table with locus that is only on an autosome or - in a pseudoautosomal region. Default is True. - :param omit_methylation: Whether to omit 'methylation_level' from the - grouping when counting variants. Default is False. - :param count_singletons: Whether to count singletons. Default is False. - :param keep_annotations: Annotations to keep in the context Table and genome - sites Table. - :param ac_cutoff: The cutoff of allele count when filtering context Table - and genome sites Table. - :param downsampling_level: The size of downsamplings will be used to count - variants. Default is 1000. + :param ht: Table returned by `prepare_ht_for_constraint_calculations`. + :param additional_grouping: Annotations other than 'context', 'ref', and 'alt'. + Default is ('methylation_level',). :param total_mu: The per-generation mutation rate. Default is 1.2e-08. - :param pops: List of populations to use for downsampling counts. If empty - Tuple is supplied, will default to '['global']'. - :param min_cov: Minimum coverage required to keep a site when calculating - the mutation rate. Default is 15. - :param max_cov: Maximum coverage required to keep a site when calculating - the mutation rate. Default is 60. - :param gerp_lower_cutoff: Minimum GERP score for variant to be included - when calculating the mutation rate. Default is -3.9885. - :param gerp_upper_cutoff: Maximum GERP score for variant to be included - when calculating the mutation rate. Default is 2.6607. :return: Mutation rate Table. """ - if not pops: - pops = ["global"] - - # Filter to autosomal sites (remove pseudoautosomal regions) between - # min_cov and max_cov. - context_ht = filter_to_autosomes( - filter_by_numeric_expr_range( - context_ht, context_ht.coverage.genomes.mean, (min_cov, max_cov) - ) - ) - genome_ht = filter_to_autosomes( - filter_by_numeric_expr_range( - genome_ht, genome_ht.coverage.genomes.mean, (min_cov, max_cov) - ) - ) - - # Filter the Table so that the most severe annotation is 'intron_variant' or - # 'intergenic_variant', and that the GERP score is between 'gerp_lower_cutoff' and - # 'gerp_upper_cutoff' (ideally these values will define the 5th and 95th - # percentile of the genome-wide distribution). - context_ht = filter_for_mu(context_ht, gerp_lower_cutoff, gerp_upper_cutoff) - genome_ht = filter_for_mu(genome_ht, gerp_lower_cutoff, gerp_upper_cutoff) - - context_ht = context_ht.select(*keep_annotations) - genome_ht = genome_ht.select(*list(keep_annotations) + ["freq", "pass_filters"]) - - # Get the frequency index of downsampling with size of `downsampling_level`. - downsampling_meta = get_downsampling_freq_indices(genome_ht.freq_meta) - downsampling_idx = hl.eval( - downsampling_meta.filter( - lambda x: x[1]["downsampling"] == str(downsampling_level) - )[0][0] - ) - freq_expr = genome_ht.freq[downsampling_idx] - - # Set up the criteria to filter out low-quality sites, and sites found in greater - # than 'ac_cutoff' copies in the downsampled set. - keep_criteria = (freq_expr.AC <= ac_cutoff) & genome_ht.pass_filters - - # Count the observed variants in the genome sites Table. - observed_ht = count_variants_by_group( - genome_ht.filter(keep_criteria).select(*list(keep_annotations) + ["freq"]), - count_downsamplings=pops, - count_singletons=count_singletons, - omit_methylation=omit_methylation, - use_table_group_by=True, - ) - - # Count possible variants in context Table, only keeping variants not in the genome - # dataset, or with AC <= 'ac_cutoff' and passing filters. - all_possible_ht = count_variants_by_group( - context_ht.anti_join(genome_ht.filter(keep_criteria, keep=False)).select( - *keep_annotations - ), - omit_methylation=omit_methylation, - use_table_group_by=True, - ) - all_possible_ht = all_possible_ht.checkpoint( - get_checkpoint_path("all_possible_summary"), - _read_if_exists=not recalculate_all_possible_summary, - overwrite=recalculate_all_possible_summary, - ) - - ht = observed_ht.annotate( - possible_variants=all_possible_ht[observed_ht.key].variant_count + # Count the observed variants in the entire Table and in each downsampling grouped + # by context, ref, alt, and 'additional_grouping'. + ht = count_observed_and_possible_by_group( + ht, + ht.compute_mu.possible_variants, + ht.compute_mu.observed_variants, + additional_grouping=additional_grouping, ) ht = ht.checkpoint(new_temp_file(prefix="constraint", extension="ht")) @@ -771,143 +850,241 @@ def calculate_mu_by_downsampling( "Total bases to use when calculating correction_factors: %f", total_bases ) - # Get the index of dowsampling with size of `downsampling_level`. - downsampling_idx = hl.eval( - downsampling_meta.map(lambda x: hl.int(x[1]["downsampling"])).index( - downsampling_level - ) - ) - # Compute the proportion observed, which represents the relative mutability of each # variant class. - ann_expr = { - "proportion_observed": ht.variant_count / ht.possible_variants, - f"proportion_observed_{downsampling_level}": ( - ht.downsampling_counts_global[downsampling_idx] / ht.possible_variants - ), - "downsamplings_frac_observed": ( - ht.downsampling_counts_global / ht.possible_variants - ), - } - - for pop in pops: - pop_counts_expr = ht[f"downsampling_counts_{pop}"] - correction_factors = ht.aggregate( - total_mu / (hl.agg.array_sum(pop_counts_expr) / total_bases), - _localize=False, - ) - downsamplings_mu_expr = ( - correction_factors * pop_counts_expr / ht.possible_variants - ) - ann_expr[f"downsamplings_mu_{'snp' if pop == 'global' else pop}"] = ( - downsamplings_mu_expr - ) - ann_expr[f"mu_snp{'' if pop == 'global' else f'_{pop}'}"] = ( - downsamplings_mu_expr[downsampling_idx] - ) - - ht = ht.annotate(**ann_expr).checkpoint( - new_temp_file(prefix="calculate_mu_by_downsampling", extension="ht") + po_expr = ht.observed_variants / ht.possible_variants + correction_factors = ht.aggregate( + total_mu / (hl.agg.array_sum(ht.observed_variants) / total_bases), + _localize=False, ) - - ht = ht.annotate_globals( - ac_cutoff=ac_cutoff, - downsampling_level=downsampling_level, - total_mu=total_mu, - min_cov=min_cov, - max_cov=max_cov, - gerp_lower_cutoff=gerp_lower_cutoff, - gerp_upper_cutoff=gerp_upper_cutoff, + mu_expr = correction_factors * ht.observed_variants / ht.possible_variants + ht = ht.annotate( + proportion_observed=po_expr, + mu=mu_expr, + mu_snp=mu_expr[ht.calculate_mu_globals.downsampling_idx], ) return annotate_mutation_type(ht) -def add_oe_lof_upper_rank_and_bin( - ht: hl.Table, use_mane_select_over_canonical: bool = True -) -> hl.Table: +# TODO: I think we decided this isn't needed right? We can just use canonical. +def filter_to_mane_select_over_canonical(ht: hl.Table) -> hl.Table: """ - Compute the rank and decile of the lof oe upper confidence interval for MANE Select or canonical ensembl transcripts. + Filter to MANE Select over canonical transcripts. - :param ht: Input Table with the value for the lof oe upper confidence interval stored in ht.lof.oe_ci.upper. - :param use_mane_select_over_canonical: Use MANE Select rather than canonical transcripts for filtering the Table. - If a gene does not have a MANE Select transcript, the canonical transcript (if available) will be used instead. Default is True. - :return: Table with anntotations added for 'upper_rank', 'upper_bin_decile'. + Filter to only ensembl transcripts of the specified transcript filter. If MANE + select is specified, and a gene does not have a MANE select transcript, use + canonical instead. + + :param ht: Table with the MANE Select and canonical annotations. + :return: Table filtered to MANE Select over canonical transcripts. """ - # Filter to only ensembl transcripts of the specified transcript filter. If MANE select is specified, and a gene - # does not have a MANE select transcript, use canonical instead. - if use_mane_select_over_canonical: - genes = ht.group_by(ht.gene_id).aggregate( - mane_present=hl.agg.any(ht.mane_select), - canonical_present=hl.agg.any(ht.canonical), + genes = ht.group_by(ht.gene_id).aggregate( + mane_present=hl.agg.any(ht.mane_select), + canonical_present=hl.agg.any(ht.canonical), + ) + genes = genes.annotate( + only_canonical=~(genes.mane_present) & (genes.canonical_present) + ) + ms_ht = ht.annotate( + _only_canonical=genes[ht.gene_id].only_canonical, + _mane_present=genes[ht.gene_id].mane_present, + ) + ms_ht = ms_ht.filter( + (ms_ht.transcript.startswith("ENST")) + & ( + (ms_ht._mane_present & ms_ht.mane_select) + | (ms_ht._only_canonical & ms_ht.canonical) ) + ) - genes = genes.annotate( - only_canonical=~(genes.mane_present) & (genes.canonical_present) - ) + return ms_ht - ms_ht = ht.annotate( - _only_canonical=genes[ht.gene_id].only_canonical, - _mane_present=genes[ht.gene_id].mane_present, - ) - total_count = ms_ht.count() - ms_ht = ms_ht.filter( - (ms_ht.transcript.startswith("ENST")) - & ( - (ms_ht._mane_present & ms_ht.mane_select) - | (ms_ht._only_canonical & ms_ht.canonical) - ) - ) - filtered_count = ms_ht.count() - logger.info( - "Retaining %d out of %d transcripts to use for rank annotations.", - filtered_count, - total_count, - ) + +# TODO: Move to gnomad_methods? +def add_oe_upper_rank_and_decile( + ht: hl.Table, + len_meta: int, + use_mane_select_over_canonical: bool = True, +) -> hl.Table: + """ + Compute the rank and decile of the oe upper confidence interval. + + :param ht: Table with the oe upper confidence interval. + :param use_mane_select_over_canonical: Use MANE Select rather than canonical + transcripts for filtering the Table when determining ranks for the lof oe upper + confidence interval. If a gene + does not have a MANE Select transcript, the canonical transcript (if available) + will be used instead. Default is True. + :return: Struct containing the rank and decile of the oe upper confidence interval. + """ + total_count = ht.count() + + if use_mane_select_over_canonical: + ms_ht = filter_to_mane_select_over_canonical(ht) else: ms_ht = ht.filter((ht.canonical) & (ht.transcript.startswith("ENST"))) - # Rank lof.oe_ci.upper in ascending order. - ms_ht = ms_ht.order_by(ms_ht.lof.oe_ci.upper).add_index(name="upper_rank") + ms_ht = ms_ht.checkpoint(new_temp_file("constraint_metrics.canonical")) - # Determine decile bins. n_transcripts = ms_ht.count() + logger.info( + "Retaining %d out of %d transcripts to use for rank annotations.", + n_transcripts, + total_count, + ) + + ms_ht = ms_ht.annotate(upper_rank=hl.empty_array(hl.tint64)) + for i in range(len_meta): + # Rank in ascending order. + ms_ht = ms_ht.order_by(ms_ht.constraint_groups[i].oe_info[0].oe_ci.upper) + ms_ht = ms_ht.add_index(name="rank") + ms_ht = ms_ht.annotate(upper_rank=ms_ht.upper_rank.append(ms_ht.rank)) + ms_ht = ms_ht.annotate( - upper_bin_decile=hl.int(ms_ht.upper_rank * 10 / n_transcripts) + upper_bin_sextile=ms_ht.upper_rank.map(lambda x: hl.int(x * 6 / n_transcripts)), + upper_bin_decile=ms_ht.upper_rank.map(lambda x: hl.int(x * 10 / n_transcripts)), ) - # Add rank and bin annotations back to original Table. + # Map rank and bin annotations back to original Table. ms_ht = ms_ht.key_by(*list(ht.key)) - ms_index = ms_ht[ht.key] - ht = ht.annotate( - lof=ht.lof.annotate( - oe_ci=ht.lof.oe_ci.annotate( - upper_rank=ms_index.upper_rank, - upper_bin_decile=ms_index.upper_bin_decile, + ms_keyed = ms_ht[ht.key] + + return ht.annotate( + constraint_groups=hl.enumerate(ht.constraint_groups).map( + lambda x: x[1].annotate( + **{ + k: ms_keyed[k][x[0]] + for k in ["upper_rank", "upper_bin_sextile", "upper_bin_decile"] + } ) ) ) - return ht + +# TODO: Move to gnomad_methods? +def build_constraint_consequence_groups( + csq_expr: hl.expr.ArrayExpression, + lof_modifier_expr: hl.expr.StringExpression, + classic_lof_annotations: Tuple = ( + "stop_gained", + "splice_donor_variant", + "splice_acceptor_variant", + ), + additional_groupings: Dict[str, Dict[str, hl.expr.BooleanExpression]] = None, + additional_grouping_combinations: List[List[str]] = None, +) -> Tuple[List[hl.expr.BooleanExpression], List[Dict[str, str]]]: + """ + Build constraint consequence groups. + + The function builds constraint groups based on the consequence expression and LoF + modifier expression. By default, the following groups are built: + + - csq_set: synonymous_variant, missense_variant + - lof: classic, hc_lc, classic_hc_lc, hc + + The resulting meta and cooresponding constraint group filters are: + + - {"csq_set": "syn"}: synonymous_variant + - {"csq_set": "mis"}: missense_variant + - {"lof": "classic"}: classic LoF annotations + - {"lof": "hc_lc"}: LoFTEE HC or LC + - {"lof": "classic_hc_lc"}: classic LoF annotations with LoFTEE HC or LC + - {"lof": "hc"}: LoF annotations with LoFTEE HC + + Additional groupings can be added to the constraint groups by specifying the + `additional_groupings` parameter, and grouping combinations can also be added + by specifying the `additional_grouping_combinations` parameter. + + :param csq_expr: Consequence expression. + :param lof_modifier_expr: LoF modifier expression. + :param classic_lof_annotations: Classic LoF Annotations used to filter the input + Table. Default is {"stop_gained", "splice_donor_variant", + "splice_acceptor_variant"}. + :param additional_groupings: Additional groupings to add to the constraint groups. + Default is None. + :param additional_grouping_combinations: Additional grouping combinations to add to + the constraint groups. Default is None. + :return: Tuple containing the constraint group filters and the meta. + """ + lof_classic_expr = hl.literal(set(classic_lof_annotations)).contains(csq_expr) + lof_hc_expr = lof_modifier_expr == "HC" + lof_hc_lc_expr = lof_hc_expr | (lof_modifier_expr == "LC") + mis_expr = csq_expr == "missense_variant" + annotation_dict = { + "csq_set": {"syn": csq_expr == "synonymous_variant", "mis": mis_expr}, + "lof": { + # Filter to classic LoF annotations. + "classic": lof_classic_expr, + # Filter to LOFTEE HC or LC. + "hc_lc": lof_hc_lc_expr, + # Filter to classic LoF annotations with LOFTEE HC or LC. + "classic_hc_lc": lof_classic_expr & lof_hc_lc_expr, + # Filter to LoF annotations with LOFTEE HC. + "hc": lof_hc_expr, + }, + } + + annotation_dict.update(additional_groupings or {}) + additional_grouping_combinations = additional_grouping_combinations or [] + + grouping_combinations = [["csq_set"], ["lof"]] + grouping_combinations.extend(additional_grouping_combinations) + + meta = generate_filter_combinations( + grouping_combinations, + {k: list(v.keys()) for k, v in annotation_dict.items()}, + ) + constraint_group_filters = [ + functools.reduce(operator.ior, [annotation_dict[k][v] for k, v in m.items()]) + for m in meta + ] + + return constraint_group_filters, meta + + +# TODO: Move to gnomad_methods? +def convert_multi_array_to_array_of_structs( + t: Union[hl.Table, hl.expr.StructExpression], + array_fields_to_combine: List[str], + new_array_field: str, +) -> hl.Table: + """ + Convert multiple arrays to an array of structs. + + :param t: Table or Struct to convert. + :param array_fields_to_combine: Array fields to combine. + :param new_array_field: Name of the new array field. + :return: Table with the array fields combined into an array of structs named + `new_array_field`. + """ + logger.warning("This function assumes that all arrays have the same length!") + return t.annotate( + **{ + new_array_field: hl.range(t[array_fields_to_combine[0]].length()).map( + lambda i: hl.struct(**{f: t[f][i] for f in array_fields_to_combine}) + ) + } + ).drop(*array_fields_to_combine) def compute_constraint_metrics( ht: hl.Table, gencode_ht: hl.Table, - keys: Tuple[str] = ("gene", "transcript", "canonical"), - classic_lof_annotations: Tuple[str] = ( + keys: Tuple = ("gene", "transcript", "canonical"), + classic_lof_annotations: Tuple = ( "stop_gained", "splice_donor_variant", "splice_acceptor_variant", ), - pops: Tuple[str] = (), + additional_groupings: Dict[str, Dict[str, hl.expr.BooleanExpression]] = None, + additional_grouping_combinations: List[List[str]] = None, expected_values: Optional[Dict[str, float]] = None, min_diff_convergence: float = 0.001, raw_z_outlier_threshold_lower_lof: float = -8.0, raw_z_outlier_threshold_lower_missense: float = -8.0, raw_z_outlier_threshold_lower_syn: float = -8.0, raw_z_outlier_threshold_upper_syn: float = 8.0, - include_os: bool = False, use_mane_select_over_canonical: bool = True, ) -> hl.Table: """ @@ -931,7 +1108,6 @@ def compute_constraint_metrics( :param classic_lof_annotations: Classic LoF Annotations used to filter the input Table. Default is {"stop_gained", "splice_donor_variant", "splice_acceptor_variant"}. - :param pops: List of populations used to compute constraint metrics. Default is (). :param expected_values: Dictionary containing the expected values for 'Null', 'Rec', and 'LI' to use as starting values. :param min_diff_convergence: Minimum iteration change in LI to consider the EM @@ -940,170 +1116,155 @@ def compute_constraint_metrics( :param raw_z_outlier_threshold_lower_missense: Value at which the raw z-score is considered an outlier for missense variants. Values below this threshold will be considered outliers. Default is -8.0. :param raw_z_outlier_threshold_lower_syn: Lower value at which the raw z-score is considered an outlier for synonymous variants. Values below this threshold will be considered outliers. Default is -8.0. :param raw_z_outlier_threshold_upper_syn: Upper value at which the raw z-score is considered an outlier for synonymous variants. Values above this threshold will be considered outliers. Default is 8.0. - :param include_os: Whether or not to include OS (other splice) as a grouping when - stratifying calculations by lof HC. :param use_mane_select_over_canonical: Use MANE Select rather than canonical transcripts for filtering the Table when determining ranks for the lof oe upper confidence interval. If a gene does not have a MANE Select transcript, the canonical transcript (if available) will be used instead. Default is True. :param gencode_ht: Table containing GENCODE annotations. :return: Table with pLI scores, observed:expected ratio, confidence interval of the observed:expected ratio, and z scores. """ - if expected_values is None: - expected_values = {"Null": 1.0, "Rec": 0.706, "LI": 0.207} - # This function aggregates over genes in all cases, as XG spans PAR and non-PAR X. - # `annotation_dict` stats the rule of filtration for each annotation. - annotation_dict = { - # Filter to classic LoF annotations with LOFTEE HC or LC. - "lof_hc_lc": hl.literal(set(classic_lof_annotations)).contains(ht.annotation) - & ((ht.modifier == "HC") | (ht.modifier == "LC")), - # Filter to LoF annotations with LOFTEE HC. - "lof": ht.modifier == "HC", - # Filter to missense variants. - "mis": ht.annotation == "missense_variant", - # Filter to probably damaging missense variants predicted by PolyPen-2. - "mis_pphen": ht.modifier == "probably_damaging", - # Filter to synonymous variants. - "syn": ht.annotation == "synonymous_variant", - } - - # Define two lists of 'annotation_dict' keys that require different computations. - # The 90% CI around obs:exp and z-scores are only computed for lof, mis, and syn. - oe_ann = ["lof", "mis", "syn"] - # pLI scores are only computed for LoF variants. - lof_ann = ["lof_hc_lc", "lof"] - - # Create dictionary with outlier z-score thresholds with annotation as key - # and list of thresholds [lower, upper] as values. - z_score_outlier_dict = { - "lof": [raw_z_outlier_threshold_lower_lof, None], - "mis": [raw_z_outlier_threshold_lower_missense, None], - "syn": [raw_z_outlier_threshold_lower_syn, raw_z_outlier_threshold_upper_syn], - } + # Build constraint groups. + constraint_group_filters_expr, meta = build_constraint_consequence_groups( + ht.annotation, + ht.modifier, + classic_lof_annotations=classic_lof_annotations, + additional_groupings=additional_groupings, + additional_grouping_combinations=additional_grouping_combinations, + ) + ht = ht.annotate(constraint_groups=constraint_group_filters_expr) + ht = ht.annotate_globals(constraint_group_meta=meta) + ht = ht.checkpoint( + new_temp_file("constraint_metrics.constraint_group_filters", "ht") + ) - if include_os: - # Filter to LoF annotations with LOFTEE HC or OS. - annotation_dict.update( - {"lof_hc_os": (ht.modifier == "HC") | (ht.modifier == "OS")} + # Group by keys and get an aggregate sum of mu_snp, observed_variants, + # possible_variants, predicted_propotion_observed, coverage_correction, and + # expected_variants for each constraint group. + ht = ht.group_by(*keys).aggregate( + constraint_groups=hl.agg.array_agg( + lambda f: hl.agg.filter(f, aggregate_expected_variants_expr(ht)), + ht.constraint_groups, ) - lof_ann.append("lof_hc_os") + ) - # Compute the observed:expected ratio. Will not compute per pop for "mis_pphen". - ht = ht.group_by(*keys).aggregate( - **{ - ann: oe_aggregation_expr( - ht, - filter_expr, - pops=() if ann == "mis_pphen" else pops, - exclude_mu_sum=True if ann == "mis_pphen" else False, - ) - for ann, filter_expr in annotation_dict.items() - } + # Add a 'no_variants' annotation indicating that there are zero observed variants + # summed across pLoF, missense, and synonymous variants. + ht = ht.annotate( + no_variants=hl.sum( + ht.constraint_groups.map(lambda x: hl.or_else(x.observed_variants[0], 0)) + ) + == 0 ) + # Filter to only rows with at least 1 obs or exp across all keys in annotation_dict. ht = ht.filter( - hl.sum( - [ - hl.or_else(ht[ann].obs, 0) + hl.or_else(ht[ann].exp, 0) - for ann in annotation_dict - ] + ~ht.no_variants + | hl.any( + ht.constraint_groups.map( + lambda x: (hl.or_else(x.expected_variants[0], 0) > 0) + ) ) - > 0 - ) - ht = ht.checkpoint( - new_temp_file(prefix="compute_constraint_metrics", extension="ht") ) - # Compute the pLI scores for LoF variants. - ann_expr = { - ann: ht[ann].annotate( - **compute_pli( - ht, - obs_expr=ht[ann].obs, - exp_expr=ht[ann].exp, - expected_values=expected_values, - min_diff_convergence=min_diff_convergence, + # Change format of arrays in constraint_groups to an array of structs. + array_fields_to_combine = [ + k + for k, v in ht.constraint_groups.dtype._element_type.items() + if isinstance(v, hl.tarray) + ] + ht = ht.annotate( + constraint_groups=ht.constraint_groups.map( + lambda x: convert_multi_array_to_array_of_structs( + x, array_fields_to_combine, "oe_info" ) ) - for ann in lof_ann - } + ) + ht = ht.checkpoint( + new_temp_file("constraint_metrics.constraint_group_filters.agg", "ht") + ) - # Add a 'no_variants' flag indicating that there are zero observed variants summed - # across pLoF, missense, and synonymous variants. - constraint_flags_expr = { - "no_variants": hl.sum([hl.or_else(ht[ann].obs, 0) for ann in oe_ann]) == 0 - } - constraint_flags = {} - for ann in oe_ann: - obs_expr = ht[ann].obs - exp_expr = ht[ann].exp - # Compute the 90% confidence interval around the observed:expected ratio. - oe_ci_expr = oe_confidence_interval(obs_expr, exp_expr) - # Compute raw z-scores. - raw_z_expr = calculate_raw_z_score(obs_expr, exp_expr) - # Add flags that define why constraint will not be calculated. - ann_constraint_flags_expr = get_constraint_flags( - exp_expr=exp_expr, - raw_z_expr=raw_z_expr, - raw_z_lower_threshold=z_score_outlier_dict[ann][0], - raw_z_upper_threshold=z_score_outlier_dict[ann][1], - flag_postfix=ann, + def _add_oe_ci_z( + oe_info: hl.expr.StructExpression, + m: Dict[str, str], + ) -> hl.expr.StructExpression: + """ + Add oe, oe_ci, and z_raw to the oe_info struct. + + :param oe_info: Struct containing the observed and expected variants. + :return: Struct containing oe, oe_ci, and z_raw. + """ + obs = oe_info.observed_variants + exp = oe_info.expected_variants + z_raw = calculate_raw_z_score(obs, exp) + z_threshold = dict( + { + "lof": (raw_z_outlier_threshold_lower_lof, None), + "mis": (raw_z_outlier_threshold_lower_missense, None), + "syn": ( + raw_z_outlier_threshold_lower_syn, + raw_z_outlier_threshold_upper_syn, + ), + } ) - constraint_flags_expr.update(ann_constraint_flags_expr) - # The constraint_flags dict is used to filter the final ht.constraint_flags - # annotation to the flags that should be considered in the z-score 'sd' - # computation of the specified ann. - constraint_flags[ann] = hl.set( - ann_constraint_flags_expr.keys() | {"no_variants"} + z_threshold = z_threshold.get( + hl.coalesce(m.get("lof"), m.get("csq_set", "None")), + (None, None), ) - # Add initial ann to ann_expr if it isn't present. - # The ann_expr dict will already have all ann in lof_ann. - if ann not in ann_expr: - ann_expr[ann] = ht[ann] - - ann_expr[ann] = ann_expr[ann].annotate( - oe_ci=oe_ci_expr, - z_raw=raw_z_expr, + flags = get_constraint_flags(exp, z_raw, z_threshold[0], z_threshold[1]) + return oe_info.annotate( + oe=divide_null(obs, exp), + oe_ci=oe_confidence_interval(obs, exp), + z_raw=z_raw, + flags=add_filters_expr(filters=flags), ) - ann_expr["constraint_flags"] = add_filters_expr(filters=constraint_flags_expr) - ht = ht.annotate(**ann_expr) - ht = ht.checkpoint( - new_temp_file(prefix="compute_constraint_metrics", extension="ht") + # Annotate with the observed:expected ratio, 95% confidence interval around the + # observed:expected ratio, and z scores for each constraint group. + ht = ht.annotate( + constraint_groups=hl.map( + lambda x, m: x.annotate( + oe_info=x.oe_info.map(lambda oe: _add_oe_ci_z(oe, m)) + ), + ht.constraint_groups, + meta, + ) ) + # ht = ht.annotate(constraint_flags=...) + ht = ht.checkpoint(new_temp_file("constraint_metrics.oe.oe_ci.z_raw", "ht")) # Add z-score 'sd' annotation to globals. ht = ht.annotate_globals( sd_raw_z=ht.aggregate( - hl.struct( - **{ - ann: calculate_raw_z_score_sd( - raw_z_expr=ht[ann].z_raw, - flag_expr=ht.constraint_flags.intersection( - constraint_flags[ann] - ), - mirror_neg_raw_z=(ann != "syn"), + hl.agg.filter( + ~ht.no_variants, + [ + calculate_raw_z_score_sd( + ht.constraint_groups[i].oe_info[0].z_raw, + ht.constraint_groups[i].oe_info[0].flags, + mirror_neg_raw_z=m.get("csq_set") != "syn", ) - for ann in oe_ann - } + for i, m in enumerate(meta) + ], ) ) ) # Compute z-score from raw z-score and standard deviations. ht = ht.annotate( - **{ - ann: ht[ann].annotate(z_score=ht[ann].z_raw / ht.sd_raw_z[ann]) - for ann in oe_ann - } + constraint_groups=hl.map( + lambda x, sd_raw_z: x.annotate(z_score=x.oe_info[0].z_raw / sd_raw_z), + ht.constraint_groups, + ht.sd_raw_z, + ) ) - ht = ht.checkpoint(new_temp_file(prefix="z_scores", extension="ht")) + # Add a rank and decile of the upper confidence interval for MANE Select or + # canonical ensembl transcripts. + ht = add_oe_upper_rank_and_decile(ht, len(meta), use_mane_select_over_canonical) - # Compute the rank and decile of the lof oe upper confidence - # interval for MANE Select or canonical ensembl transcripts. - ht = add_oe_lof_upper_rank_and_bin( - ht, use_mane_select_over_canonical=use_mane_select_over_canonical - ) + # TODO: Add back pLI computation + # Compute the observed:expected ratio. + if expected_values is None: + expected_values = {"Null": 1.0, "Rec": 0.706, "LI": 0.207} # Add transcript annotations from GENCODE. ht = add_gencode_transcript_annotations(ht, gencode_ht) @@ -1111,6 +1272,7 @@ def compute_constraint_metrics( return ht +# TODO: Move to gnomad_methods? def calculate_gerp_cutoffs(ht: hl.Table) -> Tuple[float, float]: """ Find GERP cutoffs determined by the 5% and 95% percentiles. @@ -1143,75 +1305,3 @@ def calculate_gerp_cutoffs(ht: hl.Table) -> Tuple[float, float]: cutoff_upper = list(filter(lambda i: i[1] < 0.95, zipped))[-1][0] return cutoff_lower, cutoff_upper - - -def annotate_context_ht( - ht: hl.Table, - coverage_hts: Dict[str, hl.Table], - an_hts: Dict[str, hl.Table], - methylation_ht: hl.Table, - gerp_ht: hl.Table, -) -> hl.Table: - """ - Split multiallelic sites if needed and add 'methylation', 'coverage', and 'gerp' annotation to context Table with VEP annotation. - - .. note:: - Checks for 'was_split' annotation in Table. If not present, splits - multiallelic sites. - - :param ht: Input context Table with VEP annotation. - :param coverage_hts: A Dictionary with key as one of 'exomes' or 'genomes' and - values as corresponding coverage Tables. - :param an_hts: A Dictionary with key as one of 'exomes' or 'genomes' and - values as corresponding allele number Tables. - :param methylation_ht: Methylation Table. - :param gerp_ht: Table with GERP annotation. - :return: Table with sites split and necessary annotations. - """ - # Check if context Table is split, and if not, split multiallelic sites. - if "was_split" not in list(ht.row): - ht = hl.split_multi_hts(ht) - - # Filter Table to only contigs 1-22, X, Y. - ref = get_reference_genome(ht.locus) - ht = hl.filter_intervals( - ht, [hl.parse_locus_interval(c, ref.name) for c in ref.contigs[:24]] - ) - - # If neccessary, pull out first element of coverage statistics (which includes all samples). Relevant to v4, where - # coverage stats include additional elements to stratify by ukb subset and - # platforms. - if "coverage_stats" in coverage_hts["exomes"].row: - coverage_hts["exomes"] = coverage_hts["exomes"].transmute( - **coverage_hts["exomes"].coverage_stats[0] - ) - - # Add 'methylation', 'coverage', and 'gerp' annotation. - ht = ht.annotate( - methylation=methylation_ht[ht.locus], - coverage=hl.struct( - **{loc: coverage_ht[ht.locus] for loc, coverage_ht in coverage_hts.items()} - ), - gerp=gerp_ht[ht.locus].S, - ) - ht = ht.annotate(gerp=hl.if_else(hl.is_missing(ht.gerp), 0, ht.gerp)) - - # Add allele number annotation and an_strata_sample_count global - # annotation if allele number hts are supplied. - if len(an_hts) > 0: - ht = ht.annotate( - AN=hl.struct( - **{data_type: an_ht[ht.locus].AN for data_type, an_ht in an_hts.items()} - ) - ) - - # Add strata sample count for allele number to globals. - strata_sample_counts = { - data_type: an_ht.strata_sample_count.collect()[0] - for data_type, an_ht in an_hts.items() - } - ht = ht.annotate_globals( - an_strata_sample_count=hl.struct(**strata_sample_counts) - ) - - return ht From 0ebc66d344149add434190c7c142dd6d68b8a5f8 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Tue, 11 Mar 2025 13:09:03 -0600 Subject: [PATCH 02/38] Fixes while testing --- .../pipeline/constraint_pipeline.py | 7 +- gnomad_constraint/resources/resource_utils.py | 1 + gnomad_constraint/utils/constraint.py | 85 ++++++++++++------- 3 files changed, 59 insertions(+), 34 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index ac6cee23..cc0078d8 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -524,7 +524,8 @@ def main(args): ht = res.temp_preprocess_data_ht.ht() print_global_struct(ht.apply_models_globals) ht = create_per_variant_expected_ht( - annotate_with_mu(ht, res.mutation_ht.ht().select("mu_snp")), + ht, + res.mutation_ht.ht().select("mu_snp"), res.model_plateau.he(), coverage_model=None if skip_coverage_model else res.model_coverage.he(), log10_coverage=log10_coverage, @@ -754,10 +755,10 @@ def main(args): "Lower exome coverage cutoff to use throughout the pipeline. Sites with" " coverage below this cutoff will be excluded when creating the training" " set, building and applying models, and computing constraint metrics." - " Default is 0." + " Default is None." ), type=int, - default=0, + default=None, ) preprocess_args.add_argument( "--max-af", diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 94cc51e3..7dd6a522 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -369,6 +369,7 @@ def get_apply_models( version=version, custom_vep_annotation=custom_vep_annotation, ) + # TODO: Change this name. return TableResource( f"{get_constraint_root(version, test, post_fix)}/apply_models/{custom_vep_annotation}/gnomad.v{version}.apply.per_variant_expected.ht" ) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 469c3fe0..59ce41b1 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -461,6 +461,7 @@ def get_build_calibration_model_annotation( upper_cov_cutoff=upper_cov_cutoff, skip_coverage_model=skip_coverage_model, additional_grouping_exprs={"genomic_region": genomic_region_expr}, + cpg_in_high_only=True, ) return hl.or_missing(syn_csq_expr.length() > 0, build_expr) @@ -649,6 +650,9 @@ def prepare_ht_for_constraint_calculations( **exomes_obs_pos_globals, ) + # TODO: Remove this after we have all methylation. + ht = ht.filter(ht.locus.in_autosome()) + print_global_struct(ht) return ht @@ -699,6 +703,7 @@ def create_training_set( def create_per_variant_expected_ht( ht: hl.Table, + mutation_ht: hl.Table, plateau_models: hl.StructExpression, coverage_model: Tuple[float, float], log10_coverage: bool = True, @@ -714,6 +719,7 @@ def create_per_variant_expected_ht( for more information on the expected annotations. :param ht: Table prepared using `prepare_ht_for_constraint_calculations`. + :param mutation_ht: Mutation rate Table. :param plateau_models: Plateau models for the constraint calculations. :param coverage_model: Coverage model for the constraint calculations. :param log10_coverage: Whether to use log10 coverage. Default is True. @@ -726,10 +732,12 @@ def create_per_variant_expected_ht( if filter_to_apply_variants: ht = ht.filter(hl.is_defined(ht.apply_model)) + ht = annotate_with_mu(ht, mutation_ht) + ht = ht.annotate( **apply_models( ht.mu_snp, - plateau_models.get(ht.apply_model), + plateau_models.get(ht.apply_model.model_group), ht.possible_variants, coverage_model=coverage_model, coverage_expr=ht.exomes_coverage, @@ -771,37 +779,51 @@ def aggregate_per_variant_expected_ht( :param use_mane_select: Whether to include MANE Select as a group. Default is False. :return: Table with the observed and expected counts. """ - include_canonical_group = False - include_mane_select_group = False - if custom_vep_annotation == "worst_csq_by_gene": - vep_annotation = "worst_csq_by_gene" - if use_mane_select: - raise ValueError( - "'mane_select' cannot be set to True when custom_vep_annotation is set" - " to 'worst_csq_by_gene'." - ) - else: - vep_annotation = custom_vep_annotation - include_canonical_group = True - include_mane_select_group = use_mane_select - - ht = ht.select( - "genomic_region", - *MU_GROUPING, - *additional_grouping, - *AGGREGATE_SUM_FIELDS, - "vep", - ) - - ht, groupings = annotate_exploded_vep_for_constraint_groupings( - ht=ht, - vep_annotation=vep_annotation, - include_canonical_group=include_canonical_group, - include_mane_select_group=include_mane_select_group, + # include_canonical_group = False + # include_mane_select_group = False + # if custom_vep_annotation == "worst_csq_by_gene": + # vep_annotation = "worst_csq_by_gene" + # if use_mane_select: + # raise ValueError( + # "'mane_select' cannot be set to True when custom_vep_annotation is set" + # " to 'worst_csq_by_gene'." + # ) + # else: + # vep_annotation = custom_vep_annotation + # include_canonical_group = True + # include_mane_select_group = use_mane_select + + # ht = ht.filter(hl.is_defined(ht.possible_variants)) + # ht = ht.select( + # "genomic_region", + # *MU_GROUPING, + # *additional_grouping, + # *AGGREGATE_SUM_FIELDS, + # "vep", + # ) + + # ht, groupings = annotate_exploded_vep_for_constraint_groupings( + # ht=ht, + # vep_annotation=vep_annotation, + # include_canonical_group=include_canonical_group, + # include_mane_select_group=include_mane_select_group, + # ) + # ht = annotate_with_mu(ht, mutation_ht) + # t = ht.checkpoint(new_temp_file("annotate_exploded_vep", "ht")) + ht = hl.read_table( + "gs://gnomad-tmp-4day/annotate_exploded_vep-JYgMWrtvv6cgCX9BRte85A.ht" ) - ht = annotate_with_mu(ht, mutation_ht) - ht = ht.checkpoint(new_temp_file("annotate_exploded_vep", "ht")) + groupings = [ + "annotation", + "modifier", + "gene", + "gene_id", + "transcript", + "canonical", + "mane_select", + ] + ht = ht.filter(hl.is_defined(ht.possible_variants)) ht = ht.group_by( "genomic_region", *MU_GROUPING, *groupings, *additional_grouping ).aggregate(**aggregate_expected_variants_expr(ht)) @@ -1158,7 +1180,7 @@ def compute_constraint_metrics( # Filter to only rows with at least 1 obs or exp across all keys in annotation_dict. ht = ht.filter( ~ht.no_variants - | hl.any( + & hl.any( ht.constraint_groups.map( lambda x: (hl.or_else(x.expected_variants[0], 0) > 0) ) @@ -1249,6 +1271,7 @@ def _add_oe_ci_z( ) # Compute z-score from raw z-score and standard deviations. + # TODO: Need to fix z_score ht = ht.annotate( constraint_groups=hl.map( lambda x, sd_raw_z: x.annotate(z_score=x.oe_info[0].z_raw / sd_raw_z), From e5f29aa4bf94f4ec1d1440857a39218e8aeabd51 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 12 Mar 2025 12:04:20 -0600 Subject: [PATCH 03/38] Simplify resources --- .../pipeline/constraint_pipeline.py | 41 +- gnomad_constraint/resources/resource_utils.py | 436 ++++++++---------- gnomad_constraint/utils/constraint.py | 76 +-- 3 files changed, 265 insertions(+), 288 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index cc0078d8..4e29b851 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -242,11 +242,17 @@ def get_constraint_resources( input_hts[f"{d}_sites_ht"] = constraint_res.get_sites_resource(d, version) input_hts[f"{d}_an_ht"] = all_sites_an(d) + common_params = { + "version": version, + "test": test, + "post_fix": post_fix, + } + prepare_context = PipelineStepResourceCollection( "--prepare-context-ht", output_resources={ "annotated_context_ht": constraint_res.get_annotated_context_ht( - version, test, post_fix + **common_params ) }, input_resources={"gnomAD resources": input_hts}, @@ -255,7 +261,7 @@ def get_constraint_resources( "preprocess data for downstream steps", output_resources={ "temp_preprocess_data_ht": constraint_res.get_preprocessed_ht( - version, test, post_fix + **common_params ), }, pipeline_input_steps=[prepare_context], @@ -268,23 +274,22 @@ def get_constraint_resources( calculate_mutation_rate = PipelineStepResourceCollection( "--calculate-mutation-rate", output_resources={ - "mutation_ht": constraint_res.get_mutation_ht(version, test, post_fix) + "mutation_ht": constraint_res.get_mutation_ht(**common_params) }, pipeline_input_steps=[preprocess_data], ) create_training_set = PipelineStepResourceCollection( "--create-training-set", output_resources={ - f"train_ht": constraint_res.get_training_dataset(version, test, post_fix), - f"train_tsv": constraint_res.get_training_tsv_path(version, test, post_fix), + f"train_ht": constraint_res.get_training_dataset(**common_params), + f"train_tsv": constraint_res.get_training_tsv_path(**common_params), }, pipeline_input_steps=[preprocess_data, calculate_mutation_rate], ) build_models = PipelineStepResourceCollection( "--build-models", output_resources={ - f"model_{m}": constraint_res.get_models(m, version, test, post_fix) - for m in models + f"model_{m}": constraint_res.get_models(m, **common_params) for m in models }, pipeline_input_steps=[create_training_set], ) @@ -292,7 +297,7 @@ def get_constraint_resources( "--apply-models-per-variant", output_resources={ "per_variant_apply_ht": constraint_res.get_per_variant_expected_dataset( - custom_vep_annotation, version, test, post_fix + custom_vep_annotation, **common_params ) }, pipeline_input_steps=[preprocess_data, calculate_mutation_rate, build_models], @@ -300,8 +305,8 @@ def get_constraint_resources( aggregate_per_variant_expected = PipelineStepResourceCollection( "--aggregate-per-variant-expected", output_resources={ - f"apply_ht": constraint_res.get_apply_models( - custom_vep_annotation, version, test, post_fix + f"apply_ht": constraint_res.get_aggregated_per_variant_expected( + custom_vep_annotation, **common_params ) }, pipeline_input_steps=[ @@ -314,7 +319,7 @@ def get_constraint_resources( "--compute-constraint-metrics", output_resources={ "constraint_metrics_ht": constraint_res.get_constraint_metrics_dataset( - custom_vep_annotation, version, test, post_fix + custom_vep_annotation, **common_params ) }, pipeline_input_steps=[aggregate_per_variant_expected], @@ -324,12 +329,10 @@ def get_constraint_resources( "--export-tsv", output_resources={ "constraint_metrics_tsv": constraint_res.get_constraint_tsv_path( - version, test, post_fix + **common_params ), "downsampling_constraint_metrics_tsv": ( - constraint_res.get_downsampling_constraint_tsv_path( - version, test, post_fix - ) + constraint_res.get_downsampling_constraint_tsv_path(**common_params) ), }, pipeline_input_steps=[compute_constraint_metrics], @@ -430,8 +433,9 @@ def main(args): res = resources.preprocess_data res.check_resource_existence() ht = res.annotated_context_ht.ht() + ht = filter_for_test(ht, use_gene_list=test_gene_list) if test else ht ht = prepare_ht_for_constraint_calculations( - (filter_for_test(ht, use_gene_list=test_gene_list) if test else ht), + ht, exome_coverage_metric=args.exome_coverage_metric, gen_ancs=args.genetic_ancestry_groups, include_downsamplings=args.include_downsamplings, @@ -446,6 +450,7 @@ def main(args): apply_model_low_cov_cutoff=args.pipeline_low_coverage_filter, apply_model_high_cov_cutoff=args.apply_model_high_cov_definition, skip_coverage_model=skip_coverage_model, + additional_grouping_exprs={"sfs_bin": ht.sfs_bin}, ) ht.write(res.temp_preprocess_data_ht.path, overwrite=overwrite) @@ -631,7 +636,9 @@ def main(args): finally: logger.info("Copying log to logging bucket...") - hl.copy_log(constraint_res.get_logging_path("constraint_pipeline", version)) + hl.copy_log( + constraint_res.get_logging_path("constraint_pipeline", version=version) + ) if __name__ == "__main__": diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 7dd6a522..fbba4f88 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -14,7 +14,6 @@ TableResource, VersionedTableResource, ) -from gnomad_qc.v4.resources.release import release_coverage, release_sites logging.basicConfig( format="%(asctime)s (%(name)s %(lineno)s): %(message)s", @@ -23,6 +22,7 @@ logger = logging.getLogger("constraint_pipeline") logger.setLevel(logging.INFO) +EXTENSIONS = ["ht", "tsv", "tsv.bgz", "he", "log"] VERSIONS = ["2.1.1", "4.0", "4.1"] CURRENT_VERSION = "4.1" DATA_TYPES = ["context", "exomes", "genomes"] @@ -77,6 +77,56 @@ Fields to sum (or array sum) when aggregating the expected counts Table. """ +MUTATION_TYPE_FIELDS = ( + "cpg", + "transition", + "mutation_type", + "mutation_type_model", +) +""" +Fields added by `annotate_mutation_type`. +""" + + +def check_param_scope( + version: Optional[str] = None, + model_type: Optional[str] = None, + custom_vep_annotation: Optional[str] = None, + extension: Optional[str] = None, + data_type: Optional[str] = None, +) -> Union[str, None]: + """ + Check if the specified version, genomic region, and other parameters are in the scope of the constraint pipeline. + + If version is specified, return the genome build of the version as a string. + + :param version: One of the release versions (`VERSIONS`). Default is None. + :param model_type: One of "plateau", "coverage". Default is None. + :param custom_vep_annotation: The VEP annotation used to customize the constraint + model (one of "transcript_consequences" or "worst_csq_by_gene"). Default is None. + :param extension: File extension. Default is None. + :param data_type: One of "exomes", "genomes". Default is None. + :return: Genome build of version as a string or None. + """ + if data_type and data_type not in DATA_TYPES: + raise ValueError(f"data_type must be one of: {DATA_TYPES}!") + if model_type and model_type not in MODEL_TYPES: + raise ValueError(f"model_type must be one of: {MODEL_TYPES}!") + if custom_vep_annotation and custom_vep_annotation not in CUSTOM_VEP_ANNOTATIONS: + raise ValueError( + f"custom_vep_annotation must be one of: {CUSTOM_VEP_ANNOTATIONS}!" + ) + if extension and extension not in EXTENSIONS: + raise ValueError(f"extension must be one of: {EXTENSIONS}!") + if version: + if version not in VERSIONS: + raise ValueError("The requested version doesn't exist!") + else: + if version.startswith("2"): + return "GRCh37" + else: + return "GRCh38" + def get_vep_context_ht(version: str) -> TableResource: """ @@ -95,31 +145,6 @@ def get_vep_context_ht(version: str) -> TableResource: raise ValueError("Not a valid gnomAD version -- must be either 2.1.1 or 4.x!") -def get_constraint_root( - version: str = CURRENT_VERSION, - test: bool = False, - post_fix=None, - temp: bool = False, -) -> str: - """ - Return path to constraint root folder. - - :param version: Version of constraint path to return. - :param test: Whether to use a tmp path. - :return: Root path to constraint resources. - """ - post_fix = post_fix or "" - if post_fix: - post_fix = f"_{post_fix}" - - if test: - return f"gs://gnomad-tmp/gnomad_v{version}_testing/constraint{post_fix}" - if temp: - return f"gs://gnomad-tmp/gnomad_v{version}/constraint{post_fix}" - - return f"gs://gnomad/v{version}/constraint{post_fix}" - - def get_sites_resource(data_type: str, version: str = CURRENT_VERSION) -> BaseResource: """ Return genomes or exomes sites Table. @@ -169,7 +194,7 @@ def get_methylation_ht(build: str) -> TableResource: methylation_chrx = ref_grch38.methylation_sites_chrx.ht() methylation_autosomes = ref_grch38.methylation_sites.ht() methylation_ht = methylation_autosomes.union(methylation_chrx) - tmp_path = get_checkpoint_path(f"methylation_{build}") + tmp_path = get_checkpoint_path(f"methylation_{build}").path methylation_ht.checkpoint(tmp_path, _read_if_exists=True) return TableResource(path=tmp_path) else: @@ -195,320 +220,263 @@ def get_coverage_ht( return gnomad_grch38.coverage(data_type) -def get_mutation_ht( +def get_gencode_ht(version: str) -> hl.Table: + """ + Retrieve GENCODE Table. + + :param version: gnomAD version. If version 2, GENCODE v19 will be loaded. If version 4, GENCODE v39 will be loaded. + :return: Table of GENCODE data for the specified build. + """ + if int(version[0]) == 2: + return ref_grch37.gencode.ht() + elif int(version[0]) == 4: + return ref_grch38.gencode.ht(read_args={"_n_partitions": 500}) + else: + raise ValueError("Version must be within gnomAD v2 or v4.") + + +def get_constraint_root( version: str = CURRENT_VERSION, test: bool = False, post_fix: Optional[str] = None, -) -> TableResource: + temp: bool = False, + sub_dir: Optional[str] = None, +) -> str: """ - Return mutation Table that includes the baseline mutation rate for each substitution and context. + Return path to constraint root folder. - :param version: The version of the Table. Default is CURRENT_VERSION. - :param test: Whether the Table is for testing purposes and only contains sites in - chr20, chrX, and chrY. Default is False. - released mutation rate table. - :return: Mutation rate Table. + :param version: Version of constraint path to return. Default is CURRENT_VERSION. + :param test: Whether to use a tmp path. Default is False. + :param post_fix: Postfix to append to the path. Default is None. + :param temp: Whether to use a temp path. Default is False. + :param sub_dir: Subdirectory to append to the path. Default is None. + :return: Root path to constraint resources folder. """ - check_param_scope(version) - return TableResource( - f"{get_constraint_root(version, test, post_fix)}/mutation_rate/gnomad.v{version}.mutation_rate.ht" - ) + post_fix = post_fix or "" + if post_fix: + post_fix = f"_{post_fix}" + + sub_dir = sub_dir or "" + if sub_dir: + sub_dir = f"/{sub_dir}" + + constraint_dir = f"constraint{post_fix}{sub_dir}" + + if test: + return f"gs://gnomad-tmp/gnomad_v{version}_testing/{constraint_dir}" + if temp: + return f"gs://gnomad-tmp/gnomad_v{version}/{constraint_dir}" + return f"gs://gnomad/v{version}/{constraint_dir}" -def get_annotated_context_ht( + +def get_constraint_data( + name: str, version: str = CURRENT_VERSION, test: bool = False, post_fix: Optional[str] = None, -) -> TableResource: + sub_dir: Optional[str] = None, + custom_vep_annotation: Optional[str] = None, + extension: str = "ht", +) -> Union[TableResource, str, ExpressionResource]: """ - Return TableResource of annotated context Table. + Return path, TableResource, or ExpressionResource of requested constraint data. + :param name: Name of the constraint data to retrieve. :param version: One of the release versions (`VERSIONS`). Default is `CURRENT_VERSION`. - :param test: Whether the Table is for testing purposes and only contains sites in - chr20, chrX, and chrY. Default is False. - :return: TableResource of annotated context Table. + :param test: Whether the Table is for testing purpose and only contains a subset of + the data. Default is False. + :param post_fix: Postfix to append to the root path. Default is None. + :param sub_dir: Subdirectory to append to the path. Default is None. + :param custom_vep_annotation: The VEP annotation used to customize the constraint + model (one of "transcript_consequences" or "worst_csq_by_gene"). Default is + None. + :param extension: File extension. Default is "ht". + :return: Path, TableResource, or ExpressionResource of the constraint data. """ - check_param_scope(version) - return TableResource( - f"{get_constraint_root(version, test, post_fix)}/preprocessed_data/annotated_context.ht" + check_param_scope( + version, custom_vep_annotation=custom_vep_annotation, extension=extension ) + if custom_vep_annotation: + sub_dir = f"{sub_dir}/" if sub_dir else "" + sub_dir = f"{sub_dir}{custom_vep_annotation}" -def get_preprocessed_ht( - version: str = CURRENT_VERSION, - test: bool = False, - post_fix: Optional[str] = None, -) -> TableResource: + root_dir = get_constraint_root( + version=version, + test=test, + post_fix=post_fix, + sub_dir=sub_dir, + ) + path = f"{root_dir}/gnomad.v{version}.{name}.{extension}" + + if extension == "ht": + return TableResource(path) + if extension in {"tsv", "tsv.bgz", "log"}: + return path + if extension == "he": + return ExpressionResource(path) + + +def get_mutation_ht(**kwargs) -> TableResource: """ - Return TableResource of preprocessed genome, exomes, and context Table. + Return mutation Table that includes the baseline mutation rate for each substitution and context. - The exome and genome Table will have annotations added by - `prepare_ht_for_constraint_calculations()` and VEP annotation from context Table. + :return: Mutation rate Table. + """ + return get_constraint_data("mutation_rate", sub_dir="mutation_rate", **kwargs) - The context Table will have annotations added by - `prepare_ht_for_constraint_calculations()`. - :param version: One of the release versions (`VERSIONS`). Default is - `CURRENT_VERSION`. - :param test: Whether the Table is for testing purposes and only contains sites in - chr20, chrX, and chrY. Default is False. +def get_annotated_context_ht(**kwargs) -> TableResource: + """ + Return TableResource of annotated context Table. + + :return: TableResource of annotated context Table. + """ + return get_constraint_data( + "annotated_context", sub_dir="preprocessed_data", **kwargs + ) + + +def get_preprocessed_ht(**kwargs) -> TableResource: + """ + Return TableResource of preprocessed genome, exomes, and context Table. + :return: TableResource of processed context Table. """ - check_param_scope(version) - return TableResource( - f"{get_constraint_root(version, test, post_fix, temp=True)}/preprocessed_data/gnomad.v{version}.context.preprocessed.ht" + return get_constraint_data( + "context.preprocessed", sub_dir="preprocessed_data", **kwargs ) -def get_training_dataset( - version: str = CURRENT_VERSION, - test: bool = False, - post_fix: Optional[str] = None, -) -> TableResource: +def get_training_dataset(**kwargs) -> TableResource: """ Return TableResource of training dataset with observed and possible variant count. - :param version: One of the release versions (`VERSIONS`). Default is - `CURRENT_VERSION`. - :param test: Whether the Table is for testing purpose and only contains sites in - chr20, chrX, and chrY. Default is False. :return: TableResource of training dataset. """ - check_param_scope(version) - return TableResource( - f"{get_constraint_root(version, test, post_fix)}/training_data/gnomad.v{version}.constraint_training.ht" - ) + return get_constraint_data("constraint_training", sub_dir="training_data", **kwargs) -def get_training_tsv_path( - version: str = CURRENT_VERSION, - test: bool = False, - post_fix: Optional[str] = None, -) -> str: +def get_training_tsv_path(**kwargs) -> str: """ Return tsv of training dataset with observed and possible variant count. - :param version: One of the release versions (`VERSIONS`). Default is - `CURRENT_VERSION`. - :param test: Whether the Table is for testing purpose and only contains sites in - chr20, chrX, and chrY. Default is False. :return: TSV path of training dataset. """ - check_param_scope(version) - - return f"{get_constraint_root(version, test, post_fix)}/training_data/gnomad.v{version}.constraint_training.tsv.bgz" + return get_constraint_data( + "constraint_training", sub_dir="training_data", extension="tsv.bgz", **kwargs + ) -def get_models( - model_type: str, - version: str = CURRENT_VERSION, - test: bool = False, - post_fix: Optional[str] = None, -) -> ExpressionResource: +def get_models(model_type: str, **kwargs) -> ExpressionResource: """ Return path to a HailExpression that contains desired model type. :param model_type: The type of model. One of "plateau", "coverage". Default is None. - :param version: One of the release versions (`VERSIONS`). Default is - `CURRENT_VERSION`. - :param test: Whether the Table is for testing purpose and only contains sites in - chr20, chrX, and chrY. Default is False. :return: Path to the specified model. """ - check_param_scope(version=version, model_type=model_type) - return ExpressionResource( - f"{get_constraint_root(version, test, post_fix)}/models/gnomad.v{version}.{model_type}.he" - ) + check_param_scope(model_type=model_type) + return get_constraint_data(model_type, sub_dir="models", extension="he", **kwargs) def get_per_variant_expected_dataset( - custom_vep_annotation: str = "transcript_consequences", - version: str = CURRENT_VERSION, - test: bool = False, - post_fix: Optional[str] = None, + custom_vep_annotation: str = "transcript_consequences", **kwargs ) -> TableResource: """ Return TableResource containing the expected variant counts and observed:expected ratio. :param custom_vep_annotation: The VEP annotation used to customize the constraint model (one of "transcript_consequences" or "worst_csq_by_gene"). - :param version: One of the release versions (`VERSIONS`). Default is - `CURRENT_VERSION`. - :param test: Whether the Table is for testing purpose and only contains sites in - chr20, chrX, and chrY. Default is False. :return: Path of the model. """ - check_param_scope( - version=version, + return get_constraint_data( + "per_variant_expected", + sub_dir="apply_models", custom_vep_annotation=custom_vep_annotation, - ) - return TableResource( - f"{get_constraint_root(version, test, post_fix)}/apply_models/{custom_vep_annotation}/gnomad.v{version}.per_variant_expected.ht" + **kwargs, ) -def get_apply_models( - custom_vep_annotation: str = "transcript_consequences", - version: str = CURRENT_VERSION, - test: bool = False, - post_fix: Optional[str] = None, +def get_aggregated_per_variant_expected( + custom_vep_annotation: str = "transcript_consequences", **kwargs ) -> TableResource: """ Return TableResource containing the expected variant counts and observed:expected ratio. :param custom_vep_annotation: The VEP annotation used to customize the constraint model (one of "transcript_consequences" or "worst_csq_by_gene"). - :param version: One of the release versions (`VERSIONS`). Default is - `CURRENT_VERSION`. - :param test: Whether the Table is for testing purpose and only contains sites in - chr20, chrX, and chrY. Default is False. :return: Path of the model. """ - check_param_scope( - version=version, + return get_constraint_data( + "per_variant_expected.aggregated", + sub_dir="apply_models", custom_vep_annotation=custom_vep_annotation, - ) - # TODO: Change this name. - return TableResource( - f"{get_constraint_root(version, test, post_fix)}/apply_models/{custom_vep_annotation}/gnomad.v{version}.apply.per_variant_expected.ht" + **kwargs, ) def get_constraint_metrics_dataset( - custom_vep_annotation: str = "transcript_consequences", - version: str = CURRENT_VERSION, - test: bool = False, - post_fix: Optional[str] = None, + custom_vep_annotation: str = "transcript_consequences", **kwargs ) -> TableResource: """ Return TableResource of pLI scores, observed:expected ratio, 90% confidence interval around the observed:expected ratio, and z scores. - :param version: One of the release versions (`VERSIONS`). Default is - `CURRENT_VERSION`. - :param test: Whether the Table is for testing purposes and only contains sites in - chr20, chrX, and chrY. Default is False. + :param custom_vep_annotation: The VEP annotation used to customize the constraint + model (one of "transcript_consequences" or "worst_csq_by_gene"). :return: TableResource of constraint metrics. """ - check_param_scope(version=version) - - return TableResource( - f"{get_constraint_root(version, test, post_fix)}/metrics/{custom_vep_annotation}/gnomad.v{version}.constraint_metrics.ht" + return get_constraint_data( + "constraint_metrics", + sub_dir="metrics", + custom_vep_annotation=custom_vep_annotation, + **kwargs, ) -def get_constraint_tsv_path( - version: str = CURRENT_VERSION, - test: bool = False, - post_fix: Optional[str] = None, -) -> str: +def get_constraint_tsv_path(**kwargs) -> str: """ Return tsv path of pLI scores, observed:expected ratio, 90% confidence interval around the observed:expected ratio, and z scores. - :param version: One of the release versions (`VERSIONS`). Default is - `CURRENT_VERSION`. - :param test: Whether the Table is for testing purposes. Default is False. :return: TSV path of constraint metrics. """ - check_param_scope(version=version) - - return f"{get_constraint_root(version, test, post_fix)}/metrics/tsv/gnomad.v{version}.constraint_metrics.tsv" + return get_constraint_data( + "constraint_metrics", sub_dir="metrics/tsv", extension="tsv.bgz", **kwargs + ) -def get_downsampling_constraint_tsv_path( - version: str = CURRENT_VERSION, - test: bool = False, - post_fix: Optional[str] = None, -) -> str: +def get_downsampling_constraint_tsv_path(**kwargs) -> str: """ Return tsv path of downsampling observed and expected counts. - :param version: One of the release versions (`VERSIONS`). Default is - `CURRENT_VERSION`. - :param test: Whether the Table is for testing purposes. Default is False. :return: TSV path of constraint metrics. """ - check_param_scope(version=version) - - return f"{get_constraint_root(version, test, post_fix)}/metrics/tsv/gnomad.v{version}.downsampling_constraint_metrics.tsv.bgz" - - -def check_param_scope( - version: Optional[str] = None, - data_type: Optional[str] = None, - model_type: Optional[str] = None, - custom_vep_annotation: Optional[str] = None, -) -> Union[str, None]: - """ - Check if the specified version, genomic region, and data type are in the scope of the constraint pipeline. - - If version is specified, return the genome build of the version as a string. - - :param version: One of the release versions (`VERSIONS`). Default is None. - :param data_type: One of "exomes", "genomes" or "context". Default is None. - :param model_type: One of "plateau", "coverage". Default is None. - :param custom_vep_annotation: The VEP annotation used to customize the constraint - model (one of "transcript_consequences" or "worst_csq_by_gene"). Default is None. - :return: Genome build of version as a string or None. - """ - if data_type and data_type not in DATA_TYPES: - raise ValueError(f"data_type must be one of: {DATA_TYPES}!") - if model_type and model_type not in MODEL_TYPES: - raise ValueError(f"model_type must be one of: {MODEL_TYPES}!") - if custom_vep_annotation and custom_vep_annotation not in CUSTOM_VEP_ANNOTATIONS: - raise ValueError( - f"custom_vep_annotation must be one of: {CUSTOM_VEP_ANNOTATIONS}!" - ) - if version and version not in VERSIONS: - raise ValueError("The requested version doesn't exist!") - else: - if version.startswith("2"): - return "GRCh37" - else: - return "GRCh38" + return get_constraint_data( + "constraint_metrics.downsampling", + sub_dir="metrics/tsv", + extension="tsv.bgz", + **kwargs, + ) -def get_logging_path( - name: str, version: str = CURRENT_VERSION, post_fix: Optional[str] = None -) -> str: +def get_logging_path(name: str, **kwargs) -> str: """ Create a path for Hail log files. :param name: Name of log file. - :param version: One of the release versions (`VERSIONS`). Default is - `CURRENT_VERSION`. :return: Output log path. """ - return f"{get_constraint_root(version, test=True, post_fix=post_fix)}/logging/{name}.log" - - -def get_checkpoint_path( - name: str, - version: str = CURRENT_VERSION, - mt: bool = False, - post_fix: Optional[str] = None, -) -> str: - """ - Create a checkpoint path for Table or MatrixTable. - - :param str name: Name of intermediate Table/MatrixTable. - :param version: Version of path to return. - :param bool mt: Whether path is for a MatrixTable. Default is False. - :return: Output checkpoint path. - """ - return f'{get_constraint_root(version, test=True, post_fix=post_fix)}/checkpoint_files/{name}.{"mt" if mt else "ht"}' + return get_constraint_data( + name, sub_dir="logging", extension="log", test=True, **kwargs + ) -def get_gencode_ht(version: str) -> hl.Table: +def get_checkpoint_path(name: str, **kwargs) -> TableResource: """ - Retrieve GENCODE Table. + Create a checkpoint TableResource. - :param version: gnomAD version. If version 2, GENCODE v19 will be loaded. If version 4, GENCODE v39 will be loaded. - :return: Table of GENCODE data for the specified build. + :param name: Name of intermediate Table. + :return: Output checkpoint TableResource. """ - if int(version[0]) == 2: - return ref_grch37.gencode.ht() - elif int(version[0]) == 4: - return ref_grch38.gencode.ht(read_args={"_n_partitions": 500}) - else: - raise ValueError("Version must be within gnomAD v2 or v4.") + return get_constraint_data(name, sub_dir="checkpoint_files", test=True, **kwargs) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 59ce41b1..dff7c6b0 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -38,6 +38,7 @@ CALIBRATION_GROUPING, COVERAGE_CUTOFF, MU_GROUPING, + MUTATION_TYPE_FIELDS, ) logging.basicConfig( @@ -679,10 +680,13 @@ def create_training_set( # Selecting the only fields that are needed for the training set and filtering out # the rows that are not needed, then checkpointing the Table. This is added to # help avoid memory and shuffle issues. - ht = ht.annotate(**ht.calibrate_mu) + ht = ht.transmute(**ht.calibrate_mu) ht = ht.filter(hl.is_defined(ht.build_model)) + select_fields = {*MU_GROUPING, *MUTATION_TYPE_FIELDS, *CALIBRATION_GROUPING} ht = ht.select( - *MU_GROUPING, *CALIBRATION_GROUPING, "observed_variants", "possible_variants" + *select_fields, + "observed_variants", + "possible_variants", ) ht = ht.checkpoint(new_temp_file("create_training_set", "ht")) @@ -691,7 +695,9 @@ def create_training_set( ht, ht.possible_variants, ht.observed_variants, - additional_grouping=("methylation_level",) + CALIBRATION_GROUPING, + additional_grouping=("methylation_level",) + + MUTATION_TYPE_FIELDS + + CALIBRATION_GROUPING, partition_hint=partition_hint, ) @@ -779,41 +785,38 @@ def aggregate_per_variant_expected_ht( :param use_mane_select: Whether to include MANE Select as a group. Default is False. :return: Table with the observed and expected counts. """ - # include_canonical_group = False - # include_mane_select_group = False - # if custom_vep_annotation == "worst_csq_by_gene": - # vep_annotation = "worst_csq_by_gene" - # if use_mane_select: - # raise ValueError( - # "'mane_select' cannot be set to True when custom_vep_annotation is set" - # " to 'worst_csq_by_gene'." - # ) - # else: - # vep_annotation = custom_vep_annotation - # include_canonical_group = True - # include_mane_select_group = use_mane_select - - # ht = ht.filter(hl.is_defined(ht.possible_variants)) - # ht = ht.select( - # "genomic_region", - # *MU_GROUPING, - # *additional_grouping, - # *AGGREGATE_SUM_FIELDS, - # "vep", - # ) - - # ht, groupings = annotate_exploded_vep_for_constraint_groupings( - # ht=ht, - # vep_annotation=vep_annotation, - # include_canonical_group=include_canonical_group, - # include_mane_select_group=include_mane_select_group, - # ) - # ht = annotate_with_mu(ht, mutation_ht) - # t = ht.checkpoint(new_temp_file("annotate_exploded_vep", "ht")) - ht = hl.read_table( - "gs://gnomad-tmp-4day/annotate_exploded_vep-JYgMWrtvv6cgCX9BRte85A.ht" + include_canonical_group = False + include_mane_select_group = False + if custom_vep_annotation == "worst_csq_by_gene": + vep_annotation = "worst_csq_by_gene" + if use_mane_select: + raise ValueError( + "'mane_select' cannot be set to True when custom_vep_annotation is set" + " to 'worst_csq_by_gene'." + ) + else: + vep_annotation = custom_vep_annotation + include_canonical_group = True + include_mane_select_group = use_mane_select + + ht = ht.filter(hl.is_defined(ht.possible_variants)) + ht = ht.select( + "genomic_region", + *MU_GROUPING, + *additional_grouping, + *AGGREGATE_SUM_FIELDS, + "vep", ) + ht, groupings = annotate_exploded_vep_for_constraint_groupings( + ht=ht, + vep_annotation=vep_annotation, + include_canonical_group=include_canonical_group, + include_mane_select_group=include_mane_select_group, + ) + ht = annotate_with_mu(ht, mutation_ht) + ht = ht.checkpoint(new_temp_file("annotate_exploded_vep", "ht")) + groupings = [ "annotation", "modifier", @@ -823,7 +826,6 @@ def aggregate_per_variant_expected_ht( "canonical", "mane_select", ] - ht = ht.filter(hl.is_defined(ht.possible_variants)) ht = ht.group_by( "genomic_region", *MU_GROUPING, *groupings, *additional_grouping ).aggregate(**aggregate_expected_variants_expr(ht)) From af4b9ec2b417ec9c1c1770642935c3cec07dd6dd Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 12 Mar 2025 12:07:52 -0600 Subject: [PATCH 04/38] Remove accidental sfs_bin in commit --- gnomad_constraint/pipeline/constraint_pipeline.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 4e29b851..73bc2896 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -450,7 +450,6 @@ def main(args): apply_model_low_cov_cutoff=args.pipeline_low_coverage_filter, apply_model_high_cov_cutoff=args.apply_model_high_cov_definition, skip_coverage_model=skip_coverage_model, - additional_grouping_exprs={"sfs_bin": ht.sfs_bin}, ) ht.write(res.temp_preprocess_data_ht.path, overwrite=overwrite) From dee68d3aa5a88834e533b05b76acad939019cf3e Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 12 Mar 2025 12:22:31 -0600 Subject: [PATCH 05/38] Remove changes that are specific to loeuf_all --- .../pipeline/constraint_pipeline.py | 20 ------------------- 1 file changed, 20 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 73bc2896..6956d867 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -167,17 +167,6 @@ def _build_ht_dict(ht_name: str, keep: List[str] = None): .or_missing() ) - # TODO: Make these resources. - am_ht = hl.read_table( - "gs://gnomad/v4.1/constraint/resources/alpha_missense_filters.ht" - ) - adj_r_ht = hl.read_table( - "gs://gnomad/v4.1/constraint/resources/ncc_adj_r_per_base_WG.ht" - ).key_by("locus") - multisfs_ht = hl.read_table( - "gs://gnomad/v4.1/constraint/resources/julia/constraint/multisfs.dedup.ht" - ) - am_keyed = am_ht[ht.key] ht = ht.annotate( coverage=hl.struct( exomes=ht.coverage.exomes.select("mean", "median_approx"), @@ -188,14 +177,6 @@ def _build_ht_dict(ht_name: str, keep: List[str] = None): genomes=ht.AN.genomes[0], ), genomic_region=genomic_region_expr, - alpha_missense=hl.struct( - pathogenicity=am_keyed.am_pathogenicity, - per_98=hl.or_else(am_keyed.am_per_98, False), - per_99=hl.or_else(am_keyed.am_per_99, False), - over_0_999=hl.or_else(am_keyed.am_0_999, False), - ), - adj_r=adj_r_ht[ht.locus].adj_r, - sfs_bin=multisfs_ht[ht.key].Freq_bin_9, ) return ht @@ -556,7 +537,6 @@ def main(args): ht, res.mutation_ht.ht().select("mu_snp"), custom_vep_annotation=custom_vep_annotation, - additional_grouping=("am_per_98", "am_over_0_999", "am_per_99"), use_mane_select=True, ) ht.write(res.apply_ht.path, overwrite=overwrite) From 0d50a79577ddf21992f73d63ef722b9cf79ac591 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 12 Mar 2025 13:27:13 -0600 Subject: [PATCH 06/38] Drop calibrate_mu fields in create_per_variant_expected_ht --- gnomad_constraint/utils/constraint.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index dff7c6b0..eb1dc03e 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -733,6 +733,7 @@ def create_per_variant_expected_ht( model annotation. Default is True. :return: Per-variant expected Table """ + calibrate_mu_fields = set(ht.calibrate_mu.keys()) ht = ht.annotate(**ht.calibrate_mu) if filter_to_apply_variants: @@ -759,7 +760,7 @@ def create_per_variant_expected_ht( ) ) - return ht + return ht.drop(*calibrate_mu_fields) def aggregate_per_variant_expected_ht( From 608753860eddbbdb0d74f311076c63464d204ef9 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 12 Mar 2025 15:01:12 -0600 Subject: [PATCH 07/38] Split `compute_constraint_metrics` step into two steps --- .../pipeline/constraint_pipeline.py | 59 +++++++++--- gnomad_constraint/resources/resource_utils.py | 16 ++++ gnomad_constraint/utils/constraint.py | 93 +++++++++++++------ 3 files changed, 128 insertions(+), 40 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 6956d867..100ec17a 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -296,6 +296,15 @@ def get_constraint_resources( build_models, ], ) + aggregate_by_constraint_groups = PipelineStepResourceCollection( + "--aggregate-by-constraint-groups", + output_resources={ + f"constraint_group_ht": constraint_res.get_constraint_group_ht( + custom_vep_annotation, **common_params + ) + }, + pipeline_input_steps=[aggregate_per_variant_expected], + ) compute_constraint_metrics = PipelineStepResourceCollection( "--compute-constraint-metrics", output_resources={ @@ -303,9 +312,8 @@ def get_constraint_resources( custom_vep_annotation, **common_params ) }, - pipeline_input_steps=[aggregate_per_variant_expected], + pipeline_input_steps=[aggregate_by_constraint_groups], ) - export_tsv = PipelineStepResourceCollection( "--export-tsv", output_resources={ @@ -330,6 +338,7 @@ def get_constraint_resources( "build_models": build_models, "apply_models_per_variant": apply_models_per_variant, "aggregate_per_variant_expected": aggregate_per_variant_expected, + "aggregate_by_constraint_groups": aggregate_by_constraint_groups, "compute_constraint_metrics": compute_constraint_metrics, "export_tsv": export_tsv, } @@ -547,22 +556,18 @@ def main(args): "consequence annotations, and consequence modifier annotations." ) - if args.compute_constraint_metrics: + if args.aggregate_by_constraint_groups: logger.info( - "Computing constraint metrics, including pLI scores, z scores, oe" - " ratio, and confidence interval around oe ratio..." + "Aggregating observed and expected variant counts by constraint groups..." ) - res = resources.compute_constraint_metrics + res = resources.aggregate_by_constraint_groups res.check_resource_existence() # Use new shuffle method to prevent shuffle errors. hl._set_flags(use_new_shuffle="1") - - # Compute constraint metrics. ht = res.apply_ht.ht() - compute_constraint_metrics( - ht=res.apply_ht.ht(), - gencode_ht=constraint_res.get_gencode_ht(version), + aggregate_by_constraint_groups( + ht, keys=tuple( [ i @@ -571,6 +576,25 @@ def main(args): in ["gene", "transcript", "canonical", "mane_select", "gene_id"] ] ), + ).write(res.constraint_group_ht.path, overwrite=overwrite) + hl._set_flags(use_new_shuffle=None) + logger.info("Done with aggregating by constraint groups.") + + if args.compute_constraint_metrics: + logger.info( + "Computing constraint metrics, including pLI scores, z scores, oe" + " ratio, and confidence interval around oe ratio..." + ) + res = resources.compute_constraint_metrics + res.check_resource_existence() + + # Use new shuffle method to prevent shuffle errors. + hl._set_flags(use_new_shuffle="1") + + # Compute constraint metrics. + compute_constraint_metrics( + ht=res.apply_ht.ht(), + gencode_ht=constraint_res.get_gencode_ht(version), expected_values={ "Null": args.expectation_null, "Rec": args.expectation_rec, @@ -927,6 +951,19 @@ def main(args): ) aggregate_per_variant_expected_args._group_actions.append(cov_model_type) + aggregate_by_constraint_groups_args = parser.add_argument_group( + "Aggregate by constraint groups args", + "Arguments used for aggregating by constraint groups.", + ) + aggregate_by_constraint_groups_args.add_argument( + "--aggregate-by-constraint-groups", + help=( + "Aggregate the observed and expected variant counts by constraint groups" + " to get the constraint metrics." + ), + action="store_true", + ) + compute_constraint_args = parser.add_argument_group( "Computate constraint metrics args", "Arguments used for computing constraint metrics.", diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index fbba4f88..f9c1a1ab 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -417,6 +417,22 @@ def get_aggregated_per_variant_expected( ) +def get_constraint_group_ht(custom_vep_annotation: str, **kwargs) -> TableResource: + """ + Return TableResource of constraint group Table. + + :param custom_vep_annotation: The VEP annotation used to customize the constraint + model (one of "transcript_consequences" or "worst_csq_by_gene"). + :return: TableResource of constraint group Table. + """ + return get_constraint_data( + "constraint_group", + sub_dir="apply_models", + custom_vep_annotation=custom_vep_annotation, + **kwargs, + ) + + def get_constraint_metrics_dataset( custom_vep_annotation: str = "transcript_consequences", **kwargs ) -> TableResource: diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index eb1dc03e..9ccacfe1 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -1093,9 +1093,8 @@ def convert_multi_array_to_array_of_structs( ).drop(*array_fields_to_combine) -def compute_constraint_metrics( +def aggregate_by_constraint_groups( ht: hl.Table, - gencode_ht: hl.Table, keys: Tuple = ("gene", "transcript", "canonical"), classic_lof_annotations: Tuple = ( "stop_gained", @@ -1104,27 +1103,20 @@ def compute_constraint_metrics( ), additional_groupings: Dict[str, Dict[str, hl.expr.BooleanExpression]] = None, additional_grouping_combinations: List[List[str]] = None, - expected_values: Optional[Dict[str, float]] = None, - min_diff_convergence: float = 0.001, - raw_z_outlier_threshold_lower_lof: float = -8.0, - raw_z_outlier_threshold_lower_missense: float = -8.0, - raw_z_outlier_threshold_lower_syn: float = -8.0, - raw_z_outlier_threshold_upper_syn: float = 8.0, - use_mane_select_over_canonical: bool = True, ) -> hl.Table: """ - Compute the pLI scores, observed:expected ratio, 90% confidence interval around the observed:expected ratio, and z scores for synonymous variants, missense variants, and predicted loss-of-function (pLoF) variants. + Aggregate the observed and expected variant info for synonymous variants, missense variants, and predicted loss-of-function (pLoF) variants. .. note:: + The following annotations should be present in `ht`: + - modifier - annotation - observed_variants - mu - possible_variants - expected_variants - - expected_variants_{pop} (if `pops` is specified) - - downsampling_counts_{pop} (if `pops` is specified) :param ht: Input Table with the number of expected variants (output of `get_proportion_observed()`). @@ -1133,19 +1125,12 @@ def compute_constraint_metrics( :param classic_lof_annotations: Classic LoF Annotations used to filter the input Table. Default is {"stop_gained", "splice_donor_variant", "splice_acceptor_variant"}. - :param expected_values: Dictionary containing the expected values for 'Null', - 'Rec', and 'LI' to use as starting values. - :param min_diff_convergence: Minimum iteration change in LI to consider the EM - model convergence criteria as met. Default is 0.001. - :param raw_z_outlier_threshold_lower_lof: Value at which the raw z-score is considered an outlier for lof variants. Values below this threshold will be considered outliers. Default is -8.0. - :param raw_z_outlier_threshold_lower_missense: Value at which the raw z-score is considered an outlier for missense variants. Values below this threshold will be considered outliers. Default is -8.0. - :param raw_z_outlier_threshold_lower_syn: Lower value at which the raw z-score is considered an outlier for synonymous variants. Values below this threshold will be considered outliers. Default is -8.0. - :param raw_z_outlier_threshold_upper_syn: Upper value at which the raw z-score is considered an outlier for synonymous variants. Values above this threshold will be considered outliers. Default is 8.0. - :param use_mane_select_over_canonical: Use MANE Select rather than canonical transcripts for filtering the Table when determining ranks for the lof oe upper confidence interval. - If a gene does not have a MANE Select transcript, the canonical transcript (if available) will be used instead. Default is True. - :param gencode_ht: Table containing GENCODE annotations. - :return: Table with pLI scores, observed:expected ratio, confidence interval of the - observed:expected ratio, and z scores. + :param additional_groupings: Additional groupings to add to the constraint groups. + Default is None. + :param additional_grouping_combinations: Additional grouping combinations to add to + the constraint groups. Default is None. + :return: Table with the aggregated observed and expected variant info for synonymous + variants, missense variants, and pLoF variants. """ # Build constraint groups. constraint_group_filters_expr, meta = build_constraint_consequence_groups( @@ -1162,7 +1147,7 @@ def compute_constraint_metrics( ) # Group by keys and get an aggregate sum of mu_snp, observed_variants, - # possible_variants, predicted_propotion_observed, coverage_correction, and + # possible_variants, predicted_proportion_observed, coverage_correction, and # expected_variants for each constraint group. ht = ht.group_by(*keys).aggregate( constraint_groups=hl.agg.array_agg( @@ -1203,9 +1188,58 @@ def compute_constraint_metrics( ) ) ) - ht = ht.checkpoint( - new_temp_file("constraint_metrics.constraint_group_filters.agg", "ht") - ) + + ht = ht.annotate_globals(constraint_group_meta=meta) + + return ht + + +def compute_constraint_metrics( + ht: hl.Table, + gencode_ht: hl.Table, + expected_values: Optional[Dict[str, float]] = None, + min_diff_convergence: float = 0.001, + raw_z_outlier_threshold_lower_lof: float = -8.0, + raw_z_outlier_threshold_lower_missense: float = -8.0, + raw_z_outlier_threshold_lower_syn: float = -8.0, + raw_z_outlier_threshold_upper_syn: float = 8.0, + use_mane_select_over_canonical: bool = True, +) -> hl.Table: + """ + Compute the pLI scores, observed:expected ratio, 90% confidence interval around the observed:expected ratio, and z scores for synonymous variants, missense variants, and predicted loss-of-function (pLoF) variants. + + .. note:: + The following annotations should be present in `ht`: + - modifier + - annotation + - observed_variants + - mu + - possible_variants + - expected_variants + - expected_variants_{pop} (if `pops` is specified) + - downsampling_counts_{pop} (if `pops` is specified) + + :param ht: Input Table with the number of expected variants (output of + `get_proportion_observed()`). + :param keys: The keys of the output Table, defaults to ('gene', 'transcript', + 'canonical'). + :param classic_lof_annotations: Classic LoF Annotations used to filter the input + Table. Default is {"stop_gained", "splice_donor_variant", + "splice_acceptor_variant"}. + :param expected_values: Dictionary containing the expected values for 'Null', + 'Rec', and 'LI' to use as starting values. + :param min_diff_convergence: Minimum iteration change in LI to consider the EM + model convergence criteria as met. Default is 0.001. + :param raw_z_outlier_threshold_lower_lof: Value at which the raw z-score is considered an outlier for lof variants. Values below this threshold will be considered outliers. Default is -8.0. + :param raw_z_outlier_threshold_lower_missense: Value at which the raw z-score is considered an outlier for missense variants. Values below this threshold will be considered outliers. Default is -8.0. + :param raw_z_outlier_threshold_lower_syn: Lower value at which the raw z-score is considered an outlier for synonymous variants. Values below this threshold will be considered outliers. Default is -8.0. + :param raw_z_outlier_threshold_upper_syn: Upper value at which the raw z-score is considered an outlier for synonymous variants. Values above this threshold will be considered outliers. Default is 8.0. + :param use_mane_select_over_canonical: Use MANE Select rather than canonical transcripts for filtering the Table when determining ranks for the lof oe upper confidence interval. + If a gene does not have a MANE Select transcript, the canonical transcript (if available) will be used instead. Default is True. + :param gencode_ht: Table containing GENCODE annotations. + :return: Table with pLI scores, observed:expected ratio, confidence interval of the + observed:expected ratio, and z scores. + """ def _add_oe_ci_z( oe_info: hl.expr.StructExpression, @@ -1244,6 +1278,7 @@ def _add_oe_ci_z( # Annotate with the observed:expected ratio, 95% confidence interval around the # observed:expected ratio, and z scores for each constraint group. + meta = ht.constraint_group_meta.collect() ht = ht.annotate( constraint_groups=hl.map( lambda x, m: x.annotate( From bd0d7b3348008c4ec6eaa292a0e938c7c0432b88 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 12 Mar 2025 15:34:33 -0600 Subject: [PATCH 08/38] Add `include_mu_annotations_in_grouping` option to `aggregate_per_variant_expected_ht` --- gnomad_constraint/pipeline/constraint_pipeline.py | 1 + gnomad_constraint/utils/constraint.py | 9 ++++++--- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 100ec17a..5d7173c8 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -42,6 +42,7 @@ import gnomad_constraint.resources.resource_utils as constraint_res from gnomad_constraint.utils.constraint import ( + aggregate_by_constraint_groups, aggregate_per_variant_expected_ht, calculate_gerp_cutoffs, calculate_mu_by_downsampling, diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 9ccacfe1..9e09e968 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -769,6 +769,7 @@ def aggregate_per_variant_expected_ht( additional_grouping: Tuple = (), custom_vep_annotation: str = "transcript_consequences", use_mane_select: bool = False, + include_mu_annotations_in_grouping: bool = False, ): """ Aggregate the per-variant expected Table. @@ -784,6 +785,8 @@ def aggregate_per_variant_expected_ht( :param additional_grouping: Additional fields to group by. Default is (). :param custom_vep_annotation: Custom VEP annotation to use. Default is :param use_mane_select: Whether to include MANE Select as a group. Default is False. + :param include_mu_annotations_in_grouping: Whether to include the mutation rate + key annotations in the grouping. Default is False. :return: Table with the observed and expected counts. """ include_canonical_group = False @@ -819,6 +822,7 @@ def aggregate_per_variant_expected_ht( ht = ht.checkpoint(new_temp_file("annotate_exploded_vep", "ht")) groupings = [ + *(MU_GROUPING if include_mu_annotations_in_grouping else []), "annotation", "modifier", "gene", @@ -826,10 +830,9 @@ def aggregate_per_variant_expected_ht( "transcript", "canonical", "mane_select", + *additional_grouping, ] - ht = ht.group_by( - "genomic_region", *MU_GROUPING, *groupings, *additional_grouping - ).aggregate(**aggregate_expected_variants_expr(ht)) + ht = ht.group_by(*groupings).aggregate(**aggregate_expected_variants_expr(ht)) return ht.naive_coalesce(1000) From 4cf97b2b0ddda49190c3140312825f3a9e99c6db Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 12 Mar 2025 19:05:50 -0600 Subject: [PATCH 09/38] Use the correct input Table for compute_constraint_metrics --- gnomad_constraint/pipeline/constraint_pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 5d7173c8..79444058 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -594,7 +594,7 @@ def main(args): # Compute constraint metrics. compute_constraint_metrics( - ht=res.apply_ht.ht(), + ht=res.constraint_group_ht.ht(), gencode_ht=constraint_res.get_gencode_ht(version), expected_values={ "Null": args.expectation_null, From 5bbf86424b91b96b2e2cbbef8a79a5a94d7326a8 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 12 Mar 2025 20:17:34 -0600 Subject: [PATCH 10/38] add `calibrate_mu` annotation --- gnomad_constraint/utils/constraint.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 9e09e968..4aaac9c3 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -803,10 +803,11 @@ def aggregate_per_variant_expected_ht( include_canonical_group = True include_mane_select_group = use_mane_select + ht = ht.annotate(**ht.calibrate_mu) ht = ht.filter(hl.is_defined(ht.possible_variants)) ht = ht.select( "genomic_region", - *MU_GROUPING, + *(MU_GROUPING if include_mu_annotations_in_grouping else []), *additional_grouping, *AGGREGATE_SUM_FIELDS, "vep", From 66cccec1cbac998c2dfbfc95493cdeb4aed88975 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 12 Mar 2025 22:48:42 -0600 Subject: [PATCH 11/38] Add temp methylation --- .../pipeline/constraint_pipeline.py | 1 - gnomad_constraint/resources/resource_utils.py | 26 ++++++++++++++++--- gnomad_constraint/utils/constraint.py | 4 --- 3 files changed, 23 insertions(+), 8 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 79444058..4d78122d 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -542,7 +542,6 @@ def main(args): hl._set_flags(use_new_shuffle="1") ht = res.per_variant_apply_ht.ht() - ht = ht.annotate(**{f"am_{k}": v for k, v in ht.alpha_missense.items()}) ht = aggregate_per_variant_expected_ht( ht, res.mutation_ht.ht().select("mu_snp"), diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index f9c1a1ab..388282f4 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -191,11 +191,31 @@ def get_methylation_ht(build: str) -> TableResource: if build == "GRCh37": return ref_grch37.methylation_sites elif build == "GRCh38": - methylation_chrx = ref_grch38.methylation_sites_chrx.ht() + # methylation_chrx = ref_grch38.methylation_sites_chrx.ht() + methylation_chrx_nonpar = hl.read_table( + "gs://gnomad/v4.1/constraint/resources/methylation_chrX.ht" + ) + methylation_chrx_nonpar = methylation_chrx_nonpar.select( + methylation_level=methylation_chrx_nonpar.methy_level, + ) + methylation_chrx_par = hl.read_table( + "gs://gnomad/v4.1/constraint/resources/methylation_chrX_par.ht" + ) + methylation_chrx_par = methylation_chrx_par.select( + methylation_level=methylation_chrx_par.methy_level, + ) + methylation_chry_nonpar = hl.read_table( + "gs://gnomad/v4.1/constraint/resources/methylation_chrY.ht" + ) + methylation_chry_nonpar = methylation_chry_nonpar.select( + methylation_level=methylation_chry_nonpar.methy_level, + ) methylation_autosomes = ref_grch38.methylation_sites.ht() - methylation_ht = methylation_autosomes.union(methylation_chrx) + methylation_ht = methylation_autosomes.union( + methylation_chrx_par, methylation_chrx_nonpar, methylation_chry_nonpar + ) tmp_path = get_checkpoint_path(f"methylation_{build}").path - methylation_ht.checkpoint(tmp_path, _read_if_exists=True) + methylation_ht.checkpoint(tmp_path, overwrite=True) # _read_if_exists=True) return TableResource(path=tmp_path) else: raise ValueError("Build must be one of 'GRCh37' or 'GRCh38'.") diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 4aaac9c3..f0c539a0 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -651,9 +651,6 @@ def prepare_ht_for_constraint_calculations( **exomes_obs_pos_globals, ) - # TODO: Remove this after we have all methylation. - ht = ht.filter(ht.locus.in_autosome()) - print_global_struct(ht) return ht @@ -819,7 +816,6 @@ def aggregate_per_variant_expected_ht( include_canonical_group=include_canonical_group, include_mane_select_group=include_mane_select_group, ) - ht = annotate_with_mu(ht, mutation_ht) ht = ht.checkpoint(new_temp_file("annotate_exploded_vep", "ht")) groupings = [ From e9dc5e21572406acd8fbab1e771bc5cb15c7d516 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Thu, 13 Mar 2025 10:29:32 -0600 Subject: [PATCH 12/38] Move vep explode to earlier step --- .../pipeline/constraint_pipeline.py | 4 +- gnomad_constraint/utils/constraint.py | 84 ++++++++----------- 2 files changed, 35 insertions(+), 53 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 4d78122d..b0771daf 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -524,6 +524,8 @@ def main(args): res.model_plateau.he(), coverage_model=None if skip_coverage_model else res.model_coverage.he(), log10_coverage=log10_coverage, + custom_vep_annotation=custom_vep_annotation, + use_mane_select=True, ) ht.write(res.per_variant_apply_ht.path, overwrite=overwrite) hl._set_flags(use_new_shuffle=None) @@ -545,8 +547,6 @@ def main(args): ht = aggregate_per_variant_expected_ht( ht, res.mutation_ht.ht().select("mu_snp"), - custom_vep_annotation=custom_vep_annotation, - use_mane_select=True, ) ht.write(res.apply_ht.path, overwrite=overwrite) hl._set_flags(use_new_shuffle=None) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index f0c539a0..d57266b4 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -711,6 +711,8 @@ def create_per_variant_expected_ht( coverage_model: Tuple[float, float], log10_coverage: bool = True, filter_to_apply_variants: bool = True, + custom_vep_annotation: str = "transcript_consequences", + use_mane_select: bool = False, ) -> hl.Table: """ Create the per-variant expected Table. @@ -728,13 +730,31 @@ def create_per_variant_expected_ht( :param log10_coverage: Whether to use log10 coverage. Default is True. :param filter_to_apply_variants: Whether to filter to only the rows with an apply model annotation. Default is True. - :return: Per-variant expected Table + :param custom_vep_annotation: Custom VEP annotation to use. Default is + :param use_mane_select: Whether to include MANE Select as a group. Default is False. + :return: Per-variant expected Table. """ + include_canonical_group = False + include_mane_select_group = False + if custom_vep_annotation == "worst_csq_by_gene": + vep_annotation = "worst_csq_by_gene" + if use_mane_select: + raise ValueError( + "'mane_select' cannot be set to True when custom_vep_annotation is set" + " to 'worst_csq_by_gene'." + ) + else: + vep_annotation = custom_vep_annotation + include_canonical_group = True + include_mane_select_group = use_mane_select + calibrate_mu_fields = set(ht.calibrate_mu.keys()) ht = ht.annotate(**ht.calibrate_mu) if filter_to_apply_variants: - ht = ht.filter(hl.is_defined(ht.apply_model)) + ht = ht.filter( + hl.is_defined(ht.apply_model) & hl.is_defined(ht.possible_variants) + ) ht = annotate_with_mu(ht, mutation_ht) @@ -749,23 +769,28 @@ def create_per_variant_expected_ht( log10_coverage=log10_coverage, ) ) + + ht, groupings = annotate_exploded_vep_for_constraint_groupings( + ht=ht, + vep_annotation=vep_annotation, + include_canonical_group=include_canonical_group, + include_mane_select_group=include_mane_select_group, + ) + ht = ht.annotate_globals( apply_models_globals=ht.apply_models_globals.annotate( plateau_models=plateau_models, coverage_model=coverage_model, log10_coverage=log10_coverage, + groupings=groupings, ) ) - return ht.drop(*calibrate_mu_fields) + return ht.drop(*calibrate_mu_fields).repartition(2000) def aggregate_per_variant_expected_ht( ht, - mutation_ht: hl.Table, - additional_grouping: Tuple = (), - custom_vep_annotation: str = "transcript_consequences", - use_mane_select: bool = False, include_mu_annotations_in_grouping: bool = False, ): """ @@ -778,56 +803,13 @@ def aggregate_per_variant_expected_ht( `additional_grouping` to get the observed and expected counts. :param ht: Table returned by `create_per_variant_expected_ht`. - :param mutation_ht: Mutation rate Table. - :param additional_grouping: Additional fields to group by. Default is (). - :param custom_vep_annotation: Custom VEP annotation to use. Default is - :param use_mane_select: Whether to include MANE Select as a group. Default is False. :param include_mu_annotations_in_grouping: Whether to include the mutation rate key annotations in the grouping. Default is False. :return: Table with the observed and expected counts. """ - include_canonical_group = False - include_mane_select_group = False - if custom_vep_annotation == "worst_csq_by_gene": - vep_annotation = "worst_csq_by_gene" - if use_mane_select: - raise ValueError( - "'mane_select' cannot be set to True when custom_vep_annotation is set" - " to 'worst_csq_by_gene'." - ) - else: - vep_annotation = custom_vep_annotation - include_canonical_group = True - include_mane_select_group = use_mane_select - - ht = ht.annotate(**ht.calibrate_mu) - ht = ht.filter(hl.is_defined(ht.possible_variants)) - ht = ht.select( - "genomic_region", - *(MU_GROUPING if include_mu_annotations_in_grouping else []), - *additional_grouping, - *AGGREGATE_SUM_FIELDS, - "vep", - ) - - ht, groupings = annotate_exploded_vep_for_constraint_groupings( - ht=ht, - vep_annotation=vep_annotation, - include_canonical_group=include_canonical_group, - include_mane_select_group=include_mane_select_group, - ) - ht = ht.checkpoint(new_temp_file("annotate_exploded_vep", "ht")) - groupings = [ *(MU_GROUPING if include_mu_annotations_in_grouping else []), - "annotation", - "modifier", - "gene", - "gene_id", - "transcript", - "canonical", - "mane_select", - *additional_grouping, + ht.apply_models_globals.groupings, ] ht = ht.group_by(*groupings).aggregate(**aggregate_expected_variants_expr(ht)) From f5b26f90f2e320d9ef7abf90a44a1aca2060de06 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Thu, 13 Mar 2025 22:39:52 -0600 Subject: [PATCH 13/38] fix use of groupings in `aggregate_per_variant_expected_ht` and remove repartition --- gnomad_constraint/utils/constraint.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index d57266b4..92fc8d98 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -786,7 +786,10 @@ def create_per_variant_expected_ht( ) ) - return ht.drop(*calibrate_mu_fields).repartition(2000) + tmp_path = new_temp_file(prefix="constraint", extension="ht") + ht.drop(*calibrate_mu_fields).write(tmp_path) + + return hl.read_table(tmp_path, _n_partitions=2000) def aggregate_per_variant_expected_ht( @@ -807,9 +810,14 @@ def aggregate_per_variant_expected_ht( key annotations in the grouping. Default is False. :return: Table with the observed and expected counts. """ + ht = ht.transmute(**ht.calibrate_mu) groupings = [ *(MU_GROUPING if include_mu_annotations_in_grouping else []), - ht.apply_models_globals.groupings, + *[ + g + for g in hl.eval(ht.apply_models_globals.groupings) + if g not in MU_GROUPING + ], ] ht = ht.group_by(*groupings).aggregate(**aggregate_expected_variants_expr(ht)) From af5f1edcdf337f301a33ef9b2d0c861b623fa7a8 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Fri, 14 Mar 2025 09:33:08 -0600 Subject: [PATCH 14/38] collect -> eval --- gnomad_constraint/utils/constraint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 92fc8d98..2437feee 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -1268,7 +1268,7 @@ def _add_oe_ci_z( # Annotate with the observed:expected ratio, 95% confidence interval around the # observed:expected ratio, and z scores for each constraint group. - meta = ht.constraint_group_meta.collect() + meta = hl.eval(ht.constraint_group_meta) ht = ht.annotate( constraint_groups=hl.map( lambda x, m: x.annotate( From 59045d0ceac7031f72a400f35f7e86d94f6b35dd Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Fri, 14 Mar 2025 09:46:34 -0600 Subject: [PATCH 15/38] Fix `aggregate_per_variant_expected_ht` arguments --- gnomad_constraint/pipeline/constraint_pipeline.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index b0771daf..269b39cc 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -544,10 +544,7 @@ def main(args): hl._set_flags(use_new_shuffle="1") ht = res.per_variant_apply_ht.ht() - ht = aggregate_per_variant_expected_ht( - ht, - res.mutation_ht.ht().select("mu_snp"), - ) + ht = aggregate_per_variant_expected_ht(ht) ht.write(res.apply_ht.path, overwrite=overwrite) hl._set_flags(use_new_shuffle=None) From b6cfeb05dbde7f78003fae1b1c08858be77b3d6f Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Fri, 14 Mar 2025 10:14:03 -0600 Subject: [PATCH 16/38] Add & (ht.possible_variants > 0)) filter --- gnomad_constraint/utils/constraint.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 2437feee..6ca9e2c1 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -678,7 +678,11 @@ def create_training_set( # the rows that are not needed, then checkpointing the Table. This is added to # help avoid memory and shuffle issues. ht = ht.transmute(**ht.calibrate_mu) - ht = ht.filter(hl.is_defined(ht.build_model)) + # TODO: From Konrad's script parser.add_argument('--skip_af_filter_upfront', + # help='Skip AF filter up front (to be applied later to ensure that it is not + # affecting population-specific constraint): not generally recommended', + # action='store_true') + ht = ht.filter(hl.is_defined(ht.build_model) & (ht.possible_variants > 0)) select_fields = {*MU_GROUPING, *MUTATION_TYPE_FIELDS, *CALIBRATION_GROUPING} ht = ht.select( *select_fields, @@ -752,9 +756,11 @@ def create_per_variant_expected_ht( ht = ht.annotate(**ht.calibrate_mu) if filter_to_apply_variants: - ht = ht.filter( - hl.is_defined(ht.apply_model) & hl.is_defined(ht.possible_variants) - ) + # TODO: From Konrad's script parser.add_argument('--skip_af_filter_upfront', + # help='Skip AF filter up front (to be applied later to ensure that it is not + # affecting population-specific constraint): not generally recommended', + # action='store_true') + ht = ht.filter(hl.is_defined(ht.apply_model) & (ht.possible_variants > 0)) ht = annotate_with_mu(ht, mutation_ht) From fa2412e03b86bee714d9a6b6ed44e5a0bd9320f5 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Mon, 31 Mar 2025 19:11:16 -0600 Subject: [PATCH 17/38] Full run of v4.1 AN 90% --- .../pipeline/constraint_pipeline.py | 69 +++++++++++++------ gnomad_constraint/resources/resource_utils.py | 20 ++++-- gnomad_constraint/utils/constraint.py | 23 +++++-- 3 files changed, 78 insertions(+), 34 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 269b39cc..0e8a2ac8 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -168,6 +168,19 @@ def _build_ht_dict(ht_name: str, keep: List[str] = None): .or_missing() ) + # Add annotation for SFS bin. + sfs_bin_cutoffs = [0, 1e-6, 2e-6, 4e-6, 2e-5, 5e-5, 5e-4, 5e-3, 0.5] + af_expr = ht.freq.exomes[0].AF + sfs_bin_expr = hl.case().when(hl.is_missing(af_expr), 0) + for i, af in enumerate(sfs_bin_cutoffs): + sfs_bin_expr = sfs_bin_expr.when(af_expr <= af, i) + + sfs_bin_expr = sfs_bin_expr.or_missing() + + adj_r_ht = hl.read_table( + "gs://gnomad/v4.1/constraint/resources/adj_r_per_context_methyl_genome_1kb_autosome.agg.ht" + ) + ht = ht.annotate( coverage=hl.struct( exomes=ht.coverage.exomes.select("mean", "median_approx"), @@ -178,6 +191,8 @@ def _build_ht_dict(ht_name: str, keep: List[str] = None): genomes=ht.AN.genomes[0], ), genomic_region=genomic_region_expr, + adj_r=adj_r_ht[ht.locus].adj_r[ht.context], + sfs_bin=sfs_bin_expr, ) return ht @@ -189,7 +204,8 @@ def get_constraint_resources( overwrite: bool, test: bool, models: List[str] = ["plateau", "coverage"], - post_fix: Optional[str] = None, + directory_post_fix: Optional[str] = None, + path_post_fix: Optional[str] = None, ) -> PipelineResourceCollection: """ Get PipelineResourceCollection for all resources needed in the constraint pipeline. @@ -200,7 +216,8 @@ def get_constraint_resources( :param overwrite: Whether to overwrite existing resources. :param test: Whether to use test resources. :param models: List of models to use. Default is ["plateau", "coverage"]. - :param post_fix: Optional post-fix to append to resource paths. + :param directory_post_fix: Post-fix to add to the directory path of the resources. + :param path_post_fix: Post-fix to add to the path of the resources. :return: PipelineResourceCollection containing resources for all steps of the constraint pipeline. """ @@ -227,7 +244,7 @@ def get_constraint_resources( common_params = { "version": version, "test": test, - "post_fix": post_fix, + "directory_post_fix": directory_post_fix, } prepare_context = PipelineStepResourceCollection( @@ -263,15 +280,22 @@ def get_constraint_resources( create_training_set = PipelineStepResourceCollection( "--create-training-set", output_resources={ - f"train_ht": constraint_res.get_training_dataset(**common_params), - f"train_tsv": constraint_res.get_training_tsv_path(**common_params), + f"train_ht": constraint_res.get_training_dataset( + **common_params, path_post_fix=path_post_fix + ), + f"train_tsv": constraint_res.get_training_tsv_path( + **common_params, path_post_fix=path_post_fix + ), }, pipeline_input_steps=[preprocess_data, calculate_mutation_rate], ) build_models = PipelineStepResourceCollection( "--build-models", output_resources={ - f"model_{m}": constraint_res.get_models(m, **common_params) for m in models + f"model_{m}": constraint_res.get_models( + m, **common_params, path_post_fix=path_post_fix + ) + for m in models }, pipeline_input_steps=[create_training_set], ) @@ -279,7 +303,7 @@ def get_constraint_resources( "--apply-models-per-variant", output_resources={ "per_variant_apply_ht": constraint_res.get_per_variant_expected_dataset( - custom_vep_annotation, **common_params + custom_vep_annotation, **common_params, path_post_fix=path_post_fix ) }, pipeline_input_steps=[preprocess_data, calculate_mutation_rate, build_models], @@ -288,7 +312,7 @@ def get_constraint_resources( "--aggregate-per-variant-expected", output_resources={ f"apply_ht": constraint_res.get_aggregated_per_variant_expected( - custom_vep_annotation, **common_params + custom_vep_annotation, **common_params, path_post_fix=path_post_fix ) }, pipeline_input_steps=[ @@ -301,7 +325,7 @@ def get_constraint_resources( "--aggregate-by-constraint-groups", output_resources={ f"constraint_group_ht": constraint_res.get_constraint_group_ht( - custom_vep_annotation, **common_params + custom_vep_annotation, **common_params, path_post_fix=path_post_fix ) }, pipeline_input_steps=[aggregate_per_variant_expected], @@ -310,7 +334,7 @@ def get_constraint_resources( "--compute-constraint-metrics", output_resources={ "constraint_metrics_ht": constraint_res.get_constraint_metrics_dataset( - custom_vep_annotation, **common_params + custom_vep_annotation, **common_params, path_post_fix=path_post_fix ) }, pipeline_input_steps=[aggregate_by_constraint_groups], @@ -319,7 +343,7 @@ def get_constraint_resources( "--export-tsv", output_resources={ "constraint_metrics_tsv": constraint_res.get_constraint_tsv_path( - **common_params + **common_params, path_post_fix=path_post_fix ), "downsampling_constraint_metrics_tsv": ( constraint_res.get_downsampling_constraint_tsv_path(**common_params) @@ -357,7 +381,8 @@ def main(args): version = args.version test_gene_list = args.test_gene_list test = args.test or test_gene_list - post_fix = args.post_fix + directory_post_fix = args.directory_post_fix + path_post_fix = args.path_post_fix overwrite = args.overwrite custom_vep_annotation = args.custom_vep_annotation skip_coverage_model = args.skip_coverage_model @@ -384,7 +409,8 @@ def main(args): overwrite, test, models, - post_fix, + directory_post_fix, + path_post_fix, ) try: @@ -585,12 +611,10 @@ def main(args): res = resources.compute_constraint_metrics res.check_resource_existence() - # Use new shuffle method to prevent shuffle errors. - hl._set_flags(use_new_shuffle="1") - # Compute constraint metrics. + ht = res.constraint_group_ht.ht(read_args={"_n_partitions": 500}) compute_constraint_metrics( - ht=res.constraint_group_ht.ht(), + ht=ht, gencode_ht=constraint_res.get_gencode_ht(version), expected_values={ "Null": args.expectation_null, @@ -605,7 +629,6 @@ def main(args): # ).select_globals( # "version", "apply_model_params", "constraint_meta", "sd_raw_z" ).write(res.constraint_metrics_ht.path, overwrite=overwrite) - hl._set_flags(use_new_shuffle=None) logger.info("Done with computing constraint metrics.") if args.export_tsv: @@ -657,8 +680,14 @@ def main(args): default=constraint_res.CURRENT_VERSION, ) parser.add_argument( - "--post-fix", - help="Post-fix to append to the output file names.", + "--directory-post-fix", + help="Post-fix to append to the output directory path.", + type=str, + default=None, + ) + parser.add_argument( + "--path-post-fix", + help="Post-fix to append to the output file path.", type=str, default=None, ) diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 388282f4..39c22d69 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -294,10 +294,12 @@ def get_constraint_data( name: str, version: str = CURRENT_VERSION, test: bool = False, - post_fix: Optional[str] = None, + directory_post_fix: Optional[str] = None, sub_dir: Optional[str] = None, custom_vep_annotation: Optional[str] = None, extension: str = "ht", + path_post_fix: Optional[str] = None, + temp: bool = False, ) -> Union[TableResource, str, ExpressionResource]: """ Return path, TableResource, or ExpressionResource of requested constraint data. @@ -307,12 +309,13 @@ def get_constraint_data( `CURRENT_VERSION`. :param test: Whether the Table is for testing purpose and only contains a subset of the data. Default is False. - :param post_fix: Postfix to append to the root path. Default is None. + :param directory_post_fix: Postfix to append to the root path. Default is None. :param sub_dir: Subdirectory to append to the path. Default is None. :param custom_vep_annotation: The VEP annotation used to customize the constraint model (one of "transcript_consequences" or "worst_csq_by_gene"). Default is None. :param extension: File extension. Default is "ht". + :param path_post_fix: Postfix to append to the file name. Default is None. :return: Path, TableResource, or ExpressionResource of the constraint data. """ check_param_scope( @@ -323,13 +326,18 @@ def get_constraint_data( sub_dir = f"{sub_dir}/" if sub_dir else "" sub_dir = f"{sub_dir}{custom_vep_annotation}" + path_post_fix = path_post_fix or "" + if path_post_fix: + path_post_fix = f".{path_post_fix}" + root_dir = get_constraint_root( version=version, test=test, - post_fix=post_fix, + post_fix=directory_post_fix, sub_dir=sub_dir, + temp=temp, ) - path = f"{root_dir}/gnomad.v{version}.{name}.{extension}" + path = f"{root_dir}/gnomad.v{version}.{name}{path_post_fix}.{extension}" if extension == "ht": return TableResource(path) @@ -354,9 +362,7 @@ def get_annotated_context_ht(**kwargs) -> TableResource: :return: TableResource of annotated context Table. """ - return get_constraint_data( - "annotated_context", sub_dir="preprocessed_data", **kwargs - ) + return get_constraint_data("annotated_context", temp=True, **kwargs) def get_preprocessed_ht(**kwargs) -> TableResource: diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 6ca9e2c1..274c84a8 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -381,11 +381,12 @@ def get_exomes_observed_and_possible( # set the observed and possible variant annotations to missing. Otherwise, set the # observed and possible variant annotations based on the frequency array. obs_pos_expr = hl.struct( + exomes_freq=exomes_freq_expr, **hl.or_missing( hl.is_defined(exomes_coverage_expr) & hl.or_else(hl.len(exomes_filter_expr) == 0, True), single_variant_observed_and_possible_expr(exomes_freq_expr, max_af=max_af), - ) + ), ) obs_pos_globals = hl.struct( exomes_freq_meta=exomes_freq_meta, @@ -792,10 +793,11 @@ def create_per_variant_expected_ht( ) ) - tmp_path = new_temp_file(prefix="constraint", extension="ht") - ht.drop(*calibrate_mu_fields).write(tmp_path) + # TODO: Check that this is needed + # tmp_path = new_temp_file(prefix="constraint", extension="ht") + # ht.drop(*calibrate_mu_fields).write(tmp_path) - return hl.read_table(tmp_path, _n_partitions=2000) + return ht.drop(*calibrate_mu_fields) # hl.read_table(tmp_path, _n_partitions=2000) def aggregate_per_variant_expected_ht( @@ -1164,7 +1166,7 @@ def aggregate_by_constraint_groups( # Filter to only rows with at least 1 obs or exp across all keys in annotation_dict. ht = ht.filter( ~ht.no_variants - & hl.any( + | hl.any( ht.constraint_groups.map( lambda x: (hl.or_else(x.expected_variants[0], 0) > 0) ) @@ -1240,6 +1242,7 @@ def compute_constraint_metrics( def _add_oe_ci_z( oe_info: hl.expr.StructExpression, m: Dict[str, str], + add_flags: bool = False, ) -> hl.expr.StructExpression: """ Add oe, oe_ci, and z_raw to the oe_info struct. @@ -1269,16 +1272,22 @@ def _add_oe_ci_z( oe=divide_null(obs, exp), oe_ci=oe_confidence_interval(obs, exp), z_raw=z_raw, - flags=add_filters_expr(filters=flags), + flags=( + add_filters_expr(filters=flags) if add_flags else hl.empty_set(hl.tstr) + ), ) # Annotate with the observed:expected ratio, 95% confidence interval around the # observed:expected ratio, and z scores for each constraint group. meta = hl.eval(ht.constraint_group_meta) + freq_meta_len = len(hl.eval(ht.exomes_freq_meta)) ht = ht.annotate( constraint_groups=hl.map( lambda x, m: x.annotate( - oe_info=x.oe_info.map(lambda oe: _add_oe_ci_z(oe, m)) + oe_info=[ + _add_oe_ci_z(x.oe_info[i], m, add_flags=i == 0) + for i in range(freq_meta_len) + ] ), ht.constraint_groups, meta, From 4c8f1d526a6310890451fa706b9d91488546a2f6 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 14 May 2025 09:13:12 -0600 Subject: [PATCH 18/38] Keep filtered variants in possible count --- gnomad_constraint/resources/resource_utils.py | 2 +- gnomad_constraint/utils/constraint.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 39c22d69..476c03e4 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -215,7 +215,7 @@ def get_methylation_ht(build: str) -> TableResource: methylation_chrx_par, methylation_chrx_nonpar, methylation_chry_nonpar ) tmp_path = get_checkpoint_path(f"methylation_{build}").path - methylation_ht.checkpoint(tmp_path, overwrite=True) # _read_if_exists=True) + methylation_ht.checkpoint(tmp_path, _read_if_exists=True) return TableResource(path=tmp_path) else: raise ValueError("Build must be one of 'GRCh37' or 'GRCh38'.") diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 274c84a8..831cf8b9 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -380,11 +380,12 @@ def get_exomes_observed_and_possible( # If the exome coverage is undefined or the variant does not pass the exome filters, # set the observed and possible variant annotations to missing. Otherwise, set the # observed and possible variant annotations based on the frequency array. + exomes_freq_expr = hl.or_missing(hl.len(exomes_filter_expr) == 0, exomes_freq_expr) obs_pos_expr = hl.struct( exomes_freq=exomes_freq_expr, **hl.or_missing( - hl.is_defined(exomes_coverage_expr) - & hl.or_else(hl.len(exomes_filter_expr) == 0, True), + hl.is_defined(exomes_coverage_expr), + # & hl.or_else(hl.len(exomes_filter_expr) == 0, True), single_variant_observed_and_possible_expr(exomes_freq_expr, max_af=max_af), ), ) From 4fe650bce3ee588f5979c08d972e898788255e7a Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Thu, 25 Sep 2025 10:12:43 -0600 Subject: [PATCH 19/38] changes during run --- .../pipeline/constraint_pipeline.py | 45 +- gnomad_constraint/resources/resource_utils.py | 27 +- gnomad_constraint/utils/constraint.py | 419 ++++++++++++++---- 3 files changed, 368 insertions(+), 123 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 0e8a2ac8..4959aeb0 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -178,7 +178,7 @@ def _build_ht_dict(ht_name: str, keep: List[str] = None): sfs_bin_expr = sfs_bin_expr.or_missing() adj_r_ht = hl.read_table( - "gs://gnomad/v4.1/constraint/resources/adj_r_per_context_methyl_genome_1kb_autosome.agg.ht" + "gs://gnomad/v4.1/constraint/resources/annotations/ht/adj_r_per_context_methyl_genome_1kb_autosome.agg.ht" ) ht = ht.annotate( @@ -567,12 +567,21 @@ def main(args): res.check_resource_existence() # Use new shuffle method to prevent shuffle errors. - hl._set_flags(use_new_shuffle="1") + # hl._set_flags(use_new_shuffle="1") ht = res.per_variant_apply_ht.ht() + # ht = res.per_variant_apply_ht.ht(read_args={"_n_partitions": 8000}) + ht = aggregate_per_variant_expected_ht( + ht, include_mu_annotations_in_grouping=True + ) + ht = ht.checkpoint( + "gs://gnomad/v4.1/constraint_coverage_corrected/apply_models/transcript_consequences/gnomad.v4.1.per_variant_expected.aggregated_with_mu_annotations.coverage_corrected.with_downsamplings.ht", + overwrite=overwrite, + ) + # hl._set_flags(use_new_shuffle=None) + ht = aggregate_per_variant_expected_ht(ht) ht.write(res.apply_ht.path, overwrite=overwrite) - hl._set_flags(use_new_shuffle=None) logger.info( "Done aggregating per-variant expected variant count by transcript, " @@ -588,6 +597,34 @@ def main(args): # Use new shuffle method to prevent shuffle errors. hl._set_flags(use_new_shuffle="1") + ht = hl.read_table( + "gs://gnomad/v4.1/constraint_coverage_corrected/apply_models/transcript_consequences/gnomad.v4.1.per_variant_expected.aggregated_with_mu_annotations.coverage_corrected.with_downsamplings.ht" + ) + aggregate_by_constraint_groups( + ht, + keys=tuple( + [ + i + for i in list(ht.key) + if i + in [ + "gene", + "transcript", + "canonical", + "mane_select", + "gene_id", + "context", + "ref", + "alt", + "methylation_level", + ] + ] + ), + ).write( + "gs://gnomad/v4.1/constraint_coverage_corrected/apply_models/transcript_consequences/gnomad.v4.1.constraint_group_with_mu_annotations.coverage_corrected.with_downsamplings.ht", + overwrite=overwrite, + ) + ht = res.apply_ht.ht() aggregate_by_constraint_groups( ht, @@ -612,7 +649,7 @@ def main(args): res.check_resource_existence() # Compute constraint metrics. - ht = res.constraint_group_ht.ht(read_args={"_n_partitions": 500}) + ht = res.constraint_group_ht.ht(read_args={"_n_partitions": 10000}) compute_constraint_metrics( ht=ht, gencode_ht=constraint_res.get_gencode_ht(version), diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 476c03e4..f835adc1 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -191,32 +191,7 @@ def get_methylation_ht(build: str) -> TableResource: if build == "GRCh37": return ref_grch37.methylation_sites elif build == "GRCh38": - # methylation_chrx = ref_grch38.methylation_sites_chrx.ht() - methylation_chrx_nonpar = hl.read_table( - "gs://gnomad/v4.1/constraint/resources/methylation_chrX.ht" - ) - methylation_chrx_nonpar = methylation_chrx_nonpar.select( - methylation_level=methylation_chrx_nonpar.methy_level, - ) - methylation_chrx_par = hl.read_table( - "gs://gnomad/v4.1/constraint/resources/methylation_chrX_par.ht" - ) - methylation_chrx_par = methylation_chrx_par.select( - methylation_level=methylation_chrx_par.methy_level, - ) - methylation_chry_nonpar = hl.read_table( - "gs://gnomad/v4.1/constraint/resources/methylation_chrY.ht" - ) - methylation_chry_nonpar = methylation_chry_nonpar.select( - methylation_level=methylation_chry_nonpar.methy_level, - ) - methylation_autosomes = ref_grch38.methylation_sites.ht() - methylation_ht = methylation_autosomes.union( - methylation_chrx_par, methylation_chrx_nonpar, methylation_chry_nonpar - ) - tmp_path = get_checkpoint_path(f"methylation_{build}").path - methylation_ht.checkpoint(tmp_path, _read_if_exists=True) - return TableResource(path=tmp_path) + return ref_grch38.methylation_sites else: raise ValueError("Build must be one of 'GRCh37' or 'GRCh38'.") diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 831cf8b9..2dd42b37 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -30,8 +30,9 @@ weighted_agg_sum_expr, ) from gnomad.utils.filtering import add_filters_expr -from gnomad.utils.vep import filter_vep_transcript_csqs_expr +from gnomad.utils.vep import CSQ_CODING, filter_vep_transcript_csqs_expr from hail.utils.misc import divide_null, new_temp_file +from numpy.lib import r_ from gnomad_constraint.resources.resource_utils import ( AGGREGATE_SUM_FIELDS, @@ -364,6 +365,7 @@ def get_exomes_observed_and_possible( # use the pared-down downsamplings list. downsamplings = [m["downsampling"] for m in exomes_freq_meta if "downsampling" in m] downsamplings = DOWNSAMPLINGS["v4"] if gen_ancs is None else downsamplings + downsamplings = sorted(map(int, list(set(downsamplings)))) downsamplings = downsamplings if include_downsamplings else None logger.info("The following downsamplings will be used: %s", downsamplings) @@ -804,6 +806,7 @@ def create_per_variant_expected_ht( def aggregate_per_variant_expected_ht( ht, include_mu_annotations_in_grouping: bool = False, + max_array_size: int = 200, ): """ Aggregate the per-variant expected Table. @@ -817,9 +820,9 @@ def aggregate_per_variant_expected_ht( :param ht: Table returned by `create_per_variant_expected_ht`. :param include_mu_annotations_in_grouping: Whether to include the mutation rate key annotations in the grouping. Default is False. + :param max_array_size: Maximum array size before batching. Default is 500. :return: Table with the observed and expected counts. """ - ht = ht.transmute(**ht.calibrate_mu) groupings = [ *(MU_GROUPING if include_mu_annotations_in_grouping else []), *[ @@ -828,8 +831,68 @@ def aggregate_per_variant_expected_ht( if g not in MU_GROUPING ], ] + aggregate_fields_to_sum = [ + "mu_snp", + "mu", + "observed_variants", + "possible_variants", + "predicted_proportion_observed", + "coverage_correction", + "expected_variants", + ] + + if "calibrate_mu" in ht.row: + ht = ht.annotate(**ht.calibrate_mu) + + ht = ht.filter(hl.set(CSQ_CODING).contains(ht.annotation)) + ht = ht.key_by() + ht = ht.select(*groupings, *aggregate_fields_to_sum) + # ht = ht.naive_coalesce(1000).checkpoint(new_temp_file("pre_aggregation", "ht")) + ht = ht.checkpoint(new_temp_file("pre_aggregation", "ht")) + ht = ht.group_by(*groupings).aggregate(**aggregate_expected_variants_expr(ht)) + """ + # Check array lengths to determine if we need to batch. + arrays = [ + f for f in aggregate_fields_to_sum if isinstance(ht[f], hl.ArrayExpression) + ] + array_length = len(ht.filter(hl.is_defined(ht[arrays[0]]))[arrays[0]].take(1)[0]) + logger.info(f"Array length: {array_length}") + + if not array_length > max_array_size: + # No large arrays, use standard aggregation. + ht = ht.group_by(*groupings).aggregate(**aggregate_expected_variants_expr(ht)) + else: + # Calculate number of batches needed. + num_batches = (array_length + max_array_size - 1) // max_array_size + + batches = [] + for i in range(num_batches): + start_idx = i * max_array_size + end_idx = min((i + 1) * max_array_size, array_length) + logger.info( + f"Processing batch {i+1}/{num_batches} (indices {start_idx}:{end_idx})" + ) + + _ht = ht.annotate( + **{f: ht[f][start_idx:end_idx] for f in arrays} + ).checkpoint(new_temp_file(f"batch_{i}", "ht")) + batches.append( + _ht.group_by(*groupings).aggregate( + **aggregate_expected_variants_expr( + _ht, + fields_to_sum=aggregate_fields_to_sum if i == 0 else arrays) + ).checkpoint(new_temp_file(f"batch_{i}.agg", "ht")) + ) + ht = batches[0] + batches = [_ht[ht.key] for _ht in batches[1:]] + ht = ht.annotate( + **{f: hl.flatten([ht[f]] + [_ht[f] for _ht in batches]) for f in arrays} + ) + + """ + ht = ht.checkpoint(new_temp_file("post_aggregation", "ht")) return ht.naive_coalesce(1000) @@ -892,7 +955,7 @@ def calculate_mu_by_downsampling( # TODO: I think we decided this isn't needed right? We can just use canonical. -def filter_to_mane_select_over_canonical(ht: hl.Table) -> hl.Table: +def mane_select_over_canonical_filter_expr(ht: hl.Table) -> hl.Table: """ Filter to MANE Select over canonical transcripts. @@ -910,47 +973,33 @@ def filter_to_mane_select_over_canonical(ht: hl.Table) -> hl.Table: genes = genes.annotate( only_canonical=~(genes.mane_present) & (genes.canonical_present) ) - ms_ht = ht.annotate( - _only_canonical=genes[ht.gene_id].only_canonical, - _mane_present=genes[ht.gene_id].mane_present, - ) - ms_ht = ms_ht.filter( - (ms_ht.transcript.startswith("ENST")) - & ( - (ms_ht._mane_present & ms_ht.mane_select) - | (ms_ht._only_canonical & ms_ht.canonical) - ) - ) + only_canonical_expr = genes[ht.gene_id].only_canonical + mane_present_expr = genes[ht.gene_id].mane_present - return ms_ht + return (ht.transcript.startswith("ENST")) & ( + (mane_present_expr & ht.mane_select) | (only_canonical_expr & ht.canonical) + ) # TODO: Move to gnomad_methods? def add_oe_upper_rank_and_decile( ht: hl.Table, - len_meta: int, use_mane_select_over_canonical: bool = True, ) -> hl.Table: """ Compute the rank and decile of the oe upper confidence interval. :param ht: Table with the oe upper confidence interval. - :param use_mane_select_over_canonical: Use MANE Select rather than canonical - transcripts for filtering the Table when determining ranks for the lof oe upper - confidence interval. If a gene - does not have a MANE Select transcript, the canonical transcript (if available) - will be used instead. Default is True. + :return: Struct containing the rank and decile of the oe upper confidence interval. """ total_count = ht.count() - if use_mane_select_over_canonical: - ms_ht = filter_to_mane_select_over_canonical(ht) + rank_filter_expr = mane_select_over_canonical_filter_expr(ht) else: - ms_ht = ht.filter((ht.canonical) & (ht.transcript.startswith("ENST"))) - - ms_ht = ms_ht.checkpoint(new_temp_file("constraint_metrics.canonical")) + rank_filter_expr = ht.transcript.startswith("ENST") & ht.canonical + ms_ht = ht.filter(rank_filter_expr) n_transcripts = ms_ht.count() logger.info( "Retaining %d out of %d transcripts to use for rank annotations.", @@ -958,32 +1007,97 @@ def add_oe_upper_rank_and_decile( total_count, ) - ms_ht = ms_ht.annotate(upper_rank=hl.empty_array(hl.tint64)) - for i in range(len_meta): - # Rank in ascending order. - ms_ht = ms_ht.order_by(ms_ht.constraint_groups[i].oe_info[0].oe_ci.upper) - ms_ht = ms_ht.add_index(name="rank") - ms_ht = ms_ht.annotate(upper_rank=ms_ht.upper_rank.append(ms_ht.rank)) + """ + ms_ht = ms_ht.select( + oe_ci_upper=ms_ht.constraint_groups.map( + lambda x: x.oe_info.map( + lambda y: hl.struct( + upper_ci=y.oe_ci.upper, + upper_ci_chisq=y.oe_ci_chisq.upper, + upper_ci_gamma=y.oe_ci_gamma.upper, + ) + ) + ), + ).naive_coalesce(100).checkpoint(new_temp_file("oe_ci_upper.before_rank", "ht")) + + num_constraint_groups = hl.eval(ms_ht.constraint_group_meta.length()) + num_freq_groups = hl.eval(ms_ht.exomes_freq_meta.length()) + for i in range(num_constraint_groups): + for j in range(num_freq_groups): + ms_ht = ms_ht.annotate(upper_ci_rank1=hl.struct()) + for k in ["upper_ci", "upper_ci_chisq", "upper_ci_gamma"]: + # Rank in ascending order. + ms_ht = ms_ht.order_by(ms_ht.oe_ci_upper[i][j][k]) + _rank = hl.struct(**{f'{k}_rank': hl.scan.count()}) + if "upper_ci_rank1" in ms_ht.row: + _rank = ms_ht.upper_ci_rank1.annotate(**_rank) + ms_ht = ms_ht.annotate(upper_ci_rank1=_rank) + + _rank = ms_ht.upper_ci_rank1 + if "upper_ci_rank2" in ms_ht.row: + _rank = ms_ht.upper_ci_rank2.append(_rank) + else: + _rank = [_rank] + ms_ht = ms_ht.annotate(upper_ci_rank2=_rank) + ms_ht = ms_ht.checkpoint(new_temp_file(f"constraint_group{i}.freq_groups{j}", "ht")) + + _rank = ms_ht.upper_ci_rank2 + if "upper_ci_rank3" in ms_ht.row: + _rank = ms_ht.upper_ci_rank3.append(_rank) + else: + _rank = [_rank] + ms_ht = ms_ht.annotate(upper_ci_rank3=_rank) + + ms_ht = ms_ht.checkpoint(new_temp_file("oe_ci_upper.rank", "ht")) + """ - ms_ht = ms_ht.annotate( - upper_bin_sextile=ms_ht.upper_rank.map(lambda x: hl.int(x * 6 / n_transcripts)), - upper_bin_decile=ms_ht.upper_rank.map(lambda x: hl.int(x * 10 / n_transcripts)), + ms_ht = hl.read_table( + "gs://gnomad-tmp-4day/oe_ci_upper.rank-664KaxZ6k3vyUFBpZIJaqh.ht" ) # Map rank and bin annotations back to original Table. - ms_ht = ms_ht.key_by(*list(ht.key)) - ms_keyed = ms_ht[ht.key] - - return ht.annotate( + ht = ht.annotate(upper_ci_rank3=ms_ht.key_by(*list(ht.key))[ht.key].upper_ci_rank3) + ht = ht.annotate( constraint_groups=hl.enumerate(ht.constraint_groups).map( lambda x: x[1].annotate( - **{ - k: ms_keyed[k][x[0]] - for k in ["upper_rank", "upper_bin_sextile", "upper_bin_decile"] - } + oe_info=hl.enumerate(x[1].oe_info).map( + lambda y: hl.bind( + lambda r: y[1].annotate( + **{ + f"oe_ci{k}": y[1][f"oe_ci{k}"].annotate( + **hl.or_missing( + hl.is_defined(r), + hl.struct( + upper_rank=r[f"upper_ci{k}_rank"], + upper_bin_percentile=hl.int( + r[f"upper_ci{k}_rank"] + * 100 + / n_transcripts + ), + upper_bin_decile=hl.int( + r[f"upper_ci{k}_rank"] + * 10 + / n_transcripts + ), + upper_bin_sextile=hl.int( + r[f"upper_ci{k}_rank"] + * 6 + / n_transcripts + ), + ), + ) + ) + for k in ["", "_chisq", "_gamma"] + } + ), + ht.upper_ci_rank3[x[0]][y[0]], + ) + ) ) ) - ) + ).drop("upper_ci_rank3") + + return ht # TODO: Move to gnomad_methods? @@ -1193,6 +1307,75 @@ def aggregate_by_constraint_groups( return ht +def gamma_ci( + obs: hl.expr.Int32Expression, + exp: hl.expr.Float64Expression, + alpha: float = 0.05, +) -> hl.expr.Float64Expression: + """ + Calculate the upper bound of the OE confidence interval using the Gamma distribution. + + This function uses the built-in qgamma function from the custom Hail wheel. + + :param obs: Observed count + :param exp: Expected count + :param alpha: Significance level for the confidence interval. Default is 0.05. + :return: Upper bound of the OE confidence interval + """ + # Calculate shape and scale parameters for Gamma distribution + shape = obs + hl.literal(1.0) + scale = divide_null( + hl.literal(1.0), exp + ) # Use divide_null to handle division by zero + p = hl.literal(1.0 - alpha) + + # Use the built-in qgamma function from the custom Hail wheel + # divide_null will return null if exp is 0, making the result null as well + + return hl.struct( + lower=hl.qgamma(hl.literal(alpha), shape, scale), + upper=hl.qgamma(p, shape, scale), + ) + + +def chisq_ci( + obs: hl.expr.Int32Expression, + exp: hl.expr.Float64Expression, + alpha: float = 0.05, +) -> hl.expr.StructExpression: + """ + Calculate the upper bound of the OE confidence interval using the chi-squared + distribution. + + :param obs: Observed count. + :param exp: Expected count. + :param alpha: Significance level for the confidence interval. Default is 0.05. + :return: Upper bound of the OE confidence interval. + """ + return hl.struct( + lower=hl.qchisqtail(alpha, 2 * obs, lower_tail=True) / (2 * exp), + upper=hl.qchisqtail(1 - alpha, 2 * (obs + 1), lower_tail=True) / (2 * exp), + ) + + +def calculate_oe_confidence_interval( + obs, exp, alpha=0.05, oe_upper_method: str = "gamma" +): + """ + Calculate the upper bound of the OE confidence interval. + + :param oe_expr: Array expression with observed and expected values. + :param alpha: Significance level for the OE confidence interval. Default is 0.05. + :return: Array expression with upper bound of the OE confidence interval. + """ + # Calculate upper bound of oe confidence interval. + if oe_upper_method not in ["gamma", "chisq"]: + raise ValueError(f"Invalid OE upper method: {oe_upper_method}") + + oe_upper_func = gamma_ci if oe_upper_method == "gamma" else chisq_ci + return oe_upper_func(obs, exp, alpha) + + def compute_constraint_metrics( ht: hl.Table, gencode_ht: hl.Table, @@ -1229,22 +1412,28 @@ def compute_constraint_metrics( 'Rec', and 'LI' to use as starting values. :param min_diff_convergence: Minimum iteration change in LI to consider the EM model convergence criteria as met. Default is 0.001. - :param raw_z_outlier_threshold_lower_lof: Value at which the raw z-score is considered an outlier for lof variants. Values below this threshold will be considered outliers. Default is -8.0. - :param raw_z_outlier_threshold_lower_missense: Value at which the raw z-score is considered an outlier for missense variants. Values below this threshold will be considered outliers. Default is -8.0. - :param raw_z_outlier_threshold_lower_syn: Lower value at which the raw z-score is considered an outlier for synonymous variants. Values below this threshold will be considered outliers. Default is -8.0. - :param raw_z_outlier_threshold_upper_syn: Upper value at which the raw z-score is considered an outlier for synonymous variants. Values above this threshold will be considered outliers. Default is 8.0. - :param use_mane_select_over_canonical: Use MANE Select rather than canonical transcripts for filtering the Table when determining ranks for the lof oe upper confidence interval. - If a gene does not have a MANE Select transcript, the canonical transcript (if available) will be used instead. Default is True. + :param raw_z_outlier_threshold_lower_lof: Value at which the raw z-score is + considered an outlier for lof variants. Values below this threshold will be + considered outliers. Default is -8.0. + :param raw_z_outlier_threshold_lower_missense: Value at which the raw z-score is + considered an outlier for missense variants. Values below this threshold will + be considered outliers. Default is -8.0. + :param raw_z_outlier_threshold_lower_syn: Lower value at which the raw z-score is + considered an outlier for synonymous variants. Values below this threshold will + be considered outliers. Default is -8.0. + :param raw_z_outlier_threshold_upper_syn: Upper value at which the raw z-score is + considered an outlier for synonymous variants. Values above this threshold will + be considered outliers. Default is 8.0. + :param use_mane_select_over_canonical: Use MANE Select rather than canonical + transcripts for filtering the Table when determining ranks for the lof oe + upper confidence interval. If a gene does not have a MANE Select transcript, + the canonical transcript (if available) will be used instead. Default is True. :param gencode_ht: Table containing GENCODE annotations. :return: Table with pLI scores, observed:expected ratio, confidence interval of the observed:expected ratio, and z scores. """ - def _add_oe_ci_z( - oe_info: hl.expr.StructExpression, - m: Dict[str, str], - add_flags: bool = False, - ) -> hl.expr.StructExpression: + def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: """ Add oe, oe_ci, and z_raw to the oe_info struct. @@ -1253,49 +1442,66 @@ def _add_oe_ci_z( """ obs = oe_info.observed_variants exp = oe_info.expected_variants - z_raw = calculate_raw_z_score(obs, exp) - z_threshold = dict( - { - "lof": (raw_z_outlier_threshold_lower_lof, None), - "mis": (raw_z_outlier_threshold_lower_missense, None), - "syn": ( - raw_z_outlier_threshold_lower_syn, - raw_z_outlier_threshold_upper_syn, - ), - } - ) - z_threshold = z_threshold.get( - hl.coalesce(m.get("lof"), m.get("csq_set", "None")), - (None, None), - ) - flags = get_constraint_flags(exp, z_raw, z_threshold[0], z_threshold[1]) return oe_info.annotate( oe=divide_null(obs, exp), oe_ci=oe_confidence_interval(obs, exp), - z_raw=z_raw, - flags=( - add_filters_expr(filters=flags) if add_flags else hl.empty_set(hl.tstr) + oe_ci_chisq=calculate_oe_confidence_interval( + obs, exp, oe_upper_method="chisq" ), + oe_ci_gamma=calculate_oe_confidence_interval( + obs, exp, oe_upper_method="gamma" + ), + z_raw=calculate_raw_z_score(obs, exp), ) + """ # Annotate with the observed:expected ratio, 95% confidence interval around the # observed:expected ratio, and z scores for each constraint group. - meta = hl.eval(ht.constraint_group_meta) - freq_meta_len = len(hl.eval(ht.exomes_freq_meta)) ht = ht.annotate( - constraint_groups=hl.map( - lambda x, m: x.annotate( - oe_info=[ - _add_oe_ci_z(x.oe_info[i], m, add_flags=i == 0) - for i in range(freq_meta_len) - ] - ), - ht.constraint_groups, - meta, + constraint_groups=ht.constraint_groups.map( + lambda x: x.annotate( + oe_info=x.oe_info.map(lambda oe_info: _add_oe_ci_z(oe_info)) + ) ) ) - # ht = ht.annotate(constraint_flags=...) - ht = ht.checkpoint(new_temp_file("constraint_metrics.oe.oe_ci.z_raw", "ht")) + + z_threshold = { + "lof": (raw_z_outlier_threshold_lower_lof, None), + "mis": (raw_z_outlier_threshold_lower_missense, None), + "syn": ( + raw_z_outlier_threshold_lower_syn, + raw_z_outlier_threshold_upper_syn, + ), + } + """ + meta = hl.eval(ht.constraint_group_meta) + freq_meta = hl.eval(ht.exomes_freq_meta) + syn_idx = meta.index({"csq_set": "syn"}) + mis_idx = meta.index({"csq_set": "mis"}) + lof_idx = meta.index({"lof": "hc"}) + all_freq_idx = freq_meta.index({"group": "adj"}) + """ht = ht.annotate( + constraint_groups=[ + ht.constraint_groups[i].annotate( + flags=add_filters_expr( + get_constraint_flags( + ht.constraint_groups[i].oe_info[all_freq_idx].expected_variants, + ht.constraint_groups[i].oe_info[all_freq_idx].z_raw, + z_threshold.get( + "lof" if m.get("lof") else m.get("csq_set", "None"), + (None, None), + )[0], + z_threshold.get( + "lof" if m.get("lof") else m.get("csq_set", "None"), + (None, None), + )[1], + flag_postfix="lof" if m.get("lof") else m.get("csq_set", None), + ) + ) + ) + for i, m in enumerate(meta) + ] + ).checkpoint(new_temp_file("constraint_metrics.oe.oe_ci.z_raw", "ht")) # Add z-score 'sd' annotation to globals. ht = ht.annotate_globals( @@ -1304,8 +1510,8 @@ def _add_oe_ci_z( ~ht.no_variants, [ calculate_raw_z_score_sd( - ht.constraint_groups[i].oe_info[0].z_raw, - ht.constraint_groups[i].oe_info[0].flags, + ht.constraint_groups[i].oe_info[all_freq_idx].z_raw, + ht.constraint_groups[i].flags, mirror_neg_raw_z=m.get("csq_set") != "syn", ) for i, m in enumerate(meta) @@ -1315,24 +1521,51 @@ def _add_oe_ci_z( ) # Compute z-score from raw z-score and standard deviations. - # TODO: Need to fix z_score ht = ht.annotate( constraint_groups=hl.map( - lambda x, sd_raw_z: x.annotate(z_score=x.oe_info[0].z_raw / sd_raw_z), + lambda x, sd_raw_z: x.annotate( + z_score=x.oe_info[all_freq_idx].z_raw / sd_raw_z + ), ht.constraint_groups, ht.sd_raw_z, + ), + constraint_flags=( + ht.constraint_groups[syn_idx].flags + | ht.constraint_groups[mis_idx].flags + | ht.constraint_groups[lof_idx].flags ) + ).checkpoint(new_temp_file("constraint_metrics.oe.oe_ci.z_raw.flags", "ht")) + """ + + ht = hl.read_table( + "gs://gnomad-tmp-4day/constraint_metrics.oe.oe_ci.z_raw.flags-D9jfiSqRyWZoPb6VmAzNZw.ht" ) # Add a rank and decile of the upper confidence interval for MANE Select or # canonical ensembl transcripts. - ht = add_oe_upper_rank_and_decile(ht, len(meta), use_mane_select_over_canonical) + ht = add_oe_upper_rank_and_decile(ht, use_mane_select_over_canonical).checkpoint( + new_temp_file("constraint_metrics.oe.oe_ci.z_raw.flags.rank_and_decile", "ht") + ) - # TODO: Add back pLI computation # Compute the observed:expected ratio. if expected_values is None: expected_values = {"Null": 1.0, "Rec": 0.706, "LI": 0.207} + hc_lof_expr = ht.constraint_groups[lof_idx].oe_info[all_freq_idx] + ht = ht.annotate( + **compute_pli( + ht, + obs_expr=hc_lof_expr.observed_variants, + exp_expr=hc_lof_expr.expected_variants, + expected_values=expected_values, + min_diff_convergence=min_diff_convergence, + ) + ).checkpoint( + new_temp_file( + "constraint_metrics.oe.oe_ci.z_raw.flags.rank_and_decile.pli", "ht" + ) + ) + # Add transcript annotations from GENCODE. ht = add_gencode_transcript_annotations(ht, gencode_ht) From ca2b0c6183e61d5b79e1a60639f2e3ad67a5bc69 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Mon, 16 Mar 2026 17:50:29 -0600 Subject: [PATCH 20/38] Add CLAUDE.md documentation and update .gitignore for local config files. Introduce gene quality metrics computation in the constraint pipeline and refactor resource utility constants. --- .gitignore | 4 + CLAUDE.md | 185 ++++ .../pipeline/constraint_pipeline.py | 177 +++- gnomad_constraint/resources/constants.py | 219 +++++ gnomad_constraint/resources/resource_utils.py | 217 +++-- gnomad_constraint/utils/constraint.py | 846 +++++++++++++++--- 6 files changed, 1375 insertions(+), 273 deletions(-) create mode 100644 CLAUDE.md create mode 100644 gnomad_constraint/resources/constants.py diff --git a/.gitignore b/.gitignore index 3e5b712b..dd29295e 100644 --- a/.gitignore +++ b/.gitignore @@ -106,3 +106,7 @@ ENV/ # tokens slack_creds.py + +# Claude local config +.claude/ +CLAUDE.local.md diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 00000000..bad2f4e6 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,185 @@ +# gnomad-constraint Project Reference + +## Project Overview + +Genic constraint analysis pipeline for gnomAD. Computes observed/expected ratios, pLI scores, z-scores, and confidence intervals for LoF, missense, and synonymous variants at the gene/transcript level. Current version: **v4.1** (GRCh38). Historical version 2.1.1 (GRCh37) is also supported. + +The `gnomad_constraint/experimental/proemis3d/` directory contains the ProEmis3D project for regional missense constraint visualization. + +## Code Style + +### Formatting + +Code is formatted with **black** (preview mode, line length 88), **isort** (profile `"black"`), and **autopep8** (aggressive=1, ignoring E201/E202/E203/E731). Linting uses **pylint** and **pydocstyle** (PEP 257 convention, ignoring D100/D104). Config is in `pyproject.toml`. + +```bash +# Manual formatting +black gnomad_constraint/ +isort --profile black gnomad_constraint/ +``` + +### Docstrings + +Use **Sphinx-style** (`:param:`, `:return:`) docstrings following the gnomad_methods convention (see gnomad_mnv CLAUDE.md for full examples). + +### Type Annotations + +- **All functions** must have type annotations on parameters and return values. +- Use `typing.List`, `typing.Optional`, etc. for generic types. +- For Hail expression parameters, use `hl.expr.StructExpression`, `hl.expr.BooleanExpression`, etc. +- For Hail table/matrix types, use `hl.Table`, `hl.MatrixTable`. + +## Project Structure + +| Directory | Purpose | +|-----------|---------| +| `gnomad_constraint/pipeline/constraint_pipeline.py` | Main constraint pipeline (7 steps) | +| `gnomad_constraint/pipeline/constraint_pipeline_complex.py` | Complex region constraint pipeline | +| `gnomad_constraint/utils/constraint.py` | Core utility functions (preprocessing, model building, metrics) | +| `gnomad_constraint/utils/constraint_complex.py` | Complex region constraint utilities | +| `gnomad_constraint/resources/resource_utils.py` | Resource paths, constants, `TableResource` definitions | +| `gnomad_constraint/experimental/proemis3d/` | ProEmis3D regional missense constraint | +| `gnomad_constraint/plots/` | R and Python plotting scripts | + +## Key Constants (`resource_utils.py`) + +```python +VERSIONS = ["2.1.1", "4.0", "4.1"] +CURRENT_VERSION = "4.1" +DATA_TYPES = ["context", "exomes", "genomes"] +MODEL_TYPES = ["plateau", "coverage"] +GENOMIC_REGIONS = ["autosome_par", "chrx_nonpar", "chry_nonpar"] +POPS = ("global", "afr", "amr", "eas", "nfe", "sas") +COVERAGE_CUTOFF = 40 +CUSTOM_VEP_ANNOTATIONS = ["transcript_consequences", "worst_csq_by_gene"] +``` + +Note: `MU_GROUPING = ("context", "ref", "alt", "methylation_level")` is NOT in `resource_utils.py`. If you need it, define it locally. + +## Constraint Pipeline Steps + +The main pipeline (`constraint_pipeline.py`) has these steps (each a CLI flag): + +| Step | CLI Flag | Function | +|------|----------|----------| +| 1 | `--prepare-context-ht` | Annotate VEP context with methylation, coverage, GERP | +| 2 | `--preprocess-data` | Add VEP context annotations to exome/genome tables, prepare for constraint | +| 3 | `--calculate-gerp-cutoffs` | Optional: compute GERP percentile cutoffs | +| 4 | `--calculate-mutation-rate` | Compute baseline mutation rate per substitution/context | +| 5 | `--create-training-set` | Count observed + possible variants at synonymous sites | +| 6 | `--build-models` | Build plateau and coverage regression models | +| 7 | `--apply-models` | Apply models to compute expected variant counts and o/e ratios | +| 8 | `--compute-constraint-metrics` | Compute pLI, z-scores, o/e with CIs | +| 9 | `--export-tsv` | Export constraint metrics to TSV | + +### Key Resource Paths (v4.1) + +``` +gs://gnomad/v4.1/constraint/ # Production root +gs://gnomad-tmp/gnomad_v4.1_testing/constraint/ # Test root + +# Key outputs: +.../preprocessed_data/annotated_context.ht +.../preprocessed_data/gnomad.v4.1.{context|exomes|genomes}.preprocessed.{region}.ht +.../mutation_rate/gnomad.v4.1.mutation_rate.ht +.../training_data/gnomad.v4.1.constraint_training.{region}.ht +.../models/gnomad.v4.1.{plateau|coverage}.{region}.he +.../predicted_proportion_observed/transcript_consequences/gnomad.v4.1.predicted_proportion_observed.{region}.ht +.../metrics/gnomad.v4.1.constraint_metrics.ht +``` + +### constraint_metrics Table + +Keyed by `(gene, transcript, canonical)` (and optionally `mane_select`, `gene_id`). + +Output struct per annotation category (`lof`, `mis`, `syn`): +- `.obs` — observed variant count +- `.exp` — expected variant count +- `.oe` — observed/expected ratio +- `.oe_ci` — 90% CI around o/e +- `.z_raw` — raw z-score +- `.pLI` — probability of loss-of-function intolerance (LoF only) +- `.pNull`, `.pRec` — null/recessive probabilities (LoF only) + +## Missense Score Percentile Analysis + +`gnomad_constraint/plots/determine_missense_score_percentiles.py` computes per-percentile depletion of missense variants binned by missense prediction scores. + +### Scores analyzed + +ProteinMPNN, ESM, REVEL, RASP, AM, MisFit, PolyPhen, CPT1, popEVE, EVE, MPC, CADD, GPN-MSA + +### Pipeline steps (CLI flags) + +| Step | Flag | +|------|------| +| 1 | `--preprocess-scores` | +| 2 | `--compute-percentiles` | +| 3 | `--annotate-constraint-data` | +| 4 | `--aggregate-by-transcript` | +| 5 | `--compute-cumulative` | +| 6 | `--export-percentile-summary` | +| 7 | `--export-matched-plof-summary` | + +Step 7 computes matched pLoF o/e per missense percentile bin. It computes adj_r-corrected gene-level pLoF from the per-SNV table (`--constraint-ht-path`), not from the pre-computed constraint_metrics table (which lacks adj_r for pLoF). + +## Dependencies + +- **hail** — distributed genomics framework +- **numpy**, **pandas**, **scipy** — numerical/statistical +- **gnomad** (gnomad_methods) — shared gnomAD utilities +- **gnomad_qc** — gnomAD QC pipeline resources + +## Known Gotchas + +- **MU_GROUPING not exported**: `gnomad_constraint.resources.resource_utils` does NOT export `MU_GROUPING`. It must be defined locally as `("context", "ref", "alt", "methylation_level")`. +- **v4 drops chrY/chrX**: The v4 pipeline removes `chry_nonpar` and `chrx_nonpar` from regions early in `main()`. Don't assume all 3 genomic regions are present. +- **Coverage metric**: v4 can use `"exomes_AN_percent"` instead of `"exome_coverage"`. This affects model building and application. +- **Genomes v3.1.2 for v4**: Even in v4, the genomes sites resource uses v3.1.2 (downsamplings dropped in v4). +- **Hail `hl.init` tmp_dir**: Pipeline uses `gs://gnomad-tmp-4day` as the Hail temp directory. Ensure this bucket exists and is writable. +- **`ht.get()` doesn't exist on Hail Tables**: Use `field in ht.row` to check field existence. +- **constraint_metrics.ht pLoF lacks adj_r**: The pre-computed `constraint_metrics.ht` has `lof.exp` WITHOUT the regional depletion correction (`adj_r`). To get adj_r-corrected pLoF, compute it from the per-SNV table by filtering to LOFTEE HC + `possible_variants == 1` and aggregating with `expected_variants[0] * adj_r`. +- **Table version awareness**: Undated tables (e.g., `annotate_with_oe.ht`) and dated tables (e.g., `annotate_with_oe.12_23_25.ht`) may have different schemas and column names. Always verify which version you're using. + +## Hail / Dataproc Best Practices + +- **Never use `.count()` for logging on large tables**: `count()` forces a full table materialization. On a large per-SNV table this triggers a massive Spark job and can cause shuffle failures. Use it only when the result is actually needed for computation. +- **Use `naive_coalesce()` after aggressive filters**: When filtering a large table down to a small subset (e.g., LOFTEE HC LoF from all variants), most partitions become empty. This causes shuffle skew in downstream `group_by` aggregations. Call `naive_coalesce(200)` after the filter to rebalance. +- **Per-SNV table `calibrate_mu` struct**: The per-variant expected table (`gnomad.v4.1.per_variant_expected.coverage_corrected.with_downsamplings.ht`) stores transcript-level fields (`gene`, `transcript`, `canonical`, `modifier`, `observed_variants`, `expected_variants`, `possible_variants`) inside a `calibrate_mu` struct. Flatten it with `ht = ht.annotate(**ht.calibrate_mu)` before accessing those fields. +- **Log full row field lists for debugging**: When debugging schema issues, log `list(ht.row)` (all fields), not `list(ht.row)[:20]` (truncated). Important fields like `calibrate_mu` may be beyond the first 20. +- **`order_by` destroys the key — use `add_index` to rekey cheaply**: After `ht.order_by(expr)`, the table is unkeyed. To rejoin ranked results back to the original table, call `ht.add_index("_rank_idx")` before ordering, then `rank_ht.key_by("_rank_idx")` after. Rejoining via an integer index is an O(1) lookup vs a full key scan. +- **`hl.scan.count()` for rank assignment**: After `order_by`, annotate with `hl.scan.count()` to assign 0-based ascending ranks in a single pass: `ht = ht.order_by(ht.val).annotate(rank=hl.scan.count())`. +- **Checkpoint small select-then-order tables, not the full wide table**: When computing ranks for many `(group, field)` combinations, select only the columns needed (`ht.select("_rank_idx", _val=expr)`), order, rank, checkpoint, and join results back in one pass. Avoids checkpointing the full wide table once per iteration. +- **`.count()` after checkpoint is free**: `count()` on a checkpointed table reads already-materialized metadata rather than re-executing the query. Place `count()` after a checkpoint to avoid computing the table twice. +- **Python list comprehensions over Hail arrays for indexed access**: When you need to index a Hail array with a known Python integer (e.g., `ht.constraint_groups[i]`), use a Python list comprehension rather than `hl.enumerate` + lambda. This also allows Python-time dict lookups like `rank_hts[(i, key)]` inside the expression. +- **`hl.Table.parallelize` to reconstruct a small HT from collected data**: `hl.Table.parallelize(hl.eval(ht.my_array_global), schema=ht.my_array_global.dtype.element_type).key_by(...)` reconstructs a small Hail Table from a global array without re-running any jobs. +- **Hail array elements must share the same struct schema**: All elements of a Hail array field must have identical types. You cannot annotate only `array[0]` with extra fields while leaving `array[1+]` unchanged — Hail will reject the mixed schema. Instead, promote such metadata to the parent struct level (e.g., add a `{field}_rank` struct directly on the constraint group rather than inside `oe_info[0]`). + +## Dataproc Submission + +**Important**: hailctl repackages `--pyfiles` into a temp zip using `os.walk`, which nests packages incorrectly. Use the single-zip workaround (same as gnomad_mnv): + +```bash +# Build a single zip with correct top-level package structure +cd && \ + rm -f /tmp/pyfiles.zip && \ + zip -r /tmp/pyfiles.zip gnomad_constraint/ -x '*.pyc' '*__pycache__*' && \ + cd && \ + zip -r /tmp/pyfiles.zip gnomad_qc/ -x '*.pyc' '*__pycache__*' '*.DS_Store' + +# Submit to cluster (single zip = used directly, not repackaged) +hailctl dataproc submit \ + gnomad_constraint/pipeline/constraint_pipeline.py \ + --pyfiles /tmp/pyfiles.zip \ + -- --compute-constraint-metrics --test --overwrite +``` + +## gnomad_methods / gnomad_qc API + +See the gnomad_mnv CLAUDE.md for shared API reference (`public_release`, `TableResource`, `get_gnomad_v4_vds`, etc.). Key constraint-specific imports: + +```python +from gnomad.utils.constraint import build_models, compute_pli, oe_confidence_interval +from gnomad.resources.grch38.gnomad import public_release, DOWNSAMPLINGS, all_sites_an +from gnomad_qc.resource_utils import PipelineResourceCollection, PipelineStepResourceCollection +``` diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 4959aeb0..96a94041 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -28,7 +28,6 @@ import hail as hl from gnomad.resources.grch38.gnomad import all_sites_an from gnomad.utils.constraint import ( - annotate_with_mu, assemble_constraint_context_ht, build_models, explode_downsamplings_oe, @@ -41,15 +40,23 @@ ) import gnomad_constraint.resources.resource_utils as constraint_res +from gnomad_constraint.resources.constants import ( + CURRENT_VERSION, + CUSTOM_VEP_ANNOTATIONS, + VERSIONS, +) from gnomad_constraint.utils.constraint import ( aggregate_by_constraint_groups, aggregate_per_variant_expected_ht, calculate_gerp_cutoffs, calculate_mu_by_downsampling, compute_constraint_metrics, + compute_gene_quality_metrics, create_per_variant_expected_ht, create_training_set, + flatten_release_ht, prepare_ht_for_constraint_calculations, + prepare_release_ht, print_global_struct, ) @@ -330,6 +337,18 @@ def get_constraint_resources( }, pipeline_input_steps=[aggregate_per_variant_expected], ) + compute_gene_quality_metrics_step = PipelineStepResourceCollection( + "--compute-gene-quality-metrics", + output_resources={ + "gene_quality_metrics_ht": constraint_res.get_gene_quality_metrics_ht( + version=version + ) + }, + input_resources={ + "gnomAD resources": {"exomes_sites_ht": input_hts["exomes_sites_ht"]}, + }, + pipeline_input_steps=[prepare_context], + ) compute_constraint_metrics = PipelineStepResourceCollection( "--compute-constraint-metrics", output_resources={ @@ -337,19 +356,29 @@ def get_constraint_resources( custom_vep_annotation, **common_params, path_post_fix=path_post_fix ) }, - pipeline_input_steps=[aggregate_by_constraint_groups], + pipeline_input_steps=[ + aggregate_by_constraint_groups, + compute_gene_quality_metrics_step, + ], ) - export_tsv = PipelineStepResourceCollection( - "--export-tsv", + prepare_release = PipelineStepResourceCollection( + "--prepare-release", output_resources={ - "constraint_metrics_tsv": constraint_res.get_constraint_tsv_path( - **common_params, path_post_fix=path_post_fix + "release_ht": constraint_res.get_release_constraint_ht(version=version), + }, + pipeline_input_steps=[compute_constraint_metrics], + ) + export_release_tsv = PipelineStepResourceCollection( + "--export-release-tsv", + output_resources={ + "release_tsv": constraint_res.get_release_constraint_tsv_path( + version=version ), - "downsampling_constraint_metrics_tsv": ( - constraint_res.get_downsampling_constraint_tsv_path(**common_params) + "release_downsampling_tsv": constraint_res.get_release_downsampling_tsv_path( + version=version ), }, - pipeline_input_steps=[compute_constraint_metrics], + pipeline_input_steps=[prepare_release], ) # Add all steps to the constraint pipeline resource collection. @@ -364,8 +393,10 @@ def get_constraint_resources( "apply_models_per_variant": apply_models_per_variant, "aggregate_per_variant_expected": aggregate_per_variant_expected, "aggregate_by_constraint_groups": aggregate_by_constraint_groups, + "compute_gene_quality_metrics": compute_gene_quality_metrics_step, "compute_constraint_metrics": compute_constraint_metrics, - "export_tsv": export_tsv, + "prepare_release": prepare_release, + "export_release_tsv": export_release_tsv, } ) @@ -388,7 +419,7 @@ def main(args): skip_coverage_model = args.skip_coverage_model log10_coverage = args.use_logarithmic_coverage_model - if version not in constraint_res.VERSIONS: + if version not in VERSIONS: raise ValueError("The requested version of resource Tables is not available.") if version == "2.1.1": @@ -640,6 +671,20 @@ def main(args): hl._set_flags(use_new_shuffle=None) logger.info("Done with aggregating by constraint groups.") + if args.compute_gene_quality_metrics: + logger.info("Computing per-transcript gene quality metrics...") + res = resources.compute_gene_quality_metrics + res.check_resource_existence() + + gencode_cds_ht = constraint_res.get_gencode_cds_ht(version).ht() + gene_quality_ht = compute_gene_quality_metrics( + res.annotated_context_ht.ht(), + res.exomes_sites_ht.ht(), + gencode_cds_ht, + ) + gene_quality_ht.write(res.gene_quality_metrics_ht.path, overwrite=overwrite) + logger.info("Done computing gene quality metrics.") + if args.compute_constraint_metrics: logger.info( "Computing constraint metrics, including pLI scores, z scores, oe" @@ -653,6 +698,7 @@ def main(args): compute_constraint_metrics( ht=ht, gencode_ht=constraint_res.get_gencode_ht(version), + gene_quality_metrics_ht=res.gene_quality_metrics_ht.ht(), expected_values={ "Null": args.expectation_null, "Rec": args.expectation_rec, @@ -663,36 +709,42 @@ def main(args): raw_z_outlier_threshold_lower_missense=args.raw_z_outlier_threshold_lower_missense, raw_z_outlier_threshold_lower_syn=args.raw_z_outlier_threshold_lower_syn, raw_z_outlier_threshold_upper_syn=args.raw_z_outlier_threshold_upper_syn, - # ).select_globals( - # "version", "apply_model_params", "constraint_meta", "sd_raw_z" ).write(res.constraint_metrics_ht.path, overwrite=overwrite) logger.info("Done with computing constraint metrics.") - if args.export_tsv: - res = resources.export_tsv + if args.prepare_release: + logger.info("Preparing constraint metrics Table for release...") + res = resources.prepare_release res.check_resource_existence() - logger.info("Exporting constraint tsv...") - ht = res.constraint_metrics_ht.ht() - # If downsamplings per genetic ancestry group are present, export - # downsamplings to a separate tsv and drop from the main metrics tsv. - if args.genetic_ancestry_groups: - downsampling_ht = explode_downsamplings_oe( - ht, - downsampling_meta=hl.eval(ht.apply_model_params.downsampling_meta), - ) + constraint_ht = res.constraint_metrics_ht.ht() + + release_ht = prepare_release_ht( + constraint_ht, + release_version=args.release_version, + ) + release_ht.write(res.release_ht.path, overwrite=overwrite) + logger.info("Done preparing release Table.") + + if args.export_release_tsv or args.export_release_downsampling_tsv: + res = resources.export_release_tsv + res.check_resource_existence() + release_ht = hl.read_table(res.release_ht.path) - # Drop downsampling annotations from the main metrics Table. - ht = ht.annotate( - **{ - i: ht[i].drop(*["gen_anc_exp", "gen_anc_obs"]) - for i in ["lof_hc_lc", "lof", "syn", "mis"] - } + if args.export_release_tsv: + logger.info("Exporting release TSV...") + flatten_release_ht(release_ht).export(res.release_tsv) + logger.info("Done exporting release TSV.") + + if args.export_release_downsampling_tsv: + logger.info("Exporting release downsampling TSV...") + downsampling_ht = explode_downsamplings_oe( + release_ht, + downsampling_meta=hl.eval(release_ht.downsamplings), + metrics=["syn", "mis", "lof_hc_lc", "lof"], ) - # Export separate downsampling Table. - downsampling_ht.export(res.downsampling_constraint_metrics_tsv) - ht = ht.flatten() - ht.export(res.constraint_metrics_tsv) + downsampling_ht.export(res.release_downsampling_tsv) + logger.info("Done exporting release downsampling TSV.") finally: logger.info("Copying log to logging bucket...") @@ -711,10 +763,10 @@ def main(args): "--version", help=( "Which version of the resource Tables will be used. Default is" - f" {constraint_res.CURRENT_VERSION}." + f" {CURRENT_VERSION}." ), type=str, - default=constraint_res.CURRENT_VERSION, + default=CURRENT_VERSION, ) parser.add_argument( "--directory-post-fix", @@ -1010,7 +1062,7 @@ def main(args): ), type=str, default="transcript_consequences", - choices=constraint_res.CUSTOM_VEP_ANNOTATIONS, + choices=CUSTOM_VEP_ANNOTATIONS, ) aggregate_per_variant_expected_args._group_actions.append(cov_model_type) @@ -1027,6 +1079,19 @@ def main(args): action="store_true", ) + gene_quality_args = parser.add_argument_group( + "Compute gene quality metrics args", + "Arguments used for computing per-transcript gene quality metrics.", + ) + gene_quality_args.add_argument( + "--compute-gene-quality-metrics", + help=( + "Compute per-transcript gene quality metrics (coverage, mapping quality," + " segdup, LCR) from the preprocessed context Table and gnomAD exomes" + " sites Table." + ), + action="store_true", + ) compute_constraint_args = parser.add_argument_group( "Computate constraint metrics args", "Arguments used for computing constraint metrics.", @@ -1127,9 +1192,41 @@ def main(args): type=float, default=8.0, ) - compute_constraint_args.add_argument( - "--export-tsv", - help="Export constraint metrics to tsv file.", + prepare_release_args = parser.add_argument_group( + "Prepare release args", + "Arguments used for preparing the constraint metrics Table for release.", + ) + prepare_release_args.add_argument( + "--prepare-release", + help=( + "Prepare the constraint metrics Table for public release by restructuring " + "constraint groups into named top-level fields and consolidating globals." + ), + action="store_true", + ) + prepare_release_args.add_argument( + "--release-version", + help=( + "Version string to set in the release Table globals. If not specified, " + "the existing version global is retained." + ), + type=str, + default=None, + ) + prepare_release_args.add_argument( + "--export-release-tsv", + help=( + "Flatten the release Hail Table and export it as a TSV. Output paths are" + " determined by the release resource functions." + ), + action="store_true", + ) + prepare_release_args.add_argument( + "--export-release-downsampling-tsv", + help=( + "Export per-genetic-ancestry downsampling observed and expected counts from" + " the release Hail Table as a TSV. Reads from the release HT path." + ), action="store_true", ) diff --git a/gnomad_constraint/resources/constants.py b/gnomad_constraint/resources/constants.py new file mode 100644 index 00000000..8a805b60 --- /dev/null +++ b/gnomad_constraint/resources/constants.py @@ -0,0 +1,219 @@ +"""Constants used across the constraint pipeline and release formatting.""" + +# --------------------------------------------------------------------------- +# Pipeline configuration +# --------------------------------------------------------------------------- + +EXTENSIONS = ["ht", "tsv", "tsv.bgz", "he", "log"] +"""Valid file extensions for constraint pipeline resources.""" + +VERSIONS = ["2.1.1", "4.0", "4.1"] +"""Supported gnomAD constraint pipeline versions.""" + +CURRENT_VERSION = "4.1" +"""Current default gnomAD constraint pipeline version.""" + +DATA_TYPES = ["context", "exomes", "genomes"] +"""Data types used in the constraint pipeline.""" + +MODEL_TYPES = ["plateau", "coverage"] +"""Model types used for constraint calibration.""" + +GENOMIC_REGIONS = ["autosome_par", "chrx_nonpar", "chry_nonpar"] +"""Genomic regions used to partition constraint calculations.""" + +CUSTOM_VEP_ANNOTATIONS = ["transcript_consequences", "worst_csq_by_gene"] +""" +VEP annotations used when applying models. + +"transcript_consequences" option will annotate the Table with 'annotation', 'gene', +'coverage', 'transcript', and either 'canonical' or 'mane_select' annotations using 'transcript_consequences' +VEP annotation. + +"worst_csq_by_gene" option will annotate the Table with 'annotation', 'gene', and +'coverage' annotations using 'worst_csq_by_gene' VEP annotation. +""" + +POPS = ("global", "afr", "amr", "eas", "nfe", "sas") +""" +Population labels from gnomAD. + +Abbreviations stand for: global (all populations), African-American/African, Latino, East Asian, Non-Finnish European, and South Asian. +""" + +COVERAGE_CUTOFF = 40 +""" +Minimum median exome coverage differentiating high coverage sites from low coverage sites. + +Low coverage sites require an extra calibration when computing the proportion of expected variation. +""" + +MU_GROUPING = ("context", "ref", "alt", "methylation_level") +""" +Annotations used to group variants for the mutation rate calculation. +""" + +CALIBRATION_GROUPING = ("genomic_region", "build_model", "cpg", "exomes_coverage") +""" +Annotations used to group variants for the mutation rate calibration. +""" + +AGGREGATE_SUM_FIELDS = ( + "mu_snp", + "observed_variants", + "possible_variants", + "predicted_proportion_observed", + "coverage_correction", + "expected_variants", +) +""" +Fields to sum (or array sum) when aggregating the expected counts Table. +""" + +MUTATION_TYPE_FIELDS = ( + "cpg", + "transition", + "mutation_type", + "mutation_type_model", +) +""" +Fields added by `annotate_mutation_type`. +""" + +# --------------------------------------------------------------------------- +# Frequency metadata +# --------------------------------------------------------------------------- + +ADJ_FREQ_META = {"group": "adj"} +"""Frequency metadata key for the adjusted allele frequency group.""" + +# --------------------------------------------------------------------------- +# GENCODE field renames +# --------------------------------------------------------------------------- + +GENCODE_FIELD_RENAMES = { + "transcript_id_version": "transcript_version", + "level": "transcript_level", +} +"""GENCODE field renames applied when preparing the release Table.""" + +# --------------------------------------------------------------------------- +# Release format constants +# --------------------------------------------------------------------------- + +RELEASE_KEY_ORDER = ["gene", "gene_id", "transcript", "canonical", "mane_select"] +"""Key fields for the release Table, in display order.""" + +RELEASE_SCALAR_FIELDS = [ + "transcript_version", + "transcript_type", + "transcript_level", + "chromosome", + "start_position", + "end_position", + "cds_length", + "num_coding_exons", +] +"""Scalar transcript/gene annotation fields included in the flat release Table.""" + +RELEASE_TOP_LEVEL_ANNOTATIONS = RELEASE_SCALAR_FIELDS + [ + "gene_quality_metrics", + "gene_flags", + "constraint_flags", +] +"""Non-key, non-constraint-group row fields included in the release Table.""" + +RELEASE_GROUP_NAMES = ["syn", "mis", "lof_hc_lc", "lof"] +"""Constraint group names exposed in the release Table, in display order.""" + +RELEASE_LOF_FIELDS = ["pLI", "pNull", "pRec"] +"""LoF-specific fields appended to the ``lof`` constraint group in release format.""" + +RELEASE_CI_FIELDS = ["lower", "upper"] +"""OE confidence interval sub-fields included in the release Table.""" + +RELEASE_RANK_FIELDS = [ + "upper_rank", + "upper_bin_percentile", + "upper_bin_decile", + "upper_bin_sextile", +] +"""Rank and bin sub-fields appended to the CI struct for ranked groups.""" + +RELEASE_CI_FIELDS_WITH_RANK = RELEASE_CI_FIELDS + RELEASE_RANK_FIELDS +"""CI fields including rank annotations, used for groups in RELEASE_GROUPS_WITH_RANK.""" + +RELEASE_GROUPS_WITH_RANK = ["lof_hc", "lof_hc_lc"] +"""Constraint groups for which rank and bin annotations are included.""" + +RELEASE_GROUPS_WITH_PLI = ["lof_hc", "lof_hc_lc"] +"""Constraint groups that include pLI/pNull/pRec in release format.""" + +RELEASE_GROUP_RENAMES = {"lof_hc": "lof"} +"""Internal constraint group names that are renamed for public release.""" + +RELEASE_PIPELINE_PARAM_GLOBALS = [ + ( + "calculate_mu_globals", + "calculate_mu_params", + ["freq_meta", "genetic_ancestry_groups", "downsampling_idx"], + ), + ( + "build_models_globals", + "build_models_params", + ["synonymous_transcript_filter_field", "skip_coverage_model"], + ), + ( + "apply_models_globals", + "apply_models_params", + ["skip_coverage_model", "groupings"], + ), +] +"""Pipeline parameter globals: (internal name, release name, fields to drop).""" + +RELEASE_CG_RENAME = { + "mu_snp": "mu", + "possible_variants": "possible", + "observed_variants": "obs", + "expected_variants": "exp", +} +"""Field renames applied to the constraint-group structs in release format.""" + +RELEASE_CG_SELECT = [ + "mu", + "possible", + "obs", + "exp", + "oe", + "z_raw", + "z_score", + "oe_ci", + "gen_anc_obs", + "gen_anc_exp", +] +"""Fields selected from the release constraint-group struct, in display order.""" + +# --------------------------------------------------------------------------- +# Constraint percentile threshold computation and annotation +# --------------------------------------------------------------------------- + +CONSTRAINT_SCORE_CAP = 2.0 +"""OE upper CI value above which scores are considered unconstrained and capped.""" + +PLI_EXPECTED_VALUES = {"Null": 1.0, "Rec": 0.706, "LI": 0.207} +"""Expected o/e values for the pLI model (null, recessive, loss-of-function intolerant).""" + +CONSTRAINT_METRICS = ["lof", "mis", "syn"] +"""Constraint metrics for which percentile thresholds are computed.""" + +CONSTRAINT_GRANULARITIES = { + "percentile": list(range(1, 100)), + "decile": list(range(1, 10)), + "sextile": list(range(1, 6)), +} +""" +Granularities for percentile binning. + +Keys are granularity names; values are boundary bin labels (1-indexed). +Quantile probabilities are bin / (max_bin + 1). +""" diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index f835adc1..4c783d3e 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -1,7 +1,7 @@ -"""Script containing resource utility constants, reference resources, and resources of intermediate files generated by the constraint pipeline.""" +"""Resource utility functions and resource definitions for the constraint pipeline.""" import logging -from typing import Dict, Optional, Tuple, Union +from typing import Optional, Union import gnomad.resources.grch37.gnomad as gnomad_grch37 import gnomad.resources.grch37.reference_data as ref_grch37 @@ -13,6 +13,16 @@ ExpressionResource, TableResource, VersionedTableResource, + import_gencode, +) + +from gnomad_constraint.resources.constants import ( + CURRENT_VERSION, + CUSTOM_VEP_ANNOTATIONS, + DATA_TYPES, + EXTENSIONS, + MODEL_TYPES, + VERSIONS, ) logging.basicConfig( @@ -22,71 +32,6 @@ logger = logging.getLogger("constraint_pipeline") logger.setLevel(logging.INFO) -EXTENSIONS = ["ht", "tsv", "tsv.bgz", "he", "log"] -VERSIONS = ["2.1.1", "4.0", "4.1"] -CURRENT_VERSION = "4.1" -DATA_TYPES = ["context", "exomes", "genomes"] -MODEL_TYPES = ["plateau", "coverage"] -GENOMIC_REGIONS = ["autosome_par", "chrx_nonpar", "chry_nonpar"] - -CUSTOM_VEP_ANNOTATIONS = ["transcript_consequences", "worst_csq_by_gene"] -""" -VEP annotations used when applying models. - -"transcript_consequences" option will annotate the Table with 'annotation', 'gene', -'coverage', 'transcript', and either 'canonical' or 'mane_select' annotations using 'transcript_consequences' -VEP annotation. - -"worst_csq_by_gene" option will annotate the Table with 'annotation', 'gene', and -'coverage' annotations using 'worst_csq_by_gene' VEP annotation. -""" - -POPS = ("global", "afr", "amr", "eas", "nfe", "sas") -""" -Population labels from gnomAD. - -Abbreviations stand for: global (all populations), African-American/African, Latino, East Asian, Non-Finnish European, and South Asian. -""" - -COVERAGE_CUTOFF = 40 -""" -Minimum median exome coverage differentiating high coverage sites from low coverage sites. - -Low coverage sites require an extra calibration when computing the proportion of expected variation. -""" - -MU_GROUPING = ("context", "ref", "alt", "methylation_level") -""" -Annotations used to group variants for the mutation rate calculation. -""" - -CALIBRATION_GROUPING = ("genomic_region", "build_model", "cpg", "exomes_coverage") -""" -Annotations used to group variants for the mutation rate calibration. -""" - -AGGREGATE_SUM_FIELDS = ( - "mu_snp", - "observed_variants", - "possible_variants", - "predicted_proportion_observed", - "coverage_correction", - "expected_variants", -) -""" -Fields to sum (or array sum) when aggregating the expected counts Table. -""" - -MUTATION_TYPE_FIELDS = ( - "cpg", - "transition", - "mutation_type", - "mutation_type_model", -) -""" -Fields added by `annotate_mutation_type`. -""" - def check_param_scope( version: Optional[str] = None, @@ -217,19 +162,76 @@ def get_coverage_ht( def get_gencode_ht(version: str) -> hl.Table: """ - Retrieve GENCODE Table. + Retrieve GENCODE Table with transcript version annotations. + + Re-imports the GENCODE GTF with ``include_version=True`` so that both + ``transcript_id_version`` and ``gene_id_version`` fields are present, + then checkpoints the result so subsequent calls read from the checkpoint. - :param version: gnomAD version. If version 2, GENCODE v19 will be loaded. If version 4, GENCODE v39 will be loaded. - :return: Table of GENCODE data for the specified build. + :param version: gnomAD version. If version 2, GENCODE v19 will be + loaded. If version 4, GENCODE v39 will be re-imported with version + fields and checkpointed. + :return: Table of GENCODE data with version annotations. """ if int(version[0]) == 2: return ref_grch37.gencode.ht() elif int(version[0]) == 4: - return ref_grch38.gencode.ht(read_args={"_n_partitions": 500}) + gencode_resource = ref_grch38.gencode + import_args = gencode_resource.versions[ + gencode_resource.default_version + ].import_args + ht = import_gencode(**import_args, include_version=True) + checkpoint_path = "gs://gnomad-tmp/gencode_v39_with_versions.ht" + + return ht.checkpoint(checkpoint_path, _read_if_exists=True) else: raise ValueError("Version must be within gnomAD v2 or v4.") +def get_gencode_cds_ht( + version: str = CURRENT_VERSION, +) -> TableResource: + """Build and checkpoint a per-locus GENCODE CDS transcript ID table. + + Calls :func:`get_gencode_ht` to retrieve the GENCODE table, filters to + CDS features, explodes each CDS interval into individual locus positions, + and groups by locus to produce an array of transcript IDs per position. + The result is checkpointed (read if it already exists) and returned as a + :class:`TableResource`. + + :param version: One of the release versions (`VERSIONS`). Default is + `CURRENT_VERSION`. + :return: TableResource of GENCODE CDS positions, keyed by locus with + ``transcript_id`` (array of transcript IDs whose CDS covers that + position). + """ + check_param_scope(version=version) + root = get_constraint_root(version=version, temp=True) + path = f"{root}/gencode_cds_positions.ht" + + gencode_ht = get_gencode_ht(version) + gencode_ht = gencode_ht.filter(gencode_ht.feature == "CDS").select("transcript_id") + gencode_ht = gencode_ht.annotate( + positions=hl.range( + gencode_ht.interval.start.position, + gencode_ht.interval.end.position + 1, + ) + ).explode("positions") + gencode_ht = gencode_ht.key_by( + locus=hl.locus( + gencode_ht.interval.start.contig, + gencode_ht.positions, + reference_genome="GRCh38", + ) + ).select("transcript_id") + gencode_ht = gencode_ht.group_by("locus").aggregate( + transcript_id=hl.agg.collect(gencode_ht.transcript_id) + ) + gencode_ht.checkpoint(path, _read_if_exists=True) + + return TableResource(path) + + def get_constraint_root( version: str = CURRENT_VERSION, test: bool = False, @@ -331,6 +333,45 @@ def get_mutation_ht(**kwargs) -> TableResource: return get_constraint_data("mutation_rate", sub_dir="mutation_rate", **kwargs) +def get_release_mutation_ht(version: str = CURRENT_VERSION) -> TableResource: + """ + Return TableResource for the release mutation rate Table. + + :param version: One of the release versions (`VERSIONS`). Default is + `CURRENT_VERSION`. + :return: TableResource of the release mutation rate Table. + """ + check_param_scope(version=version) + root = get_constraint_root(version=version) + return TableResource(f"{root}/release/gnomad.v{version}.mutation_rate.ht") + + +def get_release_constraint_ht(version: str = CURRENT_VERSION) -> TableResource: + """ + Return TableResource for the release constraint metrics Table. + + :param version: One of the release versions (`VERSIONS`). Default is + `CURRENT_VERSION`. + :return: TableResource of the release constraint metrics Table. + """ + check_param_scope(version=version) + root = get_constraint_root(version=version) + return TableResource(f"{root}/release/gnomad.v{version}.constraint_metrics.ht") + + +def get_release_constraint_tsv_path(version: str = CURRENT_VERSION) -> str: + """ + Return path for the release constraint metrics TSV. + + :param version: One of the release versions (`VERSIONS`). Default is + `CURRENT_VERSION`. + :return: Path of the release constraint metrics TSV. + """ + check_param_scope(version=version) + root = get_constraint_root(version=version) + return f"{root}/release/gnomad.v{version}.constraint_metrics.tsv.bgz" + + def get_annotated_context_ht(**kwargs) -> TableResource: """ Return TableResource of annotated context Table. @@ -452,29 +493,33 @@ def get_constraint_metrics_dataset( ) -def get_constraint_tsv_path(**kwargs) -> str: +def get_gene_quality_metrics_ht(version: str = CURRENT_VERSION) -> TableResource: """ - Return tsv path of pLI scores, observed:expected ratio, 90% confidence interval around the observed:expected ratio, and z scores. + Return TableResource of per-transcript gene quality metrics. + + Contains coverage and mapping quality metrics per transcript used for + annotating the release constraint Table. - :return: TSV path of constraint metrics. + :param version: One of the release versions (`VERSIONS`). Default is + `CURRENT_VERSION`. + :return: TableResource of gene quality metrics. """ - return get_constraint_data( - "constraint_metrics", sub_dir="metrics/tsv", extension="tsv.bgz", **kwargs - ) + check_param_scope(version=version) + root = get_constraint_root(version=version) + return TableResource(f"{root}/metrics/gnomad.v{version}.gene_quality_metrics.ht") -def get_downsampling_constraint_tsv_path(**kwargs) -> str: +def get_release_downsampling_tsv_path(version: str = CURRENT_VERSION) -> str: """ - Return tsv path of downsampling observed and expected counts. + Return path for the release per-genetic-ancestry downsampling TSV. - :return: TSV path of constraint metrics. + :param version: One of the release versions (`VERSIONS`). Default is + `CURRENT_VERSION`. + :return: Path of the release downsampling TSV. """ - return get_constraint_data( - "constraint_metrics.downsampling", - sub_dir="metrics/tsv", - extension="tsv.bgz", - **kwargs, - ) + check_param_scope(version=version) + root = get_constraint_root(version=version) + return f"{root}/release/gnomad.v{version}.constraint_metrics.downsampling.tsv.bgz" def get_logging_path(name: str, **kwargs) -> str: diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 2dd42b37..9a493b8f 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -16,30 +16,40 @@ annotate_mutation_type, annotate_with_mu, apply_models, - apply_plateau_models, calculate_raw_z_score, calculate_raw_z_score_sd, calibration_model_group_expr, compute_pli, count_observed_and_possible_by_group, - coverage_correction_expr, get_constraint_flags, oe_confidence_interval, - single_variant_count_expr, single_variant_observed_and_possible_expr, - weighted_agg_sum_expr, ) from gnomad.utils.filtering import add_filters_expr from gnomad.utils.vep import CSQ_CODING, filter_vep_transcript_csqs_expr from hail.utils.misc import divide_null, new_temp_file -from numpy.lib import r_ -from gnomad_constraint.resources.resource_utils import ( - AGGREGATE_SUM_FIELDS, +from gnomad_constraint.resources.constants import ( + ADJ_FREQ_META, CALIBRATION_GROUPING, + CONSTRAINT_GRANULARITIES, COVERAGE_CUTOFF, + GENCODE_FIELD_RENAMES, MU_GROUPING, MUTATION_TYPE_FIELDS, + PLI_EXPECTED_VALUES, + RELEASE_CG_RENAME, + RELEASE_CG_SELECT, + RELEASE_CI_FIELDS, + RELEASE_CI_FIELDS_WITH_RANK, + RELEASE_GROUP_NAMES, + RELEASE_GROUP_RENAMES, + RELEASE_GROUPS_WITH_PLI, + RELEASE_GROUPS_WITH_RANK, + RELEASE_KEY_ORDER, + RELEASE_LOF_FIELDS, + RELEASE_PIPELINE_PARAM_GLOBALS, + RELEASE_TOP_LEVEL_ANNOTATIONS, ) logging.basicConfig( @@ -89,15 +99,15 @@ def filter_freq_for_constraint( :return: Filtered frequency array and metadata. """ freq_meta = hl.eval(freq_meta_expr) - meta_keep = [{"group": "adj"}] + meta_keep = [ADJ_FREQ_META] if gen_ancs is not None: - meta_keep += [{"group": "adj", gen_anc_label: pop} for pop in gen_ancs] + meta_keep += [{**ADJ_FREQ_META, gen_anc_label: pop} for pop in gen_ancs] if downsamplings is not None: downsampling_pops = ["global"] + (downsampling_gen_ancs or []) meta_keep += [ - {"group": "adj", gen_anc_label: pop, "downsampling": str(ds)} + {**ADJ_FREQ_META, gen_anc_label: pop, "downsampling": str(ds)} for pop in downsampling_pops for ds in downsamplings ] @@ -200,7 +210,7 @@ def get_annotations_for_computing_mu( gen_anc_label="pop", ) downsampling_idx = genomes_freq_meta.index( - {"group": "adj", "pop": "global", "downsampling": str(downsampling_level)} + {**ADJ_FREQ_META, "pop": "global", "downsampling": str(downsampling_level)} ) # Filter to autosomal sites (remove pseudoautosomal regions). @@ -289,8 +299,8 @@ def get_exome_coverage_expr( an_meta = ht.an_globals.exomes.strata_meta # Get total AN count taking into account XX and XY samples for X and Y non-PAR. - xx_index = an_meta.index({"group": "adj", "sex": "XX"}) - xy_index = an_meta.index({"group": "adj", "sex": "XY"}) + xx_index = an_meta.index({**ADJ_FREQ_META, "sex": "XX"}) + xy_index = an_meta.index({**ADJ_FREQ_META, "sex": "XY"}) xx_an_sample_count = an_sample_count[xx_index] xy_an_sample_count = an_sample_count[xy_index] an_count = ( @@ -852,7 +862,6 @@ def aggregate_per_variant_expected_ht( ht = ht.group_by(*groupings).aggregate(**aggregate_expected_variants_expr(ht)) - """ # Check array lengths to determine if we need to batch. arrays = [ f for f in aggregate_fields_to_sum if isinstance(ht[f], hl.ArrayExpression) @@ -879,11 +888,13 @@ def aggregate_per_variant_expected_ht( **{f: ht[f][start_idx:end_idx] for f in arrays} ).checkpoint(new_temp_file(f"batch_{i}", "ht")) batches.append( - _ht.group_by(*groupings).aggregate( + _ht.group_by(*groupings) + .aggregate( **aggregate_expected_variants_expr( - _ht, - fields_to_sum=aggregate_fields_to_sum if i == 0 else arrays) - ).checkpoint(new_temp_file(f"batch_{i}.agg", "ht")) + _ht, fields_to_sum=aggregate_fields_to_sum if i == 0 else arrays + ) + ) + .checkpoint(new_temp_file(f"batch_{i}.agg", "ht")) ) ht = batches[0] batches = [_ht[ht.key] for _ht in batches[1:]] @@ -891,7 +902,6 @@ def aggregate_per_variant_expected_ht( **{f: hl.flatten([ht[f]] + [_ht[f] for _ht in batches]) for f in arrays} ) - """ ht = ht.checkpoint(new_temp_file("post_aggregation", "ht")) return ht.naive_coalesce(1000) @@ -954,7 +964,6 @@ def calculate_mu_by_downsampling( return annotate_mutation_type(ht) -# TODO: I think we decided this isn't needed right? We can just use canonical. def mane_select_over_canonical_filter_expr(ht: hl.Table) -> hl.Table: """ Filter to MANE Select over canonical transcripts. @@ -963,6 +972,12 @@ def mane_select_over_canonical_filter_expr(ht: hl.Table) -> hl.Table: select is specified, and a gene does not have a MANE select transcript, use canonical instead. + .. note:: + + In VEP 105 (used for gnomAD v4), all MANE Select transcripts are also + annotated as canonical. As a result, this function produces the same set of + transcripts as a simple canonical filter for v4 data. + :param ht: Table with the MANE Select and canonical annotations. :return: Table filtered to MANE Select over canonical transcripts. """ @@ -981,25 +996,117 @@ def mane_select_over_canonical_filter_expr(ht: hl.Table) -> hl.Table: ) +def get_transcript_filter_expr( + ht: hl.Table, + use_mane_select_over_canonical: bool = True, + mane_select_only: bool = False, +) -> hl.expr.BooleanExpression: + """ + Return a filter expression for selecting one representative transcript per gene. + + :param ht: Table with ``transcript``, ``mane_select``, ``canonical``, and + ``gene_id`` annotations. + :param use_mane_select_over_canonical: When ``True`` (default), prefer MANE + Select transcripts, falling back to canonical for genes without a MANE + Select entry. When ``False``, use canonical transcripts only. Ignored when + ``mane_select_only`` is ``True``. + :param mane_select_only: When ``True``, restrict to ENST MANE Select transcripts + only, with no canonical fallback. Default is ``False``. + :return: Boolean expression that is ``True`` for the selected transcripts. + """ + if mane_select_only: + return ht.transcript.startswith("ENST") & ht.mane_select + elif use_mane_select_over_canonical: + return mane_select_over_canonical_filter_expr(ht) + else: + return ht.transcript.startswith("ENST") & ht.canonical + + # TODO: Move to gnomad_methods? +def get_rank_and_bins( + value_expr: hl.expr.Float64Expression, + bin_granularities: Optional[Dict[str, int]] = None, +) -> hl.StructExpression: + """Rank rows by a numeric expression and assign bin labels. + + Rows are ordered ascending by ``value_expr``. Each row is assigned a + 0-based ``rank`` and a ``bin_{name}`` field for every entry in + ``bin_granularities``, computed as + ``hl.int(rank * multiplier / n_transcripts)``. + + :param value_expr: Numeric expression to rank by (ascending). + :param bin_granularities: Mapping of bin name to multiplier. Each entry + produces a ``bin_{name}`` field. Default is + ``{"percentile": 100, "decile": 10, "sextile": 6}``. + :return: Struct with ``rank`` and ``bin_{name}`` fields for each entry in + ``bin_granularities``. + """ + if bin_granularities is None: + bin_granularities = {"percentile": 100, "decile": 10, "sextile": 6} + + ht = value_expr._indices.source + source_key = list(ht.key) + n_transcripts = ht.count() + ranked_ht = ht.select(_=value_expr).order_by("_").add_index("rank") + ranked_ht = ranked_ht.select( + *source_key, + "rank", + **{ + f"bin_{name}": hl.int(ranked_ht.rank * multiplier / n_transcripts) + for name, multiplier in bin_granularities.items() + }, + ).cache() + + return ranked_ht.key_by(*source_key).cache()[ht.key] + + def add_oe_upper_rank_and_decile( ht: hl.Table, use_mane_select_over_canonical: bool = True, + mane_select_only: bool = False, + bin_granularities: Optional[Dict[str, int]] = None, ) -> hl.Table: """ Compute the rank and decile of the oe upper confidence interval. :param ht: Table with the oe upper confidence interval. - - :return: Struct containing the rank and decile of the oe upper confidence interval. + :param use_mane_select_over_canonical: Use MANE Select over canonical transcripts + for ranking, falling back to canonical when MANE Select is absent for a gene. + Default is True. Ignored when ``mane_select_only`` is True. + :param mane_select_only: Restrict ranking to ENST MANE Select transcripts only, + with no canonical fallback. Default is False. + :param bin_granularities: Mapping of bin name to multiplier used to assign each + transcript to a bin (``hl.int(rank * multiplier / n_transcripts)``). Each entry + produces a ``bin_{name}`` field. Default is + ``{"percentile": 100, "decile": 10, "sextile": 6}``. + :return: Input table with ``oe_ci_{ci}_rank`` fields added at the constraint-group + level (e.g., ``oe_ci_discretized_poisson_rank``, ``oe_ci_gamma_rank``), each + a struct with ``rank`` and ``bin_{name}`` fields for every entry in + ``bin_granularities``. Transcripts excluded from ranking have these fields set + to missing. """ + # Add an integer index so re-keying after order_by is an O(1) lookup. + ht = ht.add_index("_idx").key_by("_idx").cache() + total_count = ht.count() - if use_mane_select_over_canonical: - rank_filter_expr = mane_select_over_canonical_filter_expr(ht) - else: - rank_filter_expr = ht.transcript.startswith("ENST") & ht.canonical + ms_ht = ht.filter( + get_transcript_filter_expr(ht, use_mane_select_over_canonical, mane_select_only) + ) - ms_ht = ht.filter(rank_filter_expr) + # Extract only the first freq group (adj/all-samples) per constraint group for + # ranking. + ci_fields = ["discretized_poisson", "gamma"] + ms_ht = ( + ms_ht.select( + oe_ci_upper=ms_ht.constraint_groups.map( + lambda x: hl.struct( + **{ci: x.oe_info[0][f"oe_ci_{ci}"].upper for ci in ci_fields} + ) + ), + ) + .naive_coalesce(100) + .checkpoint(new_temp_file("oe_ci_upper.before_rank", "ht")) + ) n_transcripts = ms_ht.count() logger.info( "Retaining %d out of %d transcripts to use for rank annotations.", @@ -1007,97 +1114,98 @@ def add_oe_upper_rank_and_decile( total_count, ) - """ - ms_ht = ms_ht.select( - oe_ci_upper=ms_ht.constraint_groups.map( - lambda x: x.oe_info.map( - lambda y: hl.struct( - upper_ci=y.oe_ci.upper, - upper_ci_chisq=y.oe_ci_chisq.upper, - upper_ci_gamma=y.oe_ci_gamma.upper, - ) + # For each (constraint group, CI method), rank a minimal 2-column table and + # checkpoint it, then join all rank tables back in one pass. + num_constraint_groups = hl.eval(ms_ht.constraint_group_meta.length()) + ms_ht = ms_ht.annotate( + oe_ci_upper=[ + hl.struct( + **{ + ci: get_rank_and_bins(ms_ht.oe_ci_upper[i][ci], bin_granularities) + for ci in ci_fields + } ) - ), - ).naive_coalesce(100).checkpoint(new_temp_file("oe_ci_upper.before_rank", "ht")) + for i in range(num_constraint_groups) + ] + ).cache() + + # Annotate each constraint group with rank/bin fields at the group level (not + # inside oe_info, since all array elements must share the same struct schema). + # ht is already keyed by _idx, so ms_ht can be looked up directly. + ms_keyed = ms_ht[ht._idx] + ht = ht.annotate( + constraint_groups=hl.if_else( + hl.is_defined(ms_keyed.oe_ci_upper), + hl.map( + lambda g, r: g.annotate( + **{f"oe_ci_{ci}_rank": r[ci] for ci in ci_fields} + ), + ht.constraint_groups, + ms_keyed.oe_ci_upper, + ), + ht.constraint_groups, + ) + ) + + return ht - num_constraint_groups = hl.eval(ms_ht.constraint_group_meta.length()) - num_freq_groups = hl.eval(ms_ht.exomes_freq_meta.length()) - for i in range(num_constraint_groups): - for j in range(num_freq_groups): - ms_ht = ms_ht.annotate(upper_ci_rank1=hl.struct()) - for k in ["upper_ci", "upper_ci_chisq", "upper_ci_gamma"]: - # Rank in ascending order. - ms_ht = ms_ht.order_by(ms_ht.oe_ci_upper[i][j][k]) - _rank = hl.struct(**{f'{k}_rank': hl.scan.count()}) - if "upper_ci_rank1" in ms_ht.row: - _rank = ms_ht.upper_ci_rank1.annotate(**_rank) - ms_ht = ms_ht.annotate(upper_ci_rank1=_rank) - - _rank = ms_ht.upper_ci_rank1 - if "upper_ci_rank2" in ms_ht.row: - _rank = ms_ht.upper_ci_rank2.append(_rank) - else: - _rank = [_rank] - ms_ht = ms_ht.annotate(upper_ci_rank2=_rank) - ms_ht = ms_ht.checkpoint(new_temp_file(f"constraint_group{i}.freq_groups{j}", "ht")) - - _rank = ms_ht.upper_ci_rank2 - if "upper_ci_rank3" in ms_ht.row: - _rank = ms_ht.upper_ci_rank3.append(_rank) - else: - _rank = [_rank] - ms_ht = ms_ht.annotate(upper_ci_rank3=_rank) - - ms_ht = ms_ht.checkpoint(new_temp_file("oe_ci_upper.rank", "ht")) - """ - ms_ht = hl.read_table( - "gs://gnomad-tmp-4day/oe_ci_upper.rank-664KaxZ6k3vyUFBpZIJaqh.ht" +def compute_oe_upper_percentile_thresholds( + ht: hl.Table, + percentiles: List[float], + metric_expr: hl.expr.Float64Expression, + outlier_expr: hl.expr.BooleanExpression, + use_mane_select_over_canonical: bool = True, + mane_select_only: bool = False, + quantile_k: int = 1000, +) -> List[float]: + """ + Compute OE upper CI percentile thresholds for a single metric expression. + + Filters to a representative transcript set (controlled by + ``use_mane_select_over_canonical`` / ``mane_select_only``) and excludes + outlier transcripts, then computes approximate quantile thresholds at the + requested percentiles in a single aggregation pass. + + :param ht: Constraint metrics Table (output of ``compute_constraint_metrics``). + :param percentiles: Percentile values (0–100) at which to compute thresholds. + Pass a combined list across multiple granularities and slice the result to + avoid repeated aggregation passes. + :param metric_expr: Float expression for the metric to threshold (e.g., + ``ht.constraint_groups[i].oe_info[0].oe_ci_gamma.upper``). Must be + defined on ``ht``. + :param outlier_expr: Boolean expression that is ``True`` for transcripts to + exclude from the reference population (e.g., flagged transcripts). + :param use_mane_select_over_canonical: When ``True`` (default), prefer MANE + Select transcripts, falling back to canonical for genes without a MANE + Select entry. Ignored when ``mane_select_only`` is ``True``. + :param mane_select_only: When ``True``, restrict to ENST MANE Select + transcripts only, with no canonical fallback. Default is ``False``. + :param quantile_k: Accuracy parameter for + :func:`hail.expr.aggregators.approx_quantiles`. Default is 1000. + :return: List of float threshold values at the given percentiles. + """ + mane_filter_expr = get_transcript_filter_expr( + ht, use_mane_select_over_canonical, mane_select_only ) - # Map rank and bin annotations back to original Table. - ht = ht.annotate(upper_ci_rank3=ms_ht.key_by(*list(ht.key))[ht.key].upper_ci_rank3) - ht = ht.annotate( - constraint_groups=hl.enumerate(ht.constraint_groups).map( - lambda x: x[1].annotate( - oe_info=hl.enumerate(x[1].oe_info).map( - lambda y: hl.bind( - lambda r: y[1].annotate( - **{ - f"oe_ci{k}": y[1][f"oe_ci{k}"].annotate( - **hl.or_missing( - hl.is_defined(r), - hl.struct( - upper_rank=r[f"upper_ci{k}_rank"], - upper_bin_percentile=hl.int( - r[f"upper_ci{k}_rank"] - * 100 - / n_transcripts - ), - upper_bin_decile=hl.int( - r[f"upper_ci{k}_rank"] - * 10 - / n_transcripts - ), - upper_bin_sextile=hl.int( - r[f"upper_ci{k}_rank"] - * 6 - / n_transcripts - ), - ), - ) - ) - for k in ["", "_chisq", "_gamma"] - } - ), - ht.upper_ci_rank3[x[0]][y[0]], - ) - ) - ) + qs = [p / 100.0 for p in percentiles] + filt = mane_filter_expr & hl.is_defined(metric_expr) & ~outlier_expr + + result = ht.aggregate( + hl.struct( + thresholds=hl.agg.filter( + filt, hl.agg.approx_quantiles(metric_expr, qs, k=quantile_k) + ), + n=hl.agg.count_where(filt), ) - ).drop("upper_ci_rank3") + ) + logger.info( + "Computed percentile thresholds on %d transcripts.", + result.n, + ) - return ht + return result.thresholds # TODO: Move to gnomad_methods? @@ -1338,47 +1446,141 @@ def gamma_ci( ) -def chisq_ci( +def calculate_oe_confidence_interval( obs: hl.expr.Int32Expression, exp: hl.expr.Float64Expression, alpha: float = 0.05, ) -> hl.expr.StructExpression: - """ - Calculate the upper bound of the OE confidence interval using the chi-squared - distribution. + """Calculate the OE confidence interval using the Gamma distribution. :param obs: Observed count. :param exp: Expected count. - :param alpha: Significance level for the confidence interval. Default is 0.05. - :return: Upper bound of the OE confidence interval. + :param alpha: Significance level for the confidence interval. Default is + 0.05. + :return: Struct with ``lower`` and ``upper`` bounds. """ - return hl.struct( - lower=hl.qchisqtail(alpha, 2 * obs, lower_tail=True) / (2 * exp), - upper=hl.qchisqtail(1 - alpha, 2 * (obs + 1), lower_tail=True) / (2 * exp), + return gamma_ci(obs, exp, alpha) + + +def _compute_coverage_metrics( + ht: hl.Table, + gencode_cds_ht: hl.Table, + an_coverage_threshold: int = 90, +) -> hl.Table: + """Compute per-transcript proportion of CDS bases with adequate coverage. + + Uses the ``exomes_coverage`` field from the preprocessed context table + (AN as a percentage of total alleles) to determine what fraction of CDS + bases per transcript meet the coverage threshold. + + :param ht: Preprocessed context Hail Table with ``exomes_coverage`` (AN percent, + 0-100) per position. + :param gencode_cds_ht: GENCODE CDS positions Table keyed by locus with + ``transcript_id`` array. + :param an_coverage_threshold: Minimum ``exomes_coverage`` value (0-100) for a + position to be considered adequately covered. Default is 90. + :return: Table keyed by ``transcript`` with ``prop_bp_AN90``. + """ + # Deduplicate context table by locus (3 SNV alts per position share coverage). + ht = ht.key_by("locus").select("exomes_coverage").distinct() + + # Join CDS positions with context coverage. + ht = gencode_cds_ht.annotate( + exomes_coverage=ht[gencode_cds_ht.locus].exomes_coverage ) + ht = ht.filter(hl.is_defined(ht.exomes_coverage)) + # Explode by transcript and aggregate. + ht = ht.explode("transcript_id").cache() -def calculate_oe_confidence_interval( - obs, exp, alpha=0.05, oe_upper_method: str = "gamma" -): + return ht.group_by(transcript=ht.transcript_id).aggregate( + prop_bp_AN90=hl.agg.fraction(ht.exomes_coverage >= an_coverage_threshold), + ) + + +def _compute_site_quality_metrics( + ht: hl.Table, + gencode_cds_ht: hl.Table, +) -> hl.Table: + """Compute per-transcript mapping quality and region flag metrics. + + Computes mean AS_MQ, proportion of sites in segmental duplications, and + proportion of sites in low-complexity regions from variant sites within + CDS regions. + + :param ht: gnomAD exomes sites Hail Table. + :param gencode_cds_ht: GENCODE CDS positions Table keyed by locus with + ``transcript_id`` array. + :return: Table keyed by ``transcript`` with ``mean_AS_MQ``, + ``prop_segdup``, and ``prop_LCR``. """ - Calculate the upper bound of the OE confidence interval. + # Deduplicate sites by locus. + ht = ht.select("region_flags", AS_MQ=ht.info.AS_MQ) + ht = ht.key_by("locus").select("AS_MQ", "region_flags").distinct() + + # Join with GENCODE CDS to get per-locus transcript IDs. + ht = ht.annotate(transcript_id=gencode_cds_ht[ht.locus].transcript_id) + ht = ht.filter(hl.is_defined(ht.transcript_id)).explode("transcript_id").cache() + + return ht.group_by(transcript=ht.transcript_id).aggregate( + mean_AS_MQ=hl.agg.mean(ht.AS_MQ), + prop_segdup=hl.agg.fraction(ht.region_flags.segdup), + prop_LCR=hl.agg.fraction(ht.region_flags.lcr), + ) + - :param oe_expr: Array expression with observed and expected values. - :param alpha: Significance level for the OE confidence interval. Default is 0.05. - :return: Array expression with upper bound of the OE confidence interval. +def compute_gene_quality_metrics( + context_ht: hl.Table, + exomes_ht: hl.Table, + gencode_cds_ht: hl.Table, + an_coverage_threshold: int = 90, +) -> hl.Table: + """Compute per-transcript gene quality metrics. + + Combines coverage metrics from :func:`_compute_coverage_metrics` and + site quality metrics from :func:`_compute_site_quality_metrics` into a + single Table with release-ready fields: + + - ``gene_quality_metrics``: struct with ``exome_prop_bp_AN90``, + ``exome_mean_AS_MQ``, ``exome_prop_segdup``, ``exome_prop_LCR``. + - ``gene_flags``: set of flag strings (``low_exome_mapping_quality`` + when mean AS_MQ < 50, ``low_exome_coverage`` when + prop_bp_AN90 < 0.1). + + :param context_ht: Preprocessed context Hail Table with + ``exomes_coverage`` (AN percent, 0-100) per position. + :param exomes_ht: gnomAD exomes sites Hail Table. + :param gencode_cds_ht: GENCODE CDS positions Table keyed by locus with + ``transcript_id`` array (output of + :func:`~gnomad_constraint.resources.resource_utils.get_gencode_cds_ht`). + :param an_coverage_threshold: Minimum ``exomes_coverage`` value (0-100) + for a position to be considered adequately covered. Default is 90. + :return: Table keyed by ``transcript`` with ``gene_quality_metrics`` + and ``gene_flags``. """ - # Calculate upper bound of oe confidence interval. - if oe_upper_method not in ["gamma", "chisq"]: - raise ValueError(f"Invalid OE upper method: {oe_upper_method}") + an90_ht = _compute_coverage_metrics( + context_ht, gencode_cds_ht, an_coverage_threshold + ).cache() + sites_ht = _compute_site_quality_metrics(exomes_ht, gencode_cds_ht).cache() + + ht = an90_ht.annotate(**sites_ht[an90_ht.transcript]) + ht = ht.select( + gene_quality_metrics=hl.struct(**{f"exome_{f}": ht[f] for f in ht.row_value}), + gene_flags=add_filters_expr( + { + "low_exome_mapping_quality": ht.mean_AS_MQ < 50, + "low_exome_coverage": ht.prop_bp_AN90 < 0.1, + } + ), + ) - oe_upper_func = gamma_ci if oe_upper_method == "gamma" else chisq_ci - return oe_upper_func(obs, exp, alpha) + return ht.key_by("transcript") def compute_constraint_metrics( ht: hl.Table, gencode_ht: hl.Table, + gene_quality_metrics_ht: hl.Table, expected_values: Optional[Dict[str, float]] = None, min_diff_convergence: float = 0.001, raw_z_outlier_threshold_lower_lof: float = -8.0, @@ -1391,7 +1593,9 @@ def compute_constraint_metrics( Compute the pLI scores, observed:expected ratio, 90% confidence interval around the observed:expected ratio, and z scores for synonymous variants, missense variants, and predicted loss-of-function (pLoF) variants. .. note:: + The following annotations should be present in `ht`: + - modifier - annotation - observed_variants @@ -1429,6 +1633,9 @@ def compute_constraint_metrics( upper confidence interval. If a gene does not have a MANE Select transcript, the canonical transcript (if available) will be used instead. Default is True. :param gencode_ht: Table containing GENCODE annotations. + :param gene_quality_metrics_ht: Table keyed by transcript with + ``gene_quality_metrics`` and ``gene_flags`` fields (output of + :func:`compute_gene_quality_metrics`). :return: Table with pLI scores, observed:expected ratio, confidence interval of the observed:expected ratio, and z scores. """ @@ -1444,17 +1651,13 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: exp = oe_info.expected_variants return oe_info.annotate( oe=divide_null(obs, exp), - oe_ci=oe_confidence_interval(obs, exp), - oe_ci_chisq=calculate_oe_confidence_interval( - obs, exp, oe_upper_method="chisq" - ), + oe_ci_discretized_poisson=oe_confidence_interval(obs, exp), oe_ci_gamma=calculate_oe_confidence_interval( obs, exp, oe_upper_method="gamma" ), z_raw=calculate_raw_z_score(obs, exp), ) - """ # Annotate with the observed:expected ratio, 95% confidence interval around the # observed:expected ratio, and z scores for each constraint group. ht = ht.annotate( @@ -1473,14 +1676,13 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: raw_z_outlier_threshold_upper_syn, ), } - """ meta = hl.eval(ht.constraint_group_meta) freq_meta = hl.eval(ht.exomes_freq_meta) syn_idx = meta.index({"csq_set": "syn"}) mis_idx = meta.index({"csq_set": "mis"}) lof_idx = meta.index({"lof": "hc"}) - all_freq_idx = freq_meta.index({"group": "adj"}) - """ht = ht.annotate( + all_freq_idx = freq_meta.index(ADJ_FREQ_META) + ht = ht.annotate( constraint_groups=[ ht.constraint_groups[i].annotate( flags=add_filters_expr( @@ -1533,13 +1735,8 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: ht.constraint_groups[syn_idx].flags | ht.constraint_groups[mis_idx].flags | ht.constraint_groups[lof_idx].flags - ) + ), ).checkpoint(new_temp_file("constraint_metrics.oe.oe_ci.z_raw.flags", "ht")) - """ - - ht = hl.read_table( - "gs://gnomad-tmp-4day/constraint_metrics.oe.oe_ci.z_raw.flags-D9jfiSqRyWZoPb6VmAzNZw.ht" - ) # Add a rank and decile of the upper confidence interval for MANE Select or # canonical ensembl transcripts. @@ -1547,9 +1744,43 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: new_temp_file("constraint_metrics.oe.oe_ci.z_raw.flags.rank_and_decile", "ht") ) + # Compute OE upper CI percentile thresholds and annotate bins. + # Look up constraint group indices for each metric from the global meta. + meta = hl.eval(ht.constraint_group_meta) + metric_group_idx = { + "syn": next(i for i, m in enumerate(meta) if m == {"csq_set": "syn"}), + "mis": next(i for i, m in enumerate(meta) if m == {"csq_set": "mis"}), + "lof": next(i for i, m in enumerate(meta) if m == {"lof": "hc"}), + } + outlier_expr = ht.constraint_flags.length() > 0 + + # Combine all granularity boundary quantiles for a single aggregate pass per metric. + all_qs = [] + gran_slices: Dict[str, slice] = {} + for gran_name, bins in CONSTRAINT_GRANULARITIES.items(): + n_bins = len(bins) + 1 + start = len(all_qs) + all_qs.extend(b / n_bins * 100 for b in bins) + gran_slices[gran_name] = slice(start, len(all_qs)) + + # Call once per metric and assemble thresholds dict. + thresholds = {} + for metric, idx in metric_group_idx.items(): + vals = compute_oe_upper_percentile_thresholds( + ht, + percentiles=all_qs, + metric_expr=ht.constraint_groups[idx].oe_info[0].oe_ci_gamma.upper, + outlier_expr=outlier_expr, + mane_select_only=True, + ) + for gran_name, sl in gran_slices.items(): + thresholds[(gran_name, metric)] = list(vals[sl]) + + ht = annotate_constraint_percentile_bins(ht, thresholds, metric_group_idx) + # Compute the observed:expected ratio. if expected_values is None: - expected_values = {"Null": 1.0, "Rec": 0.706, "LI": 0.207} + expected_values = PLI_EXPECTED_VALUES hc_lof_expr = ht.constraint_groups[lof_idx].oe_info[all_freq_idx] ht = ht.annotate( @@ -1566,6 +1797,9 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: ) ) + # Add per-transcript gene quality metrics and flags. + ht = ht.annotate(**gene_quality_metrics_ht[ht.transcript]) + # Add transcript annotations from GENCODE. ht = add_gencode_transcript_annotations(ht, gencode_ht) @@ -1605,3 +1839,321 @@ def calculate_gerp_cutoffs(ht: hl.Table) -> Tuple[float, float]: cutoff_upper = list(filter(lambda i: i[1] < 0.95, zipped))[-1][0] return cutoff_lower, cutoff_upper + + +def _restructure_release_rows( + ht: hl.Table, + field_names: List[str], + all_freq_idx: int, + gen_anc_ds_indices: Dict[str, List[int]], +) -> hl.Table: + """ + Restructure ``constraint_groups`` into named top-level release fields. + + For each constraint group, builds a flat release struct by: + + - Flattening the adjusted-frequency ``oe_info`` entry onto the group + struct, keeping ``oe`` and ``z_raw`` under their original names and + overriding ``oe_ci`` with ``oe_ci_gamma``. + - Applying ``RELEASE_CG_RENAME`` to rename group-level and ``oe_info`` + fields (e.g. ``mu_snp`` -> ``mu``, ``observed_variants`` -> ``obs``). + - When downsampling data is present, adding ``gen_anc_obs`` / + ``gen_anc_exp`` structs keyed by genetic ancestry with arrays of + values ordered by downsampling level. + + Annotates the Table with one top-level field per group (applying + ``RELEASE_GROUP_RENAMES``, e.g. ``lof_hc`` -> ``lof``), trims the + ``oe_ci`` struct to ranked or unranked CI fields depending on the group, + and adds ``pLI`` / ``pNull`` / ``pRec`` for the LoF group. Finally + selects release row fields, re-keys, and filters out transcripts with + no possible variants in any group. + + :param ht: Table with ``constraint_groups`` and associated annotations. + :param field_names: Internal name for each constraint group, derived + from ``constraint_group_meta``. + :param all_freq_idx: Index into ``oe_info`` for the adjusted allele + frequency group. + :param gen_anc_ds_indices: Mapping from genetic ancestry label to list + of ``oe_info`` indices for its downsampling entries. Empty dict when + no downsampling data is present. + :return: Table with named top-level constraint group structs, release + row fields selected, re-keyed, and filtered. + """ + add_ds_fields = ( + ["observed_variants", "expected_variants"] if gen_anc_ds_indices else [] + ) + + cg_expr = ht.constraint_groups.map( + lambda cg: cg.annotate( + **cg.oe_info[all_freq_idx], + oe_ci=cg.oe_info[all_freq_idx].oe_ci_gamma, + # Per-genetic-ancestry downsampling obs/exp arrays, one value per + # downsampling level, keyed by genetic ancestry. + **{ + f"gen_anc_{RELEASE_CG_RENAME[f]}": hl.struct( + **{ + gen_anc: hl.array([cg.oe_info[j][f] for j in indices]) + for gen_anc, indices in gen_anc_ds_indices.items() + } + ) + for f in add_ds_fields + }, + ) + ) + cg_select = ( + RELEASE_CG_SELECT + if gen_anc_ds_indices + else [f for f in RELEASE_CG_SELECT if not f.startswith("gen_anc_")] + ) + cg_expr = cg_expr.map( + lambda cg: cg.annotate( + **{RELEASE_CG_RENAME[k]: cg[k] for k in RELEASE_CG_RENAME} + ).select(*cg_select) + ) + + # Build a top-level release field for each constraint group, applying + # group renames (e.g. lof_hc -> lof), trimming oe_ci sub-fields, and + # adding pLI/pNull/pRec for the LoF group. + cg_fields = { + RELEASE_GROUP_RENAMES.get(name, name): cg_expr[i] + .annotate( + oe_ci=cg_expr[i].oe_ci.select( + *( + RELEASE_CI_FIELDS_WITH_RANK + if name in RELEASE_GROUPS_WITH_RANK + else RELEASE_CI_FIELDS + ) + ), + **{ + k: ht[k] + for k in (RELEASE_LOF_FIELDS if name in RELEASE_GROUPS_WITH_PLI else []) + }, + ) + .select() + for i, name in enumerate(field_names) + } + ht = ht.annotate(**cg_fields) + ht = ht.select(*RELEASE_TOP_LEVEL_ANNOTATIONS, *RELEASE_GROUP_NAMES) + + available_keys = [k for k in RELEASE_KEY_ORDER if k in ht.key] + if list(ht.key) != available_keys: + ht = ht.key_by(*available_keys) + + ht = ht.filter(hl.any([ht[k].possible != 0 for k in RELEASE_GROUP_NAMES])) + return ht + + +def _restructure_release_globals( + ht: hl.Table, + field_names: List[str], + freq_meta: List[Dict], + gen_anc_ds_indices: Dict[str, List[int]], + sd_raw_z_arr: List[float], + release_version: Optional[str], +) -> hl.Table: + """ + Restructure globals for public release. + + Replaces internal globals with a clean release set: + + - Pipeline parameter globals are renamed and stripped of internal-only + fields via ``RELEASE_PIPELINE_PARAM_GLOBALS``. + - ``sd_raw_z`` is converted from an ordered array (one entry per + constraint group) to a named struct keyed by release group name, + retaining only groups in ``RELEASE_GROUP_NAMES``. + - When downsampling data is present, a ``downsamplings`` struct is added + keyed by genetic ancestry, with arrays of integer downsampling levels + matching the order of ``gen_anc_obs`` / ``gen_anc_exp`` in the rows. + - ``max_af`` is preserved unchanged if present. + - ``version`` is set to ``release_version`` if provided, otherwise + carried over from the existing global. + + :param ht: Table whose globals are being restructured. + :param field_names: Internal name for each constraint group (parallel to + ``sd_raw_z_arr``), used to map array positions to release group names. + :param freq_meta: Evaluated ``exomes_freq_meta`` global, used to extract + downsampling levels for each genetic ancestry. + :param gen_anc_ds_indices: Mapping from genetic ancestry label to list + of ``oe_info`` indices for its downsampling entries. Empty dict when + no downsampling data is present. + :param sd_raw_z_arr: Evaluated ``sd_raw_z`` global array, parallel to + ``field_names``. + :param release_version: Version string for the ``version`` global. When + *None*, the existing ``version`` global is retained if present. + :return: Table with release-formatted globals. + """ + sd_raw_z_name_map = {n: RELEASE_GROUP_RENAMES.get(n, n) for n in field_names} + sd_raw_z_struct = hl.struct( + **{ + sd_raw_z_name_map[field_names[i]]: sd_raw_z_arr[i] + for i in range(len(field_names)) + if sd_raw_z_name_map[field_names[i]] in RELEASE_GROUP_NAMES + } + ) + + global_kwargs = {} + if release_version is not None: + global_kwargs["version"] = release_version + elif "version" in ht.globals: + global_kwargs["version"] = ht.globals.version + + for src, dest, drop_fields in RELEASE_PIPELINE_PARAM_GLOBALS: + if src in ht.globals: + global_kwargs[dest] = ht.globals[src].drop(*drop_fields) + + if gen_anc_ds_indices: + global_kwargs["downsamplings"] = hl.struct( + **{ + gen_anc: [int(freq_meta[j]["downsampling"]) for j in indices] + for gen_anc, indices in gen_anc_ds_indices.items() + } + ) + + if "max_af" in ht.globals: + global_kwargs["max_af"] = ht.globals.max_af + + global_kwargs["sd_raw_z"] = sd_raw_z_struct + return ht.select_globals(**global_kwargs) + + +def prepare_release_ht( + ht: hl.Table, + release_version: Optional[str] = None, +) -> hl.Table: + """ + Prepare the constraint metrics Table for public release. + + Computes shared metadata needed by both restructuring steps, then + delegates row and global restructuring to + :func:`_restructure_release_rows` and + :func:`_restructure_release_globals`. + + The internal ``constraint_groups`` schema has: + + - Group-level fields: ``mu_snp``, ``possible_variants``, ``z_score``. + - Per-frequency ``oe_info`` array (one entry per ``exomes_freq_meta`` + element): ``observed_variants``, ``expected_variants``, ``oe``, + ``oe_ci_gamma``, ``z_raw``. + + The release schema exposes one top-level struct per group + (``syn``, ``mis``, ``lof_hc_lc``, ``lof``; ``lof_hc`` is renamed to + ``lof``), with fields ``mu``, ``possible``, ``obs``, ``exp``, ``oe``, + ``oe_ci``, ``z_raw``, ``z_score``, and optionally ``gen_anc_obs`` / + ``gen_anc_exp`` when downsampling data is present. + + :param ht: Internal constraint metrics Table (output of + ``compute_constraint_metrics``). Expected to already contain GENCODE + transcript annotations (``transcript_id_version``, ``level``, etc.) + and gene quality metric annotations (``gene_quality_metrics``, + ``gene_flags``). + :param release_version: Version string for the ``version`` global. When + *None*, the existing ``version`` global is retained if present. + :return: Release-formatted Table. + """ + ht = ht.rename(GENCODE_FIELD_RENAMES) + + constraint_meta = hl.eval(ht.constraint_group_meta) + freq_meta = hl.eval(ht.exomes_freq_meta) + all_freq_idx = freq_meta.index(ADJ_FREQ_META) + + field_names = [ + "_".join(f"{k}_{v}" for k, v in m.items()).replace("csq_set_", "") + for m in constraint_meta + ] + logger.info("Release constraint group field names: %s", field_names) + + gen_anc_ds_indices: Dict[str, List[int]] = {} + if "downsamplings" in ht.globals: + for j, m in enumerate(freq_meta): + gen_anc = m.get("gen_anc") + if gen_anc is not None and "downsampling" in m: + gen_anc_ds_indices.setdefault(gen_anc, []).append(j) + + # Evaluate sd_raw_z before the row select (globals persist through it). + sd_raw_z_arr = hl.eval(ht.sd_raw_z) + + ht = _restructure_release_rows(ht, field_names, all_freq_idx, gen_anc_ds_indices) + ht = _restructure_release_globals( + ht, field_names, freq_meta, gen_anc_ds_indices, sd_raw_z_arr, release_version + ) + return ht + + +def flatten_release_ht(ht: hl.Table) -> hl.Table: + """ + Flatten the release constraint metrics Table for TSV export. + + Drops per-genetic-ancestry downsampling fields (``gen_anc_obs``, + ``gen_anc_exp``) when present and calls :meth:`~hail.Table.flatten` + to expand nested struct fields using ``.`` as the separator + (e.g. ``lof.obs``, ``lof.oe_ci.upper``). + + :param ht: Release-format constraint metrics Table (output of + :func:`prepare_release_ht`). + :return: Flat Table suitable for :meth:`~hail.Table.export`. + """ + # Drop struct/array fields not suitable for flat TSV export. + drop_fields = [f for f in ["gen_anc_obs", "gen_anc_exp"] if f in ht.row] + if drop_fields: + ht = ht.drop(*drop_fields) + + return ht.flatten() + + +def annotate_constraint_percentile_bins( + ht: hl.Table, + thresholds: Dict[Tuple[str, str], List[float]], + metric_group_idx: Dict[str, int], +) -> hl.Table: + """ + Annotate each transcript with its percentile bin for all metric/granularity combinations. + + Annotates ``constraint_bins.{granularity}.{metric}`` for each combination. + Bin 0 is the most constrained (value below all thresholds); bin N equals + the number of boundaries the value exceeds. + + :param ht: Constraint metrics Table with a ``constraint_groups`` array field. + :param thresholds: Mapping of ``(granularity, metric)`` to an ordered list + of threshold values, as produced by + :func:`compute_oe_upper_percentile_thresholds`. + :param metric_group_idx: Mapping of metric name to its index in + ``constraint_groups`` (e.g. ``{"lof": 5, "mis": 1, "syn": 0}``). + :return: Annotated Table with an added ``constraint_bins`` struct field. + """ + logger.info( + "Annotating bins for %d (granularity, metric) combinations.", + len(thresholds), + ) + + metric_exprs = { + metric: ht.constraint_groups[idx].oe_info[0].oe_ci_gamma.upper + for metric, idx in metric_group_idx.items() + } + miss = hl.missing(hl.tint32) + + def _bin_expr( + value_expr: hl.expr.Float64Expression, + threshold_list: List[float], + ) -> hl.expr.Int32Expression: + arr = hl.literal(threshold_list) + return hl.sum(arr.map(lambda t: hl.int(value_expr >= t))) + + return ht.annotate( + constraint_bins=hl.struct( + **{ + gran: hl.struct( + **{ + metric: hl.if_else( + hl.is_defined(metric_exprs[metric]), + _bin_expr(metric_exprs[metric], thresholds[(gran, metric)]), + miss, + ) + for metric in metric_group_idx + } + ) + for gran in CONSTRAINT_GRANULARITIES + } + ) + ) + + return ht.select(**flat) From c3930ea1baac619cd3238b03472baf70e2fb9529 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Mon, 16 Mar 2026 20:06:05 -0600 Subject: [PATCH 21/38] Refactor constraint pipeline resources and constants. Introduce classic loss-of-function annotations and streamline resource collection for the constraint pipeline. Remove deprecated code and improve clarity in resource utility functions. --- .../pipeline/constraint_pipeline.py | 258 +----------- gnomad_constraint/resources/constants.py | 8 + gnomad_constraint/resources/resource_utils.py | 194 ++++++++- gnomad_constraint/utils/constraint.py | 385 +++++++++--------- 4 files changed, 402 insertions(+), 443 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 96a94041..2f13c433 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -23,10 +23,9 @@ import argparse import logging -from typing import List, Optional +from typing import List import hail as hl -from gnomad.resources.grch38.gnomad import all_sites_an from gnomad.utils.constraint import ( assemble_constraint_context_ht, build_models, @@ -34,15 +33,13 @@ ) from gnomad.utils.reference_genome import get_reference_genome from gnomad.utils.vep import update_loftee_end_trunc_filter -from gnomad_qc.resource_utils import ( - PipelineResourceCollection, - PipelineStepResourceCollection, -) +from gnomad_qc.resource_utils import PipelineResourceCollection import gnomad_constraint.resources.resource_utils as constraint_res from gnomad_constraint.resources.constants import ( CURRENT_VERSION, CUSTOM_VEP_ANNOTATIONS, + RELEASE_KEY_ORDER, VERSIONS, ) from gnomad_constraint.utils.constraint import ( @@ -205,204 +202,6 @@ def _build_ht_dict(ht_name: str, keep: List[str] = None): return ht -def get_constraint_resources( - version: str, - custom_vep_annotation: str, - overwrite: bool, - test: bool, - models: List[str] = ["plateau", "coverage"], - directory_post_fix: Optional[str] = None, - path_post_fix: Optional[str] = None, -) -> PipelineResourceCollection: - """ - Get PipelineResourceCollection for all resources needed in the constraint pipeline. - - :param version: Version of constraint resources to use. - :param custom_vep_annotation: Custom VEP annotation to use for applying models - resources. - :param overwrite: Whether to overwrite existing resources. - :param test: Whether to use test resources. - :param models: List of models to use. Default is ["plateau", "coverage"]. - :param directory_post_fix: Post-fix to add to the directory path of the resources. - :param path_post_fix: Post-fix to add to the path of the resources. - :return: PipelineResourceCollection containing resources for all steps of the - constraint pipeline. - """ - # Initialize constraint pipeline resource collection. - constraint_pipeline = PipelineResourceCollection( - pipeline_name="constraint", - overwrite=overwrite, - ) - - # Create resource collection for each step of the constraint pipeline. - context_res = constraint_res.get_vep_context_ht(version) - context_build = get_reference_genome(context_res.ht().locus).name - - # Make dictionary for prepare_context input Tables. - input_hts = { - "context_ht": context_res, - "methylation_ht": constraint_res.get_methylation_ht(context_build), - } - for d in ["exomes", "genomes"]: - input_hts[f"{d}_coverage_ht"] = constraint_res.get_coverage_ht(d, version) - input_hts[f"{d}_sites_ht"] = constraint_res.get_sites_resource(d, version) - input_hts[f"{d}_an_ht"] = all_sites_an(d) - - common_params = { - "version": version, - "test": test, - "directory_post_fix": directory_post_fix, - } - - prepare_context = PipelineStepResourceCollection( - "--prepare-context-ht", - output_resources={ - "annotated_context_ht": constraint_res.get_annotated_context_ht( - **common_params - ) - }, - input_resources={"gnomAD resources": input_hts}, - ) - preprocess_data = PipelineStepResourceCollection( - "preprocess data for downstream steps", - output_resources={ - "temp_preprocess_data_ht": constraint_res.get_preprocessed_ht( - **common_params - ), - }, - pipeline_input_steps=[prepare_context], - ) - calculate_gerp_cutoffs = PipelineStepResourceCollection( - "--calculate-gerp-cutoffs", - output_resources={}, - pipeline_input_steps=[prepare_context], - ) - calculate_mutation_rate = PipelineStepResourceCollection( - "--calculate-mutation-rate", - output_resources={ - "mutation_ht": constraint_res.get_mutation_ht(**common_params) - }, - pipeline_input_steps=[preprocess_data], - ) - create_training_set = PipelineStepResourceCollection( - "--create-training-set", - output_resources={ - f"train_ht": constraint_res.get_training_dataset( - **common_params, path_post_fix=path_post_fix - ), - f"train_tsv": constraint_res.get_training_tsv_path( - **common_params, path_post_fix=path_post_fix - ), - }, - pipeline_input_steps=[preprocess_data, calculate_mutation_rate], - ) - build_models = PipelineStepResourceCollection( - "--build-models", - output_resources={ - f"model_{m}": constraint_res.get_models( - m, **common_params, path_post_fix=path_post_fix - ) - for m in models - }, - pipeline_input_steps=[create_training_set], - ) - apply_models_per_variant = PipelineStepResourceCollection( - "--apply-models-per-variant", - output_resources={ - "per_variant_apply_ht": constraint_res.get_per_variant_expected_dataset( - custom_vep_annotation, **common_params, path_post_fix=path_post_fix - ) - }, - pipeline_input_steps=[preprocess_data, calculate_mutation_rate, build_models], - ) - aggregate_per_variant_expected = PipelineStepResourceCollection( - "--aggregate-per-variant-expected", - output_resources={ - f"apply_ht": constraint_res.get_aggregated_per_variant_expected( - custom_vep_annotation, **common_params, path_post_fix=path_post_fix - ) - }, - pipeline_input_steps=[ - apply_models_per_variant, - calculate_mutation_rate, - build_models, - ], - ) - aggregate_by_constraint_groups = PipelineStepResourceCollection( - "--aggregate-by-constraint-groups", - output_resources={ - f"constraint_group_ht": constraint_res.get_constraint_group_ht( - custom_vep_annotation, **common_params, path_post_fix=path_post_fix - ) - }, - pipeline_input_steps=[aggregate_per_variant_expected], - ) - compute_gene_quality_metrics_step = PipelineStepResourceCollection( - "--compute-gene-quality-metrics", - output_resources={ - "gene_quality_metrics_ht": constraint_res.get_gene_quality_metrics_ht( - version=version - ) - }, - input_resources={ - "gnomAD resources": {"exomes_sites_ht": input_hts["exomes_sites_ht"]}, - }, - pipeline_input_steps=[prepare_context], - ) - compute_constraint_metrics = PipelineStepResourceCollection( - "--compute-constraint-metrics", - output_resources={ - "constraint_metrics_ht": constraint_res.get_constraint_metrics_dataset( - custom_vep_annotation, **common_params, path_post_fix=path_post_fix - ) - }, - pipeline_input_steps=[ - aggregate_by_constraint_groups, - compute_gene_quality_metrics_step, - ], - ) - prepare_release = PipelineStepResourceCollection( - "--prepare-release", - output_resources={ - "release_ht": constraint_res.get_release_constraint_ht(version=version), - }, - pipeline_input_steps=[compute_constraint_metrics], - ) - export_release_tsv = PipelineStepResourceCollection( - "--export-release-tsv", - output_resources={ - "release_tsv": constraint_res.get_release_constraint_tsv_path( - version=version - ), - "release_downsampling_tsv": constraint_res.get_release_downsampling_tsv_path( - version=version - ), - }, - pipeline_input_steps=[prepare_release], - ) - - # Add all steps to the constraint pipeline resource collection. - constraint_pipeline.add_steps( - { - "prepare_context": prepare_context, - "preprocess_data": preprocess_data, - "calculate_gerp_cutoffs": calculate_gerp_cutoffs, - "calculate_mutation_rate": calculate_mutation_rate, - "create_training_set": create_training_set, - "build_models": build_models, - "apply_models_per_variant": apply_models_per_variant, - "aggregate_per_variant_expected": aggregate_per_variant_expected, - "aggregate_by_constraint_groups": aggregate_by_constraint_groups, - "compute_gene_quality_metrics": compute_gene_quality_metrics_step, - "compute_constraint_metrics": compute_constraint_metrics, - "prepare_release": prepare_release, - "export_release_tsv": export_release_tsv, - } - ) - - return constraint_pipeline - - def main(args): """Execute the constraint pipeline.""" hl.init( @@ -434,7 +233,7 @@ def main(args): models = ["plateau", "coverage"] if not skip_coverage_model else ["plateau"] # Construct resources with paths for intermediate Tables generated in the pipeline. - resources = get_constraint_resources( + resources = constraint_res.get_constraint_resources( version, custom_vep_annotation, overwrite, @@ -598,21 +397,12 @@ def main(args): res.check_resource_existence() # Use new shuffle method to prevent shuffle errors. - # hl._set_flags(use_new_shuffle="1") + hl._set_flags(use_new_shuffle="1") ht = res.per_variant_apply_ht.ht() - # ht = res.per_variant_apply_ht.ht(read_args={"_n_partitions": 8000}) - ht = aggregate_per_variant_expected_ht( - ht, include_mu_annotations_in_grouping=True - ) - ht = ht.checkpoint( - "gs://gnomad/v4.1/constraint_coverage_corrected/apply_models/transcript_consequences/gnomad.v4.1.per_variant_expected.aggregated_with_mu_annotations.coverage_corrected.with_downsamplings.ht", - overwrite=overwrite, - ) - # hl._set_flags(use_new_shuffle=None) - ht = aggregate_per_variant_expected_ht(ht) ht.write(res.apply_ht.path, overwrite=overwrite) + hl._set_flags(use_new_shuffle=None) logger.info( "Done aggregating per-variant expected variant count by transcript, " @@ -628,45 +418,11 @@ def main(args): # Use new shuffle method to prevent shuffle errors. hl._set_flags(use_new_shuffle="1") - ht = hl.read_table( - "gs://gnomad/v4.1/constraint_coverage_corrected/apply_models/transcript_consequences/gnomad.v4.1.per_variant_expected.aggregated_with_mu_annotations.coverage_corrected.with_downsamplings.ht" - ) - aggregate_by_constraint_groups( - ht, - keys=tuple( - [ - i - for i in list(ht.key) - if i - in [ - "gene", - "transcript", - "canonical", - "mane_select", - "gene_id", - "context", - "ref", - "alt", - "methylation_level", - ] - ] - ), - ).write( - "gs://gnomad/v4.1/constraint_coverage_corrected/apply_models/transcript_consequences/gnomad.v4.1.constraint_group_with_mu_annotations.coverage_corrected.with_downsamplings.ht", - overwrite=overwrite, - ) ht = res.apply_ht.ht() aggregate_by_constraint_groups( ht, - keys=tuple( - [ - i - for i in list(ht.key) - if i - in ["gene", "transcript", "canonical", "mane_select", "gene_id"] - ] - ), + keys=tuple([i for i in list(ht.key) if i in RELEASE_KEY_ORDER]), ).write(res.constraint_group_ht.path, overwrite=overwrite) hl._set_flags(use_new_shuffle=None) logger.info("Done with aggregating by constraint groups.") diff --git a/gnomad_constraint/resources/constants.py b/gnomad_constraint/resources/constants.py index 8a805b60..ee6b36a6 100644 --- a/gnomad_constraint/resources/constants.py +++ b/gnomad_constraint/resources/constants.py @@ -48,6 +48,13 @@ Low coverage sites require an extra calibration when computing the proportion of expected variation. """ +CLASSIC_LOF_ANNOTATIONS = ( + "stop_gained", + "splice_donor_variant", + "splice_acceptor_variant", +) +"""Classic loss-of-function VEP annotations.""" + MU_GROUPING = ("context", "ref", "alt", "methylation_level") """ Annotations used to group variants for the mutation rate calculation. @@ -60,6 +67,7 @@ AGGREGATE_SUM_FIELDS = ( "mu_snp", + "mu", "observed_variants", "possible_variants", "predicted_proportion_observed", diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 4c783d3e..6607ac3e 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -1,13 +1,14 @@ """Resource utility functions and resource definitions for the constraint pipeline.""" import logging -from typing import Optional, Union +from typing import List, Optional, Union import gnomad.resources.grch37.gnomad as gnomad_grch37 import gnomad.resources.grch37.reference_data as ref_grch37 import gnomad.resources.grch38.gnomad as gnomad_grch38 import gnomad.resources.grch38.reference_data as ref_grch38 import hail as hl +from gnomad.resources.grch38.gnomad import all_sites_an from gnomad.resources.resource_utils import ( BaseResource, ExpressionResource, @@ -15,6 +16,11 @@ VersionedTableResource, import_gencode, ) +from gnomad.utils.reference_genome import get_reference_genome +from gnomad_qc.resource_utils import ( + PipelineResourceCollection, + PipelineStepResourceCollection, +) from gnomad_constraint.resources.constants import ( CURRENT_VERSION, @@ -542,3 +548,189 @@ def get_checkpoint_path(name: str, **kwargs) -> TableResource: :return: Output checkpoint TableResource. """ return get_constraint_data(name, sub_dir="checkpoint_files", test=True, **kwargs) + + +def get_constraint_resources( + version: str, + custom_vep_annotation: str, + overwrite: bool, + test: bool, + models: List[str] = ["plateau", "coverage"], + directory_post_fix: Optional[str] = None, + path_post_fix: Optional[str] = None, +) -> PipelineResourceCollection: + """ + Get PipelineResourceCollection for all resources needed in the constraint pipeline. + + :param version: Version of constraint resources to use. + :param custom_vep_annotation: Custom VEP annotation to use for applying models + resources. + :param overwrite: Whether to overwrite existing resources. + :param test: Whether to use test resources. + :param models: List of models to use. Default is ["plateau", "coverage"]. + :param directory_post_fix: Post-fix to add to the directory path of the resources. + :param path_post_fix: Post-fix to add to the path of the resources. + :return: PipelineResourceCollection containing resources for all steps of the + constraint pipeline. + """ + # Initialize constraint pipeline resource collection. + constraint_pipeline = PipelineResourceCollection( + pipeline_name="constraint", + overwrite=overwrite, + ) + + # Create resource collection for each step of the constraint pipeline. + context_res = get_vep_context_ht(version) + context_build = get_reference_genome(context_res.ht().locus).name + + # Make dictionary for prepare_context input Tables. + input_hts = { + "context_ht": context_res, + "methylation_ht": get_methylation_ht(context_build), + } + for d in ["exomes", "genomes"]: + input_hts[f"{d}_coverage_ht"] = get_coverage_ht(d, version) + input_hts[f"{d}_sites_ht"] = get_sites_resource(d, version) + input_hts[f"{d}_an_ht"] = all_sites_an(d) + + common_params = { + "version": version, + "test": test, + "directory_post_fix": directory_post_fix, + } + + prepare_context = PipelineStepResourceCollection( + "--prepare-context-ht", + output_resources={ + "annotated_context_ht": get_annotated_context_ht(**common_params) + }, + input_resources={"gnomAD resources": input_hts}, + ) + preprocess_data = PipelineStepResourceCollection( + "preprocess data for downstream steps", + output_resources={ + "temp_preprocess_data_ht": get_preprocessed_ht(**common_params), + }, + pipeline_input_steps=[prepare_context], + ) + calculate_gerp_cutoffs = PipelineStepResourceCollection( + "--calculate-gerp-cutoffs", + output_resources={}, + pipeline_input_steps=[prepare_context], + ) + calculate_mutation_rate = PipelineStepResourceCollection( + "--calculate-mutation-rate", + output_resources={"mutation_ht": get_mutation_ht(**common_params)}, + pipeline_input_steps=[preprocess_data], + ) + create_training_set = PipelineStepResourceCollection( + "--create-training-set", + output_resources={ + "train_ht": get_training_dataset( + **common_params, path_post_fix=path_post_fix + ), + "train_tsv": get_training_tsv_path( + **common_params, path_post_fix=path_post_fix + ), + }, + pipeline_input_steps=[preprocess_data, calculate_mutation_rate], + ) + build_models = PipelineStepResourceCollection( + "--build-models", + output_resources={ + f"model_{m}": get_models(m, **common_params, path_post_fix=path_post_fix) + for m in models + }, + pipeline_input_steps=[create_training_set], + ) + apply_models_per_variant = PipelineStepResourceCollection( + "--apply-models-per-variant", + output_resources={ + "per_variant_apply_ht": get_per_variant_expected_dataset( + custom_vep_annotation, **common_params, path_post_fix=path_post_fix + ) + }, + pipeline_input_steps=[preprocess_data, calculate_mutation_rate, build_models], + ) + aggregate_per_variant_expected = PipelineStepResourceCollection( + "--aggregate-per-variant-expected", + output_resources={ + "apply_ht": get_aggregated_per_variant_expected( + custom_vep_annotation, **common_params, path_post_fix=path_post_fix + ) + }, + pipeline_input_steps=[ + apply_models_per_variant, + calculate_mutation_rate, + build_models, + ], + ) + aggregate_by_constraint_groups = PipelineStepResourceCollection( + "--aggregate-by-constraint-groups", + output_resources={ + "constraint_group_ht": get_constraint_group_ht( + custom_vep_annotation, **common_params, path_post_fix=path_post_fix + ) + }, + pipeline_input_steps=[aggregate_per_variant_expected], + ) + compute_gene_quality_metrics_step = PipelineStepResourceCollection( + "--compute-gene-quality-metrics", + output_resources={ + "gene_quality_metrics_ht": get_gene_quality_metrics_ht(version=version) + }, + input_resources={ + "gnomAD resources": {"exomes_sites_ht": input_hts["exomes_sites_ht"]}, + }, + pipeline_input_steps=[prepare_context], + ) + compute_constraint_metrics = PipelineStepResourceCollection( + "--compute-constraint-metrics", + output_resources={ + "constraint_metrics_ht": get_constraint_metrics_dataset( + custom_vep_annotation, **common_params, path_post_fix=path_post_fix + ) + }, + pipeline_input_steps=[ + aggregate_by_constraint_groups, + compute_gene_quality_metrics_step, + ], + ) + prepare_release = PipelineStepResourceCollection( + "--prepare-release", + output_resources={ + "release_ht": get_release_constraint_ht(version=version), + }, + pipeline_input_steps=[compute_constraint_metrics], + ) + export_release_tsv = PipelineStepResourceCollection( + "--export-release-tsv", + output_resources={ + "release_tsv": get_release_constraint_tsv_path(version=version), + "release_downsampling_tsv": get_release_downsampling_tsv_path( + version=version + ), + }, + pipeline_input_steps=[prepare_release], + ) + + # Add all steps to the constraint pipeline resource collection. + constraint_pipeline.add_steps( + { + "prepare_context": prepare_context, + "preprocess_data": preprocess_data, + "calculate_gerp_cutoffs": calculate_gerp_cutoffs, + "calculate_mutation_rate": calculate_mutation_rate, + "create_training_set": create_training_set, + "build_models": build_models, + "apply_models_per_variant": apply_models_per_variant, + "aggregate_per_variant_expected": aggregate_per_variant_expected, + "aggregate_by_constraint_groups": aggregate_by_constraint_groups, + "compute_gene_quality_metrics": compute_gene_quality_metrics_step, + "compute_constraint_metrics": compute_constraint_metrics, + "prepare_release": prepare_release, + "export_release_tsv": export_release_tsv, + } + ) + + return constraint_pipeline diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 9a493b8f..2bf4fd0e 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -31,7 +31,9 @@ from gnomad_constraint.resources.constants import ( ADJ_FREQ_META, + AGGREGATE_SUM_FIELDS, CALIBRATION_GROUPING, + CLASSIC_LOF_ANNOTATIONS, CONSTRAINT_GRANULARITIES, COVERAGE_CUTOFF, GENCODE_FIELD_RENAMES, @@ -397,7 +399,6 @@ def get_exomes_observed_and_possible( exomes_freq=exomes_freq_expr, **hl.or_missing( hl.is_defined(exomes_coverage_expr), - # & hl.or_else(hl.len(exomes_filter_expr) == 0, True), single_variant_observed_and_possible_expr(exomes_freq_expr, max_af=max_af), ), ) @@ -806,17 +807,12 @@ def create_per_variant_expected_ht( ) ) - # TODO: Check that this is needed - # tmp_path = new_temp_file(prefix="constraint", extension="ht") - # ht.drop(*calibrate_mu_fields).write(tmp_path) - - return ht.drop(*calibrate_mu_fields) # hl.read_table(tmp_path, _n_partitions=2000) + return ht.drop(*calibrate_mu_fields) def aggregate_per_variant_expected_ht( ht, include_mu_annotations_in_grouping: bool = False, - max_array_size: int = 200, ): """ Aggregate the per-variant expected Table. @@ -830,7 +826,6 @@ def aggregate_per_variant_expected_ht( :param ht: Table returned by `create_per_variant_expected_ht`. :param include_mu_annotations_in_grouping: Whether to include the mutation rate key annotations in the grouping. Default is False. - :param max_array_size: Maximum array size before batching. Default is 500. :return: Table with the observed and expected counts. """ groupings = [ @@ -841,68 +836,16 @@ def aggregate_per_variant_expected_ht( if g not in MU_GROUPING ], ] - aggregate_fields_to_sum = [ - "mu_snp", - "mu", - "observed_variants", - "possible_variants", - "predicted_proportion_observed", - "coverage_correction", - "expected_variants", - ] - if "calibrate_mu" in ht.row: ht = ht.annotate(**ht.calibrate_mu) ht = ht.filter(hl.set(CSQ_CODING).contains(ht.annotation)) - ht = ht.key_by() - ht = ht.select(*groupings, *aggregate_fields_to_sum) - # ht = ht.naive_coalesce(1000).checkpoint(new_temp_file("pre_aggregation", "ht")) + ht = ht.key_by().select(*groupings, *AGGREGATE_SUM_FIELDS) ht = ht.checkpoint(new_temp_file("pre_aggregation", "ht")) ht = ht.group_by(*groupings).aggregate(**aggregate_expected_variants_expr(ht)) - - # Check array lengths to determine if we need to batch. - arrays = [ - f for f in aggregate_fields_to_sum if isinstance(ht[f], hl.ArrayExpression) - ] - array_length = len(ht.filter(hl.is_defined(ht[arrays[0]]))[arrays[0]].take(1)[0]) - logger.info(f"Array length: {array_length}") - - if not array_length > max_array_size: - # No large arrays, use standard aggregation. - ht = ht.group_by(*groupings).aggregate(**aggregate_expected_variants_expr(ht)) - else: - # Calculate number of batches needed. - num_batches = (array_length + max_array_size - 1) // max_array_size - - batches = [] - for i in range(num_batches): - start_idx = i * max_array_size - end_idx = min((i + 1) * max_array_size, array_length) - logger.info( - f"Processing batch {i+1}/{num_batches} (indices {start_idx}:{end_idx})" - ) - - _ht = ht.annotate( - **{f: ht[f][start_idx:end_idx] for f in arrays} - ).checkpoint(new_temp_file(f"batch_{i}", "ht")) - batches.append( - _ht.group_by(*groupings) - .aggregate( - **aggregate_expected_variants_expr( - _ht, fields_to_sum=aggregate_fields_to_sum if i == 0 else arrays - ) - ) - .checkpoint(new_temp_file(f"batch_{i}.agg", "ht")) - ) - ht = batches[0] - batches = [_ht[ht.key] for _ht in batches[1:]] - ht = ht.annotate( - **{f: hl.flatten([ht[f]] + [_ht[f] for _ht in batches]) for f in arrays} - ) - ht = ht.checkpoint(new_temp_file("post_aggregation", "ht")) + return ht.naive_coalesce(1000) @@ -1212,11 +1155,7 @@ def compute_oe_upper_percentile_thresholds( def build_constraint_consequence_groups( csq_expr: hl.expr.ArrayExpression, lof_modifier_expr: hl.expr.StringExpression, - classic_lof_annotations: Tuple = ( - "stop_gained", - "splice_donor_variant", - "splice_acceptor_variant", - ), + classic_lof_annotations: Tuple = CLASSIC_LOF_ANNOTATIONS, additional_groupings: Dict[str, Dict[str, hl.expr.BooleanExpression]] = None, additional_grouping_combinations: List[List[str]] = None, ) -> Tuple[List[hl.expr.BooleanExpression], List[Dict[str, str]]]: @@ -1317,11 +1256,7 @@ def convert_multi_array_to_array_of_structs( def aggregate_by_constraint_groups( ht: hl.Table, keys: Tuple = ("gene", "transcript", "canonical"), - classic_lof_annotations: Tuple = ( - "stop_gained", - "splice_donor_variant", - "splice_acceptor_variant", - ), + classic_lof_annotations: Tuple = CLASSIC_LOF_ANNOTATIONS, additional_groupings: Dict[str, Dict[str, hl.expr.BooleanExpression]] = None, additional_grouping_combinations: List[List[str]] = None, ) -> hl.Table: @@ -1430,16 +1365,14 @@ def gamma_ci( :param alpha: Significance level for the confidence interval. Default is 0.05. :return: Upper bound of the OE confidence interval """ - # Calculate shape and scale parameters for Gamma distribution + # Calculate shape and scale parameters for Gamma distribution. shape = obs + hl.literal(1.0) - scale = divide_null( - hl.literal(1.0), exp - ) # Use divide_null to handle division by zero + # Use divide_null to handle division by zero. + scale = divide_null(hl.literal(1.0), exp) p = hl.literal(1.0 - alpha) - # Use the built-in qgamma function from the custom Hail wheel - # divide_null will return null if exp is 0, making the result null as well - + # Use the built-in qgamma function from the custom Hail wheel. + # divide_null will return null if exp is 0, making the result null as well. return hl.struct( lower=hl.qgamma(hl.literal(alpha), shape, scale), upper=hl.qgamma(p, shape, scale), @@ -1577,110 +1510,54 @@ def compute_gene_quality_metrics( return ht.key_by("transcript") -def compute_constraint_metrics( +def _annotate_oe_ci_z( ht: hl.Table, - gencode_ht: hl.Table, - gene_quality_metrics_ht: hl.Table, - expected_values: Optional[Dict[str, float]] = None, - min_diff_convergence: float = 0.001, - raw_z_outlier_threshold_lower_lof: float = -8.0, - raw_z_outlier_threshold_lower_missense: float = -8.0, - raw_z_outlier_threshold_lower_syn: float = -8.0, - raw_z_outlier_threshold_upper_syn: float = 8.0, - use_mane_select_over_canonical: bool = True, + z_thresholds: Dict[str, Tuple[Optional[float], Optional[float]]], ) -> hl.Table: """ - Compute the pLI scores, observed:expected ratio, 90% confidence interval around the observed:expected ratio, and z scores for synonymous variants, missense variants, and predicted loss-of-function (pLoF) variants. + Annotate constraint groups with OE ratio, confidence intervals, z-scores, and flags. - .. note:: - - The following annotations should be present in `ht`: + For each constraint group's ``oe_info`` entries, adds: - - modifier - - annotation - - observed_variants - - mu - - possible_variants - - expected_variants - - expected_variants_{pop} (if `pops` is specified) - - downsampling_counts_{pop} (if `pops` is specified) + - ``oe`` — observed / expected ratio. + - ``oe_ci_discretized_poisson`` — discretized Poisson CI. + - ``oe_ci_gamma`` — gamma-distribution CI. + - ``z_raw`` — raw z-score. - :param ht: Input Table with the number of expected variants (output of - `get_proportion_observed()`). - :param keys: The keys of the output Table, defaults to ('gene', 'transcript', - 'canonical'). - :param classic_lof_annotations: Classic LoF Annotations used to filter the input - Table. Default is {"stop_gained", "splice_donor_variant", - "splice_acceptor_variant"}. - :param expected_values: Dictionary containing the expected values for 'Null', - 'Rec', and 'LI' to use as starting values. - :param min_diff_convergence: Minimum iteration change in LI to consider the EM - model convergence criteria as met. Default is 0.001. - :param raw_z_outlier_threshold_lower_lof: Value at which the raw z-score is - considered an outlier for lof variants. Values below this threshold will be - considered outliers. Default is -8.0. - :param raw_z_outlier_threshold_lower_missense: Value at which the raw z-score is - considered an outlier for missense variants. Values below this threshold will - be considered outliers. Default is -8.0. - :param raw_z_outlier_threshold_lower_syn: Lower value at which the raw z-score is - considered an outlier for synonymous variants. Values below this threshold will - be considered outliers. Default is -8.0. - :param raw_z_outlier_threshold_upper_syn: Upper value at which the raw z-score is - considered an outlier for synonymous variants. Values above this threshold will - be considered outliers. Default is 8.0. - :param use_mane_select_over_canonical: Use MANE Select rather than canonical - transcripts for filtering the Table when determining ranks for the lof oe - upper confidence interval. If a gene does not have a MANE Select transcript, - the canonical transcript (if available) will be used instead. Default is True. - :param gencode_ht: Table containing GENCODE annotations. - :param gene_quality_metrics_ht: Table keyed by transcript with - ``gene_quality_metrics`` and ``gene_flags`` fields (output of - :func:`compute_gene_quality_metrics`). - :return: Table with pLI scores, observed:expected ratio, confidence interval of the - observed:expected ratio, and z scores. - """ - - def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: - """ - Add oe, oe_ci, and z_raw to the oe_info struct. - - :param oe_info: Struct containing the observed and expected variants. - :return: Struct containing oe, oe_ci, and z_raw. - """ - obs = oe_info.observed_variants - exp = oe_info.expected_variants - return oe_info.annotate( - oe=divide_null(obs, exp), - oe_ci_discretized_poisson=oe_confidence_interval(obs, exp), - oe_ci_gamma=calculate_oe_confidence_interval( - obs, exp, oe_upper_method="gamma" - ), - z_raw=calculate_raw_z_score(obs, exp), - ) + Then adds per-group ``flags`` based on z-score outlier thresholds. - # Annotate with the observed:expected ratio, 95% confidence interval around the - # observed:expected ratio, and z scores for each constraint group. + :param ht: Table with ``constraint_groups`` array. + :param z_thresholds: Mapping from constraint category (``"lof"``, ``"mis"``, + ``"syn"``) to ``(lower, upper)`` raw z-score outlier thresholds. + :return: Table with OE, CI, z-score, and flag annotations. + """ ht = ht.annotate( constraint_groups=ht.constraint_groups.map( lambda x: x.annotate( - oe_info=x.oe_info.map(lambda oe_info: _add_oe_ci_z(oe_info)) + oe_info=x.oe_info.map( + lambda oe_info: oe_info.annotate( + oe=divide_null( + oe_info.observed_variants, oe_info.expected_variants + ), + oe_ci_discretized_poisson=oe_confidence_interval( + oe_info.observed_variants, oe_info.expected_variants + ), + oe_ci_gamma=calculate_oe_confidence_interval( + oe_info.observed_variants, + oe_info.expected_variants, + oe_upper_method="gamma", + ), + z_raw=calculate_raw_z_score( + oe_info.observed_variants, oe_info.expected_variants + ), + ) + ) ) ) ) - z_threshold = { - "lof": (raw_z_outlier_threshold_lower_lof, None), - "mis": (raw_z_outlier_threshold_lower_missense, None), - "syn": ( - raw_z_outlier_threshold_lower_syn, - raw_z_outlier_threshold_upper_syn, - ), - } meta = hl.eval(ht.constraint_group_meta) freq_meta = hl.eval(ht.exomes_freq_meta) - syn_idx = meta.index({"csq_set": "syn"}) - mis_idx = meta.index({"csq_set": "mis"}) - lof_idx = meta.index({"lof": "hc"}) all_freq_idx = freq_meta.index(ADJ_FREQ_META) ht = ht.annotate( constraint_groups=[ @@ -1689,11 +1566,11 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: get_constraint_flags( ht.constraint_groups[i].oe_info[all_freq_idx].expected_variants, ht.constraint_groups[i].oe_info[all_freq_idx].z_raw, - z_threshold.get( + z_thresholds.get( "lof" if m.get("lof") else m.get("csq_set", "None"), (None, None), )[0], - z_threshold.get( + z_thresholds.get( "lof" if m.get("lof") else m.get("csq_set", "None"), (None, None), )[1], @@ -1703,9 +1580,29 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: ) for i, m in enumerate(meta) ] - ).checkpoint(new_temp_file("constraint_metrics.oe.oe_ci.z_raw", "ht")) + ) + + return ht + + +def _compute_z_scores(ht: hl.Table) -> hl.Table: + """ + Compute normalized z-scores and union per-group constraint flags. + + Computes the standard deviation of raw z-scores (stored as a global), normalizes + each group's raw z-score by its standard deviation, and unions the syn, mis, and + lof flags into a single ``constraint_flags`` set. + + :param ht: Table output by :func:`_annotate_oe_ci_z`. + :return: Table with ``z_score`` and ``constraint_flags`` annotations. + """ + meta = hl.eval(ht.constraint_group_meta) + freq_meta = hl.eval(ht.exomes_freq_meta) + syn_idx = meta.index({"csq_set": "syn"}) + mis_idx = meta.index({"csq_set": "mis"}) + lof_idx = meta.index({"lof": "hc"}) + all_freq_idx = freq_meta.index(ADJ_FREQ_META) - # Add z-score 'sd' annotation to globals. ht = ht.annotate_globals( sd_raw_z=ht.aggregate( hl.agg.filter( @@ -1722,7 +1619,6 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: ) ) - # Compute z-score from raw z-score and standard deviations. ht = ht.annotate( constraint_groups=hl.map( lambda x, sd_raw_z: x.annotate( @@ -1736,16 +1632,30 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: | ht.constraint_groups[mis_idx].flags | ht.constraint_groups[lof_idx].flags ), - ).checkpoint(new_temp_file("constraint_metrics.oe.oe_ci.z_raw.flags", "ht")) - - # Add a rank and decile of the upper confidence interval for MANE Select or - # canonical ensembl transcripts. - ht = add_oe_upper_rank_and_decile(ht, use_mane_select_over_canonical).checkpoint( - new_temp_file("constraint_metrics.oe.oe_ci.z_raw.flags.rank_and_decile", "ht") ) - # Compute OE upper CI percentile thresholds and annotate bins. - # Look up constraint group indices for each metric from the global meta. + return ht + + +def _compute_percentile_bins( + ht: hl.Table, + use_mane_select_over_canonical: bool = True, +) -> hl.Table: + """ + Add OE upper CI rank, decile, and percentile bin annotations. + + Adds rank and decile annotations via :func:`add_oe_upper_rank_and_decile`, + then computes percentile thresholds across all granularities defined in + ``CONSTRAINT_GRANULARITIES`` and annotates bins via + :func:`annotate_constraint_percentile_bins`. + + :param ht: Table output by :func:`_compute_z_scores`. + :param use_mane_select_over_canonical: Use MANE Select rather than canonical + transcripts for filtering when determining ranks. Default is True. + :return: Table with rank, decile, and percentile bin annotations. + """ + ht = add_oe_upper_rank_and_decile(ht, use_mane_select_over_canonical) + meta = hl.eval(ht.constraint_group_meta) metric_group_idx = { "syn": next(i for i, m in enumerate(meta) if m == {"csq_set": "syn"}), @@ -1754,7 +1664,6 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: } outlier_expr = ht.constraint_flags.length() > 0 - # Combine all granularity boundary quantiles for a single aggregate pass per metric. all_qs = [] gran_slices: Dict[str, slice] = {} for gran_name, bins in CONSTRAINT_GRANULARITIES.items(): @@ -1763,7 +1672,6 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: all_qs.extend(b / n_bins * 100 for b in bins) gran_slices[gran_name] = slice(start, len(all_qs)) - # Call once per metric and assemble thresholds dict. thresholds = {} for metric, idx in metric_group_idx.items(): vals = compute_oe_upper_percentile_thresholds( @@ -1776,14 +1684,34 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: for gran_name, sl in gran_slices.items(): thresholds[(gran_name, metric)] = list(vals[sl]) - ht = annotate_constraint_percentile_bins(ht, thresholds, metric_group_idx) + return annotate_constraint_percentile_bins(ht, thresholds, metric_group_idx) + + +def _compute_pli_scores( + ht: hl.Table, + expected_values: Optional[Dict[str, float]] = None, + min_diff_convergence: float = 0.001, +) -> hl.Table: + """ + Compute pLI, pNull, and pRec scores for the HC LoF constraint group. - # Compute the observed:expected ratio. + :param ht: Table output by :func:`_compute_percentile_bins`. + :param expected_values: Dictionary containing the expected OE values for 'Null', + 'Rec', and 'LI' to use as starting values. Default is ``PLI_EXPECTED_VALUES``. + :param min_diff_convergence: Minimum iteration change in LI to consider the EM + model convergence criteria as met. Default is 0.001. + :return: Table with pLI, pNull, and pRec annotations. + """ if expected_values is None: expected_values = PLI_EXPECTED_VALUES + meta = hl.eval(ht.constraint_group_meta) + freq_meta = hl.eval(ht.exomes_freq_meta) + lof_idx = meta.index({"lof": "hc"}) + all_freq_idx = freq_meta.index(ADJ_FREQ_META) + hc_lof_expr = ht.constraint_groups[lof_idx].oe_info[all_freq_idx] - ht = ht.annotate( + return ht.annotate( **compute_pli( ht, obs_expr=hc_lof_expr.observed_variants, @@ -1791,12 +1719,89 @@ def _add_oe_ci_z(oe_info: hl.expr.StructExpression) -> hl.expr.StructExpression: expected_values=expected_values, min_diff_convergence=min_diff_convergence, ) - ).checkpoint( - new_temp_file( - "constraint_metrics.oe.oe_ci.z_raw.flags.rank_and_decile.pli", "ht" - ) ) + +def compute_constraint_metrics( + ht: hl.Table, + gencode_ht: hl.Table, + gene_quality_metrics_ht: hl.Table, + expected_values: Optional[Dict[str, float]] = None, + min_diff_convergence: float = 0.001, + raw_z_outlier_threshold_lower_lof: float = -8.0, + raw_z_outlier_threshold_lower_missense: float = -8.0, + raw_z_outlier_threshold_lower_syn: float = -8.0, + raw_z_outlier_threshold_upper_syn: float = 8.0, + use_mane_select_over_canonical: bool = True, +) -> hl.Table: + """ + Compute constraint metrics for synonymous, missense, and pLoF variants. + + Orchestrates the following steps: + + 1. Annotate OE ratios, confidence intervals, raw z-scores, and per-group flags + (:func:`_annotate_oe_ci_z`). + 2. Normalize z-scores and union constraint flags (:func:`_compute_z_scores`). + 3. Add OE upper CI rank, decile, and percentile bins + (:func:`_compute_percentile_bins`). + 4. Compute pLI / pNull / pRec scores (:func:`_compute_pli_scores`). + 5. Annotate with gene quality metrics and GENCODE transcript annotations. + + .. note:: + + The following annotations should be present in `ht`: + + - modifier + - annotation + - observed_variants + - mu + - possible_variants + - expected_variants + + :param ht: Input Table with the number of expected variants (output of + ``aggregate_by_constraint_groups``). + :param gencode_ht: Table containing GENCODE annotations. + :param gene_quality_metrics_ht: Table keyed by transcript with + ``gene_quality_metrics`` and ``gene_flags`` fields (output of + :func:`compute_gene_quality_metrics`). + :param expected_values: Dictionary containing the expected OE values for 'Null', + 'Rec', and 'LI' to use as starting values. + :param min_diff_convergence: Minimum iteration change in LI to consider the EM + model convergence criteria as met. Default is 0.001. + :param raw_z_outlier_threshold_lower_lof: Lower raw z-score outlier threshold for + LoF variants. Default is -8.0. + :param raw_z_outlier_threshold_lower_missense: Lower raw z-score outlier threshold + for missense variants. Default is -8.0. + :param raw_z_outlier_threshold_lower_syn: Lower raw z-score outlier threshold for + synonymous variants. Default is -8.0. + :param raw_z_outlier_threshold_upper_syn: Upper raw z-score outlier threshold for + synonymous variants. Default is 8.0. + :param use_mane_select_over_canonical: Use MANE Select rather than canonical + transcripts for filtering when determining ranks. Default is True. + :return: Table with pLI scores, OE ratios, confidence intervals, z-scores, + percentile bins, gene quality metrics, and GENCODE annotations. + """ + z_thresholds = { + "lof": (raw_z_outlier_threshold_lower_lof, None), + "mis": (raw_z_outlier_threshold_lower_missense, None), + "syn": ( + raw_z_outlier_threshold_lower_syn, + raw_z_outlier_threshold_upper_syn, + ), + } + + ht = _annotate_oe_ci_z(ht, z_thresholds) + ht = ht.checkpoint(new_temp_file("constraint_metrics.oe_ci_z", "ht")) + + ht = _compute_z_scores(ht) + ht = ht.checkpoint(new_temp_file("constraint_metrics.z_scores", "ht")) + + ht = _compute_percentile_bins(ht, use_mane_select_over_canonical) + ht = ht.checkpoint(new_temp_file("constraint_metrics.percentile_bins", "ht")) + + ht = _compute_pli_scores(ht, expected_values, min_diff_convergence) + ht = ht.checkpoint(new_temp_file("constraint_metrics.pli", "ht")) + # Add per-transcript gene quality metrics and flags. ht = ht.annotate(**gene_quality_metrics_ht[ht.transcript]) @@ -2155,5 +2160,3 @@ def _bin_expr( } ) ) - - return ht.select(**flat) From 82ec428010734fecf54afd44cd331cbb9edde147 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 18 Mar 2026 07:30:58 -0600 Subject: [PATCH 22/38] Refactor constraint utility functions and streamline imports. Remove unused `print_global_struct` function and replace `aggregate_expected_variants_expr` with `aggregate_constraint_metrics_expr` in the aggregation process. Introduce new utility functions for better organization and clarity. --- .../pipeline/constraint_pipeline.py | 4 +- gnomad_constraint/utils/constraint.py | 399 ++---------------- 2 files changed, 41 insertions(+), 362 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 2f13c433..c4842d31 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -29,8 +29,10 @@ from gnomad.utils.constraint import ( assemble_constraint_context_ht, build_models, + calculate_gerp_cutoffs, explode_downsamplings_oe, ) +from gnomad.utils.file_utils import print_global_struct from gnomad.utils.reference_genome import get_reference_genome from gnomad.utils.vep import update_loftee_end_trunc_filter from gnomad_qc.resource_utils import PipelineResourceCollection @@ -45,7 +47,6 @@ from gnomad_constraint.utils.constraint import ( aggregate_by_constraint_groups, aggregate_per_variant_expected_ht, - calculate_gerp_cutoffs, calculate_mu_by_downsampling, compute_constraint_metrics, compute_gene_quality_metrics, @@ -54,7 +55,6 @@ flatten_release_ht, prepare_ht_for_constraint_calculations, prepare_release_ht, - print_global_struct, ) logging.basicConfig( diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 2bf4fd0e..70625441 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -1,32 +1,40 @@ """Script containing utility functions used in the constraint pipeline.""" -import functools import logging -import operator from typing import Dict, List, Optional, Tuple, Union import hail as hl -import numpy as np -from gnomad.assessment.summary_stats import generate_filter_combinations from gnomad.resources.grch38.gnomad import DOWNSAMPLINGS from gnomad.utils.constraint import ( add_gencode_transcript_annotations, - aggregate_expected_variants_expr, + aggregate_constraint_metrics_expr, annotate_exploded_vep_for_constraint_groupings, annotate_mutation_type, annotate_with_mu, apply_models, + build_constraint_consequence_groups, + calculate_gerp_cutoffs, calculate_raw_z_score, calculate_raw_z_score_sd, calibration_model_group_expr, + compute_oe_upper_percentile_thresholds, compute_pli, count_observed_and_possible_by_group, get_constraint_flags, oe_confidence_interval, + rank_and_assign_bins, single_variant_observed_and_possible_expr, ) +from gnomad.utils.file_utils import ( + convert_multi_array_to_array_of_structs, + print_global_struct, +) from gnomad.utils.filtering import add_filters_expr -from gnomad.utils.vep import CSQ_CODING, filter_vep_transcript_csqs_expr +from gnomad.utils.vep import ( + CSQ_CODING, + filter_vep_transcript_csqs_expr, + mane_select_over_canonical_filter_expr, +) from hail.utils.misc import divide_null, new_temp_file from gnomad_constraint.resources.constants import ( @@ -483,36 +491,6 @@ def get_build_calibration_model_annotation( return hl.or_missing(syn_csq_expr.length() > 0, build_expr) -# TODO: We don't really need this, I just found int helpful to look over the -# chosen parameters. -def print_global_struct(t: Union[hl.Table, hl.Struct, hl.StructExpression]) -> None: - """ - Print the global struct. - - :param t: Table with globals or globals struct to print. - :return: None - """ - if isinstance(t, hl.Table): - t = t.globals - if isinstance(t, hl.StructExpression): - t = hl.eval(t) - - def _get_pretty_print_globals(global_struct: hl.Struct, level: int = 1) -> str: - output = "" - level_tab = "".join([" "] * level) - for k, v in global_struct.items(): - if isinstance(v, hl.Struct): - v = f"\n{_get_pretty_print_globals(v, level + 1)}" - - output += f"{level_tab}{k}: {v}\n" - - return output - - logger.info( - "\nThe following parameters were used: \n%s", _get_pretty_print_globals(t) - ) - - def prepare_ht_for_constraint_calculations( ht: hl.Table, exome_coverage_metric: str = "median", @@ -843,7 +821,7 @@ def aggregate_per_variant_expected_ht( ht = ht.key_by().select(*groupings, *AGGREGATE_SUM_FIELDS) ht = ht.checkpoint(new_temp_file("pre_aggregation", "ht")) - ht = ht.group_by(*groupings).aggregate(**aggregate_expected_variants_expr(ht)) + ht = ht.group_by(*groupings).aggregate(**aggregate_constraint_metrics_expr(ht)) ht = ht.checkpoint(new_temp_file("post_aggregation", "ht")) return ht.naive_coalesce(1000) @@ -907,38 +885,6 @@ def calculate_mu_by_downsampling( return annotate_mutation_type(ht) -def mane_select_over_canonical_filter_expr(ht: hl.Table) -> hl.Table: - """ - Filter to MANE Select over canonical transcripts. - - Filter to only ensembl transcripts of the specified transcript filter. If MANE - select is specified, and a gene does not have a MANE select transcript, use - canonical instead. - - .. note:: - - In VEP 105 (used for gnomAD v4), all MANE Select transcripts are also - annotated as canonical. As a result, this function produces the same set of - transcripts as a simple canonical filter for v4 data. - - :param ht: Table with the MANE Select and canonical annotations. - :return: Table filtered to MANE Select over canonical transcripts. - """ - genes = ht.group_by(ht.gene_id).aggregate( - mane_present=hl.agg.any(ht.mane_select), - canonical_present=hl.agg.any(ht.canonical), - ) - genes = genes.annotate( - only_canonical=~(genes.mane_present) & (genes.canonical_present) - ) - only_canonical_expr = genes[ht.gene_id].only_canonical - mane_present_expr = genes[ht.gene_id].mane_present - - return (ht.transcript.startswith("ENST")) & ( - (mane_present_expr & ht.mane_select) | (only_canonical_expr & ht.canonical) - ) - - def get_transcript_filter_expr( ht: hl.Table, use_mane_select_over_canonical: bool = True, @@ -947,6 +893,9 @@ def get_transcript_filter_expr( """ Return a filter expression for selecting one representative transcript per gene. + Operates on an exploded, transcript-keyed table (one row per gene/transcript + pair) — not on VEP ``transcript_consequences`` arrays. + :param ht: Table with ``transcript``, ``mane_select``, ``canonical``, and ``gene_id`` annotations. :param use_mane_select_over_canonical: When ``True`` (default), prefer MANE @@ -960,49 +909,13 @@ def get_transcript_filter_expr( if mane_select_only: return ht.transcript.startswith("ENST") & ht.mane_select elif use_mane_select_over_canonical: - return mane_select_over_canonical_filter_expr(ht) + return mane_select_over_canonical_filter_expr( + ht.transcript, ht.mane_select, ht.canonical, ht.gene_id + ) else: return ht.transcript.startswith("ENST") & ht.canonical -# TODO: Move to gnomad_methods? -def get_rank_and_bins( - value_expr: hl.expr.Float64Expression, - bin_granularities: Optional[Dict[str, int]] = None, -) -> hl.StructExpression: - """Rank rows by a numeric expression and assign bin labels. - - Rows are ordered ascending by ``value_expr``. Each row is assigned a - 0-based ``rank`` and a ``bin_{name}`` field for every entry in - ``bin_granularities``, computed as - ``hl.int(rank * multiplier / n_transcripts)``. - - :param value_expr: Numeric expression to rank by (ascending). - :param bin_granularities: Mapping of bin name to multiplier. Each entry - produces a ``bin_{name}`` field. Default is - ``{"percentile": 100, "decile": 10, "sextile": 6}``. - :return: Struct with ``rank`` and ``bin_{name}`` fields for each entry in - ``bin_granularities``. - """ - if bin_granularities is None: - bin_granularities = {"percentile": 100, "decile": 10, "sextile": 6} - - ht = value_expr._indices.source - source_key = list(ht.key) - n_transcripts = ht.count() - ranked_ht = ht.select(_=value_expr).order_by("_").add_index("rank") - ranked_ht = ranked_ht.select( - *source_key, - "rank", - **{ - f"bin_{name}": hl.int(ranked_ht.rank * multiplier / n_transcripts) - for name, multiplier in bin_granularities.items() - }, - ).cache() - - return ranked_ht.key_by(*source_key).cache()[ht.key] - - def add_oe_upper_rank_and_decile( ht: hl.Table, use_mane_select_over_canonical: bool = True, @@ -1064,7 +977,9 @@ def add_oe_upper_rank_and_decile( oe_ci_upper=[ hl.struct( **{ - ci: get_rank_and_bins(ms_ht.oe_ci_upper[i][ci], bin_granularities) + ci: rank_and_assign_bins( + ms_ht.oe_ci_upper[i][ci], bin_granularities + ) for ci in ci_fields } ) @@ -1093,166 +1008,6 @@ def add_oe_upper_rank_and_decile( return ht -def compute_oe_upper_percentile_thresholds( - ht: hl.Table, - percentiles: List[float], - metric_expr: hl.expr.Float64Expression, - outlier_expr: hl.expr.BooleanExpression, - use_mane_select_over_canonical: bool = True, - mane_select_only: bool = False, - quantile_k: int = 1000, -) -> List[float]: - """ - Compute OE upper CI percentile thresholds for a single metric expression. - - Filters to a representative transcript set (controlled by - ``use_mane_select_over_canonical`` / ``mane_select_only``) and excludes - outlier transcripts, then computes approximate quantile thresholds at the - requested percentiles in a single aggregation pass. - - :param ht: Constraint metrics Table (output of ``compute_constraint_metrics``). - :param percentiles: Percentile values (0–100) at which to compute thresholds. - Pass a combined list across multiple granularities and slice the result to - avoid repeated aggregation passes. - :param metric_expr: Float expression for the metric to threshold (e.g., - ``ht.constraint_groups[i].oe_info[0].oe_ci_gamma.upper``). Must be - defined on ``ht``. - :param outlier_expr: Boolean expression that is ``True`` for transcripts to - exclude from the reference population (e.g., flagged transcripts). - :param use_mane_select_over_canonical: When ``True`` (default), prefer MANE - Select transcripts, falling back to canonical for genes without a MANE - Select entry. Ignored when ``mane_select_only`` is ``True``. - :param mane_select_only: When ``True``, restrict to ENST MANE Select - transcripts only, with no canonical fallback. Default is ``False``. - :param quantile_k: Accuracy parameter for - :func:`hail.expr.aggregators.approx_quantiles`. Default is 1000. - :return: List of float threshold values at the given percentiles. - """ - mane_filter_expr = get_transcript_filter_expr( - ht, use_mane_select_over_canonical, mane_select_only - ) - - qs = [p / 100.0 for p in percentiles] - filt = mane_filter_expr & hl.is_defined(metric_expr) & ~outlier_expr - - result = ht.aggregate( - hl.struct( - thresholds=hl.agg.filter( - filt, hl.agg.approx_quantiles(metric_expr, qs, k=quantile_k) - ), - n=hl.agg.count_where(filt), - ) - ) - logger.info( - "Computed percentile thresholds on %d transcripts.", - result.n, - ) - - return result.thresholds - - -# TODO: Move to gnomad_methods? -def build_constraint_consequence_groups( - csq_expr: hl.expr.ArrayExpression, - lof_modifier_expr: hl.expr.StringExpression, - classic_lof_annotations: Tuple = CLASSIC_LOF_ANNOTATIONS, - additional_groupings: Dict[str, Dict[str, hl.expr.BooleanExpression]] = None, - additional_grouping_combinations: List[List[str]] = None, -) -> Tuple[List[hl.expr.BooleanExpression], List[Dict[str, str]]]: - """ - Build constraint consequence groups. - - The function builds constraint groups based on the consequence expression and LoF - modifier expression. By default, the following groups are built: - - - csq_set: synonymous_variant, missense_variant - - lof: classic, hc_lc, classic_hc_lc, hc - - The resulting meta and cooresponding constraint group filters are: - - - {"csq_set": "syn"}: synonymous_variant - - {"csq_set": "mis"}: missense_variant - - {"lof": "classic"}: classic LoF annotations - - {"lof": "hc_lc"}: LoFTEE HC or LC - - {"lof": "classic_hc_lc"}: classic LoF annotations with LoFTEE HC or LC - - {"lof": "hc"}: LoF annotations with LoFTEE HC - - Additional groupings can be added to the constraint groups by specifying the - `additional_groupings` parameter, and grouping combinations can also be added - by specifying the `additional_grouping_combinations` parameter. - - :param csq_expr: Consequence expression. - :param lof_modifier_expr: LoF modifier expression. - :param classic_lof_annotations: Classic LoF Annotations used to filter the input - Table. Default is {"stop_gained", "splice_donor_variant", - "splice_acceptor_variant"}. - :param additional_groupings: Additional groupings to add to the constraint groups. - Default is None. - :param additional_grouping_combinations: Additional grouping combinations to add to - the constraint groups. Default is None. - :return: Tuple containing the constraint group filters and the meta. - """ - lof_classic_expr = hl.literal(set(classic_lof_annotations)).contains(csq_expr) - lof_hc_expr = lof_modifier_expr == "HC" - lof_hc_lc_expr = lof_hc_expr | (lof_modifier_expr == "LC") - mis_expr = csq_expr == "missense_variant" - annotation_dict = { - "csq_set": {"syn": csq_expr == "synonymous_variant", "mis": mis_expr}, - "lof": { - # Filter to classic LoF annotations. - "classic": lof_classic_expr, - # Filter to LOFTEE HC or LC. - "hc_lc": lof_hc_lc_expr, - # Filter to classic LoF annotations with LOFTEE HC or LC. - "classic_hc_lc": lof_classic_expr & lof_hc_lc_expr, - # Filter to LoF annotations with LOFTEE HC. - "hc": lof_hc_expr, - }, - } - - annotation_dict.update(additional_groupings or {}) - additional_grouping_combinations = additional_grouping_combinations or [] - - grouping_combinations = [["csq_set"], ["lof"]] - grouping_combinations.extend(additional_grouping_combinations) - - meta = generate_filter_combinations( - grouping_combinations, - {k: list(v.keys()) for k, v in annotation_dict.items()}, - ) - constraint_group_filters = [ - functools.reduce(operator.ior, [annotation_dict[k][v] for k, v in m.items()]) - for m in meta - ] - - return constraint_group_filters, meta - - -# TODO: Move to gnomad_methods? -def convert_multi_array_to_array_of_structs( - t: Union[hl.Table, hl.expr.StructExpression], - array_fields_to_combine: List[str], - new_array_field: str, -) -> hl.Table: - """ - Convert multiple arrays to an array of structs. - - :param t: Table or Struct to convert. - :param array_fields_to_combine: Array fields to combine. - :param new_array_field: Name of the new array field. - :return: Table with the array fields combined into an array of structs named - `new_array_field`. - """ - logger.warning("This function assumes that all arrays have the same length!") - return t.annotate( - **{ - new_array_field: hl.range(t[array_fields_to_combine[0]].length()).map( - lambda i: hl.struct(**{f: t[f][i] for f in array_fields_to_combine}) - ) - } - ).drop(*array_fields_to_combine) - - def aggregate_by_constraint_groups( ht: hl.Table, keys: Tuple = ("gene", "transcript", "canonical"), @@ -1307,7 +1062,7 @@ def aggregate_by_constraint_groups( # expected_variants for each constraint group. ht = ht.group_by(*keys).aggregate( constraint_groups=hl.agg.array_agg( - lambda f: hl.agg.filter(f, aggregate_expected_variants_expr(ht)), + lambda f: hl.agg.filter(f, aggregate_constraint_metrics_expr(ht)), ht.constraint_groups, ) ) @@ -1350,51 +1105,6 @@ def aggregate_by_constraint_groups( return ht -def gamma_ci( - obs: hl.expr.Int32Expression, - exp: hl.expr.Float64Expression, - alpha: float = 0.05, -) -> hl.expr.Float64Expression: - """ - Calculate the upper bound of the OE confidence interval using the Gamma distribution. - - This function uses the built-in qgamma function from the custom Hail wheel. - - :param obs: Observed count - :param exp: Expected count - :param alpha: Significance level for the confidence interval. Default is 0.05. - :return: Upper bound of the OE confidence interval - """ - # Calculate shape and scale parameters for Gamma distribution. - shape = obs + hl.literal(1.0) - # Use divide_null to handle division by zero. - scale = divide_null(hl.literal(1.0), exp) - p = hl.literal(1.0 - alpha) - - # Use the built-in qgamma function from the custom Hail wheel. - # divide_null will return null if exp is 0, making the result null as well. - return hl.struct( - lower=hl.qgamma(hl.literal(alpha), shape, scale), - upper=hl.qgamma(p, shape, scale), - ) - - -def calculate_oe_confidence_interval( - obs: hl.expr.Int32Expression, - exp: hl.expr.Float64Expression, - alpha: float = 0.05, -) -> hl.expr.StructExpression: - """Calculate the OE confidence interval using the Gamma distribution. - - :param obs: Observed count. - :param exp: Expected count. - :param alpha: Significance level for the confidence interval. Default is - 0.05. - :return: Struct with ``lower`` and ``upper`` bounds. - """ - return gamma_ci(obs, exp, alpha) - - def _compute_coverage_metrics( ht: hl.Table, gencode_cds_ht: hl.Table, @@ -1540,12 +1250,14 @@ def _annotate_oe_ci_z( oe_info.observed_variants, oe_info.expected_variants ), oe_ci_discretized_poisson=oe_confidence_interval( - oe_info.observed_variants, oe_info.expected_variants + oe_info.observed_variants, + oe_info.expected_variants, + method="poisson", ), - oe_ci_gamma=calculate_oe_confidence_interval( + oe_ci_gamma=oe_confidence_interval( oe_info.observed_variants, oe_info.expected_variants, - oe_upper_method="gamma", + method="gamma", ), z_raw=calculate_raw_z_score( oe_info.observed_variants, oe_info.expected_variants @@ -1664,13 +1376,13 @@ def _compute_percentile_bins( } outlier_expr = ht.constraint_flags.length() > 0 + gran_percentiles: Dict[str, List[float]] = {} all_qs = [] - gran_slices: Dict[str, slice] = {} for gran_name, bins in CONSTRAINT_GRANULARITIES.items(): n_bins = len(bins) + 1 - start = len(all_qs) - all_qs.extend(b / n_bins * 100 for b in bins) - gran_slices[gran_name] = slice(start, len(all_qs)) + pcts = [b / n_bins * 100 for b in bins] + gran_percentiles[gran_name] = pcts + all_qs.extend(pcts) thresholds = {} for metric, idx in metric_group_idx.items(): @@ -1679,10 +1391,12 @@ def _compute_percentile_bins( percentiles=all_qs, metric_expr=ht.constraint_groups[idx].oe_info[0].oe_ci_gamma.upper, outlier_expr=outlier_expr, - mane_select_only=True, + transcript_filter_expr=get_transcript_filter_expr( + ht, mane_select_only=True + ), ) - for gran_name, sl in gran_slices.items(): - thresholds[(gran_name, metric)] = list(vals[sl]) + for gran_name, pcts in gran_percentiles.items(): + thresholds[(gran_name, metric)] = [vals[p] for p in pcts] return annotate_constraint_percentile_bins(ht, thresholds, metric_group_idx) @@ -1811,41 +1525,6 @@ def compute_constraint_metrics( return ht -# TODO: Move to gnomad_methods? -def calculate_gerp_cutoffs(ht: hl.Table) -> Tuple[float, float]: - """ - Find GERP cutoffs determined by the 5% and 95% percentiles. - - :param ht: Input Table. - :return: Tuple containing values determining the 5-95th percentile of the GERP score. - """ - # Aggregate histogram of GERP values from -12.3 to 6.17 (-12.3 to 6.17 is the range - # of GERP values where 6.17 is the most conserved). - summary_hist = ht.aggregate(hl.struct(gerp=hl.agg.hist(ht.gerp, -12.3, 6.17, 100))) - - # Get cumulative sum of the hist array and add value of n_smaller to every value in - # the cumulative sum array. - cumulative_data = ( - np.cumsum(summary_hist.gerp.bin_freq) + summary_hist.gerp.n_smaller - ) - - # Append final value to the cumulative sum array (value added is last value of the - # array plus n_larger). - np.append(cumulative_data, [cumulative_data[-1] + summary_hist.gerp.n_larger]) - - # Get zip of (bin_edge, value in cumulative sum array divided by max value in - # cumulative sum array). - zipped = zip(summary_hist.gerp.bin_edges, cumulative_data / max(cumulative_data)) - - # Define lower and upper GERP cutoffs based on 5th and 95th percentiles. - cutoff_lower = list(filter(lambda i: i[1] > 0.05, zipped))[0][0] - - zipped = zip(summary_hist.gerp.bin_edges, cumulative_data / max(cumulative_data)) - cutoff_upper = list(filter(lambda i: i[1] < 0.95, zipped))[-1][0] - - return cutoff_lower, cutoff_upper - - def _restructure_release_rows( ht: hl.Table, field_names: List[str], From 81614d3ad3a6f001fb78698dc9cf4160822dbfc6 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 18 Mar 2026 08:49:18 -0600 Subject: [PATCH 23/38] Refactor constraint utility functions to improve clarity and organization. Replace deprecated functions with updated counterparts, streamline percentile threshold calculations, and enhance the handling of observed and possible variants in the constraint pipeline. --- gnomad_constraint/utils/constraint.py | 126 +++++++------------------- 1 file changed, 31 insertions(+), 95 deletions(-) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 70625441..caaa1488 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -1,29 +1,29 @@ """Script containing utility functions used in the constraint pipeline.""" import logging -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple import hail as hl from gnomad.resources.grch38.gnomad import DOWNSAMPLINGS from gnomad.utils.constraint import ( add_gencode_transcript_annotations, aggregate_constraint_metrics_expr, + annotate_bins_by_threshold, annotate_exploded_vep_for_constraint_groupings, annotate_mutation_type, annotate_with_mu, apply_models, build_constraint_consequence_groups, - calculate_gerp_cutoffs, calculate_raw_z_score, calculate_raw_z_score_sd, calibration_model_group_expr, - compute_oe_upper_percentile_thresholds, + compute_percentile_thresholds, compute_pli, count_observed_and_possible_by_group, get_constraint_flags, oe_confidence_interval, - rank_and_assign_bins, - single_variant_observed_and_possible_expr, + rank_array_element_metrics, + variant_observed_and_possible_expr, ) from gnomad.utils.file_utils import ( convert_multi_array_to_array_of_structs, @@ -252,7 +252,7 @@ def get_annotations_for_computing_mu( obs_pos_expr = hl.struct( genomes_freq=genomes_freq_expr, **hl.or_missing( - keep_expr, single_variant_observed_and_possible_expr(genomes_freq_expr) + keep_expr, variant_observed_and_possible_expr(genomes_freq_expr) ), ) obs_pos_globals = hl.struct( @@ -407,7 +407,7 @@ def get_exomes_observed_and_possible( exomes_freq=exomes_freq_expr, **hl.or_missing( hl.is_defined(exomes_coverage_expr), - single_variant_observed_and_possible_expr(exomes_freq_expr, max_af=max_af), + variant_observed_and_possible_expr(exomes_freq_expr, max_af=max_af), ), ) obs_pos_globals = hl.struct( @@ -925,6 +925,10 @@ def add_oe_upper_rank_and_decile( """ Compute the rank and decile of the oe upper confidence interval. + Thin wrapper around :func:`rank_array_element_metrics` that extracts the + discretized Poisson and gamma upper CI values from each constraint group's + first oe_info element. + :param ht: Table with the oe upper confidence interval. :param use_mane_select_over_canonical: Use MANE Select over canonical transcripts for ranking, falling back to canonical when MANE Select is absent for a gene. @@ -941,72 +945,20 @@ def add_oe_upper_rank_and_decile( ``bin_granularities``. Transcripts excluded from ranking have these fields set to missing. """ - # Add an integer index so re-keying after order_by is an O(1) lookup. - ht = ht.add_index("_idx").key_by("_idx").cache() - - total_count = ht.count() - ms_ht = ht.filter( - get_transcript_filter_expr(ht, use_mane_select_over_canonical, mane_select_only) - ) - - # Extract only the first freq group (adj/all-samples) per constraint group for - # ranking. ci_fields = ["discretized_poisson", "gamma"] - ms_ht = ( - ms_ht.select( - oe_ci_upper=ms_ht.constraint_groups.map( - lambda x: hl.struct( - **{ci: x.oe_info[0][f"oe_ci_{ci}"].upper for ci in ci_fields} - ) - ), - ) - .naive_coalesce(100) - .checkpoint(new_temp_file("oe_ci_upper.before_rank", "ht")) - ) - n_transcripts = ms_ht.count() - logger.info( - "Retaining %d out of %d transcripts to use for rank annotations.", - n_transcripts, - total_count, - ) - - # For each (constraint group, CI method), rank a minimal 2-column table and - # checkpoint it, then join all rank tables back in one pass. - num_constraint_groups = hl.eval(ms_ht.constraint_group_meta.length()) - ms_ht = ms_ht.annotate( - oe_ci_upper=[ - hl.struct( - **{ - ci: rank_and_assign_bins( - ms_ht.oe_ci_upper[i][ci], bin_granularities - ) - for ci in ci_fields - } - ) - for i in range(num_constraint_groups) - ] - ).cache() - # Annotate each constraint group with rank/bin fields at the group level (not - # inside oe_info, since all array elements must share the same struct schema). - # ht is already keyed by _idx, so ms_ht can be looked up directly. - ms_keyed = ms_ht[ht._idx] - ht = ht.annotate( - constraint_groups=hl.if_else( - hl.is_defined(ms_keyed.oe_ci_upper), - hl.map( - lambda g, r: g.annotate( - **{f"oe_ci_{ci}_rank": r[ci] for ci in ci_fields} - ), - ht.constraint_groups, - ms_keyed.oe_ci_upper, - ), - ht.constraint_groups, - ) + return rank_array_element_metrics( + ht, + array_field="constraint_groups", + element_value_fn=lambda x: { + f"oe_ci_{ci}": x.oe_info[0][f"oe_ci_{ci}"].upper for ci in ci_fields + }, + filter_fn=lambda t: get_transcript_filter_expr( + t, use_mane_select_over_canonical, mane_select_only + ), + bin_granularities=bin_granularities, ) - return ht - def aggregate_by_constraint_groups( ht: hl.Table, @@ -1386,7 +1338,7 @@ def _compute_percentile_bins( thresholds = {} for metric, idx in metric_group_idx.items(): - vals = compute_oe_upper_percentile_thresholds( + vals = compute_percentile_thresholds( ht, percentiles=all_qs, metric_expr=ht.constraint_groups[idx].oe_info[0].oe_ci_gamma.upper, @@ -1792,6 +1744,9 @@ def annotate_constraint_percentile_bins( """ Annotate each transcript with its percentile bin for all metric/granularity combinations. + Thin wrapper around :func:`annotate_bins_by_threshold` that extracts the + gamma upper CI value from each constraint group's first oe_info element. + Annotates ``constraint_bins.{granularity}.{metric}`` for each combination. Bin 0 is the most constrained (value below all thresholds); bin N equals the number of boundaries the value exceeds. @@ -1799,7 +1754,7 @@ def annotate_constraint_percentile_bins( :param ht: Constraint metrics Table with a ``constraint_groups`` array field. :param thresholds: Mapping of ``(granularity, metric)`` to an ordered list of threshold values, as produced by - :func:`compute_oe_upper_percentile_thresholds`. + :func:`compute_percentile_thresholds`. :param metric_group_idx: Mapping of metric name to its index in ``constraint_groups`` (e.g. ``{"lof": 5, "mis": 1, "syn": 0}``). :return: Annotated Table with an added ``constraint_bins`` struct field. @@ -1813,29 +1768,10 @@ def annotate_constraint_percentile_bins( metric: ht.constraint_groups[idx].oe_info[0].oe_ci_gamma.upper for metric, idx in metric_group_idx.items() } - miss = hl.missing(hl.tint32) - - def _bin_expr( - value_expr: hl.expr.Float64Expression, - threshold_list: List[float], - ) -> hl.expr.Int32Expression: - arr = hl.literal(threshold_list) - return hl.sum(arr.map(lambda t: hl.int(value_expr >= t))) - return ht.annotate( - constraint_bins=hl.struct( - **{ - gran: hl.struct( - **{ - metric: hl.if_else( - hl.is_defined(metric_exprs[metric]), - _bin_expr(metric_exprs[metric], thresholds[(gran, metric)]), - miss, - ) - for metric in metric_group_idx - } - ) - for gran in CONSTRAINT_GRANULARITIES - } - ) + return annotate_bins_by_threshold( + ht, + metric_exprs=metric_exprs, + thresholds=thresholds, + granularities=list(CONSTRAINT_GRANULARITIES), ) From d44caafd1458e26f2625b39e09a0cc542d093957 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 18 Mar 2026 09:17:07 -0600 Subject: [PATCH 24/38] Refactor constraint pipeline to enhance context preparation and filtering. Introduce new utility functions for filtering test data and reading adjacency data, while updating the context preparation process to streamline annotations for coverage, allele number, and site frequency spectrum. Remove deprecated functions and improve overall code organization. --- .../pipeline/constraint_pipeline.py | 168 +++--------------- gnomad_constraint/resources/constants.py | 7 + gnomad_constraint/resources/resource_utils.py | 53 ++++++ gnomad_constraint/utils/constraint.py | 89 ++++++++++ 4 files changed, 174 insertions(+), 143 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index c4842d31..f8b62241 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -23,19 +23,15 @@ import argparse import logging -from typing import List import hail as hl from gnomad.utils.constraint import ( - assemble_constraint_context_ht, build_models, calculate_gerp_cutoffs, explode_downsamplings_oe, ) from gnomad.utils.file_utils import print_global_struct from gnomad.utils.reference_genome import get_reference_genome -from gnomad.utils.vep import update_loftee_end_trunc_filter -from gnomad_qc.resource_utils import PipelineResourceCollection import gnomad_constraint.resources.resource_utils as constraint_res from gnomad_constraint.resources.constants import ( @@ -44,6 +40,7 @@ RELEASE_KEY_ORDER, VERSIONS, ) +from gnomad_constraint.resources.resource_utils import filter_for_test, get_adj_r_ht from gnomad_constraint.utils.constraint import ( aggregate_by_constraint_groups, aggregate_per_variant_expected_ht, @@ -53,6 +50,7 @@ create_per_variant_expected_ht, create_training_set, flatten_release_ht, + prepare_context_ht, prepare_ht_for_constraint_calculations, prepare_release_ht, ) @@ -65,143 +63,6 @@ logger.setLevel(logging.INFO) -def filter_for_test( - ht: hl.Table, - use_gene_list: bool = False, -) -> hl.Table: - """ - Filter `ht` to chr20, chrX, and chrY or a gene list for testing. - - :param ht: Table to filter. - :param use_gene_list: Whether to use a gene list for testing instead of all of - chr20, chrX, and chrY for testing. - :return: Filtered Table for testing. - """ - rg = get_reference_genome(ht.locus) - if use_gene_list: - if rg == "GRCh37": - keep_regions = [ - "20:49505585-49547958", # ADNP - "20:853296-896977", # ANGPT4 - "X:13752832-13787480", # OFD1 - "X:57313139-57515629", # FAAH2 - "Y:2803112-2850547", # ZFY - ] - else: - keep_regions = [ - "chr20:50888916-50931437", # ADNP - "chr20:869900-916334", # ANGPT4 - "chrX:13734743-13777955", # OFD1 - "chrX:57286706-57489193", # FAAH2 - "chrY:2935281-2982506", # ZFY - ] - keep = [hl.parse_locus_interval(c, reference_genome=rg) for c in keep_regions] - else: - keep = [ - hl.parse_locus_interval(c, reference_genome=rg) - for c in [rg.contigs[19], rg.x_contigs[0], rg.y_contigs[0]] - ] - logger.info("Filtering the context HT to chr20, chrX, and chrY for testing...") - - ht = hl.filter_intervals(ht, keep) - - return ht - - -def run_prepare_context( - resources: PipelineResourceCollection, - test: bool = False, - test_gene_list: bool = False, -) -> hl.Table: - """ - Annotate the context Table with coverage, AN, and frequency annotations. - - Uses `assemble_constraint_context_ht` to annotate the context Table with annotations - that are used in downstream steps of the constraint pipeline. - - :param resources: PipelineResourceCollection containing resources for the constraint - pipeline. - :param test: Whether to filter the context Table to only chr20, chrX, and chrY for - testing. - :param test_gene_list: Whether to filter the context Table to a gene list for - testing. - :return: Annotated context Table. - """ - # We use naive_coalesce on the context Table because it has a large number of - # partitions which caused some issues with Hail 0.2.133. 5000 partitions was a - # number that worked well for the context Table in the past. - ht = resources.context_ht.ht().naive_coalesce(5000) - - if test: - ht = filter_for_test(ht, use_gene_list=test_gene_list) - - def _build_ht_dict(ht_name: str, keep: List[str] = None): - dts = ["exomes", "genomes"] - hts = {d: getattr(resources, f"{d}_{ht_name}_ht").ht() for d in dts} - return {d: t.select(*keep) for d, t in hts.items()} if keep else hts - - # There was a bug in the GERP cutoffs used to filter transcripts with the - # "END_TRUNC" filter in the LOFTEE VEP plugin resulting in some transcripts - # being considered "HC" when they should have been "LC". We use the - # `update_loftee_end_trunc_filter` function to correct this issue. - ht = ht.annotate( - vep=ht.vep.annotate( - transcript_consequences=update_loftee_end_trunc_filter( - ht.vep.transcript_consequences - ) - ) - ) - ht = assemble_constraint_context_ht( - ht, - coverage_hts=_build_ht_dict("coverage"), - an_hts=_build_ht_dict("an"), - freq_hts=_build_ht_dict("sites", ["freq"]), - filter_hts=_build_ht_dict("sites", ["filters"]), - methylation_ht=resources.methylation_ht.ht(), - gerp_ht=constraint_res.get_gerp_ht(get_reference_genome(ht.locus).name), - transformation_funcs=None, - ) - - # Add annotation for exome coverage and genomic region (autosome/PAR, X non-PAR, - # Y non-PAR). - genomic_region_expr = ( - hl.case() - .when(ht.locus.in_autosome_or_par(), "autosome_or_par") - .when(ht.locus.in_x_nonpar(), "chrx_nonpar") - .when(ht.locus.in_y_nonpar(), "chry_nonpar") - .or_missing() - ) - - # Add annotation for SFS bin. - sfs_bin_cutoffs = [0, 1e-6, 2e-6, 4e-6, 2e-5, 5e-5, 5e-4, 5e-3, 0.5] - af_expr = ht.freq.exomes[0].AF - sfs_bin_expr = hl.case().when(hl.is_missing(af_expr), 0) - for i, af in enumerate(sfs_bin_cutoffs): - sfs_bin_expr = sfs_bin_expr.when(af_expr <= af, i) - - sfs_bin_expr = sfs_bin_expr.or_missing() - - adj_r_ht = hl.read_table( - "gs://gnomad/v4.1/constraint/resources/annotations/ht/adj_r_per_context_methyl_genome_1kb_autosome.agg.ht" - ) - - ht = ht.annotate( - coverage=hl.struct( - exomes=ht.coverage.exomes.select("mean", "median_approx"), - genomes=ht.coverage.genomes.select("mean", "median_approx"), - ), - AN=hl.struct( - exomes=ht.AN.exomes[0], - genomes=ht.AN.genomes[0], - ), - genomic_region=genomic_region_expr, - adj_r=adj_r_ht[ht.locus].adj_r[ht.context], - sfs_bin=sfs_bin_expr, - ) - - return ht - - def main(args): """Execute the constraint pipeline.""" hl.init( @@ -250,7 +111,28 @@ def main(args): ) res = resources.prepare_context res.check_resource_existence() - ht = run_prepare_context(res, test=test, test_gene_list=test_gene_list) + + # We use naive_coalesce on the context Table because it has a large + # number of partitions which caused issues with Hail 0.2.133. + ht = res.context_ht.ht().naive_coalesce(5000) + if test: + ht = filter_for_test(ht, use_gene_list=test_gene_list) + + dts = ["exomes", "genomes"] + ht = prepare_context_ht( + ht, + coverage_hts={d: getattr(res, f"{d}_coverage_ht").ht() for d in dts}, + an_hts={d: getattr(res, f"{d}_an_ht").ht() for d in dts}, + freq_hts={ + d: getattr(res, f"{d}_sites_ht").ht().select("freq") for d in dts + }, + filter_hts={ + d: getattr(res, f"{d}_sites_ht").ht().select("filters") for d in dts + }, + methylation_ht=res.methylation_ht.ht(), + gerp_ht=constraint_res.get_gerp_ht(get_reference_genome(ht.locus).name), + adj_r_ht=get_adj_r_ht(), + ) ht.write(res.annotated_context_ht.path, overwrite) logger.info("Done annotating the VEP context Table.") @@ -422,7 +304,7 @@ def main(args): ht = res.apply_ht.ht() aggregate_by_constraint_groups( ht, - keys=tuple([i for i in list(ht.key) if i in RELEASE_KEY_ORDER]), + keys=tuple(k for k in ht.key if k in RELEASE_KEY_ORDER), ).write(res.constraint_group_ht.path, overwrite=overwrite) hl._set_flags(use_new_shuffle=None) logger.info("Done with aggregating by constraint groups.") diff --git a/gnomad_constraint/resources/constants.py b/gnomad_constraint/resources/constants.py index ee6b36a6..f291b334 100644 --- a/gnomad_constraint/resources/constants.py +++ b/gnomad_constraint/resources/constants.py @@ -41,6 +41,13 @@ Abbreviations stand for: global (all populations), African-American/African, Latino, East Asian, Non-Finnish European, and South Asian. """ +SFS_BIN_CUTOFFS = (0, 1e-6, 2e-6, 4e-6, 2e-5, 5e-5, 5e-4, 5e-3, 0.5) +"""Allele frequency upper bounds defining site frequency spectrum bins. + +Variants with missing frequency are assigned bin 0. Otherwise, each variant is +assigned the index of the first cutoff its AF falls at or below. +""" + COVERAGE_CUTOFF = 40 """ Minimum median exome coverage differentiating high coverage sites from low coverage sites. diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 6607ac3e..cc4a3101 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -550,6 +550,59 @@ def get_checkpoint_path(name: str, **kwargs) -> TableResource: return get_constraint_data(name, sub_dir="checkpoint_files", test=True, **kwargs) +def filter_for_test( + ht: hl.Table, + use_gene_list: bool = False, +) -> hl.Table: + """ + Filter ``ht`` to chr20, chrX, and chrY or a gene list for testing. + + :param ht: Table to filter. + :param use_gene_list: Whether to use a gene list for testing instead of all of + chr20, chrX, and chrY for testing. + :return: Filtered Table for testing. + """ + rg = get_reference_genome(ht.locus) + if use_gene_list: + if rg == "GRCh37": + keep_regions = [ + "20:49505585-49547958", # ADNP + "20:853296-896977", # ANGPT4 + "X:13752832-13787480", # OFD1 + "X:57313139-57515629", # FAAH2 + "Y:2803112-2850547", # ZFY + ] + else: + keep_regions = [ + "chr20:50888916-50931437", # ADNP + "chr20:869900-916334", # ANGPT4 + "chrX:13734743-13777955", # OFD1 + "chrX:57286706-57489193", # FAAH2 + "chrY:2935281-2982506", # ZFY + ] + keep = [hl.parse_locus_interval(c, reference_genome=rg) for c in keep_regions] + else: + keep = [ + hl.parse_locus_interval(c, reference_genome=rg) + for c in [rg.contigs[19], rg.x_contigs[0], rg.y_contigs[0]] + ] + logger.info("Filtering the context HT to chr20, chrX, and chrY for testing...") + + return hl.filter_intervals(ht, keep) + + +def get_adj_r_ht() -> hl.Table: + """ + Read the adj_r per-context methylation genome 1kb autosome aggregate Table. + + :return: Table with adj_r annotation keyed by locus. + """ + return hl.read_table( + "gs://gnomad/v4.1/constraint/resources/annotations/ht/" + "adj_r_per_context_methyl_genome_1kb_autosome.agg.ht" + ) + + def get_constraint_resources( version: str, custom_vep_annotation: str, diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index caaa1488..7c95dcff 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -13,6 +13,7 @@ annotate_mutation_type, annotate_with_mu, apply_models, + assemble_constraint_context_ht, build_constraint_consequence_groups, calculate_raw_z_score, calculate_raw_z_score_sd, @@ -34,6 +35,7 @@ CSQ_CODING, filter_vep_transcript_csqs_expr, mane_select_over_canonical_filter_expr, + update_loftee_end_trunc_filter, ) from hail.utils.misc import divide_null, new_temp_file @@ -60,6 +62,7 @@ RELEASE_LOF_FIELDS, RELEASE_PIPELINE_PARAM_GLOBALS, RELEASE_TOP_LEVEL_ANNOTATIONS, + SFS_BIN_CUTOFFS, ) logging.basicConfig( @@ -70,6 +73,92 @@ logger.setLevel(logging.INFO) +def prepare_context_ht( + ht: hl.Table, + coverage_hts: Dict[str, hl.Table], + an_hts: Dict[str, hl.Table], + freq_hts: Dict[str, hl.Table], + filter_hts: Dict[str, hl.Table], + methylation_ht: hl.Table, + gerp_ht: hl.Table, + adj_r_ht: hl.Table, + sfs_bin_cutoffs: Tuple[float, ...] = SFS_BIN_CUTOFFS, +) -> hl.Table: + """ + Annotate the context Table with coverage, AN, frequency, and constraint annotations. + + Applies the LOFTEE END_TRUNC filter fix, assembles the constraint context + Table via :func:`assemble_constraint_context_ht`, then adds genomic region, + SFS bin, adj_r, and coverage/AN reshaping annotations. + + :param ht: VEP context Table. + :param coverage_hts: Dict mapping data type ("exomes", "genomes") to coverage + Tables. + :param an_hts: Dict mapping data type to allele number Tables. + :param freq_hts: Dict mapping data type to frequency Tables (with ``freq`` + field). + :param filter_hts: Dict mapping data type to filter Tables (with ``filters`` + field). + :param methylation_ht: Methylation sites Table. + :param gerp_ht: GERP scores Table. + :param adj_r_ht: Table with adj_r annotation keyed by locus. + :param sfs_bin_cutoffs: Allele frequency upper bounds defining site frequency + spectrum bins. Default is ``SFS_BIN_CUTOFFS``. + :return: Annotated context Table. + """ + # There was a bug in the GERP cutoffs used to filter transcripts with the + # "END_TRUNC" filter in the LOFTEE VEP plugin resulting in some transcripts + # being considered "HC" when they should have been "LC". We use the + # `update_loftee_end_trunc_filter` function to correct this issue. + ht = ht.annotate( + vep=ht.vep.annotate( + transcript_consequences=update_loftee_end_trunc_filter( + ht.vep.transcript_consequences + ) + ) + ) + ht = assemble_constraint_context_ht( + ht, + coverage_hts=coverage_hts, + an_hts=an_hts, + freq_hts=freq_hts, + filter_hts=filter_hts, + methylation_ht=methylation_ht, + gerp_ht=gerp_ht, + transformation_funcs=None, + ) + + # Add annotation for genomic region (autosome/PAR, X non-PAR, Y non-PAR). + genomic_region_expr = ( + hl.case() + .when(ht.locus.in_autosome_or_par(), "autosome_or_par") + .when(ht.locus.in_x_nonpar(), "chrx_nonpar") + .when(ht.locus.in_y_nonpar(), "chry_nonpar") + .or_missing() + ) + + # Add annotation for SFS bin. + af_expr = ht.freq.exomes[0].AF + sfs_bin_expr = hl.case().when(hl.is_missing(af_expr), 0) + for i, af in enumerate(sfs_bin_cutoffs): + sfs_bin_expr = sfs_bin_expr.when(af_expr <= af, i) + sfs_bin_expr = sfs_bin_expr.or_missing() + + return ht.annotate( + coverage=hl.struct( + exomes=ht.coverage.exomes.select("mean", "median_approx"), + genomes=ht.coverage.genomes.select("mean", "median_approx"), + ), + AN=hl.struct( + exomes=ht.AN.exomes[0], + genomes=ht.AN.genomes[0], + ), + genomic_region=genomic_region_expr, + adj_r=adj_r_ht[ht.locus].adj_r[ht.context], + sfs_bin=sfs_bin_expr, + ) + + # TODO: For now I am leaving this here instead of moving to gnomad_methods because # there is another PR in gnomad_methods that might change the way this function is # implemented. From 43aefd5e2e70e44df1ad34adac1af9b46f3707b9 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 18 Mar 2026 09:22:09 -0600 Subject: [PATCH 25/38] Add synonymous adjacency resource to constraint pipeline. Introduce `get_syn_adj_r_ht` function for reading synonymous DNM adjacency data and update context preparation to include synonymous adjacencies in annotations. Enhance overall clarity and organization of the constraint pipeline. --- gnomad_constraint/pipeline/constraint_pipeline.py | 7 ++++++- gnomad_constraint/resources/resource_utils.py | 12 ++++++++++++ gnomad_constraint/utils/constraint.py | 5 ++++- 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index f8b62241..f2dd67c2 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -40,7 +40,11 @@ RELEASE_KEY_ORDER, VERSIONS, ) -from gnomad_constraint.resources.resource_utils import filter_for_test, get_adj_r_ht +from gnomad_constraint.resources.resource_utils import ( + filter_for_test, + get_adj_r_ht, + get_syn_adj_r_ht, +) from gnomad_constraint.utils.constraint import ( aggregate_by_constraint_groups, aggregate_per_variant_expected_ht, @@ -132,6 +136,7 @@ def main(args): methylation_ht=res.methylation_ht.ht(), gerp_ht=constraint_res.get_gerp_ht(get_reference_genome(ht.locus).name), adj_r_ht=get_adj_r_ht(), + syn_adj_r_ht=get_syn_adj_r_ht(), ) ht.write(res.annotated_context_ht.path, overwrite) diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index cc4a3101..4763cab5 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -603,6 +603,18 @@ def get_adj_r_ht() -> hl.Table: ) +def get_syn_adj_r_ht() -> hl.Table: + """ + Read the synonymous DNM adj_r per-context methylation genome 1kb autosome Table. + + :return: Table with syn_adj_r annotation keyed by locus. + """ + return hl.read_table( + "gs://gnomad/v4.1/constraint/resources/annotations/ht/" + "adj_r_syn_dnm_per_context_methyl_genome_1kb_autosome.ht" + ) + + def get_constraint_resources( version: str, custom_vep_annotation: str, diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 7c95dcff..6f679aa2 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -82,6 +82,7 @@ def prepare_context_ht( methylation_ht: hl.Table, gerp_ht: hl.Table, adj_r_ht: hl.Table, + syn_adj_r_ht: hl.Table, sfs_bin_cutoffs: Tuple[float, ...] = SFS_BIN_CUTOFFS, ) -> hl.Table: """ @@ -89,7 +90,7 @@ def prepare_context_ht( Applies the LOFTEE END_TRUNC filter fix, assembles the constraint context Table via :func:`assemble_constraint_context_ht`, then adds genomic region, - SFS bin, adj_r, and coverage/AN reshaping annotations. + SFS bin, adj_r, syn_adj_r, and coverage/AN reshaping annotations. :param ht: VEP context Table. :param coverage_hts: Dict mapping data type ("exomes", "genomes") to coverage @@ -102,6 +103,7 @@ def prepare_context_ht( :param methylation_ht: Methylation sites Table. :param gerp_ht: GERP scores Table. :param adj_r_ht: Table with adj_r annotation keyed by locus. + :param syn_adj_r_ht: Table with synonymous DNM adj_r annotation keyed by locus. :param sfs_bin_cutoffs: Allele frequency upper bounds defining site frequency spectrum bins. Default is ``SFS_BIN_CUTOFFS``. :return: Annotated context Table. @@ -155,6 +157,7 @@ def prepare_context_ht( ), genomic_region=genomic_region_expr, adj_r=adj_r_ht[ht.locus].adj_r[ht.context], + syn_adj_r=syn_adj_r_ht[ht.locus].adj_r[ht.context], sfs_bin=sfs_bin_expr, ) From c3fcf2f6d395f76d46560219b7e5fe9541e172b4 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 18 Mar 2026 10:18:29 -0600 Subject: [PATCH 26/38] Update constraint utility functions to enhance clarity in frequency filtering. Replace population references with genetic ancestry group terminology and improve documentation for downsampling logic. Streamline annotations for observed and possible variants in the constraint pipeline. --- gnomad_constraint/utils/constraint.py | 247 +++++++++++++------------- 1 file changed, 122 insertions(+), 125 deletions(-) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 6f679aa2..73d55740 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -177,16 +177,19 @@ def filter_freq_for_constraint( Filter the frequency array for constraint calculations. The frequency array is filtered to include only adj frequencies for - the populations in `pops` and the downsamplings in `downsamplings`, for the - populations in `downsampling_pops`. + the genetic ancestry groups in ``gen_ancs`` and the downsamplings in + ``downsamplings``, for the genetic ancestry groups in + ``downsampling_gen_ancs``. No matter the input, the frequency array is always filtered to include the "adj" frequency for the full dataset. - If `downsamplings` is None, no downsamplings are included. If `downsamplings` is - provided, and `downsampling_pops` is None, only the "global" downsampling is - included. If `downsampling_pops` is provided, the downsamplings for the populations - in `downsampling_pops` are included as well as the "global" downsampling. + If ``downsamplings`` is None, no downsamplings are included. If + ``downsamplings`` is provided, and ``downsampling_gen_ancs`` is None, only + the "global" downsampling is included. If ``downsampling_gen_ancs`` is + provided, the downsamplings for the genetic ancestry groups in + ``downsampling_gen_ancs`` are included as well as the "global" + downsampling. :param freq_expr: Frequency array. :param freq_meta_expr: Frequency metadata array. @@ -204,13 +207,13 @@ def filter_freq_for_constraint( meta_keep = [ADJ_FREQ_META] if gen_ancs is not None: - meta_keep += [{**ADJ_FREQ_META, gen_anc_label: pop} for pop in gen_ancs] + meta_keep += [{**ADJ_FREQ_META, gen_anc_label: gen_anc} for gen_anc in gen_ancs] if downsamplings is not None: - downsampling_pops = ["global"] + (downsampling_gen_ancs or []) + downsampling_gen_ancs = ["global"] + (downsampling_gen_ancs or []) meta_keep += [ - {**ADJ_FREQ_META, gen_anc_label: pop, "downsampling": str(ds)} - for pop in downsampling_pops + {**ADJ_FREQ_META, gen_anc_label: gen_anc, "downsampling": str(ds)} + for gen_anc in downsampling_gen_ancs for ds in downsamplings ] @@ -245,10 +248,10 @@ def get_annotations_for_computing_mu( requested genetic ancestries and downsampling level. - observed_variants: This annotation is an array, where each element corresponds to whether the variant is observed in the genomes dataset for the - frequency group at the corresponding index in the `genomes_freq` array. - Must PASS genome filters, have AC <= 'ac_cutoff' at the specified - `downsampling_level`, and have a genome mean coverage >= `min_cov` - and <= `max_cov`. The boolean value is stored as an integer (0 or 1). + frequency group at the corresponding index in the ``genomes_freq`` array. + Must PASS genome filters, have AC <= ``ac_cutoff`` at the specified + ``downsampling_level``, and have a genome mean coverage >= ``min_cov`` + and <= ``max_cov``. The boolean value is stored as an integer (0 or 1). - possible_variants: Whether the variant is considered a possible variant in the genomes dataset. This includes variants not in the genome dataset (genome AF undefined), or also considered in the observed variant set. The boolean @@ -260,7 +263,7 @@ def get_annotations_for_computing_mu( - Is autosomal. - Has a most severe transcript consequence of: "intron_variant" or "intergenic_variant". - - Is at a site with GERP > `gerp_lower_cutoff` and < `gerp_upper_cutoff`. + - Is at a site with GERP > ``gerp_lower_cutoff`` and < ``gerp_upper_cutoff``. The function also returns a struct of the mutation rate globals: @@ -278,7 +281,7 @@ def get_annotations_for_computing_mu( .. note:: - Values for `gerp_lower_cutoff` and `gerp_upper_cutoff` default to -3.9885 and + Values for ``gerp_lower_cutoff`` and ``gerp_upper_cutoff`` default to -3.9885 and 2.6607, respectively. These values were precalculated on the GRCh37 context table and define the 5th and 95th percentiles. @@ -303,6 +306,8 @@ def get_annotations_for_computing_mu( :return: Tuple containing the observed and possible variant annotations and the globals. """ + # Always include the global downsampling; gen_ancs only controls which + # per-ancestry downsamplings are included. genomes_freq_expr, genomes_freq_meta = filter_freq_for_constraint( genomes_freq_expr, genomes_freq_meta, @@ -327,9 +332,10 @@ def get_annotations_for_computing_mu( keep_expr &= (gerp_expr > gerp_lower_cutoff) & (gerp_expr < gerp_upper_cutoff) # Filter so that the most severe annotation is 'intron_variant' or - # 'intergenic_variant' - keep_expr &= (most_severe_consequence_expr == "intron_variant") | ( - most_severe_consequence_expr == "intergenic_variant" + # 'intergenic_variant'. + keep_expr &= hl.any( + most_severe_consequence_expr == c + for c in ["intron_variant", "intergenic_variant"] ) # Set up the criteria to keep high-quality sites, and sites found in less than or @@ -370,16 +376,16 @@ def get_exome_coverage_expr( """ Get the exome coverage expression based on the specified metric. - The requested `exome_coverage_metric` is extracted from the exome coverage - annotations in the input `ht`: + The requested ``exome_coverage_metric`` is extracted from the exome coverage + annotations in the input ``ht``: - "median": the expression returned is "median_approx" if it exists in - `ht.coverage.exomes`, otherwise "median". - - "AN": the expression returned is the exomes allele number (`ht.AN.exomes`). + ``ht.coverage.exomes``, otherwise "median". + - "AN": the expression returned is the exomes allele number (``ht.AN.exomes``). - "AN_percent": the expression returned is the percent of samples with a - non-missing genotype, which is the exomes allele number (`ht.AN.exomes`) + non-missing genotype, which is the exomes allele number (``ht.AN.exomes``) divided by the total number of alleles in the exomes dataset (pulled from - `ht.an_globals.exomes.strata_sample_count` * 2) multiplied by 100. + ``ht.an_globals.exomes.strata_sample_count`` * 2) multiplied by 100. :param ht: Input Table with exome coverage information. :param exome_coverage_metric: Metric to use for exome coverage. One of ["median", @@ -409,7 +415,7 @@ def get_exome_coverage_expr( hl.case() .when(ht.locus.in_x_nonpar(), (xx_an_sample_count * 2) + xy_an_sample_count) .when(ht.locus.in_y_nonpar(), xy_an_sample_count) - .default(an_sample_count[0] * 2) + .default(an_sample_count[0] * 2) # Index 0 is adj (all samples). ) cov_expr = hl.int((ht.AN.exomes / an_count) * 100) @@ -441,7 +447,7 @@ def get_exomes_observed_and_possible( - observed_variants: This annotation is an array, where each element corresponds to whether the variant is observed in the exomes dataset for the frequency - group at the corresponding index in the `exomes_freq` array and has an + group at the corresponding index in the ``exomes_freq`` array and has an AF <= 0.001. The boolean value is stored as an integer (0 or 1). - possible_variants: Whether the variant is considered a possible variant in the exomes dataset. This includes variants not in the exome dataset (exome AF @@ -482,7 +488,7 @@ def get_exomes_observed_and_possible( logger.info("The following downsamplings will be used: %s", downsamplings) # Filter frequency array for computing the observed expression on all requested - # populations and downsamplings. + # genetic ancestry groups and downsamplings. exomes_freq_expr, exomes_freq_meta = filter_freq_for_constraint( exomes_freq_expr, exomes_freq_meta, @@ -522,14 +528,14 @@ def get_build_calibration_model_annotation( high_cov_cutoff: int = COVERAGE_CUTOFF, upper_cov_cutoff: Optional[int] = None, skip_coverage_model: bool = False, -) -> Tuple[hl.expr.StructExpression, hl.expr.StructExpression]: +) -> hl.expr.StructExpression: """ - Get the annotation and globals for building the calibration models. + Get the annotation for building the calibration models. The build model grouping is set to missing if the variant is not a "synonymous_variant" in a canonical or MANE Select transcript (depending on - `synonymous_transcript_filter_field`). Otherwise, it is a struct with the following - fields detailed in `calibration_model_group_expr`. + ``synonymous_transcript_filter_field``). Otherwise, it is a struct with the + following fields detailed in ``calibration_model_group_expr``. :param exomes_coverage_expr: Exome coverage expression. :param transcript_csq_expr: Transcript consequences expression. @@ -545,7 +551,7 @@ def get_build_calibration_model_annotation( None. :param skip_coverage_model: Whether the coverage model should be skipped during the build models step. Default is False. - :return: Tuple containing the build model expression and the global parameters. + :return: Build model struct expression, or missing if no synonymous transcripts. """ # Determine the canonical and mane_select parameters for # 'filter_vep_transcript_csqs_expr' based on 'synonymous_transcript_filter_field'. @@ -610,10 +616,10 @@ def prepare_ht_for_constraint_calculations( required for the constraint calculations. Please see the following functions for more information on the annotations generated: - - `get_annotations_for_computing_mu` - - `get_exomes_observed_and_possible` - - `get_build_calibration_model_annotation` - - `get_apply_calibration_model_annotation` + - ``get_annotations_for_computing_mu`` + - ``get_exomes_observed_and_possible`` + - ``get_build_calibration_model_annotation`` + - ``calibration_model_group_expr`` (for apply model annotations) :param ht: Annotated context Table. :param exome_coverage_metric: Metric to use for exome coverage. One of ["median", @@ -643,13 +649,13 @@ def prepare_ht_for_constraint_calculations( :param build_model_upper_cov_cutoff: Upper coverage cutoff for the build models step. Default is None. :param apply_model_low_cov_cutoff: Low coverage cutoff for the apply models step. - Default is COVERAGE_CUTOFF. + Default is None. :param apply_model_high_cov_cutoff: High coverage cutoff for the apply models step. Default is COVERAGE_CUTOFF. :param skip_coverage_model: Whether the coverage model should be skipped during the build and apply models steps. Default is False. :param synonymous_transcript_filter_field: Field used to filter to variants with a - transcript consequence of "synonymous_variant". Default is "canonical". + transcript consequence of "synonymous_variant". Default is "mane_select". :return: Table with the computed annotations. """ # Get the annotations relevant for computing the mutation rate. @@ -744,17 +750,18 @@ def prepare_ht_for_constraint_calculations( def create_training_set( ht: hl.Table, mutation_ht: hl.Table, - partition_hint=100, + partition_hint: int = 100, ) -> hl.Table: """ Create the training set for the constraint model. - The input `ht` should be prepared using `prepare_ht_for_constraint_calculations`. - The `ht` is filtered to include only the rows that have a build model annotation. - The observed and possible variants are counted by group and annotated with the - mutation rate. The Table is then checkpointed to avoid memory and shuffle issues. + The input ``ht`` should be prepared using + ``prepare_ht_for_constraint_calculations``. The ``ht`` is filtered to include only + the rows that have a build model annotation. The observed and possible variants are + counted by group and annotated with the mutation rate. The Table is then + checkpointed to avoid memory and shuffle issues. - :param ht: Table prepared using `prepare_ht_for_constraint_calculations`. + :param ht: Table prepared using ``prepare_ht_for_constraint_calculations``. :param mutation_ht: Mutation rate Table. :param partition_hint: Partition hint for the Table. Default is 100. :return: Training set Table. @@ -806,13 +813,14 @@ def create_per_variant_expected_ht( """ Create the per-variant expected Table. - The input `ht` should be prepared using `prepare_ht_for_constraint_calculations`. - The `ht` is filtered to include only the rows that have an apply model annotation ( - if `filter_to_apply_variants` is True). The Table is then annotated with the - expected number of variants using `apply_models`. See the function `apply_models` - for more information on the expected annotations. + The input ``ht`` should be prepared using + ``prepare_ht_for_constraint_calculations``. The ``ht`` is filtered to include only + the rows that have an apply model annotation (if ``filter_to_apply_variants`` is + True). The Table is then annotated with the expected number of variants using + ``apply_models``. See the function ``apply_models`` for more information on the + expected annotations. - :param ht: Table prepared using `prepare_ht_for_constraint_calculations`. + :param ht: Table prepared using ``prepare_ht_for_constraint_calculations``. :param mutation_ht: Mutation rate Table. :param plateau_models: Plateau models for the constraint calculations. :param coverage_model: Coverage model for the constraint calculations. @@ -820,22 +828,17 @@ def create_per_variant_expected_ht( :param filter_to_apply_variants: Whether to filter to only the rows with an apply model annotation. Default is True. :param custom_vep_annotation: Custom VEP annotation to use. Default is + "transcript_consequences". :param use_mane_select: Whether to include MANE Select as a group. Default is False. :return: Per-variant expected Table. """ - include_canonical_group = False - include_mane_select_group = False - if custom_vep_annotation == "worst_csq_by_gene": - vep_annotation = "worst_csq_by_gene" - if use_mane_select: - raise ValueError( - "'mane_select' cannot be set to True when custom_vep_annotation is set" - " to 'worst_csq_by_gene'." - ) - else: - vep_annotation = custom_vep_annotation - include_canonical_group = True - include_mane_select_group = use_mane_select + if custom_vep_annotation == "worst_csq_by_gene" and use_mane_select: + raise ValueError( + "'mane_select' cannot be set to True when custom_vep_annotation is set" + " to 'worst_csq_by_gene'." + ) + include_canonical_group = custom_vep_annotation != "worst_csq_by_gene" + include_mane_select_group = include_canonical_group and use_mane_select calibrate_mu_fields = set(ht.calibrate_mu.keys()) ht = ht.annotate(**ht.calibrate_mu) @@ -863,7 +866,7 @@ def create_per_variant_expected_ht( ht, groupings = annotate_exploded_vep_for_constraint_groupings( ht=ht, - vep_annotation=vep_annotation, + vep_annotation=custom_vep_annotation, include_canonical_group=include_canonical_group, include_mane_select_group=include_mane_select_group, ) @@ -881,19 +884,17 @@ def create_per_variant_expected_ht( def aggregate_per_variant_expected_ht( - ht, + ht: hl.Table, include_mu_annotations_in_grouping: bool = False, -): +) -> hl.Table: """ Aggregate the per-variant expected Table. - The input `ht` should be the Table returned by `create_per_variant_expected_ht`. - The Table is exploded by the VEP annotation and aggregated by "genomic_region", - "context", "ref", "alt", "methylation_level", groupings returned by - `annotate_exploded_vep_for_constraint_groupings` and fields in - `additional_grouping` to get the observed and expected counts. + The input ``ht`` should be the Table returned by ``create_per_variant_expected_ht``. + The Table is aggregated by the groupings stored in + ``apply_models_globals.groupings`` to get the observed and expected counts. - :param ht: Table returned by `create_per_variant_expected_ht`. + :param ht: Table returned by ``create_per_variant_expected_ht``. :param include_mu_annotations_in_grouping: Whether to include the mutation rate key annotations in the grouping. Default is False. :return: Table with the observed and expected counts. @@ -933,13 +934,13 @@ def calculate_mu_by_downsampling( - ref - the reference allele. - alt - the alternate base. - methylation_level - methylation_level. - - downsampling_counts_{pop} - variant counts in downsamplings for populations - in `pops`. + - downsampling_counts_{gen_anc} - variant counts in downsamplings for genetic + ancestry groups in ``gen_ancs``. - mu_snp - SNP mutation rate. - - annotations added by `annotate_mutation_type`. + - annotations added by ``annotate_mutation_type``. - :param ht: Table returned by `prepare_ht_for_constraint_calculations`. - :param additional_grouping: Annotations other than 'context', 'ref', and 'alt'. + :param ht: Table returned by ``prepare_ht_for_constraint_calculations``. + :param additional_grouping: Annotations other than "context", "ref", and "alt". Default is ('methylation_level',). :param total_mu: The per-generation mutation rate. Default is 1.2e-08. :return: Mutation rate Table. @@ -1008,14 +1009,14 @@ def get_transcript_filter_expr( return ht.transcript.startswith("ENST") & ht.canonical -def add_oe_upper_rank_and_decile( +def add_oe_upper_rank_and_bins( ht: hl.Table, use_mane_select_over_canonical: bool = True, mane_select_only: bool = False, bin_granularities: Optional[Dict[str, int]] = None, ) -> hl.Table: """ - Compute the rank and decile of the oe upper confidence interval. + Compute the rank and bins of the oe upper confidence interval. Thin wrapper around :func:`rank_array_element_metrics` that extracts the discretized Poisson and gamma upper CI values from each constraint group's @@ -1056,15 +1057,17 @@ def aggregate_by_constraint_groups( ht: hl.Table, keys: Tuple = ("gene", "transcript", "canonical"), classic_lof_annotations: Tuple = CLASSIC_LOF_ANNOTATIONS, - additional_groupings: Dict[str, Dict[str, hl.expr.BooleanExpression]] = None, - additional_grouping_combinations: List[List[str]] = None, + additional_groupings: Optional[ + Dict[str, Dict[str, hl.expr.BooleanExpression]] + ] = None, + additional_grouping_combinations: Optional[List[List[str]]] = None, ) -> hl.Table: """ - Aggregate the observed and expected variant info for synonymous variants, missense variants, and predicted loss-of-function (pLoF) variants. + Aggregate observed and expected variant info for synonymous, missense, and pLoF variants. .. note:: - The following annotations should be present in `ht`: + The following annotations should be present in ``ht``: - modifier - annotation @@ -1073,8 +1076,8 @@ def aggregate_by_constraint_groups( - possible_variants - expected_variants - :param ht: Input Table with the number of expected variants (output of - `get_proportion_observed()`). + :param ht: Input Table with observed and expected variant counts (output of the + apply models step). :param keys: The keys of the output Table, defaults to ('gene', 'transcript', 'canonical'). :param classic_lof_annotations: Classic LoF Annotations used to filter the input @@ -1144,8 +1147,6 @@ def aggregate_by_constraint_groups( ) ) - ht = ht.annotate_globals(constraint_group_meta=meta) - return ht @@ -1273,10 +1274,10 @@ def _annotate_oe_ci_z( For each constraint group's ``oe_info`` entries, adds: - - ``oe`` — observed / expected ratio. - - ``oe_ci_discretized_poisson`` — discretized Poisson CI. - - ``oe_ci_gamma`` — gamma-distribution CI. - - ``z_raw`` — raw z-score. + - ``oe`` — observed / expected ratio. + - ``oe_ci_discretized_poisson`` — discretized Poisson CI. + - ``oe_ci_gamma`` — gamma-distribution CI. + - ``z_raw`` — raw z-score. Then adds per-group ``flags`` based on z-score outlier thresholds. @@ -1398,9 +1399,9 @@ def _compute_percentile_bins( use_mane_select_over_canonical: bool = True, ) -> hl.Table: """ - Add OE upper CI rank, decile, and percentile bin annotations. + Add OE upper CI rank and percentile bin annotations. - Adds rank and decile annotations via :func:`add_oe_upper_rank_and_decile`, + Adds rank and bin annotations via :func:`add_oe_upper_rank_and_bins`, then computes percentile thresholds across all granularities defined in ``CONSTRAINT_GRANULARITIES`` and annotates bins via :func:`annotate_constraint_percentile_bins`. @@ -1410,7 +1411,7 @@ def _compute_percentile_bins( transcripts for filtering when determining ranks. Default is True. :return: Table with rank, decile, and percentile bin annotations. """ - ht = add_oe_upper_rank_and_decile(ht, use_mane_select_over_canonical) + ht = add_oe_upper_rank_and_bins(ht, use_mane_select_over_canonical) meta = hl.eval(ht.constraint_group_meta) metric_group_idx = { @@ -1507,17 +1508,12 @@ def compute_constraint_metrics( .. note:: - The following annotations should be present in `ht`: - - - modifier - - annotation - - observed_variants - - mu - - possible_variants - - expected_variants + The input ``ht`` should be the output of + :func:`aggregate_by_constraint_groups`, which has a + ``constraint_groups`` array, ``constraint_group_meta``, + ``exomes_freq_meta``, and ``no_variants`` annotations. - :param ht: Input Table with the number of expected variants (output of - ``aggregate_by_constraint_groups``). + :param ht: Table output by :func:`aggregate_by_constraint_groups`. :param gencode_ht: Table containing GENCODE annotations. :param gene_quality_metrics_ht: Table keyed by transcript with ``gene_quality_metrics`` and ``gene_flags`` fields (output of @@ -1580,14 +1576,14 @@ def _restructure_release_rows( For each constraint group, builds a flat release struct by: - - Flattening the adjusted-frequency ``oe_info`` entry onto the group - struct, keeping ``oe`` and ``z_raw`` under their original names and - overriding ``oe_ci`` with ``oe_ci_gamma``. - - Applying ``RELEASE_CG_RENAME`` to rename group-level and ``oe_info`` - fields (e.g. ``mu_snp`` -> ``mu``, ``observed_variants`` -> ``obs``). - - When downsampling data is present, adding ``gen_anc_obs`` / - ``gen_anc_exp`` structs keyed by genetic ancestry with arrays of - values ordered by downsampling level. + - Flattening the adjusted-frequency ``oe_info`` entry onto the group + struct, keeping ``oe`` and ``z_raw`` under their original names and + overriding ``oe_ci`` with ``oe_ci_gamma``. + - Applying ``RELEASE_CG_RENAME`` to rename group-level and ``oe_info`` + fields (e.g. ``mu_snp`` -> ``mu``, ``observed_variants`` -> ``obs``). + - When downsampling data is present, adding ``gen_anc_obs`` / + ``gen_anc_exp`` structs keyed by genetic ancestry with arrays of + values ordered by downsampling level. Annotates the Table with one top-level field per group (applying ``RELEASE_GROUP_RENAMES``, e.g. ``lof_hc`` -> ``lof``), trims the @@ -1657,7 +1653,7 @@ def _restructure_release_rows( for k in (RELEASE_LOF_FIELDS if name in RELEASE_GROUPS_WITH_PLI else []) }, ) - .select() + .select() # Drop intermediate fields, keep only annotated release fields. for i, name in enumerate(field_names) } ht = ht.annotate(**cg_fields) @@ -1668,6 +1664,7 @@ def _restructure_release_rows( ht = ht.key_by(*available_keys) ht = ht.filter(hl.any([ht[k].possible != 0 for k in RELEASE_GROUP_NAMES])) + return ht @@ -1684,17 +1681,17 @@ def _restructure_release_globals( Replaces internal globals with a clean release set: - - Pipeline parameter globals are renamed and stripped of internal-only - fields via ``RELEASE_PIPELINE_PARAM_GLOBALS``. - - ``sd_raw_z`` is converted from an ordered array (one entry per - constraint group) to a named struct keyed by release group name, - retaining only groups in ``RELEASE_GROUP_NAMES``. - - When downsampling data is present, a ``downsamplings`` struct is added - keyed by genetic ancestry, with arrays of integer downsampling levels - matching the order of ``gen_anc_obs`` / ``gen_anc_exp`` in the rows. - - ``max_af`` is preserved unchanged if present. - - ``version`` is set to ``release_version`` if provided, otherwise - carried over from the existing global. + - Pipeline parameter globals are renamed and stripped of internal-only + fields via ``RELEASE_PIPELINE_PARAM_GLOBALS``. + - ``sd_raw_z`` is converted from an ordered array (one entry per + constraint group) to a named struct keyed by release group name, + retaining only groups in ``RELEASE_GROUP_NAMES``. + - When downsampling data is present, a ``downsamplings`` struct is added + keyed by genetic ancestry, with arrays of integer downsampling levels + matching the order of ``gen_anc_obs`` / ``gen_anc_exp`` in the rows. + - ``max_af`` is preserved unchanged if present. + - ``version`` is set to ``release_version`` if provided, otherwise + carried over from the existing global. :param ht: Table whose globals are being restructured. :param field_names: Internal name for each constraint group (parallel to From 2cf6808d6d978c153f05c2888a5feeb9ed8efd55 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 18 Mar 2026 10:21:09 -0600 Subject: [PATCH 27/38] Update gnomAD constraint pipeline version to 4.1.1 and add it to the supported versions list. Adjust current version constant accordingly. --- gnomad_constraint/resources/constants.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gnomad_constraint/resources/constants.py b/gnomad_constraint/resources/constants.py index f291b334..3559c771 100644 --- a/gnomad_constraint/resources/constants.py +++ b/gnomad_constraint/resources/constants.py @@ -7,10 +7,10 @@ EXTENSIONS = ["ht", "tsv", "tsv.bgz", "he", "log"] """Valid file extensions for constraint pipeline resources.""" -VERSIONS = ["2.1.1", "4.0", "4.1"] +VERSIONS = ["2.1.1", "4.0", "4.1", "4.1.1"] """Supported gnomAD constraint pipeline versions.""" -CURRENT_VERSION = "4.1" +CURRENT_VERSION = "4.1.1" """Current default gnomAD constraint pipeline version.""" DATA_TYPES = ["context", "exomes", "genomes"] From 05f4c34af096984ef40ee02bf654c1159f33077d Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 18 Mar 2026 11:08:14 -0600 Subject: [PATCH 28/38] Enhance constraint pipeline by introducing aggregated model application. Add functionality to compute gene quality metrics and streamline aggregation processes. Implement new command-line arguments for applying models on aggregated data and using aggregated expected outputs. Refactor utility functions for clarity and organization. --- .../pipeline/constraint_pipeline.py | 89 +++++-- gnomad_constraint/resources/resource_utils.py | 33 ++- gnomad_constraint/utils/constraint.py | 234 ++++++++++++++---- 3 files changed, 292 insertions(+), 64 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index f2dd67c2..2fd480d1 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -51,6 +51,7 @@ calculate_mu_by_downsampling, compute_constraint_metrics, compute_gene_quality_metrics, + create_aggregated_expected_ht, create_per_variant_expected_ht, create_training_set, flatten_release_ht, @@ -142,6 +143,20 @@ def main(args): logger.info("Done annotating the VEP context Table.") + if args.compute_gene_quality_metrics: + logger.info("Computing per-transcript gene quality metrics...") + res = resources.compute_gene_quality_metrics + res.check_resource_existence() + + gencode_cds_ht = constraint_res.get_gencode_cds_ht(version).ht() + gene_quality_ht = compute_gene_quality_metrics( + res.annotated_context_ht.ht(), + res.exomes_sites_ht.ht(), + gencode_cds_ht, + ) + gene_quality_ht.write(res.gene_quality_metrics_ht.path, overwrite=overwrite) + logger.info("Done computing gene quality metrics.") + if args.calculate_gerp_cutoffs: logger.warning( "Calculating new GERP cutoffs to be used instead of" @@ -296,38 +311,54 @@ def main(args): "consequence annotations, and consequence modifier annotations." ) + if args.apply_models_aggregated: + logger.info("Aggregating counts and applying models on aggregated data...") + res = resources.apply_models_aggregated + res.check_resource_existence() + + hl._set_flags(use_new_shuffle="1") + + ht = res.temp_preprocess_data_ht.ht() + print_global_struct(ht.apply_models_globals) + ht = create_aggregated_expected_ht( + ht, + res.mutation_ht.ht().select("mu_snp"), + res.model_plateau.he(), + coverage_model=( + None if skip_coverage_model else res.model_coverage.he() + ), + log10_coverage=log10_coverage, + custom_vep_annotation=custom_vep_annotation, + use_mane_select=True, + ) + ht.write(res.aggregated_expected_ht.path, overwrite=overwrite) + hl._set_flags(use_new_shuffle=None) + + logger.info("Done with aggregated model application.") + if args.aggregate_by_constraint_groups: logger.info( "Aggregating observed and expected variant counts by constraint groups..." ) - res = resources.aggregate_by_constraint_groups - res.check_resource_existence() - # Use new shuffle method to prevent shuffle errors. hl._set_flags(use_new_shuffle="1") - ht = res.apply_ht.ht() + if args.use_aggregated_expected: + res = resources.apply_models_aggregated + ht = res.aggregated_expected_ht.ht() + else: + res = resources.aggregate_by_constraint_groups + res.check_resource_existence() + ht = res.apply_ht.ht() + + out_res = resources.aggregate_by_constraint_groups aggregate_by_constraint_groups( ht, keys=tuple(k for k in ht.key if k in RELEASE_KEY_ORDER), - ).write(res.constraint_group_ht.path, overwrite=overwrite) + ).write(out_res.constraint_group_ht.path, overwrite=overwrite) hl._set_flags(use_new_shuffle=None) logger.info("Done with aggregating by constraint groups.") - if args.compute_gene_quality_metrics: - logger.info("Computing per-transcript gene quality metrics...") - res = resources.compute_gene_quality_metrics - res.check_resource_existence() - - gencode_cds_ht = constraint_res.get_gencode_cds_ht(version).ht() - gene_quality_ht = compute_gene_quality_metrics( - res.annotated_context_ht.ht(), - res.exomes_sites_ht.ht(), - gencode_cds_ht, - ) - gene_quality_ht.write(res.gene_quality_metrics_ht.path, overwrite=overwrite) - logger.info("Done computing gene quality metrics.") - if args.compute_constraint_metrics: logger.info( "Computing constraint metrics, including pLI scores, z scores, oe" @@ -664,6 +695,17 @@ def main(args): ), action="store_true", ) + parser.add_argument( + "--apply-models-aggregated", + help=( + "Apply plateau and coverage models and aggregate to constraint groups" + " in a single step, without writing per-variant intermediates. This is" + " an alternative to running --apply-models-per-variant," + " --aggregate-per-variant-expected, and" + " --aggregate-by-constraint-groups separately." + ), + action="store_true", + ) aggregate_per_variant_expected_args = parser.add_argument_group( "Aggregate per variant expected args", @@ -721,6 +763,15 @@ def main(args): ), action="store_true", ) + aggregate_by_constraint_groups_args.add_argument( + "--use-aggregated-expected", + help=( + "Read from the --apply-models-aggregated output instead of the" + " --aggregate-per-variant-expected output as input for" + " --aggregate-by-constraint-groups." + ), + action="store_true", + ) gene_quality_args = parser.add_argument_group( "Compute gene quality metrics args", diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 4763cab5..b0a967a1 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -465,6 +465,27 @@ def get_aggregated_per_variant_expected( ) +def get_aggregated_expected( + custom_vep_annotation: str = "transcript_consequences", **kwargs +) -> TableResource: + """ + Return TableResource for the aggregated expected variant counts Table. + + This is the output of the aggregated model application path, where counts + are aggregated before applying models (as opposed to the per-variant path). + + :param custom_vep_annotation: The VEP annotation used to customize the constraint + model (one of "transcript_consequences" or "worst_csq_by_gene"). + :return: TableResource of the aggregated expected Table. + """ + return get_constraint_data( + "aggregated_expected", + sub_dir="apply_models", + custom_vep_annotation=custom_vep_annotation, + **kwargs, + ) + + def get_constraint_group_ht(custom_vep_annotation: str, **kwargs) -> TableResource: """ Return TableResource of constraint group Table. @@ -739,6 +760,15 @@ def get_constraint_resources( }, pipeline_input_steps=[aggregate_per_variant_expected], ) + apply_models_aggregated = PipelineStepResourceCollection( + "--apply-models-aggregated", + output_resources={ + "aggregated_expected_ht": get_aggregated_expected( + custom_vep_annotation, **common_params, path_post_fix=path_post_fix + ) + }, + pipeline_input_steps=[preprocess_data, calculate_mutation_rate, build_models], + ) compute_gene_quality_metrics_step = PipelineStepResourceCollection( "--compute-gene-quality-metrics", output_resources={ @@ -783,6 +813,7 @@ def get_constraint_resources( constraint_pipeline.add_steps( { "prepare_context": prepare_context, + "compute_gene_quality_metrics": compute_gene_quality_metrics_step, "preprocess_data": preprocess_data, "calculate_gerp_cutoffs": calculate_gerp_cutoffs, "calculate_mutation_rate": calculate_mutation_rate, @@ -791,7 +822,7 @@ def get_constraint_resources( "apply_models_per_variant": apply_models_per_variant, "aggregate_per_variant_expected": aggregate_per_variant_expected, "aggregate_by_constraint_groups": aggregate_by_constraint_groups, - "compute_gene_quality_metrics": compute_gene_quality_metrics_step, + "apply_models_aggregated": apply_models_aggregated, "compute_constraint_metrics": compute_constraint_metrics, "prepare_release": prepare_release, "export_release_tsv": export_release_tsv, diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 73d55740..cea31fc9 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -800,37 +800,24 @@ def create_training_set( return ht -def create_per_variant_expected_ht( +def _prepare_ht_for_apply_models( ht: hl.Table, - mutation_ht: hl.Table, - plateau_models: hl.StructExpression, - coverage_model: Tuple[float, float], - log10_coverage: bool = True, - filter_to_apply_variants: bool = True, custom_vep_annotation: str = "transcript_consequences", use_mane_select: bool = False, -) -> hl.Table: +) -> Tuple[hl.Table, List[str]]: """ - Create the per-variant expected Table. + Prepare a preprocessed Table for model application. - The input ``ht`` should be prepared using - ``prepare_ht_for_constraint_calculations``. The ``ht`` is filtered to include only - the rows that have an apply model annotation (if ``filter_to_apply_variants`` is - True). The Table is then annotated with the expected number of variants using - ``apply_models``. See the function ``apply_models`` for more information on the - expected annotations. + Promotes ``calibrate_mu`` fields, filters to rows with a defined apply model + annotation and positive possible variant count, and explodes VEP annotations + to per-transcript rows. :param ht: Table prepared using ``prepare_ht_for_constraint_calculations``. - :param mutation_ht: Mutation rate Table. - :param plateau_models: Plateau models for the constraint calculations. - :param coverage_model: Coverage model for the constraint calculations. - :param log10_coverage: Whether to use log10 coverage. Default is True. - :param filter_to_apply_variants: Whether to filter to only the rows with an apply - model annotation. Default is True. :param custom_vep_annotation: Custom VEP annotation to use. Default is - "transcript_consequences". - :param use_mane_select: Whether to include MANE Select as a group. Default is False. - :return: Per-variant expected Table. + ``"transcript_consequences"``. + :param use_mane_select: Whether to include MANE Select as a group. Default is + False. + :return: Tuple of (prepared Table, list of VEP grouping field names). """ if custom_vep_annotation == "worst_csq_by_gene" and use_mane_select: raise ValueError( @@ -840,19 +827,38 @@ def create_per_variant_expected_ht( include_canonical_group = custom_vep_annotation != "worst_csq_by_gene" include_mane_select_group = include_canonical_group and use_mane_select - calibrate_mu_fields = set(ht.calibrate_mu.keys()) ht = ht.annotate(**ht.calibrate_mu) + ht = ht.filter(hl.is_defined(ht.apply_model) & (ht.possible_variants > 0)) - if filter_to_apply_variants: - # TODO: From Konrad's script parser.add_argument('--skip_af_filter_upfront', - # help='Skip AF filter up front (to be applied later to ensure that it is not - # affecting population-specific constraint): not generally recommended', - # action='store_true') - ht = ht.filter(hl.is_defined(ht.apply_model) & (ht.possible_variants > 0)) + ht, groupings = annotate_exploded_vep_for_constraint_groupings( + ht=ht, + vep_annotation=custom_vep_annotation, + include_canonical_group=include_canonical_group, + include_mane_select_group=include_mane_select_group, + ) - ht = annotate_with_mu(ht, mutation_ht) + return ht, groupings - ht = ht.annotate( + +def _apply_constraint_models( + ht: hl.Table, + plateau_models: hl.StructExpression, + coverage_model: Tuple[float, float], + log10_coverage: bool = True, +) -> hl.Table: + """ + Apply plateau and coverage models to a Table with ``mu_snp`` and ``apply_model``. + + :param ht: Table with ``mu_snp``, ``possible_variants``, ``exomes_coverage``, + and ``apply_model`` fields. + :param plateau_models: Plateau models for the constraint calculations. + :param coverage_model: Coverage model for the constraint calculations. + :param log10_coverage: Whether to use log10 coverage. Default is True. + :return: Table annotated with model outputs (``mu``, + ``predicted_proportion_observed``, ``expected_variants``, + ``coverage_correction``). + """ + return ht.annotate( **apply_models( ht.mu_snp, plateau_models.get(ht.apply_model.model_group), @@ -864,20 +870,80 @@ def create_per_variant_expected_ht( ) ) - ht, groupings = annotate_exploded_vep_for_constraint_groupings( - ht=ht, - vep_annotation=custom_vep_annotation, - include_canonical_group=include_canonical_group, - include_mane_select_group=include_mane_select_group, - ) - ht = ht.annotate_globals( - apply_models_globals=ht.apply_models_globals.annotate( - plateau_models=plateau_models, - coverage_model=coverage_model, - log10_coverage=log10_coverage, - groupings=groupings, +def _annotate_apply_models_globals( + ht: hl.Table, + plateau_models: hl.StructExpression, + coverage_model: Tuple[float, float], + log10_coverage: bool, + groupings: List[str], +) -> hl.Table: + """ + Annotate the Table with model parameters in ``apply_models_globals``. + + If the Table already has ``apply_models_globals``, the new fields are added + to the existing struct. Otherwise a new struct is created. + + :param ht: Input Table. + :param plateau_models: Plateau models used. + :param coverage_model: Coverage model used. + :param log10_coverage: Whether log10 coverage was used. + :param groupings: List of grouping field names. + :return: Table with updated ``apply_models_globals`` global. + """ + model_params = hl.struct( + plateau_models=plateau_models, + coverage_model=coverage_model, + log10_coverage=log10_coverage, + groupings=groupings, + ) + if "apply_models_globals" in ht.globals: + ht = ht.annotate_globals( + apply_models_globals=ht.apply_models_globals.annotate(**model_params) ) + else: + ht = ht.annotate_globals(apply_models_globals=model_params) + return ht + + +def create_per_variant_expected_ht( + ht: hl.Table, + mutation_ht: hl.Table, + plateau_models: hl.StructExpression, + coverage_model: Tuple[float, float], + log10_coverage: bool = True, + custom_vep_annotation: str = "transcript_consequences", + use_mane_select: bool = False, +) -> hl.Table: + """ + Create the per-variant expected Table. + + The input ``ht`` should be prepared using + ``prepare_ht_for_constraint_calculations``. The ``ht`` is filtered to include only + the rows that have an apply model annotation. The Table is then annotated with the + expected number of variants using ``apply_models``. See the function + ``apply_models`` for more information on the expected annotations. + + :param ht: Table prepared using ``prepare_ht_for_constraint_calculations``. + :param mutation_ht: Mutation rate Table. + :param plateau_models: Plateau models for the constraint calculations. + :param coverage_model: Coverage model for the constraint calculations. + :param log10_coverage: Whether to use log10 coverage. Default is True. + :param custom_vep_annotation: Custom VEP annotation to use. Default is + ``"transcript_consequences"``. + :param use_mane_select: Whether to include MANE Select as a group. Default is False. + :return: Per-variant expected Table. + """ + calibrate_mu_fields = set(ht.calibrate_mu.keys()) + + ht, groupings = _prepare_ht_for_apply_models( + ht, custom_vep_annotation, use_mane_select + ) + + ht = annotate_with_mu(ht, mutation_ht) + ht = _apply_constraint_models(ht, plateau_models, coverage_model, log10_coverage) + ht = _annotate_apply_models_globals( + ht, plateau_models, coverage_model, log10_coverage, groupings ) return ht.drop(*calibrate_mu_fields) @@ -920,6 +986,86 @@ def aggregate_per_variant_expected_ht( return ht.naive_coalesce(1000) +def create_aggregated_expected_ht( + ht: hl.Table, + mutation_ht: hl.Table, + plateau_models: hl.StructExpression, + coverage_model: Tuple[float, float], + log10_coverage: bool = True, + custom_vep_annotation: str = "transcript_consequences", + use_mane_select: bool = False, + partition_hint: int = 100, +) -> hl.Table: + """ + Create aggregated expected variant counts by first aggregating, then applying models. + + Unlike :func:`create_per_variant_expected_ht`, which applies models per-variant and + then aggregates, this function first aggregates observed and possible variant counts + by VEP groupings and coverage, then applies plateau and coverage models on the + aggregated counts. The output is compatible with + :func:`aggregate_by_constraint_groups`. + + The steps are: + + 1. Explode VEP annotations to get per-transcript rows. + 2. Aggregate observed and possible counts by VEP groupings, coverage, and + mutation rate context (using ``count_observed_and_possible_by_group``). + 3. Annotate with mutation rate and apply plateau/coverage models on the + aggregated counts. + 4. Aggregate by VEP groupings only (summing model outputs across coverage + and context groups). + + :param ht: Table prepared using ``prepare_ht_for_constraint_calculations``. + :param mutation_ht: Mutation rate Table. + :param plateau_models: Plateau models for the constraint calculations. + :param coverage_model: Coverage model for the constraint calculations. + :param log10_coverage: Whether to use log10 coverage. Default is True. + :param custom_vep_annotation: Custom VEP annotation to use. Default is + ``"transcript_consequences"``. + :param use_mane_select: Whether to include MANE Select as a group. Default is + False. + :param partition_hint: Target number of partitions for aggregation. Default is 100. + :return: Table with aggregated expected variant counts, compatible with + ``aggregate_by_constraint_groups``. + """ + ht, groupings = _prepare_ht_for_apply_models( + ht, custom_vep_annotation, use_mane_select + ) + + # Filter to coding consequences. + ht = ht.filter(hl.set(CSQ_CODING).contains(ht.annotation)) + + # Aggregate observed and possible counts by VEP groupings, coverage, and mutation + # rate context. The additional_grouping includes VEP groupings (gene, transcript, + # annotation, etc.) plus the apply_model fields needed for model application. + vep_groupings = tuple(g for g in groupings if g not in MU_GROUPING) + ht = count_observed_and_possible_by_group( + ht, + ht.possible_variants, + ht.observed_variants, + additional_grouping=vep_groupings + + ("exomes_coverage", "apply_model") + + tuple(f for f in MUTATION_TYPE_FIELDS if f not in MU_GROUPING), + partition_hint=partition_hint, + ) + + # Annotate with mutation rate and apply models on the aggregated counts. + ht = annotate_with_mu(ht, mutation_ht) + ht = _apply_constraint_models(ht, plateau_models, coverage_model, log10_coverage) + ht = ht.checkpoint(new_temp_file("aggregated_apply_models", "ht")) + + # Aggregate by VEP groupings only, summing model outputs across coverage and + # context groups to produce the same schema as aggregate_per_variant_expected_ht. + ht = ht.key_by().select(*vep_groupings, *AGGREGATE_SUM_FIELDS) + ht = ht.group_by(*vep_groupings).aggregate(**aggregate_constraint_metrics_expr(ht)) + + ht = _annotate_apply_models_globals( + ht, plateau_models, coverage_model, log10_coverage, list(vep_groupings) + ) + + return ht + + # TODO: Move this up after review in this location. def calculate_mu_by_downsampling( ht: hl.Table, From 2ecbb62367981cb11c868fe805c879f57aafb47e Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 18 Mar 2026 11:15:54 -0600 Subject: [PATCH 29/38] Update documentation for gnomad-constraint project, including enhancements to the README and CLAUDE.md files. Revise project overview, structure, and pipeline steps to reflect the current version 4.1.1. Add details on new command-line arguments and resource paths, while improving clarity and organization throughout the documentation. --- CLAUDE.md | 103 +++--------------------------------------------------- README.md | 95 ++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 99 insertions(+), 99 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index bad2f4e6..d044e294 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,10 +1,4 @@ -# gnomad-constraint Project Reference - -## Project Overview - -Genic constraint analysis pipeline for gnomAD. Computes observed/expected ratios, pLI scores, z-scores, and confidence intervals for LoF, missense, and synonymous variants at the gene/transcript level. Current version: **v4.1** (GRCh38). Historical version 2.1.1 (GRCh37) is also supported. - -The `gnomad_constraint/experimental/proemis3d/` directory contains the ProEmis3D project for regional missense constraint visualization. +# gnomad-constraint Claude Reference ## Code Style @@ -29,107 +23,20 @@ Use **Sphinx-style** (`:param:`, `:return:`) docstrings following the gnomad_met - For Hail expression parameters, use `hl.expr.StructExpression`, `hl.expr.BooleanExpression`, etc. - For Hail table/matrix types, use `hl.Table`, `hl.MatrixTable`. -## Project Structure - -| Directory | Purpose | -|-----------|---------| -| `gnomad_constraint/pipeline/constraint_pipeline.py` | Main constraint pipeline (7 steps) | -| `gnomad_constraint/pipeline/constraint_pipeline_complex.py` | Complex region constraint pipeline | -| `gnomad_constraint/utils/constraint.py` | Core utility functions (preprocessing, model building, metrics) | -| `gnomad_constraint/utils/constraint_complex.py` | Complex region constraint utilities | -| `gnomad_constraint/resources/resource_utils.py` | Resource paths, constants, `TableResource` definitions | -| `gnomad_constraint/experimental/proemis3d/` | ProEmis3D regional missense constraint | -| `gnomad_constraint/plots/` | R and Python plotting scripts | - -## Key Constants (`resource_utils.py`) +## Key Constants (`constants.py`) ```python -VERSIONS = ["2.1.1", "4.0", "4.1"] -CURRENT_VERSION = "4.1" +VERSIONS = ["2.1.1", "4.0", "4.1", "4.1.1"] +CURRENT_VERSION = "4.1.1" DATA_TYPES = ["context", "exomes", "genomes"] MODEL_TYPES = ["plateau", "coverage"] GENOMIC_REGIONS = ["autosome_par", "chrx_nonpar", "chry_nonpar"] POPS = ("global", "afr", "amr", "eas", "nfe", "sas") COVERAGE_CUTOFF = 40 CUSTOM_VEP_ANNOTATIONS = ["transcript_consequences", "worst_csq_by_gene"] +MU_GROUPING = ("context", "ref", "alt", "methylation_level") ``` -Note: `MU_GROUPING = ("context", "ref", "alt", "methylation_level")` is NOT in `resource_utils.py`. If you need it, define it locally. - -## Constraint Pipeline Steps - -The main pipeline (`constraint_pipeline.py`) has these steps (each a CLI flag): - -| Step | CLI Flag | Function | -|------|----------|----------| -| 1 | `--prepare-context-ht` | Annotate VEP context with methylation, coverage, GERP | -| 2 | `--preprocess-data` | Add VEP context annotations to exome/genome tables, prepare for constraint | -| 3 | `--calculate-gerp-cutoffs` | Optional: compute GERP percentile cutoffs | -| 4 | `--calculate-mutation-rate` | Compute baseline mutation rate per substitution/context | -| 5 | `--create-training-set` | Count observed + possible variants at synonymous sites | -| 6 | `--build-models` | Build plateau and coverage regression models | -| 7 | `--apply-models` | Apply models to compute expected variant counts and o/e ratios | -| 8 | `--compute-constraint-metrics` | Compute pLI, z-scores, o/e with CIs | -| 9 | `--export-tsv` | Export constraint metrics to TSV | - -### Key Resource Paths (v4.1) - -``` -gs://gnomad/v4.1/constraint/ # Production root -gs://gnomad-tmp/gnomad_v4.1_testing/constraint/ # Test root - -# Key outputs: -.../preprocessed_data/annotated_context.ht -.../preprocessed_data/gnomad.v4.1.{context|exomes|genomes}.preprocessed.{region}.ht -.../mutation_rate/gnomad.v4.1.mutation_rate.ht -.../training_data/gnomad.v4.1.constraint_training.{region}.ht -.../models/gnomad.v4.1.{plateau|coverage}.{region}.he -.../predicted_proportion_observed/transcript_consequences/gnomad.v4.1.predicted_proportion_observed.{region}.ht -.../metrics/gnomad.v4.1.constraint_metrics.ht -``` - -### constraint_metrics Table - -Keyed by `(gene, transcript, canonical)` (and optionally `mane_select`, `gene_id`). - -Output struct per annotation category (`lof`, `mis`, `syn`): -- `.obs` — observed variant count -- `.exp` — expected variant count -- `.oe` — observed/expected ratio -- `.oe_ci` — 90% CI around o/e -- `.z_raw` — raw z-score -- `.pLI` — probability of loss-of-function intolerance (LoF only) -- `.pNull`, `.pRec` — null/recessive probabilities (LoF only) - -## Missense Score Percentile Analysis - -`gnomad_constraint/plots/determine_missense_score_percentiles.py` computes per-percentile depletion of missense variants binned by missense prediction scores. - -### Scores analyzed - -ProteinMPNN, ESM, REVEL, RASP, AM, MisFit, PolyPhen, CPT1, popEVE, EVE, MPC, CADD, GPN-MSA - -### Pipeline steps (CLI flags) - -| Step | Flag | -|------|------| -| 1 | `--preprocess-scores` | -| 2 | `--compute-percentiles` | -| 3 | `--annotate-constraint-data` | -| 4 | `--aggregate-by-transcript` | -| 5 | `--compute-cumulative` | -| 6 | `--export-percentile-summary` | -| 7 | `--export-matched-plof-summary` | - -Step 7 computes matched pLoF o/e per missense percentile bin. It computes adj_r-corrected gene-level pLoF from the per-SNV table (`--constraint-ht-path`), not from the pre-computed constraint_metrics table (which lacks adj_r for pLoF). - -## Dependencies - -- **hail** — distributed genomics framework -- **numpy**, **pandas**, **scipy** — numerical/statistical -- **gnomad** (gnomad_methods) — shared gnomAD utilities -- **gnomad_qc** — gnomAD QC pipeline resources - ## Known Gotchas - **MU_GROUPING not exported**: `gnomad_constraint.resources.resource_utils` does NOT export `MU_GROUPING`. It must be defined locally as `("context", "ref", "alt", "methylation_level")`. diff --git a/README.md b/README.md index 3ecdf7e8..92263fb1 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,105 @@ # gnomad-constraint -This repo contains code for generating constraint metrics for gnomAD. An overview of the pipeline and functions that are used can be found in [/gnomad_constraint/flowchart/constraint_pipeline_v4.pdf](./gnomad_constraint/flowchart/constraint_pipeline_v4.pdf). Note that many functions are imported from the [gnomad_methods](https://github.com/broadinstitute/gnomad_methods) repo. +Genic constraint analysis pipeline for gnomAD. Computes observed/expected ratios, pLI scores, z-scores, and confidence intervals for LoF, missense, and synonymous variants at the gene/transcript level. Current version: **v4.1.1** (GRCh38). Historical version 2.1.1 (GRCh37) is also supported. +An overview of the pipeline and functions that are used can be found in [/gnomad_constraint/flowchart/constraint_pipeline_v4.pdf](./gnomad_constraint/flowchart/constraint_pipeline_v4.pdf). Note that many functions are imported from the [gnomad_methods](https://github.com/broadinstitute/gnomad_methods) repo. +The `gnomad_constraint/experimental/proemis3d/` directory contains the ProEmis3D project for regional missense constraint visualization. +## Project Structure +| Directory | Purpose | +|-----------|---------| +| `gnomad_constraint/pipeline/constraint_pipeline.py` | Main constraint pipeline | +| `gnomad_constraint/pipeline/constraint_pipeline_complex.py` | Complex region constraint pipeline | +| `gnomad_constraint/utils/constraint.py` | Core utility functions (preprocessing, model building, metrics) | +| `gnomad_constraint/utils/constraint_complex.py` | Complex region constraint utilities | +| `gnomad_constraint/resources/resource_utils.py` | Resource paths and `TableResource` definitions | +| `gnomad_constraint/resources/constants.py` | Pipeline constants | +| `gnomad_constraint/experimental/proemis3d/` | ProEmis3D regional missense constraint | +| `gnomad_constraint/plots/` | R and Python plotting scripts | +## Constraint Pipeline Steps +The main pipeline (`constraint_pipeline.py`) has these steps (each a CLI flag): +| Step | CLI Flag | Function | +|------|----------|----------| +| 1 | `--prepare-context-ht` | Annotate VEP context with methylation, coverage, GERP | +| 2 | `--preprocess-data` | Add VEP context annotations to exome/genome tables, prepare for constraint | +| 3 | `--calculate-gerp-cutoffs` | Optional: compute GERP percentile cutoffs | +| 4 | `--calculate-mutation-rate` | Compute baseline mutation rate per substitution/context | +| 5 | `--create-training-set` | Count observed + possible variants at synonymous sites | +| 6 | `--build-models` | Build plateau and coverage regression models | +| 7a | `--apply-models-per-variant` | Per-variant: apply models to compute expected counts per variant | +| 7b | `--apply-models-aggregated` | Aggregated: aggregate counts first, then apply models (alternative to 7a+8+9) | +| 8 | `--aggregate-per-variant-expected` | Aggregate per-variant expected counts by transcript/consequence | +| 9 | `--aggregate-by-constraint-groups` | Group aggregated counts into constraint groups (lof, mis, syn) | +| 10 | `--compute-gene-quality-metrics` | Compute per-transcript gene quality metrics (coverage, MQ, segdup, LCR) | +| 11 | `--compute-constraint-metrics` | Compute pLI, z-scores, o/e with CIs | +| 12 | `--prepare-release` | Format constraint metrics for public release | +| 13 | `--export-release-tsv` | Export release table to TSV | +| 14 | `--export-release-downsampling-tsv` | Export downsampling constraint metrics to TSV | +Steps 7a/8/9 (per-variant path) and 7b (aggregated path) are alternative ways to go from models to constraint groups. The per-variant path applies coverage correction per-variant then aggregates; the aggregated path aggregates counts first then applies models. Use `--use-aggregated-expected` with `--aggregate-by-constraint-groups` to read from the aggregated path output instead of the per-variant path. +### Key Resource Paths (v4.1.1) + +``` +gs://gnomad/v4.1.1/constraint/ # Production root +gs://gnomad-tmp/gnomad_v4.1.1_testing/constraint/ # Test root + +# Key outputs: +.../preprocessed_data/annotated_context.ht +.../preprocessed_data/gnomad.v4.1.1.{context|exomes|genomes}.preprocessed.{region}.ht +.../mutation_rate/gnomad.v4.1.1.mutation_rate.ht +.../training_data/gnomad.v4.1.1.constraint_training.{region}.ht +.../models/gnomad.v4.1.1.{plateau|coverage}.{region}.he +.../apply_models/transcript_consequences/gnomad.v4.1.1.per_variant_expected.{region}.ht +.../apply_models/transcript_consequences/gnomad.v4.1.1.aggregated_expected.{region}.ht +.../apply_models/transcript_consequences/gnomad.v4.1.1.apply.{region}.ht +.../constraint_groups/transcript_consequences/gnomad.v4.1.1.constraint_groups.{region}.ht +.../metrics/gnomad.v4.1.1.constraint_metrics.ht +``` + +### constraint_metrics Table + +Keyed by `(gene, transcript, canonical)` (and optionally `mane_select`, `gene_id`). + +Output struct per annotation category (`lof`, `mis`, `syn`): +- `.obs` — observed variant count +- `.exp` — expected variant count +- `.oe` — observed/expected ratio +- `.oe_ci` — 90% CI around o/e +- `.z_raw` — raw z-score +- `.pLI` — probability of loss-of-function intolerance (LoF only) +- `.pNull`, `.pRec` — null/recessive probabilities (LoF only) + +## Missense Score Percentile Analysis + +`gnomad_constraint/plots/determine_missense_score_percentiles.py` computes per-percentile depletion of missense variants binned by missense prediction scores. + +### Scores analyzed + +ProteinMPNN, ESM, REVEL, RASP, AM, MisFit, PolyPhen, CPT1, popEVE, EVE, MPC, CADD, GPN-MSA + +### Pipeline steps (CLI flags) + +| Step | Flag | +|------|------| +| 1 | `--preprocess-scores` | +| 2 | `--compute-percentiles` | +| 3 | `--annotate-constraint-data` | +| 4 | `--aggregate-by-transcript` | +| 5 | `--compute-cumulative` | +| 6 | `--export-percentile-summary` | +| 7 | `--export-matched-plof-summary` | + +Step 7 computes matched pLoF o/e per missense percentile bin. It computes adj_r-corrected gene-level pLoF from the per-SNV table (`--constraint-ht-path`), not from the pre-computed constraint_metrics table (which lacks adj_r for pLoF). + +## Dependencies + +- **hail** — distributed genomics framework +- **numpy**, **pandas**, **scipy** — numerical/statistical +- **gnomad** (gnomad_methods) — shared gnomAD utilities +- **gnomad_qc** — gnomAD QC pipeline resources From 9833ec3db57dc2070b00401c99d8121fbf2d9485 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 18 Mar 2026 11:18:43 -0600 Subject: [PATCH 30/38] Add PCSK9 regions to filtering criteria in resource utility functions. Update keep_regions for both GRCh37 and alternative chromosome formats to include specific genomic intervals for PCSK9. --- gnomad_constraint/resources/resource_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index b0a967a1..26ba314a 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -587,6 +587,7 @@ def filter_for_test( if use_gene_list: if rg == "GRCh37": keep_regions = [ + "1:55505149-55530526", # PCSK9 "20:49505585-49547958", # ADNP "20:853296-896977", # ANGPT4 "X:13752832-13787480", # OFD1 @@ -595,6 +596,7 @@ def filter_for_test( ] else: keep_regions = [ + "chr1:55039447-55064852", # PCSK9 "chr20:50888916-50931437", # ADNP "chr20:869900-916334", # ANGPT4 "chrX:13734743-13777955", # OFD1 From d2d21db74eca7db495d52aa4c6c76c797b49eebc Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Wed, 18 Mar 2026 13:23:59 -0600 Subject: [PATCH 31/38] Enhance gnomAD constraint pipeline by updating resource handling and filtering. Exclude non-Python files from zip builds to reduce size. Introduce a mapping for sites versioning and refactor gene quality metrics computation to utilize updated data structures. Improve clarity in utility functions and adjust documentation accordingly. --- CLAUDE.md | 5 ++- .../pipeline/constraint_pipeline.py | 12 +++++- gnomad_constraint/resources/constants.py | 8 ++++ gnomad_constraint/resources/resource_utils.py | 38 ++++++++++--------- gnomad_constraint/utils/constraint.py | 25 +++++++++--- 5 files changed, 61 insertions(+), 27 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index d044e294..c8c440cf 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -68,11 +68,12 @@ MU_GROUPING = ("context", "ref", "alt", "methylation_level") ```bash # Build a single zip with correct top-level package structure +# Exclude non-Python files (.RData, renv/, images, notebooks) to keep zip small cd && \ rm -f /tmp/pyfiles.zip && \ - zip -r /tmp/pyfiles.zip gnomad_constraint/ -x '*.pyc' '*__pycache__*' && \ + zip -r /tmp/pyfiles.zip gnomad_constraint/ -x '*.pyc' '*__pycache__*' '*.RData' '*/renv/*' '*.DS_Store' '*.png' '*.pdf' '*.ipynb' && \ cd && \ - zip -r /tmp/pyfiles.zip gnomad_qc/ -x '*.pyc' '*__pycache__*' '*.DS_Store' + zip -r /tmp/pyfiles.zip gnomad_qc/ -x '*.pyc' '*__pycache__*' '*.DS_Store' '*.ipynb' # Submit to cluster (single zip = used directly, not repackaged) hailctl dataproc submit \ diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 2fd480d1..d443b814 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -149,9 +149,17 @@ def main(args): res.check_resource_existence() gencode_cds_ht = constraint_res.get_gencode_cds_ht(version).ht() + exomes_sites_ht = res.exomes_sites_ht.ht() + if test: + gencode_cds_ht = filter_for_test( + gencode_cds_ht, use_gene_list=test_gene_list + ) + exomes_sites_ht = filter_for_test( + exomes_sites_ht, use_gene_list=test_gene_list + ) gene_quality_ht = compute_gene_quality_metrics( - res.annotated_context_ht.ht(), - res.exomes_sites_ht.ht(), + res.temp_preprocess_data_ht.ht(), + exomes_sites_ht, gencode_cds_ht, ) gene_quality_ht.write(res.gene_quality_metrics_ht.path, overwrite=overwrite) diff --git a/gnomad_constraint/resources/constants.py b/gnomad_constraint/resources/constants.py index 3559c771..2b08b7e0 100644 --- a/gnomad_constraint/resources/constants.py +++ b/gnomad_constraint/resources/constants.py @@ -13,6 +13,14 @@ CURRENT_VERSION = "4.1.1" """Current default gnomAD constraint pipeline version.""" +SITES_VERSION_MAP = { + "2.1.1": "2.1.1", + "4.0": "4.0", + "4.1": "4.1", + "4.1.1": "4.1", +} +"""Map from constraint pipeline version to gnomAD sites release version.""" + DATA_TYPES = ["context", "exomes", "genomes"] """Data types used in the constraint pipeline.""" diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 26ba314a..2796c51a 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -28,6 +28,7 @@ DATA_TYPES, EXTENSIONS, MODEL_TYPES, + SITES_VERSION_MAP, VERSIONS, ) @@ -105,15 +106,16 @@ def get_sites_resource(data_type: str, version: str = CURRENT_VERSION) -> BaseRe :return: Genome or exomes sites Table. """ build = check_param_scope(version=version, data_type=data_type) + sites_version = SITES_VERSION_MAP[version] if build == "GRCh37": - return gnomad_grch37.public_release(data_type).versions[version] + return gnomad_grch37.public_release(data_type).versions[sites_version] elif int(version[0]) == 4: # Continue to use v3.1.2 for genomes as downsamplings are dropped in v4 # versions. if data_type == "genomes": return gnomad_grch38.public_release(data_type).versions["3.1.2"] else: - return gnomad_grch38.public_release(data_type).versions[version] + return gnomad_grch38.public_release(data_type).versions[sites_version] else: raise ValueError( "The sites resource has not been defined for the specified version!" @@ -384,7 +386,9 @@ def get_annotated_context_ht(**kwargs) -> TableResource: :return: TableResource of annotated context Table. """ - return get_constraint_data("annotated_context", temp=True, **kwargs) + return get_constraint_data( + "annotated_context", sub_dir="preprocessed_data", **kwargs + ) def get_preprocessed_ht(**kwargs) -> TableResource: @@ -628,13 +632,13 @@ def get_adj_r_ht() -> hl.Table: def get_syn_adj_r_ht() -> hl.Table: """ - Read the synonymous DNM adj_r per-context methylation genome 1kb autosome Table. + Read the aggregated synonymous DNM adj_r per-context methylation genome 1kb autosome Table. - :return: Table with syn_adj_r annotation keyed by locus. + :return: Table with adj_r dict annotation keyed by interval. """ return hl.read_table( "gs://gnomad/v4.1/constraint/resources/annotations/ht/" - "adj_r_syn_dnm_per_context_methyl_genome_1kb_autosome.ht" + "adj_r_syn_dnm_per_context_methyl_genome_1kb_autosome.agg.ht" ) @@ -701,6 +705,16 @@ def get_constraint_resources( }, pipeline_input_steps=[prepare_context], ) + compute_gene_quality_metrics_step = PipelineStepResourceCollection( + "--compute-gene-quality-metrics", + output_resources={ + "gene_quality_metrics_ht": get_gene_quality_metrics_ht(version=version) + }, + add_input_resources={ + "gnomAD resources": {"exomes_sites_ht": input_hts["exomes_sites_ht"]}, + }, + pipeline_input_steps=[preprocess_data], + ) calculate_gerp_cutoffs = PipelineStepResourceCollection( "--calculate-gerp-cutoffs", output_resources={}, @@ -771,16 +785,6 @@ def get_constraint_resources( }, pipeline_input_steps=[preprocess_data, calculate_mutation_rate, build_models], ) - compute_gene_quality_metrics_step = PipelineStepResourceCollection( - "--compute-gene-quality-metrics", - output_resources={ - "gene_quality_metrics_ht": get_gene_quality_metrics_ht(version=version) - }, - input_resources={ - "gnomAD resources": {"exomes_sites_ht": input_hts["exomes_sites_ht"]}, - }, - pipeline_input_steps=[prepare_context], - ) compute_constraint_metrics = PipelineStepResourceCollection( "--compute-constraint-metrics", output_resources={ @@ -815,8 +819,8 @@ def get_constraint_resources( constraint_pipeline.add_steps( { "prepare_context": prepare_context, - "compute_gene_quality_metrics": compute_gene_quality_metrics_step, "preprocess_data": preprocess_data, + "compute_gene_quality_metrics": compute_gene_quality_metrics_step, "calculate_gerp_cutoffs": calculate_gerp_cutoffs, "calculate_mutation_rate": calculate_mutation_rate, "create_training_set": create_training_set, diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index cea31fc9..4d9c9337 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -334,8 +334,10 @@ def get_annotations_for_computing_mu( # Filter so that the most severe annotation is 'intron_variant' or # 'intergenic_variant'. keep_expr &= hl.any( - most_severe_consequence_expr == c - for c in ["intron_variant", "intergenic_variant"] + [ + most_severe_consequence_expr == c + for c in ["intron_variant", "intergenic_variant"] + ] ) # Set up the criteria to keep high-quality sites, and sites found in less than or @@ -1039,11 +1041,13 @@ def create_aggregated_expected_ht( # rate context. The additional_grouping includes VEP groupings (gene, transcript, # annotation, etc.) plus the apply_model fields needed for model application. vep_groupings = tuple(g for g in groupings if g not in MU_GROUPING) + mu_extra = tuple(f for f in MU_GROUPING if f not in ("context", "ref", "alt")) ht = count_observed_and_possible_by_group( ht, ht.possible_variants, ht.observed_variants, additional_grouping=vep_groupings + + mu_extra + ("exomes_coverage", "apply_model") + tuple(f for f in MUTATION_TYPE_FIELDS if f not in MU_GROUPING), partition_hint=partition_hint, @@ -1054,6 +1058,15 @@ def create_aggregated_expected_ht( ht = _apply_constraint_models(ht, plateau_models, coverage_model, log10_coverage) ht = ht.checkpoint(new_temp_file("aggregated_apply_models", "ht")) + # Scale per-subgroup rate/factor fields by possible_variants so that summing + # across subgroups matches the per-variant path (which sums one copy per variant). + ht = ht.annotate( + mu_snp=ht.mu_snp * ht.possible_variants, + predicted_proportion_observed=ht.predicted_proportion_observed + * ht.possible_variants, + coverage_correction=ht.coverage_correction * ht.possible_variants, + ) + # Aggregate by VEP groupings only, summing model outputs across coverage and # context groups to produce the same schema as aggregate_per_variant_expected_ht. ht = ht.key_by().select(*vep_groupings, *AGGREGATE_SUM_FIELDS) @@ -1364,7 +1377,7 @@ def _compute_site_quality_metrics( def compute_gene_quality_metrics( - context_ht: hl.Table, + ht: hl.Table, exomes_ht: hl.Table, gencode_cds_ht: hl.Table, an_coverage_threshold: int = 90, @@ -1381,8 +1394,8 @@ def compute_gene_quality_metrics( when mean AS_MQ < 50, ``low_exome_coverage`` when prop_bp_AN90 < 0.1). - :param context_ht: Preprocessed context Hail Table with - ``exomes_coverage`` (AN percent, 0-100) per position. + :param ht: Preprocessed constraint Hail Table with ``exomes_coverage`` + field (output of :func:`prepare_ht_for_constraint_calculations`). :param exomes_ht: gnomAD exomes sites Hail Table. :param gencode_cds_ht: GENCODE CDS positions Table keyed by locus with ``transcript_id`` array (output of @@ -1393,7 +1406,7 @@ def compute_gene_quality_metrics( and ``gene_flags``. """ an90_ht = _compute_coverage_metrics( - context_ht, gencode_cds_ht, an_coverage_threshold + ht, gencode_cds_ht, an_coverage_threshold ).cache() sites_ht = _compute_site_quality_metrics(exomes_ht, gencode_cds_ht).cache() From 64229acd8df157c6960b7acfd751e0a7e1bf7b4a Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Mon, 23 Mar 2026 10:55:52 -0600 Subject: [PATCH 32/38] Refactor gene quality metrics computation in the constraint pipeline to include segmental duplication and low-complexity region metrics. Update related utility functions for improved clarity and efficiency. Enhance resource handling by integrating new genomic interval data for more accurate quality assessments. --- .../pipeline/constraint_pipeline.py | 51 +-- gnomad_constraint/resources/constants.py | 47 ++- gnomad_constraint/utils/constraint.py | 298 +++++++++++++++--- 3 files changed, 306 insertions(+), 90 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index d443b814..f0b76ee5 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -25,6 +25,7 @@ import logging import hail as hl +from gnomad.resources.grch38.reference_data import lcr_intervals, seg_dup_intervals from gnomad.utils.constraint import ( build_models, calculate_gerp_cutoffs, @@ -143,28 +144,6 @@ def main(args): logger.info("Done annotating the VEP context Table.") - if args.compute_gene_quality_metrics: - logger.info("Computing per-transcript gene quality metrics...") - res = resources.compute_gene_quality_metrics - res.check_resource_existence() - - gencode_cds_ht = constraint_res.get_gencode_cds_ht(version).ht() - exomes_sites_ht = res.exomes_sites_ht.ht() - if test: - gencode_cds_ht = filter_for_test( - gencode_cds_ht, use_gene_list=test_gene_list - ) - exomes_sites_ht = filter_for_test( - exomes_sites_ht, use_gene_list=test_gene_list - ) - gene_quality_ht = compute_gene_quality_metrics( - res.temp_preprocess_data_ht.ht(), - exomes_sites_ht, - gencode_cds_ht, - ) - gene_quality_ht.write(res.gene_quality_metrics_ht.path, overwrite=overwrite) - logger.info("Done computing gene quality metrics.") - if args.calculate_gerp_cutoffs: logger.warning( "Calculating new GERP cutoffs to be used instead of" @@ -212,6 +191,30 @@ def main(args): logger.info("Done preprocessing the context Table.") + if args.compute_gene_quality_metrics: + logger.info("Computing per-transcript gene quality metrics...") + res = resources.compute_gene_quality_metrics + res.check_resource_existence() + + gencode_cds_ht = constraint_res.get_gencode_cds_ht(version).ht() + exomes_sites_ht = res.exomes_sites_ht.ht() + if test: + gencode_cds_ht = filter_for_test( + gencode_cds_ht, use_gene_list=test_gene_list + ) + exomes_sites_ht = filter_for_test( + exomes_sites_ht, use_gene_list=test_gene_list + ) + gene_quality_ht = compute_gene_quality_metrics( + res.temp_preprocess_data_ht.ht(), + exomes_sites_ht, + gencode_cds_ht, + seg_dup_intervals.ht(), + lcr_intervals.ht(), + ) + gene_quality_ht.write(res.gene_quality_metrics_ht.path, overwrite=overwrite) + logger.info("Done computing gene quality metrics.") + if args.calculate_mutation_rate: logger.info("Calculating mutation rate...") res = resources.calculate_mutation_rate @@ -404,14 +407,14 @@ def main(args): release_ht = prepare_release_ht( constraint_ht, release_version=args.release_version, - ) + ).naive_coalesce(1000) release_ht.write(res.release_ht.path, overwrite=overwrite) logger.info("Done preparing release Table.") if args.export_release_tsv or args.export_release_downsampling_tsv: res = resources.export_release_tsv res.check_resource_existence() - release_ht = hl.read_table(res.release_ht.path) + release_ht = res.release_ht.ht() if args.export_release_tsv: logger.info("Exporting release TSV...") diff --git a/gnomad_constraint/resources/constants.py b/gnomad_constraint/resources/constants.py index 2b08b7e0..c83c4216 100644 --- a/gnomad_constraint/resources/constants.py +++ b/gnomad_constraint/resources/constants.py @@ -120,6 +120,28 @@ } """GENCODE field renames applied when preparing the release Table.""" +# --------------------------------------------------------------------------- +# Constraint percentile threshold computation and annotation +# --------------------------------------------------------------------------- + +PLI_EXPECTED_VALUES = {"Null": 1.0, "Rec": 0.706, "LI": 0.207} +"""Expected o/e values for the pLI model (null, recessive, loss-of-function intolerant).""" + +CONSTRAINT_METRICS = ["lof", "mis", "syn"] +"""Constraint metrics for which percentile thresholds are computed.""" + +CONSTRAINT_GRANULARITIES = { + "percentile": list(range(1, 100)), + "decile": list(range(1, 10)), + "sextile": list(range(1, 6)), +} +""" +Granularities for percentile binning. + +Keys are granularity names; values are boundary bin labels (1-indexed). +Quantile probabilities are bin / (max_bin + 1). +""" + # --------------------------------------------------------------------------- # Release format constants # --------------------------------------------------------------------------- @@ -215,28 +237,3 @@ "gen_anc_exp", ] """Fields selected from the release constraint-group struct, in display order.""" - -# --------------------------------------------------------------------------- -# Constraint percentile threshold computation and annotation -# --------------------------------------------------------------------------- - -CONSTRAINT_SCORE_CAP = 2.0 -"""OE upper CI value above which scores are considered unconstrained and capped.""" - -PLI_EXPECTED_VALUES = {"Null": 1.0, "Rec": 0.706, "LI": 0.207} -"""Expected o/e values for the pLI model (null, recessive, loss-of-function intolerant).""" - -CONSTRAINT_METRICS = ["lof", "mis", "syn"] -"""Constraint metrics for which percentile thresholds are computed.""" - -CONSTRAINT_GRANULARITIES = { - "percentile": list(range(1, 100)), - "decile": list(range(1, 10)), - "sextile": list(range(1, 6)), -} -""" -Granularities for percentile binning. - -Keys are granularity names; values are boundary bin labels (1-indexed). -Quantile probabilities are bin / (max_bin + 1). -""" diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 4d9c9337..50d07c15 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -967,6 +967,10 @@ def aggregate_per_variant_expected_ht( key annotations in the grouping. Default is False. :return: Table with the observed and expected counts. """ + # Build the grouping key: optionally include mutation rate annotations + # (context, ref, alt, methylation_level) for finer-grained output, + # plus the VEP-derived groupings (annotation, gene, transcript, etc.) + # stored in the apply_models globals. groupings = [ *(MU_GROUPING if include_mu_annotations_in_grouping else []), *[ @@ -975,13 +979,25 @@ def aggregate_per_variant_expected_ht( if g not in MU_GROUPING ], ] + + # The per-variant table from create_per_variant_expected_ht nests + # transcript-level fields (gene, transcript, canonical, annotation, etc.) + # inside a `calibrate_mu` struct. Promote them to top-level so they can + # be used as grouping keys. if "calibrate_mu" in ht.row: ht = ht.annotate(**ht.calibrate_mu) + # Keep only coding consequences (e.g. synonymous, missense, LoF) — drops + # non-coding VEP annotations to improve computation time and memory. ht = ht.filter(hl.set(CSQ_CODING).contains(ht.annotation)) + + # Narrow to just the grouping keys and the fields we need to sum, + # dropping everything else to minimize shuffle size. ht = ht.key_by().select(*groupings, *AGGREGATE_SUM_FIELDS) ht = ht.checkpoint(new_temp_file("pre_aggregation", "ht")) + # Sum observed_variants, expected_variants, possible_variants, mu_snp, + # etc. within each (transcript, consequence, ...) group. ht = ht.group_by(*groupings).aggregate(**aggregate_constraint_metrics_expr(ht)) ht = ht.checkpoint(new_temp_file("post_aggregation", "ht")) @@ -1112,7 +1128,6 @@ def calculate_mu_by_downsampling( ht.compute_mu.observed_variants, additional_grouping=additional_grouping, ) - ht = ht.checkpoint(new_temp_file(prefix="constraint", extension="ht")) total_bases = ht.aggregate(hl.agg.sum(ht.possible_variants)) // 3 @@ -1209,6 +1224,7 @@ def add_oe_upper_rank_and_bins( t, use_mane_select_over_canonical, mane_select_only ), bin_granularities=bin_granularities, + rank_field_prefix="upper_", ) @@ -1349,30 +1365,65 @@ def _compute_site_quality_metrics( ht: hl.Table, gencode_cds_ht: hl.Table, ) -> hl.Table: - """Compute per-transcript mapping quality and region flag metrics. - - Computes mean AS_MQ, proportion of sites in segmental duplications, and - proportion of sites in low-complexity regions from variant sites within - CDS regions. + """Compute per-transcript mean mapping quality from SNV sites in CDS. :param ht: gnomAD exomes sites Hail Table. :param gencode_cds_ht: GENCODE CDS positions Table keyed by locus with ``transcript_id`` array. - :return: Table keyed by ``transcript`` with ``mean_AS_MQ``, - ``prop_segdup``, and ``prop_LCR``. + :return: Table keyed by ``transcript`` with ``mean_AS_MQ``. """ - # Deduplicate sites by locus. - ht = ht.select("region_flags", AS_MQ=ht.info.AS_MQ) - ht = ht.key_by("locus").select("AS_MQ", "region_flags").distinct() + # Extract allele-specific mapping quality from the sites info struct. + ht = ht.select(AS_MQ=ht.info.AS_MQ) + + # Restrict to SNVs — indels don't have comparable AS_MQ values. + ht = ht.filter(hl.is_snp(ht.alleles[0], ht.alleles[1])) - # Join with GENCODE CDS to get per-locus transcript IDs. + # Look up which transcripts overlap each variant's locus via the + # GENCODE CDS table, then keep only sites inside annotated CDS regions. + # Explode so each (locus, transcript) pair is a separate row. ht = ht.annotate(transcript_id=gencode_cds_ht[ht.locus].transcript_id) ht = ht.filter(hl.is_defined(ht.transcript_id)).explode("transcript_id").cache() + # Average AS_MQ across all SNV sites in each transcript's CDS. + # Used downstream to flag transcripts with low mapping quality. return ht.group_by(transcript=ht.transcript_id).aggregate( mean_AS_MQ=hl.agg.mean(ht.AS_MQ), - prop_segdup=hl.agg.fraction(ht.region_flags.segdup), - prop_LCR=hl.agg.fraction(ht.region_flags.lcr), + ) + + +def _compute_region_flag_metrics( + gencode_cds_ht: hl.Table, + seg_dup_intervals_ht: hl.Table, + lcr_intervals_ht: hl.Table, +) -> hl.Table: + """Compute per-transcript fraction of CDS bases in segdup and LCR regions. + + :param gencode_cds_ht: GENCODE CDS positions Table keyed by locus with + ``transcript_id`` array. + :param seg_dup_intervals_ht: Segmental duplication intervals Table keyed + by locus/interval. + :param lcr_intervals_ht: Low-complexity region intervals Table keyed by + locus/interval. + :return: Table keyed by ``transcript`` with ``prop_segdup`` and + ``prop_LCR``. + """ + # For each CDS base position, check whether it falls within a segmental + # duplication or low-complexity region interval. The interval tables are + # keyed by locus, so a defined lookup means the base is inside the region. + ht = gencode_cds_ht.annotate( + in_segdup=hl.is_defined(seg_dup_intervals_ht[gencode_cds_ht.locus]), + in_lcr=hl.is_defined(lcr_intervals_ht[gencode_cds_ht.locus]), + ) + + # Explode so each (locus, transcript) pair is a separate row, then + # compute the fraction of CDS bases in each region per transcript. + # These proportions are used downstream to flag transcripts with high + # segdup/LCR overlap. + ht = ht.explode("transcript_id").cache() + + return ht.group_by(transcript=ht.transcript_id).aggregate( + prop_segdup=hl.agg.fraction(ht.in_segdup), + prop_LCR=hl.agg.fraction(ht.in_lcr), ) @@ -1380,19 +1431,22 @@ def compute_gene_quality_metrics( ht: hl.Table, exomes_ht: hl.Table, gencode_cds_ht: hl.Table, + seg_dup_intervals_ht: hl.Table, + lcr_intervals_ht: hl.Table, an_coverage_threshold: int = 90, ) -> hl.Table: """Compute per-transcript gene quality metrics. - Combines coverage metrics from :func:`_compute_coverage_metrics` and - site quality metrics from :func:`_compute_site_quality_metrics` into a + Combines coverage metrics from :func:`_compute_coverage_metrics`, + mapping quality from :func:`_compute_site_quality_metrics`, and + region flag metrics from :func:`_compute_region_flag_metrics` into a single Table with release-ready fields: - - ``gene_quality_metrics``: struct with ``exome_prop_bp_AN90``, - ``exome_mean_AS_MQ``, ``exome_prop_segdup``, ``exome_prop_LCR``. - - ``gene_flags``: set of flag strings (``low_exome_mapping_quality`` - when mean AS_MQ < 50, ``low_exome_coverage`` when - prop_bp_AN90 < 0.1). + - ``gene_quality_metrics``: struct with ``exome_prop_bp_AN90``, + ``exome_mean_AS_MQ``, ``exome_prop_segdup``, ``exome_prop_LCR``. + - ``gene_flags``: set of flag strings (``low_exome_mapping_quality`` + when mean AS_MQ < 50, ``low_exome_coverage`` when + prop_bp_AN90 < 0.1). :param ht: Preprocessed constraint Hail Table with ``exomes_coverage`` field (output of :func:`prepare_ht_for_constraint_calculations`). @@ -1400,17 +1454,36 @@ def compute_gene_quality_metrics( :param gencode_cds_ht: GENCODE CDS positions Table keyed by locus with ``transcript_id`` array (output of :func:`~gnomad_constraint.resources.resource_utils.get_gencode_cds_ht`). + :param seg_dup_intervals_ht: Segmental duplication intervals Table. + :param lcr_intervals_ht: Low-complexity region intervals Table. :param an_coverage_threshold: Minimum ``exomes_coverage`` value (0-100) for a position to be considered adequately covered. Default is 90. :return: Table keyed by ``transcript`` with ``gene_quality_metrics`` and ``gene_flags``. """ - an90_ht = _compute_coverage_metrics( - ht, gencode_cds_ht, an_coverage_threshold - ).cache() + # Compute three independent per-transcript metric tables, each keyed by + # transcript. These use different source data (preprocessed context HT + # for coverage, exomes sites HT for mapping quality, GENCODE CDS + + # interval tables for region overlap) so they can run in parallel. + an90_ht = _compute_coverage_metrics(ht, gencode_cds_ht, an_coverage_threshold) + an90_ht = an90_ht.cache() sites_ht = _compute_site_quality_metrics(exomes_ht, gencode_cds_ht).cache() + region_ht = _compute_region_flag_metrics( + gencode_cds_ht, seg_dup_intervals_ht, lcr_intervals_ht + ).cache() + + # Join the three metric tables on transcript into a single row per + # transcript with all four metrics. + ht = an90_ht.annotate( + **sites_ht[an90_ht.transcript], + **region_ht[an90_ht.transcript], + ) - ht = an90_ht.annotate(**sites_ht[an90_ht.transcript]) + # Package all metrics into a single `gene_quality_metrics` struct, + # prefixing each field with "exome_" for the release schema. Then + # derive `gene_flags` — a set of string flags for transcripts that + # fail quality thresholds (mean AS_MQ < 50 or < 10% of CDS bases + # passing allele number 90th percentile). ht = ht.select( gene_quality_metrics=hl.struct(**{f"exome_{f}": ht[f] for f in ht.row_value}), gene_flags=add_filters_expr( @@ -1445,6 +1518,14 @@ def _annotate_oe_ci_z( ``"syn"``) to ``(lower, upper)`` raw z-score outlier thresholds. :return: Table with OE, CI, z-score, and flag annotations. """ + # For every constraint group (syn, mis, lof_hc, lof_hc_lc) and every + # frequency slice within each group (adj, per-genetic-ancestry + # downsamplings), compute: + # - oe: observed/expected ratio (null when expected is 0) + # - oe_ci_discretized_poisson: confidence interval via discretized Poisson + # - oe_ci_gamma: confidence interval via gamma distribution (used for + # ranking in the release; more stable at low counts) + # - z_raw: raw z-score measuring deviation from expected ht = ht.annotate( constraint_groups=ht.constraint_groups.map( lambda x: x.annotate( @@ -1472,6 +1553,11 @@ def _annotate_oe_ci_z( ) ) + # Add per-group flags based on the adj-frequency z-score. Each group + # gets flags like "no_exp_{csq}" (expected == 0) or "z_raw_{csq}" + # (raw z-score outside the outlier thresholds). The z_thresholds dict + # maps category names (lof, mis, syn) to (lower, upper) bounds; both + # LoF groups (hc and hc_lc) use the "lof" thresholds. meta = hl.eval(ht.constraint_group_meta) freq_meta = hl.eval(ht.exomes_freq_meta) all_freq_idx = freq_meta.index(ADJ_FREQ_META) @@ -1512,6 +1598,8 @@ def _compute_z_scores(ht: hl.Table) -> hl.Table: :param ht: Table output by :func:`_annotate_oe_ci_z`. :return: Table with ``z_score`` and ``constraint_flags`` annotations. """ + # Resolve constraint group indices at Python time so we can reference + # specific groups (syn, mis, lof_hc) by position in the array. meta = hl.eval(ht.constraint_group_meta) freq_meta = hl.eval(ht.exomes_freq_meta) syn_idx = meta.index({"csq_set": "syn"}) @@ -1519,6 +1607,12 @@ def _compute_z_scores(ht: hl.Table) -> hl.Table: lof_idx = meta.index({"lof": "hc"}) all_freq_idx = freq_meta.index(ADJ_FREQ_META) + # Compute the standard deviation of raw z-scores across all transcripts + # (excluding those with no variants). This produces one SD per constraint + # group, stored as a global array parallel to constraint_groups. + # For non-synonymous groups, negative z-scores are mirrored to build a + # symmetric distribution (constrained genes skew the left tail); for syn, + # the distribution is already roughly symmetric so no mirroring is needed. ht = ht.annotate_globals( sd_raw_z=ht.aggregate( hl.agg.filter( @@ -1535,6 +1629,10 @@ def _compute_z_scores(ht: hl.Table) -> hl.Table: ) ) + # Normalize each group's raw z-score by its SD to produce the final + # z_score. Also union the per-group flags from syn, mis, and lof_hc + # into a single transcript-level constraint_flags set (used downstream + # to exclude outliers from percentile threshold computation). ht = ht.annotate( constraint_groups=hl.map( lambda x, sd_raw_z: x.annotate( @@ -1570,16 +1668,27 @@ def _compute_percentile_bins( transcripts for filtering when determining ranks. Default is True. :return: Table with rank, decile, and percentile bin annotations. """ + # Assign each transcript a dense rank (0-based) for its gamma OE upper + # CI within each constraint group, plus rank-based decile/sextile bins. + # Only MANE Select (or canonical) transcripts are ranked. ht = add_oe_upper_rank_and_bins(ht, use_mane_select_over_canonical) + # Map the three metric categories we compute thresholds for to their + # indices in the constraint_groups array. meta = hl.eval(ht.constraint_group_meta) metric_group_idx = { "syn": next(i for i, m in enumerate(meta) if m == {"csq_set": "syn"}), "mis": next(i for i, m in enumerate(meta) if m == {"csq_set": "mis"}), "lof": next(i for i, m in enumerate(meta) if m == {"lof": "hc"}), } + # Transcripts with any constraint flag are excluded from threshold + # computation (but still assigned bins afterward). outlier_expr = ht.constraint_flags.length() > 0 + # Convert CONSTRAINT_GRANULARITIES bin boundaries into quantile + # probabilities. E.g. decile bins [1..9] with 10 total bins become + # quantile probs [10, 20, ..., 90]. Collect all unique probs into + # all_qs so we can compute them in a single aggregation pass per metric. gran_percentiles: Dict[str, List[float]] = {} all_qs = [] for gran_name, bins in CONSTRAINT_GRANULARITIES.items(): @@ -1588,6 +1697,10 @@ def _compute_percentile_bins( gran_percentiles[gran_name] = pcts all_qs.extend(pcts) + # For each metric (syn, mis, lof), compute approximate quantile + # thresholds on the gamma OE upper CI across MANE Select transcripts + # (excluding outliers). Then slice the result by granularity to get + # the bin-edge values for percentile, decile, and sextile bins. thresholds = {} for metric, idx in metric_group_idx.items(): vals = compute_percentile_thresholds( @@ -1602,6 +1715,24 @@ def _compute_percentile_bins( for gran_name, pcts in gran_percentiles.items(): thresholds[(gran_name, metric)] = [vals[p] for p in pcts] + # Store thresholds as a global so they survive through to release. + # Structure: percentile_thresholds.{metric}.{granularity} = array + ht = ht.annotate_globals( + percentile_thresholds=hl.struct( + **{ + metric: hl.struct( + **{ + gran_name: thresholds[(gran_name, metric)] + for gran_name in gran_percentiles + } + ) + for metric in metric_group_idx + } + ) + ) + + # Annotate each transcript with its threshold-based bin assignment + # for every (granularity, metric) combination. return annotate_constraint_percentile_bins(ht, thresholds, metric_group_idx) @@ -1623,11 +1754,20 @@ def _compute_pli_scores( if expected_values is None: expected_values = PLI_EXPECTED_VALUES + # Locate the HC LoF constraint group and its adj-frequency oe_info entry + # to extract observed and expected variant counts for the EM model. meta = hl.eval(ht.constraint_group_meta) freq_meta = hl.eval(ht.exomes_freq_meta) lof_idx = meta.index({"lof": "hc"}) all_freq_idx = freq_meta.index(ADJ_FREQ_META) + # Run the EM algorithm (via compute_pli) to classify each transcript + # into three categories based on its observed vs expected HC LoF count: + # - pNull: probability of being unconstrained (OE ~ 1.0) + # - pRec: probability of being recessive-lethal (OE ~ 0.706) + # - pLI: probability of being LoF-intolerant (OE ~ 0.207) + # The result is annotated as top-level fields (pLI, pNull, pRec), not + # inside constraint_groups, since they only apply to HC LoF. hc_lof_expr = ht.constraint_groups[lof_idx].oe_info[all_freq_idx] return ht.annotate( **compute_pli( @@ -1694,6 +1834,10 @@ def compute_constraint_metrics( :return: Table with pLI scores, OE ratios, confidence intervals, z-scores, percentile bins, gene quality metrics, and GENCODE annotations. """ + # Map each consequence category to its (lower, upper) raw z-score + # outlier bounds. LoF and missense are one-sided (only lower bound); + # synonymous is two-sided since both depletion and enrichment are + # biologically meaningful. z_thresholds = { "lof": (raw_z_outlier_threshold_lower_lof, None), "mis": (raw_z_outlier_threshold_lower_missense, None), @@ -1703,22 +1847,33 @@ def compute_constraint_metrics( ), } + # Compute OE ratios, two flavors of confidence intervals + # (Poisson + gamma), raw z-scores, and per-group outlier flags. ht = _annotate_oe_ci_z(ht, z_thresholds) ht = ht.checkpoint(new_temp_file("constraint_metrics.oe_ci_z", "ht")) + # Compute per-group SD of raw z, normalize to final z_score, + # and union per-group flags into a single constraint_flags set. ht = _compute_z_scores(ht) ht = ht.checkpoint(new_temp_file("constraint_metrics.z_scores", "ht")) + # Rank transcripts by gamma OE upper CI and assign + # percentile/decile/sextile bins. Thresholds are computed on MANE + # Select transcripts and stored as globals. ht = _compute_percentile_bins(ht, use_mane_select_over_canonical) ht = ht.checkpoint(new_temp_file("constraint_metrics.percentile_bins", "ht")) + # Run the EM algorithm to compute pLI/pNull/pRec from HC LoF + # observed vs expected counts. ht = _compute_pli_scores(ht, expected_values, min_diff_convergence) ht = ht.checkpoint(new_temp_file("constraint_metrics.pli", "ht")) - # Add per-transcript gene quality metrics and flags. + # Join per-transcript gene quality metrics (coverage, mapping + # quality, segdup/LCR overlap) and gene-level flags. ht = ht.annotate(**gene_quality_metrics_ht[ht.transcript]) - # Add transcript annotations from GENCODE. + # Add transcript-level annotations from GENCODE (gene name, + # biotype, CDS length, coding exon count, etc.). ht = add_gencode_transcript_annotations(ht, gencode_ht) return ht @@ -1762,16 +1917,25 @@ def _restructure_release_rows( :return: Table with named top-level constraint group structs, release row fields selected, re-keyed, and filtered. """ + # Only build per-genetic-ancestry downsampling fields when downsampling + # data is present (i.e. when the pipeline was run with downsamplings). add_ds_fields = ( ["observed_variants", "expected_variants"] if gen_anc_ds_indices else [] ) + # Flatten the internal constraint_groups array structure into a + # release-friendly form. For each group: + # 1. Promote adj-frequency oe_info fields (oe, z_raw, obs, exp, CIs) + # to the group level. + # 2. Replace oe_ci with the gamma CI and attach rank/bin annotations + # from oe_ci_gamma_rank. + # 3. If downsamplings exist, build gen_anc_obs/gen_anc_exp structs + # keyed by genetic ancestry, each containing an array of values + # ordered by downsampling level. cg_expr = ht.constraint_groups.map( lambda cg: cg.annotate( **cg.oe_info[all_freq_idx], - oe_ci=cg.oe_info[all_freq_idx].oe_ci_gamma, - # Per-genetic-ancestry downsampling obs/exp arrays, one value per - # downsampling level, keyed by genetic ancestry. + oe_ci=cg.oe_info[all_freq_idx].oe_ci_gamma.annotate(**cg.oe_ci_gamma_rank), **{ f"gen_anc_{RELEASE_CG_RENAME[f]}": hl.struct( **{ @@ -1783,6 +1947,10 @@ def _restructure_release_rows( }, ) ) + + # Apply field renames (mu_snp -> mu, observed_variants -> obs, etc.) + # and select only the fields included in the release schema. Drop + # gen_anc_* fields when there are no downsamplings. cg_select = ( RELEASE_CG_SELECT if gen_anc_ds_indices @@ -1794,12 +1962,14 @@ def _restructure_release_rows( ).select(*cg_select) ) - # Build a top-level release field for each constraint group, applying - # group renames (e.g. lof_hc -> lof), trimming oe_ci sub-fields, and - # adding pLI/pNull/pRec for the LoF group. + # Explode the array into named top-level fields (syn, mis, lof_hc_lc, + # lof), applying group renames (lof_hc -> lof). For each group: + # - Trim oe_ci to just lower/upper for unranked groups, or add + # rank + bin fields for ranked groups (lof, lof_hc_lc). + # - Attach pLI/pNull/pRec (top-level fields from _compute_pli_scores) + # to the LoF groups. cg_fields = { - RELEASE_GROUP_RENAMES.get(name, name): cg_expr[i] - .annotate( + RELEASE_GROUP_RENAMES.get(name, name): cg_expr[i].annotate( oe_ci=cg_expr[i].oe_ci.select( *( RELEASE_CI_FIELDS_WITH_RANK @@ -1812,16 +1982,25 @@ def _restructure_release_rows( for k in (RELEASE_LOF_FIELDS if name in RELEASE_GROUPS_WITH_PLI else []) }, ) - .select() # Drop intermediate fields, keep only annotated release fields. for i, name in enumerate(field_names) } ht = ht.annotate(**cg_fields) - ht = ht.select(*RELEASE_TOP_LEVEL_ANNOTATIONS, *RELEASE_GROUP_NAMES) - available_keys = [k for k in RELEASE_KEY_ORDER if k in ht.key] - if list(ht.key) != available_keys: - ht = ht.key_by(*available_keys) + # Re-key to the canonical release key order, then select only release + # fields: scalar annotations (cds_length, gene_quality_metrics, + # constraint_flags, etc.) and the named constraint group structs. + if list(ht.key) != RELEASE_KEY_ORDER: + ht = ht.key_by(*RELEASE_KEY_ORDER) + + top_level_select = [ + k + for k in RELEASE_TOP_LEVEL_ANNOTATIONS + RELEASE_GROUP_NAMES + if k not in ht.key + ] + ht = ht.select(*top_level_select) + # Drop transcripts with zero possible variants across all constraint + # groups — these have no meaningful constraint estimates. ht = ht.filter(hl.any([ht[k].possible != 0 for k in RELEASE_GROUP_NAMES])) return ht @@ -1866,6 +2045,11 @@ def _restructure_release_globals( *None*, the existing ``version`` global is retained if present. :return: Table with release-formatted globals. """ + # Convert sd_raw_z from a positional array (parallel to + # constraint_groups) to a named struct keyed by release group name + # (syn, mis, lof, lof_hc_lc). Internal names like "lof_hc" are + # remapped via RELEASE_GROUP_RENAMES; groups not in RELEASE_GROUP_NAMES + # are dropped. sd_raw_z_name_map = {n: RELEASE_GROUP_RENAMES.get(n, n) for n in field_names} sd_raw_z_struct = hl.struct( **{ @@ -1875,16 +2059,26 @@ def _restructure_release_globals( } ) + # Build the release globals dict. select_globals at the end replaces + # all internal globals with just these. global_kwargs = {} + + # Set the release version string. if release_version is not None: global_kwargs["version"] = release_version elif "version" in ht.globals: global_kwargs["version"] = ht.globals.version + # Rename pipeline parameter globals (e.g. calculate_mu_globals -> + # calculate_mu_params) and strip internal-only sub-fields like + # freq_meta, genetic_ancestry_groups, downsampling_idx, etc. for src, dest, drop_fields in RELEASE_PIPELINE_PARAM_GLOBALS: if src in ht.globals: global_kwargs[dest] = ht.globals[src].drop(*drop_fields) + # When downsamplings are present, record the integer downsampling + # levels per genetic ancestry so consumers can interpret the + # gen_anc_obs/gen_anc_exp arrays in the row data. if gen_anc_ds_indices: global_kwargs["downsamplings"] = hl.struct( **{ @@ -1897,6 +2091,12 @@ def _restructure_release_globals( global_kwargs["max_af"] = ht.globals.max_af global_kwargs["sd_raw_z"] = sd_raw_z_struct + + # Carry through the OE upper CI threshold values used for percentile/ + # decile/sextile bin assignment, if computed. + if "percentile_thresholds" in ht.globals: + global_kwargs["percentile_thresholds"] = ht.globals.percentile_thresholds + return ht.select_globals(**global_kwargs) @@ -1934,18 +2134,26 @@ def prepare_release_ht( *None*, the existing ``version`` global is retained if present. :return: Release-formatted Table. """ + # Rename GENCODE fields to release names (e.g. transcript_id_version -> + # transcript_version, level -> transcript_level). ht = ht.rename(GENCODE_FIELD_RENAMES) + # Evaluate globals at Python time to drive the restructuring logic. constraint_meta = hl.eval(ht.constraint_group_meta) freq_meta = hl.eval(ht.exomes_freq_meta) all_freq_idx = freq_meta.index(ADJ_FREQ_META) + # Derive human-readable field names from the constraint group metadata + # dicts (e.g. {"csq_set": "syn"} -> "syn", {"lof": "hc"} -> "lof_hc"). field_names = [ "_".join(f"{k}_{v}" for k, v in m.items()).replace("csq_set_", "") for m in constraint_meta ] logger.info("Release constraint group field names: %s", field_names) + # Build a mapping from genetic ancestry to its oe_info indices for + # downsampling entries, so the row restructuring can assemble + # gen_anc_obs/gen_anc_exp arrays in the correct order. gen_anc_ds_indices: Dict[str, List[int]] = {} if "downsamplings" in ht.globals: for j, m in enumerate(freq_meta): @@ -1953,10 +2161,18 @@ def prepare_release_ht( if gen_anc is not None and "downsampling" in m: gen_anc_ds_indices.setdefault(gen_anc, []).append(j) - # Evaluate sd_raw_z before the row select (globals persist through it). + # Materialize sd_raw_z now — it's an array global that needs to be + # passed as Python values to the globals restructuring step. Must be + # evaluated before _restructure_release_rows drops the internal globals. sd_raw_z_arr = hl.eval(ht.sd_raw_z) + # Restructure rows: flatten constraint_groups array into named + # top-level structs (syn, mis, lof, lof_hc_lc) with release field names. ht = _restructure_release_rows(ht, field_names, all_freq_idx, gen_anc_ds_indices) + + # Restructure globals: replace internal pipeline globals with clean + # release globals (version, pipeline params, sd_raw_z, downsamplings, + # percentile_thresholds). ht = _restructure_release_globals( ht, field_names, freq_meta, gen_anc_ds_indices, sd_raw_z_arr, release_version ) From 17fdd0ed26ca388f3676a8dcd8ea95b48838b944 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Tue, 24 Mar 2026 12:48:51 -0600 Subject: [PATCH 33/38] Implement LoF OE CI upper bin thresholds export in the constraint pipeline. Add utility function to retrieve thresholds path and update resource handling to include new TSV export functionality. Refactor related code for clarity and consistency. --- .../pipeline/constraint_pipeline.py | 5 +++ gnomad_constraint/resources/constants.py | 2 +- gnomad_constraint/resources/resource_utils.py | 14 +++++++ gnomad_constraint/utils/constraint.py | 38 +++++++++++++++++-- 4 files changed, 55 insertions(+), 4 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index f0b76ee5..0f632fe6 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -56,6 +56,7 @@ create_per_variant_expected_ht, create_training_set, flatten_release_ht, + lof_bin_thresholds_to_ht, prepare_context_ht, prepare_ht_for_constraint_calculations, prepare_release_ht, @@ -421,6 +422,10 @@ def main(args): flatten_release_ht(release_ht).export(res.release_tsv) logger.info("Done exporting release TSV.") + logger.info("Exporting LoF OE CI upper bin thresholds TSV...") + lof_bin_thresholds_to_ht(release_ht).export(res.lof_threshold_tsv) + logger.info("Done exporting LoF threshold TSV.") + if args.export_release_downsampling_tsv: logger.info("Exporting release downsampling TSV...") downsampling_ht = explode_downsamplings_oe( diff --git a/gnomad_constraint/resources/constants.py b/gnomad_constraint/resources/constants.py index c83c4216..ed3bc169 100644 --- a/gnomad_constraint/resources/constants.py +++ b/gnomad_constraint/resources/constants.py @@ -188,7 +188,7 @@ RELEASE_CI_FIELDS_WITH_RANK = RELEASE_CI_FIELDS + RELEASE_RANK_FIELDS """CI fields including rank annotations, used for groups in RELEASE_GROUPS_WITH_RANK.""" -RELEASE_GROUPS_WITH_RANK = ["lof_hc", "lof_hc_lc"] +RELEASE_GROUPS_WITH_RANK = ["lof_hc"] """Constraint groups for which rank and bin annotations are included.""" RELEASE_GROUPS_WITH_PLI = ["lof_hc", "lof_hc_lc"] diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 2796c51a..642fa324 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -540,6 +540,19 @@ def get_gene_quality_metrics_ht(version: str = CURRENT_VERSION) -> TableResource return TableResource(f"{root}/metrics/gnomad.v{version}.gene_quality_metrics.ht") +def get_lof_threshold_tsv_path(version: str = CURRENT_VERSION) -> str: + """ + Return path for the LoF OE CI upper bin thresholds TSV. + + :param version: One of the release versions (`VERSIONS`). Default is + `CURRENT_VERSION`. + :return: Path of the LoF threshold TSV. + """ + check_param_scope(version=version) + root = get_constraint_root(version=version) + return f"{root}/release/gnomad.v{version}.loeuf_percentile_thresholds.tsv" + + def get_release_downsampling_tsv_path(version: str = CURRENT_VERSION) -> str: """ Return path for the release per-genetic-ancestry downsampling TSV. @@ -811,6 +824,7 @@ def get_constraint_resources( "release_downsampling_tsv": get_release_downsampling_tsv_path( version=version ), + "lof_threshold_tsv": get_lof_threshold_tsv_path(version=version), }, pipeline_input_steps=[prepare_release], ) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 50d07c15..9dc452fb 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -2092,10 +2092,12 @@ def _restructure_release_globals( global_kwargs["sd_raw_z"] = sd_raw_z_struct - # Carry through the OE upper CI threshold values used for percentile/ - # decile/sextile bin assignment, if computed. + # Carry through only the LoF OE upper CI threshold values used for + # percentile/decile/sextile bin assignment, renamed for clarity. if "percentile_thresholds" in ht.globals: - global_kwargs["percentile_thresholds"] = ht.globals.percentile_thresholds + global_kwargs["loeuf_percentile_thresholds"] = ( + ht.globals.percentile_thresholds.lof + ) return ht.select_globals(**global_kwargs) @@ -2197,9 +2199,39 @@ def flatten_release_ht(ht: hl.Table) -> hl.Table: if drop_fields: ht = ht.drop(*drop_fields) + # Reorder key fields to match RELEASE_KEY_ORDER before flattening so + # the TSV columns appear in the expected order (flatten drops the key + # and may emit fields in internal storage order rather than key order). + key_fields = list(ht.key) + other_fields = [f for f in ht.row if f not in ht.key] + ht = ht.key_by().select(*key_fields, *other_fields) + return ht.flatten() +def lof_bin_thresholds_to_ht(release_ht: hl.Table) -> hl.Table: + """ + Convert the LoF OE CI upper bin thresholds global into a flat Table. + + Creates a Table with one row per (granularity, bin) pair, suitable for + TSV export. + + :param release_ht: Release-format constraint metrics Table with a + ``loeuf_percentile_thresholds`` global. + :return: Unkeyed Table with ``granularity``, ``bin``, and ``threshold`` + fields. + """ + thresholds = hl.eval(release_ht.globals.loeuf_percentile_thresholds) + rows = [] + for gran in thresholds: + for i, val in enumerate(thresholds[gran]): + rows.append(hl.Struct(granularity=gran, bin=i + 1, threshold=val)) + return hl.Table.parallelize( + rows, + hl.tstruct(granularity=hl.tstr, bin=hl.tint32, threshold=hl.tfloat64), + ) + + def annotate_constraint_percentile_bins( ht: hl.Table, thresholds: Dict[Tuple[str, str], List[float]], From 76e5ac247d48cc096800232918cd7c269afbd3d6 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Tue, 24 Mar 2026 13:33:26 -0600 Subject: [PATCH 34/38] Add functionality to prepare and export mutation rate Table for release in the constraint pipeline. Introduce new command-line argument for mutation rate preparation, update resource handling, and implement utility function for mutation rate Table preparation. Enhance logging for export processes. --- .../pipeline/constraint_pipeline.py | 26 ++++++++++ gnomad_constraint/resources/resource_utils.py | 22 +++++++++ gnomad_constraint/utils/constraint.py | 48 ++++++++++++++++++- 3 files changed, 95 insertions(+), 1 deletion(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 0f632fe6..4e4238ca 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -60,6 +60,7 @@ prepare_context_ht, prepare_ht_for_constraint_calculations, prepare_release_ht, + prepare_release_mutation_ht, ) logging.basicConfig( @@ -412,6 +413,22 @@ def main(args): release_ht.write(res.release_ht.path, overwrite=overwrite) logger.info("Done preparing release Table.") + if args.prepare_release_mutation_rate: + logger.info("Preparing mutation rate Table for release...") + res = resources.prepare_release_mutation_rate + res.check_resource_existence() + + mutation_ht = res.mutation_ht.ht() + release_mutation_ht = prepare_release_mutation_ht( + mutation_ht, + release_version=args.release_version, + ) + release_mutation_ht.write(res.release_mutation_ht.path, overwrite=overwrite) + + logger.info("Exporting release mutation rate TSV...") + release_mutation_ht.export(res.release_mutation_tsv) + logger.info("Done preparing and exporting release mutation rate Table.") + if args.export_release_tsv or args.export_release_downsampling_tsv: res = resources.export_release_tsv res.check_resource_existence() @@ -923,6 +940,15 @@ def main(args): type=str, default=None, ) + prepare_release_args.add_argument( + "--prepare-release-mutation-rate", + help=( + "Prepare the mutation rate Table for public release by selecting" + " the scalar mutation rate (mu), trinucleotide-class flags, and" + " restructuring globals. Also exports a TSV." + ), + action="store_true", + ) prepare_release_args.add_argument( "--export-release-tsv", help=( diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 642fa324..3a1c0daf 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -354,6 +354,19 @@ def get_release_mutation_ht(version: str = CURRENT_VERSION) -> TableResource: return TableResource(f"{root}/release/gnomad.v{version}.mutation_rate.ht") +def get_release_mutation_tsv_path(version: str = CURRENT_VERSION) -> str: + """ + Return path for the release mutation rate TSV. + + :param version: One of the release versions (`VERSIONS`). Default is + `CURRENT_VERSION`. + :return: Path of the release mutation rate TSV. + """ + check_param_scope(version=version) + root = get_constraint_root(version=version) + return f"{root}/release/gnomad.v{version}.mutation_rate.tsv" + + def get_release_constraint_ht(version: str = CURRENT_VERSION) -> TableResource: """ Return TableResource for the release constraint metrics Table. @@ -817,6 +830,14 @@ def get_constraint_resources( }, pipeline_input_steps=[compute_constraint_metrics], ) + prepare_release_mutation_rate = PipelineStepResourceCollection( + "--prepare-release-mutation-rate", + output_resources={ + "release_mutation_ht": get_release_mutation_ht(version=version), + "release_mutation_tsv": get_release_mutation_tsv_path(version=version), + }, + pipeline_input_steps=[calculate_mutation_rate], + ) export_release_tsv = PipelineStepResourceCollection( "--export-release-tsv", output_resources={ @@ -845,6 +866,7 @@ def get_constraint_resources( "apply_models_aggregated": apply_models_aggregated, "compute_constraint_metrics": compute_constraint_metrics, "prepare_release": prepare_release, + "prepare_release_mutation_rate": prepare_release_mutation_rate, "export_release_tsv": export_release_tsv, } ) diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 9dc452fb..87e3409c 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -1,7 +1,7 @@ """Script containing utility functions used in the constraint pipeline.""" import logging -from typing import Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple import hail as hl from gnomad.resources.grch38.gnomad import DOWNSAMPLINGS @@ -2181,6 +2181,52 @@ def prepare_release_ht( return ht +def prepare_release_mutation_ht( + ht: hl.Table, + release_version: Optional[str] = None, +) -> hl.Table: + """ + Prepare the mutation rate Table for public release. + + Selects the per-context mutation rate (``mu_snp`` renamed to ``mu``), + trinucleotide-class flags (``cpg``, ``transition``, ``mutation_type``), + and restructures the globals to drop internal pipeline bookkeeping + fields. + + :param ht: Mutation rate Table produced by + :func:`calculate_mu_by_downsampling`. + :param release_version: Version string for the ``version`` global. + When *None*, the existing ``version`` global is retained if + present. + :return: Release-formatted mutation rate Table. + """ + # Keep only the scalar mutation rate and the trinucleotide-class flags. + ht = ht.select( + mu=ht.mu_snp, + cpg=ht.cpg, + transition=ht.transition, + mutation_type=ht.mutation_type, + ) + + # Restructure globals. + global_kwargs: Dict[str, Any] = {} + if release_version is not None: + global_kwargs["version"] = release_version + elif "version" in ht.globals: + global_kwargs["version"] = ht.globals.version + + # Rename calculate_mu_globals → calculate_mu_params, dropping + # internal-only sub-fields (freq_meta, genetic_ancestry_groups, + # downsampling_idx). Only keep calculate_mu_globals; the other + # pipeline globals (build_models, apply_models) are not relevant + # to the mutation rate release. + src, dest, drop_fields = RELEASE_PIPELINE_PARAM_GLOBALS[0] + if src in ht.globals: + global_kwargs[dest] = ht.globals[src].drop(*drop_fields) + + return ht.select_globals(**global_kwargs) + + def flatten_release_ht(ht: hl.Table) -> hl.Table: """ Flatten the release constraint metrics Table for TSV export. From 5bc7d3a7029f2bf1a91e87b72a6cf43eaa1f1ad1 Mon Sep 17 00:00:00 2001 From: Ruchit10 Date: Wed, 15 Apr 2026 16:34:16 -0400 Subject: [PATCH 35/38] Fix missing setuptools version in pyproject --- pyproject.toml | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 244c8233..f2136aa7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,3 +1,15 @@ +[build-system] +requires = ["setuptools>=45"] +build-backend = "setuptools.build_meta" + +[project] +name = "gnomad-constraint" +version = "0.1.0" +requires-python = ">=3.9" + +[tool.setuptools.packages.find] +include = ["gnomad_constraint*"] + [tool.pydocstyle] convention = "pep257" match = ".*\\.py" From 237902aa49ef329b651165585233afe8cd740c83 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Fri, 15 May 2026 09:25:55 -0600 Subject: [PATCH 36/38] Add README for gnomad_constraint resources, detailing pipeline outputs, path conventions, and shared structures for version 4.1.1. --- gnomad_constraint/resources/README.md | 714 ++++++++++++++++++++++++++ 1 file changed, 714 insertions(+) create mode 100644 gnomad_constraint/resources/README.md diff --git a/gnomad_constraint/resources/README.md b/gnomad_constraint/resources/README.md new file mode 100644 index 00000000..3d7fe7c1 --- /dev/null +++ b/gnomad_constraint/resources/README.md @@ -0,0 +1,714 @@ +# gnomad_constraint resources + +This package wires every Hail Table produced by the constraint pipeline to a +canonical GCS path. The functions live in [resource_utils.py](resource_utils.py) +and the constants they consume live in [constants.py](constants.py). + +This document is the schema reference for every **pipeline output** of the v4.1.1 +constraint pipeline. Schemas were captured by running `ht.describe()` against +the live tables on GCS. For upstream **input** resources (VEP context, sites, +coverage, methylation, GERP, GENCODE, `all_sites_an`) see +[`gnomad.resources.grch38`](https://github.com/broadinstitute/gnomad_methods/tree/main/gnomad/resources/grch38). + +> **Scope.** This README documents v4.1.1. Older versions (`2.1.1`, `4.0`, +> `4.1`) follow the same general layout but partition some artifacts by +> genomic region (e.g. `*.autosome_par.ht`, `*.chrx_nonpar.ht`, +> `*.chry_nonpar.ht`). v4 drops chrX/chrY early in the pipeline. + +--- + +## Path conventions + +Built by [`get_constraint_root`](resource_utils.py) and +[`get_constraint_data`](resource_utils.py): + +``` +gs://gnomad/v{version}/constraint/ +├── preprocessed_data/ +│ ├── gnomad.v{version}.annotated_context.ht +│ └── gnomad.v{version}.context.preprocessed.ht +├── mutation_rate/ +│ └── gnomad.v{version}.mutation_rate.ht +├── training_data/ +│ ├── gnomad.v{version}.constraint_training.ht +│ └── gnomad.v{version}.constraint_training.tsv.bgz +├── models/ +│ ├── gnomad.v{version}.plateau.he +│ └── gnomad.v{version}.coverage.he +├── apply_models/{transcript_consequences|worst_csq_by_gene}/ +│ ├── gnomad.v{version}.per_variant_expected.ht +│ ├── gnomad.v{version}.per_variant_expected.aggregated.ht +│ ├── gnomad.v{version}.aggregated_expected.ht +│ └── gnomad.v{version}.constraint_group.ht +├── metrics/ +│ ├── gnomad.v{version}.gene_quality_metrics.ht +│ └── {transcript_consequences|worst_csq_by_gene}/ +│ └── gnomad.v{version}.constraint_metrics.ht +└── release/ + ├── gnomad.v{version}.constraint_metrics.ht + ├── gnomad.v{version}.constraint_metrics.tsv.bgz + ├── gnomad.v{version}.constraint_metrics.downsampling.tsv.bgz + ├── gnomad.v{version}.loeuf_percentile_thresholds.tsv + ├── gnomad.v{version}.mutation_rate.ht + └── gnomad.v{version}.mutation_rate.tsv +``` + +Tests use `gs://gnomad-tmp/gnomad_v{version}_testing/constraint/...` (`test=True`) +and intermediate checkpoints use `gs://gnomad-tmp/gnomad_v{version}/constraint/...` +(`temp=True`). + +--- + +## Shared structures + +Several tables carry identical **pipeline-parameter globals** that capture the +arguments used to produce them. They are written by +[`calculate_mu_by_downsampling`](../utils/constraint.py), [`build_models`](../utils/constraint.py), +and the apply-models steps: + +### `calculate_mu_globals` — `struct` +Parameters passed to `calculate_mu_by_downsampling`. +- `freq_meta: array>` — frequency metadata entries (one per + freq-array index): `{group: adj}` for the global adj entry plus one + `{downsampling, group, pop}` entry per downsampling × genetic-ancestry-group. +- `ac_cutoff: int32` — variants with `AC > ac_cutoff` are excluded from mu + calculation (typical value: 5). +- `min_cov: int32` — minimum mean exome coverage filter for mu sites (15). +- `max_cov: int32` — maximum mean exome coverage filter (60). +- `gerp_lower_cutoff: float64` — lower GERP bound for the mu site set (−3.9885). +- `gerp_upper_cutoff: float64` — upper GERP bound (2.6607). +- `genetic_ancestry_groups: array` — gen-anc groups used for downsampling + (`global, afr, amr, eas, nfe, sas`). +- `downsampling_level: int32` — the downsampling size used for the canonical + mu calculation (1000). +- `downsampling_idx: int32` — index into the parallel freq arrays whose entry + was used as `mu_snp` (the global×1000 entry). +- `most_severe_consequence: array` — VEP `most_severe_consequence` values + retained as putatively neutral for mu fitting + (`intron_variant`, `intergenic_variant`). + +### `build_models_globals` — `struct` +Parameters passed to the plateau/coverage model fit. +- `synonymous_transcript_filter_field: str` — transcript filter used when + selecting synonymous training variants (e.g. `canonical`). +- `low_cov_cutoff: int32` — coverage below this is treated as the low-coverage + regime (19). +- `high_cov_cutoff: int32` — coverage at or above this is the high-coverage + regime that the plateau model is fit to (90). +- `upper_cov_cutoff: int32` — upper bound for the coverage model fit (nullable). +- `skip_coverage_model: bool` — whether the coverage model fit was skipped. + +### `apply_models_globals` — `struct` +Carries the previous two plus the fitted models (after `--apply-models-*`). +- `low_cov_cutoff`, `high_cov_cutoff`, `skip_coverage_model` — as above. +- `plateau_models: dict>>` + — fitted plateau slope/intercept per (cpg, genomic_region). Each value is + an array parallel to `mutation_rate.mu` (one fit per downsampling). +- `coverage_model: array` — `[intercept, slope]` of the coverage + correction regression. +- `log10_coverage: bool` — whether the coverage model is in log10 space. +- `groupings: tuple(str×7)` or `array` — the field names used to group + variants when applying models (annotation, modifier, gene, gene_id, + transcript, canonical, mane_select). Stored as a tuple on the per-variant + path, as an array on the aggregated path. + +### Other shared globals +- `exomes_freq_meta: array>` — frequency metadata of the exomes + freq array (one entry per `(downsampling, gen_anc)`). +- `genetic_ancestry_groups: array` — copy of the gen-anc groups list. +- `downsamplings: array` — ordered list of downsampling sizes the freq + array covers (e.g. `[10, 20, ..., 1000, ...]`). +- `max_af: float64` — AF cap for the "observed" variant filter. + +--- + +# Pipeline output tables + +Each section below lists the **canonical path**, the **key**, the **globals** +and the **row fields** with sub-struct expansion. Producing function is the +`--` flag from [constraint_pipeline.py](../pipeline/constraint_pipeline.py). + +--- + +## 1. `annotated_context.ht` + +The fully annotated universe of all possible single-nucleotide substitutions +in the genome, produced by `--prepare-context-ht` / +[`prepare_ht_for_constraint_calculations`](../utils/constraint.py). + +- **Path:** `gs://gnomad/v{version}/constraint/preprocessed_data/gnomad.v{version}.annotated_context.ht` +- **Resource fn:** [`get_annotated_context_ht`](resource_utils.py) +- **Key:** `locus`, `alleles` + +### Globals +- `grange: array` — methylation-level grange bin edges (10 bins). +- `vep_help: str` — captured `vep --help` text from the run that annotated VEP. +- `vep_config: str` — JSON-serialized VEP runner config (cache version, plugins, + command). +- `version: str` — gnomAD version string. +- `an_globals: struct{exomes, genomes}` — AN strata metadata (from + [`all_sites_an`](https://github.com/broadinstitute/gnomad_methods)) for each data type: + - `strata_sample_count: array` — sample count for each strata index. + - `strata_meta: array>` — strata key (e.g. `{group: adj, gen_anc: afr}`). +- `freq_globals: struct{exomes, genomes}` — frequency strata metadata: + - `exomes.freq_meta_sample_count: array` + - `exomes.freq_meta: array>` + - `genomes.freq_meta: array>` + +### Rows +- `locus: locus` +- `alleles: array` — `[ref, alt]` of one SNV. +- `context: str` — trinucleotide context (ref-strand-collapsed via + [`collapse_strand`](../utils/constraint.py): rows with G/T ref are reverse-complemented and + `was_flipped` is set true). +- `vep: struct` — minimal VEP output kept for constraint: + - `most_severe_consequence: str` + - `transcript_consequences: array` — per-transcript VEP entries with + `transcript_id, gene_id, gene_symbol, biotype, most_severe_consequence, + mane_select, canonical, lof, lof_flags, sift_score, polyphen_score, + domains, uniprot_isoform, amino_acids, codons`. +- `ref: str`, `alt: str` — ref-strand-collapsed alleles. +- `was_flipped: bool` — `true` if the original ref was G or T and + ref/alt/context were reverse-complemented. +- `transition: bool` — purine↔purine or pyrimidine↔pyrimidine. +- `cpg: bool` — variant is in a CpG dinucleotide context. +- `mutation_type: str` — coarse mutation type (e.g. `CpG`, `non-CpG transition`, + `transversion`). +- `mutation_type_model: str` — mutation-type label used as the model group. +- `methylation_level: int32` — discretized methylation level (0–`len(grange)-1`). +- `gerp: float64` — GERP RS score at the locus. +- `coverage: struct{exomes, genomes}` — coverage from gnomAD coverage HT: + - `
.mean: float64`, `
.median_approx: int32`. +- `AN: struct{exomes: int64, genomes: int64}` — global-adj allele number from + [`all_sites_an`](https://github.com/broadinstitute/gnomad_methods). +- `freq: struct{exomes, genomes}` — per-downsampling allele freqs from the + sites table. Each entry is `array` + parallel to `freq_globals.
.freq_meta`. (Note: exomes + `homozygote_count` is `int64`, genomes `int32`.) +- `filters: struct{exomes: set, genomes: set}` — site filters + inherited from the corresponding release sites table. +- `genomic_region: str` — `autosome_or_par`, `chrx_nonpar`, or `chry_nonpar`. +- `adj_r: float64` — per-context regional depletion correction (see + `adj_r.ht`). +- `syn_adj_r: float64` — synonymous-DNM variant of `adj_r` (see `syn_adj_r.ht`). +- `sfs_bin: int32` — site-frequency-spectrum bin assigned via + [`annotate_sfs_bin`](../utils/constraint.py) from `SFS_BIN_CUTOFFS` (0 if AF + is missing; otherwise index of first cutoff `af_expr <= cutoff` is true). + +--- + +## 2. `context.preprocessed.ht` + +Context joined with exomes/genomes frequency, coverage, AN, and the +pre-computed mu inputs (`compute_mu`, `calibrate_mu` structs). Produced by the +`preprocess_data` step / [`preprocess_data`](../utils/constraint.py) — the +shared upstream input for nearly every downstream step. + +- **Path:** `gs://gnomad/v{version}/constraint/preprocessed_data/gnomad.v{version}.context.preprocessed.ht` +- **Resource fn:** [`get_preprocessed_ht`](resource_utils.py) +- **Key:** `locus`, `alleles` + +### Globals +`calculate_mu_globals`, `build_models_globals`, `apply_models_globals`, +`exomes_freq_meta`, `genetic_ancestry_groups`, `downsamplings`, `max_af` — +see [Shared structures](#shared-structures). + +### Rows +Carries every annotated_context row field *except* `freq` (which is split +into `compute_mu.genomes_freq` / `calibrate_mu.exomes_freq`), plus: +- `exomes_coverage: int32` — `coverage.exomes.median_approx` clamped to the + `[low_cov_cutoff, high_cov_cutoff]` model bands. +- `compute_mu: struct` — fields used to fit the mutation-rate model on the + *genome* SNVs: + - `genomes_freq: array` — genomes freq + array (parallel to `exomes_freq_meta`). + - `observed_variants: array` — observed counts (one per freq entry). + - `possible_variants: int32` — count of possible SNVs at this site (always 1 + for context rows; used by aggregation). +- `calibrate_mu: struct` — fields used to *apply* mu (i.e. calibrate by + comparison to *exome* SNVs): + - `exomes_freq: array` + - `observed_variants: array` + - `possible_variants: int32` + - `build_model: struct{high_or_low_coverage: str, model_group: struct{cpg: bool, genomic_region: str}}` + — which plateau model bucket this row was fit into. + - `apply_model: struct{high_or_low_coverage: str, model_group: struct{cpg: bool, genomic_region: str}}` + — which plateau model bucket this row will be evaluated against. + +--- + +## 3. `mutation_rate.ht` + +Per-context mutation rate, produced by `--calculate-mutation-rate` / +[`calculate_mu_by_downsampling`](../utils/constraint.py). + +- **Path:** `gs://gnomad/v{version}/constraint/mutation_rate/gnomad.v{version}.mutation_rate.ht` +- **Resource fn:** [`get_mutation_ht`](resource_utils.py) +- **Key:** `context`, `ref`, `alt`, `methylation_level` (the `MU_GROUPING`) + +### Globals +`calculate_mu_globals`, `build_models_globals`, `apply_models_globals`, +`exomes_freq_meta`, `genetic_ancestry_groups`, `downsamplings`, `max_af` — +see [Shared structures](#shared-structures). + +### Rows +- `context: str`, `ref: str`, `alt: str`, `methylation_level: int32` — key. +- `observed_variants: array` — observed counts per freq-meta entry. +- `possible_variants: int64` — count of possible SNVs of this trinucleotide + context/substitution. +- `proportion_observed: array` — `observed / possible` per freq entry. +- `mu: array` — scaled mutation rate per freq entry. +- `mu_snp: float64` — scalar mutation rate at the canonical + `(global, downsampling=downsampling_level)` index (typically global×1000). +- `transition: bool`, `cpg: bool`, `mutation_type: str`, `mutation_type_model: str` + — mutation-type annotations propagated from + [`annotate_mutation_type`](../utils/constraint.py). + +--- + +## 4. `constraint_training.ht` + +Training set for plateau / coverage model fit. Produced by +`--create-training-set` from synonymous high-coverage SNVs. + +- **Path:** `gs://gnomad/v{version}/constraint/training_data/gnomad.v{version}.constraint_training.ht` +- **Resource fn:** [`get_training_dataset`](resource_utils.py) +- **Key:** `context, ref, alt, methylation_level, cpg, transition, mutation_type, mutation_type_model, genomic_region, build_model, exomes_coverage` + +### Globals +Same shared pipeline parameter globals as `mutation_rate.ht`. + +### Rows +- `context`, `ref`, `alt`, `methylation_level` — substitution / methylation key. +- `cpg`, `transition`, `mutation_type`, `mutation_type_model` — mutation-type annotations. +- `genomic_region: str` — `autosome_or_par` / `chrx_nonpar` / `chry_nonpar`. +- `build_model: struct{high_or_low_coverage: str, model_group: struct{cpg: bool, genomic_region: str}}` + — the (coverage-band × cpg × region) bucket whose plateau model this row trains. +- `exomes_coverage: int32` — clamped median exome coverage bin. +- `observed_variants: array` — aggregated observed SNV counts per freq entry. +- `possible_variants: int64` — aggregated possible SNV count. +- `mu_snp: float64` — joined-in mu_snp for this (context, ref, alt, methylation). + +A `tsv.bgz` mirror of this table is written to the same directory. + +--- + +## 5. `plateau.he` and `coverage.he` + +Fitted Hail Expressions produced by `--build-models` / +[`build_models`](gnomad.utils.constraint:build_models). These are *not* tables — they are +pickled Hail expressions reread by `apply_models_globals`. + +- **Paths:** + - `gs://gnomad/v{version}/constraint/models/gnomad.v{version}.plateau.he` + - `gs://gnomad/v{version}/constraint/models/gnomad.v{version}.coverage.he` +- **Resource fn:** [`get_models`](resource_utils.py) + +### Types and meaning +- `plateau.he : dict>>` + — for each (cpg, genomic_region) bucket, a list of `[slope, intercept]` + pairs parallel to `mutation_rate.mu` (one entry per freq-meta index). The + inner ordering follows `freq_meta`, so element `downsampling_idx` is the + canonical fit. +- `coverage.he : array` — `[intercept, slope]` of the linear + regression mapping (log10) median exome coverage to the multiplicative + coverage correction in the low-coverage regime. + +These two expressions are read back into `apply_models_globals.plateau_models` +and `apply_models_globals.coverage_model` on every downstream apply step. + +--- + +## 6. `per_variant_expected.ht` + +Per-SNV expected counts, produced by `--apply-models-per-variant`. One row per +locus×allele pair, exploded through `transcript_consequences` (or per-gene if +the `worst_csq_by_gene` variant is built). + +- **Path:** `gs://gnomad/v{version}/constraint/apply_models/{vep_annot}/gnomad.v{version}.per_variant_expected.ht` +- **Resource fn:** [`get_per_variant_expected_dataset`](resource_utils.py) +- **Key:** `locus`, `alleles` + +### Globals +`calculate_mu_globals`, `build_models_globals`, `apply_models_globals` (now +populated with `plateau_models` and `coverage_model`), `exomes_freq_meta`, +`genetic_ancestry_groups`, `downsamplings`, `max_af`. + +### Rows +All preprocessed-context fields (`locus, alleles, context, ref, alt, +was_flipped, transition, cpg, mutation_type, mutation_type_model, +methylation_level, gerp, coverage, AN, filters, genomic_region, adj_r, +syn_adj_r, sfs_bin, exomes_coverage, compute_mu, calibrate_mu`) plus the +applied-model fields: +- `annotation: str` — VEP `most_severe_consequence` for this transcript row. +- `modifier: str` — finer-grained consequence sub-class (e.g. LOFTEE HC/LC, + `missense_variant` modifier, etc.). +- `gene: str`, `gene_id: str`, `transcript: str` — VEP transcript identity. +- `canonical: bool`, `mane_select: bool` — transcript flags. +- `mu_snp: float64` — joined per-context mu (scalar). +- `mu: float64` — `mu_snp * possible_variants` (the per-variant mu mass). +- `predicted_proportion_observed: array` — plateau-applied predicted + observed proportion per freq entry (one element per `freq_meta` row). +- `expected_variants: array` — `predicted_proportion_observed * + possible * coverage_correction` per freq entry. +- `coverage_correction: float64` — coverage-model multiplicative correction + evaluated at this row's `exomes_coverage`. + +--- + +## 7. `per_variant_expected.aggregated.ht` + +Sum of `per_variant_expected.ht` over variants within each +`(annotation, modifier, gene, gene_id, transcript, canonical, mane_select)` +tuple. Produced by `--aggregate-per-variant-expected`. + +- **Path:** `gs://gnomad/v{version}/constraint/apply_models/{vep_annot}/gnomad.v{version}.per_variant_expected.aggregated.ht` +- **Resource fn:** [`get_aggregated_per_variant_expected`](resource_utils.py) +- **Key:** `annotation, modifier, gene, gene_id, transcript, canonical, mane_select` + +### Globals +Identical pipeline-parameter globals as `per_variant_expected.ht`. + +### Rows +- `annotation, modifier, gene, gene_id, transcript, canonical, mane_select` — key. +- `mu_snp: float64` — summed `mu_snp * possible_variants` (i.e. total per-context + mu mass over the bucket; despite the name, this is no longer a per-SNV rate). +- `mu: float64` — summed per-variant `mu` (`mu_snp * possible` over bucket). +- `observed_variants: array` — summed observed counts per freq entry. +- `possible_variants: int64` — total possible SNV count in the bucket. +- `predicted_proportion_observed: array` — `Σ predicted_proportion_observed + · possible` (note: weighted by `possible`, divide by `possible_variants` to + recover the bucket-mean fraction). +- `coverage_correction: float64` — `Σ coverage_correction · possible`. +- `expected_variants: array` — summed expected counts per freq entry. + +--- + +## 8. `aggregated_expected.ht` + +Alternative path: aggregate *first*, then apply models. Produced by +`--apply-models-aggregated` and has the same schema as +`per_variant_expected.aggregated.ht`. + +- **Path:** `gs://gnomad/v{version}/constraint/apply_models/{vep_annot}/gnomad.v{version}.aggregated_expected.ht` +- **Resource fn:** [`get_aggregated_expected`](resource_utils.py) +- **Key:** identical to §7. + +### Globals +Same as §7, except `apply_models_globals.groupings` is `array` (rather +than `tuple(str×7)`) because aggregation happens before per-variant grouping. + +### Rows +Identical to §7. Use this table to compare aggregate-before-apply vs. +apply-then-aggregate (§7) results; the canonical downstream input is §7. + +--- + +## 9. `constraint_group.ht` + +Pre-metric per-transcript table organized into constraint groups (`syn`, `mis`, +`lof_hc`, `lof_hc_lc`, plus any additional groupings). Produced by +`--aggregate-by-constraint-groups` / +[`aggregate_by_constraint_groups`](../utils/constraint.py). + +- **Path:** `gs://gnomad/v{version}/constraint/apply_models/{vep_annot}/gnomad.v{version}.constraint_group.ht` +- **Resource fn:** [`get_constraint_group_ht`](resource_utils.py) +- **Key:** `gene, gene_id, transcript, canonical, mane_select` + +### Globals +Pipeline-parameter globals (same as §7) plus: +- `constraint_group_meta: array>` — one dict per + `constraint_groups[i]` describing the group's filter (e.g. `{annotation: + synonymous_variant}`, `{annotation: missense_variant}`, + `{annotation: lof, modifier: HC}`, `{annotation: lof, modifier: HC_LC}`). + Position in this array matches position in the row-level + `constraint_groups` array. + +### Rows +- `gene, gene_id, transcript, canonical, mane_select` — key. +- `constraint_groups: array` — one element per metadata entry in + `constraint_group_meta`, each carrying summed counts: + - `mu_snp: float64` — `Σ mu_snp · possible` over variants in this group. + - `mu: float64` — `Σ mu` (i.e. `Σ mu_snp · possible`). + - `possible_variants: int64` — total possible SNVs. + - `coverage_correction: float64` — `Σ coverage_correction · possible`. + - `oe_info: array` + — one entry per freq-meta index. `observed_variants` is the summed + observed count, `expected_variants` is the summed predicted expected + count, and `predicted_proportion_observed` is the `Σ ppo·possible` sum. +- `no_variants: bool` — true if **every** constraint group has zero observed + variants at the global-adj index (used to drop these rows in downstream + metrics steps). + +--- + +## 10. `constraint_metrics.ht` (internal) + +Per-transcript metrics: pLI, OE confidence intervals (two estimators), +raw and standardized z-scores, percentile/decile/sextile bins. Produced by +`--compute-constraint-metrics`. This is the *internal* metrics table; the +public release table (§12) is a renamed/flattened projection. + +- **Path:** `gs://gnomad/v{version}/constraint/metrics/{vep_annot}/gnomad.v{version}.constraint_metrics.ht` +- **Resource fn:** [`get_constraint_metrics_dataset`](resource_utils.py) +- **Key:** `gene, gene_id, transcript, canonical, mane_select` + +### Globals +Pipeline-parameter globals (as §9) plus: +- `constraint_group_meta` — see §9. +- `sd_raw_z: array` — per-constraint-group standard deviation of the + raw z-statistic across all transcripts (parallel to `constraint_group_meta`). + Used to scale `z_raw → z_score`. +- `percentile_thresholds: struct` — observed quantile cutoffs of the + `oe_ci.upper` distribution. One sub-struct per metric: + - `syn`, `mis`, `lof: struct{percentile: array, decile: array, sextile: array}` + — each granularity's boundary values (e.g. + `percentile_thresholds.lof.decile[0]` = 10th-percentile of LOEUF upper). + +### Rows +- `gene, gene_id, transcript, canonical, mane_select` — key. +- `constraint_groups: array` — one element per group (parallel to + `constraint_group_meta`). Each element extends §9's struct: + - `mu_snp`, `mu`, `possible_variants`, `coverage_correction` — as in §9. + - `oe_info: array` — extended with metric fields: + - `observed_variants: int64`, `predicted_proportion_observed: float64`, + `expected_variants: float64` + - `oe: float64` — `observed / expected` ratio. + - `oe_ci_discretized_poisson: struct{lower: float64, upper: float64}` — OE + CI using a discretized-Poisson estimator. + - `oe_ci_gamma: struct{lower: float64, upper: float64}` — OE CI using a + Gamma estimator (the LOEUF estimator). + - `z_raw: float64` — `(observed - expected) / sqrt(expected)`. + - `flags: set` — per-group QC flags (e.g. low observed, low expected). + - `z_score: float64` — `z_raw / sd_raw_z[group_idx]`, taken at the + canonical (global-adj) freq index. + - `oe_ci_discretized_poisson_rank: struct` and `oe_ci_gamma_rank: struct` — + rank/bin annotations for the `upper` bound of each CI: + - `upper_rank: int64` — 0-based ascending rank within the group. + - `upper_bin_percentile: int32` — percentile bin (1–99). + - `upper_bin_decile: int32` — decile bin (1–9). + - `upper_bin_sextile: int32` — sextile bin (1–5). +- `no_variants: bool` — as §9. +- `constraint_flags: set` — transcript-level QC flags + (`no_variants`, `not_in_gencode`, `outlier_*`, etc.). +- `constraint_bins: struct{percentile, decile, sextile}` — convenience bins + per metric, where each is `struct{syn: int32, mis: int32, lof: int32}`. +- `pLI: float64`, `pNull: float64`, `pRec: float64` — pLI / pRec / pNull scores + derived from LOEUF observed / expected via + [`compute_pli`](https://github.com/broadinstitute/gnomad_methods). +- `gene_quality_metrics: struct` — joined-in from §11: + - `exome_prop_bp_AN90: float64` — fraction of transcript CDS bp with + exome AN ≥ 90% of max. + - `exome_mean_AS_MQ: float64` — mean AS_MQ over the transcript. + - `exome_prop_segdup: float64` — fraction of transcript CDS in segmental + duplications. + - `exome_prop_LCR: float64` — fraction in low-complexity regions. +- `gene_flags: set` — joined-in from §11 (e.g. low coverage, high segdup). +- `level: str`, `transcript_type: str`, `chromosome: str`, + `start_position: int32`, `end_position: int32`, + `gene_id_version: str`, `transcript_id_version: str`, + `cds_length: int64`, `num_coding_exons: int64` — GENCODE annotations. + +--- + +## 11. `gene_quality_metrics.ht` + +Per-transcript coverage / mapping-quality and region-overlap metrics. Produced +by `--compute-gene-quality-metrics`. + +- **Path:** `gs://gnomad/v{version}/constraint/metrics/gnomad.v{version}.gene_quality_metrics.ht` +- **Resource fn:** [`get_gene_quality_metrics_ht`](resource_utils.py) +- **Key:** `transcript` + +### Globals +*(none)* + +### Rows +- `transcript: str` — key. +- `gene_quality_metrics: struct{exome_prop_bp_AN90, exome_mean_AS_MQ, exome_prop_segdup, exome_prop_LCR}` + — see §10 for the individual fields. +- `gene_flags: set` — transcript-level flags derived from the above + (e.g. `low_exome_coverage`, `high_segdup_overlap`). + +--- + +## 12. `release/constraint_metrics.ht` + +Public-release flattened constraint table. One row per +`(gene, gene_id, transcript, canonical, mane_select)`. Produced by +`--prepare-release` from §10 by renaming/selecting fields per the +`RELEASE_*` constants in [constants.py](constants.py). + +- **Path:** `gs://gnomad/v{version}/constraint/release/gnomad.v{version}.constraint_metrics.ht` +- **Resource fn:** [`get_release_constraint_ht`](resource_utils.py) +- **Key:** `gene, gene_id, transcript, canonical, mane_select` +- A flat **TSV** (`constraint_metrics.tsv.bgz`), a **downsampling TSV** + (`constraint_metrics.downsampling.tsv.bgz`), and a **LOEUF percentile + thresholds TSV** (`loeuf_percentile_thresholds.tsv`) are exported alongside. + +### Globals +- `version: str` — gnomAD release version (e.g. `4.1.1`). +- `calculate_mu_params: struct` — release-cleaned subset of + `calculate_mu_globals` (drops `freq_meta`, `genetic_ancestry_groups`, + `downsampling_idx`): + - `ac_cutoff, min_cov, max_cov, gerp_lower_cutoff, gerp_upper_cutoff, + downsampling_level, most_severe_consequence`. +- `build_models_params: struct{low_cov_cutoff, high_cov_cutoff, upper_cov_cutoff}` + — release-cleaned `build_models_globals` (drops + `synonymous_transcript_filter_field`, `skip_coverage_model`). +- `apply_models_params: struct{low_cov_cutoff, high_cov_cutoff, plateau_models, coverage_model, log10_coverage}` + — release-cleaned `apply_models_globals` (drops `skip_coverage_model`, + `groupings`). +- `downsamplings: struct{global, afr, amr, eas, nfe, sas}` — each is + `array` of downsampling sizes present for that gen-anc group + (parallel to the per-row `gen_anc_obs.` / `gen_anc_exp.` + arrays). +- `max_af: float64` — AF cap for observed. +- `sd_raw_z: struct{syn, mis, lof_hc_lc, lof}` — per-release-group `sd_raw_z` + (re-keyed from positional array to named struct via `RELEASE_GROUP_RENAMES`, + which maps internal `lof_hc → lof` for release). +- `loeuf_percentile_thresholds: struct{percentile, decile, sextile}` — LOEUF + upper-bound thresholds at each granularity (copied from + `percentile_thresholds.lof` in §10). + +### Rows +Key + GENCODE annotations + per-group structs: + +- `gene, gene_id, transcript, canonical, mane_select` — key. +- `transcript_version: str` (renamed from `transcript_id_version`), + `transcript_type: str`, + `transcript_level: str` (renamed from `level`), + `chromosome: str`, `start_position: int32`, `end_position: int32`, + `cds_length: int64`, `num_coding_exons: int64`. +- `gene_quality_metrics: struct{exome_prop_bp_AN90, exome_mean_AS_MQ, exome_prop_segdup, exome_prop_LCR}` + — see §10. +- `gene_flags: set`, `constraint_flags: set` — propagated from §10/§11. + +Each constraint group is exposed as a named struct (group names from +`RELEASE_GROUP_NAMES`): + +- `syn: struct` — synonymous constraint (fields below). +- `mis: struct` — missense constraint. +- `lof_hc_lc: struct` — pLoF (LOFTEE HC + LC). +- `lof: struct` — pLoF (LOFTEE HC only; renamed from internal `lof_hc`). This + is the canonical LOEUF group. + +**Common sub-fields** on every `{syn,mis,lof_hc_lc,lof}` struct +(`RELEASE_CG_SELECT` order): + +- `mu: float64` — total mu mass for this group (renamed from `mu_snp` per + `RELEASE_CG_RENAME`). +- `possible: int64` — total possible SNVs (from `possible_variants`). +- `obs: int64` — observed SNV count at the global-adj index. +- `exp: float64` — expected SNV count at the global-adj index. +- `oe: float64` — `obs / exp`. +- `z_raw: float64` — raw z-statistic. +- `z_score: float64` — `z_raw / sd_raw_z[]`. +- `oe_ci: struct{lower: float64, upper: float64}` — OE confidence interval. + The estimator (Poisson vs Gamma) depends on the group: + - `syn`, `mis`, `lof_hc_lc` use the discretized-Poisson CI. + - `lof` uses the Gamma CI (the LOEUF), and **additionally** carries rank + fields on `oe_ci` (the only group in `RELEASE_GROUPS_WITH_RANK`): + - `upper_rank: int64` — 0-based rank by `oe_ci.upper` ascending. + - `upper_bin_percentile: int32` — percentile bin (1–99). + - `upper_bin_decile: int32` — decile bin (1–9; LOEUF decile). + - `upper_bin_sextile: int32` — sextile bin (1–5). +- `gen_anc_obs: struct{global, afr, amr, eas, nfe, sas}` — per-genetic-ancestry + observed counts. Each field is `array` parallel to + `downsamplings.`. +- `gen_anc_exp: struct{global, afr, amr, eas, nfe, sas}` — per-genetic-ancestry + expected counts. Each field is `array` parallel to + `downsamplings.`. + +**Additional sub-fields on `lof_hc_lc` and `lof`** (`RELEASE_GROUPS_WITH_PLI`): +- `pLI: float64`, `pNull: float64`, `pRec: float64` — probability of being + loss-of-function intolerant / null / recessive (Lek et al. 2016). + +--- + +## 13. `release/mutation_rate.ht` + +Public-release mutation rate. Produced by `--prepare-release-mutation-rate` +from §3 by selecting/renaming. + +- **Path:** `gs://gnomad/v{version}/constraint/release/gnomad.v{version}.mutation_rate.ht` +- **Resource fn:** [`get_release_mutation_ht`](resource_utils.py) +- **Key:** `context, ref, alt, methylation_level` +- A TSV mirror is exported to `release/gnomad.v{version}.mutation_rate.tsv`. + +### Globals +- `version: str`. +- `calculate_mu_params: struct{ac_cutoff, min_cov, max_cov, gerp_lower_cutoff, + gerp_upper_cutoff, downsampling_level, most_severe_consequence}` — + release-cleaned `calculate_mu_globals`. + +### Rows +- `context, ref, alt, methylation_level` — key. +- `mu: float64` — scalar mu at the canonical downsampling index (renamed from + `mu_snp`). +- `cpg: bool`, `transition: bool`, `mutation_type: str` — mutation-type tags. + +--- + +# Constraint-specific input tables + +These are constraint-pipeline-owned input tables (not part of `gnomad_methods`) +that are read by `prepare_ht_for_constraint_calculations` to annotate +`annotated_context.ht`. + +## `adj_r_per_context_methyl_genome_1kb_autosome.agg.ht` + +Per-context regional-depletion correction. Aggregated to 1kb autosomal +intervals over the genome reference; values vary by trinucleotide context. + +- **Path:** `gs://gnomad/v4.1/constraint/resources/annotations/ht/adj_r_per_context_methyl_genome_1kb_autosome.agg.ht` +- **Resource fn:** [`get_adj_r_ht`](resource_utils.py) +- **Key:** `interval` + +### Globals +*(none)* + +### Rows +- `interval: interval>` — 1kb autosomal interval. +- `adj_r: dict` — context-keyed correction. Look up with + `adj_r_ht[ht.locus].adj_r[ht.context]`. + +## `adj_r_syn_dnm_per_context_methyl_genome_1kb_autosome.agg.ht` + +Identical schema to `adj_r`, but computed from a synonymous de-novo-mutation +constraint baseline rather than the general per-context baseline. + +- **Path:** `gs://gnomad/v4.1/constraint/resources/annotations/ht/adj_r_syn_dnm_per_context_methyl_genome_1kb_autosome.agg.ht` +- **Resource fn:** [`get_syn_adj_r_ht`](resource_utils.py) +- **Key:** `interval` + +### Globals +*(none)* + +### Rows +- `interval: interval>`. +- `adj_r: dict` — context-keyed correction. Joined into + `annotated_context.syn_adj_r` as `adj_r_ht[ht.locus].adj_r[ht.context]`. + +--- + +## Regenerating this README + +The schemas above were captured against live GCS tables for version `4.1.1`. +To refresh after a pipeline rerun: + +```bash +source /Users/jgoodric/miniconda3/etc/profile.d/conda.sh && conda activate hail +export PATH="/Users/jgoodric/google-cloud-sdk/bin:$PATH" +python -c " +import hail as hl +hl.init(quiet=True, idempotent=True) +ht = hl.read_table('') +ht.describe() +print(hl.eval(ht.globals)) +" +``` + +The path conventions are sourced from [`get_constraint_data`](resource_utils.py) +and [`get_constraint_root`](resource_utils.py). When in doubt, instantiate a +resource via the corresponding `get_*` function and read `.ht().path`. From a1ab860f3beede4a13f08063cdadf29e8a584df2 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:00:05 -0600 Subject: [PATCH 37/38] Correct documented rank bin ranges to 0-based in resources README Bins are computed as hl.int(rank * multiplier / n_rows) over 0-based ranks, so they run from 0 to multiplier - 1. The README documented them as 1-based. Assisted-by: ClaudeCode:claude-opus-5[1m] --- gnomad_constraint/resources/README.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/gnomad_constraint/resources/README.md b/gnomad_constraint/resources/README.md index 3d7fe7c1..4a733755 100644 --- a/gnomad_constraint/resources/README.md +++ b/gnomad_constraint/resources/README.md @@ -485,9 +485,9 @@ Pipeline-parameter globals (as §9) plus: - `oe_ci_discretized_poisson_rank: struct` and `oe_ci_gamma_rank: struct` — rank/bin annotations for the `upper` bound of each CI: - `upper_rank: int64` — 0-based ascending rank within the group. - - `upper_bin_percentile: int32` — percentile bin (1–99). - - `upper_bin_decile: int32` — decile bin (1–9). - - `upper_bin_sextile: int32` — sextile bin (1–5). + - `upper_bin_percentile: int32` — percentile bin (0–99). + - `upper_bin_decile: int32` — decile bin (0–9). + - `upper_bin_sextile: int32` — sextile bin (0–5). - `no_variants: bool` — as §9. - `constraint_flags: set` — transcript-level QC flags (`no_variants`, `not_in_gencode`, `outlier_*`, etc.). @@ -610,9 +610,9 @@ Each constraint group is exposed as a named struct (group names from - `lof` uses the Gamma CI (the LOEUF), and **additionally** carries rank fields on `oe_ci` (the only group in `RELEASE_GROUPS_WITH_RANK`): - `upper_rank: int64` — 0-based rank by `oe_ci.upper` ascending. - - `upper_bin_percentile: int32` — percentile bin (1–99). - - `upper_bin_decile: int32` — decile bin (1–9; LOEUF decile). - - `upper_bin_sextile: int32` — sextile bin (1–5). + - `upper_bin_percentile: int32` — percentile bin (0–99). + - `upper_bin_decile: int32` — decile bin (0–9; LOEUF decile). + - `upper_bin_sextile: int32` — sextile bin (0–5). - `gen_anc_obs: struct{global, afr, amr, eas, nfe, sas}` — per-genetic-ancestry observed counts. Each field is `array` parallel to `downsamplings.`. From 0470d69a3941f4d5df55aec33d17424ac9d1fb92 Mon Sep 17 00:00:00 2001 From: jkgoodrich <33063077+jkgoodrich@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:36:54 -0600 Subject: [PATCH 38/38] Split rank and bin annotations into a separate constraint metrics phase Rank and percentile bin annotations are no longer added inside compute_constraint_metrics. They are applied by the now-public compute_constraint_percentile_bins, which runs as a second phase over an intermediate constraint_metrics_pre_rank Table. The ranking moved from third to last in the orchestration; nothing downstream of it reads the rank or bin fields, so the annotations are additive and order-independent. This allows the ranking to be recomputed without rerunning the metrics, needed to reissue v4.1.1 as v4.1.2 with corrected bins (broadinstitute/gnomad_production#2202). --skip-pre-rank-metrics reuses the existing pre-rank Table and recomputes only the ranking, and --use-mane-select-over-canonical exposes the transcript filter that the ranking phase now takes directly. With --skip-pre-rank-metrics the step declares the pre-rank Table as an input and drops its upstream pipeline inputs, since the aggregated constraint group and gene quality metrics Tables are not read in that mode and need not exist for the target version. Without this the input existence check fails on a missing constraint_group.ht, which --overwrite does not suppress. Register 4.1.2 in VERSIONS and SITES_VERSION_MAP. CURRENT_VERSION is left at 4.1.1, so runs pass --version 4.1.2 explicitly. Assisted-by: ClaudeCode:claude-opus-5[1m] --- .../pipeline/constraint_pipeline.py | 67 ++++++++++++---- gnomad_constraint/resources/constants.py | 3 +- gnomad_constraint/resources/resource_utils.py | 78 ++++++++++++++++--- gnomad_constraint/utils/constraint.py | 28 +++---- 4 files changed, 130 insertions(+), 46 deletions(-) diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 4e4238ca..b696ecf2 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -51,6 +51,7 @@ aggregate_per_variant_expected_ht, calculate_mu_by_downsampling, compute_constraint_metrics, + compute_constraint_percentile_bins, compute_gene_quality_metrics, create_aggregated_expected_ht, create_per_variant_expected_ht, @@ -110,6 +111,7 @@ def main(args): models, directory_post_fix, path_post_fix, + skip_pre_rank_metrics=args.skip_pre_rank_metrics, ) try: @@ -380,22 +382,36 @@ def main(args): res = resources.compute_constraint_metrics res.check_resource_existence() - # Compute constraint metrics. - ht = res.constraint_group_ht.ht(read_args={"_n_partitions": 10000}) - compute_constraint_metrics( - ht=ht, - gencode_ht=constraint_res.get_gencode_ht(version), - gene_quality_metrics_ht=res.gene_quality_metrics_ht.ht(), - expected_values={ - "Null": args.expectation_null, - "Rec": args.expectation_rec, - "LI": args.expectation_li, - }, - min_diff_convergence=args.min_diff_convergence, - raw_z_outlier_threshold_lower_lof=args.raw_z_outlier_threshold_lower_lof, - raw_z_outlier_threshold_lower_missense=args.raw_z_outlier_threshold_lower_missense, - raw_z_outlier_threshold_lower_syn=args.raw_z_outlier_threshold_lower_syn, - raw_z_outlier_threshold_upper_syn=args.raw_z_outlier_threshold_upper_syn, + # Compute constraint metrics, excluding rank and bin annotations. + if args.skip_pre_rank_metrics: + logger.info( + "Skipping metrics computation, reusing %s.", + res.pre_rank_constraint_metrics_ht.path, + ) + else: + ht = res.constraint_group_ht.ht(read_args={"_n_partitions": 10000}) + compute_constraint_metrics( + ht=ht, + gencode_ht=constraint_res.get_gencode_ht(version), + gene_quality_metrics_ht=res.gene_quality_metrics_ht.ht(), + expected_values={ + "Null": args.expectation_null, + "Rec": args.expectation_rec, + "LI": args.expectation_li, + }, + min_diff_convergence=args.min_diff_convergence, + raw_z_outlier_threshold_lower_lof=args.raw_z_outlier_threshold_lower_lof, + raw_z_outlier_threshold_lower_missense=args.raw_z_outlier_threshold_lower_missense, + raw_z_outlier_threshold_lower_syn=args.raw_z_outlier_threshold_lower_syn, + raw_z_outlier_threshold_upper_syn=args.raw_z_outlier_threshold_upper_syn, + ).write(res.pre_rank_constraint_metrics_ht.path, overwrite=overwrite) + + # Add rank and bin annotations as a separate phase so they can be + # recomputed without rerunning the metrics above. + logger.info("Adding rank and percentile bin annotations...") + compute_constraint_percentile_bins( + res.pre_rank_constraint_metrics_ht.ht(), + use_mane_select_over_canonical=args.use_mane_select_over_canonical, ).write(res.constraint_metrics_ht.path, overwrite=overwrite) logger.info("Done with computing constraint metrics.") @@ -840,6 +856,25 @@ def main(args): type=int, default=1000, ) + compute_constraint_args.add_argument( + "--skip-pre-rank-metrics", + help=( + "Skip computing the constraint metrics and reuse the existing pre-rank" + " Table, recomputing only the rank and percentile bin annotations. Used to" + " reissue a release with corrected ranks without rerunning the pipeline." + ), + action="store_true", + ) + compute_constraint_args.add_argument( + "--use-mane-select-over-canonical", + help=( + "Use MANE Select rather than canonical transcripts when determining which" + " transcripts to rank, falling back to canonical for genes without a MANE" + " Select transcript." + ), + action=argparse.BooleanOptionalAction, + default=True, + ) compute_constraint_args.add_argument( "--min-diff-convergence", help=( diff --git a/gnomad_constraint/resources/constants.py b/gnomad_constraint/resources/constants.py index ed3bc169..5dd638c0 100644 --- a/gnomad_constraint/resources/constants.py +++ b/gnomad_constraint/resources/constants.py @@ -7,7 +7,7 @@ EXTENSIONS = ["ht", "tsv", "tsv.bgz", "he", "log"] """Valid file extensions for constraint pipeline resources.""" -VERSIONS = ["2.1.1", "4.0", "4.1", "4.1.1"] +VERSIONS = ["2.1.1", "4.0", "4.1", "4.1.1", "4.1.2"] """Supported gnomAD constraint pipeline versions.""" CURRENT_VERSION = "4.1.1" @@ -18,6 +18,7 @@ "4.0": "4.0", "4.1": "4.1", "4.1.1": "4.1", + "4.1.2": "4.1", } """Map from constraint pipeline version to gnomAD sites release version.""" diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 3a1c0daf..bbb92843 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -537,6 +537,28 @@ def get_constraint_metrics_dataset( ) +def get_pre_rank_constraint_metrics_dataset( + custom_vep_annotation: str = "transcript_consequences", **kwargs +) -> TableResource: + """ + Return TableResource of constraint metrics before rank and bin annotations are added. + + This is the output of :func:`compute_constraint_metrics` and the input to + :func:`compute_constraint_percentile_bins`. Keeping it as its own dataset + allows the ranking to be recomputed without rerunning the metrics. + + :param custom_vep_annotation: The VEP annotation used to customize the constraint + model (one of "transcript_consequences" or "worst_csq_by_gene"). + :return: TableResource of constraint metrics without rank annotations. + """ + return get_constraint_data( + "constraint_metrics_pre_rank", + sub_dir="metrics", + custom_vep_annotation=custom_vep_annotation, + **kwargs, + ) + + def get_gene_quality_metrics_ht(version: str = CURRENT_VERSION) -> TableResource: """ Return TableResource of per-transcript gene quality metrics. @@ -676,6 +698,7 @@ def get_constraint_resources( models: List[str] = ["plateau", "coverage"], directory_post_fix: Optional[str] = None, path_post_fix: Optional[str] = None, + skip_pre_rank_metrics: bool = False, ) -> PipelineResourceCollection: """ Get PipelineResourceCollection for all resources needed in the constraint pipeline. @@ -811,18 +834,49 @@ def get_constraint_resources( }, pipeline_input_steps=[preprocess_data, calculate_mutation_rate, build_models], ) - compute_constraint_metrics = PipelineStepResourceCollection( - "--compute-constraint-metrics", - output_resources={ - "constraint_metrics_ht": get_constraint_metrics_dataset( - custom_vep_annotation, **common_params, path_post_fix=path_post_fix - ) - }, - pipeline_input_steps=[ - aggregate_by_constraint_groups, - compute_gene_quality_metrics_step, - ], - ) + if skip_pre_rank_metrics: + # Only the ranking is rerun, over an existing pre-rank Table. The + # upstream metrics inputs are not read, and the pre-rank Table becomes + # an input rather than an output. + compute_constraint_metrics = PipelineStepResourceCollection( + "--compute-constraint-metrics --skip-pre-rank-metrics", + input_resources={ + "pre-rank constraint metrics": { + "pre_rank_constraint_metrics_ht": ( + get_pre_rank_constraint_metrics_dataset( + custom_vep_annotation, + **common_params, + path_post_fix=path_post_fix, + ) + ) + } + }, + output_resources={ + "constraint_metrics_ht": get_constraint_metrics_dataset( + custom_vep_annotation, **common_params, path_post_fix=path_post_fix + ), + }, + ) + else: + compute_constraint_metrics = PipelineStepResourceCollection( + "--compute-constraint-metrics", + output_resources={ + "pre_rank_constraint_metrics_ht": ( + get_pre_rank_constraint_metrics_dataset( + custom_vep_annotation, + **common_params, + path_post_fix=path_post_fix, + ) + ), + "constraint_metrics_ht": get_constraint_metrics_dataset( + custom_vep_annotation, **common_params, path_post_fix=path_post_fix + ), + }, + pipeline_input_steps=[ + aggregate_by_constraint_groups, + compute_gene_quality_metrics_step, + ], + ) prepare_release = PipelineStepResourceCollection( "--prepare-release", output_resources={ diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 87e3409c..f435601b 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -1651,7 +1651,7 @@ def _compute_z_scores(ht: hl.Table) -> hl.Table: return ht -def _compute_percentile_bins( +def compute_constraint_percentile_bins( ht: hl.Table, use_mane_select_over_canonical: bool = True, ) -> hl.Table: @@ -1663,7 +1663,7 @@ def _compute_percentile_bins( ``CONSTRAINT_GRANULARITIES`` and annotates bins via :func:`annotate_constraint_percentile_bins`. - :param ht: Table output by :func:`_compute_z_scores`. + :param ht: Table output by :func:`compute_constraint_metrics`. :param use_mane_select_over_canonical: Use MANE Select rather than canonical transcripts for filtering when determining ranks. Default is True. :return: Table with rank, decile, and percentile bin annotations. @@ -1744,7 +1744,7 @@ def _compute_pli_scores( """ Compute pLI, pNull, and pRec scores for the HC LoF constraint group. - :param ht: Table output by :func:`_compute_percentile_bins`. + :param ht: Table output by :func:`_compute_z_scores`. :param expected_values: Dictionary containing the expected OE values for 'Null', 'Rec', and 'LI' to use as starting values. Default is ``PLI_EXPECTED_VALUES``. :param min_diff_convergence: Minimum iteration change in LI to consider the EM @@ -1790,7 +1790,6 @@ def compute_constraint_metrics( raw_z_outlier_threshold_lower_missense: float = -8.0, raw_z_outlier_threshold_lower_syn: float = -8.0, raw_z_outlier_threshold_upper_syn: float = 8.0, - use_mane_select_over_canonical: bool = True, ) -> hl.Table: """ Compute constraint metrics for synonymous, missense, and pLoF variants. @@ -1800,10 +1799,13 @@ def compute_constraint_metrics( 1. Annotate OE ratios, confidence intervals, raw z-scores, and per-group flags (:func:`_annotate_oe_ci_z`). 2. Normalize z-scores and union constraint flags (:func:`_compute_z_scores`). - 3. Add OE upper CI rank, decile, and percentile bins - (:func:`_compute_percentile_bins`). - 4. Compute pLI / pNull / pRec scores (:func:`_compute_pli_scores`). - 5. Annotate with gene quality metrics and GENCODE transcript annotations. + 3. Compute pLI / pNull / pRec scores (:func:`_compute_pli_scores`). + 4. Annotate with gene quality metrics and GENCODE transcript annotations. + + Rank, decile, and percentile bin annotations are *not* added here. They are + applied separately by :func:`compute_constraint_percentile_bins`, which + takes the output of this function. Keeping them in a separate phase means + the ranking can be recomputed without rerunning the metrics. .. note:: @@ -1829,10 +1831,8 @@ def compute_constraint_metrics( synonymous variants. Default is -8.0. :param raw_z_outlier_threshold_upper_syn: Upper raw z-score outlier threshold for synonymous variants. Default is 8.0. - :param use_mane_select_over_canonical: Use MANE Select rather than canonical - transcripts for filtering when determining ranks. Default is True. :return: Table with pLI scores, OE ratios, confidence intervals, z-scores, - percentile bins, gene quality metrics, and GENCODE annotations. + gene quality metrics, and GENCODE annotations. """ # Map each consequence category to its (lower, upper) raw z-score # outlier bounds. LoF and missense are one-sided (only lower bound); @@ -1857,12 +1857,6 @@ def compute_constraint_metrics( ht = _compute_z_scores(ht) ht = ht.checkpoint(new_temp_file("constraint_metrics.z_scores", "ht")) - # Rank transcripts by gamma OE upper CI and assign - # percentile/decile/sextile bins. Thresholds are computed on MANE - # Select transcripts and stored as globals. - ht = _compute_percentile_bins(ht, use_mane_select_over_canonical) - ht = ht.checkpoint(new_temp_file("constraint_metrics.percentile_bins", "ht")) - # Run the EM algorithm to compute pLI/pNull/pRec from HC LoF # observed vs expected counts. ht = _compute_pli_scores(ht, expected_values, min_diff_convergence)