Merge branch 'master' into GH-1054-add-vcf-export

Author: adilraza99

malariagen_data/anoph/fst.py

@@ -360,7 +360,7 @@ def plot_fst_gwss(
     )
     def average_fst(
         self,
-        region: base_params.region,
+        region: base_params.regions,
         cohort1_query: base_params.sample_query,
         cohort2_query: base_params.sample_query,
         sample_query_options: Optional[base_params.sample_query_options] = None,
@@ -435,7 +435,7 @@ def average_fst(
     )
     def pairwise_average_fst(
         self,
-        region: base_params.region,
+        region: base_params.regions,
         cohorts: base_params.cohorts,
         sample_sets: Optional[base_params.sample_sets] = None,
         sample_query: Optional[base_params.sample_query] = None,

tests/anoph/conftest.py

@@ -334,7 +334,7 @@ def simulate_exons(
         # keep things simple for now.
         if strand == "-":
             # Take exons in reverse order.
-            exons == exons[::-1]
+            exons = exons[::-1]
         for exon_ix, exon in enumerate(exons):
             first_exon = exon_ix == 0
             last_exon = exon_ix == len(exons) - 1
@@ -646,8 +646,8 @@ def simulate_cnv_hmm(zarr_path, metadata_path, contigs, contig_sizes, rng):
     # - sample_is_high_variance [1D array] [bool] [True or False for n_samples]
     # - samples [1D array] [str]
 
-    # Get a random probability for a sample being high variance, between 0 and 1.
-    p_variance = rng.random()
+    # Keep high variance sample prevalence stable for deterministic tests.
+    p_variance = 0.1
 
     # Open a zarr at the specified path.
     root = zarr.open(zarr_path, mode="w")
@@ -862,8 +862,8 @@ def simulate_cnv_discordant_read_calls(
     # - sample_is_high_variance [1D array] [bool] [True or False for n_samples]
     # - samples [1D array] [str for n_samples]
 
-    # Get a random probability for a sample being high variance, between 0 and 1.
-    p_variance = rng.random()
+    # Keep high variance sample prevalence stable for deterministic tests.
+    p_variance = 0.1
 
     # Get a random probability for choosing allele 1, between 0 and 1.
     p_allele = rng.random()

tests/anoph/test_fst.py

@@ -368,3 +368,50 @@ def test_pairwise_average_fst_with_bad_cohorts(fixture, api: AnophelesFstAnalysi
     # Run function under test.
     with pytest.raises(ValueError):
         api.pairwise_average_fst(**fst_params)
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_average_fst_with_list_of_regions(fixture, api: AnophelesFstAnalysis):
+    # Set up test parameters.
+    all_sample_sets = api.sample_sets()["sample_set"].to_list()
+    all_countries = api.sample_metadata()["country"].dropna().unique().tolist()
+    countries = random.sample(all_countries, 2)
+    cohort1_query = f"country == {countries[0]!r}"
+    cohort2_query = f"country == {countries[1]!r}"
+    fst_params = dict(
+        region=random.sample(api.contigs, 2),
+        sample_sets=all_sample_sets,
+        cohort1_query=cohort1_query,
+        cohort2_query=cohort2_query,
+        site_mask=random.choice(api.site_mask_ids),
+        min_cohort_size=1,
+        n_jack=random.randint(10, 200),
+    )
+
+    # Run function under test.
+    fst, se = api.average_fst(**fst_params)
+
+    # Checks.
+    assert isinstance(fst, float)
+    assert isinstance(se, float)
+    assert 0 <= fst <= 1
+    assert 0 <= se <= 1
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_pairwise_average_fst_with_list_of_regions(fixture, api: AnophelesFstAnalysis):
+    # Set up test parameters.
+    all_sample_sets = api.sample_sets()["sample_set"].to_list()
+    region = random.sample(api.contigs, 2)
+    site_mask = random.choice(api.site_mask_ids)
+    fst_params = dict(
+        region=region,
+        cohorts="country",
+        sample_sets=all_sample_sets,
+        site_mask=site_mask,
+        min_cohort_size=1,
+        n_jack=random.randint(10, 200),
+    )
+
+    # Run checks.
+    check_pairwise_average_fst(api=api, fst_params=fst_params)

tests/anoph/test_simulated_cnv_data.py

@@ -0,0 +1,81 @@
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import zarr
+
+from .conftest import (
+    Gff3Simulator,
+    simulate_cnv_discordant_read_calls,
+    simulate_cnv_hmm,
+)
+
+
+def _write_sample_metadata(path: Path, n_samples: int = 100) -> None:
+    df_samples = pd.DataFrame({"sample_id": [f"S{i:04d}" for i in range(n_samples)]})
+    df_samples.to_csv(path, index=False)
+
+
+def test_simulate_cnv_hmm_limits_high_variance_fraction(tmp_path):
+    zarr_path = tmp_path / "cnv_hmm.zarr"
+    metadata_path = tmp_path / "samples.csv"
+    _write_sample_metadata(metadata_path)
+
+    simulate_cnv_hmm(
+        zarr_path=zarr_path,
+        metadata_path=metadata_path,
+        contigs=("2L",),
+        contig_sizes={"2L": 10_000},
+        rng=np.random.default_rng(0),
+    )
+
+    root = zarr.open(zarr_path, mode="r")
+    high_variance_fraction = np.mean(root["sample_is_high_variance"][:])
+    assert high_variance_fraction < 0.3
+
+
+def test_simulate_cnv_discordant_read_calls_limits_high_variance_fraction(tmp_path):
+    zarr_path = tmp_path / "cnv_discordant.zarr"
+    metadata_path = tmp_path / "samples.csv"
+    _write_sample_metadata(metadata_path)
+
+    simulate_cnv_discordant_read_calls(
+        zarr_path=zarr_path,
+        metadata_path=metadata_path,
+        contigs=("2L",),
+        contig_sizes={"2L": 10_000},
+        rng=np.random.default_rng(0),
+    )
+
+    root = zarr.open(zarr_path, mode="r")
+    high_variance_fraction = np.mean(root["sample_is_high_variance"][:])
+    assert high_variance_fraction < 0.3
+
+
+def test_simulate_exons_on_minus_strand_reverses_feature_order():
+    sim = Gff3Simulator(
+        contig_sizes={"2L": 10_000},
+        rng=np.random.default_rng(0),
+        n_exons_low=3,
+        n_exons_high=3,
+        intron_size_low=10,
+        intron_size_high=10,
+        exon_size_low=100,
+        exon_size_high=100,
+    )
+    rows = list(
+        sim.simulate_exons(
+            contig="2L",
+            strand="-",
+            gene_ix=0,
+            transcript_ix=0,
+            transcript_id="transcript-2L-0-0",
+            transcript_start=1,
+            transcript_end=1_000,
+        )
+    )
+    cds_and_utrs = [
+        row for row in rows if row[2] in {sim.utr5_type, sim.utr3_type, sim.cds_type}
+    ]
+    starts = [row[3] for row in cds_and_utrs]
+    assert starts == sorted(starts, reverse=True)