# Copyright (C) 2020 NumS Development Team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
from nums.core.array.blockarray import BlockArray, Block
from nums.core.compute.compute_manager import ComputeManager
from nums.core.grid.grid import ArrayGrid
[docs]class NumsRandomState:
def __init__(self, cm: ComputeManager, seed):
self._cm = cm
self._rng = self._cm.get_rng(seed)
[docs] def seed(self, seed=None):
# New RNG based on given seed.
self._rng = self._cm.get_rng(seed)
[docs] def numpy(self):
# pylint: disable = import-outside-toplevel
from nums.core.compute.numpy_compute import block_rng
return block_rng(*self._rng.new_block_rng_params())
[docs] def random(self, shape=None, block_shape=None, dtype=None):
if dtype is None:
dtype = np.float64
assert isinstance(dtype, type)
return self._sample_basic("random", shape, block_shape, dtype, (dtype,))
[docs] def integers(
self, low, high=None, shape=None, block_shape=None, dtype=None, endpoint=False
):
if dtype is None:
dtype = np.int64
assert isinstance(dtype, type)
return self._sample_basic(
"integers", shape, block_shape, dtype, (low, high, dtype, endpoint)
)
[docs] def normal(self, loc=0.0, scale=1.0, shape=None, block_shape=None, dtype=None):
return self._sample_basic("normal", shape, block_shape, dtype, (loc, scale))
[docs] def beta(self, a, b, shape=None, block_shape=None, dtype=None):
return self._sample_basic("beta", shape, block_shape, dtype, (a, b))
[docs] def binomial(self, n, p, shape=None, block_shape=None, dtype=None):
return self._sample_basic("binomial", shape, block_shape, dtype, (n, p))
[docs] def chisquare(self, df, shape=None, block_shape=None, dtype=None):
return self._sample_basic("chisquare", shape, block_shape, dtype, (df,))
[docs] def exponential(self, scale=1.0, shape=None, block_shape=None, dtype=None):
return self._sample_basic("exponential", shape, block_shape, dtype, (scale,))
[docs] def f(self, dfnum, dfden, shape=None, block_shape=None, dtype=None):
return self._sample_basic("f", shape, block_shape, dtype, (dfnum, dfden))
[docs] def gamma(self, k, theta, shape=None, block_shape=None, dtype=None):
return self._sample_basic("gamma", shape, block_shape, dtype, (k, theta))
[docs] def geometric(self, p, shape=None, block_shape=None, dtype=None):
return self._sample_basic("geometric", shape, block_shape, dtype, (p,))
[docs] def gumbel(self, loc=0.0, scale=1.0, shape=None, block_shape=None, dtype=None):
return self._sample_basic("gumbel", shape, block_shape, dtype, (loc, scale))
[docs] def hypergeometric(
self, ngood, nbad, nsample, shape=None, block_shape=None, dtype=None
):
return self._sample_basic(
"hypergeometric", shape, block_shape, dtype, (ngood, nbad, nsample)
)
[docs] def laplace(self, loc=0.0, scale=1.0, shape=None, block_shape=None, dtype=None):
return self._sample_basic("laplace", shape, block_shape, dtype, (loc, scale))
[docs] def logistic(self, loc=0.0, scale=1.0, shape=None, block_shape=None, dtype=None):
return self._sample_basic("logistic", shape, block_shape, dtype, (loc, scale))
[docs] def lognormal(self, mean=0.0, sigma=1.0, shape=None, block_shape=None, dtype=None):
return self._sample_basic("lognormal", shape, block_shape, dtype, (mean, sigma))
[docs] def logseries(self, p, shape=None, block_shape=None, dtype=None):
return self._sample_basic("logseries", shape, block_shape, dtype, (p,))
[docs] def negative_binomial(self, n, p, shape=None, block_shape=None, dtype=None):
return self._sample_basic(
"negative_binomial", shape, block_shape, dtype, (n, p)
)
[docs] def noncentral_chisquare(self, df, nonc, shape=None, block_shape=None, dtype=None):
return self._sample_basic(
"noncentral_chisquare", shape, block_shape, dtype, (df, nonc)
)
[docs] def noncentral_f(
self, dfnum, dfden, nonc, shape=None, block_shape=None, dtype=None
):
return self._sample_basic(
"noncentral_f", shape, block_shape, dtype, (dfnum, dfden, nonc)
)
[docs] def pareto(self, a, shape=None, block_shape=None, dtype=None):
return self._sample_basic("pareto", shape, block_shape, dtype, (a,))
[docs] def poisson(self, lam=1.0, shape=None, block_shape=None, dtype=None):
return self._sample_basic("poisson", shape, block_shape, dtype, (lam,))
[docs] def power(self, a, shape=None, block_shape=None, dtype=None):
return self._sample_basic("power", shape, block_shape, dtype, (a,))
[docs] def rayleigh(self, scale=1.0, shape=None, block_shape=None, dtype=None):
return self._sample_basic("rayleigh", shape, block_shape, dtype, (scale,))
[docs] def standard_cauchy(self, shape=None, block_shape=None, dtype=None):
return self._sample_basic("standard_cauchy", shape, block_shape, dtype, ())
[docs] def standard_t(self, df, shape=None, block_shape=None, dtype=None):
return self._sample_basic("standard_t", shape, block_shape, dtype, (df,))
# TODO Restrict params to scalars.
[docs] def triangular(self, left, mode, right, shape=None, block_shape=None, dtype=None):
return self._sample_basic(
"triangular", shape, block_shape, dtype, (left, mode, right)
)
[docs] def vonmises(self, mu, kappa, shape=None, block_shape=None, dtype=None):
return self._sample_basic("vonmises", shape, block_shape, dtype, (mu, kappa))
[docs] def wald(self, mean, scale, shape=None, block_shape=None, dtype=None):
return self._sample_basic("wald", shape, block_shape, dtype, (mean, scale))
[docs] def weibull(self, a, shape=None, block_shape=None, dtype=None):
return self._sample_basic("weibull", shape, block_shape, dtype, (a,))
[docs] def zipf(self, a, shape=None, block_shape=None, dtype=None):
return self._sample_basic("zipf", shape, block_shape, dtype, (a,))
# TODO (hme): Add multivariate samplers.
def _sample_basic(
self, rfunc_name, shape, block_shape, dtype, rfunc_args
) -> BlockArray:
if shape is None:
assert block_shape is None
shape = ()
block_shape = ()
else:
assert block_shape is not None
if dtype is None:
dtype = np.float64
assert isinstance(dtype, type)
grid: ArrayGrid = ArrayGrid(shape, block_shape, dtype=dtype.__name__)
ba: BlockArray = BlockArray(grid, self._cm)
for grid_entry in ba.grid.get_entry_iterator():
rng_params = list(self._rng.new_block_rng_params())
# Size and dtype to begin with.
this_block_shape = grid.get_block_shape(grid_entry)
size = int(np.product(this_block_shape))
# Inconsistent param orderings.
if rfunc_name == "random":
rfunc_args_final = tuple([size] + list(rfunc_args))
elif rfunc_name == "integers":
# rfunc_args == (low, high, dtype, endpoint)
rfunc_args_final = tuple(
list(rfunc_args[:2]) + [size] + list(rfunc_args[2:])
)
else:
rfunc_args_final = tuple(list(rfunc_args) + [size])
block: Block = ba.blocks[grid_entry]
block.oid = self._cm.random_block(
rng_params,
rfunc_name,
rfunc_args_final,
this_block_shape,
dtype,
syskwargs={"grid_entry": grid_entry, "grid_shape": grid.grid_shape},
)
return ba
[docs] def permutation(self, size, block_size):
shape = (size,)
block_shape = (block_size,)
grid: ArrayGrid = ArrayGrid(
shape=shape, block_shape=shape, dtype=np.int64.__name__
)
ba = BlockArray(grid, self._cm)
for grid_entry in ba.grid.get_entry_iterator():
rng_params = list(self._rng.new_block_rng_params())
block: Block = ba.blocks[grid_entry]
block.oid = self._cm.permutation(
rng_params,
size,
syskwargs={"grid_entry": grid_entry, "grid_shape": grid.grid_shape},
)
return ba.reshape(block_shape=block_shape)