Code to handle arbitrary number of for loops in Python/Numba

Question

I have a function compiled under the 'njit' framework in Numba. It looks like this:

import numpy as np
from numba import njit, types, prange
from numba.typed import List

@njit(cache=CACHE_FLAG)
def find_combinations(target, *arrays):
    """
    Find all combinations of element indices from 7 arrays containing segment lengths that sum to the target update time.
    Returns both the lengths that sum up to the target and the corresponding indices.

    Args:
        target (int): The target sum.
        arrays (tuple of lists): Each list contains segment lengths.

    Returns:
        Tuple of two lists:
            1. Lengths list - contains tuples of segment lengths that sum up to the target.
            2. Indices list - contains tuples of indices in the original lists corresponding to the lengths.
    """
    lengths_list = List()
    indices_list = List()
    for i in prange(len(arrays[0])):
        sum_i = arrays[0][i]
        if sum_i > target:
            continue
        for j in range(len(arrays[1])):
            sum_j = sum_i + arrays[1][j]
            if sum_j > target:
                continue
            for k in range(len(arrays[2])):
                sum_k = sum_j + arrays[2][k]
                if sum_k > target:
                    continue
                for l in range(len(arrays[3])):
                    sum_l = sum_k + arrays[3][l]
                    if sum_l > target:
                        continue
                    for m in range(len(arrays[4])):
                        sum_m = sum_l + arrays[4][m]
                        if sum_m > target:
                            continue
                        for n in range(len(arrays[5])):
                            sum_n = sum_m + arrays[5][n]
                            if sum_n > target:
                                continue
                            for o in range(len(arrays[6])):
                                total = sum_n + arrays[6][o]
                                if total == target:
                                    lengths_list.append(
                                        (
                                            arrays[0][i],
                                            arrays[1][j],
                                            arrays[2][k],
                                            arrays[3][l],
                                            arrays[4][m],
                                            arrays[5][n],
                                            arrays[6][o],
                                        )
                                    )
                                    indices_list.append((i, j, k, l, m, n, o))

    return lengths_list, indices_list

The function works as expected but is designed to work for an 'arrays' input of length 7. The outputs lengths list and indices list contain some number of tuples of length 7 (= length of input). I want to re-write this function to be able to handle inputs of arbitrary lengths.

I tried to rewrite this function using recursion like this:

import numpy as np
from numba import njit, types, prange
from numba.typed import List

@njit
def find_combinations_recursive(arrays, target, current_sum, current_indices, current_lengths, lengths_list, indices_list, depth):
    # Base case: when we've processed all arrays
    if depth == len(arrays):
        if current_sum == target:
            # lengths_list.append(tuple(current_lengths)) # tuple does not work in the Pythonic way-- cannot convert iterable to tuple in Numba?
            # indices_list.append(tuple(current_indices))
            lengths_list.append(current_lengths)
            indices_list.append(current_indices)
        return
    
    # Recursive case: iterate over the current array
    for i in range(len(arrays[depth])):
        new_sum = current_sum + arrays[depth][i]
        if new_sum > target:
            continue
        current_indices.append(i)
        # current_lengths.extend(List([arrays[depth][i]]))
        current_lengths.append(arrays[depth][i])
        find_combinations_recursive(
            arrays, target, new_sum, current_indices, current_lengths, lengths_list, indices_list, depth + 1
        )

@njit(cache=True)
def find_combinations(target, arrays):
    # lengths_list = List()
    # indices_list = List()
    # lengths_list = List.empty_list(tuple([np.int64] * len(arrays)))
    # indices_list = List.empty_list(tuple([np.int64] * len(arrays)))
    lengths_list = List([List([1]*len(arrays))])
    # print(lengths_list)
    # indices_list = List(tuple([1] * len(arrays)))
    indices_list = List([List([1]*len(arrays))])
    # print(indices_list)
    # import pdb; pdb.set_trace()
    k = lengths_list.pop()
    x = indices_list.pop()
    current_indices = List.empty_list(types.int64)
    current_lengths = List.empty_list(types.int64)
    find_combinations_recursive(arrays, target, 0, current_indices, current_lengths, lengths_list, indices_list, 0)
    return lengths_list, indices_list

However, while the function runs it does not produce the expected output-- for some reason, lengths_list does not contain the expected elements and seems to contain a single list of ints rather than multiple lists of ints of length 7. Also, I had a lot of trouble with Numba's List() data structure and eventually resorted to a sort of hacky way of initializing where I give it an initial value to set the type and then pop this value out to obtain an empty list of desired type. What do I need to fix to get this all to work coherently?

itertools.combinations already solves the problem of computing combinations of an arbitrary number of sets. I don't know how or if Numba can optimize it, though. — chepner
– chepner, Commented Aug 27, 2024 at 14:45

ken · Accepted Answer · 2024-09-01 18:01:38Z

Your recursive implementation was almost complete, except for one bug.

Here is the problem:

    # Recursive case: iterate over the current array
    for i in range(len(arrays[depth])):
        new_sum = current_sum + arrays[depth][i]
        if new_sum > target:
            continue
        current_indices.append(i)  # <-- HERE!

Since you are adding to current_indices in-place, it keeps growing longer with each loop iteration. But, this list should grow longer with each recursive call, not with each loop iteration. This was your only mistake.

Regarding how to create a list of lists, although there isn't any documentation that I know of, it's quite simple. You can do it like this:

indices_list = List.empty_list(List.empty_list(types.int64))

Or you can even use another typed.List like this:

current_indices = List.empty_list(types.int64)
indices_list = List.empty_list(current_indices)

Here is the code with the problem fixed. Although it wasn't part of your request, I modified the code to return only the indices. Since you can easily retrieve the values from arrays using the indices, separating these two concerns makes the code simpler.

import numpy as np
from numba import njit, types
from numba.typed import List


@njit
def find_combinations_recursive(arrays, target, current_sum, current_indices, indices_list, depth):
    if depth == len(arrays):
        if current_sum == target:
            indices_list.append(current_indices)
        return

    for i in range(len(arrays[depth])):
        new_sum = current_sum + arrays[depth][i]
        if new_sum > target:
            continue

        # This is an in-place operation, and it will grow longer for each iteration.
        # current_indices.append(i)

        # This is a copy operation, and a new list is created for each iteration.
        new_indices = current_indices.copy()
        new_indices.append(i)

        find_combinations_recursive(arrays, target, new_sum, new_indices, indices_list, depth + 1)


@njit(cache=True)
def find_combinations(target, arrays):
    current_indices = List.empty_list(types.int64)

    # To create a list of lists, like this.
    indices_list = List.empty_list(List.empty_list(types.int64))

    # Or you can do it like this.
    # indices_list = List.empty_list(current_indices)

    find_combinations_recursive(arrays, target, 0, current_indices, indices_list, 0)
    return indices_list


@njit
def get_lengths_list(arrays, indices_list):
    """Retrieve values from arrays using lists of indices."""
    lengths_list = []
    for indices in indices_list:
        lengths_list.append([arrays[i][indices[i]] for i in range(len(arrays))])
    return lengths_list


def main():
    np.random.seed(0)
    arrays = np.random.randint(0, 100, (7, 7))
    target = arrays.max(axis=1).sum()
    print("arrays:")
    print(arrays)
    print(f"{target=}")
    indices_list = find_combinations(target, arrays)
    lengths_list = get_lengths_list(arrays, indices_list)
    print(f"{indices_list=}")
    print(f"{lengths_list=}")


if __name__ == '__main__':
    main()

Result:

arrays:
[[44 47 64 67 67  9 83]
 [21 36 87 70 88 88 12]
 [58 65 39 87 46 88 81]
 [37 25 77 72  9 20 80]
 [69 79 47 64 82 99 88]
 [49 29 19 19 14 39 32]
 [65  9 57 32 31 74 23]]
target=np.int64(561)
indices_list=ListType[ListType[int64]]([[6, 4, 5, 6, 5, 0, 5], [6, 5, 5, 6, 5, 0, 5]])
lengths_list=[[83, 88, 88, 80, 99, 49, 74], [83, 88, 88, 80, 99, 49, 74]]

Collectives™ on Stack Overflow

Code to handle arbitrary number of for loops in Python/Numba

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related