mathhook_core/functions/special/
polygamma.rs

1use super::{digamma, digamma_numerical};
2use crate::{Expression, Number};
3
4/// Polygamma function ψ^(n)(z)
5///
6/// The polygamma function is the (n+1)-th derivative of ln(Γ(z)):
7/// ψ^(n)(z) = d^(n+1)/dz^(n+1) ln(Γ(z))
8///
9/// # Special Cases
10///
11/// - ψ^(0)(z) = ψ(z) (digamma)
12/// - ψ^(1)(z) = trigamma
13/// - ψ^(2)(z) = tetragamma
14///
15/// # Mathematical Properties
16///
17/// - ψ^(1)(1) = π²/6 (trigamma at 1)
18/// - ψ^(n)(z+1) = ψ^(n)(z) + (-1)^n · n! / z^(n+1)
19///
20/// # Arguments
21///
22/// * `n` - Order of derivative (0 = digamma, 1 = trigamma, etc.)
23/// * `z` - Argument
24///
25/// # Examples
26///
27/// ```rust
28/// use mathhook_core::{Expression, Number};
29/// use mathhook_core::functions::special::polygamma;
30///
31/// let result = polygamma(0, &Expression::Number(Number::Integer(1)));
32/// let trigamma = polygamma(1, &Expression::Number(Number::Integer(1)));
33/// ```
34pub fn polygamma(n: i32, z: &Expression) -> Expression {
35    if n == 0 {
36        return digamma(z);
37    }
38
39    if n == 1 {
40        match z {
41            Expression::Number(Number::Integer(1)) => Expression::div(
42                Expression::pow(Expression::pi(), Expression::integer(2)),
43                Expression::integer(6),
44            ),
45            Expression::Number(Number::Float(x)) => {
46                let result = polygamma_numerical(1, *x);
47                Expression::Number(Number::Float(result))
48            }
49            _ => Expression::function(
50                "polygamma",
51                vec![Expression::Number(Number::Integer(n as i64)), z.clone()],
52            ),
53        }
54    } else {
55        match z {
56            Expression::Number(Number::Float(x)) => {
57                let result = polygamma_numerical(n, *x);
58                Expression::Number(Number::Float(result))
59            }
60            _ => Expression::function(
61                "polygamma",
62                vec![Expression::Number(Number::Integer(n as i64)), z.clone()],
63            ),
64        }
65    }
66}
67
68/// Numerically evaluates the polygamma function
69///
70/// Uses the series formula: ψ^(n)(z) = (-1)^(n+1) n! ∑_{k=0}^∞ 1/(z+k)^(n+1)
71///
72/// # Arguments
73///
74/// * `n` - Order of derivative
75/// * `z` - Value to evaluate at
76///
77/// # Returns
78///
79/// ψ^(n)(z) value
80fn polygamma_numerical(n: i32, z: f64) -> f64 {
81    if n < 0 {
82        return f64::NAN;
83    }
84
85    if n == 0 {
86        return digamma_numerical(z);
87    }
88
89    if z.is_nan() || z.is_infinite() {
90        return f64::NAN;
91    }
92
93    if z <= 0.0 && (z - z.round()).abs() < 1e-10 {
94        return f64::NAN;
95    }
96
97    let sign = if (n + 1) % 2 == 0 { 1.0 } else { -1.0 };
98
99    let mut sum = 0.0;
100    for k in 0..1000 {
101        let term = 1.0 / f64::powi(z + k as f64, n + 1);
102        sum += term;
103        if term.abs() < 1e-15 * sum.abs() {
104            break;
105        }
106    }
107
108    let mut factorial = 1.0;
109    for i in 1..=n {
110        factorial *= i as f64;
111    }
112
113    sign * factorial * sum
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use crate::{Expression, Number};
120
121    #[test]
122    fn test_polygamma_zero_is_digamma() {
123        let z = Expression::Number(Number::Integer(2));
124        let poly_0 = polygamma(0, &z);
125        let dig = digamma(&z);
126        assert_eq!(poly_0, dig);
127    }
128
129    #[test]
130    fn test_polygamma_trigamma_special_value() {
131        let result = polygamma(1, &Expression::Number(Number::Integer(1)));
132        let result_str = result.to_string();
133        assert!(
134            result_str.contains("Pi") && (result_str.contains('6') || result_str.contains("Pow")),
135            "ψ^(1)(1) = π²/6"
136        );
137    }
138}
mathhook_core/functions/special/polygamma.rs

mathhook_core/functions/special/
polygamma.rs
