diff --git a/mpmath/functions/theta.py b/mpmath/functions/theta.py index 900011cf..3e1df0d0 100644 --- a/mpmath/functions/theta.py +++ b/mpmath/functions/theta.py @@ -404,193 +404,12 @@ def _djacobi_theta2(ctx, z, q, nd): else: return (-1)**(1 + nd//2) * s -@defun -def _jacobi_theta3(ctx, z, q): - extra1 = 10 - extra2 = 20 - MIN = 2 - if z == ctx.zero: - if not ctx._im(q): - wp = ctx.prec + extra1 - x = ctx.to_fixed(ctx._re(q), wp) - s = x - a = b = x - x2 = (x*x) >> wp - while abs(a) > MIN: - b = (b*x2) >> wp - a = (a*b) >> wp - s += a - s = (1 << wp) + (s << 1) - s = ctx.ldexp(s, -wp) - return s - else: - wp = ctx.prec + extra1 - xre = ctx.to_fixed(ctx._re(q), wp) - xim = ctx.to_fixed(ctx._im(q), wp) - x2re = (xre*xre - xim*xim) >> wp - x2im = (xre*xim) >> (wp - 1) - sre = are = bre = xre - sim = aim = bim = xim - while are**2 + aim**2 > MIN: - bre, bim = (bre * x2re - bim * x2im) >> wp, \ - (bre * x2im + bim * x2re) >> wp - are, aim = (are * bre - aim * bim) >> wp, \ - (are * bim + aim * bre) >> wp - sre += are - sim += aim - sre = (1 << wp) + (sre << 1) - sim = (sim << 1) - sre = ctx.ldexp(sre, -wp) - sim = ctx.ldexp(sim, -wp) - s = ctx.mpc(sre, sim) - return s - else: - if (not ctx._im(q)) and (not ctx._im(z)): - s = 0 - wp = ctx.prec + extra1 - x = ctx.to_fixed(ctx._re(q), wp) - a = (1 << wp) - b = x - x2 = (x*x) >> wp - c1, s1 = ctx.cos_sin(ctx._re(z)*2, prec=wp) - c1 = ctx.to_fixed(c1, wp) - s1 = ctx.to_fixed(s1, wp) - cn = c1 - sn = s1 - s += (a * cn) >> wp - while True: - b = (b*x2) >> wp - a = (a*b) >> wp - if abs(a) <= MIN: - break - cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp - s += (a * cn) >> wp - s = (s << 1) - s = ctx.ldexp(s, -wp) - return 1 + s*q - # case z real, q complex - elif not ctx._im(z): - wp = ctx.prec + extra2 - xre = ctx.to_fixed(ctx._re(q), wp) - xim = ctx.to_fixed(ctx._im(q), wp) - x2re = (xre*xre - xim*xim) >> wp - x2im = (xre*xim) >> (wp - 1) - are = (1 << wp) - aim = 0 - bre = xre - bim = xim - c1, s1 = ctx.cos_sin(ctx._re(z)*2, prec=wp) - c1 = ctx.to_fixed(c1, wp) - s1 = ctx.to_fixed(s1, wp) - cn = c1 - sn = s1 - sre = (are * cn) >> wp - sim = (aim * cn) >> wp - while True: - bre, bim = (bre * x2re - bim * x2im) >> wp, \ - (bre * x2im + bim * x2re) >> wp - are, aim = (are * bre - aim * bim) >> wp, \ - (are * bim + aim * bre) >> wp - if are**2 + aim**2 <= MIN: - break - cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp - sre += (are * cn) >> wp - sim += (aim * cn) >> wp - sre = (sre << 1) - sim = (sim << 1) - sre = ctx.ldexp(sre, -wp) - sim = ctx.ldexp(sim, -wp) - s = ctx.mpc(sre, sim) - return 1 + s*q - #case z complex, q real - elif not ctx._im(q): - wp = ctx.prec + extra2 - x = ctx.to_fixed(ctx._re(q), wp) - a = (1 << wp) - b = x - x2 = (x*x) >> wp - c1, s1 = ctx.cos_sin(2*z, prec=wp) - cnre = c1re = ctx.to_fixed(ctx._re(c1), wp) - cnim = c1im = ctx.to_fixed(ctx._im(c1), wp) - snre = s1re = ctx.to_fixed(ctx._re(s1), wp) - snim = s1im = ctx.to_fixed(ctx._im(s1), wp) - sre = (a * cnre) >> wp - sim = (a * cnim) >> wp - i = 1 - while True: - i+=1 - b = (b*x2) >> wp - a = (a*b) >> wp - if abs(a) <= MIN: - break - t1 = (cnre*c1re - cnim*c1im - snre*s1re + snim*s1im) >> wp - t2 = (cnre*c1im + cnim*c1re - snre*s1im - snim*s1re) >> wp - t3 = (snre*c1re - snim*c1im + cnre*s1re - cnim*s1im) >> wp - t4 = (snre*c1im + snim*c1re + cnre*s1im + cnim*s1re) >> wp - cnre = t1 - cnim = t2 - snre = t3 - snim = t4 - sre += (a * cnre) >> wp - sim += (a * cnim) >> wp - sre = (sre << 1) - sim = (sim << 1) - sre = ctx.ldexp(sre, -wp) - sim = ctx.ldexp(sim, -wp) - s = ctx.mpc(sre, sim) - return 1 + s*q - # case z and q complex - else: - wp = ctx.prec + extra2 - xre = ctx.to_fixed(ctx._re(q), wp) - xim = ctx.to_fixed(ctx._im(q), wp) - x2re = (xre*xre - xim*xim) >> wp - x2im = (xre*xim) >> (wp - 1) - are = (1 << wp) - aim = 0 - bre = xre - bim = xim - # cos(2*z), sin(2*z) with z complex - c1, s1 = ctx.cos_sin(2*z, prec=wp) - cnre = c1re = ctx.to_fixed(ctx._re(c1), wp) - cnim = c1im = ctx.to_fixed(ctx._im(c1), wp) - snre = s1re = ctx.to_fixed(ctx._re(s1), wp) - snim = s1im = ctx.to_fixed(ctx._im(s1), wp) - sre = (are * cnre - aim * cnim) >> wp - sim = (aim * cnre + are * cnim) >> wp - while True: - bre, bim = (bre * x2re - bim * x2im) >> wp, \ - (bre * x2im + bim * x2re) >> wp - are, aim = (are * bre - aim * bim) >> wp, \ - (are * bim + aim * bre) >> wp - if are**2 + aim**2 <= MIN: - break - t1 = (cnre*c1re - cnim*c1im - snre*s1re + snim*s1im) >> wp - t2 = (cnre*c1im + cnim*c1re - snre*s1im - snim*s1re) >> wp - t3 = (snre*c1re - snim*c1im + cnre*s1re - cnim*s1im) >> wp - t4 = (snre*c1im + snim*c1re + cnre*s1im + cnim*s1re) >> wp - cnre = t1 - cnim = t2 - snre = t3 - snim = t4 - sre += (are * cnre - aim * cnim) >> wp - sim += (aim * cnre + are * cnim) >> wp - sre = (sre << 1) - sim = (sim << 1) - sre = ctx.ldexp(sre, -wp) - sim = ctx.ldexp(sim, -wp) - s = ctx.mpc(sre, sim) - return 1 + s*q - @defun def _djacobi_theta3(ctx, z, q, nd): - """nd=1,2,3 order of the derivative with respect to z""" - if not nd: - return ctx._jacobi_theta3(z, q) MIN = 2 extra1 = 10 extra2 = 20 - if (not ctx._im(q)) and (not ctx._im(z)): + if not ctx._im(q) and not ctx._im(z): s = 0 wp = ctx.prec + extra1 x = ctx.to_fixed(ctx._re(q), wp) @@ -602,7 +421,7 @@ def _djacobi_theta3(ctx, z, q, nd): s1 = ctx.to_fixed(s1, wp) cn = c1 sn = s1 - if (nd&1): + if nd&1: s += (a * sn) >> wp else: s += (a * cn) >> wp @@ -636,7 +455,7 @@ def _djacobi_theta3(ctx, z, q, nd): s1 = ctx.to_fixed(s1, wp) cn = c1 sn = s1 - if (nd&1): + if nd&1: sre = (are * sn) >> wp sim = (aim * sn) >> wp else: @@ -663,7 +482,7 @@ def _djacobi_theta3(ctx, z, q, nd): sre = ctx.ldexp(sre, -wp) sim = ctx.ldexp(sim, -wp) s = ctx.mpc(sre, sim)*q - #case z complex, q real + # case z complex, q real elif not ctx._im(q): wp = ctx.prec + extra2 x = ctx.to_fixed(ctx._re(q), wp) @@ -675,7 +494,7 @@ def _djacobi_theta3(ctx, z, q, nd): cnim = c1im = ctx.to_fixed(ctx._im(c1), wp) snre = s1re = ctx.to_fixed(ctx._re(s1), wp) snim = s1im = ctx.to_fixed(ctx._im(s1), wp) - if (nd&1): + if nd&1: sre = (a * snre) >> wp sim = (a * snim) >> wp else: @@ -695,7 +514,7 @@ def _djacobi_theta3(ctx, z, q, nd): cnim = t2 snre = t3 snim = t4 - if (nd&1): + if nd&1: sre += (a * snre * n**nd) >> wp sim += (a * snim * n**nd) >> wp else: @@ -723,7 +542,7 @@ def _djacobi_theta3(ctx, z, q, nd): cnim = c1im = ctx.to_fixed(ctx._im(c1), wp) snre = s1re = ctx.to_fixed(ctx._re(s1), wp) snim = s1im = ctx.to_fixed(ctx._im(s1), wp) - if (nd&1): + if nd&1: sre = (are * snre - aim * snim) >> wp sim = (aim * snre + are * snim) >> wp else: @@ -745,7 +564,7 @@ def _djacobi_theta3(ctx, z, q, nd): cnim = t2 snre = t3 snim = t4 - if(nd&1): + if nd&1: sre += ((are * snre - aim * snim) * n**nd) >> wp sim += ((aim * snre + are * snim) * n**nd) >> wp else: @@ -757,10 +576,10 @@ def _djacobi_theta3(ctx, z, q, nd): sre = ctx.ldexp(sre, -wp) sim = ctx.ldexp(sim, -wp) s = ctx.mpc(sre, sim)*q - if (nd&1): + if nd&1: return (-1)**(nd//2) * s else: - return (-1)**(1 + nd//2) * s + return (-1)**(1 + nd//2) * s + (ctx.zero if nd else ctx.one) @defun def _djacobi_theta2a(ctx, z, q, nd):