Merge pull request #63 from JSorngard/lambert_w

Implement the Lambert W function as a wrapper around the `lambert_w` crate.

Cargo.toml

@@ -39,6 +39,7 @@ serde = { version = "1.0", features = ["derive"], optional = true }
 json = { version = "0.12", optional = true }
 arrow2 = { version = "0.18", features = ["io_parquet", "io_parquet_compression"], optional = true }
 num-complex = { version = "0.4", optional = true }
+lambert_w = { version = "0.3.0", default-features = false, features = ["24bits", "50bits"] }
 
 [package.metadata.docs.rs]
 rustdoc-args = [ "--html-in-header", "katex-header.html", "--cfg", "docsrs"]

src/prelude/mod.rs

@@ -193,7 +193,10 @@ pub use crate::util::{api::*, low_level::*, non_macro::*, print::*, useful::*, w
 pub use crate::statistics::{dist::*, ops::*, rand::*, stat::*};
 
 #[allow(unused_imports)]
-pub use crate::special::function::*;
+pub use crate::special::function::{
+    beta, erf, erfc, gamma, gaussian, inc_beta, inc_gamma, inv_erf, inv_erfc, inv_inc_gamma,
+    inv_inv_beta, ln_gamma, phi, poch,
+};
 
 #[allow(unused_imports)]
 pub use crate::numerical::{
@@ -206,7 +209,7 @@ pub use crate::numerical::{
     utils::*,
 };
 
-pub use simpler::{eigen, integrate, chebyshev_polynomial, cubic_hermite_spline};
+pub use simpler::{eigen, integrate, chebyshev_polynomial, cubic_hermite_spline, lambert_w0, lambert_w_m1};
 
 #[allow(unused_imports)]
 pub use crate::statistics::stat::Metric::*;

src/prelude/simpler.rs

@@ -152,6 +152,40 @@ pub fn cubic_hermite_spline(node_x: &[f64], node_y: &[f64]) -> anyhow::Result<Cu
     spline::cubic_hermite_spline(node_x, node_y, Quadratic)
 }
 
+use crate::special::function::{
+    lambert_w0 as lambert_w0_flex, lambert_wm1 as lambert_w_m1_flex, LambertWAccuracyMode,
+};
+
+/// The principal branch of the Lambert W function, W_0(`z`).
+///
+/// Returns [`NAN`](f64::NAN) if the given input is smaller than -1/e (≈ -0.36787944117144233).
+///
+/// Accurate to 50 bits.
+/// 
+/// Wrapper of `lambert_w_0` function of `lambert_w` crate.
+///
+/// # Reference
+///
+/// [Toshio Fukushima, Precise and fast computation of Lambert W function by piecewise minimax rational function approximation with variable transformation](https://www.researchgate.net/publication/346309410_Precise_and_fast_computation_of_Lambert_W_function_by_piecewise_minimax_rational_function_approximation_with_variable_transformation)
+pub fn lambert_w0(z: f64) -> f64 {
+    lambert_w0_flex(z, LambertWAccuracyMode::Precise)
+}
+
+/// The secondary branch of the Lambert W function, W_-1(`z`).
+///
+/// Returns [`NAN`](f64::NAN) if the given input is positive or smaller than -1/e (≈ -0.36787944117144233).
+///
+/// Accurate to 50 bits.
+/// 
+/// Wrapper of `lambert_w_m1` function of `lambert_w` crate.
+///
+/// # Reference
+///
+/// [Toshio Fukushima, Precise and fast computation of Lambert W function by piecewise minimax rational function approximation with variable transformation](https://www.researchgate.net/publication/346309410_Precise_and_fast_computation_of_Lambert_W_function_by_piecewise_minimax_rational_function_approximation_with_variable_transformation)
+pub fn lambert_w_m1(z: f64) -> f64 {
+    lambert_w_m1_flex(z, LambertWAccuracyMode::Precise)
+}
+
 /// Simple handle parquet
 #[cfg(feature="parquet")]
 pub trait SimpleParquet: Sized {

src/special/function.rs

@@ -126,3 +126,52 @@ pub fn phi(x: f64) -> f64 {
 //         special_fun::unsafe_cephes_double::hyp2f1(a, b, c, x)
 //     }
 // }
+
+/// The principal branch of the Lambert W function, W_0(`z`).
+///
+/// Returns [`NAN`](f64::NAN) if the given input is smaller than -1/e (≈ -0.36787944117144233).
+///
+/// Use [`Precise`](LambertWAccuracyMode::Precise) for 50 bits of accuracy and the [`Simple`](LambertWAccuracyMode::Simple) mode 
+/// for only 24 bits, but with faster execution time.
+/// 
+/// Wrapper of the `lambert_w_0` and `sp_lambert_w_0` functions of the `lambert_w` crate.
+///
+/// # Reference
+///
+/// [Toshio Fukushima, Precise and fast computation of Lambert W function by piecewise minimax rational function approximation with variable transformation](https://www.researchgate.net/publication/346309410_Precise_and_fast_computation_of_Lambert_W_function_by_piecewise_minimax_rational_function_approximation_with_variable_transformation)
+pub fn lambert_w0(z: f64, mode: LambertWAccuracyMode) -> f64 {
+    match mode {
+        LambertWAccuracyMode::Precise => lambert_w::lambert_w_0(z),
+        LambertWAccuracyMode::Simple => lambert_w::sp_lambert_w_0(z),
+    }
+    .unwrap_or(f64::NAN)
+}
+
+/// The secondary branch of the Lambert W function, W_-1(`z`).
+///
+/// Returns [`NAN`](f64::NAN) if the given input is positive or smaller than -1/e (≈ -0.36787944117144233).
+///
+/// Use [`Precise`](LambertWAccuracyMode::Precise) for 50 bits of accuracy and the [`Simple`](LambertWAccuracyMode::Simple) mode 
+/// for only 24 bits, but with faster execution time.
+/// 
+/// Wrapper of the `lambert_w_m1` and `sp_lambert_w_m1` functions of the `lambert_w` crate.
+/// 
+/// # Reference
+///
+/// [Toshio Fukushima, Precise and fast computation of Lambert W function by piecewise minimax rational function approximation with variable transformation](https://www.researchgate.net/publication/346309410_Precise_and_fast_computation_of_Lambert_W_function_by_piecewise_minimax_rational_function_approximation_with_variable_transformation)
+pub fn lambert_wm1(z: f64, mode: LambertWAccuracyMode) -> f64 {
+    match mode {
+        LambertWAccuracyMode::Precise => lambert_w::lambert_w_m1(z),
+        LambertWAccuracyMode::Simple => lambert_w::sp_lambert_w_m1(z),
+    }
+    .unwrap_or(f64::NAN)
+}
+
+/// Decides the accuracy mode of the Lambert W functions.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum LambertWAccuracyMode {
+    /// Faster, 24 bits of accuracy.
+    Simple,
+    /// Slower, 50 bits of accuracy.
+    Precise,
+}