traccc/utils/prob.hpp

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * This file includes code from the ROOT (https://github.com/root-project/root)
0004  * and Cephes Library (http://www.netlib.org/cephes)
0005  *
0006  * ROOT is licensed under the GNU Lesser General Public License v2.1
0007  *
0008  */
0009
0010 #pragma once
0011
0012 // Project include(s).
0013 #include "traccc/definitions/common.hpp"
0014 #include "traccc/definitions/math.hpp"
0015 #include "traccc/definitions/qualifiers.hpp"
0016
0017 namespace traccc {
0018
0019 /* Logarithm of gamma function */
0020 /* A[]: Stirling's formula expansion of log gamma
0021  * B[], C[]: log gamma function between 2 and 3
0022  */
0023
0024 template <typename scalar_t>
0025 struct log_gamma {
0026   TRACCC_HOST_DEVICE static scalar_t A(unsigned int i) {
0027     switch (i) {
0028       case 0:
0029         return A0;
0030       case 1:
0031         return A1;
0032       case 2:
0033         return A2;
0034       case 3:
0035         return A3;
0036       case 4:
0037         return A4;
0038     }
0039     assert(false);
0040     return 0.f;
0041   }
0042
0043   TRACCC_HOST_DEVICE static scalar_t B(unsigned int i) {
0044     switch (i) {
0045       case 0:
0046         return B0;
0047       case 1:
0048         return B1;
0049       case 2:
0050         return B2;
0051       case 3:
0052         return B3;
0053       case 4:
0054         return B4;
0055       case 5:
0056         return B5;
0057     }
0058     assert(false);
0059     return 0.f;
0060   }
0061
0062   TRACCC_HOST_DEVICE static scalar_t C(unsigned int i) {
0063     switch (i) {
0064       case 0:
0065         return C0;
0066       case 1:
0067         return C1;
0068       case 2:
0069         return C2;
0070       case 3:
0071         return C3;
0072       case 4:
0073         return C4;
0074       case 5:
0075         return C5;
0076     }
0077     assert(false);
0078     return 0.f;
0079   }
0080
0081   static constexpr scalar_t kMAXLGM = static_cast<scalar_t>(2.556348e305);
0082   static constexpr scalar_t kMACHEP = 1.11022302462515654042363166809e-16f;
0083   static constexpr scalar_t kMAXLOG = 709.782712893383973096206318587f;
0084   static constexpr scalar_t kBig = 4.503599627370496e15f;
0085   static constexpr scalar_t kBiginv = 2.22044604925031308085e-16f;
0086   /* log( sqrt( 2*pi ) ) */
0087   static constexpr scalar_t LS2PI = 0.91893853320467274178f;
0088
0089  private:
0090   static constexpr scalar_t A0 = 8.11614167470508450300E-4f;
0091   static constexpr scalar_t A1 = -5.95061904284301438324E-4f;
0092   static constexpr scalar_t A2 = 7.93650340457716943945E-4f;
0093   static constexpr scalar_t A3 = -2.77777777730099687205E-3f;
0094   static constexpr scalar_t A4 = 8.33333333333331927722E-2f;
0095
0096   static constexpr scalar_t B0 = -1.37825152569120859100E3f;
0097   static constexpr scalar_t B1 = -3.88016315134637840924E4f;
0098   static constexpr scalar_t B2 = -3.31612992738871184744E5f;
0099   static constexpr scalar_t B3 = -1.16237097492762307383E6f;
0100   static constexpr scalar_t B4 = -1.72173700820839662146E6f;
0101   static constexpr scalar_t B5 = -8.53555664245765465627E5f;
0102
0103   static constexpr scalar_t C0 = -3.51815701436523470549E2f;
0104   static constexpr scalar_t C1 = -1.70642106651881159223E4f;
0105   static constexpr scalar_t C2 = -2.20528590553854454839E5f;
0106   static constexpr scalar_t C3 = -1.13933444367982507207E6f;
0107   static constexpr scalar_t C4 = -2.53252307177582951285E6f;
0108   static constexpr scalar_t C5 = -2.01889141433532773231E6f;
0109 };
0110
0111 // Forward declarations of functions
0112 template <typename scalar_t>
0113 TRACCC_HOST_DEVICE scalar_t chisquared_cdf_c(const scalar_t x,
0114                                              const scalar_t r);
0115
0116 template <typename scalar_t>
0117 TRACCC_HOST_DEVICE inline scalar_t igam(const scalar_t a, const scalar_t x);
0118
0119 template <typename scalar_t>
0120 TRACCC_HOST_DEVICE inline scalar_t igam_impl(const scalar_t a,
0121                                              const scalar_t x);
0122
0123 template <typename scalar_t>
0124 TRACCC_HOST_DEVICE inline scalar_t igamc(const scalar_t a, const scalar_t x);
0125
0126 template <typename scalar_t>
0127 TRACCC_HOST_DEVICE inline scalar_t igamc_impl(const scalar_t a,
0128                                               const scalar_t x);
0129
0130 template <typename scalar_t>
0131 TRACCC_HOST_DEVICE inline scalar_t lgam(scalar_t x);
0132
0133 template <typename scalar_t>
0134 TRACCC_HOST_DEVICE inline scalar_t lgam_impl(scalar_t x);
0135
0136 template <typename scalar_t>
0137 TRACCC_HOST_DEVICE inline scalar_t Polynomialeval_A(const scalar_t x,
0138                                                     const unsigned int N);
0139 template <typename scalar_t>
0140 TRACCC_HOST_DEVICE inline scalar_t Polynomialeval_B(const scalar_t x,
0141                                                     const unsigned int N);
0142
0143 template <typename scalar_t>
0144 TRACCC_HOST_DEVICE inline scalar_t Polynomialeval_C(const scalar_t x,
0145                                                     const unsigned int N);
0146
0147 template <typename scalar_t>
0148 TRACCC_HOST_DEVICE inline scalar_t prob(const scalar_t chi2,
0149                                         const scalar_t ndf) {
0150   if (ndf <= 0)
0151     return 0;  // Set CL to zero in case ndf<=0
0152
0153   if (chi2 <= 0) {
0154     if (chi2 < 0)
0155       return 0;
0156     else
0157       return 1;
0158   }
0159
0160   const scalar_t ret_val = chisquared_cdf_c(chi2, ndf);
0161   assert(ret_val >= 0.f && ret_val <= 1.f);
0162   return ret_val;
0163 }
0164
0165 template <typename scalar_t>
0166 TRACCC_HOST_DEVICE inline scalar_t chisquared_cdf_c(const scalar_t x,
0167                                                     const scalar_t r) {
0168   scalar_t retval = igamc(0.5f * r, 0.5f * x);
0169   return static_cast<scalar_t>(retval);
0170 }
0171
0172 template <typename scalar_t>
0173 TRACCC_HOST_DEVICE inline scalar_t igam(const scalar_t a, const scalar_t x) {
0174   // LM: for negative values returns 1.0 instead of zero
0175   // This is correct if a is a negative integer since Gamma(-n) = +/- inf
0176   if (a <= 0)
0177     return 1.0;
0178
0179   if (x <= 0)
0180     return 0.0;
0181
0182   // for (x > 1) && (x > a where a > 0)
0183   if ((x > 1.0) && (x > a))
0184     return (1.0 - igamc_impl<scalar_t>(a, x));
0185
0186   // for (0 < x < 1) || ( x < a where a > 0)
0187   return igam_impl<scalar_t>(a, x);
0188 }
0189
0190 // for (0 < x < 1) || ( x < a where a > 0)
0191 template <typename scalar_t>
0192 TRACCC_HOST_DEVICE inline scalar_t igam_impl(const scalar_t a,
0193                                              const scalar_t x) {
0194   scalar_t ans, ax, c, r;
0195
0196   /* Compute  x**a * exp(-x) / gamma(a)  */
0197   ax = a * math::log(x) - x - lgam(a);
0198   if (ax < -log_gamma<scalar_t>::kMAXLOG)
0199     return (0.0f);
0200
0201   ax = std::exp(ax);
0202
0203   /* power series */
0204   r = a;
0205   c = 1.0f;
0206   ans = 1.0f;
0207
0208   do {
0209     r += 1.0f;
0210     c *= x / r;
0211     ans += c;
0212   } while (c / ans > log_gamma<scalar_t>::kMACHEP);
0213
0214   return (ans * ax / a);
0215 }
0216
0217 template <typename scalar_t>
0218 TRACCC_HOST_DEVICE inline scalar_t igamc(const scalar_t a, const scalar_t x) {
0219   // LM: for negative values returns 0.0
0220   // This is correct if a is a negative integer since Gamma(-n) = +/- inf
0221   if (a <= 0.f)
0222     return 0.0f;
0223
0224   if (x <= 0.f)
0225     return 1.0f;
0226
0227   // for (0 < x < 1) || (x < a where a > 0)
0228   if ((x < 1.0f) || (x < a))
0229     return (1.0f - igam_impl<scalar_t>(a, x));
0230
0231   // for (x > 1) && (x > a where a > 0)
0232   return igamc_impl<scalar_t>(a, x);
0233 }
0234
0235 // for (x > 1) && (x > a where a > 0)
0236 template <typename scalar_t>
0237 TRACCC_HOST_DEVICE inline scalar_t igamc_impl(const scalar_t a,
0238                                               const scalar_t x) {
0239   scalar_t ans, ax, c, yc, r, t, y, z;
0240   scalar_t pk, pkm1, pkm2, qk, qkm1, qkm2;
0241
0242   ax = a * std::log(x) - x - lgam(a);
0243   if (ax < -log_gamma<scalar_t>::kMAXLOG)
0244     return (0.0f);
0245
0246   ax = std::exp(ax);
0247
0248   /* continued fraction */
0249   y = 1.0f - a;
0250   z = x + y + 1.0f;
0251   c = 0.0f;
0252   pkm2 = 1.0f;
0253   qkm2 = x;
0254   pkm1 = x + 1.0f;
0255   qkm1 = z * x;
0256   ans = pkm1 / qkm1;
0257
0258   do {
0259     c += 1.0f;
0260     y += 1.0f;
0261     z += 2.0f;
0262     yc = y * c;
0263     pk = pkm1 * z - pkm2 * yc;
0264     qk = qkm1 * z - qkm2 * yc;
0265     if (qk != 0.f) {
0266       r = pk / qk;
0267       t = std::abs((ans - r) / r);
0268       ans = r;
0269     } else {
0270       t = 1.0f;
0271     }
0272     pkm2 = pkm1;
0273     pkm1 = pk;
0274     qkm2 = qkm1;
0275     qkm1 = qk;
0276     if (std::abs(pk) > log_gamma<scalar_t>::kBig) {
0277       pkm2 *= log_gamma<scalar_t>::kBiginv;
0278       pkm1 *= log_gamma<scalar_t>::kBiginv;
0279       qkm2 *= log_gamma<scalar_t>::kBiginv;
0280       qkm1 *= log_gamma<scalar_t>::kBiginv;
0281     }
0282   } while (t > log_gamma<scalar_t>::kMACHEP);
0283
0284   return (ans * ax);
0285 }
0286
0287 template <typename scalar_t>
0288 TRACCC_HOST_DEVICE inline scalar_t lgam(scalar_t x) {
0289   scalar_t p, q, u, w, z;
0290   int i;
0291
0292   [[maybe_unused]] int sgngam = 1;
0293
0294   if (x >= std::numeric_limits<scalar_t>::infinity())
0295     return (std::numeric_limits<scalar_t>::infinity());
0296
0297   if (x < -34.0f) {
0298     q = -x;
0299
0300     // For x > 34
0301     w = lgam_impl<scalar_t>(q);
0302     p = std::floor(q);
0303     if (p == q)  //_unur_FP_same(p,q)
0304       return (std::numeric_limits<scalar_t>::infinity());
0305     i = static_cast<int>(p);
0306     if ((i & 1) == 0)
0307       sgngam = -1;
0308     else
0309       sgngam = 1;
0310     z = q - p;
0311     if (z > 0.5f) {
0312       p += 1.0f;
0313       z = p - q;
0314     }
0315     z = q * std::sin(constant<scalar_t>::pi * z);
0316     if (z == 0)
0317       return (std::numeric_limits<scalar_t>::infinity());
0318     /* z = log(ROOT::Math::Pi()) - log( z ) - w;*/
0319     z = std::log(constant<scalar_t>::pi) - math::log(z) - w;
0320     return (z);
0321   }
0322
0323   if (x < 13.0f) {
0324     z = 1.0f;
0325     p = 0.0f;
0326     u = x;
0327     while (u >= 3.0f) {
0328       p -= 1.0f;
0329       u = x + p;
0330       z *= u;
0331     }
0332     while (u < 2.0f) {
0333       if (u == 0.f)
0334         return (std::numeric_limits<scalar_t>::infinity());
0335       z /= u;
0336       p += 1.0f;
0337       u = x + p;
0338     }
0339     if (z < 0.0f) {
0340       sgngam = -1;
0341       z = -z;
0342     } else
0343       sgngam = 1;
0344     if (u == static_cast<scalar_t>(2.0))
0345       return (std::log(z));
0346     p -= 2.0f;
0347     x = x + p;
0348     p = x * Polynomialeval_B(x, 5) / Polynomialeval_C(x, 6);
0349     return (std::log(z) + p);
0350   }
0351
0352   return lgam_impl<scalar_t>(x);
0353 }
0354
0355 template <typename scalar_t>
0356 TRACCC_HOST_DEVICE inline scalar_t lgam_impl(scalar_t x) {
0357   scalar_t p, q;
0358   scalar_t sgngam = 1.f;
0359
0360   if (x > log_gamma<scalar_t>::kMAXLGM)
0361     return (sgngam * std::numeric_limits<scalar_t>::infinity());
0362
0363   q = (x - 0.5f) * std::log(x) - x + log_gamma<scalar_t>::LS2PI;
0364   if (x > 1.0e8f)
0365     return (q);
0366
0367   p = 1.0f / (x * x);
0368   if (x >= 1000.0f)
0369     q +=
0370         ((7.9365079365079365079365e-4f * p - 2.7777777777777777777778e-3f) * p +
0371          0.0833333333333333333333f) /
0372         x;
0373   else
0374     q += Polynomialeval_A(p, 4) / x;
0375   return (q);
0376 }
0377
0378 template <typename scalar_t>
0379 TRACCC_HOST_DEVICE inline scalar_t Polynomialeval_A(const scalar_t x,
0380                                                     const unsigned int N) {
0381   if (N == 0)
0382     return log_gamma<scalar_t>::A(0);
0383   else {
0384     scalar_t pom = log_gamma<scalar_t>::A(0);
0385     for (unsigned int i = 1; i <= N; i++)
0386       pom = pom * x + log_gamma<scalar_t>::A(i);
0387     return pom;
0388   }
0389 }
0390
0391 template <typename scalar_t>
0392 TRACCC_HOST_DEVICE inline scalar_t Polynomialeval_B(const scalar_t x,
0393                                                     const unsigned int N) {
0394   if (N == 0)
0395     return log_gamma<scalar_t>::B(0);
0396   else {
0397     scalar_t pom = log_gamma<scalar_t>::B(0);
0398     for (unsigned int i = 1; i <= N; i++)
0399       pom = pom * x + log_gamma<scalar_t>::B(i);
0400     return pom;
0401   }
0402 }
0403
0404 template <typename scalar_t>
0405 TRACCC_HOST_DEVICE inline scalar_t Polynomialeval_C(const scalar_t x,
0406                                                     const unsigned int N) {
0407   if (N == 0)
0408     return log_gamma<scalar_t>::C(0);
0409   else {
0410     scalar_t pom = x + log_gamma<scalar_t>::C(0);
0411     for (unsigned int i = 1; i < N; i++)
0412       pom = pom * x + log_gamma<scalar_t>::C(i);
0413     return pom;
0414   }
0415 }
0416
0417 }  // namespace traccc