phoenicyan/combinadics: fast combinations calculation in jax

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.

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]

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 combinadic digits are between $0$ and $n-1$ (just like all digits in ordinary base $10$ are between $0$ and $9$). The value of $k$ determines the number of terms in the combinadic.

Here’s an example of how a combinadic is calculated. Suppose you are working with ($n=7, k=4$) combinations, and $m = 8$. You want the combinadic of 8 because, as it turns out, the combinadic can be converted to combination element [8].

The combinadic of 8 will have the form:

$$
8 = \dbinom{c_1}{4} + \dbinom{c_2}{3}+\dbinom{c_3}{2} + \dbinom{c_4}{1}
$$

The first step is to determine the value of $c_1$. We try $c_1$ = $6$ (the largest number less than $n = 7$) and get $\dbinom{6}{4} = 15$, which is too large because we’re over $8$. Next, we try $c_1 = 5$ and get $\dbinom{5}{4} = 5$, which is less than $8$, so bingo, $c_1 = 5$.

At this point we have used up $5$ of the original number $m=8$ so we have $3$ left to account for. To determine the value of $c_2$, we try $4$ (the largest number less than the $5$ we got for $c_1$), but get $\dbinom{4}{3} = 4$, which is barely too large. Working down we get to $c_2 = 3$ and $\dbinom{3}{3} = 1$, so $c_2 = 3$.

We used up $1$ of the remaining $3$ we had to account for, so we have $2$ left to consume. Using the same ideas we’ll get $c_3 = 2$ with $\dbinom{2}{2} = 1$, so we have $1$ left to account for. And then we’ll find that $c_4 = 1$ because $\dbinom{1}{1} = 1$. Putting our four $c$ values together we conclude that the combinadic of $m=8$ for $(n=7, k=4)$ combinations is ( $5$ $3$ $2$ $1$ ).

Suppose $(n=7, k=4)$. There are $\dbinom{7}{4} = 35$ combination elements, indexed from $0$ to $34$. The dual lexicographic indexes are the ones on opposite ends so to speak: indexes $0$ and $34$ are duals, indexes $1$ and $33$ are duals, indexes $2$ and $32$ are duals, and so forth. Notice that each pair of dual indexes sum to $34$, so if you know any index it is easy to find its dual.

Now, continuing the first example above for the number $m=27$ with $(n=7, k=4)$ suppose you are able to find the combinadic of $27$ and get ( $6$ $5$ $2$ $1$ ). Now suppose you subtract each digit in the combinadic from $n-1 = 6$ and get ( $0$ $1$ $4$ $5$ ). Amazingly, this gives you the combination element $[7]$, the dual index of $27$. Putting these ideas together you have an elegant algorithm to determine an arbitrarily specified combination element for given $n$ and $k$ values. To find the combination element for index $m$, first find its dual and call it $x$. Next, find the combinadic of $x$. Then subtract each digit of the combinadic of $x$ from $n-1$ and the result is the mth lexicographic combination element.

The table below shows the relationships among $m$, the dual of $m$, combination element $[m]$, the combinadic of $m$, and $(n-1) – c_i$ for $(n=5, k=3)$.

m dual(m) Element(m) combinadic(m) (n-1) - ci
==============================================
[0]  9    { 0 1 2 }   ( 2 1 0 )     ( 2 3 4 )
[1]  8    { 0 1 3 }   ( 3 1 0 )     ( 1 3 4 )
[2]  7    { 0 1 4 }   ( 3 2 0 )     ( 1 2 4 )
[3]  6    { 0 2 3 }   ( 3 2 1 )     ( 1 2 3 )
[4]  5    { 0 2 4 }   ( 4 1 0 )     ( 0 3 4 )
[5]  4    { 0 3 4 }   ( 4 2 0 )     ( 0 2 4 )
[6]  3    { 1 2 3 }   ( 4 2 1 )     ( 0 2 3 )
[7]  2    { 1 2 4 }   ( 4 3 0 )     ( 0 1 4 )
[8]  1    { 1 3 4 }   ( 4 3 1 )     ( 0 1 3 )
[9]  0    { 2 3 4 }   ( 4 3 2 )     ( 0 1 2 )

64-bit numbers
Performance of a single GPU