# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
# All rights reserved. Use of this source code is governed by
# a BSD-style license that can be found in the LICENSE file.
import numpy as np
from .utils import Logger, Environment, AlignedU64, AlignedF64
from .dist import distribute_uniform
from .timing import function_timer
from ._libtoast import (
rng_dist_uint64,
rng_dist_uniform_01,
rng_dist_uniform_11,
rng_dist_normal,
rng_multi_dist_uint64,
rng_multi_dist_uniform_01,
rng_multi_dist_uniform_11,
rng_multi_dist_normal,
)
[docs]@function_timer
def random(samples, key=(0, 0), counter=(0, 0), sampler="gaussian", threads=False):
"""Generate random samples from a distribution for one stream.
This returns values from a single stream drawn from the specified
distribution. The starting state is specified by the two key values and
the two counter values. The second value of the "counter" is used to
represent the sample index. If the serial option is enabled, only a
single thread will be used. Otherwise the stream generation is divided
equally between OpenMP threads.
Args:
samples (int): The number of samples to return.
key (tuple): Two uint64 values which (along with the counter) define
the starting state of the generator.
counter (tuple): Two uint64 values which (along with the key) define
the starting state of the generator.
sampler (string): The distribution to sample from. Allowed values are
"gaussian", "uniform_01", "uniform_m11", and "uniform_uint64".
threads (bool): If True, use OpenMP threads to generate the stream
in parallel. NOTE: this may actually run slower for short streams
and many threads.
Returns:
(Aligned array): The random values of appropriate type for the sampler.
"""
env = Environment.get()
nthread = env.max_threads()
log = Logger.get()
ret = None
if (not threads) or (samples < nthread):
# Run serially
if sampler == "gaussian":
ret = AlignedF64(samples)
rng_dist_normal(key[0], key[1], counter[0], counter[1], ret)
elif sampler == "uniform_01":
ret = AlignedF64(samples)
rng_dist_uniform_01(key[0], key[1], counter[0], counter[1], ret)
elif sampler == "uniform_m11":
ret = AlignedF64(samples)
rng_dist_uniform_11(key[0], key[1], counter[0], counter[1], ret)
elif sampler == "uniform_uint64":
ret = AlignedU64(samples)
rng_dist_uint64(key[0], key[1], counter[0], counter[1], ret)
else:
msg = "Undefined sampler. Choose among: gaussian, uniform_01,\
uniform_m11, uniform_uint64"
log.error(msg)
raise ValueError(msg)
else:
# We are using threads, divide the samples up.
dst = distribute_uniform(samples, nthread)
k1 = AlignedU64(nthread)
k2 = AlignedU64(nthread)
c1 = AlignedU64(nthread)
c2 = AlignedU64(nthread)
k1[:] = np.array([key[0] for x in dst], dtype=np.uint64)
k2[:] = np.array([key[1] for x in dst], dtype=np.uint64)
c1[:] = np.array([counter[0] for x in dst], dtype=np.uint64)
c2[:] = np.array([counter[1] + x[0] for x in dst], dtype=np.uint64)
lengths = [x[1] for x in dst]
if sampler == "gaussian":
chunks = rng_multi_dist_normal(k1, k2, c1, c2, lengths)
ret = AlignedF64(samples)
for t in range(nthread):
ret[dst[t][0] : dst[t][0] + dst[t][1]] = chunks[t]
elif sampler == "uniform_01":
chunks = rng_multi_dist_uniform_01(k1, k2, c1, c2, lengths)
ret = AlignedF64(samples)
for t in range(nthread):
ret[dst[t][0] : dst[t][0] + dst[t][1]] = chunks[t]
elif sampler == "uniform_m11":
chunks = rng_multi_dist_uniform_11(k1, k2, c1, c2, lengths)
ret = AlignedF64(samples)
for t in range(nthread):
ret[dst[t][0] : dst[t][0] + dst[t][1]] = chunks[t]
elif sampler == "uniform_uint64":
chunks = rng_multi_dist_uint64(k1, k2, c1, c2, lengths)
ret = AlignedU64(samples)
for t in range(nthread):
ret[dst[t][0] : dst[t][0] + dst[t][1]] = chunks[t]
else:
msg = "Undefined sampler. Choose among: gaussian, uniform_01,\
uniform_m11, uniform_uint64"
log.error(msg)
raise ValueError(msg)
return ret
[docs]@function_timer
def random_multi(samples, keys, counters, sampler="gaussian"):
"""Generate random samples from multiple streams.
Given multiple streams, each specified by a pair of key values and a pair
of counter values, generate some number of samples from each. The number
of samples is specified independently for each stream. The generation of
the streams is run with multiple threads.
NOTE: if you just want threaded generation of a single stream, use the
"threads" option to the random() function.
Args:
samples (list): The number of samples to return for each stream
keys (list): A tuple of integer values for each stream, which (along
with the counter) define the starting state of the generator for
that stream.
counters (list): A tuple of integer values for each stream, which
(along with the key) define the starting state of the generator
for that stream.
sampler (string): The distribution to sample from. Allowed values are
"gaussian", "uniform_01", "uniform_m11", and "uniform_uint64".
Returns:
(list): The random samples for each stream.
"""
log = Logger.get()
k1 = AlignedU64(len(keys))
k2 = AlignedU64(len(keys))
c1 = AlignedU64(len(counters))
c2 = AlignedU64(len(counters))
k1[:] = np.array([x[0] for x in keys], dtype=np.uint64)
k2[:] = np.array([x[1] for x in keys], dtype=np.uint64)
c1[:] = np.array([x[0] for x in counters], dtype=np.uint64)
c2[:] = np.array([x[1] for x in counters], dtype=np.uint64)
ret = None
if sampler == "gaussian":
ret = rng_multi_dist_normal(k1, k2, c1, c2, samples)
elif sampler == "uniform_01":
ret = rng_multi_dist_uniform_01(k1, k2, c1, c2, samples)
elif sampler == "uniform_m11":
ret = rng_multi_dist_uniform_11(k1, k2, c1, c2, samples)
elif sampler == "uniform_uint64":
ret = rng_multi_dist_uint64(k1, k2, c1, c2, samples)
else:
msg = "Undefined sampler. Choose among: gaussian, uniform_01,\
uniform_m11, uniform_uint64"
log.error(msg)
raise ValueError(msg)
return ret