Re-implement next_k_array; add k_array_rank

quantecon/game_theory/support_enumeration.py

@@ -14,6 +14,7 @@
 import numpy as np
 from numba import jit
 from ..util.numba import _numba_linalg_solve
+from ..util.combinatorics import next_k_array
 
 
 def support_enumeration(g):
@@ -106,8 +107,8 @@ def _support_enumeration_gen(payoff_matrix0, payoff_matrix1):
                                                                actions)):
                             out[p][supp] = action[:-1]
                         yield out
-                _next_k_array(supps[1])
-            _next_k_array(supps[0])
+                next_k_array(supps[1])
+            next_k_array(supps[0])
 
 
 @jit(nopython=True, cache=True)
@@ -180,84 +181,3 @@ def _indiff_mixed_action(payoff_matrix, own_supp, opp_supp, A, out):
             if payoff > val:
                 return False
     return True
-
-
-@jit(nopython=True, cache=True)
-def _next_k_combination(x):
-    """
-    Find the next k-combination, as described by an integer in binary
-    representation with the k set bits, by "Gosper's hack".
-
-    Copy-paste from en.wikipedia.org/wiki/Combinatorial_number_system
-
-    Parameters
-    ----------
-    x : int
-        Integer with k set bits.
-
-    Returns
-    -------
-    int
-        Smallest integer > x with k set bits.
-
-    """
-    u = x & -x
-    v = u + x
-    return v + (((v ^ x) // u) >> 2)
-
-
-@jit(nopython=True, cache=True)
-def _next_k_array(a):
-    """
-    Given an array `a` of k distinct nonnegative integers, return the
-    next k-array in lexicographic ordering of the descending sequences
-    of the elements. `a` is modified in place.
-
-    Parameters
-    ----------
-    a : ndarray(int, ndim=1)
-        Array of length k.
-
-    Returns
-    -------
-    a : ndarray(int, ndim=1)
-        View of `a`.
-
-    Examples
-    --------
-    Enumerate all the subsets with k elements of the set {0, ..., n-1}.
-
-    >>> n, k = 4, 2
-    >>> a = np.arange(k)
-    >>> while a[-1] < n:
-    ...     print(a)
-    ...     a = _next_k_array(a)
-    ...
-    [0 1]
-    [0 2]
-    [1 2]
-    [0 3]
-    [1 3]
-    [2 3]
-
-    """
-    k = len(a)
-    if k == 0:
-        return a
-
-    x = 0
-    for i in range(k):
-        x += (1 << a[i])
-
-    x = _next_k_combination(x)
-
-    pos = 0
-    for i in range(k):
-        while x & 1 == 0:
-            x = x >> 1
-            pos += 1
-        a[i] = pos
-        x = x >> 1
-        pos += 1
-
-    return a

quantecon/util/combinatorics.py

@@ -0,0 +1,122 @@
+"""
+Useful routines for combinatorics
+
+"""
+from scipy.special import comb
+from numba import jit
+
+from .numba import comb_jit
+
+
+@jit(nopython=True, cache=True)
+def next_k_array(a):
+    """
+    Given an array `a` of k distinct nonnegative integers, sorted in
+    ascending order, return the next k-array in the lexicographic
+    ordering of the descending sequences of the elements [1]_. `a` is
+    modified in place.
+
+    Parameters
+    ----------
+    a : ndarray(int, ndim=1)
+        Array of length k.
+
+    Returns
+    -------
+    a : ndarray(int, ndim=1)
+        View of `a`.
+
+    Examples
+    --------
+    Enumerate all the subsets with k elements of the set {0, ..., n-1}.
+
+    >>> n, k = 4, 2
+    >>> a = np.arange(k)
+    >>> while a[-1] < n:
+    ...     print(a)
+    ...     a = next_k_array(a)
+    ...
+    [0 1]
+    [0 2]
+    [1 2]
+    [0 3]
+    [1 3]
+    [2 3]
+
+    References
+    ----------
+    .. [1] `Combinatorial number system
+       <https://en.wikipedia.org/wiki/Combinatorial_number_system>`_,
+       Wikipedia.
+
+    """
+    # Logic taken from Algotirhm T in D. Knuth, The Art of Computer
+    # Programming, Section 7.2.1.3 "Generating All Combinations".
+    k = len(a)
+    if k == 1 or a[0] + 1 < a[1]:
+        a[0] += 1
+        return a
+
+    a[0] = 0
+    i = 1
+    x = a[i] + 1
+
+    while i < k-1 and x == a[i+1]:
+        i += 1
+        a[i-1] = i - 1
+        x = a[i] + 1
+    a[i] = x
+
+    return a
+
+
+def k_array_rank(a):
+    """
+    Given an array `a` of k distinct nonnegative integers, sorted in
+    ascending order, return its ranking in the lexicographic ordering of
+    the descending sequences of the elements [1]_.
+
+    Parameters
+    ----------
+    a : ndarray(int, ndim=1)
+        Array of length k.
+
+    Returns
+    -------
+    idx : scalar(int)
+        Ranking of `a`.
+
+    References
+    ----------
+    .. [1] `Combinatorial number system
+       <https://en.wikipedia.org/wiki/Combinatorial_number_system>`_,
+       Wikipedia.
+
+    """
+    k = len(a)
+    idx = int(a[0])  # Convert to Python int
+    for i in range(1, k):
+        idx += comb(a[i], i+1, exact=True)
+    return idx
+
+
+@jit(nopython=True, cache=True)
+def k_array_rank_jit(a):
+    """
+    Numba jit version of `k_array_rank`.
+
+    Notes
+    -----
+    An incorrect value will be returned without warning or error if
+    overflow occurs during the computation. It is the user's
+    responsibility to ensure that the rank of the input array fits
+    within the range of possible values of `np.intp`; a sufficient
+    condition for it is `scipy.special.comb(a[-1]+1, len(a), exact=True)
+    <= np.iinfo(np.intp).max`.
+
+    """
+    k = len(a)
+    idx = a[0]
+    for i in range(1, k):
+        idx += comb_jit(a[i], i+1)
+    return idx

quantecon/util/tests/test_combinatorics.py

@@ -0,0 +1,70 @@
+"""
+Tests for util/combinatorics.py
+
+"""
+import numpy as np
+from numpy.testing import assert_array_equal
+from nose.tools import eq_
+import scipy.special
+from quantecon.util.combinatorics import (
+    next_k_array, k_array_rank, k_array_rank_jit
+)
+
+
+class TestKArray:
+    def setUp(self):
+        self.k_arrays = np.array(
+            [[0, 1, 2],
+             [0, 1, 3],
+             [0, 2, 3],
+             [1, 2, 3],
+             [0, 1, 4],
+             [0, 2, 4],
+             [1, 2, 4],
+             [0, 3, 4],
+             [1, 3, 4],
+             [2, 3, 4],
+             [0, 1, 5],
+             [0, 2, 5],
+             [1, 2, 5],
+             [0, 3, 5],
+             [1, 3, 5],
+             [2, 3, 5],
+             [0, 4, 5],
+             [1, 4, 5],
+             [2, 4, 5],
+             [3, 4, 5]]
+        )
+        self.L, self.k = self.k_arrays.shape
+
+    def test_next_k_array(self):
+        k_arrays_computed = np.empty((self.L, self.k), dtype=int)
+        k_arrays_computed[0] = np.arange(self.k)
+        for i in range(1, self.L):
+            k_arrays_computed[i] = k_arrays_computed[i-1]
+            next_k_array(k_arrays_computed[i])
+        assert_array_equal(k_arrays_computed, self.k_arrays)
+
+    def test_k_array_rank(self):
+        for i in range(self.L):
+            eq_(k_array_rank(self.k_arrays[i]), i)
+
+    def test_k_array_rank_jit(self):
+        for i in range(self.L):
+            eq_(k_array_rank_jit(self.k_arrays[i]), i)
+
+
+def test_k_array_rank_arbitrary_precision():
+    n, k = 100, 50
+    a = np.arange(n-k, n)
+    eq_(k_array_rank(a), scipy.special.comb(n, k, exact=True)-1)
+
+
+if __name__ == '__main__':
+    import sys
+    import nose
+
+    argv = sys.argv[:]
+    argv.append('--verbose')
+    argv.append('--nocapture')
+    nose.main(argv=argv, defaultTest=__file__)