Added polygamma function.

This commit is contained in:
Till Hoffmann 2016-04-01 14:35:21 +01:00
parent dd5d390daf
commit 57239f4a81
9 changed files with 118 additions and 5 deletions

View File

@ -77,6 +77,7 @@ struct default_packet_traits
HasLGamma = 0,
HasDiGamma = 0,
HasZeta = 0,
HasPolygamma = 0,
HasErf = 0,
HasErfc = 0,
HasIGamma = 0,
@ -456,6 +457,10 @@ Packet pdigamma(const Packet& a) { using numext::digamma; return digamma(a); }
template<typename Packet> EIGEN_DECLARE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS
Packet pzeta(const Packet& x, const Packet& q) { using numext::zeta; return zeta(x, q); }
/** \internal \returns the polygamma function (coeff-wise) */
template<typename Packet> EIGEN_DECLARE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS
Packet ppolygamma(const Packet& n, const Packet& x) { using numext::polygamma; return polygamma(n, x); }
/** \internal \returns the erf(\a a) (coeff-wise) */
template<typename Packet> EIGEN_DECLARE_FUNCTION_ALLOWING_MULTIPLE_DEFINITIONS
Packet perf(const Packet& a) { using numext::erf; return erf(a); }

View File

@ -52,6 +52,7 @@ namespace Eigen
EIGEN_ARRAY_DECLARE_GLOBAL_UNARY(lgamma,scalar_lgamma_op)
EIGEN_ARRAY_DECLARE_GLOBAL_UNARY(digamma,scalar_digamma_op)
EIGEN_ARRAY_DECLARE_GLOBAL_UNARY(zeta,scalar_zeta_op)
EIGEN_ARRAY_DECLARE_GLOBAL_UNARY(polygamma,scalar_polygamma_op)
EIGEN_ARRAY_DECLARE_GLOBAL_UNARY(erf,scalar_erf_op)
EIGEN_ARRAY_DECLARE_GLOBAL_UNARY(erfc,scalar_erfc_op)
EIGEN_ARRAY_DECLARE_GLOBAL_UNARY(exp,scalar_exp_op)

View File

@ -736,7 +736,7 @@ struct zeta_retval {
template <typename Scalar>
struct zeta_impl {
EIGEN_DEVICE_FUNC
static Scalar run(Scalar x) {
static Scalar run(Scalar x, Scalar q) {
EIGEN_STATIC_ASSERT((internal::is_same<Scalar, Scalar>::value == false),
THIS_TYPE_IS_NOT_SUPPORTED);
return Scalar(0);
@ -905,6 +905,50 @@ struct zeta_impl {
#endif // EIGEN_HAS_C99_MATH
/****************************************************************************
* Implementation of polygamma function *
****************************************************************************/
template <typename Scalar>
struct polygamma_retval {
typedef Scalar type;
};
#ifndef EIGEN_HAS_C99_MATH
template <typename Scalar>
struct polygamma_impl {
EIGEN_DEVICE_FUNC
static Scalar run(Scalar n, Scalar x) {
EIGEN_STATIC_ASSERT((internal::is_same<Scalar, Scalar>::value == false),
THIS_TYPE_IS_NOT_SUPPORTED);
return Scalar(0);
}
};
#else
template <typename Scalar>
struct polygamma_impl {
EIGEN_DEVICE_FUNC
static Scalar run(Scalar n, Scalar x) {
Scalar zero = 0.0, one = 1.0;
Scalar nplus = n + one;
// Just return the digamma function for n = 1
if (n == zero) {
return digamma_impl<Scalar>::run(x);
}
// Use the same implementation as scipy
else {
Scalar factorial = numext::exp(lgamma_impl<Scalar>::run(nplus));
return numext::pow(-one, nplus) * factorial * zeta_impl<Scalar>::run(nplus, x);
}
}
};
#endif // EIGEN_HAS_C99_MATH
} // end namespace internal
namespace numext {
@ -927,6 +971,12 @@ zeta(const Scalar& x, const Scalar& q) {
return EIGEN_MATHFUNC_IMPL(zeta, Scalar)::run(x, q);
}
template <typename Scalar>
EIGEN_DEVICE_FUNC inline EIGEN_MATHFUNC_RETVAL(polygamma, Scalar)
polygamma(const Scalar& n, const Scalar& x) {
return EIGEN_MATHFUNC_IMPL(polygamma, Scalar)::run(n, x);
}
template <typename Scalar>
EIGEN_DEVICE_FUNC inline EIGEN_MATHFUNC_RETVAL(erf, Scalar)
erf(const Scalar& x) {

View File

@ -93,17 +93,31 @@ double2 pdigamma<double2>(const double2& a)
}
template<> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
float4 pzeta<float4>(const float4& a)
float4 pzeta<float4>(const float4& x, const float4& q)
{
using numext::zeta;
return make_float4(zeta(a.x), zeta(a.y), zeta(a.z), zeta(a.w));
return make_float4(zeta(x.x, q.x), zeta(x.y, q.y), zeta(x.z, q.z), zeta(x.w, q.w));
}
template<> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
double2 pzeta<double2>(const double2& a)
double2 pzeta<double2>(const double2& x, const double2& q)
{
using numext::zeta;
return make_double2(zeta(a.x), zeta(a.y));
return make_double2(zeta(x.x, q.x), zeta(x.y, q.y));
}
template<> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
float4 ppolygamma<float4>(const float4& n, const float4& x)
{
using numext::polygamma;
return make_float4(polygamma(n.x, x.x), polygamma(n.y, x.y), polygamma(n.z, x.z), polygamma(n.w, x.w));
}
template<> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
double2 ppolygamma<double2>(const double2& n, const double2& x)
{
using numext::polygamma;
return make_double2(polygamma(n.x, x.x), polygamma(n.y, x.y));
}
template<> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE

View File

@ -41,6 +41,7 @@ template<> struct packet_traits<float> : default_packet_traits
HasLGamma = 1,
HasDiGamma = 1,
HasZeta = 1,
HasPolygamma = 1,
HasErf = 1,
HasErfc = 1,
HasIgamma = 1,

View File

@ -471,6 +471,28 @@ struct functor_traits<scalar_zeta_op<Scalar> >
};
};
/** \internal
* \brief Template functor to compute the polygamma function.
* \sa class CwiseUnaryOp, Cwise::polygamma()
*/
template<typename Scalar> struct scalar_polygamma_op {
EIGEN_EMPTY_STRUCT_CTOR(scalar_polygamma_op)
EIGEN_DEVICE_FUNC inline const Scalar operator() (const Scalar& n, const Scalar& x) const {
using numext::polygamma; return polygamma(n, x);
}
typedef typename packet_traits<Scalar>::type Packet;
EIGEN_DEVICE_FUNC inline Packet packetOp(const Packet& n, const Packet& x) const { return internal::ppolygamma(n, x); }
};
template<typename Scalar>
struct functor_traits<scalar_polygamma_op<Scalar> >
{
enum {
// Guesstimate
Cost = 10 * NumTraits<Scalar>::MulCost + 5 * NumTraits<Scalar>::AddCost,
PacketAccess = packet_traits<Scalar>::HasPolygamma
};
};
/** \internal
* \brief Template functor to compute the Gauss error function of a
* scalar

View File

@ -24,6 +24,7 @@ typedef CwiseUnaryOp<internal::scalar_cosh_op<Scalar>, const Derived> CoshReturn
typedef CwiseUnaryOp<internal::scalar_lgamma_op<Scalar>, const Derived> LgammaReturnType;
typedef CwiseUnaryOp<internal::scalar_digamma_op<Scalar>, const Derived> DigammaReturnType;
typedef CwiseUnaryOp<internal::scalar_zeta_op<Scalar>, const Derived> ZetaReturnType;
typedef CwiseUnaryOp<internal::scalar_polygamma_op<Scalar>, const Derived> PolygammaReturnType;
typedef CwiseUnaryOp<internal::scalar_erf_op<Scalar>, const Derived> ErfReturnType;
typedef CwiseUnaryOp<internal::scalar_erfc_op<Scalar>, const Derived> ErfcReturnType;
typedef CwiseUnaryOp<internal::scalar_pow_op<Scalar>, const Derived> PowReturnType;
@ -338,6 +339,14 @@ zeta() const
return ZetaReturnType(derived());
}
/** \returns an expression of the coefficient-wise polygamma function.
*/
inline const PolygammaReturnType
polygamma() const
{
return PolygammaReturnType(derived());
}
/** \returns an expression of the coefficient-wise Gauss error
* function of *this.
*

View File

@ -331,6 +331,11 @@ template<typename ArrayType> void array_real(const ArrayType& m)
VERIFY_IS_APPROX(numext::zeta(Scalar(3), Scalar(-2.5)), RealScalar(0.054102025820864097));
VERIFY_IS_EQUAL(numext::zeta(Scalar(1), Scalar(1.2345)), // The second scalar does not matter
std::numeric_limits<RealScalar>::infinity());
// Check the polygamma against scipy.special.polygamma
VERIFY_IS_APPROX(numext::polygamma(Scalar(1), Scalar(2)), RealScalar(0.644934066848));
VERIFY_IS_APPROX(numext::polygamma(Scalar(1), Scalar(3)), RealScalar(0.394934066848));
VERIFY_IS_APPROX(numext::polygamma(Scalar(1), Scalar(25.5)), RealScalar(0.0399946696496));
{
// Test various propreties of igamma & igammac. These are normalized

View File

@ -138,6 +138,12 @@ class TensorBase<Derived, ReadOnlyAccessors>
zeta() const {
return unaryExpr(internal::scalar_zeta_op<Scalar>());
}
EIGEN_DEVICE_FUNC
EIGEN_STRONG_INLINE const TensorCwiseUnaryOp<internal::scalar_polygamma_op<Scalar>, const Derived>
polygamma() const {
return unaryExpr(internal::scalar_polygamma_op<Scalar>());
}
EIGEN_DEVICE_FUNC
EIGEN_STRONG_INLINE const TensorCwiseUnaryOp<internal::scalar_erf_op<Scalar>, const Derived>