Possibility for Numba integration?

[jleem-iso]

Hi, would you guys be able to integrate Numba into the compute_aa_CDR3_pgen capabilities in this package? This may speed up things by a magnitude, which makes all the difference when you have datasets with more than 1e4 sequences (e.g. from OAS).

[kmayerb]

Hi ideasbyjin,

This is an aside to your main request: but we (an outside group) use OLGA for pgen estimation as part of tcrdist3. You'll see that we use parmap to spread out the computation over multiple cores. This has allowed us to complete 1e4-1e5 pgen estimates in a reasonable amount of time, until such time that a numba or other accelerated version of OLGA exists.

Best,
K

[jleem-iso]

Hi Koshlan,

Thanks – as an alternative I've mixed up pandarallel in a similar vein. That still doesn't get around the fact that the base function itself is slow – you can imagine that even if you had 96 cores at your disposal, you're still looking at almost half a day or so for large datasets / repeating the analyses. Numba integration is pretty simple (it only really takes a @numba.njit decorator) on a function, but the trouble is that there's lots of bits and pieces in the OLGA code base that use numpy and vanilla Python calls on objects that Numba won't recognise. I might have a crack at a PR when time allows.

Thanks,
Jin

[giulioisac]

Hello,

Thanks for the comment. The parallel implementation of pgen is implemented in our SONIA package which wraps up OLGA. The SONIA package has additional functions but if you are only interested in pgen evaluation you can look at the 'sonia-evaluate' command with the '--pgen' option.

As Jin correctly pointed out, numba integration would require considerable modifications of multiple parts of the code and I do not see it happening soon from our side. However, we would be happy to integrate any modification from the community!

Best,
Giulio

[dweb0]

Hi. I know this is an old thread, but as someone who has found this tool really useful and was also struggling to handle big data
(especially IGH), I decided to have a go at the numba implementation and thought I'd share here. Here's a link to the fork: dweb0/OLGA.

There is a README with more details, but in short, if using the command line tool olga-compute_pgen, all you have to do is add the flag --fast_pgen, or for SONIA/soNNia users, add one line of code to wrap the GenerationProbabilityV(D)J object.

I tested all human models with datasets of 500k sequences, and can confirm the results match the original to within 9 ULPs
and 99.9% <= 4 ULPs (more info on that in the README). In these test datasets, the speedup was as follows:

IGH: 100x faster
TRB: 75x faster
TRA, IGK, IGL: 15-20x faster

Most of that is just due to numba, but I also tried to make some micro optimizations like avoiding temp allocations,
caching variables, etc. There are still a few functions left untranslated, so we could probably squeeze out a bit more performance, but the slowest ones should all be taken care of (compute_Pi_R was the slowest by far).

If any of the devs find this useful, I'm happy to have it merged with the main branch, otherwise, people can just directly install the fork with pip install git+https://github.com/dweb0/OLGA.

Best,
Daniel

[giulioisac]

Hi @dweb0, that looks very promising. I will have a look at it, and if all looks fine, we can merge it into the main branch. Thank you for the wonderful work!

[dweb0]

@giulioisac sorry for the late reply. I'll keep an eye out from now on if you have any questions. But that sounds great, thanks! Just a couple notes:

• The rational behind the current implementation as a FastPgen class was mainly to develop in a separate module and provide a wrapper, without modifying any of the original code, and to simply overwrite compute_CDR3_pgen with a monkey patch. But this is may not the best architecture. It's also not pickable in its state. Full issue here.

• The coding style of kernels.py is C-ish with a focus on performance, with short and generic names when possible, so that we could reuse a lot of the helper functions, given that there were a lot of similar patterns (e.g. zero_nt_bias_insVD/DJ/VJ -> znb). As a result, some of the variable names aren't as precise as the original. And since a lot of the biological context is lost on me (was really just thinking in terms of data structures), I'm sure some variables could be named better (e.g. _fill_from_left).

[giulioisac]

Hello @dweb0. I have integrated your changes, made your class pickable, and now fast pgen works with SoNNia multiprocessing as well. I have pushed to PyPI new versions of the Olga and SoNNia packages. Also, Olga now uses FastPgen by default in the CLI. Thank you so much for the great work!

[dweb0]

@giulioisac Wow very nice, can tell you've been busy. Smart way to handling the pickling / multiprocessing. Also appreciate you adding the additional unit tests. And np, happy to contribute!