# Source code for mars.tensor.random.wald

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2021 Alibaba Group Holding Ltd.
#
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#
# Unless required by applicable law or agreed to in writing, software
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and

import numpy as np

from ... import opcodes as OperandDef
from ...serialization.serializables import AnyField
from ..utils import gen_random_seeds
from .core import TensorRandomOperandMixin, handle_array, TensorDistribution

class TensorWald(TensorDistribution, TensorRandomOperandMixin):
_input_fields_ = ["mean", "scale"]
_op_type_ = OperandDef.RAND_WALD

_fields_ = "mean", "scale", "size"
mean = AnyField("mean")
scale = AnyField("scale")
_func_name = "wald"

def __call__(self, mean, scale, chunk_size=None):
return self.new_tensor([mean, scale], None, raw_chunk_size=chunk_size)

[docs]def wald(random_state, mean, scale, size=None, chunk_size=None, gpu=None, dtype=None):
r"""
Draw samples from a Wald, or inverse Gaussian, distribution.

As the scale approaches infinity, the distribution becomes more like a
Gaussian. Some references claim that the Wald is an inverse Gaussian
with mean equal to 1, but this is by no means universal.

The inverse Gaussian distribution was first studied in relationship to
Brownian motion. In 1956 M.C.K. Tweedie used the name inverse Gaussian
because there is an inverse relationship between the time to cover a
unit distance and distance covered in unit time.

Parameters
----------
mean : float or array_like of floats
Distribution mean, should be > 0.
scale : float or array_like of floats
Scale parameter, should be >= 0.
size : int or tuple of ints, optional
Output shape.  If the given shape is, e.g., (m, n, k), then
m * n * k samples are drawn.  If size is None (default),
a single value is returned if mean and scale are both scalars.
Otherwise, np.broadcast(mean, scale).size samples are drawn.
chunk_size : int or tuple of int or tuple of ints, optional
Desired chunk size on each dimension
gpu : bool, optional
Allocate the tensor on GPU if True, False as default
dtype : data-type, optional
Data-type of the returned tensor.

Returns
-------
out : Tensor or scalar
Drawn samples from the parameterized Wald distribution.

Notes
-----
The probability density function for the Wald distribution is

.. math:: P(x;mean,scale) = \sqrt{\frac{scale}{2\pi x^3}}e^
\frac{-scale(x-mean)^2}{2\cdotp mean^2x}

As noted above the inverse Gaussian distribution first arise
from attempts to model Brownian motion. It is also a
competitor to the Weibull for use in reliability modeling and
modeling stock returns and interest rate processes.

References
----------
.. [1] Brighton Webs Ltd., Wald Distribution,
http://www.brighton-webs.co.uk/distributions/wald.asp
.. [2] Chhikara, Raj S., and Folks, J. Leroy, "The Inverse Gaussian
Distribution: Theory : Methodology, and Applications", CRC Press,
1988.
.. [3] Wikipedia, "Wald distribution"
http://en.wikipedia.org/wiki/Wald_distribution

Examples
--------
Draw values from the distribution and plot the histogram:

>>> import matplotlib.pyplot as plt
>>> import mars.tensor as mt
>>> h = plt.hist(mt.random.wald(3, 2, 100000).execute(), bins=200, normed=True)
>>> plt.show()
"""
if dtype is None:
dtype = (
np.random.RandomState()
.wald(handle_array(mean), handle_array(scale), size=(0,))
.dtype
)
size = random_state._handle_size(size)
seed = gen_random_seeds(1, random_state.to_numpy())[0]
op = TensorWald(size=size, seed=seed, gpu=gpu, dtype=dtype)
return op(mean, scale, chunk_size=chunk_size)