Многопоточная генерация

Четыре основных распределения (random, standard_normal, standard_exponential и standard_gamma) все позволяют заполнять существующие массивы с использованием ключевого аргумента out. Существующие массивы должны быть непрерывными и корректными (доступными для записи и выровненными). В обычных условиях массивы, созданные с использованием общих конструкторов, таких как numpy.empty, будут соответствовать этим требованиям.

В этом примере используется Python 3 concurrent.futures для заполнения массива с помощью нескольких потоков. Потоки долгоживущие, поэтому повторные вызовы не требуют дополнительных затрат на создание потоков. Базовый BitGenerator — PCG64, который быстрый, имеет большой период и поддерживает использование PCG64.jumped для возврата нового генератора при одновременном продвижении состояния. Сгенерированные случайные числа воспроизводимы в том смысле, что один и тот же seed приведет к тем же результатам.

from numpy.random import Generator, PCG64
import multiprocessing
import concurrent.futures
import numpy as np

class MultithreadedRNG(object):
    def __init__(self, n, seed=None, threads=None):
        rg = PCG64(seed)
        if threads is None:
            threads = multiprocessing.cpu_count()
        self.threads = threads

        self._random_generators = [rg]
        last_rg = rg
        for _ in range(0, threads-1):
            new_rg = last_rg.jumped()
            self._random_generators.append(new_rg)
            last_rg = new_rg

        self.n = n
        self.executor = concurrent.futures.ThreadPoolExecutor(threads)
        self.values = np.empty(n)
        self.step = np.ceil(n / threads).astype(np.int_)

    def fill(self):
        def _fill(random_state, out, first, last):
            random_state.standard_normal(out=out[first:last])

        futures = {}
        for i in range(self.threads):
            args = (_fill,
                    self._random_generators[i],
                    self.values,
                    i * self.step,
                    (i + 1) * self.step)
            futures[self.executor.submit(*args)] = i
        concurrent.futures.wait(futures)

    def __del__(self):
        self.executor.shutdown(False)

Многопоточный генератор случайных чисел может использоваться для заполнения массива. Атрибут values показывает нулевое значение до заполнения и случайное значение после.

In [2]: mrng = MultithreadedRNG(10000000, seed=0)
...: print(mrng.values[-1])
0.0

In [3]: mrng.fill()
    ...: print(mrng.values[-1])
3.296046120254392

Время, необходимое для генерации с использованием нескольких потоков, можно сравнить со временем, необходимым для генерации с использованием одного потока.

In [4]: print(mrng.threads)
    ...: %timeit mrng.fill()

4
32.8 ms ± 2.71 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Вызов в одном потоке напрямую использует BitGenerator.

In [5]: values = np.empty(10000000)
    ...: rg = Generator(PCG64())
    ...: %timeit rg.standard_normal(out=values)

99.6 ms ± 222 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

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

In [6]: rg = Generator(PCG64())
    ...: %timeit rg.standard_normal(10000000)

125 ms ± 309 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)