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..c8c440cf --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,93 @@ +# gnomad-constraint Claude Reference + +## 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`. + +## Key Constants (`constants.py`) + +```python +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") +``` + +## 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 +# 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__*' '*.RData' '*/renv/*' '*.DS_Store' '*.png' '*.pdf' '*.ipynb' && \ + cd && \ + 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 \ + 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/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 diff --git a/gnomad_constraint/pipeline/constraint_pipeline.py b/gnomad_constraint/pipeline/constraint_pipeline.py index 77163354..b696ecf2 100644 --- a/gnomad_constraint/pipeline/constraint_pipeline.py +++ b/gnomad_constraint/pipeline/constraint_pipeline.py @@ -23,29 +23,45 @@ import argparse import logging -from typing import List 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.utils.reference_genome import get_reference_genome -from gnomad_qc.resource_utils import ( - PipelineResourceCollection, - PipelineStepResourceCollection, +from gnomad.resources.grch38.reference_data import lcr_intervals, seg_dup_intervals +from gnomad.utils.constraint import ( + build_models, + calculate_gerp_cutoffs, + explode_downsamplings_oe, ) -from hail.utils.misc import new_temp_file +from gnomad.utils.file_utils import print_global_struct +from gnomad.utils.reference_genome import get_reference_genome 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.resources.resource_utils import ( + filter_for_test, + get_adj_r_ht, + get_syn_adj_r_ht, +) from gnomad_constraint.utils.constraint import ( - add_vep_context_annotations, - annotate_context_ht, - apply_models, - calculate_gerp_cutoffs, + aggregate_by_constraint_groups, + aggregate_per_variant_expected_ht, calculate_mu_by_downsampling, compute_constraint_metrics, - create_observed_and_possible_ht, + compute_constraint_percentile_bins, + compute_gene_quality_metrics, + create_aggregated_expected_ht, + 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, + prepare_release_mutation_ht, ) logging.basicConfig( @@ -56,294 +72,46 @@ logger.setLevel(logging.INFO) -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. - """ - 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 %s HT to chr20, chrX, and chrY for testing...", - data_type, - ) - ht = hl.filter_intervals(ht, keep) - - 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"], -) -> 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"]. - :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 - prepare_context = PipelineStepResourceCollection( - "--prepare-context-ht", - output_resources={ - "annotated_context_ht": constraint_res.get_annotated_context_ht( - version, use_v2_release_context_ht, test - ) - }, - 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, - }, - }, - ) - # 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", - 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") - }, - 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], - ) - calculate_mutation_rate = PipelineStepResourceCollection( - "--calculate-mutation-rate", - output_resources={ - "mutation_ht": constraint_res.get_mutation_ht( - version, test, use_v2_release_mutation_ht - ) - }, - 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 - }, - }, - 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) - for m in models - for r in regions - }, - pipeline_input_steps=[create_training_set], - ) - apply_models = PipelineStepResourceCollection( - "--apply-models", - output_resources={ - f"apply_{r}_ht": constraint_res.get_predicted_proportion_observed_dataset( - custom_vep_annotation, version, r, test - ) - for r in regions - }, - pipeline_input_steps=[preprocess_data, 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 - ) - }, - pipeline_input_steps=[apply_models], - ) - - export_tsv = PipelineStepResourceCollection( - "--export-tsv", - output_resources={ - "constraint_metrics_tsv": constraint_res.get_constraint_tsv_path( - version, test - ), - "downsampling_constraint_metrics_tsv": ( - constraint_res.get_downsampling_constraint_tsv_path(version, test) - ), - }, - pipeline_input_steps=[compute_constraint_metrics], - ) - - # 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": apply_models, - "compute_constraint_metrics": compute_constraint_metrics, - "export_tsv": export_tsv, - } - ) - - return constraint_pipeline - - def main(args): """Execute the constraint pipeline.""" hl.init( 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 + directory_post_fix = args.directory_post_fix + path_post_fix = args.path_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: + if version not in 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( + resources = constraint_res.get_constraint_resources( version, - use_v2_release_mutation_ht, - args.use_v2_release_context_ht, custom_vep_annotation, overwrite, test, models, + directory_post_fix, + path_post_fix, + skip_pre_rank_metrics=args.skip_pre_rank_metrics, ) try: @@ -353,85 +121,32 @@ 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 {} + # 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(), + syn_adj_r_ht=get_syn_adj_r_ht(), ) + ht.write(res.annotated_context_ht.path, overwrite) - 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() - - # 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,158 +155,224 @@ 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 = filter_for_test(ht, use_gene_list=test_gene_list) if test else ht + ht = prepare_ht_for_constraint_calculations( + 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.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 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( + 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.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( + "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 = 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, " + "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( - "Done computing expected variant count and observed:expected ratio." + "Aggregating observed and expected variant counts by constraint groups..." ) + # Use new shuffle method to prevent shuffle errors. + hl._set_flags(use_new_shuffle="1") + + 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(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_constraint_metrics: logger.info( @@ -601,75 +382,98 @@ 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") - ) + # 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.") - # Compute constraint metrics. - compute_constraint_metrics( - ht=union_ht, - gencode_ht=constraint_res.get_gencode_ht(version), - pops=pops, - keys=tuple( - [ - i - for i in list(union_ht.key) - if i - in ["gene", "transcript", "canonical", "mane_select", "gene_id"] - ] - ), - 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, - # 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 + if args.prepare_release: + logger.info("Preparing constraint metrics Table for release...") + res = resources.prepare_release + res.check_resource_existence() + + constraint_ht = res.constraint_metrics_ht.ht() + + 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.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, ) - logger.info("Done with computing constraint metrics.") + release_mutation_ht.write(res.release_mutation_ht.path, overwrite=overwrite) - if args.export_tsv: - res = resources.export_tsv + 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() - logger.info("Exporting constraint tsv...") + release_ht = res.release_ht.ht() - 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: - downsampling_ht = explode_downsamplings_oe( - ht, - downsampling_meta=hl.eval(ht.apply_model_params.downsampling_meta), - ) + 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.") - # 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"] - } + 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( + 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...") - 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__": @@ -682,10 +486,22 @@ 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", + 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, ) parser.add_argument( "--test", @@ -704,11 +520,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 +529,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 +541,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, + preprocess_args = parser.add_argument_group( + "Preprocess data args", + "All arguments used for preprocessing data for downstream steps.", ) - - mutation_rate_args = parser.add_argument_group( - "Calculate mutation rate args", - "Arguments used for calculating the mutation rate.", - ) - - 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 +559,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 +568,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 +578,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 +588,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 None." + ), + type=int, + default=None, + ) + 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 +693,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 +727,50 @@ def main(args): ), action="store_true", ) - build_models_args.add_argument( - "--upper-cov-cutoff", + cov_model_type = build_models_args.add_argument( + "--use-logarithmic-coverage-model", 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." + "Use a logarithmic model for low coverage sites when building and applying " + "the coverage model." ), - type=int, - default=100, + action="store_true", ) - build_models_args.add_argument( - "--high-cov-definition", + parser.add_argument( + "--apply-models-per-variant", 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." + "Apply plateau and coverage models to variants in exome sites Table and" + " context Table to compute expected variant counts per variant." ), - 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-aggregated", 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 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." ), - 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 +779,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 +788,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" @@ -983,13 +796,45 @@ 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) + + 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", + ) + 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", ) - 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) + 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.", @@ -1005,14 +850,31 @@ 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( + "--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=( @@ -1024,7 +886,6 @@ def main(args): type=float, default=0.001, ) - compute_constraint_args.add_argument( "--expectation-null", help=( @@ -1063,7 +924,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 +954,52 @@ def main(args): type=float, default=8.0, ) - parser.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( + "--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=( + "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", ) - - compute_constraint_args._group_actions.append(populations) args = parser.parse_args() main(args) diff --git a/gnomad_constraint/resources/README.md b/gnomad_constraint/resources/README.md new file mode 100644 index 00000000..4a733755 --- /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 (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.). +- `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 (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.`. +- `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`. diff --git a/gnomad_constraint/resources/constants.py b/gnomad_constraint/resources/constants.py new file mode 100644 index 00000000..5dd638c0 --- /dev/null +++ b/gnomad_constraint/resources/constants.py @@ -0,0 +1,240 @@ +"""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", "4.1.1", "4.1.2"] +"""Supported gnomAD constraint pipeline versions.""" + +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", + "4.1.2": "4.1", +} +"""Map from constraint pipeline version to gnomAD sites release 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. +""" + +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. + +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. +""" + +CALIBRATION_GROUPING = ("genomic_region", "build_model", "cpg", "exomes_coverage") +""" +Annotations used to group variants for the mutation rate calibration. +""" + +AGGREGATE_SUM_FIELDS = ( + "mu_snp", + "mu", + "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.""" + +# --------------------------------------------------------------------------- +# 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 +# --------------------------------------------------------------------------- + +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"] +"""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.""" diff --git a/gnomad_constraint/resources/resource_utils.py b/gnomad_constraint/resources/resource_utils.py index 3f497786..bbb92843 100644 --- a/gnomad_constraint/resources/resource_utils.py +++ b/gnomad_constraint/resources/resource_utils.py @@ -1,20 +1,36 @@ -"""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 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, TableResource, 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, + CUSTOM_VEP_ANNOTATIONS, + DATA_TYPES, + EXTENSIONS, + MODEL_TYPES, + SITES_VERSION_MAP, + VERSIONS, ) -from gnomad_qc.v4.resources.release import release_coverage, release_sites logging.basicConfig( format="%(asctime)s (%(name)s %(lineno)s): %(message)s", @@ -23,40 +39,47 @@ logger = logging.getLogger("constraint_pipeline") logger.setLevel(logging.INFO) -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. -""" +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. -COVERAGE_CUTOFF = 40 -""" -Minimum median exome coverage differentiating high coverage sites from low coverage sites. + If version is specified, return the genome build of the version as a string. -Low coverage sites require an extra calibration when computing the proportion of expected variation. -""" + :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" -# VEP context Table. def get_vep_context_ht(version: str) -> TableResource: """ Return VEP context Table corresponding to specified gnomAD version. @@ -74,21 +97,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) -> 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. - """ - return ( - f"gs://gnomad-tmp/gnomad_v{version}_testing/constraint" - if test - else f"gs://gnomad/v{version}/constraint" - ) - - def get_sites_resource(data_type: str, version: str = CURRENT_VERSION) -> BaseResource: """ Return genomes or exomes sites Table. @@ -98,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!" @@ -135,12 +144,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_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) - return TableResource(path=tmp_path) + return ref_grch38.methylation_sites else: raise ValueError("Build must be one of 'GRCh37' or 'GRCh38'.") @@ -164,317 +168,761 @@ def get_coverage_ht( return gnomad_grch38.coverage(data_type) -def get_mutation_ht( - version: str = CURRENT_VERSION, - test: bool = False, - use_v2_release_mutation_ht: bool = False, -) -> TableResource: +def get_gencode_ht(version: str) -> hl.Table: """ - Return mutation Table that includes the baseline mutation rate for each substitution and context. + Retrieve GENCODE Table with transcript version annotations. - :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. + 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 re-imported with version + fields and checkpointed. + :return: Table of GENCODE data with version annotations. """ - 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", - ) + if int(version[0]) == 2: + return ref_grch37.gencode.ht() + elif int(version[0]) == 4: + 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: - check_param_scope(version) - return TableResource( - f"{get_constraint_root(version, test)}/mutation_rate/gnomad.v{version}.mutation_rate.ht" - ) + raise ValueError("Version must be within gnomAD v2 or v4.") -def get_annotated_context_ht( +def get_gencode_cds_ht( version: str = CURRENT_VERSION, - use_v2_context_ht: bool = False, - test: bool = False, ) -> TableResource: - """ - Return TableResource of annotated context Table. + """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`. - :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. + :return: TableResource of GENCODE CDS positions, keyed by locus with + ``transcript_id`` (array of transcript IDs whose CDS covers that + position). """ - 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=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, ) - - check_param_scope(version) - return TableResource( - f"{get_constraint_root(version, test)}/preprocessed_data/annotated_context.ht" + ).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_preprocessed_ht( - data_type: str, + +def get_constraint_root( version: str = CURRENT_VERSION, - genomic_region: str = "autosome_par", test: bool = False, -) -> TableResource: + post_fix: Optional[str] = None, + temp: bool = False, + sub_dir: Optional[str] = None, +) -> str: """ - Return TableResource of preprocessed genome, exomes, and context Table. + Return path to constraint root folder. + + :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. + """ + 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}" - The exome and genome Table will have annotations added by - `prepare_ht_for_constraint_calculations()` and VEP annotation from context Table. + 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}" - The context Table will have annotations added by - `prepare_ht_for_constraint_calculations()`. + return f"gs://gnomad/v{version}/{constraint_dir}" - :param data_type: One of "exomes", "genomes" or "context. + +def get_constraint_data( + name: str, + version: str = CURRENT_VERSION, + test: bool = False, + 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. + + :param name: Name of the constraint data to retrieve. :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. + :param test: Whether the Table is for testing purpose and only contains a subset of + the data. Default is False. + :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(version, genomic_region, data_type) - return TableResource( - f"{get_constraint_root(version, test)}/preprocessed_data/gnomad.v{version}.{data_type}.preprocessed.{genomic_region}.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_training_dataset( - version: str = CURRENT_VERSION, - genomic_region: str = "autosome_par", - test: bool = False, -) -> TableResource: + 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=directory_post_fix, + sub_dir=sub_dir, + temp=temp, + ) + path = f"{root_dir}/gnomad.v{version}.{name}{path_post_fix}.{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 training dataset with observed and possible variant count. + Return mutation Table that includes the baseline mutation rate for each substitution and context. + + :return: Mutation rate Table. + """ + 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`. - :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. + :return: TableResource of the release mutation rate Table. """ - check_param_scope(version, genomic_region) - return TableResource( - f"{get_constraint_root(version, test)}/training_data/gnomad.v{version}.constraint_training.{genomic_region}.ht" - ) + check_param_scope(version=version) + root = get_constraint_root(version=version) + return TableResource(f"{root}/release/gnomad.v{version}.mutation_rate.ht") -def get_training_tsv_path( - version: str = CURRENT_VERSION, - genomic_region: str = "autosome_par", - test: bool = False, -) -> str: +def get_release_mutation_tsv_path(version: str = CURRENT_VERSION) -> str: """ - Return tsv of training dataset with observed and possible variant count. + Return path for the release mutation rate TSV. :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. + :return: Path of the release mutation rate TSV. """ - check_param_scope(version, genomic_region) + check_param_scope(version=version) + root = get_constraint_root(version=version) + return f"{root}/release/gnomad.v{version}.mutation_rate.tsv" - return f"{get_constraint_root(version, test)}/training_data/gnomad.v{version}.constraint_training.{genomic_region}.tsv.bgz" +def get_release_constraint_ht(version: str = CURRENT_VERSION) -> TableResource: + """ + Return TableResource for the release constraint metrics Table. -def get_models( - model_type: str, - version: str = CURRENT_VERSION, - genomic_region: str = "autosome_par", - test: bool = False, -) -> ExpressionResource: + :param version: One of the release versions (`VERSIONS`). Default is + `CURRENT_VERSION`. + :return: TableResource of the release constraint metrics Table. """ - Return path to a HailExpression that contains desired model type. + 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 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. + :return: Path of the release constraint metrics TSV. """ - check_param_scope( - version=version, genomic_region=genomic_region, model_type=model_type + 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. + + :return: TableResource of annotated context Table. + """ + return get_constraint_data( + "annotated_context", sub_dir="preprocessed_data", **kwargs ) - return ExpressionResource( - f"{get_constraint_root(version, test)}/models/gnomad.v{version}.{model_type}.{genomic_region}.he" + + +def get_preprocessed_ht(**kwargs) -> TableResource: + """ + Return TableResource of preprocessed genome, exomes, and context Table. + + :return: TableResource of processed context Table. + """ + return get_constraint_data( + "context.preprocessed", sub_dir="preprocessed_data", **kwargs ) -def get_predicted_proportion_observed_dataset( - custom_vep_annotation: str = "transcript_consequences", - version: str = CURRENT_VERSION, - genomic_region: str = "autosome_par", - test: bool = False, +def get_training_dataset(**kwargs) -> TableResource: + """ + Return TableResource of training dataset with observed and possible variant count. + + :return: TableResource of training dataset. + """ + return get_constraint_data("constraint_training", sub_dir="training_data", **kwargs) + + +def get_training_tsv_path(**kwargs) -> str: + """ + Return tsv of training dataset with observed and possible variant count. + + :return: TSV path of training dataset. + """ + return get_constraint_data( + "constraint_training", sub_dir="training_data", extension="tsv.bgz", **kwargs + ) + + +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. + :return: Path to the specified model. + """ + 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", **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 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, + return get_constraint_data( + "per_variant_expected", + sub_dir="apply_models", custom_vep_annotation=custom_vep_annotation, + **kwargs, ) - return TableResource( - f"{get_constraint_root(version, test)}/predicted_proportion_observed/{custom_vep_annotation}/gnomad.v{version}.predicted_proportion_observed.{genomic_region}.ht" + + +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"). + :return: Path of the model. + """ + return get_constraint_data( + "per_variant_expected.aggregated", + sub_dir="apply_models", + custom_vep_annotation=custom_vep_annotation, + **kwargs, + ) + + +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. + + :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( - version: str = CURRENT_VERSION, - test: bool = False, + 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 get_constraint_data( + "constraint_metrics", + sub_dir="metrics", + custom_vep_annotation=custom_vep_annotation, + **kwargs, + ) + - return TableResource( - f"{get_constraint_root(version, test)}/metrics/gnomad.v{version}.constraint_metrics.ht" +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_constraint_tsv_path( - version: str = CURRENT_VERSION, - test: bool = False, -) -> 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. :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. + :return: TableResource of gene quality metrics. """ check_param_scope(version=version) + root = get_constraint_root(version=version) + return TableResource(f"{root}/metrics/gnomad.v{version}.gene_quality_metrics.ht") - return f"{get_constraint_root(version, test)}/metrics/tsv/gnomad.v{version}.constraint_metrics.tsv" - -def get_downsampling_constraint_tsv_path( - version: str = CURRENT_VERSION, - test: bool = False, -) -> str: +def get_lof_threshold_tsv_path(version: str = CURRENT_VERSION) -> str: """ - Return tsv path of downsampling observed and expected counts. + Return path for the LoF OE CI upper bin thresholds TSV. :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. + :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" - return f"{get_constraint_root(version, test)}/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, -) -> Union[str, None]: +def get_release_downsampling_tsv_path(version: str = CURRENT_VERSION) -> str: """ - Check if the specified version, genomic region, and data type are in the scope of the constraint pipeline. + Return path for the release per-genetic-ancestry downsampling TSV. - 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. + :param version: One of the release versions (`VERSIONS`). Default is + `CURRENT_VERSION`. + :return: Path of the release downsampling TSV. """ - 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: - 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" + 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, version: str = CURRENT_VERSION) -> 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)}/logging/{name}.log" + return get_constraint_data( + name, sub_dir="logging", extension="log", test=True, **kwargs + ) -def get_checkpoint_path( - name: str, version: str = CURRENT_VERSION, mt: bool = False -) -> str: +def get_checkpoint_path(name: str, **kwargs) -> TableResource: """ - Create a checkpoint path for Table or MatrixTable. + Create a checkpoint TableResource. - :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. + :param name: Name of intermediate Table. + :return: Output checkpoint TableResource. """ - return f'{get_constraint_root(version, test=True)}/checkpoint_files/{name}.{"mt" if mt else "ht"}' + return get_constraint_data(name, sub_dir="checkpoint_files", test=True, **kwargs) -def get_gencode_ht(version: str) -> hl.Table: +def filter_for_test( + ht: hl.Table, + use_gene_list: bool = False, +) -> hl.Table: """ - Retrieve GENCODE Table. + Filter ``ht`` to chr20, chrX, and chrY or a gene list for testing. - :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 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. """ - if int(version[0]) == 2: - return ref_grch37.gencode.ht() - elif int(version[0]) == 4: - return ref_grch38.gencode.ht() + rg = get_reference_genome(ht.locus) + 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 + "X:57313139-57515629", # FAAH2 + "Y:2803112-2850547", # ZFY + ] + else: + keep_regions = [ + "chr1:55039447-55064852", # PCSK9 + "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: - raise ValueError("Version must be within gnomAD v2 or v4.") + 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_syn_adj_r_ht() -> hl.Table: + """ + Read the aggregated synonymous DNM adj_r per-context methylation genome 1kb autosome Table. + + :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.agg.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, + skip_pre_rank_metrics: bool = False, +) -> 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], + ) + 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={}, + 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], + ) + 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], + ) + 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={ + "release_ht": get_release_constraint_ht(version=version), + }, + 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={ + "release_tsv": get_release_constraint_tsv_path(version=version), + "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], + ) + + # Add all steps to the constraint pipeline resource collection. + constraint_pipeline.add_steps( + { + "prepare_context": prepare_context, + "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, + "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, + "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, + } + ) + + return constraint_pipeline diff --git a/gnomad_constraint/utils/constraint.py b/gnomad_constraint/utils/constraint.py index 07d710b1..f435601b 100644 --- a/gnomad_constraint/utils/constraint.py +++ b/gnomad_constraint/utils/constraint.py @@ -1,43 +1,68 @@ """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 -import numpy as np +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, + assemble_constraint_context_ht, + build_constraint_consequence_groups, calculate_raw_z_score, calculate_raw_z_score_sd, - collapse_strand, - compute_expected_variants, + calibration_model_group_expr, + compute_percentile_thresholds, compute_pli, - count_variants_by_group, + count_observed_and_possible_by_group, get_constraint_flags, - get_downsampling_freq_indices, - oe_aggregation_expr, oe_confidence_interval, - trimer_from_heptamer, + rank_array_element_metrics, + variant_observed_and_possible_expr, ) -from gnomad.utils.filtering import ( - add_filters_expr, - filter_by_numeric_expr_range, - filter_for_mu, - filter_to_autosomes, +from gnomad.utils.file_utils import ( + convert_multi_array_to_array_of_structs, + print_global_struct, ) -from gnomad.utils.reference_genome import get_reference_genome +from gnomad.utils.filtering import add_filters_expr from gnomad.utils.vep import ( - add_most_severe_csq_to_tc_within_vep_root, - filter_vep_transcript_csqs, + CSQ_CODING, + filter_vep_transcript_csqs_expr, + mane_select_over_canonical_filter_expr, + update_loftee_end_trunc_filter, ) -from hail.utils.misc import new_temp_file - -from gnomad_constraint.resources.resource_utils import ( +from hail.utils.misc import divide_null, new_temp_file + +from gnomad_constraint.resources.constants import ( + ADJ_FREQ_META, + AGGREGATE_SUM_FIELDS, + CALIBRATION_GROUPING, + CLASSIC_LOF_ANNOTATIONS, + CONSTRAINT_GRANULARITIES, COVERAGE_CUTOFF, - get_checkpoint_path, + 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, + SFS_BIN_CUTOFFS, ) logging.basicConfig( @@ -48,1170 +73,2241 @@ logger.setLevel(logging.INFO) -def add_vep_context_annotations( - ht: hl.Table, annotated_context_ht: hl.Table +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, + syn_adj_r_ht: hl.Table, + sfs_bin_cutoffs: Tuple[float, ...] = SFS_BIN_CUTOFFS, ) -> hl.Table: """ - Add annotations from VEP context Table to gnomAD data. + 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, syn_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 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. + """ + # 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, + ) - Function adds the following annotations: - - context - - methylation - - coverage - - gerp + # 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() + ) - Function drops `a_index`, `was_split`, and`colocated_variants` annotations from - gnomAD data. + # 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() - .. note:: - Function expects that multiallelic variants in the VEP context Table have been - split. + 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], + syn_adj_r=syn_adj_r_ht[ht.locus].adj_r[ht.context], + sfs_bin=sfs_bin_expr, + ) - 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. +# 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]]]: """ - 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 + Filter the frequency array for constraint calculations. + + The frequency array is filtered to include only adj frequencies for + 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_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. + :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 = [ADJ_FREQ_META] + + if gen_ancs is not None: + meta_keep += [{**ADJ_FREQ_META, gen_anc_label: gen_anc} for gen_anc in gen_ancs] + + if downsamplings is not None: + downsampling_gen_ancs = ["global"] + (downsampling_gen_ancs or []) + meta_keep += [ + {**ADJ_FREQ_META, gen_anc_label: gen_anc, "downsampling": str(ds)} + for gen_anc in downsampling_gen_ancs + for ds in downsamplings + ] + 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]) -def prepare_ht_for_constraint_calculations( + return freq_expr, meta_keep + + +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]: + """ + 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:: + + 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. + """ + # 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, + gen_ancs=None, + downsamplings=[downsampling_level], + downsampling_gen_ancs=gen_ancs, + gen_anc_label="pop", + ) + downsampling_idx = genomes_freq_meta.index( + {**ADJ_FREQ_META, "pop": "global", "downsampling": str(downsampling_level)} + ) + + # Filter to autosomal sites (remove pseudoautosomal regions). + keep_expr = locus_expr.in_autosome() + + # 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) + + # 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) + + # 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"] + ] + ) + + # 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, variant_observed_and_possible_expr(genomes_freq_expr) + ), + ) + 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"], + ) + + return obs_pos_expr, obs_pos_globals + + +def get_exome_coverage_expr( ht: hl.Table, - require_exome_coverage: bool = True, - coverage_metric: str = "exome_coverage", -) -> hl.Table: + exome_coverage_metric: str = "AN_percent", +) -> hl.expr.Int32Expression: """ - 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) - - # Add annotations for 'ref' and 'alt'. - ht = ht.annotate(ref=ht.alleles[0], alt=ht.alleles[1]) - - # 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}}}")) - - # Annotate mutation type (such as "CpG", "non-CpG transition", "transversion") and - # collapse strands to deduplicate the context. - ht = annotate_mutation_type(collapse_strand(ht)) - - # 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" - ) - - # 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)) + 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). + 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({**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 = ( + 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) # Index 0 is adj (all samples). + ) + + cov_expr = hl.int((ht.AN.exomes / an_count) * 100) else: raise ValueError( - "No 'methylation_level' or 'MEAN' found in 'methylation' annotation." + f"Exome coverage metric must be one of ['median', 'AN', 'AN_percent'], not {exome_coverage_metric}" ) - # 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) + logger.info("Setting 'exome_coverage' to %s", exome_coverage_metric) + + return cov_expr + + +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, +) -> Tuple[hl.expr.StructExpression, hl.expr.StructExpression]: + """ + 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 = sorted(map(int, list(set(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 + # genetic ancestry groups 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. + 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), + variant_observed_and_possible_expr(exomes_freq_expr, max_af=max_af), ), - exome_coverage=ht.coverage.exomes[exome_median_cov_field], + ) + 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, ) - # 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 - # 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 - ), + +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, +) -> hl.expr.StructExpression: + """ + 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``. + + :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: 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'. + 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'" ) - # Add most_severe_consequence annotation to 'transcript_consequences' within the - # vep root annotation. - ht = add_most_severe_csq_to_tc_within_vep_root(ht) + # 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, + ) + + # 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}, + cpg_in_high_only=True, + ) + + return hl.or_missing(syn_csq_expr.length() > 0, build_expr) + - if require_exome_coverage: - # Filter out locus with undefined coverage_metric. - ht = ht.filter(hl.is_defined(ht[coverage_metric])) +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`` + - ``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", + "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 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 "mane_select". + :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, + ) + + # 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, + ) + + # 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, + ) + + # 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}, + ) + + # 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 create_observed_and_possible_ht( - exome_ht: hl.Table, - context_ht: hl.Table, +def create_training_set( + ht: hl.Table, mutation_ht: hl.Table, - 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: """ - Count the observed variants and possible variants by substitution, context, methylation level, and additional `grouping`. + 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.transmute(**ht.calibrate_mu) + # 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, + "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",) + + MUTATION_TYPE_FIELDS + + CALIBRATION_GROUPING, + partition_hint=partition_hint, + ) - 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 + # Annotate with mutation rate. + ht = annotate_with_mu(ht, mutation_ht) - 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. + return ht - 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. - """ - 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 - ) - 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] - - # 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) - ) - if filter_coverage_over_0: - keep_criteria &= exome_ht[coverage_metric] > 0 - - keep_annotations += grouping - - logger.info("Setting keep annotations to %s", keep_annotations) - - # Keep variants that satisfy the criteria above. - filtered_exome_ht = exome_ht.filter(keep_criteria) - - # Filter context ht to sites with defined exome coverage. - context_ht = context_ht.filter(hl.is_defined(context_ht[coverage_metric])) - - # 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 + +def _prepare_ht_for_apply_models( + ht: hl.Table, + custom_vep_annotation: str = "transcript_consequences", + use_mane_select: bool = False, +) -> Tuple[hl.Table, List[str]]: + """ + Prepare a preprocessed Table for model application. + + 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 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: Tuple of (prepared Table, list of VEP grouping field names). + """ + 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'." ) - # 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, - ) + include_canonical_group = custom_vep_annotation != "worst_csq_by_gene" + include_mane_select_group = include_canonical_group and use_mane_select - # 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) + ht = ht.annotate(**ht.calibrate_mu) + ht = ht.filter(hl.is_defined(ht.apply_model) & (ht.possible_variants > 0)) - # 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) + 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, ) - # 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, + return ht, groupings + + +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), + ht.possible_variants, + coverage_model=coverage_model, + coverage_expr=ht.exomes_coverage, + model_group_expr=ht.apply_model, + log10_coverage=log10_coverage, + ) ) - 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")) - # Annotate the Table with 'cpg' and 'mutation_type' (one of "CpG", "non-CpG - # transition", or "transversion"). - ht = annotate_mutation_type(ht) +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. - if global_annotation: + :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( - **{global_annotation: hl.struct(max_af=max_af, pops=pops)} + apply_models_globals=ht.apply_models_globals.annotate(**model_params) ) - + else: + ht = ht.annotate_globals(apply_models_globals=model_params) return ht -def apply_models( - exome_ht: hl.Table, - context_ht: hl.Table, +def create_per_variant_expected_ht( + ht: hl.Table, mutation_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, + custom_vep_annotation: str = "transcript_consequences", + use_mane_select: bool = False, ) -> 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. - """ - # Filter context ht to sites with defined exome coverage_metric. - context_ht = context_ht.filter(hl.is_defined(context_ht[coverage_metric])) - - if low_coverage_filter is not None: - context_ht = context_ht.filter( - context_ht[coverage_metric] >= low_coverage_filter - ) - exome_ht = exome_ht.filter(exome_ht[coverage_metric] >= low_coverage_filter) - - # Add necessary constraint annotations for grouping. - 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'." - ) + 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()) - else: - vep_annotation = "transcript_consequences" - 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, groupings = _prepare_ht_for_apply_models( + ht, custom_vep_annotation, use_mane_select ) - exome_ht, grouping = annotate_exploded_vep_for_constraint_groupings( - ht=exome_ht, - coverage_expr=exome_ht[coverage_metric], - 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 = _apply_constraint_models(ht, plateau_models, coverage_model, log10_coverage) + ht = _annotate_apply_models_globals( + ht, plateau_models, coverage_model, log10_coverage, groupings ) - # 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) + return ht.drop(*calibrate_mu_fields) + + +def aggregate_per_variant_expected_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 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 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. + """ + # 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 []), + *[ + g + for g in hl.eval(ht.apply_models_globals.groupings) + 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")) + + 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) + 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, + ) + + # 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")) + + # 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) + 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, + additional_grouping: Tuple[str] = ("methylation_level",), + total_mu: float = 1.2e-08, +) -> hl.Table: + """ + Calculate mutation rate. + + The returned Table includes the following annotations: + - context - trinucleotide genomic context. + - ref - the reference allele. + - alt - the alternate base. + - methylation_level - methylation_level. + - 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``. + + :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. + """ + # 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")) + + total_bases = ht.aggregate(hl.agg.sum(ht.possible_variants)) // 3 + logger.info( + "Total bases to use when calculating correction_factors: %f", total_bases + ) + + # Compute the proportion observed, which represents the relative mutability of each + # variant class. + 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, + ) + 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 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. + + 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 + 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.transcript, ht.mane_select, ht.canonical, ht.gene_id + ) else: - cov_value = ht.coverage + return ht.transcript.startswith("ENST") & ht.canonical + + +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 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 + 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. + 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. + """ + ci_fields = ["discretized_poisson", "gamma"] + + 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, + rank_field_prefix="upper_", + ) + + +def aggregate_by_constraint_groups( + ht: hl.Table, + keys: Tuple = ("gene", "transcript", "canonical"), + classic_lof_annotations: Tuple = CLASSIC_LOF_ANNOTATIONS, + additional_groupings: Optional[ + Dict[str, Dict[str, hl.expr.BooleanExpression]] + ] = None, + additional_grouping_combinations: Optional[List[List[str]]] = None, +) -> hl.Table: + """ + Aggregate observed and expected variant info for synonymous, missense, and pLoF variants. + + .. note:: + + The following annotations should be present in ``ht``: + + - modifier + - annotation + - observed_variants + - mu + - possible_variants + - expected_variants + + :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 + 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: 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( + 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") + ) - 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 + # Group by keys and get an aggregate sum of mu_snp, observed_variants, + # 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( + lambda f: hl.agg.filter(f, aggregate_constraint_metrics_expr(ht)), + ht.constraint_groups, ) ) - # 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, + # 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 ) - 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, + + # 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( + ht.constraint_groups.map( + lambda x: (hl.or_else(x.expected_variants[0], 0) > 0) ) ) + ) - # 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 + # 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" ) - ] + ) + ) - # Remove coverage from grouping. - grouping = list(grouping) - grouping.remove("coverage") + return ht + + +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() - # Aggregate the sum aggregators grouped by `grouping`. - ht = ( - ht.group_by(*grouping) - .partition_hint(expected_variant_partition_hint) - .aggregate(**agg_expr) + # 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)) - # TODO: Remove repartition once partition_hint bugs are resolved. - ht = ht.repartition(expected_variant_partition_hint) + # Explode by transcript and aggregate. + ht = ht.explode("transcript_id").cache() - # 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", - ) + return ht.group_by(transcript=ht.transcript_id).aggregate( + prop_bp_AN90=hl.agg.fraction(ht.exomes_coverage >= an_coverage_threshold), ) - # Compute the observed:expected ratio. - return ht.annotate(obs_exp=ht.observed_variants / ht.expected_variants) -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, - 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, +def _compute_site_quality_metrics( + ht: hl.Table, + gencode_cds_ht: hl.Table, ) -> hl.Table: + """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``. """ - Calculate mutation rate using the downsampling with size specified by `downsampling_level` in genome sites Table. + # 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])) + + # 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), + ) - 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)`. - The returned Table includes the following annotations: - - context - trinucleotide genomic context. - - ref - the reference allele. - - alt - the alternate base. - - methylation_level - methylation_level. - - downsampling_counts_{pop} - variant counts in downsamplings for populations - in `pops`. - - 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 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. +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``. """ - 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) - ) + # 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]), ) - # 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"]) + # 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() - # 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] + 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), ) - 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, +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`, + 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). + + :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 + :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``. + """ + # 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], ) - # 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 + # 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( + { + "low_exome_mapping_quality": ht.mean_AS_MQ < 50, + "low_exome_coverage": ht.prop_bp_AN90 < 0.1, + } ), - 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 - ) + return ht.key_by("transcript") - ht = ht.checkpoint(new_temp_file(prefix="constraint", extension="ht")) - total_bases = ht.aggregate(hl.agg.sum(ht.possible_variants)) // 3 - logger.info( - "Total bases to use when calculating correction_factors: %f", total_bases - ) +def _annotate_oe_ci_z( + ht: hl.Table, + z_thresholds: Dict[str, Tuple[Optional[float], Optional[float]]], +) -> hl.Table: + """ + Annotate constraint groups with OE ratio, confidence intervals, z-scores, and flags. - # 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 - ) - ) + For each constraint group's ``oe_info`` entries, adds: - # 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 - ), - } + - ``oe`` — observed / expected ratio. + - ``oe_ci_discretized_poisson`` — discretized Poisson CI. + - ``oe_ci_gamma`` — gamma-distribution CI. + - ``z_raw`` — raw z-score. - 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] + Then adds per-group ``flags`` based on z-score outlier thresholds. + + :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. + """ + # 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( + 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, + method="poisson", + ), + oe_ci_gamma=oe_confidence_interval( + oe_info.observed_variants, + oe_info.expected_variants, + method="gamma", + ), + z_raw=calculate_raw_z_score( + oe_info.observed_variants, oe_info.expected_variants + ), + ) + ) + ) ) + ) - ht = ht.annotate(**ann_expr).checkpoint( - new_temp_file(prefix="calculate_mu_by_downsampling", extension="ht") + # 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) + 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_thresholds.get( + "lof" if m.get("lof") else m.get("csq_set", "None"), + (None, None), + )[0], + z_thresholds.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) + ] ) + 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. + """ + # 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"}) + mis_idx = meta.index({"csq_set": "mis"}) + 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( - 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, + sd_raw_z=ht.aggregate( + hl.agg.filter( + ~ht.no_variants, + [ + calculate_raw_z_score_sd( + 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) + ], + ) + ) ) - return annotate_mutation_type(ht) + # 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( + 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 + ), + ) + + return ht -def add_oe_lof_upper_rank_and_bin( - ht: hl.Table, use_mane_select_over_canonical: bool = True +def compute_constraint_percentile_bins( + ht: hl.Table, + use_mane_select_over_canonical: bool = True, ) -> hl.Table: """ - Compute the rank and decile of the lof oe upper confidence interval for MANE Select or canonical ensembl transcripts. + Add OE upper CI rank and percentile bin annotations. - :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'. + 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`. + + :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. """ - # 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), + # 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(): + n_bins = len(bins) + 1 + pcts = [b / n_bins * 100 for b in 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( + ht, + percentiles=all_qs, + metric_expr=ht.constraint_groups[idx].oe_info[0].oe_ci_gamma.upper, + outlier_expr=outlier_expr, + transcript_filter_expr=get_transcript_filter_expr( + ht, mane_select_only=True + ), ) + for gran_name, pcts in gran_percentiles.items(): + thresholds[(gran_name, metric)] = [vals[p] for p in pcts] - genes = genes.annotate( - only_canonical=~(genes.mane_present) & (genes.canonical_present) + # 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 + } ) + ) - 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, - ) - else: - ms_ht = ht.filter((ht.canonical) & (ht.transcript.startswith("ENST"))) + # Annotate each transcript with its threshold-based bin assignment + # for every (granularity, metric) combination. + return annotate_constraint_percentile_bins(ht, thresholds, metric_group_idx) - # 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") - # Determine decile bins. - n_transcripts = ms_ht.count() - ms_ht = ms_ht.annotate( - upper_bin_decile=hl.int(ms_ht.upper_rank * 10 / n_transcripts) - ) +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. - # Add 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, - ) + :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 + 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 + + # 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( + 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, ) ) - return ht - def compute_constraint_metrics( ht: hl.Table, gencode_ht: hl.Table, - keys: Tuple[str] = ("gene", "transcript", "canonical"), - classic_lof_annotations: Tuple[str] = ( - "stop_gained", - "splice_donor_variant", - "splice_acceptor_variant", - ), - pops: Tuple[str] = (), + 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, - include_os: bool = False, - 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. + 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. 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:: - 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 pops: List of populations used to compute constraint metrics. Default is (). - :param expected_values: Dictionary containing the expected values for 'Null', + 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: 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 + :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: 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 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. + :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. + :return: Table with pLI scores, OE ratios, confidence intervals, z-scores, + gene quality metrics, and GENCODE annotations. """ - 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", + # 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), + "syn": ( + raw_z_outlier_threshold_lower_syn, + raw_z_outlier_threshold_upper_syn, + ), } - # 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], - } + # 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")) - if include_os: - # Filter to LoF annotations with LOFTEE HC or OS. - annotation_dict.update( - {"lof_hc_os": (ht.modifier == "HC") | (ht.modifier == "OS")} - ) - lof_ann.append("lof_hc_os") + # 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")) - # 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() - } + # 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")) + + # 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-level annotations from GENCODE (gene name, + # biotype, CDS length, coding exon count, etc.). + ht = add_gencode_transcript_annotations(ht, gencode_ht) + + return ht + + +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. + """ + # 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 [] ) - # 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 - ] + + # 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.annotate(**cg.oe_ci_gamma_rank), + **{ + 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 + }, ) - > 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, - ) + + # 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 + 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) + ) + + # 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( + 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 []) + }, ) - for ann in lof_ann + for i, name in enumerate(field_names) } + ht = ht.annotate(**cg_fields) - # 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, - ) - 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"} - ) - # 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, - ) + # 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) - 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") - ) + 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) - # 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"), - ) - for ann in oe_ann - } - ) - ) - ) + # 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])) - # Compute z-score from raw z-score and standard deviations. - ht = ht.annotate( + 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. + """ + # 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( **{ - ann: ht[ann].annotate(z_score=ht[ann].z_raw / ht.sd_raw_z[ann]) - for ann in oe_ann + 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 } ) - ht = ht.checkpoint(new_temp_file(prefix="z_scores", extension="ht")) + # 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( + **{ + gen_anc: [int(freq_meta[j]["downsampling"]) for j in indices] + for gen_anc, indices in gen_anc_ds_indices.items() + } + ) - # 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 - ) + if "max_af" in ht.globals: + global_kwargs["max_af"] = ht.globals.max_af - # Add transcript annotations from GENCODE. - ht = add_gencode_transcript_annotations(ht, gencode_ht) + global_kwargs["sd_raw_z"] = sd_raw_z_struct - return ht + # 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["loeuf_percentile_thresholds"] = ( + ht.globals.percentile_thresholds.lof + ) + return ht.select_globals(**global_kwargs) -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. +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. """ - # 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))) + # 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): + 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) + + # 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 + ) + return ht - # 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 + +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, ) - # 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]) + # 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) - # 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)) + return ht.select_globals(**global_kwargs) - # 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] +def flatten_release_ht(ht: hl.Table) -> hl.Table: + """ + Flatten the release constraint metrics Table for TSV export. - return cutoff_lower, cutoff_upper + 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) -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: + # 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: """ - Split multiallelic sites if needed and add 'methylation', 'coverage', and 'gerp' annotation to context Table with VEP annotation. + Convert the LoF OE CI upper bin thresholds global into a flat Table. - .. 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] - ) + Creates a Table with one row per (granularity, bin) pair, suitable for + TSV export. - # 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, + :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), ) - 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) - ) +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. + + 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. + + :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_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), + ) - return ht + metric_exprs = { + metric: ht.constraint_groups[idx].oe_info[0].oe_ci_gamma.upper + for metric, idx in metric_group_idx.items() + } + + return annotate_bins_by_threshold( + ht, + metric_exprs=metric_exprs, + thresholds=thresholds, + granularities=list(CONSTRAINT_GRANULARITIES), + ) 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"