Jax: Fast Combinations Calculation

Summary

Combinadics A fast combinations calculation in jax. Idea of combinadic implementation is from https://jamesmccaffrey.wordpress.com/2022/06/28/generating-the-mth-lexicographical-element-of-a-combination-using-the-combinadic and some useful information can be found here: https://en.wikipedia.org/wiki/Combinatorial_number_system. Below I copied and aggregated some of the details. Introduction The following code demostrates the combinations calculation in numpy and via combinadics: # setup n = 4 k = 3 totalcount = math . comb ( n , k ) # numpy print ( f"Calculate combinations \" { n } choose { k } \" in numpy:" ) for comb in itertools . combinations ( np . arange ( start = 0 , stop = n , dtype = jnp . int32 ), k ): print ( comb ) # combinadics print ( "Calculate via combinadics:" ) actual = n - 1 - calculateMth ( n , k , totalcount - 1 - jnp . arange ( start = 0 , stop = n , dtype = jnp . int32 ),) for comb in actual : print ( comb ) And the output from execution of the code is: Calculate combinations "4 choose 3" in numpy: (0, 1, 2) (0, 1, 3) (0, 2, 3) (1, 2, 3) Calculate via combinadics: [0 1 2] [0 1 3] [0 2 3] [1 2 3] A bit of theory You can think of a combinadic as an alternate representation of an integer. Consider the integer $859$ . It can be represented as the sum of powers of $10$ as $$ 859 = 8 \times 10^2 + 5 \times 10^1 + 9 \times 10^0 $$ The combinadic of an integer is its representation based on a variable base corresponding to the values of the binomial coefficient $\dbinom{n}{k}$ . For example if ( $n=7, k=4$ ) then the integer $27$ can be represented as $$ 27 = \dbinom{{6}}{4} + \dbinom{5}{3} + \dbinom{2}{2} + \dbinom{1}{1} = 15 + 10 + 1 + 1 $$ With ( $n=7, k=4$ ), any number $m$ between $0$ and $34$ (the total number of combination elements for $n$ and $k$ ) can be uniquely represented as $$ m = \dbinom{c_1}{4} + \dbinom{c_2}{3}+\dbinom{c_3}{2} + \dbinom{c_4}{1} $$ where $n > c_1 > c_2 > c_3 > c_4$ . Notice that $n$ is analogous to the base because all ...