Ускорьте присвоение плотностей вероятностей с учетом двух значений в Python 3

Question

Для некоторых моих исследований мне нужно назначить плотность вероятности, учитывая значение, среднее значение и стандартное отклонение, за исключением того, что мне нужно сделать это около 40 миллионов раз, поэтому ускорение этого кода становится критически важным для работы в продуктивная мода.

У меня есть только 10 значений для тестирования (матрица значений = 10x1), но я хочу назначить вероятность для каждого из этих значений, учитывая всего 4 миллиона усеченных нормальных распределений на значение, каждое из которых имеет различные средние значения (all_means = 4 миллиона x 10 матрицы), и то же самое стандартное отклонение (ошибка = 1 значение). Код, который я использовал для этого, приведен ниже:

import scipy.stats as ss

all_probabilities =[]

for row in all_means:

    temp_row = []
    for i in range(len(row)):

        # Isolate key values
        mean = row[i]
        error = 0.05
        value = values[i]

        # Create truncated normal distribution and calculate PMF
        a, b = 0, np.inf
        mu, sigma = float(mean), float(error)
        alpha, beta = ((a-mu)/sigma), ((b-mu)/sigma)
        prob = ss.truncnorm.pdf(float(value), alpha, beta, loc=mu, scale=sigma)
        temp_row.extend([prob])

    all_probabilities.extend([temp_row])

Один l oop занимает в среднем 5 мс, но для этого 4 миллиона раз означает, что этот раздел кода будет займет около 5 часов. Я предполагаю, что ограничивающими факторами являются вызов ss.truncnorm.pdf и использование extends. Последнее я могу обойти, предварительно выделив матрицу вероятностей, но первое я не вижу обходного пути.

Для большего контекста этот бит кода является частью алгоритма, который использует этот код в среднем 5 раз (хотя и с быстрым уменьшением количества распространяемых тестов), поэтому любые советы по ускорению этого кода будут быть огромной помощью.

Извиняюсь, если это тривиально, я относительно новичок в оптимизации кода и не смог найти что-то конкретно по этой проблеме.

Answer 1

Вы можете избежать внутреннего l oop, так как scipy.stats.truncnorm можно определить как вектор случайных величин, т. Е.

import numpy as np
from scipy.stats import truncnorm

all_probabilities = []
a, b = 0, np.inf
error = 0.05

for row in all_means:

    alpha, beta = ((a-row )/error), ((b-row )/error)

    # vectorized truncnorm
    rv_tn = truncnorm(alpha, beta, loc=row, scale=error)

    # predict vector
    prob = rv_tn.pdf(values)

    all_probabilities.extend(prob)