trigamma.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_PRIM_FUN_TRIGAMMA_HPP
#define STAN_MATH_PRIM_FUN_TRIGAMMA_HPP
 
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/fun/constants.hpp>
#include <stan/math/prim/fun/floor.hpp>
#include <stan/math/prim/fun/inv.hpp>
#include <stan/math/prim/fun/inv_square.hpp>
#include <stan/math/prim/fun/sin.hpp>
#include <stan/math/prim/fun/square.hpp>
#include <stan/math/prim/functor/apply_scalar_unary.hpp>
#include <cmath>
 
// Reference:
//   BE Schneider,
//   Algorithm AS 121:
//   Trigamma Function,
//   Applied Statistics,
//   Volume 27, Number 1, pages 97-99, 1978.
 
namespace stan {
namespace math {
 
template <typename T>
inline T trigamma_impl(const T& x) {
  double small = 0.0001;
  double large = 5.0;
  T value;
  T y;
  T z;
 
  // bernoulli numbers
  double b2 = inv(6.0);
  double b4 = -inv(30.0);
  double b6 = inv(42.0);
  double b8 = -inv(30.0);
 
  // negative integers and zero return positive infinity
  // see http://mathworld.wolfram.com/PolygammaFunction.html
  if (x <= 0.0 && floor(x) == x) {
    value = positive_infinity();
    return value;
  }
 
  // negative non-integers: use the reflection formula
  // see http://mathworld.wolfram.com/PolygammaFunction.html
  if (x <= 0 && floor(x) != x) {
    value = -trigamma_impl(-x + 1.0) + square(pi() / sin(-pi() * x));
    return value;
  }
 
  // small value approximation if x <= small.
  if (x <= small) {
    return inv_square(x);
  }
 
  // use recurrence relation until x >= large
  // see http://mathworld.wolfram.com/PolygammaFunction.html
  z = x;
  value = 0.0;
  while (z < large) {
    value += inv_square(z);
    z += 1.0;
  }
 
  // asymptotic expansion as a Laurent series if x >= large
  // see http://en.wikipedia.org/wiki/Trigamma_function
  y = inv_square(z);
  value += 0.5 * y + (1.0 + y * (b2 + y * (b4 + y * (b6 + y * b8)))) / z;
 
  return value;
}
 
inline double trigamma(double u) { return trigamma_impl(u); }
 
inline double trigamma(int u) { return trigamma(static_cast<double>(u)); }
 
struct trigamma_fun {
  template <typename T>
  static inline auto fun(const T& x) {
    return trigamma(x);
  }
};
 
template <typename T,
          require_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
inline auto trigamma(const T& x) {
  return apply_scalar_unary<trigamma_fun, T>::apply(x);
}
 
}  // namespace math
}  // namespace stan
 
#endif