sorAI locus CNNs

These are the 384 locus-specific convolutional networks used by sorAI to predict whether a fern sample is diploid or polyploid from GoFlag 408 target-capture reads.

Each file checkpoints/best_model_probe_3_<index>.pt is the network for one locus. probe_index_to_id.csv maps model_index to the GoFlag locus id. The random forest that combines these locus probabilities is in the GitHub repository, not in this model repo.

On the held-out fern test set (34 samples) the random forest that reads these probabilities has accuracy 0.765 and AUROC 0.870. The models were trained on ferns and do not estimate a chromosome number.

Download

hf download Fahimehr140/sorai --local-dir sorai-cnns

Then, with tokenized locus alignments named L<probe_id>.json:

sorai predict --loci-dir sample_loci/ --cnn-dir sorai-cnns/checkpoints

Code, citation, and the random forest: https://github.com/Fahimehgit/sorai