Source code for xorbits._mars.tensor.random.wald

# Copyright 2022-2023 XProbe Inc.
# derived from copyright 1999-2021 Alibaba Group Holding Ltd.
#
# 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 ... import opcodes as OperandDef
from ...serialization.serializables import AnyField
from ..utils import gen_random_seeds
from .core import TensorDistribution, TensorRandomOperandMixin, handle_array


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)