Ускорить продукт itertools

Я использую itertools.product для генерации всех возможных вариаций 4 элементов длины 13. 4 и 13 могут быть произвольными, но на самом деле я получаю 4^13 результатов, что много. Мне нужен результат в виде массива Numpy, и в настоящее время я делаю следующее:

  c = it.product([1,-1,np.complex(0,1), np.complex(0,-1)], repeat=length)
  sendbuf = np.array(list(c))

С некоторым простым профилирующим кодом, вставленным между ними, похоже, что первая строка в значительной степени мгновенная, тогда как преобразование в список, а затем в массив Numpy занимает около 3 часов. Есть ли способ сделать это быстрее? Это, наверное, что-то действительно очевидное, что я пропускаю.

Спасибо!

6 ответов

Решение

NumPy эквивалент itertools.product() является numpy.indices(), но это только даст вам произведение диапазонов вида 0,...,k-1:

numpy.rollaxis(numpy.indices((2, 3, 3)), 0, 4)
array([[[[0, 0, 0],
         [0, 0, 1],
         [0, 0, 2]],

        [[0, 1, 0],
         [0, 1, 1],
         [0, 1, 2]],

        [[0, 2, 0],
         [0, 2, 1],
         [0, 2, 2]]],


       [[[1, 0, 0],
         [1, 0, 1],
         [1, 0, 2]],

        [[1, 1, 0],
         [1, 1, 1],
         [1, 1, 2]],

        [[1, 2, 0],
         [1, 2, 1],
         [1, 2, 2]]]])

Для вашего особого случая вы можете использовать

a = numpy.indices((4,)*13)
b = 1j ** numpy.rollaxis(a, 0, 14)

(Это не будет работать в 32-битной системе, потому что массив слишком велик. Однако, экстраполируя на размер, который я могу протестировать, он должен работать менее чем за минуту.)

Эйдт: Просто чтобы упомянуть это: призыв к numpy.rollaxis() является более или менее косметическим, чтобы получить тот же результат, что и itertools.product(), Если вам не важен порядок индексов, вы можете просто его опустить (но в любом случае это дешево, если у вас нет последующих операций, которые преобразуют ваш массив в непрерывный массив).

EDIT2: чтобы получить точный аналог

numpy.array(list(itertools.product(some_list, repeat=some_length)))

ты можешь использовать

numpy.array(some_list)[numpy.rollaxis(
    numpy.indices((len(some_list),) * some_length), 0, some_length + 1)
    .reshape(-1, some_length)]

Это стало совершенно нечитаемым - просто скажите мне, должен ли я объяснить это дальше:)

Первая строка кажется мгновенной, потому что никакой фактической операции не происходит. Генераторный объект создается только тогда, когда вы выполняете его итерацию в процессе работы. Как вы сказали, вы получаете 4^13 = 67108864 числа, все они рассчитаны и доступны во время вашего list вызов. Я вижу, что np.array принимает только список или кортеж, поэтому вы можете попробовать создать кортеж из вашего итератора и передать его в np.array, чтобы увидеть, есть ли разница в производительности, и это не влияет на общую производительность вашей программы., Это можно определить, только попробовав ваш вариант использования, хотя есть некоторые моменты, которые говорят, что кортеж немного быстрее.

Чтобы попробовать с кортежем, вместо списка просто сделайте

sendbuf = np.array(tuple(c))

Вы можете ускорить процесс, пропустив преобразование в список:

numpy.fromiter(c, count=…)  # Using count also speeds things up, but it's optional

С помощью этой функции массив NumPy сначала выделяется, а затем инициализируется элемент за элементом, без необходимости проходить дополнительный этап построения списка.

PS: fromiter() не обрабатывает кортежи, возвращаемые product(), так что это может быть не решение, пока. Если fromiter() сделал ручку dtype=object, это должно работать, хотя.

PPS: Как отметил Джо Кингтон, это можно сделать, поместив кортежи в структурированный массив. Тем не менее, это не всегда дает ускорение.

Пусть numpy.meshgrid сделает всю работу:

length = 13
x = [1, -1, 1j, -1j]
mesh = numpy.meshgrid(*([x] * length))
result = numpy.vstack([y.flat for y in mesh]).T

на моем ноутбуке это занимает ~2 минуты

Возможно, вы захотите попробовать совершенно другой подход: сначала создайте пустой массив нужного размера:

result = np.empty((4**length, length), dtype=complex)

затем используйте способности нарезки NumPy, чтобы заполнить массив самостоятельно:

# Set up of the last "digit":
result[::4, length-1] = 1
result[1::4, length-1] = -1
result[2::4, length-1] = 1j
result[3::4, length-1] = -1j

Вы можете сделать аналогичные вещи для других "цифр" (то есть элементов result[:, 2], result[:, 1] и result[:, 0]). Все это, безусловно, можно поместить в цикл, который повторяется по каждой цифре.

Транспонирование всей операции (np.empty((length, 4**length)…)) стоит попробовать, так как это может привести к увеличению скорости (за счет лучшего использования кеша памяти).

Вероятно, не оптимизирован, но гораздо менее зависим от преобразования типов Python:

ints = [1,2,3,4]
repeat = 3

def prod(ints, repeat):
    w = repeat
    l = len(ints)
    h = l**repeat
    ints = np.array(ints)
    A = np.empty((h,w), dtype=int)
    rng = np.arange(h)
    for i in range(w):
        x = l**i
        idx = np.mod(rng,l*x)/x
        A[:,i] = ints[idx]
    return A   
Другие вопросы по тегам