sorAI - a suite of machine learning models to predict plant polyploidy from target capture data

Premise: Polyploidy is a persistent feature of land plant evolution, yet it is often difficult to determine the ploidy of individual plants. The growing abundance of target capture data and other high-throughput sequence datasets across plant species has generated interest in methods to infer ploidy from these data. We leverage chromosome count data and target capture data from the GoFlag probe set in ferns to train new machine learning models to infer if a sample is from a diploid or polyploid genome. Methods and Results: We trained a convolutional neural network (CNN) for each GoFlag probe locus and developed ensemble models that use two loci and a random forest model that leverages all the loci to classify fern samples as diploid or polyploid. We tested the performance of these models using held-out test datasets of GoFlag target capture data across ferns and an independent test set from the genus Lygodium. We also performed experiments to evaluate the features of loci and sequences most associated with model performance and tested the performance of the models as sequence depth decreased. Conclusions: These models, collectively named sorAI, outperform the frequency-based methods nQuire and nQuack on held-out test datasets of GoFlag target capture data across ferns. These models are capable of leveraging positional information and covariation among nucleotide sites beyond individual-site allele frequencies, and they perform well even with low sequence coverage. The random forest model, which uses all loci, appears to be robust to varying dataset features and achieves an accuracy of approximately 71-78% across test datasets and sequencing depths. sorAI is available at https://github.com/Fahimehgit/sorai, and the trained models are available at https://huggingface.co/Fahimehr140/sorai.