Exploring learned SNP representations for improved classification of ambiguous taxa

[AswaniSahoo]

Hi @jonbrenas @ahernank,

I've been contributing to malariagen-data-python (PRs #895 and #969, both merged) and have been studying the classifier codebase and the Ag3/Af1 data it operates on. A few observations and questions:

1. Feature representation
The current Random Forest baseline uses raw genotype arrays as input. I've been thinking about whether a learned representation layer for example, a lightweight embedding that captures local linkage disequilibrium structure between adjacent SNP positions could improve generalization on under-represented taxa. In my own work building a Vision Transformer for weather forecasting on ERA5 data, positional encodings that respected spatial relationships significantly improved prediction quality. The analogy to SNP positions along a chromosome seems potentially valuable here.

2. Handling taxonomic ambiguity
For samples near species boundaries (e.g., gambiae/coluzzii hybrids, or arabiensis in sympatric zones), hard classification seems insufficient. Would calibrated probability outputs (e.g., via Platt scaling or temperature-adjusted softmax) be preferred, so researchers can flag uncertain samples for follow-up?

3. Integration with malariagen-data-python
Since the classifier needs to consume data from the API (snp_calls(), sample_metadata()), would it be useful to have the model packaged as a method within the existing API something like ag3.predict_taxon(sample_sets=..., region=...)?

I'd be interested in taking on a piece of thispotentially starting with a comparison notebook benchmarking the RF baseline against a small neural approach on the existing training data. Would that be useful?

[AswaniSahoo]

Following up, I went ahead and put together the comparison notebook I mentioned:

Notebook: taxon_classifier_exploration.ipynb

Using the 2,784-sample dataset from gsoc-mosquito-taxon-classifier (simulated genotype features matching the real Ag3 GT structure — diploid calls encoded as 0/1/2, 500 SNPs), here's what I found:

Metric	Random Forest	MLP (PyTorch)
Overall Accuracy	94.25%	99.10%
Macro F1	0.773	0.973
gcx3 F1 (n=13 test)	0.000	0.880
coluzzii→gambiae leakage	9.9%	0.0%

Two critical RF failure modes emerged:

1. Rare taxon collapse — gcx3 (n=65 total, 2.3% of data) was completely absorbed into gambiae (0% recall). The MLP with balanced class weights recovered 84.6% recall.

2. Sister species confusion — 14 coluzzii samples (9.9%) leaked into gambiae, consistent with only ~2% of SNPs being diagnostic between them. The MLP learned to separate these using the weak-signal SNPs that RF discards.

I also ran a confidence calibration analysis — the RF clusters predictions at 0.5–0.8 confidence with little separation between correct and incorrect. The MLP shows >0.99 on correct predictions vs. 0.85 on errors, which could support threshold-based flagging for ambiguous samples (relevant to the Platt scaling question above).

The RF feature importance analysis recovered 7/10 known diagnostic SNPs in its top 20 features, validating that the model finds biologically meaningful signal.

Next step would be re-running on real Ag3 genotype data once I have data access — the architectural findings should transfer. I'd also like to explore the SNP embedding idea (point 1 above) to see if learned positional representations capture LD structure better than raw genotype arrays.

Full results with confusion matrices, training curves, and confidence distributions are in the notebook.

[AswaniSahoo]

Update: I got access to the real Ag3 API and re-ran the full pipeline on actual genotype data. Results below supersede my earlier simulated-data findings.

Notebook: taxon_classifier_exploration.ipynb

Using real SNP genotype data from malariagen_data Ag3 API (sample_sets="3.0", region 3L:1-100000) — 2,781 samples, 60,467 SNPs, diploid calls encoded as 0/1/2 , here's what I found:

Metric	Random Forest	MLP (PyTorch)
Overall Accuracy	95.5%	99.8%
Macro F1	~0.93	~0.99
pwani F1 (n≈13 test)	partial failure	near-perfect
bissau→gambiae leakage	23.5%	0.0%
Training Time	~2s	~30s

The real data confirmed and sharpened the patterns from simulated data, with one major surprise:

1. The hardest pair isn't gambiae/coluzzii - Under RF, the dominant failure mode was bissau → gambiae leakage (23.5% misclassification), not the sister species pair. The MLP with balanced class weights resolved this completely.

2. Rare taxon recovery - pwani (2.3% of data, formerly labeled gcx3 in simulated runs) was partially lost by RF. The MLP recovered near-perfect recall using balanced CrossEntropyLoss weights.

3. Sister species separation - coluzzii/gambiae, which share >98% of their genome, are cleanly separated by the MLP (F1 ~1.00 vs RF ~0.96), confirming the network learns weak-signal SNP interactions that RF discards.

Confidence calibration remains strong on real data - the MLP shows high confidence (>0.99) on correct predictions and noticeably lower confidence on errors, supporting threshold-based flagging for ambiguous samples (relevant to the Platt scaling / calibrated probability outputs question above).

Feature importance now points to actual genomic positions on chromosome 3L. The RF identified biologically meaningful SNP coordinates as its top features, validating that the classification signal is genuine and not a simulation artifact.

Class distribution (real Ag3 data):

Taxon	Samples	% of total
gambiae	~1,476	53.0%
coluzzii	~708	25.4%
arabiensis	~368	13.2%
bissau	~170	6.1%
pwani	~64	2.3%

What changed from simulated → real data:

• Feature space grew from 500 simulated SNPs to 60,467 real SNPs
• Class imbalance patterns held - gambiae still dominates at 53%
• RF accuracy improved slightly (94.25% → 95.5%), MLP improved (99.10% → 99.8%)
• The dominant RF failure mode shifted from rare-taxon collapse (gcx3) to bissau→gambiae leakage - a pattern the simulation didn't fully capture

Next steps:

• Expand beyond the 3L:1-100000 region to genome-wide SNPs
• Explore the FASTQ-based approach (k-mer features instead of genotype arrays) for the GSoC project scope
• SNP embeddings to capture linkage disequilibrium structure
• Temperature scaling for formal probability calibration

Full results with confusion matrices, training curves, per-class comparisons, and confidence distributions are in the notebook. The notebook supports both real Ag3 data (with Google Cloud auth) and simulated genotype fallback (no auth needed) , both pathways feed into the same ML pipeline.