apply numba to itertools import product

1@nb.njit(nb.int32[:,:](nb.int32[:]))
2def cproduct_idx(sizes: np.ndarray):
3    """Generates ids tuples for a cartesian product"""
4    assert len(sizes) >= 2
5    tuples_count  = np.prod(sizes)
6    tuples = np.zeros((tuples_count, len(sizes)), dtype=np.int32)
7    tuple_idx = 0
8    # stores the current combination
9    current_tuple = np.zeros(len(sizes))
10    while tuple_idx < tuples_count:
11        tuples[tuple_idx] = current_tuple
12        current_tuple[0] += 1
13        # using a condition here instead of including this in the inner loop
14        # to gain a bit of speed: this is going to be tested each iteration,
15        # and starting a loop to have it end right away is a bit silly
16        if current_tuple[0] == sizes[0]:
17            # the reset to 0 and subsequent increment amount to carrying
18            # the number to the higher "power"
19            current_tuple[0] = 0
20            current_tuple[1] += 1
21            for i in range(1, len(sizes) - 1):
22                if current_tuple[i] == sizes[i]:
23                    # same as before, but in a loop, since this is going
24                    # to get run less often
25                    current_tuple[i + 1] += 1
26                    current_tuple[i] = 0
27                else:
28                    break
29        tuple_idx += 1
30    return tuples
31
32@nb.njit
33def cartesian_product(*arrays):
34    sizes = [len(a) for a in arrays]
35    sizes = np.asarray(sizes, dtype=np.int32)
36    tuples_count  = np.prod(sizes)
37    array_ids = cproduct_idx(sizes)
38    tuples = np.zeros((tuples_count, len(sizes)))
39    for i in range(len(arrays)):
40        tuples[:, i] = arrays[i][array_ids[:, i]]
41    return tuples
42
43@nb.njit
44def cartesian_product_repeat(array, repeat):
45    sizes = [len(array) for _ in range(repeat)]
46    sizes = np.asarray(sizes, dtype=np.int32)
47    tuples_count  = np.prod(sizes)
48    array_ids = cproduct_idx(sizes)
49    tuples = np.zeros((tuples_count, len(sizes)))
50    for i in range(repeat):
51        tuples[:, i] = array[array_ids[:, i]]
52    return tuples
53