Package: torchmers
torchmersEfficient k-mer counting implemented with PyTorch primitives and out-of-the-box GPU support.Features:⚡Blazingly fast.Count 9-mers over the entire human genome in just about 30 seconds on a desktop GPU.⚡Flexible.User friendly Python API. CLI soon to come.⚡GPU acceleration.Supports Nvidia GPUs and most other PyTorch hardware backends.⚡Deep Learning ready.Easy integration into PyTorch-based deep learning models.⚡Format agnostic.Works on any kind of sequence. Built-in tokenizers include 4 letter DNA or 20 letter amino acid codes. Natural language can be used as well.MotivationK-mer frequency analysis is a simple yet highly important tool for genomics. Practitioners typically have to choose between two options: (1) implementations in scripting languages such as Python that are slow but flexible, or (2) tools in compiled languages that are efficient but might have limiting APIs.Why can't we get the best of both worlds? While PyTorch and similar frameworks have been developed specifically for deep learning applications, but after all they offer highly efficient implementations of numerical operations optimized for vector operations on CPUs and numerous hardware accelerators. With a bit of hacking, these primitive operations can be utilized for compute-intensive problems in genomics such as k-mer counting, which are typically not treated as numerical problems.How does it work?torchmersis based on two primitive PyTorch operations:unfoldandscatter_add.unfoldextracts sliding local blocks from a tensor, similar to as found in a conv layer. Each value within a block is offset depending on its position and all values within a block are summed up to produce a distinct value for each possible k-mer.scatter_addis then simply used to add the k-mer indices to a count tensor. Some additional hacking is required for handling padding or varying sequence length, but this is a topic for another time.UsageInstallationpipinstalltorchmersFunctional APIThe main k-mer counting logic of the package is implemented by thecount_k_mersfunction. It expects sequences to be already tokenized, or rathernumericalizedand to be provided as either NumPy or PyTorch arrays. To convert a biological sequence into this format, aTokenizercan be used.fromtorchmersimportTokenizer,count_k_merssequence='CGCTATAAAAGGGC'tokenizer=Tokenizer.from_name('DNA')tokens=tokenizer.encode(sequence)counts=count_k_mers(tokens,k=5)GPU-supporttorchmersuses PyTorch primitives under the hood. As long as PyTorch implements theunfoldandscatter_addoperations, among others, for some hardware backend, it can in principle be used for accelerated k-mer counting. This includes, of course, Nvidia GPUs, TPUs and possibly also Apple M CPUs, although I don't have hardware to test this in the latter case.Thecount_k_merswill simply use the device of the sequence tensor it has been provided, but this can also be overwritten with thedeviceargument:# Both options are equivalenttokens=torch.tensor(tokenizer.encode(sequence),device='cuda')counts=count_k_mers(tokens,k=5,device='cuda')PyTorch Module APItorchmersprovides a PyTorchModulewhich can be integrated into neural networks to encode (DNA) sequences into k-mer spectra.importtorchimporttorch.nnasnnfromtorchmers.modulesimportKMerFrequencyEncodersequences=torch.randint(0,3,(32,256))encoder=KMerFrequencyEncoder(k=5),encoder(sequences)# [32, 4 ** 5]Note thatKMerFrequencyEncoderexpects input sequences to be numericalized and formatted as PyTorch tensors.PaddingIf sequences have been padded, a tensor with sequences lengths of shape[batch_size]can be passed asseq_lensargument toKMerFrequencyEncoder.forward(). Alternatively, a boolean mask of shape[batch_size, seq_len]is supported  as well, where tokens marked withFalsewill not be considered for k-mer counting.ExamplesTheexamples/folder contains Jupyter notebooks with different usage examples.BenchmarksNote:The "python native" implementation is certainly not highly optimized but should represent a simple Python baseline. The GPU benchmarks have been created on a RTX 4090.ContributingContributions are welcomed. Please run tests withpython -m pytest tests/to test your work before submitting.Future workThis project started out of frustration with slow on-the-fly computation of k-mer frequency spectra for deep learning applications when done in pure Python. I don't have very many other personal use cases, but I would be glad if users could provide me with their feature requests.