use crate::{FloatExt, StableError, XResult, random::STABLE_PAR_THRESHOLD};
use rand::{Rng, prelude::*};
use rand_distr::{Exp1, uniform::SampleUniform};
use rand_xoshiro::Xoshiro256PlusPlus;
use rayon::prelude::*;
#[derive(Debug, Clone, Copy)]
pub(crate) struct StableConstants<T: FloatExt = f64> {
inv_alpha: T,
one_minus_alpha_div_alpha: T,
b: T,
s: T,
}
impl<T: FloatExt> StableConstants<T> {
#[inline]
pub fn new(alpha: T, beta: T) -> Self {
let inv_alpha = T::one() / alpha;
let one_minus_alpha_div_alpha = (T::one() - alpha) * inv_alpha;
let tmp = beta * (alpha * T::FRAC_PI_2()).tan();
let b = tmp.atan() * inv_alpha;
let s = (T::one() + tmp * tmp).powf(T::from(0.5).unwrap() * inv_alpha);
Self {
inv_alpha,
b,
s,
one_minus_alpha_div_alpha,
}
}
}
#[derive(Debug, Clone)]
pub struct StandardStable<T: FloatExt = f64> {
alpha: T,
beta: T,
}
impl<T: FloatExt> StandardStable<T> {
pub fn new(alpha: T, beta: T) -> XResult<Self> {
if alpha <= T::zero() || alpha > T::from(2).unwrap() || alpha.is_nan() {
return Err(StableError::InvalidIndex.into());
}
if !(-T::one()..=T::one()).contains(&beta) {
return Err(StableError::InvalidSkewness.into());
}
Ok(Self { alpha, beta })
}
pub fn get_index(&self) -> T {
self.alpha
}
pub fn get_skewness(&self) -> T {
self.beta
}
pub fn samples(&self, n: usize) -> XResult<Vec<T>>
where
T: SampleUniform,
Exp1: Distribution<T>,
{
standard_rands(self.alpha, self.beta, n)
}
}
pub(crate) fn sample_standard_alpha<T: FloatExt + SampleUniform, R: Rng + ?Sized>(
alpha: T,
beta: T,
rng: &mut R,
) -> T
where
Exp1: Distribution<T>,
{
let constants = StableConstants::new(alpha, beta);
sample_standard_alpha_with_constants(&constants, alpha, rng)
}
#[inline]
pub(crate) fn sample_standard_alpha_with_constants<T, R: Rng + ?Sized>(
c: &StableConstants<T>,
alpha: T,
rng: &mut R,
) -> T
where
T: FloatExt + SampleUniform,
Exp1: Distribution<T>,
{
let v = rng.random_range(-T::FRAC_PI_2()..T::FRAC_PI_2());
let w: T = rng.sample(Exp1);
let v_plus_b = v + c.b;
let cos_v = v.cos();
let c1 = alpha * v_plus_b.sin() / cos_v.powf(c.inv_alpha);
let c2 = ((v - alpha * v_plus_b).cos() / w).powf(c.one_minus_alpha_div_alpha);
c.s * c1 * c2
}
#[inline]
pub(crate) fn sample_standard_alpha_one<T, R: Rng + ?Sized>(_alpha: T, beta: T, rng: &mut R) -> T
where
T: FloatExt + SampleUniform,
Exp1: Distribution<T>,
{
let v = rng.random_range(-T::FRAC_PI_2()..T::FRAC_PI_2());
let w: T = rng.sample(Exp1);
let half_pi_plus_beta_v = T::FRAC_PI_2() + beta * v;
let c1 = half_pi_plus_beta_v * v.tan();
let c2 = ((T::FRAC_PI_2() * w * v.cos()) / half_pi_plus_beta_v).ln() * beta;
(c1 - c2) * T::FRAC_2_PI()
}
impl<T: FloatExt + SampleUniform> Distribution<T> for StandardStable<T>
where
Exp1: Distribution<T>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T {
if (self.alpha - T::one()).abs() > T::epsilon() {
sample_standard_alpha(self.alpha, self.beta, rng)
} else {
sample_standard_alpha_one(self.alpha, self.beta, rng)
}
}
}
#[derive(Debug, Clone, Copy)]
pub struct Stable<T: FloatExt = f64> {
alpha: T,
beta: T,
sigma: T,
mu: T,
}
impl<T: FloatExt> From<&Stable<T>> for StandardStable<T> {
fn from(stable: &Stable<T>) -> Self {
StandardStable::new(stable.alpha, stable.beta).unwrap()
}
}
impl<T: FloatExt> Stable<T> {
pub fn new(alpha: T, beta: T, sigma: T, mu: T) -> XResult<Self> {
if alpha <= T::zero() || alpha > T::from(2).unwrap() || alpha.is_nan() {
return Err(StableError::InvalidIndex.into());
}
if !(-T::one()..=T::one()).contains(&beta) {
return Err(StableError::InvalidSkewness.into());
}
if sigma <= T::zero() || sigma.is_nan() {
return Err(StableError::InvalidScale.into());
}
if mu.is_nan() {
return Err(StableError::InvalidLocation.into());
}
Ok(Self {
alpha,
beta,
sigma,
mu,
})
}
pub fn get_index(&self) -> T {
self.alpha
}
pub fn get_skewness(&self) -> T {
self.beta
}
pub fn get_scale(&self) -> T {
self.sigma
}
pub fn get_location(&self) -> T {
self.mu
}
pub fn samples(&self, n: usize) -> XResult<Vec<T>>
where
T: SampleUniform,
Exp1: Distribution<T>,
{
rands(self.alpha, self.beta, self.sigma, self.mu, n)
}
}
impl<T: FloatExt + SampleUniform> Distribution<T> for Stable<T>
where
Exp1: Distribution<T>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T {
let standard = StandardStable::from(self);
let r = rng.sample(standard);
if self.alpha != T::one() {
self.sigma * r + self.mu
} else {
self.sigma * r
+ self.mu
+ T::from(2).unwrap() * self.beta * self.sigma * self.sigma.ln() / T::PI()
}
}
}
#[derive(Debug, Clone, Copy)]
pub struct StandardSkewStable<T: FloatExt = f64>(pub T);
impl<T: FloatExt> StandardSkewStable<T> {
pub fn new(alpha: T) -> XResult<Self> {
if alpha <= T::zero() || alpha >= T::one() || alpha.is_nan() {
return Err(StableError::InvalidSkewIndex.into());
}
Ok(Self(alpha))
}
pub fn get_index(&self) -> T {
self.0
}
pub fn samples(&self, n: usize) -> XResult<Vec<T>>
where
T: SampleUniform,
Exp1: Distribution<T>,
{
skew_rands(self.0, n)
}
}
impl<T: FloatExt + SampleUniform> Distribution<T> for StandardSkewStable<T>
where
Exp1: Distribution<T>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T {
let alpha = self.0;
if alpha <= T::zero() || alpha >= T::one() || alpha.is_nan() {
panic!("Invalid skew index");
}
sample_standard_alpha(self.0, T::one(), rng)
}
}
#[derive(Debug, Clone, Copy)]
pub struct SymmetricStandardStable<T: FloatExt = f64>(pub T);
impl<T: FloatExt> SymmetricStandardStable<T> {
pub fn new(alpha: T) -> XResult<Self> {
if alpha <= T::zero() || alpha >= T::from(2).unwrap() || alpha.is_nan() {
return Err(StableError::InvalidSkewIndex.into());
}
Ok(Self(alpha))
}
pub fn get_index(&self) -> T {
self.0
}
pub fn samples(&self, n: usize) -> XResult<Vec<T>>
where
T: SampleUniform,
Exp1: Distribution<T>,
{
sym_standard_rands(self.0, n)
}
}
impl<T: FloatExt + SampleUniform> Distribution<T> for SymmetricStandardStable<T>
where
Exp1: Distribution<T>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T {
let alpha = self.0;
if alpha <= T::zero() || alpha > T::from(2).unwrap() || alpha.is_nan() {
panic!("Invalid stability index");
}
if (alpha - T::one()).abs() > T::epsilon() {
let inv_alpha = T::one() / alpha;
let one_minus_alpha_div_alpha = (T::one() - alpha) * inv_alpha;
sample_sym_standard_alpha_with_constants(
inv_alpha,
one_minus_alpha_div_alpha,
alpha,
rng,
)
} else {
sample_sym_standard_alpha_one(rng)
}
}
}
pub fn standard_rand<T: FloatExt + SampleUniform>(alpha: T, beta: T) -> XResult<T>
where
Exp1: Distribution<T>,
{
let standard = StandardStable::new(alpha, beta)?;
let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng());
Ok(rng.sample(standard))
}
pub fn standard_rands<T>(alpha: T, beta: T, n: usize) -> XResult<Vec<T>>
where
T: FloatExt + SampleUniform,
Exp1: Distribution<T>,
{
if alpha <= T::zero() || alpha > T::from(2).unwrap() || alpha.is_nan() {
return Err(StableError::InvalidIndex.into());
}
if !(-T::one()..=T::one()).contains(&beta) {
return Err(StableError::InvalidSkewness.into());
}
if (alpha - T::one()).abs() < T::epsilon() {
if n <= STABLE_PAR_THRESHOLD {
let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng());
Ok((0..n)
.map(|_| sample_standard_alpha_one(alpha, beta, &mut rng))
.collect())
} else {
Ok((0..n)
.into_par_iter()
.map_init(
|| Xoshiro256PlusPlus::from_rng(&mut rand::rng()),
|r, _| sample_standard_alpha_one(alpha, beta, r),
)
.collect())
}
} else {
let constants = StableConstants::new(alpha, beta);
if n <= STABLE_PAR_THRESHOLD {
let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng());
Ok((0..n)
.map(|_| sample_standard_alpha_with_constants(&constants, alpha, &mut rng))
.collect())
} else {
Ok((0..n)
.into_par_iter()
.map_init(
|| Xoshiro256PlusPlus::from_rng(&mut rand::rng()),
|r, _| sample_standard_alpha_with_constants(&constants, alpha, r),
)
.collect())
}
}
}
pub fn rand<T: FloatExt + SampleUniform>(alpha: T, beta: T, sigma: T, mu: T) -> XResult<T>
where
Exp1: Distribution<T>,
{
let levy = Stable::new(alpha, beta, sigma, mu)?;
let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng());
Ok(rng.sample(levy))
}
pub fn rands<T>(alpha: T, beta: T, sigma: T, mu: T, n: usize) -> XResult<Vec<T>>
where
T: FloatExt + SampleUniform,
Exp1: Distribution<T>,
{
let two = T::from(2).unwrap();
if alpha <= T::zero() || alpha > two || alpha.is_nan() {
return Err(StableError::InvalidIndex.into());
}
if !(-T::one()..=T::one()).contains(&beta) {
return Err(StableError::InvalidSkewness.into());
}
if sigma <= T::zero() || sigma.is_nan() {
return Err(StableError::InvalidScale.into());
}
if mu.is_nan() {
return Err(StableError::InvalidLocation.into());
}
if (alpha - T::one()).abs() < T::epsilon() {
let correction = two * beta * sigma * sigma.ln() / T::PI();
if n <= STABLE_PAR_THRESHOLD {
let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng());
Ok((0..n)
.map(|_| {
let std_sample = sample_standard_alpha_one(alpha, beta, &mut rng);
sigma * std_sample + mu + correction
})
.collect())
} else {
Ok((0..n)
.into_par_iter()
.map_init(
|| Xoshiro256PlusPlus::from_rng(&mut rand::rng()),
|r, _| {
let std_sample = sample_standard_alpha_one(alpha, beta, r);
sigma * std_sample + mu + correction
},
)
.collect())
}
} else {
let constants = StableConstants::new(alpha, beta);
if n <= STABLE_PAR_THRESHOLD {
let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng());
Ok((0..n)
.map(|_| {
let std_sample =
sample_standard_alpha_with_constants(&constants, alpha, &mut rng);
sigma * std_sample + mu
})
.collect())
} else {
Ok((0..n)
.into_par_iter()
.map_init(
|| Xoshiro256PlusPlus::from_rng(&mut rand::rng()),
|r, _| {
let std_sample = sample_standard_alpha_with_constants(&constants, alpha, r);
sigma * std_sample + mu
},
)
.collect())
}
}
}
pub fn skew_rand<T: FloatExt + SampleUniform>(alpha: T) -> XResult<T>
where
Exp1: Distribution<T>,
{
let skew = StandardSkewStable::new(alpha)?;
let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng());
Ok(rng.sample(skew))
}
pub fn skew_rands<T>(alpha: T, n: usize) -> XResult<Vec<T>>
where
T: FloatExt + SampleUniform,
Exp1: Distribution<T>,
{
if alpha <= T::zero() || alpha >= T::one() || alpha.is_nan() {
return Err(StableError::InvalidSkewIndex.into());
}
let constants = StableConstants::new(alpha, T::one());
if n <= STABLE_PAR_THRESHOLD {
let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng());
Ok((0..n)
.map(|_| sample_standard_alpha_with_constants(&constants, alpha, &mut rng))
.collect())
} else {
Ok((0..n)
.into_par_iter()
.map_init(
|| Xoshiro256PlusPlus::from_rng(&mut rand::rng()),
|r, _| sample_standard_alpha_with_constants(&constants, alpha, r),
)
.collect())
}
}
#[inline]
pub(crate) fn sample_sym_standard_alpha_with_constants<T, R: Rng + ?Sized>(
inv_alpha: T,
one_minus_alpha_div_alpha: T,
alpha: T,
rng: &mut R,
) -> T
where
T: FloatExt + SampleUniform,
Exp1: Distribution<T>,
{
let v = rng.random_range(-T::FRAC_PI_2()..T::FRAC_PI_2());
let w: T = rng.sample(Exp1);
let cos_v = v.cos();
let c1 = alpha * v.sin() / cos_v.powf(inv_alpha);
let c2 = ((v - alpha * v).cos() / w).powf(one_minus_alpha_div_alpha);
c1 * c2
}
#[inline]
pub(crate) fn sample_sym_standard_alpha_with_stable_constants<T, R: Rng + ?Sized>(
c: &StableConstants<T>,
alpha: T,
rng: &mut R,
) -> T
where
T: FloatExt + SampleUniform,
Exp1: Distribution<T>,
{
sample_sym_standard_alpha_with_constants(c.inv_alpha, c.one_minus_alpha_div_alpha, alpha, rng)
}
#[inline]
pub(crate) fn sample_sym_standard_alpha_one<T, R: Rng + ?Sized>(rng: &mut R) -> T
where
T: FloatExt + SampleUniform,
Exp1: Distribution<T>,
{
let v = rng.random_range(-T::FRAC_PI_2()..T::FRAC_PI_2());
v.tan()
}
pub fn sym_standard_rand<T: FloatExt + SampleUniform>(alpha: T) -> XResult<T>
where
Exp1: Distribution<T>,
{
let sym = SymmetricStandardStable::new(alpha)?;
let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng());
Ok(rng.sample(sym))
}
pub fn sym_standard_rands<T>(alpha: T, n: usize) -> XResult<Vec<T>>
where
T: FloatExt + SampleUniform,
Exp1: Distribution<T>,
{
if alpha <= T::zero() || alpha > T::from(2).unwrap() || alpha.is_nan() {
return Err(StableError::InvalidIndex.into());
}
if (alpha - T::one()).abs() < T::epsilon() {
if n <= STABLE_PAR_THRESHOLD {
let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng());
Ok((0..n)
.map(|_| sample_sym_standard_alpha_one(&mut rng))
.collect())
} else {
Ok((0..n)
.into_par_iter()
.map_init(
|| Xoshiro256PlusPlus::from_rng(&mut rand::rng()),
|r, _| sample_sym_standard_alpha_one(r),
)
.collect())
}
} else {
let inv_alpha = T::one() / alpha;
let one_minus_alpha_div_alpha = (T::one() - alpha) * inv_alpha;
if n <= STABLE_PAR_THRESHOLD {
let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng());
Ok((0..n)
.map(|_| {
sample_sym_standard_alpha_with_constants(
inv_alpha,
one_minus_alpha_div_alpha,
alpha,
&mut rng,
)
})
.collect())
} else {
Ok((0..n)
.into_par_iter()
.map_init(
|| Xoshiro256PlusPlus::from_rng(&mut rand::rng()),
|r, _| {
sample_sym_standard_alpha_with_constants(
inv_alpha,
one_minus_alpha_div_alpha,
alpha,
r,
)
},
)
.collect())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use num_traits::Float;
use rand::rng;
#[test]
fn test_sample_standard_alpha() {
let alpha = 0.7;
let beta = 1.0;
let mut rng = rng();
let standard = StandardStable::new(alpha, beta).unwrap();
let r = rng.sample(standard);
assert!(r.is_finite());
let standard = StandardStable::new(alpha as f32, beta as f32).unwrap();
let r = rng.sample(standard);
assert!(r.is_finite());
}
#[test]
fn test_sample_symmetric_standard_alpha() {
let alpha = 0.7;
let mut rng = rng();
let r = rng.sample(SymmetricStandardStable::new(alpha).unwrap());
assert!(r.is_finite());
}
#[test]
fn test_sample_symmetric_standard_alpha_rands() {
let alpha = 0.7;
let n = 10;
let r = sym_standard_rands(alpha, n).unwrap();
assert!(r.iter().all(|&x| x.is_finite()));
}
#[test]
fn test_sample_skew_standard_alpha() {
let alpha = 0.7;
let mut rng = rng();
let r = rng.sample(StandardSkewStable::new(alpha).unwrap());
assert!(r > 0.0);
}
#[test]
fn test_sample_skew_standard_alpha_rands() {
let alpha = 0.7;
let n = 10;
let r = skew_rands(alpha, n).unwrap();
assert!(r.iter().all(|&x| x > 0.0));
}
}