Coverage for /usr/lib/python3/dist-packages/mpmath/functions/expintegrals.py: 15%

1from .functions import defun, defun_wrapped 

2 

3@defun_wrapped 

4def _erf_complex(ctx, z): 

5 z2 = ctx.square_exp_arg(z, -1) 

6 #z2 = -z**2 

7 v = (2/ctx.sqrt(ctx.pi))*z * ctx.hyp1f1((1,2),(3,2), z2) 

8 if not ctx._re(z): 

9 v = ctx._im(v)*ctx.j 

10 return v 

11 

12@defun_wrapped 

13def _erfc_complex(ctx, z): 

14 if ctx.re(z) > 2: 

15 z2 = ctx.square_exp_arg(z) 

16 nz2 = ctx.fneg(z2, exact=True) 

17 v = ctx.exp(nz2)/ctx.sqrt(ctx.pi) * ctx.hyperu((1,2),(1,2), z2) 

18 else: 

19 v = 1 - ctx._erf_complex(z) 

20 if not ctx._re(z): 

21 v = 1+ctx._im(v)*ctx.j 

22 return v 

23 

24@defun 

25def erf(ctx, z): 

26 z = ctx.convert(z) 

27 if ctx._is_real_type(z): 

28 try: 

29 return ctx._erf(z) 

30 except NotImplementedError: 

31 pass 

32 if ctx._is_complex_type(z) and not z.imag: 

33 try: 

34 return type(z)(ctx._erf(z.real)) 

35 except NotImplementedError: 

36 pass 

37 return ctx._erf_complex(z) 

38 

39@defun 

40def erfc(ctx, z): 

41 z = ctx.convert(z) 

42 if ctx._is_real_type(z): 

43 try: 

44 return ctx._erfc(z) 

45 except NotImplementedError: 

46 pass 

47 if ctx._is_complex_type(z) and not z.imag: 

48 try: 

49 return type(z)(ctx._erfc(z.real)) 

50 except NotImplementedError: 

51 pass 

52 return ctx._erfc_complex(z) 

53 

54@defun 

55def square_exp_arg(ctx, z, mult=1, reciprocal=False): 

56 prec = ctx.prec*4+20 

57 if reciprocal: 

58 z2 = ctx.fmul(z, z, prec=prec) 

59 z2 = ctx.fdiv(ctx.one, z2, prec=prec) 

60 else: 

61 z2 = ctx.fmul(z, z, prec=prec) 

62 if mult != 1: 

63 z2 = ctx.fmul(z2, mult, exact=True) 

64 return z2 

65 

66@defun_wrapped 

67def erfi(ctx, z): 

68 if not z: 

69 return z 

70 z2 = ctx.square_exp_arg(z) 

71 v = (2/ctx.sqrt(ctx.pi)*z) * ctx.hyp1f1((1,2), (3,2), z2) 

72 if not ctx._re(z): 

73 v = ctx._im(v)*ctx.j 

74 return v 

75 

76@defun_wrapped 

77def erfinv(ctx, x): 

78 xre = ctx._re(x) 

79 if (xre != x) or (xre < -1) or (xre > 1): 

80 return ctx.bad_domain("erfinv(x) is defined only for -1 <= x <= 1") 

81 x = xre 

82 #if ctx.isnan(x): return x 

83 if not x: return x 

84 if x == 1: return ctx.inf 

85 if x == -1: return ctx.ninf 

86 if abs(x) < 0.9: 

87 a = 0.53728*x**3 + 0.813198*x 

88 else: 

89 # An asymptotic formula 

90 u = ctx.ln(2/ctx.pi/(abs(x)-1)**2) 

91 a = ctx.sign(x) * ctx.sqrt(u - ctx.ln(u))/ctx.sqrt(2) 

92 ctx.prec += 10 

93 return ctx.findroot(lambda t: ctx.erf(t)-x, a) 

94 

95@defun_wrapped 

96def npdf(ctx, x, mu=0, sigma=1): 

97 sigma = ctx.convert(sigma) 

98 return ctx.exp(-(x-mu)**2/(2*sigma**2)) / (sigma*ctx.sqrt(2*ctx.pi)) 

99 

100@defun_wrapped 

101def ncdf(ctx, x, mu=0, sigma=1): 

102 a = (x-mu)/(sigma*ctx.sqrt(2)) 

103 if a < 0: 

104 return ctx.erfc(-a)/2 

105 else: 

106 return (1+ctx.erf(a))/2 

107 

108@defun_wrapped 

109def betainc(ctx, a, b, x1=0, x2=1, regularized=False): 

110 if x1 == x2: 

111 v = 0 

112 elif not x1: 

113 if x1 == 0 and x2 == 1: 

114 v = ctx.beta(a, b) 

115 else: 

116 v = x2**a * ctx.hyp2f1(a, 1-b, a+1, x2) / a 

117 else: 

118 m, d = ctx.nint_distance(a) 

119 if m <= 0: 

120 if d < -ctx.prec: 

121 h = +ctx.eps 

122 ctx.prec *= 2 

123 a += h 

124 elif d < -4: 

125 ctx.prec -= d 

126 s1 = x2**a * ctx.hyp2f1(a,1-b,a+1,x2) 

127 s2 = x1**a * ctx.hyp2f1(a,1-b,a+1,x1) 

128 v = (s1 - s2) / a 

129 if regularized: 

130 v /= ctx.beta(a,b) 

131 return v 

132 

133@defun 

134def gammainc(ctx, z, a=0, b=None, regularized=False): 

135 regularized = bool(regularized) 

136 z = ctx.convert(z) 

137 if a is None: 

138 a = ctx.zero 

139 lower_modified = False 

140 else: 

141 a = ctx.convert(a) 

142 lower_modified = a != ctx.zero 

143 if b is None: 

144 b = ctx.inf 

145 upper_modified = False 

146 else: 

147 b = ctx.convert(b) 

148 upper_modified = b != ctx.inf 

149 # Complete gamma function 

150 if not (upper_modified or lower_modified): 

151 if regularized: 

152 if ctx.re(z) < 0: 

153 return ctx.inf 

154 elif ctx.re(z) > 0: 

155 return ctx.one 

156 else: 

157 return ctx.nan 

158 return ctx.gamma(z) 

159 if a == b: 

160 return ctx.zero 

161 # Standardize 

162 if ctx.re(a) > ctx.re(b): 

163 return -ctx.gammainc(z, b, a, regularized) 

164 # Generalized gamma 

165 if upper_modified and lower_modified: 

166 return +ctx._gamma3(z, a, b, regularized) 

167 # Upper gamma 

168 elif lower_modified: 

169 return ctx._upper_gamma(z, a, regularized) 

170 # Lower gamma 

171 elif upper_modified: 

172 return ctx._lower_gamma(z, b, regularized) 

173 

174@defun 

175def _lower_gamma(ctx, z, b, regularized=False): 

176 # Pole 

177 if ctx.isnpint(z): 

178 return type(z)(ctx.inf) 

179 G = [z] * regularized 

180 negb = ctx.fneg(b, exact=True) 

181 def h(z): 

182 T1 = [ctx.exp(negb), b, z], [1, z, -1], [], G, [1], [1+z], b 

183 return (T1,) 

184 return ctx.hypercomb(h, [z]) 

185 

186@defun 

187def _upper_gamma(ctx, z, a, regularized=False): 

188 # Fast integer case, when available 

189 if ctx.isint(z): 

190 try: 

191 if regularized: 

192 # Gamma pole 

193 if ctx.isnpint(z): 

194 return type(z)(ctx.zero) 

195 orig = ctx.prec 

196 try: 

197 ctx.prec += 10 

198 return ctx._gamma_upper_int(z, a) / ctx.gamma(z) 

199 finally: 

200 ctx.prec = orig 

201 else: 

202 return ctx._gamma_upper_int(z, a) 

203 except NotImplementedError: 

204 pass 

205 # hypercomb is unable to detect the exact zeros, so handle them here 

206 if z == 2 and a == -1: 

207 return (z+a)*0 

208 if z == 3 and (a == -1-1j or a == -1+1j): 

209 return (z+a)*0 

210 nega = ctx.fneg(a, exact=True) 

211 G = [z] * regularized 

212 # Use 2F0 series when possible; fall back to lower gamma representation 

213 try: 

214 def h(z): 

215 r = z-1 

216 return [([ctx.exp(nega), a], [1, r], [], G, [1, -r], [], 1/nega)] 

217 return ctx.hypercomb(h, [z], force_series=True) 

218 except ctx.NoConvergence: 

219 def h(z): 

220 T1 = [], [1, z-1], [z], G, [], [], 0 

221 T2 = [-ctx.exp(nega), a, z], [1, z, -1], [], G, [1], [1+z], a 

222 return T1, T2 

223 return ctx.hypercomb(h, [z]) 

224 

225@defun 

226def _gamma3(ctx, z, a, b, regularized=False): 

227 pole = ctx.isnpint(z) 

228 if regularized and pole: 

229 return ctx.zero 

230 try: 

231 ctx.prec += 15 

232 # We don't know in advance whether it's better to write as a difference 

233 # of lower or upper gamma functions, so try both 

234 T1 = ctx.gammainc(z, a, regularized=regularized) 

235 T2 = ctx.gammainc(z, b, regularized=regularized) 

236 R = T1 - T2 

237 if ctx.mag(R) - max(ctx.mag(T1), ctx.mag(T2)) > -10: 

238 return R 

239 if not pole: 

240 T1 = ctx.gammainc(z, 0, b, regularized=regularized) 

241 T2 = ctx.gammainc(z, 0, a, regularized=regularized) 

242 R = T1 - T2 

243 # May be ok, but should probably at least print a warning 

244 # about possible cancellation 

245 if 1: #ctx.mag(R) - max(ctx.mag(T1), ctx.mag(T2)) > -10: 

246 return R 

247 finally: 

248 ctx.prec -= 15 

249 raise NotImplementedError 

250 

251@defun_wrapped 

252def expint(ctx, n, z): 

253 if ctx.isint(n) and ctx._is_real_type(z): 

254 try: 

255 return ctx._expint_int(n, z) 

256 except NotImplementedError: 

257 pass 

258 if ctx.isnan(n) or ctx.isnan(z): 

259 return z*n 

260 if z == ctx.inf: 

261 return 1/z 

262 if z == 0: 

263 # integral from 1 to infinity of t^n 

264 if ctx.re(n) <= 1: 

265 # TODO: reasonable sign of infinity 

266 return type(z)(ctx.inf) 

267 else: 

268 return ctx.one/(n-1) 

269 if n == 0: 

270 return ctx.exp(-z)/z 

271 if n == -1: 

272 return ctx.exp(-z)*(z+1)/z**2 

273 return z**(n-1) * ctx.gammainc(1-n, z) 

274 

275@defun_wrapped 

276def li(ctx, z, offset=False): 

277 if offset: 

278 if z == 2: 

279 return ctx.zero 

280 return ctx.ei(ctx.ln(z)) - ctx.ei(ctx.ln2) 

281 if not z: 

282 return z 

283 if z == 1: 

284 return ctx.ninf 

285 return ctx.ei(ctx.ln(z)) 

286 

287@defun 

288def ei(ctx, z): 

289 try: 

290 return ctx._ei(z) 

291 except NotImplementedError: 

292 return ctx._ei_generic(z) 

293 

294@defun_wrapped 

295def _ei_generic(ctx, z): 

296 # Note: the following is currently untested because mp and fp 

297 # both use special-case ei code 

298 if z == ctx.inf: 

299 return z 

300 if z == ctx.ninf: 

301 return ctx.zero 

302 if ctx.mag(z) > 1: 

303 try: 

304 r = ctx.one/z 

305 v = ctx.exp(z)*ctx.hyper([1,1],[],r, 

306 maxterms=ctx.prec, force_series=True)/z 

307 im = ctx._im(z) 

308 if im > 0: 

309 v += ctx.pi*ctx.j 

310 if im < 0: 

311 v -= ctx.pi*ctx.j 

312 return v 

313 except ctx.NoConvergence: 

314 pass 

315 v = z*ctx.hyp2f2(1,1,2,2,z) + ctx.euler 

316 if ctx._im(z): 

317 v += 0.5*(ctx.log(z) - ctx.log(ctx.one/z)) 

318 else: 

319 v += ctx.log(abs(z)) 

320 return v 

321 

322@defun 

323def e1(ctx, z): 

324 try: 

325 return ctx._e1(z) 

326 except NotImplementedError: 

327 return ctx.expint(1, z) 

328 

329@defun 

330def ci(ctx, z): 

331 try: 

332 return ctx._ci(z) 

333 except NotImplementedError: 

334 return ctx._ci_generic(z) 

335 

336@defun_wrapped 

337def _ci_generic(ctx, z): 

338 if ctx.isinf(z): 

339 if z == ctx.inf: return ctx.zero 

340 if z == ctx.ninf: return ctx.pi*1j 

341 jz = ctx.fmul(ctx.j,z,exact=True) 

342 njz = ctx.fneg(jz,exact=True) 

343 v = 0.5*(ctx.ei(jz) + ctx.ei(njz)) 

344 zreal = ctx._re(z) 

345 zimag = ctx._im(z) 

346 if zreal == 0: 

347 if zimag > 0: v += ctx.pi*0.5j 

348 if zimag < 0: v -= ctx.pi*0.5j 

349 if zreal < 0: 

350 if zimag >= 0: v += ctx.pi*1j 

351 if zimag < 0: v -= ctx.pi*1j 

352 if ctx._is_real_type(z) and zreal > 0: 

353 v = ctx._re(v) 

354 return v 

355 

356@defun 

357def si(ctx, z): 

358 try: 

359 return ctx._si(z) 

360 except NotImplementedError: 

361 return ctx._si_generic(z) 

362 

363@defun_wrapped 

364def _si_generic(ctx, z): 

365 if ctx.isinf(z): 

366 if z == ctx.inf: return 0.5*ctx.pi 

367 if z == ctx.ninf: return -0.5*ctx.pi 

368 # Suffers from cancellation near 0 

369 if ctx.mag(z) >= -1: 

370 jz = ctx.fmul(ctx.j,z,exact=True) 

371 njz = ctx.fneg(jz,exact=True) 

372 v = (-0.5j)*(ctx.ei(jz) - ctx.ei(njz)) 

373 zreal = ctx._re(z) 

374 if zreal > 0: 

375 v -= 0.5*ctx.pi 

376 if zreal < 0: 

377 v += 0.5*ctx.pi 

378 if ctx._is_real_type(z): 

379 v = ctx._re(v) 

380 return v 

381 else: 

382 return z*ctx.hyp1f2((1,2),(3,2),(3,2),-0.25*z*z) 

383 

384@defun_wrapped 

385def chi(ctx, z): 

386 nz = ctx.fneg(z, exact=True) 

387 v = 0.5*(ctx.ei(z) + ctx.ei(nz)) 

388 zreal = ctx._re(z) 

389 zimag = ctx._im(z) 

390 if zimag > 0: 

391 v += ctx.pi*0.5j 

392 elif zimag < 0: 

393 v -= ctx.pi*0.5j 

394 elif zreal < 0: 

395 v += ctx.pi*1j 

396 return v 

397 

398@defun_wrapped 

399def shi(ctx, z): 

400 # Suffers from cancellation near 0 

401 if ctx.mag(z) >= -1: 

402 nz = ctx.fneg(z, exact=True) 

403 v = 0.5*(ctx.ei(z) - ctx.ei(nz)) 

404 zimag = ctx._im(z) 

405 if zimag > 0: v -= 0.5j*ctx.pi 

406 if zimag < 0: v += 0.5j*ctx.pi 

407 return v 

408 else: 

409 return z * ctx.hyp1f2((1,2),(3,2),(3,2),0.25*z*z) 

410 

411@defun_wrapped 

412def fresnels(ctx, z): 

413 if z == ctx.inf: 

414 return ctx.mpf(0.5) 

415 if z == ctx.ninf: 

416 return ctx.mpf(-0.5) 

417 return ctx.pi*z**3/6*ctx.hyp1f2((3,4),(3,2),(7,4),-ctx.pi**2*z**4/16) 

418 

419@defun_wrapped 

420def fresnelc(ctx, z): 

421 if z == ctx.inf: 

422 return ctx.mpf(0.5) 

423 if z == ctx.ninf: 

424 return ctx.mpf(-0.5) 

425 return z*ctx.hyp1f2((1,4),(1,2),(5,4),-ctx.pi**2*z**4/16)