dvcompute_rand 2.0.0

// Copyright (c) 2020-2022  David Sorokin <davsor@mail.ru>, based in Yoshkar-Ola, Russia
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

#[cfg(feature="dist_mode")]
use std::mem;

#[cfg(feature="dist_mode")]
use libc::*;

use rand::prelude::*;

/// Specifies a random number generator from 0 to 1.
pub type GeneratorBox = Box<dyn FnMut() -> f64>;

/// It represents a raw trait object.
#[cfg(feature="dist_mode")]
#[repr(C)]
#[derive(Copy, Clone)]
struct GeneratorTraitObject {

    field1: *mut c_void,
    field2: *mut c_void
}

/// A C-friendly representaton of the generator.
#[cfg(feature="dist_mode")]
#[repr(C)]
pub struct GeneratorRepr {

    /// Delete the object.
    delete: unsafe extern "C" fn(obj: *mut GeneratorTraitObject),

    /// The callback.
    callback: unsafe extern "C" fn(obj: *mut GeneratorTraitObject) -> f64,

    /// The trait object.
    trait_object: GeneratorTraitObject
}

#[cfg(feature="dist_mode")]
impl Drop for GeneratorRepr {

    fn drop(&mut self) {
        unsafe {
            (self.delete)(&mut self.trait_object);
        }
    }
}

#[cfg(feature="dist_mode")]
impl GeneratorRepr {

    /// Convert to a C-friendly representation.
    #[inline]
    pub fn into_repr(gen: GeneratorBox) -> GeneratorRepr {
        unsafe {
            GeneratorRepr {
                delete: delete_generator_repr,
                callback: call_generator_repr,
                trait_object: mem::transmute(gen)
            }
        }
    }

    /// Call the representation.
    #[inline]
    pub fn call(&mut self) -> f64 {
        unsafe {
            (self.callback)(&mut self.trait_object)
        }
    }
}

/// Call the generator representation.
#[cfg(feature="dist_mode")]
unsafe extern "C" fn call_generator_repr(gen: *mut GeneratorTraitObject) -> f64 {
    let mut gen: GeneratorBox = mem::transmute(*gen);
    let x = gen();
    mem::forget(gen);
    x
}

/// Delete the generator representation.
#[cfg(feature="dist_mode")]
unsafe extern "C" fn delete_generator_repr(gen: *mut GeneratorTraitObject) {
    let _: GeneratorBox = mem::transmute(*gen);
}

/// A thin wrapper for representaton of the generator.
#[cfg(any(feature="seq_mode", feature="wasm_mode"))]
pub struct GeneratorRepr {

    /// The generator itself.
    gen: GeneratorBox
}

#[cfg(any(feature="seq_mode", feature="wasm_mode"))]
impl GeneratorRepr {

    /// Convert to the representation.
    #[inline]
    pub fn into_repr(gen: GeneratorBox) -> GeneratorRepr {
        GeneratorRepr { gen: gen }
    }

    /// Call the representation.
    #[inline]
    pub fn call(&mut self) -> f64 {
        (self.gen)()
    }
}

/// Create an uniform random number generator with minimum value 0 and maximum 1.
pub fn uniform_generator() -> GeneratorRepr {
    let mut rng = rand::thread_rng();
    let gen = Box::new(move || {
        rng.gen()
    });
    GeneratorRepr::into_repr(gen)
}

/// Create an uniform random number generator with the specified seed.
pub fn uniform_generator_with_seed(seed: [u8; 32]) -> GeneratorRepr {
    let mut rng = StdRng::from_seed(seed);
    let gen = Box::new(move || {
        rng.gen()
    });
    GeneratorRepr::into_repr(gen)
}

/// Return the normal number generator with mean 0 and deviation 1.
pub fn normal_generator(gen: GeneratorRepr) -> GeneratorRepr {
    let mut gen  = gen;
    let mut next = 0.0;
    let mut flag = false;
    let gen = Box::new(move || {
        if flag {
            flag = false;
            next
        } else {
            let mut xi1 = 0.0;
            let mut xi2 = 0.0;
            let mut psi = 0.0;
            loop {
                if psi >= 1.0 || psi == 0.0 {
                    let g1 = gen.call();
                    let g2 = gen.call();
                    xi1 = 2.0 * g1 - 1.0;
                    xi2 = 2.0 * g2 - 1.0;
                    psi = xi1 * xi1 + xi2 * xi2;
                } else {
                    psi = (- 2.0 * psi.ln() / psi).sqrt();
                    break;
                }
            }
            flag = true;
            next  = xi2 * psi;
            xi1 * psi
        }
    });
    GeneratorRepr::into_repr(gen)
}

/// Generate an uniform random number with the specified minimum and maximum.
#[inline]
pub fn generate_uniform(gen: &mut GeneratorRepr, min: f64, max: f64) -> f64 {
    let x = gen.call();
    min + x * (max - min)
}

/// Generate an integer uniform random number with the specified minimum and maximum.
pub fn generate_int_uniform(gen: &mut GeneratorRepr, min: isize, max: isize) -> isize {
    let x = gen.call();
    let min2 = (min as f64) - 0.5;
    let max2 = (max as f64) + 0.5;
    let z = (min2 + x * (max2 - min2)).round() as isize;
    if z < min {
        min
    } else if z > max {
        max
    } else {
        z
    }
}

/// Generate the triangular random number by the specified minimum, median and maximum.
pub fn generate_triangular(gen: &mut GeneratorRepr, min: f64, median: f64, max: f64) -> f64 {
    let x = gen.call();
    if x <= (median - min) / (max - min) {
        min + ((median - min) * (max - min) * x).sqrt()
    } else {
        max - ((max - median) * (max - min) * (1.0 - x)).sqrt()
    }
}

/// Generate a normal random number by the specified generator, mean and deviation.
#[inline]
pub fn generate_normal(gen: &mut GeneratorRepr, mu: f64, nu: f64) -> f64 {
    let x = gen.call();
    mu + nu * x
}

/// Generate the lognormal random number derived from a normal distribution with
/// the specified generator, mean and deviation.
#[inline]
pub fn generate_log_normal(gen: &mut GeneratorRepr, mu: f64, nu: f64) -> f64 {
    let x = gen.call();
    (mu + nu * x).exp()
}

/// Return the exponential random number with the specified mean.
#[inline]
pub fn generate_exponential(gen: &mut GeneratorRepr, mu: f64) -> f64 {
    let x = gen.call();
    - x.ln() * mu
}

/// Return the Erlang random number.
pub fn generate_erlang(gen: &mut GeneratorRepr, scale: f64, shape: isize) -> f64 {
    assert!(shape >= 0, "Negative shape");
    let mut y = 1.0;
    for _ in 0..shape {
        let x = gen.call();
        y *= x;
    }
    - y.ln() * scale
}

/// Generate the Poisson random number with the specified mean.
pub fn generate_poisson(gen: &mut GeneratorRepr, mu: f64) -> isize {
    let mut prob = gen.call();
    let mut prod = (- mu).exp();
    let mut acc  = 0;
    loop {
        if prob <= prod {
            return acc;
        } else {
            let prob2 = prob - prod;
            let prod2 = prod * mu / ((acc + 1) as f64);
            let acc2  = acc + 1;
            prob = prob2;
            prod = prod2;
            acc  = acc2;
        }
    }
}

/// Generate a binomial random number with the specified probability and number of trials.
pub fn generate_binomial(gen: &mut GeneratorRepr, prob: f64, trials: isize) -> isize {
    assert!(trials >= 0, "Negative number of trials");
    let mut acc = 0;
    for _ in 0..trials {
        let x = gen.call();
        if x <= prob {
            acc += 1;
        }
    }
    acc
}

/// Generate a random number from the Gamma distribution using Marsaglia and Tsang method.
#[allow(unconditional_recursion)]
pub fn generate_gamma(genn: &mut GeneratorRepr, genu: &mut GeneratorRepr, kappa: f64, theta: f64) -> f64 {
    if kappa <= 0.0 {
        panic!("The shape parameter (kappa) must be positive")
    } else if kappa > 1.0 {
        let d = kappa - 1.0 / 3.0;
        let c = 1.0 / (9.0 * d).sqrt();
        loop {
            let z = genn.call();
            if z > - (1.0 / c) {
                let v = (1.0 + c * z).powi(3);
                let u = genu.call();
                if u.ln() <= 0.5 * z * z + d - d * v + d * v.ln() {
                    return d * v * theta
                }
            }
        }
    } else {
        let x = generate_gamma(genn, genu, 1.0 + kappa, theta);
        let u = genu.call();
        x * u.powf(1.0 / kappa)
    }
}

/// Generate a random number from the Beta distribution.
#[inline]
pub fn generate_beta(genn: &mut GeneratorRepr, genu: &mut GeneratorRepr, alpha: f64, beta: f64) -> f64 {
    let g1 = generate_gamma(genn, genu, alpha, 1.0);
    let g2 = generate_gamma(genn, genu, beta, 1.0);
    g1 / (g1 + g2)
}

/// Generate a random number from the Weibull distribution.
#[inline]
pub fn generate_weibull(gen: &mut GeneratorRepr, alpha: f64, beta: f64) -> f64 {
    let x = gen.call();
    beta * (- x.ln()).powf(1.0 / alpha)
}

/// Generate a random value from the specified discrete distribution.
pub fn generate_discrete<'a, T>(gen: &mut GeneratorRepr, dpdf: &'a [(T, f64)]) -> &'a T {
    assert!(dpdf.len() > 0, "Empty discrete probability density");
    let x = gen.call();
    let mut acc  = 0.0;
    let mut dpdf = dpdf;
    loop {
        if dpdf.len() == 1 {
            let a = &dpdf[0].0;
            return a;
        } else {
            let p = dpdf[0].1;
            if x <= acc + p {
                let a = &dpdf[0].0;
                return a;
            } else {
                acc += p;
                dpdf = &dpdf[1..];
            }
        }
    }
}