Функция оптимизации Python Numba со словарем наборов

Превращение в Нумбу:

Это было утомительно, так как Numba накладывает много ограничений. Конечно, было бы лучше опубликовать этот код пошагово, однако тогда ответ будет слишком длинным.
Короче говоря, чтобы преобразовать код в Numba и использовать его полную оптимизацию с помощью nopython, вам нужно следовать примерно этим правилам. .
Отказ от ответственности: я не эксперт по нумбе, и этот раздел основан только на наблюдениях:

• Никаких объектов, кроме основных и пустых. Это означает отсутствие наборов и словарей.
• Функции могут возвращать только одно значение, а не массив.
• Все должно быть напечатано или достаточно просто, чтобы Numba автоматически определял тип.

Есть, вероятно, больше. Я не слишком уверен. Так что в основном, как следует из названия, если вы идете по маршруту nopython, забудьте, что вы находитесь в python. Либо код должен быть очень простым (как в нескольких циклах с простым вычислением), либо вам придется зарезать код до забвения. (Или есть какой-то другой трюк, который я еще не нашел.)
Редактировать: я нашел трюк - numba теперь поддерживает словарь даже в nojit, если вы заранее объявляете все типы. Он также поддерживает наборы. Однако он не может хранить наборы в словаре, поэтому в вашем случае вам придется сначала хранить наборы в массиве, а затем использовать словарь как своего рода таблицу перевода из id в индекс. Не невозможно, но определенно не просто.

Код, который скорее всего работает. Однако часть его устарела, так как я отказался от попыток избавиться от последнего [] в пользу массива numpy. Кроме того, он не оптимизирован, и я не измерял производительность.

import random
import numpy as np
from numba import jit
n_positions = 100
position_coverage = 10
min_dist = 5
n_ids = 5
near_keys = random.sample(range(n_positions), n_ids)


@jit(nopython=True)
    def simulate(near_keys):
    to_pos = -1
    _ixs_to = np.zeros(1)
    to_return = []
    for id in sorted(near_keys, key=lambda x: random.random()):
        if id < n_positions:
            near = np.arange(id, id + position_coverage)
        else:
            near = np.arange(id, n_positions)

        print(id, near)
        if to_pos == -1:
            _ixs_to = near
        else:
            to_close = []
            a = int(n_positions - position_coverage + min_dist)
            to_pos = int(to_pos)
        if to_pos < a:
            pos_values = np.arange(to_pos, to_pos + min_dist)
        else:
            pos_values = np.arange(to_pos, n_positions)

        for i in pos_values:
            if not i in to_close:
                to_close.append(i)

        _ixs_to = np.zeros(1, dtype=np.int64)
        for i in near:
            if not i in to_close:
                _ixs_to = np.append(_ixs_to, i)
        if _ixs_to[0] == 0:
            _ixs_to = _ixs_to[1:]

        if _ixs_to.ndim == 0:
            return [-1] # failed no reachable positions
        np.random.rand()
        to_pos = np.random.choice(_ixs_to)
        to_return.append(id)
        to_return.append(to_pos)
        return to_return

if __name__ == '__main__':
    print(simulate(near_keys))

Так что же случилось с вашим довольно красивым, хотя, может быть, излишне сложным кодом?

Во-первых, я заменил наборы генерирующими функциями:
near_dict = {x: set(range(x, x + position_coverage)) if x < n_positions - position_coverage else set(range(x, n_positions)) for x in range(n_positions)}
стал

def get_near(x):
    if x < n - position_coverage:
        return set(range(x, x + position_coverage))
    else:
        return set(range(x, n_positions)

Затем я переписал объединения и различия множеств так, чтобы они не требовались. Это было сделано путем перебора массивов и проверки значений для их удаления или добавления. Обратите внимание, что это та часть, где это, скорее всего, медленнее, если только операция объединения в чистом питоне не медленнее, чем неоптимизированное наивное объединение, которое я сделал. Если вы хотите оптимизировать код, это одно из мест, с которого можно начать.

К сожалению, хотя njit numba позволяет использовать массивы параметров, он не разрешить возвращаемое значение функции быть массивом. Поскольку нам нужны массивы, следующим шагом было встраивание любой функции, которая возвращала массив в исходном коде.

Читая документы, я заметил еще одну вещь: Numba не позволяет создавать массивы в nopython и в таком случае возвращается к объектному режиму. Это означает, что вам нужно будет перерабатывать массивы вместо их воссоздания в коде выше, чтобы полностью использовать этот режим.

Последним шагом было сгладить все несоответствия типов и убедиться, что все работает вместе. В результате получается код, который Numba nopython хочет запускать, но, к сожалению, довольно запутанный и не использующий практически никаких удобств Python. Он также ужасно неоптимизирован. Однако, если вы хотите применить numba и впоследствии получить дешевое распараллеливание, которое оно обеспечивает, это, скорее всего, необходимо. Если кто-то, более знакомый с numba, поправит меня и укажет направление, более удобное для пользователя, я с радостью уберу этот ответ как вводящий в заблуждение.

Подводя итог, хотя Numba, безусловно, является мощным инструментом, я не слишком уверен, разумно ли использовать его для вашего проекта, поскольку ваша основная проблема заключается не в простой повторяющейся математической операции, а в повторяющейся итерации логики, которая немного сложнее, чем предпочитает Нумба.

Несколько советов:

• Я бы сначала занялся профилированием и оптимизацией вашего кода. Легче профилировать код до того, как вы начнете делать параллельные махинации. Если этого окажется недостаточно, это самое время попробовать и:
• посмотрите в многопроцессорную библиотеку python, если такого ускорения будет достаточно для ты. Если этого недостаточно, то:
• попробуйте переписать свой уже оптимизированный код для numba и, наконец, если этого все еще недостаточно:
• попробуйте C, Rust, C++, Java или что-то подобное. Вы можете быть удивлены.

Заключительное слово: Хотя я достаточно уверен в том, как заменил наборы функциями, которые затем встроил, я не совсем уверен, что все сделал правильно. Рекомендуется осторожность.