Improve plot_njt errors for insufficient data

Signed-off-by: Aryan-SINGH-GIT <aryansingh12oct2005@gmail.com>

Author: Aryan-SINGH-GIT

malariagen_data/anoph/distance.py

@@ -482,6 +482,20 @@ def plot_njt(
             count_sort = True
             distance_sort = False
 
+        # Ensure we have enough samples for a tree.
+        # If we have 0 samples, `biallelic_snp_calls` or `snp_calls` should have already raised "No samples found".
+        # However, if we have 1 sample, it might pass through until here, where it would cause a failure in njt.
+        df_samples = self.sample_metadata(
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_query_options=sample_query_options,
+            sample_indices=sample_indices,
+        )
+        if 0 < len(df_samples) < 2:
+            raise ValueError(
+                f"Not enough samples for neighbour-joining tree. Found {len(df_samples)}, needed at least 2."
+            )
+
         # Compute neighbour-joining tree.
         Z, samples, n_snps_used = self.njt(
             region=region,

malariagen_data/anoph/snp_data.py

@@ -1886,7 +1886,9 @@ 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. Found {ds_out.sizes['variants']}, needed {n_snps}."
+                    )
 
         return ds_out
 

tests/anoph/test_distance_errors.py

@@ -0,0 +1,67 @@
+import pytest
+from malariagen_data import ag3 as _ag3
+from malariagen_data.anoph.distance import AnophelesDistanceAnalysis
+
+
+@pytest.fixture
+def ag3_sim_api(ag3_sim_fixture):
+    return AnophelesDistanceAnalysis(
+        url=ag3_sim_fixture.url,
+        public_url=ag3_sim_fixture.url,
+        config_path=_ag3.CONFIG_PATH,
+        major_version_number=_ag3.MAJOR_VERSION_NUMBER,
+        major_version_path=_ag3.MAJOR_VERSION_PATH,
+        pre=True,
+        aim_metadata_dtype={
+            "aim_species_fraction_arab": "float64",
+            "aim_species_fraction_colu": "float64",
+            "aim_species_fraction_colu_no2l": "float64",
+            "aim_species_gambcolu_arabiensis": object,
+            "aim_species_gambiae_coluzzii": object,
+            "aim_species": object,
+        },
+        gff_gene_type="gene",
+        gff_gene_name_attribute="Name",
+        gff_default_attributes=("ID", "Parent", "Name", "description"),
+        default_site_mask="gamb_colu_arab",
+        results_cache=ag3_sim_fixture.results_cache_path.as_posix(),
+        taxon_colors=_ag3.TAXON_COLORS,
+        virtual_contigs=_ag3.VIRTUAL_CONTIGS,
+    )
+
+
+def test_plot_njt_no_samples(ag3_sim_api):
+    # Test with a query matching no samples.
+    with pytest.raises(ValueError) as e:
+        ag3_sim_api.plot_njt(
+            region="2L", n_snps=10, sample_query="sex_call == 'Impossible_Value'"
+        )
+    assert "No samples found for query" in str(
+        e.value
+    ) or "No relevant samples found" in str(e.value)
+
+
+def test_plot_njt_not_enough_snps(ag3_sim_api):
+    # Request more SNPs than available in the region
+    with pytest.raises(ValueError) as e:
+        ag3_sim_api.plot_njt(region="2L", n_snps=10000000, sample_query=None)
+    assert "Not enough SNPs." in str(e.value)
+    assert "Found" in str(e.value)
+    assert "needed 10000000" in str(e.value)
+
+
+def test_plot_njt_one_sample(ag3_sim_api):
+    # Test with a query that returns only 1 sample.
+    # This should trigger the minimum sample check in plot_njt.
+
+    # First, find a sample so we can query for just one
+    df_samples = ag3_sim_api.sample_metadata()
+    sample_id = df_samples.iloc[0]["sample_id"]
+
+    with pytest.raises(ValueError) as e:
+        ag3_sim_api.plot_njt(
+            region="2L", n_snps=10, sample_query=f"sample_id == '{sample_id}'"
+        )
+    assert "Not enough samples for neighbour-joining tree" in str(e.value)
+    assert "Found 1" in str(e.value)
+    assert "needed at least 2" in str(e.value)