gammainc_asy.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119
  1. """
  2. Precompute coefficients of Temme's asymptotic expansion for gammainc.
  3. This takes about 8 hours to run on a 2.3 GHz Macbook Pro with 4GB ram.
  4. Sources:
  5. [1] NIST, "Digital Library of Mathematical Functions",
  6. https://dlmf.nist.gov/
  7. """
  8. from __future__ import division, print_function, absolute_import
  9. import os
  10. from scipy.special._precompute.utils import lagrange_inversion
  11. try:
  12. import mpmath as mp
  13. except ImportError:
  14. pass
  15. def compute_a(n):
  16. """a_k from DLMF 5.11.6"""
  17. a = [mp.sqrt(2)/2]
  18. for k in range(1, n):
  19. ak = a[-1]/k
  20. for j in range(1, len(a)):
  21. ak -= a[j]*a[-j]/(j + 1)
  22. ak /= a[0]*(1 + mp.mpf(1)/(k + 1))
  23. a.append(ak)
  24. return a
  25. def compute_g(n):
  26. """g_k from DLMF 5.11.3/5.11.5"""
  27. a = compute_a(2*n)
  28. g = []
  29. for k in range(n):
  30. g.append(mp.sqrt(2)*mp.rf(0.5, k)*a[2*k])
  31. return g
  32. def eta(lam):
  33. """Function from DLMF 8.12.1 shifted to be centered at 0."""
  34. if lam > 0:
  35. return mp.sqrt(2*(lam - mp.log(lam + 1)))
  36. elif lam < 0:
  37. return -mp.sqrt(2*(lam - mp.log(lam + 1)))
  38. else:
  39. return 0
  40. def compute_alpha(n):
  41. """alpha_n from DLMF 8.12.13"""
  42. coeffs = mp.taylor(eta, 0, n - 1)
  43. return lagrange_inversion(coeffs)
  44. def compute_d(K, N):
  45. """d_{k, n} from DLMF 8.12.12"""
  46. M = N + 2*K
  47. d0 = [-mp.mpf(1)/3]
  48. alpha = compute_alpha(M + 2)
  49. for n in range(1, M):
  50. d0.append((n + 2)*alpha[n+2])
  51. d = [d0]
  52. g = compute_g(K)
  53. for k in range(1, K):
  54. dk = []
  55. for n in range(M - 2*k):
  56. dk.append((-1)**k*g[k]*d[0][n] + (n + 2)*d[k-1][n+2])
  57. d.append(dk)
  58. for k in range(K):
  59. d[k] = d[k][:N]
  60. return d
  61. header = \
  62. r"""/* This file was automatically generated by _precomp/gammainc.py.
  63. * Do not edit it manually!
  64. */
  65. #ifndef IGAM_H
  66. #define IGAM_H
  67. #define K {}
  68. #define N {}
  69. static const double d[K][N] =
  70. {{"""
  71. footer = \
  72. r"""
  73. #endif
  74. """
  75. def main():
  76. print(__doc__)
  77. K = 25
  78. N = 25
  79. with mp.workdps(50):
  80. d = compute_d(K, N)
  81. fn = os.path.join(os.path.dirname(__file__), '..', 'cephes', 'igam.h')
  82. with open(fn + '.new', 'w') as f:
  83. f.write(header.format(K, N))
  84. for k, row in enumerate(d):
  85. row = map(lambda x: mp.nstr(x, 17, min_fixed=0, max_fixed=0), row)
  86. f.write('{')
  87. f.write(", ".join(row))
  88. if k < K - 1:
  89. f.write('},\n')
  90. else:
  91. f.write('}};\n')
  92. f.write(footer)
  93. os.rename(fn + '.new', fn)
  94. if __name__ == "__main__":
  95. main()