123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672 |
- """
- This module complements the math and cmath builtin modules by providing
- fast machine precision versions of some additional functions (gamma, ...)
- and wrapping math/cmath functions so that they can be called with either
- real or complex arguments.
- """
- import operator
- import math
- import cmath
- # Irrational (?) constants
- pi = 3.1415926535897932385
- e = 2.7182818284590452354
- sqrt2 = 1.4142135623730950488
- sqrt5 = 2.2360679774997896964
- phi = 1.6180339887498948482
- ln2 = 0.69314718055994530942
- ln10 = 2.302585092994045684
- euler = 0.57721566490153286061
- catalan = 0.91596559417721901505
- khinchin = 2.6854520010653064453
- apery = 1.2020569031595942854
- logpi = 1.1447298858494001741
- def _mathfun_real(f_real, f_complex):
- def f(x, **kwargs):
- if type(x) is float:
- return f_real(x)
- if type(x) is complex:
- return f_complex(x)
- try:
- x = float(x)
- return f_real(x)
- except (TypeError, ValueError):
- x = complex(x)
- return f_complex(x)
- f.__name__ = f_real.__name__
- return f
- def _mathfun(f_real, f_complex):
- def f(x, **kwargs):
- if type(x) is complex:
- return f_complex(x)
- try:
- return f_real(float(x))
- except (TypeError, ValueError):
- return f_complex(complex(x))
- f.__name__ = f_real.__name__
- return f
- def _mathfun_n(f_real, f_complex):
- def f(*args, **kwargs):
- try:
- return f_real(*(float(x) for x in args))
- except (TypeError, ValueError):
- return f_complex(*(complex(x) for x in args))
- f.__name__ = f_real.__name__
- return f
- # Workaround for non-raising log and sqrt in Python 2.5 and 2.4
- # on Unix system
- try:
- math.log(-2.0)
- def math_log(x):
- if x <= 0.0:
- raise ValueError("math domain error")
- return math.log(x)
- def math_sqrt(x):
- if x < 0.0:
- raise ValueError("math domain error")
- return math.sqrt(x)
- except (ValueError, TypeError):
- math_log = math.log
- math_sqrt = math.sqrt
- pow = _mathfun_n(operator.pow, lambda x, y: complex(x)**y)
- log = _mathfun_n(math_log, cmath.log)
- sqrt = _mathfun(math_sqrt, cmath.sqrt)
- exp = _mathfun_real(math.exp, cmath.exp)
- cos = _mathfun_real(math.cos, cmath.cos)
- sin = _mathfun_real(math.sin, cmath.sin)
- tan = _mathfun_real(math.tan, cmath.tan)
- acos = _mathfun(math.acos, cmath.acos)
- asin = _mathfun(math.asin, cmath.asin)
- atan = _mathfun_real(math.atan, cmath.atan)
- cosh = _mathfun_real(math.cosh, cmath.cosh)
- sinh = _mathfun_real(math.sinh, cmath.sinh)
- tanh = _mathfun_real(math.tanh, cmath.tanh)
- floor = _mathfun_real(math.floor,
- lambda z: complex(math.floor(z.real), math.floor(z.imag)))
- ceil = _mathfun_real(math.ceil,
- lambda z: complex(math.ceil(z.real), math.ceil(z.imag)))
- cos_sin = _mathfun_real(lambda x: (math.cos(x), math.sin(x)),
- lambda z: (cmath.cos(z), cmath.sin(z)))
- cbrt = _mathfun(lambda x: x**(1./3), lambda z: z**(1./3))
- def nthroot(x, n):
- r = 1./n
- try:
- return float(x) ** r
- except (ValueError, TypeError):
- return complex(x) ** r
- def _sinpi_real(x):
- if x < 0:
- return -_sinpi_real(-x)
- n, r = divmod(x, 0.5)
- r *= pi
- n %= 4
- if n == 0: return math.sin(r)
- if n == 1: return math.cos(r)
- if n == 2: return -math.sin(r)
- if n == 3: return -math.cos(r)
- def _cospi_real(x):
- if x < 0:
- x = -x
- n, r = divmod(x, 0.5)
- r *= pi
- n %= 4
- if n == 0: return math.cos(r)
- if n == 1: return -math.sin(r)
- if n == 2: return -math.cos(r)
- if n == 3: return math.sin(r)
- def _sinpi_complex(z):
- if z.real < 0:
- return -_sinpi_complex(-z)
- n, r = divmod(z.real, 0.5)
- z = pi*complex(r, z.imag)
- n %= 4
- if n == 0: return cmath.sin(z)
- if n == 1: return cmath.cos(z)
- if n == 2: return -cmath.sin(z)
- if n == 3: return -cmath.cos(z)
- def _cospi_complex(z):
- if z.real < 0:
- z = -z
- n, r = divmod(z.real, 0.5)
- z = pi*complex(r, z.imag)
- n %= 4
- if n == 0: return cmath.cos(z)
- if n == 1: return -cmath.sin(z)
- if n == 2: return -cmath.cos(z)
- if n == 3: return cmath.sin(z)
- cospi = _mathfun_real(_cospi_real, _cospi_complex)
- sinpi = _mathfun_real(_sinpi_real, _sinpi_complex)
- def tanpi(x):
- try:
- return sinpi(x) / cospi(x)
- except OverflowError:
- if complex(x).imag > 10:
- return 1j
- if complex(x).imag < 10:
- return -1j
- raise
- def cotpi(x):
- try:
- return cospi(x) / sinpi(x)
- except OverflowError:
- if complex(x).imag > 10:
- return -1j
- if complex(x).imag < 10:
- return 1j
- raise
- INF = 1e300*1e300
- NINF = -INF
- NAN = INF-INF
- EPS = 2.2204460492503131e-16
- _exact_gamma = (INF, 1.0, 1.0, 2.0, 6.0, 24.0, 120.0, 720.0, 5040.0, 40320.0,
- 362880.0, 3628800.0, 39916800.0, 479001600.0, 6227020800.0, 87178291200.0,
- 1307674368000.0, 20922789888000.0, 355687428096000.0, 6402373705728000.0,
- 121645100408832000.0, 2432902008176640000.0)
- _max_exact_gamma = len(_exact_gamma)-1
- # Lanczos coefficients used by the GNU Scientific Library
- _lanczos_g = 7
- _lanczos_p = (0.99999999999980993, 676.5203681218851, -1259.1392167224028,
- 771.32342877765313, -176.61502916214059, 12.507343278686905,
- -0.13857109526572012, 9.9843695780195716e-6, 1.5056327351493116e-7)
- def _gamma_real(x):
- _intx = int(x)
- if _intx == x:
- if _intx <= 0:
- #return (-1)**_intx * INF
- raise ZeroDivisionError("gamma function pole")
- if _intx <= _max_exact_gamma:
- return _exact_gamma[_intx]
- if x < 0.5:
- # TODO: sinpi
- return pi / (_sinpi_real(x)*_gamma_real(1-x))
- else:
- x -= 1.0
- r = _lanczos_p[0]
- for i in range(1, _lanczos_g+2):
- r += _lanczos_p[i]/(x+i)
- t = x + _lanczos_g + 0.5
- return 2.506628274631000502417 * t**(x+0.5) * math.exp(-t) * r
- def _gamma_complex(x):
- if not x.imag:
- return complex(_gamma_real(x.real))
- if x.real < 0.5:
- # TODO: sinpi
- return pi / (_sinpi_complex(x)*_gamma_complex(1-x))
- else:
- x -= 1.0
- r = _lanczos_p[0]
- for i in range(1, _lanczos_g+2):
- r += _lanczos_p[i]/(x+i)
- t = x + _lanczos_g + 0.5
- return 2.506628274631000502417 * t**(x+0.5) * cmath.exp(-t) * r
- gamma = _mathfun_real(_gamma_real, _gamma_complex)
- def rgamma(x):
- try:
- return 1./gamma(x)
- except ZeroDivisionError:
- return x*0.0
- def factorial(x):
- return gamma(x+1.0)
- def arg(x):
- if type(x) is float:
- return math.atan2(0.0,x)
- return math.atan2(x.imag,x.real)
- # XXX: broken for negatives
- def loggamma(x):
- if type(x) not in (float, complex):
- try:
- x = float(x)
- except (ValueError, TypeError):
- x = complex(x)
- try:
- xreal = x.real
- ximag = x.imag
- except AttributeError: # py2.5
- xreal = x
- ximag = 0.0
- # Reflection formula
- # http://functions.wolfram.com/GammaBetaErf/LogGamma/16/01/01/0003/
- if xreal < 0.0:
- if abs(x) < 0.5:
- v = log(gamma(x))
- if ximag == 0:
- v = v.conjugate()
- return v
- z = 1-x
- try:
- re = z.real
- im = z.imag
- except AttributeError: # py2.5
- re = z
- im = 0.0
- refloor = floor(re)
- if im == 0.0:
- imsign = 0
- elif im < 0.0:
- imsign = -1
- else:
- imsign = 1
- return (-pi*1j)*abs(refloor)*(1-abs(imsign)) + logpi - \
- log(sinpi(z-refloor)) - loggamma(z) + 1j*pi*refloor*imsign
- if x == 1.0 or x == 2.0:
- return x*0
- p = 0.
- while abs(x) < 11:
- p -= log(x)
- x += 1.0
- s = 0.918938533204672742 + (x-0.5)*log(x) - x
- r = 1./x
- r2 = r*r
- s += 0.083333333333333333333*r; r *= r2
- s += -0.0027777777777777777778*r; r *= r2
- s += 0.00079365079365079365079*r; r *= r2
- s += -0.0005952380952380952381*r; r *= r2
- s += 0.00084175084175084175084*r; r *= r2
- s += -0.0019175269175269175269*r; r *= r2
- s += 0.0064102564102564102564*r; r *= r2
- s += -0.02955065359477124183*r
- return s + p
- _psi_coeff = [
- 0.083333333333333333333,
- -0.0083333333333333333333,
- 0.003968253968253968254,
- -0.0041666666666666666667,
- 0.0075757575757575757576,
- -0.021092796092796092796,
- 0.083333333333333333333,
- -0.44325980392156862745,
- 3.0539543302701197438,
- -26.456212121212121212]
- def _digamma_real(x):
- _intx = int(x)
- if _intx == x:
- if _intx <= 0:
- raise ZeroDivisionError("polygamma pole")
- if x < 0.5:
- x = 1.0-x
- s = pi*cotpi(x)
- else:
- s = 0.0
- while x < 10.0:
- s -= 1.0/x
- x += 1.0
- x2 = x**-2
- t = x2
- for c in _psi_coeff:
- s -= c*t
- if t < 1e-20:
- break
- t *= x2
- return s + math_log(x) - 0.5/x
- def _digamma_complex(x):
- if not x.imag:
- return complex(_digamma_real(x.real))
- if x.real < 0.5:
- x = 1.0-x
- s = pi*cotpi(x)
- else:
- s = 0.0
- while abs(x) < 10.0:
- s -= 1.0/x
- x += 1.0
- x2 = x**-2
- t = x2
- for c in _psi_coeff:
- s -= c*t
- if abs(t) < 1e-20:
- break
- t *= x2
- return s + cmath.log(x) - 0.5/x
- digamma = _mathfun_real(_digamma_real, _digamma_complex)
- # TODO: could implement complex erf and erfc here. Need
- # to find an accurate method (avoiding cancellation)
- # for approx. 1 < abs(x) < 9.
- _erfc_coeff_P = [
- 1.0000000161203922312,
- 2.1275306946297962644,
- 2.2280433377390253297,
- 1.4695509105618423961,
- 0.66275911699770787537,
- 0.20924776504163751585,
- 0.045459713768411264339,
- 0.0063065951710717791934,
- 0.00044560259661560421715][::-1]
- _erfc_coeff_Q = [
- 1.0000000000000000000,
- 3.2559100272784894318,
- 4.9019435608903239131,
- 4.4971472894498014205,
- 2.7845640601891186528,
- 1.2146026030046904138,
- 0.37647108453729465912,
- 0.080970149639040548613,
- 0.011178148899483545902,
- 0.00078981003831980423513][::-1]
- def _polyval(coeffs, x):
- p = coeffs[0]
- for c in coeffs[1:]:
- p = c + x*p
- return p
- def _erf_taylor(x):
- # Taylor series assuming 0 <= x <= 1
- x2 = x*x
- s = t = x
- n = 1
- while abs(t) > 1e-17:
- t *= x2/n
- s -= t/(n+n+1)
- n += 1
- t *= x2/n
- s += t/(n+n+1)
- n += 1
- return 1.1283791670955125739*s
- def _erfc_mid(x):
- # Rational approximation assuming 0 <= x <= 9
- return exp(-x*x)*_polyval(_erfc_coeff_P,x)/_polyval(_erfc_coeff_Q,x)
- def _erfc_asymp(x):
- # Asymptotic expansion assuming x >= 9
- x2 = x*x
- v = exp(-x2)/x*0.56418958354775628695
- r = t = 0.5 / x2
- s = 1.0
- for n in range(1,22,4):
- s -= t
- t *= r * (n+2)
- s += t
- t *= r * (n+4)
- if abs(t) < 1e-17:
- break
- return s * v
- def erf(x):
- """
- erf of a real number.
- """
- x = float(x)
- if x != x:
- return x
- if x < 0.0:
- return -erf(-x)
- if x >= 1.0:
- if x >= 6.0:
- return 1.0
- return 1.0 - _erfc_mid(x)
- return _erf_taylor(x)
- def erfc(x):
- """
- erfc of a real number.
- """
- x = float(x)
- if x != x:
- return x
- if x < 0.0:
- if x < -6.0:
- return 2.0
- return 2.0-erfc(-x)
- if x > 9.0:
- return _erfc_asymp(x)
- if x >= 1.0:
- return _erfc_mid(x)
- return 1.0 - _erf_taylor(x)
- gauss42 = [\
- (0.99839961899006235, 0.0041059986046490839),
- (-0.99839961899006235, 0.0041059986046490839),
- (0.9915772883408609, 0.009536220301748501),
- (-0.9915772883408609,0.009536220301748501),
- (0.97934250806374812, 0.014922443697357493),
- (-0.97934250806374812, 0.014922443697357493),
- (0.96175936533820439,0.020227869569052644),
- (-0.96175936533820439, 0.020227869569052644),
- (0.93892355735498811, 0.025422959526113047),
- (-0.93892355735498811,0.025422959526113047),
- (0.91095972490412735, 0.030479240699603467),
- (-0.91095972490412735, 0.030479240699603467),
- (0.87802056981217269,0.03536907109759211),
- (-0.87802056981217269, 0.03536907109759211),
- (0.8402859832618168, 0.040065735180692258),
- (-0.8402859832618168,0.040065735180692258),
- (0.7979620532554873, 0.044543577771965874),
- (-0.7979620532554873, 0.044543577771965874),
- (0.75127993568948048,0.048778140792803244),
- (-0.75127993568948048, 0.048778140792803244),
- (0.70049459055617114, 0.052746295699174064),
- (-0.70049459055617114,0.052746295699174064),
- (0.64588338886924779, 0.056426369358018376),
- (-0.64588338886924779, 0.056426369358018376),
- (0.58774459748510932, 0.059798262227586649),
- (-0.58774459748510932, 0.059798262227586649),
- (0.5263957499311922, 0.062843558045002565),
- (-0.5263957499311922, 0.062843558045002565),
- (0.46217191207042191, 0.065545624364908975),
- (-0.46217191207042191, 0.065545624364908975),
- (0.39542385204297503, 0.067889703376521934),
- (-0.39542385204297503, 0.067889703376521934),
- (0.32651612446541151, 0.069862992492594159),
- (-0.32651612446541151, 0.069862992492594159),
- (0.25582507934287907, 0.071454714265170971),
- (-0.25582507934287907, 0.071454714265170971),
- (0.18373680656485453, 0.072656175243804091),
- (-0.18373680656485453, 0.072656175243804091),
- (0.11064502720851986, 0.073460813453467527),
- (-0.11064502720851986, 0.073460813453467527),
- (0.036948943165351772, 0.073864234232172879),
- (-0.036948943165351772, 0.073864234232172879)]
- EI_ASYMP_CONVERGENCE_RADIUS = 40.0
- def ei_asymp(z, _e1=False):
- r = 1./z
- s = t = 1.0
- k = 1
- while 1:
- t *= k*r
- s += t
- if abs(t) < 1e-16:
- break
- k += 1
- v = s*exp(z)/z
- if _e1:
- if type(z) is complex:
- zreal = z.real
- zimag = z.imag
- else:
- zreal = z
- zimag = 0.0
- if zimag == 0.0 and zreal > 0.0:
- v += pi*1j
- else:
- if type(z) is complex:
- if z.imag > 0:
- v += pi*1j
- if z.imag < 0:
- v -= pi*1j
- return v
- def ei_taylor(z, _e1=False):
- s = t = z
- k = 2
- while 1:
- t = t*z/k
- term = t/k
- if abs(term) < 1e-17:
- break
- s += term
- k += 1
- s += euler
- if _e1:
- s += log(-z)
- else:
- if type(z) is float or z.imag == 0.0:
- s += math_log(abs(z))
- else:
- s += cmath.log(z)
- return s
- def ei(z, _e1=False):
- typez = type(z)
- if typez not in (float, complex):
- try:
- z = float(z)
- typez = float
- except (TypeError, ValueError):
- z = complex(z)
- typez = complex
- if not z:
- return -INF
- absz = abs(z)
- if absz > EI_ASYMP_CONVERGENCE_RADIUS:
- return ei_asymp(z, _e1)
- elif absz <= 2.0 or (typez is float and z > 0.0):
- return ei_taylor(z, _e1)
- # Integrate, starting from whichever is smaller of a Taylor
- # series value or an asymptotic series value
- if typez is complex and z.real > 0.0:
- zref = z / absz
- ref = ei_taylor(zref, _e1)
- else:
- zref = EI_ASYMP_CONVERGENCE_RADIUS * z / absz
- ref = ei_asymp(zref, _e1)
- C = (zref-z)*0.5
- D = (zref+z)*0.5
- s = 0.0
- if type(z) is complex:
- _exp = cmath.exp
- else:
- _exp = math.exp
- for x,w in gauss42:
- t = C*x+D
- s += w*_exp(t)/t
- ref -= C*s
- return ref
- def e1(z):
- # hack to get consistent signs if the imaginary part if 0
- # and signed
- typez = type(z)
- if type(z) not in (float, complex):
- try:
- z = float(z)
- typez = float
- except (TypeError, ValueError):
- z = complex(z)
- typez = complex
- if typez is complex and not z.imag:
- z = complex(z.real, 0.0)
- # end hack
- return -ei(-z, _e1=True)
- _zeta_int = [\
- -0.5,
- 0.0,
- 1.6449340668482264365,1.2020569031595942854,1.0823232337111381915,
- 1.0369277551433699263,1.0173430619844491397,1.0083492773819228268,
- 1.0040773561979443394,1.0020083928260822144,1.0009945751278180853,
- 1.0004941886041194646,1.0002460865533080483,1.0001227133475784891,
- 1.0000612481350587048,1.0000305882363070205,1.0000152822594086519,
- 1.0000076371976378998,1.0000038172932649998,1.0000019082127165539,
- 1.0000009539620338728,1.0000004769329867878,1.0000002384505027277,
- 1.0000001192199259653,1.0000000596081890513,1.0000000298035035147,
- 1.0000000149015548284]
- _zeta_P = [-3.50000000087575873, -0.701274355654678147,
- -0.0672313458590012612, -0.00398731457954257841,
- -0.000160948723019303141, -4.67633010038383371e-6,
- -1.02078104417700585e-7, -1.68030037095896287e-9,
- -1.85231868742346722e-11][::-1]
- _zeta_Q = [1.00000000000000000, -0.936552848762465319,
- -0.0588835413263763741, -0.00441498861482948666,
- -0.000143416758067432622, -5.10691659585090782e-6,
- -9.58813053268913799e-8, -1.72963791443181972e-9,
- -1.83527919681474132e-11][::-1]
- _zeta_1 = [3.03768838606128127e-10, -1.21924525236601262e-8,
- 2.01201845887608893e-7, -1.53917240683468381e-6,
- -5.09890411005967954e-7, 0.000122464707271619326,
- -0.000905721539353130232, -0.00239315326074843037,
- 0.084239750013159168, 0.418938517907442414, 0.500000001921884009]
- _zeta_0 = [-3.46092485016748794e-10, -6.42610089468292485e-9,
- 1.76409071536679773e-7, -1.47141263991560698e-6, -6.38880222546167613e-7,
- 0.000122641099800668209, -0.000905894913516772796, -0.00239303348507992713,
- 0.0842396947501199816, 0.418938533204660256, 0.500000000000000052]
- def zeta(s):
- """
- Riemann zeta function, real argument
- """
- if not isinstance(s, (float, int)):
- try:
- s = float(s)
- except (ValueError, TypeError):
- try:
- s = complex(s)
- if not s.imag:
- return complex(zeta(s.real))
- except (ValueError, TypeError):
- pass
- raise NotImplementedError
- if s == 1:
- raise ValueError("zeta(1) pole")
- if s >= 27:
- return 1.0 + 2.0**(-s) + 3.0**(-s)
- n = int(s)
- if n == s:
- if n >= 0:
- return _zeta_int[n]
- if not (n % 2):
- return 0.0
- if s <= 0.0:
- return 2.**s*pi**(s-1)*_sinpi_real(0.5*s)*_gamma_real(1-s)*zeta(1-s)
- if s <= 2.0:
- if s <= 1.0:
- return _polyval(_zeta_0,s)/(s-1)
- return _polyval(_zeta_1,s)/(s-1)
- z = _polyval(_zeta_P,s) / _polyval(_zeta_Q,s)
- return 1.0 + 2.0**(-s) + 3.0**(-s) + 4.0**(-s)*z
|