Skip to main content

torsh_special/
lib.rs

1//! Special mathematical functions for ToRSh
2
3#![allow(clippy::result_large_err)]
4
5// Version information
6pub const VERSION: &str = env!("CARGO_PKG_VERSION");
7pub const VERSION_MAJOR: u32 = 0;
8pub const VERSION_MINOR: u32 = 1;
9pub const VERSION_PATCH: u32 = 0;
10
11use torsh_core::Result;
12use torsh_tensor::Tensor;
13
14/// Convenience type alias for Results in this crate
15pub type TorshResult<T> = Result<T>;
16
17pub mod advanced;
18pub mod advanced_special;
19pub mod airy;
20pub mod bessel;
21pub mod complex;
22pub mod constants;
23pub mod coulomb;
24pub mod elliptic;
25pub mod error_functions;
26pub mod error_handling;
27pub mod exponential_integrals;
28pub mod fast_approximations;
29pub mod gamma;
30pub mod hypergeometric;
31pub mod lambert_w;
32pub mod lommel;
33pub mod lookup_tables;
34pub mod mathieu;
35pub mod numerical_accuracy_tests;
36pub mod orthogonal_polynomials;
37pub mod scirs2_integration;
38pub mod simd_optimizations;
39pub mod smart_caching;
40pub mod spheroidal;
41pub mod statistical;
42pub mod trigonometric;
43pub mod utils;
44pub mod visualization;
45
46#[cfg(test)]
47pub mod test_fix;
48
49// Re-exports
50pub use bessel::*;
51// Use scirs2_integration versions for better PyTorch compatibility
52pub use scirs2_integration::{
53    bessel_i0_scirs2, bessel_i1_scirs2, bessel_j0_scirs2, bessel_j1_scirs2, bessel_jn_scirs2,
54    bessel_k0_scirs2, bessel_k1_scirs2, bessel_y0_scirs2, bessel_y1_scirs2, bessel_yn_scirs2, beta,
55    digamma, erf, erfc, erfcx, erfinv, fresnel, fresnel_c, fresnel_s, gamma, lgamma, polygamma,
56    sinc,
57};
58// Export elliptic functions
59pub use elliptic::{
60    elliptic_e, elliptic_e_incomplete, elliptic_f, elliptic_k, jacobi_cn, jacobi_dn, jacobi_sn,
61    theta_1, theta_2, theta_3, theta_4, weierstrass_p, weierstrass_sigma, weierstrass_zeta,
62};
63// Export exponential integrals
64pub use exponential_integrals::{
65    cosine_integral, exponential_integral_e1, exponential_integral_ei, exponential_integral_en,
66    hyperbolic_cosine_integral, hyperbolic_sine_integral, logarithmic_integral, sine_integral,
67};
68// Export hypergeometric functions
69pub use hypergeometric::{
70    appell_f1, hypergeometric_1f1, hypergeometric_2f1, hypergeometric_pfq, hypergeometric_u,
71    meijer_g,
72};
73// Export orthogonal polynomials
74pub use orthogonal_polynomials::{
75    chebyshev_t, chebyshev_u, gegenbauer_c, hermite_h, hermite_he, jacobi_p, laguerre_l,
76    laguerre_l_associated, legendre_p, legendre_p_associated,
77};
78// Export advanced functions
79pub use advanced::{barnes_g, dirichlet_eta, hurwitz_zeta, polylogarithm, riemann_zeta};
80// Export advanced special functions
81pub use advanced_special::{
82    dawson, kelvin_bei, kelvin_ber, kelvin_kei, kelvin_ker, parabolic_cylinder_d, spence, struve_h,
83    struve_l, voigt_profile,
84};
85// Export Airy functions
86pub use airy::{airy_ai, airy_ai_prime, airy_bi, airy_bi_prime};
87// Export Coulomb wave functions
88pub use coulomb::{coulomb_f, coulomb_g, coulomb_sigma};
89// Export Lommel functions
90pub use lommel::{lommel_S, lommel_s, lommel_u, lommel_v};
91// Export Mathieu functions
92pub use mathieu::{mathieu_Ce, mathieu_Se, mathieu_a, mathieu_b, mathieu_ce, mathieu_se};
93// Export Spheroidal Wave functions
94pub use spheroidal::{
95    oblate_angular, oblate_angular_tensor, oblate_radial, oblate_radial_tensor, prolate_angular,
96    prolate_angular_tensor, prolate_radial, prolate_radial_tensor, spheroidal_eigenvalue,
97};
98// Export Lambert W functions
99pub use lambert_w::{
100    lambert_w, lambert_w_complex, lambert_w_derivative, lambert_w_principal, lambert_w_secondary,
101};
102// Export Lambert W applications
103pub use lambert_w::applications::{
104    solve_exponential_equation, solve_x_to_x_equals_c, tree_function, wright_omega,
105};
106// Export statistical functions
107pub use statistical::{
108    chi_squared_cdf, f_distribution_cdf, incomplete_beta, normal_cdf, normal_cdf_general,
109    normal_pdf, normal_pdf_general, student_t_cdf,
110};
111// Export complex functions
112pub use complex::{
113    complex_airy_ai_c64,
114    complex_airy_bi_c64,
115    complex_bessel_j_c32,
116    complex_bessel_j_c64,
117    complex_bessel_y_c32,
118    complex_bessel_y_c64,
119    complex_beta_c32,
120    complex_beta_c64,
121    complex_erf_c32,
122    complex_erf_c64,
123    complex_erfc_c32,
124    complex_erfc_c64,
125    complex_gamma_c32,
126    complex_gamma_c64,
127    complex_incomplete_gamma_c64,
128    complex_log_principal,
129    complex_polygamma_c32,
130    // New complex functions
131    complex_polygamma_c64,
132    complex_pow_principal,
133    complex_sqrt_principal,
134    complex_zeta_c32,
135    complex_zeta_c64,
136};
137// Export enhanced error handling
138pub use error_handling::{
139    error_recovery::{clamp_to_finite, gradual_clamp, replace_problematic_values},
140    safe_functions::{
141        safe_bessel_j0, safe_bessel_j1, safe_bessel_y0, safe_bessel_y1, safe_erf, safe_erfc,
142        safe_gamma, safe_lgamma,
143    },
144    DomainConstraints, InputValidation,
145};
146// Export performance optimizations
147pub use fast_approximations::{
148    atanh_fast, bessel_j0_fast, cos_fast, cosh_fast, erf_fast, erfc_fast, exp_fast, gamma_fast,
149    log_fast, sin_fast, sinh_fast, tanh_fast,
150};
151pub use lookup_tables::{
152    bessel_j0_optimized, erf_optimized, factorial, gamma_optimized, POLY_COEFFS,
153};
154pub use simd_optimizations::{erf_simd, exp_family_simd, gamma_simd, ExpVariant};
155// Export smart caching functionality
156pub use smart_caching::{
157    cache_stats, cached_compute, clear_cache, function_ids, CacheStats, FloatKey, SmartCache,
158};
159// Export non-overlapping functions from other modules
160pub use trigonometric::sinc as sinc_unnormalized;
161// Keep error_functions and gamma modules available for reference
162// but don't re-export to avoid conflicts
163
164/// Exponential minus one (exp(x) - 1) for better numerical stability
165pub fn expm1(tensor: &Tensor<f32>) -> TorshResult<Tensor<f32>> {
166    // expm1(x) = exp(x) - 1
167    let exp_tensor = tensor.exp()?;
168    let ones = tensor.ones_like()?;
169    exp_tensor.sub(&ones)
170}
171
172/// Natural logarithm of one plus x (log(1 + x)) for better numerical stability  
173pub fn log1p(tensor: &Tensor<f32>) -> TorshResult<Tensor<f32>> {
174    // log1p(x) = log(1 + x)
175    let ones = tensor.ones_like()?;
176    let one_plus_x = ones.add_op(tensor)?;
177    one_plus_x.log()
178}
179
180/// Hyperbolic sine (using existing implementation)
181pub fn sinh(tensor: &Tensor<f32>) -> TorshResult<Tensor<f32>> {
182    tensor.sinh()
183}
184
185/// Hyperbolic cosine (using existing implementation)
186pub fn cosh(tensor: &Tensor<f32>) -> TorshResult<Tensor<f32>> {
187    tensor.cosh()
188}
189
190/// Hyperbolic tangent (using existing implementation)
191pub fn tanh_special(tensor: &Tensor<f32>) -> TorshResult<Tensor<f32>> {
192    tensor.tanh()
193}
194
195/// Inverse hyperbolic sine
196pub fn asinh(tensor: &Tensor<f32>) -> TorshResult<Tensor<f32>> {
197    // asinh(x) = log(x + sqrt(x^2 + 1))
198    let x_squared = tensor.mul_op(tensor)?;
199    let ones = tensor.ones_like()?;
200    let under_sqrt = x_squared.add_op(&ones)?;
201    let sqrt_part = under_sqrt.sqrt()?;
202    let sum = tensor.add_op(&sqrt_part)?;
203    sum.log()
204}
205
206/// Inverse hyperbolic cosine
207pub fn acosh(tensor: &Tensor<f32>) -> TorshResult<Tensor<f32>> {
208    // acosh(x) = log(x + sqrt(x^2 - 1))
209    let x_squared = tensor.mul_op(tensor)?;
210    let ones = tensor.ones_like()?;
211    let under_sqrt = x_squared.sub(&ones)?;
212    let sqrt_part = under_sqrt.sqrt()?;
213    let sum = tensor.add_op(&sqrt_part)?;
214    sum.log()
215}
216
217/// Inverse hyperbolic tangent
218pub fn atanh(tensor: &Tensor<f32>) -> TorshResult<Tensor<f32>> {
219    // atanh(x) = 0.5 * log((1 + x) / (1 - x))
220    let ones = tensor.ones_like()?;
221    let one_plus_x = ones.add_op(tensor)?;
222    let one_minus_x = ones.sub(tensor)?;
223    let ratio = one_plus_x.div(&one_minus_x)?;
224    let log_ratio = ratio.log()?;
225    log_ratio.mul_scalar(0.5)
226}
227
228// Export visualization tools
229pub use visualization::{
230    analyze_function_behavior, benchmark_optimization_levels, compare_function_accuracy,
231    generate_ascii_plot, print_accuracy_report, print_analysis_report, AccuracyComparison,
232    FunctionAnalysis, Monotonicity, PlotData,
233};
234// Export mathematical constants
235pub use constants::{
236    APERY, CATALAN, DOTTIE, EULER_GAMMA, FEIGENBAUM_ALPHA, FEIGENBAUM_DELTA, GLAISHER_KINKELIN,
237    GOLDEN_RATIO, GOLDEN_RATIO_RECIPROCAL, KHINCHIN, LEVY, LN_2PI, LN_PI, MILLS, PI_SQUARED_OVER_6,
238    PLASTIC_NUMBER, RAMANUJAN, SQRT_2PI,
239};
240// Export utility constants
241pub use constants::utility::{
242    HALF_LN_2PI, INV_2PI, INV_SQRT_2PI, LN_4PI, LOG_SQRT_2PI, PI_OVER_4, SQRT_PI_OVER_2,
243    SQRT_PI_OVER_8, STIRLING_CONSTANT, THREE_PI_OVER_4, TWO_OVER_SQRT_PI,
244};
245// Export specialized constant collections for advanced users
246pub use constants::{bernoulli, fractions, physics, zeta_values};
247// Export utility functions
248pub use utils::{
249    chebyshev_expansion, chebyshev_expansion_tensor, continued_fraction, continued_fraction_tensor,
250    double_factorial, factorial_safe, pade_approximant, pade_approximant_tensor,
251};
252
253/// Prelude module for convenient imports
254pub mod prelude {
255    pub use crate::{
256        advanced::*, advanced_special::*, airy::*, bessel::*, coulomb::*, elliptic::*,
257        exponential_integrals::*, hypergeometric::*, lambert_w::*, lommel::*, mathieu::*,
258        orthogonal_polynomials::*, scirs2_integration::*, spheroidal::*,
259    };
260}