Merge pull request #1077 from Tanisha127/feature/775-heterozygosity-support

feat: add public sample_count_het() and cohort_heterozygosity() methods (#775)

Author: jonbrenas

malariagen_data/anoph/heterozygosity.py

@@ -708,3 +708,120 @@ def plot_roh(
             return None
         else:
             return fig_all
+
+    @_check_types
+    @doc(
+        summary="Return windowed heterozygosity for a single sample over a genome region.",
+        returns="""
+            A DataFrame with one row per window. Columns are:
+            `sample_id` is the identifier of the sample,
+            `window_start` is the start position of the window,
+            `window_stop` is the stop position of the window,
+            `heterozygosity` is the heterozygosity in the window.
+        """,
+    )
+    def sample_count_het(
+        self,
+        sample: base_params.sample,
+        region: base_params.region,
+        window_size: het_params.window_size = het_params.window_size_default,
+        site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
+        sample_set: Optional[base_params.sample_set] = None,
+        chunks: base_params.chunks = base_params.native_chunks,
+        inline_array: base_params.inline_array = base_params.inline_array_default,
+    ) -> pd.DataFrame:
+        # Normalise region parameter.
+        region_prepped: Region = _parse_single_region(self, region)
+        del region
+
+        # Compute windowed heterozygosity using private method.
+        sample_id, sample_set, windows, counts = self._sample_count_het(
+            sample=sample,
+            region=region_prepped,
+            site_mask=site_mask,
+            window_size=window_size,
+            sample_set=sample_set,
+            chunks=chunks,
+            inline_array=inline_array,
+        )
+
+        # Build and return a DataFrame.
+        df = pd.DataFrame(
+            {
+                "sample_id": sample_id,
+                "window_start": windows[:, 0],
+                "window_stop": windows[:, 1],
+                "heterozygosity": counts / window_size,
+            }
+        )
+
+        return df
+
+    @_check_types
+    @doc(
+        summary="Compute mean heterozygosity for multiple cohorts over a genome region.",
+        returns="""
+            A DataFrame with one row per cohort. Columns are:
+            `cohort` is the cohort identifier,
+            `n_samples` is the number of samples in the cohort,
+            `mean_heterozygosity` is the mean heterozygosity across all samples and windows in the cohort.
+        """,
+    )
+    def cohort_heterozygosity(
+        self,
+        region: base_params.region,
+        cohorts: base_params.cohorts,
+        sample_sets: Optional[base_params.sample_sets] = None,
+        sample_query: Optional[base_params.sample_query] = None,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
+        window_size: het_params.window_size = het_params.window_size_default,
+        site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
+        chunks: base_params.chunks = base_params.native_chunks,
+        inline_array: base_params.inline_array = base_params.inline_array_default,
+    ) -> pd.DataFrame:
+        # Normalise region parameter.
+        region_prepped: Region = _parse_single_region(self, region)
+        del region
+
+        # Set up cohort queries.
+        cohort_queries = self._setup_cohort_queries(
+            cohorts=cohorts,
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_query_options=sample_query_options,
+            cohort_size=None,
+            min_cohort_size=1,
+        )
+
+        # Compute heterozygosity for each cohort.
+        results = []
+        for cohort_label, cohort_query in cohort_queries.items():
+            # Get samples in this cohort.
+            df_cohort_samples = self.sample_metadata(
+                sample_sets=sample_sets,
+                sample_query=cohort_query,
+            )
+            n_samples = len(df_cohort_samples)
+
+            # Compute heterozygosity for each sample and take the mean.
+            het_values = []
+            for sample_id in df_cohort_samples["sample_id"]:
+                df_het = self.sample_count_het(
+                    sample=sample_id,
+                    region=region_prepped,
+                    window_size=window_size,
+                    site_mask=site_mask,
+                    chunks=chunks,
+                    inline_array=inline_array,
+                )
+                het_values.append(df_het["heterozygosity"].mean())
+
+            results.append(
+                {
+                    "cohort": cohort_label,
+                    "n_samples": n_samples,
+                    "mean_heterozygosity": np.mean(het_values),
+                }
+            )
+
+        return pd.DataFrame(results)

tests/anoph/test_heterozygosity.py

@@ -205,3 +205,58 @@ def test_plot_roh(fixture, api: AnophelesHetAnalysis):
 
     # Check results.
     assert isinstance(fig, bokeh.models.GridPlot)
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_sample_count_het(fixture, api: AnophelesHetAnalysis):
+    # Set up test parameters.
+    all_sample_sets = api.sample_sets()["sample_set"].to_list()
+    sample_set = random.choice(all_sample_sets)
+    df_samples = api.sample_metadata(sample_sets=sample_set)
+    sample = random.choice(df_samples["sample_id"].to_list())
+
+    het_params = dict(
+        sample=sample,
+        region=random.choice(api.contigs),
+        sample_set=sample_set,
+        window_size=20_000,
+    )
+
+    # Run function under test.
+    df = api.sample_count_het(**het_params)
+
+    # Check results.
+    assert isinstance(df, pd.DataFrame)
+    expected_columns = ["sample_id", "window_start", "window_stop", "heterozygosity"]
+    for col in expected_columns:
+        assert col in df.columns
+    assert len(df) > 0
+    assert (df["heterozygosity"] >= 0).all()
+    assert (df["heterozygosity"] <= 1).all()
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_cohort_heterozygosity(fixture, api: AnophelesHetAnalysis):
+    # Set up test parameters.
+    all_sample_sets = api.sample_sets()["sample_set"].to_list()
+    sample_set = random.choice(all_sample_sets)
+
+    cohort_params = dict(
+        region=random.choice(api.contigs),
+        cohorts="taxon",
+        sample_sets=sample_set,
+        window_size=20_000,
+    )
+
+    # Run function under test.
+    df = api.cohort_heterozygosity(**cohort_params)
+
+    # Check results.
+    assert isinstance(df, pd.DataFrame)
+    expected_columns = ["cohort", "n_samples", "mean_heterozygosity"]
+    for col in expected_columns:
+        assert col in df.columns
+    assert len(df) > 0
+    assert (df["n_samples"] > 0).all()
+    assert (df["mean_heterozygosity"] >= 0).all()
+    assert (df["mean_heterozygosity"] <= 1).all()