1 #ifndef STAN_MATH_PRIM_FUN_TRIGAMMA_HPP
2 #define STAN_MATH_PRIM_FUN_TRIGAMMA_HPP
3
4 #include <stan/math/prim/meta.hpp>
5 #include <stan/math/prim/fun/constants.hpp>
6 #include <stan/math/prim/fun/floor.hpp>
7 #include <stan/math/prim/fun/inv.hpp>
8 #include <stan/math/prim/fun/inv_square.hpp>
9 #include <stan/math/prim/fun/sin.hpp>
10 #include <stan/math/prim/fun/square.hpp>
11 #include <stan/math/prim/functor/apply_scalar_unary.hpp>
12 #include <cmath>
13
14 // Reference:
15 // BE Schneider,
16 // Algorithm AS 121:
17 // Trigamma Function,
18 // Applied Statistics,
19 // Volume 27, Number 1, pages 97-99, 1978.
20
21 namespace stan {
22 namespace math {
23
24 /**
25 * Return the trigamma function applied to the argument. The
26 * trigamma function returns the second derivative of the
27 * log Gamma function evaluated at the specified argument.
28 * This base templated version is used in prim, fwd, and rev
29 * implementations.
30 *
31 * @tparam T scalar argument type
32 * @param x argument
33 * @return second derivative of log Gamma function at argument
34 *
35 */
36 template <typename T>
trigamma_impl(const T & x)37 inline T trigamma_impl(const T& x) {
38 using std::floor;
39 using std::sin;
40
41 double small = 0.0001;
42 double large = 5.0;
43 T value;
44 T y;
45 T z;
46
47 // bernoulli numbers
48 double b2 = inv(6.0);
49 double b4 = -inv(30.0);
50 double b6 = inv(42.0);
51 double b8 = -inv(30.0);
52
53 // negative integers and zero return positive infinity
54 // see http://mathworld.wolfram.com/PolygammaFunction.html
55 if (x <= 0.0 && floor(x) == x) {
56 value = positive_infinity();
57 return value;
58 }
59
60 // negative non-integers: use the reflection formula
61 // see http://mathworld.wolfram.com/PolygammaFunction.html
62 if (x <= 0 && floor(x) != x) {
63 value = -trigamma_impl(-x + 1.0) + square(pi() / sin(-pi() * x));
64 return value;
65 }
66
67 // small value approximation if x <= small.
68 if (x <= small) {
69 return inv_square(x);
70 }
71
72 // use recurrence relation until x >= large
73 // see http://mathworld.wolfram.com/PolygammaFunction.html
74 z = x;
75 value = 0.0;
76 while (z < large) {
77 value += inv_square(z);
78 z += 1.0;
79 }
80
81 // asymptotic expansion as a Laurent series if x >= large
82 // see http://en.wikipedia.org/wiki/Trigamma_function
83 y = inv_square(z);
84 value += 0.5 * y + (1.0 + y * (b2 + y * (b4 + y * (b6 + y * b8)))) / z;
85
86 return value;
87 }
88
89 /**
90 * Return the second derivative of the log Gamma function
91 * evaluated at the specified argument.
92 *
93 \f[
94 \mbox{trigamma}(x) =
95 \begin{cases}
96 \textrm{error} & \mbox{if } x\in \{\dots, -3, -2, -1, 0\}\\
97 \Psi_1(x) & \mbox{if } x\not\in \{\dots, -3, -2, -1, 0\}\\[6pt]
98 \textrm{NaN} & \mbox{if } x = \textrm{NaN}
99 \end{cases}
100 \f]
101
102 \f[
103 \frac{\partial\, \mbox{trigamma}(x)}{\partial x} =
104 \begin{cases}
105 \textrm{error} & \mbox{if } x\in \{\dots, -3, -2, -1, 0\}\\
106 \frac{\partial\, \Psi_1(x)}{\partial x} & \mbox{if } x\not\in \{\dots, -3,
107 -2, -1, 0\}\\[6pt] \textrm{NaN} & \mbox{if } x = \textrm{NaN} \end{cases} \f]
108
109 \f[
110 \Psi_1(x)=\sum_{n=0}^\infty \frac{1}{(x+n)^2}
111 \f]
112
113 \f[
114 \frac{\partial \, \Psi_1(x)}{\partial x} = -2\sum_{n=0}^\infty
115 \frac{1}{(x+n)^3} \f]
116 *
117 * @param[in] u argument
118 * @return second derivative of log Gamma function at argument
119 */
trigamma(double u)120 inline double trigamma(double u) { return trigamma_impl(u); }
121
122 /**
123 * Return the second derivative of the log Gamma function
124 * evaluated at the specified argument.
125 *
126 * @param[in] u argument
127 * @return second derivative of log Gamma function at argument
128 */
trigamma(int u)129 inline double trigamma(int u) { return trigamma(static_cast<double>(u)); }
130
131 /**
132 * Structure to wrap `trigamma()` so it can be vectorized.
133 */
134 struct trigamma_fun {
135 /**
136 * Return the trigamma() function applied to
137 * the argument.
138 *
139 * @tparam T type of argument
140 * @param x argument
141 * @return trigamma applied to argument.
142 */
143 template <typename T>
funstan::math::trigamma_fun144 static inline T fun(const T& x) {
145 return trigamma(x);
146 }
147 };
148
149 /**
150 * Return the elementwise application of `trigamma()` to
151 * specified argument container. The return type promotes the
152 * underlying scalar argument type to double if it is an integer,
153 * and otherwise is the argument type.
154 *
155 * @tparam T type of container
156 * @param x container
157 * @return elementwise trigamma of container elements
158 */
159 template <typename T,
160 require_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
trigamma(const T & x)161 inline auto trigamma(const T& x) {
162 return apply_scalar_unary<trigamma_fun, T>::apply(x);
163 }
164
165 } // namespace math
166 } // namespace stan
167
168 #endif
169