Merge pull request malariagen#921 from bleedblack1/fix/778-njt-descriptive-errors

Fix malariagen#778: descriptive errors in njt/plot_njt for insufficient data

Author: jonbrenas

malariagen_data/anoph/distance.py

@@ -365,24 +365,43 @@ def _njt(
         from scipy.spatial.distance import squareform  # type: ignore
 
         # Compute pairwise distances.
-        dist, samples, n_snps = self.biallelic_diplotype_pairwise_distances(
-            region=region,
-            n_snps=n_snps,
-            metric=metric,
-            sample_sets=sample_sets,
-            sample_indices=sample_indices,
-            site_mask=site_mask,
-            site_class=site_class,
-            inline_array=inline_array,
-            chunks=chunks,
-            cohort_size=cohort_size,
-            min_cohort_size=min_cohort_size,
-            max_cohort_size=max_cohort_size,
-            random_seed=random_seed,
-            max_missing_an=max_missing_an,
-            min_minor_ac=min_minor_ac,
-            thin_offset=thin_offset,
-        )
+        try:
+            dist, samples, n_snps_used = self.biallelic_diplotype_pairwise_distances(
+                region=region,
+                n_snps=n_snps,
+                metric=metric,
+                sample_sets=sample_sets,
+                sample_indices=sample_indices,
+                site_mask=site_mask,
+                site_class=site_class,
+                inline_array=inline_array,
+                chunks=chunks,
+                cohort_size=cohort_size,
+                min_cohort_size=min_cohort_size,
+                max_cohort_size=max_cohort_size,
+                random_seed=random_seed,
+                max_missing_an=max_missing_an,
+                min_minor_ac=min_minor_ac,
+                thin_offset=thin_offset,
+            )
+
+        except ValueError as e:
+            raise ValueError(
+                f"Unable to construct neighbour-joining tree. {e} "
+                f"This could be because the selected region does not "
+                f"contain enough polymorphic SNPs for the given sample "
+                f"sets and query parameters."
+            ) from e
+
+        # Validate enough samples for a tree.
+        n_samples = len(samples)
+        if n_samples < 3:
+            raise ValueError(
+                f"Not enough samples to construct a neighbour-joining tree. "
+                f"A minimum of 3 samples is required, but only {n_samples} "
+                f"were found for the given region and sample sets."
+            )
+
         D = squareform(dist)
 
         # anjl supports passing in a progress bar function to get progress on the

malariagen_data/anoph/snp_data.py

@@ -1903,7 +1903,12 @@ def biallelic_snp_calls(
                     ds_out = ds_out.isel(variants=loc_thin)
 
                 elif ds_out.sizes["variants"] < n_snps:
-                    raise ValueError("Not enough SNPs.")
+                    raise ValueError(
+                        f"Not enough SNPs. Requested {n_snps} SNPs but only "
+                        f"{ds_out.sizes['variants']} were found in the selected "
+                        f"region after applying filters. Try using a larger region, "
+                        f"relaxing site filters, or reducing the n_snps parameter."
+                    )
 
         return ds_out
 

tests/anoph/test_distance.py

@@ -269,3 +269,20 @@ def test_plot_njt(fixture, api: AnophelesDistanceAnalysis):
             **data_params,
         )
         assert isinstance(fig, go.Figure)
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_njt_not_enough_snps(fixture, api: AnophelesDistanceAnalysis):
+    all_sample_sets = api.sample_sets()["sample_set"].to_list()
+    with pytest.raises(
+        ValueError,
+        match="Unable to construct neighbour-joining tree|Not enough SNPs",
+    ):
+        api.njt(
+            region=random.choice(api.contigs),
+            n_snps=1_000_000_000,  # impossibly high to guarantee failure
+            sample_sets=random.sample(all_sample_sets, 1),
+            site_mask=random.choice((None,) + api.site_mask_ids),
+            min_minor_ac=pca_params.min_minor_ac_default,
+            max_missing_an=pca_params.max_missing_an_default,
+        )