Merge pull request #1034 from broadinstitute/sample-qc-filtered-callrate

jklugherz · web-flow · commit 48537ac41ecf · 2025-03-17T15:11:47.000-04:00
add sample qc task and filter_flags
diff --git a/v03_pipeline/lib/methods/sample_qc.py b/v03_pipeline/lib/methods/sample_qc.py
@@ -0,0 +1,66 @@
+import hail as hl
+from gnomad.sample_qc.pipeline import filter_rows_for_qc
+
+from v03_pipeline.lib.model import SampleType
+
+GNOMAD_FILTER_MIN_AF = 0.001
+GNOMAD_FILTER_MIN_CALLRATE = 0.99
+
+CALLRATE_LOW_THRESHOLD = 0.85
+CONTAMINATION_UPPER_THRESHOLD = 5
+WES_COVERAGE_LOW_THRESHOLD = 85
+WGS_CALLRATE_LOW_THRESHOLD = 30
+
+
+def call_sample_qc(
+    mt: hl.MatrixTable,
+    tdr_metrics_ht: hl.Table,
+    sample_type: SampleType,
+):
+    mt = mt.annotate_entries(
+        GT=hl.case()
+        .when(mt.GT.is_diploid(), hl.call(mt.GT[0], mt.GT[1], phased=False))
+        .when(mt.GT.is_haploid(), hl.call(mt.GT[0], phased=False))
+        .default(hl.missing(hl.tcall)),
+    )
+    mt = annotate_filtered_callrate(mt)
+    return annotate_filter_flags(mt, tdr_metrics_ht, sample_type)
+
+
+def annotate_filtered_callrate(mt: hl.MatrixTable) -> hl.MatrixTable:
+    filtered_mt = filter_rows_for_qc(
+        mt,
+        min_af=GNOMAD_FILTER_MIN_AF,
+        min_callrate=GNOMAD_FILTER_MIN_CALLRATE,
+        bi_allelic_only=True,
+        snv_only=True,
+        apply_hard_filters=False,
+        min_inbreeding_coeff_threshold=None,
+        min_hardy_weinberg_threshold=None,
+    )
+    callrate_ht = filtered_mt.select_cols(
+        filtered_callrate=hl.agg.fraction(hl.is_defined(filtered_mt.GT)),
+    ).cols()
+    return mt.annotate_cols(**callrate_ht[mt.col_key])
+
+
+def annotate_filter_flags(
+    mt: hl.MatrixTable,
+    tdr_metrics_ht: hl.Table,
+    sample_type: SampleType,
+) -> hl.MatrixTable:
+    mt = mt.annotate_cols(**tdr_metrics_ht[mt.col_key])
+    flags = {
+        'callrate': mt.filtered_callrate < CALLRATE_LOW_THRESHOLD,
+        'contamination': mt.contamination_rate > CONTAMINATION_UPPER_THRESHOLD,
+    }
+    if sample_type == SampleType.WES:
+        flags['coverage'] = mt.percent_bases_at_20x < WES_COVERAGE_LOW_THRESHOLD
+    else:
+        flags['coverage'] = mt.mean_coverage < WGS_CALLRATE_LOW_THRESHOLD
+
+    return mt.annotate_cols(
+        filter_flags=hl.array(
+            [hl.or_missing(filter_cond, name) for name, filter_cond in flags.items()],
+        ).filter(hl.is_defined),
+    )
diff --git a/v03_pipeline/lib/misc/io.py b/v03_pipeline/lib/misc/io.py
@@ -244,6 +244,24 @@ def import_imputed_sex(imputed_sex_path: str) -> hl.Table:
     return ht.key_by(ht.s)
 
 
+def import_tdr_qc_metrics(file_path: str) -> hl.Table:
+    ht = hl.import_table(
+        file_path,
+        types={
+            'contamination_rate': hl.tfloat32,
+            'percent_bases_at_20x': hl.tfloat32,
+            'mean_coverage': hl.tfloat32,
+        },
+    )
+    ht = ht.select(
+        s=ht.collaborator_sample_id,
+        contamination_rate=ht.contamination_rate,
+        percent_bases_at_20x=ht.percent_bases_at_20x,
+        mean_coverage=ht.mean_coverage,
+    )
+    return ht.key_by(ht.s)
+
+
 def import_remap(remap_path: str) -> hl.Table:
     ht = hl.import_table(remap_path)
     ht = ht.select(
diff --git a/v03_pipeline/lib/paths.py b/v03_pipeline/lib/paths.py
@@ -219,6 +219,22 @@ def relatedness_check_tsv_path(
     )
 
 
+def sample_qc_json_path(
+    reference_genome: ReferenceGenome,
+    dataset_type: DatasetType,
+    callset_path: str,
+) -> str:
+    return os.path.join(
+        pipeline_prefix(
+            Env.LOADING_DATASETS_DIR,
+            reference_genome,
+            dataset_type,
+        ),
+        'sample_qc',
+        f'{_callset_path_hash(callset_path)}.json',
+    )
+
+
 def remapped_and_subsetted_callset_path(
     reference_genome: ReferenceGenome,
     dataset_type: DatasetType,
diff --git a/v03_pipeline/lib/tasks/write_metadata_for_run.py b/v03_pipeline/lib/tasks/write_metadata_for_run.py
@@ -4,9 +4,11 @@
 import luigi
 import luigi.util
 
+from v03_pipeline.lib.model import FeatureFlag
 from v03_pipeline.lib.paths import (
     metadata_for_run_path,
     relatedness_check_tsv_path,
+    sample_qc_json_path,
 )
 from v03_pipeline.lib.tasks.base.base_loading_run_params import (
     BaseLoadingRunParams,
@@ -54,6 +56,7 @@ def run(self) -> None:
                 self.dataset_type,
                 self.callset_path,
             ),
+            'sample_qc': {},
         }
         for remapped_and_subsetted_callset in self.input():
             callset_mt = hl.read_matrix_table(remapped_and_subsetted_callset.path)
@@ -67,6 +70,20 @@ def run(self) -> None:
                     **collected_globals['failed_family_samples'][key],
                     **metadata_json['failed_family_samples'][key],
                 }
-
+        if (
+            FeatureFlag.EXPECT_TDR_METRICS
+            and not self.skip_expect_tdr_metrics
+            and self.dataset_type.expect_tdr_metrics(
+                self.reference_genome,
+            )
+        ):
+            with open(
+                sample_qc_json_path(
+                    self.reference_genome,
+                    self.dataset_type,
+                    self.callset_path,
+                ),
+            ) as f:
+                metadata_json['sample_qc'] = json.load(f)
         with self.output().open('w') as f:
             json.dump(metadata_json, f)
diff --git a/v03_pipeline/lib/tasks/write_metadata_for_run_test.py b/v03_pipeline/lib/tasks/write_metadata_for_run_test.py
@@ -1,20 +1,38 @@
 import json
+from unittest import mock
+from unittest.mock import Mock
 
 import luigi.worker
 
 from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
 from v03_pipeline.lib.paths import relatedness_check_tsv_path
 from v03_pipeline.lib.tasks.write_metadata_for_run import WriteMetadataForRunTask
+from v03_pipeline.lib.test.mock_complete_task import MockCompleteTask
 from v03_pipeline.lib.test.mocked_dataroot_testcase import MockedDatarootTestCase
 
 TEST_VCF = 'v03_pipeline/var/test/callsets/1kg_30variants.vcf'
 TEST_REMAP_2 = 'v03_pipeline/var/test/remaps/test_remap_2.tsv'
 TEST_PEDIGREE_3 = 'v03_pipeline/var/test/pedigrees/test_pedigree_3.tsv'
 TEST_PEDIGREE_4 = 'v03_pipeline/var/test/pedigrees/test_pedigree_4.tsv'
+TEST_SAMPLE_QC_JSON = 'v03_pipeline/var/test/sample_qc_1.json'
 
 
 class WriteMetadataForRunTaskTest(MockedDatarootTestCase):
-    def test_write_metadata_for_run_task(self) -> None:
+    @mock.patch(
+        'v03_pipeline.lib.tasks.write_metadata_for_run.sample_qc_json_path',
+        lambda *_: TEST_SAMPLE_QC_JSON,
+    )
+    @mock.patch('v03_pipeline.lib.tasks.write_metadata_for_run.FeatureFlag')
+    @mock.patch(
+        'v03_pipeline.lib.tasks.write_imported_callset.WriteTDRMetricsFilesTask',
+    )
+    def test_write_metadata_for_run_task(
+        self,
+        write_tdr_metrics_task: Mock,
+        mock_ff: Mock,
+    ) -> None:
+        mock_ff.EXPECT_TDR_METRICS = True
+        write_tdr_metrics_task.return_value = MockCompleteTask()
         worker = luigi.worker.Worker()
         write_metadata_for_run_task = WriteMetadataForRunTask(
             reference_genome=ReferenceGenome.GRCh38,
@@ -77,5 +95,11 @@ def test_write_metadata_for_run_task(self) -> None:
                         DatasetType.SNV_INDEL,
                         TEST_VCF,
                     ),
+                    'sample_qc': {
+                        'HG00731': {'filter_flags': ['coverage', 'contamination']},
+                        'HG00732': {'filter_flags': ['coverage']},
+                        'HG00733': {'filter_flags': ['contamination']},
+                        'NA19675': {'filter_flags': []},
+                    },
                 },
             )
diff --git a/v03_pipeline/lib/tasks/write_remapped_and_subsetted_callset.py b/v03_pipeline/lib/tasks/write_remapped_and_subsetted_callset.py
@@ -29,6 +29,7 @@
 from v03_pipeline.lib.tasks.write_relatedness_check_tsv import (
     WriteRelatednessCheckTsvTask,
 )
+from v03_pipeline.lib.tasks.write_sample_qc_json import WriteSampleQCJsonTask
 from v03_pipeline.lib.tasks.write_sex_check_table import WriteSexCheckTableTask
 from v03_pipeline.lib.tasks.write_validation_errors_for_run import (
     with_persisted_validation_errors,
@@ -83,6 +84,17 @@ def requires(self) -> list[luigi.Task]:
                 self.clone(WriteRelatednessCheckTsvTask),
                 self.clone(WriteSexCheckTableTask),
             ]
+        if (
+            FeatureFlag.EXPECT_TDR_METRICS
+            and not self.skip_expect_tdr_metrics
+            and self.dataset_type.expect_tdr_metrics(
+                self.reference_genome,
+            )
+        ):
+            requirements = [
+                *requirements,
+                self.clone(WriteSampleQCJsonTask),
+            ]
         return requirements
 
     @with_persisted_validation_errors
diff --git a/v03_pipeline/lib/tasks/write_remapped_and_subsetted_callset_test.py b/v03_pipeline/lib/tasks/write_remapped_and_subsetted_callset_test.py
@@ -94,6 +94,7 @@ def test_write_remapped_and_subsetted_callset_task(
             project_pedigree_paths=[TEST_PEDIGREE_3],
             project_i=0,
             skip_validation=True,
+            skip_expect_tdr_metrics=True,
         )
         worker.add(wrsc_task)
         worker.run()
@@ -138,6 +139,7 @@ def test_write_remapped_and_subsetted_callset_task_failed_sex_check_family(
             project_pedigree_paths=[TEST_PEDIGREE_4],
             project_i=0,
             skip_validation=True,
+            skip_expect_tdr_metrics=True,
         )
         worker.add(wrsc_task)
         worker.run()
@@ -203,6 +205,7 @@ def test_write_remapped_and_subsetted_callset_task_all_families_failed(
             project_pedigree_paths=[TEST_PEDIGREE_7],
             project_i=0,
             skip_validation=True,
+            skip_expect_tdr_metrics=True,
         )
         worker.add(wrsc_task)
         worker.run()
diff --git a/v03_pipeline/lib/tasks/write_sample_qc_json.py b/v03_pipeline/lib/tasks/write_sample_qc_json.py
@@ -0,0 +1,60 @@
+import json
+from collections import defaultdict
+
+import hail as hl
+import hailtop.fs as hfs
+import luigi
+import luigi.util
+
+from v03_pipeline.lib.methods.sample_qc import call_sample_qc
+from v03_pipeline.lib.misc.io import import_tdr_qc_metrics
+from v03_pipeline.lib.paths import sample_qc_json_path, tdr_metrics_dir
+from v03_pipeline.lib.tasks.base.base_loading_run_params import BaseLoadingRunParams
+from v03_pipeline.lib.tasks.files import GCSorLocalTarget
+from v03_pipeline.lib.tasks.validate_callset import ValidateCallsetTask
+from v03_pipeline.lib.tasks.write_tdr_metrics_files import WriteTDRMetricsFilesTask
+
+
+@luigi.util.inherits(BaseLoadingRunParams)
+class WriteSampleQCJsonTask(luigi.Task):
+    def output(self) -> luigi.Target:
+        return GCSorLocalTarget(
+            sample_qc_json_path(
+                self.reference_genome,
+                self.dataset_type,
+                self.callset_path,
+            ),
+        )
+
+    def requires(self):
+        return [self.clone(ValidateCallsetTask), self.clone(WriteTDRMetricsFilesTask)]
+
+    def run(self):
+        callset_mt = hl.read_matrix_table(self.input()[0].path)
+
+        tdr_metrics_ht = None
+        for tdr_metrics_file in hfs.ls(
+            tdr_metrics_dir(self.reference_genome, self.dataset_type),
+        ):
+            if not tdr_metrics_ht:
+                tdr_metrics_ht = import_tdr_qc_metrics(tdr_metrics_file.path)
+                continue
+            tdr_metrics_ht = tdr_metrics_ht.union(
+                import_tdr_qc_metrics(tdr_metrics_file.path),
+            )
+
+        callset_mt = call_sample_qc(
+            callset_mt,
+            tdr_metrics_ht,
+            self.sample_type,
+        )
+        ht = callset_mt.cols()
+        sample_qc_dict = defaultdict(dict)
+        for row in ht.flatten().collect():
+            r = dict(row)
+            sample_id = r.pop('s')
+            for field, value in r.items():
+                sample_qc_dict[sample_id][field] = value
+
+        with self.output().open('w') as f:
+            json.dump(sample_qc_dict, f)
diff --git a/v03_pipeline/lib/tasks/write_sample_qc_json_test.py b/v03_pipeline/lib/tasks/write_sample_qc_json_test.py
diff --git a/v03_pipeline/lib/tasks/write_variant_annotations_vcf_test.py b/v03_pipeline/lib/tasks/write_variant_annotations_vcf_test.py
diff --git a/v03_pipeline/var/test/sample_qc_1.json b/v03_pipeline/var/test/sample_qc_1.json
diff --git a/v03_pipeline/var/test/tdr_metrics.tsv b/v03_pipeline/var/test/tdr_metrics.tsv
