Быстрый векторизованный полином в питоне

Question

В настоящее время я использую NumPy для выполнения следующей задачи: у меня большая сетка значений, и мне нужно взять полиномиальную выборку в каждой точке.Вектор вероятности для многочлена будет варьироваться от узла сетки к узлу сетки, поэтому функция многочлена NumPy мне не совсем подходит, поскольку она выполняет все свои отрисовки из одного и того же распределения.Итерации по каждому сайту кажутся невероятно неэффективными, и мне было интересно, есть ли способ сделать это эффективно в NumPy.Похоже, что-то подобное возможно (и быстро), если я использую Theano (см. этот ответ ), но для этого потребуется довольно много переписать, чего в идеале я бы хотел избежать.Можно ли эффективно векторизовать многочленную выборку в базовом NumPy?

Редактировать: Код Уоррена легко модифицировать, чтобы разрешить различное число на разных сайтах, поскольку я обнаружил, что мне нужно: все, что нужно сделать, это передатьполный count массив и удалить первый, например, так:

import numpy as np


def multinomial_rvs(count, p):
    """
    Sample from the multinomial distribution with multiple p vectors.

    * count must be an (n-1)-dimensional numpy array.
    * p must an n-dimensional numpy array, n >= 1.  The last axis of p
      holds the sequence of probabilities for a multinomial distribution.

    The return value has the same shape as p.
    """
    out = np.zeros(p.shape, dtype=int)
    ps = p.cumsum(axis=-1)
    # Conditional probabilities
    with np.errstate(divide='ignore', invalid='ignore'):
        condp = p / ps
    condp[np.isnan(condp)] = 0.0
    for i in range(p.shape[-1]-1, 0, -1):
        binsample = np.random.binomial(count, condp[..., i])
        out[..., i] = binsample
        count -= binsample
    out[..., 0] = count
    return out

Answer 1

Вот один из способов, которым вы можете это сделать.Он не полностью векторизован, но цикл Python превышает значения p.Если длина вашего p вектора (ов) не слишком велика, это может быть достаточно быстрым для вас.

Полиномиальное распределение реализовано с использованием повторных вызовов np.random.binomial, который реализует передачу своих аргументов.

import numpy as np


def multinomial_rvs(n, p):
    """
    Sample from the multinomial distribution with multiple p vectors.

    * n must be a scalar.
    * p must an n-dimensional numpy array, n >= 1.  The last axis of p
      holds the sequence of probabilities for a multinomial distribution.

    The return value has the same shape as p.
    """
    count = np.full(p.shape[:-1], n)
    out = np.zeros(p.shape, dtype=int)
    ps = p.cumsum(axis=-1)
    # Conditional probabilities
    with np.errstate(divide='ignore', invalid='ignore'):
        condp = p / ps
    condp[np.isnan(condp)] = 0.0
    for i in range(p.shape[-1]-1, 0, -1):
        binsample = np.random.binomial(count, condp[..., i])
        out[..., i] = binsample
        count -= binsample
    out[..., 0] = count
    return out

Вот пример, где «сетка» имеет форму (2, 3), а полиномиальное распределение является четырехмерным (то есть каждый p вектор имеет длину 4).

In [182]: p = np.array([[[0.25, 0.25, 0.25, 0.25], 
     ...:                [0.01, 0.02, 0.03, 0.94], 
     ...:                [0.75, 0.15, 0.05, 0.05]], 
     ...:               [[0.01, 0.99, 0.00, 0.00], 
     ...:                [1.00, 0.00, 0.00, 0.00], 
     ...:                [0.00, 0.25, 0.25, 0.50]]])                                                                                                 

In [183]: sample = multinomial_rvs(1000, p)                                     

In [184]: sample                                                                
Out[184]: 
array([[[ 249,  260,  233,  258],
        [   3,   21,   33,  943],
        [ 766,  131,   55,   48]],

       [[   5,  995,    0,    0],
        [1000,    0,    0,    0],
        [   0,  273,  243,  484]]])

In [185]: sample.sum(axis=-1)                                                   
Out[185]: 
array([[1000, 1000, 1000],
       [1000, 1000, 1000]])

---

В комментарии вы сказали: «Вектор p имеет вид: p = [p_s, (1-p_s) / 4, (1-p_s) / 4, (1-p_s) / 4»), (1-p_s) / 4], причем p_s варьируется от сайта к сайту. "Вот как вы можете использовать вышеуказанную функцию, учитывая массив, содержащий значения p_s.

Сначала создайте некоторые данные для примера:

In [73]: p_s = np.random.beta(4, 2, size=(2, 3))                                                                                                        

In [74]: p_s                                                                                                                                            
Out[74]: 
array([[0.61662208, 0.6072323 , 0.62208711],
       [0.86848938, 0.58959038, 0.47565799]])

Создайте массив p, содержащийполиномиальные вероятности по формуле p = [p_s, (1-p_s)/4, (1-p_s)/4, (1-p_s)/4, (1-p_s)/4]:

In [75]: p = np.expand_dims(p_s, -1) * np.array([1, -0.25, -0.25, -0.25, -0.25]) + np.array([0, 0.25, 0.25, 0.25, 0.25])                                

In [76]: p                                                                                                                                              
Out[76]: 
array([[[0.61662208, 0.09584448, 0.09584448, 0.09584448, 0.09584448],
        [0.6072323 , 0.09819192, 0.09819192, 0.09819192, 0.09819192],
        [0.62208711, 0.09447822, 0.09447822, 0.09447822, 0.09447822]],

       [[0.86848938, 0.03287765, 0.03287765, 0.03287765, 0.03287765],
        [0.58959038, 0.1026024 , 0.1026024 , 0.1026024 , 0.1026024 ],
        [0.47565799, 0.1310855 , 0.1310855 , 0.1310855 , 0.1310855 ]]])

Теперь сделайте то же самое, что и раньше, чтобы сгенерировать выборку (измените значение 1000 на то, что подходит для вашей задачи):

In [77]: sample = multinomial_rvs(1000, p)                                                                                                              

In [78]: sample                                                                                                                                         
Out[78]: 
array([[[618,  92, 112,  88,  90],
        [597, 104, 103, 101,  95],
        [605, 100,  95,  98, 102]],

       [[863,  32,  43,  27,  35],
        [602, 107, 108,  94,  89],
        [489, 130, 129, 129, 123]]])

In [79]: sample.sum(axis=-1)                                                                                                                            
Out[79]: 
array([[1000, 1000, 1000],
       [1000, 1000, 1000]])

Answer 2

Возможное решение

В принципе, вы можете сделать это с numba, так как поддерживается распространение multinomial.

Numba позволяет вам просто украсить свой numpy (и, что еще более важно, стандартный Pythonфункции) с numba.njit декоратором для получения значительного прироста производительности.

Проверьте их документацию , чтобы увидеть этот подход более подробно.Особенно 2.7.4, поскольку речь идет о поддержке np.random (также поддерживается многочленное распределение).

Недостаток : size аргумент в настоящее время не поддерживается.Вы можете вызывать np.random.multinomial несколько раз во вложенном цикле, и все же это должно быть быстрее, если оформлено с помощью numba.njit.

Последнее, но не менее важное: вы можете распараллелить внешний цикл с аргументами numba.prange и parallelвышеупомянутому декоратору.

Тест производительности

Первый тест:

• непараллелизованная numba с сигнатурой типов
• вообще не numba

Код для тестирования:

import sys
from functools import wraps
from time import time

import numba
import numpy as np


def timing(function):
    @wraps(function)
    def wrap(*args, **kwargs):
        start = time()
        result = function(*args, **kwargs)
        end = time()
        print(f"Time elapsed: {end - start}", file=sys.stderr)
        return result

    return wrap


@timing
@numba.njit(numba.int64(numba.int64[:, :], numba.int64))
def my_multinomial(probabilities, output):
    experiments: int = 5000
    output_array = []
    for i in numba.prange(probabilities.shape[0]):
        probability = probabilities[i] / np.sum(probabilities[i])
        result = np.random.multinomial(experiments, pvals=probability)
        if i % output == 0:
            output_array.append(result)

    return output_array[-1][-1]


if __name__ == "__main__":
    np.random.seed(0)
    probabilities = np.random.randint(low=1, high=100, size=(10000, 1000))
    for _ in range(5):
        output = my_multinomial(probabilities, np.random.randint(low=3000, high=10000))

Результаты:

Непараллелизованная нумба с типами подписи

Time elapsed: 1.0510437488555908
Time elapsed: 1.0691254138946533
Time elapsed: 1.065258264541626
Time elapsed: 1.0559568405151367
Time elapsed: 1.0446960926055908

Никакой нумбы вообще

Time elapsed: 0.9460861682891846
Time elapsed: 0.9581060409545898
Time elapsed: 0.9654934406280518
Time elapsed: 0.9708254337310791
Time elapsed: 0.9757359027862549

Как видно, numba не поможет в этом случае (на самом деле это ухудшает производительность).Результаты согласуются для переменных размеров входного массива.

Второй тест

• распараллеленная нумба без сигнатуры типов
• вообще никакой нумбы

Код для теста:

import sys
from functools import wraps
from time import time

import numba
import numpy as np


def timing(function):
    @wraps(function)
    def wrap(*args, **kwargs):
        start = time()
        result = function(*args, **kwargs)
        end = time()
        print(f"Time elapsed: {end - start}", file=sys.stderr)
        return result

    return wrap


@timing
@numba.njit(parallel=True)
def my_multinomial(probabilities, output):
    experiments: int = 5000
    for i in range(probabilities.shape[0]):
        probability = probabilities[i] / np.sum(probabilities[i])
        result = np.random.multinomial(experiments, pvals=probability)
        if i % output == 0:
            print(result)


if __name__ == "__main__":
    np.random.seed(0)
    probabilities = np.random.randint(low=1, high=100, size=(10000, 1000))
    for _ in range(5):
        my_multinomial(probabilities, np.random.randint(low=3000, high=10000))

Результаты:

Распараллеленная нумба без подписи типов:

Time elapsed: 1.0705969333648682                                                                                                                                                          
Time elapsed: 0.18749785423278809                                                                                                                                                         
Time elapsed: 0.1877145767211914                                                                                                                                                          
Time elapsed: 0.18813610076904297                                                                                                                                                         
Time elapsed: 0.18747472763061523 

Никакой нумбы вообще

Time elapsed: 1.0142333507537842                                                                                                                                                          
Time elapsed: 1.0311956405639648                                                                                                                                                          
Time elapsed: 1.022024154663086                                                                                                                                                           
Time elapsed: 1.0191617012023926                                                                                                                                                          
Time elapsed: 1.0144879817962646

Частичные выводы

Как правильно указано в комментариях max9111 , я слишком рано пришел к своим выводам.Кажется, что распараллеливание (если возможно) было бы самой большой помощью в вашем случае, в то время как numba (по крайней мере, в этом все еще простом и не слишком всеобъемлющем тесте) не создает больших улучшений.

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