github-featured/mpmath--mpmath
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Drop _jacobi_theta2()

Browse Source
This commit is contained in:
Sergey B Kirpichev
2026-06-22 12:20:16 +03:00
parent 1cfd36f49e
commit 9e890f67cd
1 changed files with 12 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mpmath/functions/theta.py
+12 -225
View File
@@ -1,220 +1,15 @@
from .functions import defun, defun_wrapped
@defun
def _jacobi_theta2(ctx, z, q):
extra1 = 10
extra2 = 20
def _djacobi_theta2(ctx, z, q, nd):
# the loops below break when the fixed precision quantities
# a and b go to zero;
# right shifting small negative numbers by wp one obtains -1, not zero,
# so the condition a**2 + b**2 > MIN is used to break the loops.
MIN = 2
if z == ctx.zero:
if (not ctx._im(q)):
wp = ctx.prec + extra1
x = ctx.to_fixed(ctx._re(q), wp)
x2 = (x*x) >> wp
a = b = x2
s = x2
while abs(a) > MIN:
b = (b*x2) >> wp
a = (a*b) >> wp
s += a
s = (1 << (wp+1)) + (s << 1)
s = ctx.ldexp(s, -wp)
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)
are = bre = x2re
aim = bim = x2im
sre = (1<<wp) + are
sim = aim
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 = (sre << 1)
sim = (sim << 1)
sre = ctx.ldexp(sre, -wp)
sim = ctx.ldexp(sim, -wp)
s = ctx.mpc(sre, sim)
else:
if (not ctx._im(q)) and (not ctx._im(z)):
wp = ctx.prec + extra1
x = ctx.to_fixed(ctx._re(q), wp)
x2 = (x*x) >> wp
a = b = x2
c1, s1 = ctx.cos_sin(ctx._re(z), prec=wp)
cn = c1 = ctx.to_fixed(c1, wp)
sn = s1 = ctx.to_fixed(s1, wp)
c2 = (c1*c1 - s1*s1) >> wp
s2 = (c1 * s1) >> (wp - 1)
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
s = c1 + ((a * cn) >> wp)
while abs(a) > MIN:
b = (b*x2) >> wp
a = (a*b) >> wp
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
s += (a * cn) >> wp
s = (s << 1)
s = ctx.ldexp(s, -wp)
s *= ctx.nthroot(q, 4)
return s
# 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 = bre = x2re
aim = bim = x2im
c1, s1 = ctx.cos_sin(ctx._re(z), prec=wp)
cn = c1 = ctx.to_fixed(c1, wp)
sn = s1 = ctx.to_fixed(s1, wp)
c2 = (c1*c1 - s1*s1) >> wp
s2 = (c1 * s1) >> (wp - 1)
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
sre = c1 + ((are * cn) >> wp)
sim = ((aim * cn) >> wp)
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
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> 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)
#case z complex, q real
elif not ctx._im(q):
wp = ctx.prec + extra2
x = ctx.to_fixed(ctx._re(q), wp)
x2 = (x*x) >> wp
a = b = x2
c1, s1 = ctx.cos_sin(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)
#c2 = (c1*c1 - s1*s1) >> wp
c2re = (c1re*c1re - c1im*c1im - s1re*s1re + s1im*s1im) >> wp
c2im = (c1re*c1im - s1re*s1im) >> (wp - 1)
#s2 = (c1 * s1) >> (wp - 1)
s2re = (c1re*s1re - c1im*s1im) >> (wp - 1)
s2im = (c1re*s1im + c1im*s1re) >> (wp - 1)
#cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
cnre = t1
cnim = t2
snre = t3
snim = t4
sre = c1re + ((a * cnre) >> wp)
sim = c1im + ((a * cnim) >> wp)
while abs(a) > MIN:
b = (b*x2) >> wp
a = (a*b) >> wp
t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> 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)
# 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 = bre = x2re
aim = bim = x2im
# cos(z), sin(z) with z complex
c1, s1 = ctx.cos_sin(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)
c2re = (c1re*c1re - c1im*c1im - s1re*s1re + s1im*s1im) >> wp
c2im = (c1re*c1im - s1re*s1im) >> (wp - 1)
s2re = (c1re*s1re - c1im*s1im) >> (wp - 1)
s2im = (c1re*s1im + c1im*s1re) >> (wp - 1)
t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
cnre = t1
cnim = t2
snre = t3
snim = t4
n = 1
termre = c1re
termim = c1im
sre = c1re + ((are * cnre - aim * cnim) >> wp)
sim = c1im + ((are * cnim + aim * cnre) >> wp)
n = 3
termre = ((are * cnre - aim * cnim) >> wp)
termim = ((are * cnim + aim * cnre) >> wp)
sre = c1re + ((are * cnre - aim * cnim) >> wp)
sim = c1im + ((are * cnim + aim * cnre) >> wp)
n = 5
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
#cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
t4 = (snre*c2im + snim*c2re + cnre*s2im + cnim*s2re) >> wp
cnre = t1
cnim = t2
snre = t3
snim = t4
termre = ((are * cnre - aim * cnim) >> wp)
termim = ((aim * cnre + are * cnim) >> wp)
sre += ((are * cnre - aim * cnim) >> wp)
sim += ((aim * cnre + are * cnim) >> wp)
n += 2
sre = (sre << 1)
sim = (sim << 1)
sre = ctx.ldexp(sre, -wp)
sim = ctx.ldexp(sim, -wp)
s = ctx.mpc(sre, sim)
s *= ctx.nthroot(q, 4)
return s
@defun
def _djacobi_theta2(ctx, z, q, nd):
if not nd:
return ctx._jacobi_theta2(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):
wp = ctx.prec + extra1
x = ctx.to_fixed(ctx._re(q), wp)
x2 = (x*x) >> wp
@@ -225,7 +20,7 @@ def _djacobi_theta2(ctx, z, q, nd):
c2 = (c1*c1 - s1*s1) >> wp
s2 = (c1 * s1) >> (wp - 1)
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
if (nd&1):
if nd&1:
s = s1 + ((a * sn * 3**nd) >> wp)
else:
s = c1 + ((a * cn * 3**nd) >> wp)
@@ -241,7 +36,7 @@ def _djacobi_theta2(ctx, z, q, nd):
n += 1
s = -(s << 1)
s = ctx.ldexp(s, -wp)
# case z real, q complex
# case z real, q complex
elif not ctx._im(z):
wp = ctx.prec + extra2
xre = ctx.to_fixed(ctx._re(q), wp)
@@ -256,7 +51,7 @@ def _djacobi_theta2(ctx, z, q, nd):
c2 = (c1*c1 - s1*s1) >> wp
s2 = (c1 * s1) >> (wp - 1)
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
if (nd&1):
if nd&1:
sre = s1 + ((are * sn * 3**nd) >> wp)
sim = ((aim * sn * 3**nd) >> wp)
else:
@@ -270,7 +65,7 @@ def _djacobi_theta2(ctx, z, q, nd):
(are * bim + aim * bre) >> wp
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
if (nd&1):
if nd&1:
sre += ((are * sn * n**nd) >> wp)
sim += ((aim * sn * n**nd) >> wp)
else:
@@ -282,7 +77,7 @@ def _djacobi_theta2(ctx, z, q, nd):
sre = ctx.ldexp(sre, -wp)
sim = ctx.ldexp(sim, -wp)
s = ctx.mpc(sre, sim)
#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)
@@ -293,13 +88,10 @@ def _djacobi_theta2(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)
#c2 = (c1*c1 - s1*s1) >> wp
c2re = (c1re*c1re - c1im*c1im - s1re*s1re + s1im*s1im) >> wp
c2im = (c1re*c1im - s1re*s1im) >> (wp - 1)
#s2 = (c1 * s1) >> (wp - 1)
s2re = (c1re*s1re - c1im*s1im) >> (wp - 1)
s2im = (c1re*s1im + c1im*s1re) >> (wp - 1)
#cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
@@ -308,7 +100,7 @@ def _djacobi_theta2(ctx, z, q, nd):
cnim = t2
snre = t3
snim = t4
if (nd&1):
if nd&1:
sre = s1re + ((a * snre * 3**nd) >> wp)
sim = s1im + ((a * snim * 3**nd) >> wp)
else:
@@ -326,7 +118,7 @@ def _djacobi_theta2(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:
@@ -347,7 +139,6 @@ def _djacobi_theta2(ctx, z, q, nd):
x2im = (xre*xim) >> (wp - 1)
are = bre = x2re
aim = bim = x2im
# cos(2*z), sin(2*z) with z complex
c1, s1 = ctx.cos_sin(z, prec=wp)
cnre = c1re = ctx.to_fixed(ctx._re(c1), wp)
cnim = c1im = ctx.to_fixed(ctx._im(c1), wp)
@@ -365,7 +156,7 @@ def _djacobi_theta2(ctx, z, q, nd):
cnim = t2
snre = t3
snim = t4
if (nd&1):
if nd&1:
sre = s1re + (((are * snre - aim * snim) * 3**nd) >> wp)
sim = s1im + (((are * snim + aim * snre)* 3**nd) >> wp)
else:
@@ -377,7 +168,6 @@ def _djacobi_theta2(ctx, z, q, nd):
(bre * x2im + bim * x2re) >> wp
are, aim = (are * bre - aim * bim) >> wp, \
(are * bim + aim * bre) >> wp
#cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
t1 = (cnre*c2re - cnim*c2im - snre*s2re + snim*s2im) >> wp
t2 = (cnre*c2im + cnim*c2re - snre*s2im - snim*s2re) >> wp
t3 = (snre*c2re - snim*c2im + cnre*s2re - cnim*s2im) >> wp
@@ -386,7 +176,7 @@ def _djacobi_theta2(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:
@@ -399,10 +189,7 @@ def _djacobi_theta2(ctx, z, q, nd):
sim = ctx.ldexp(sim, -wp)
s = ctx.mpc(sre, sim)
s *= ctx.nthroot(q, 4)
if (nd&1):
return (-1)**(nd//2) * s
else:
return (-1)**(1 + nd//2) * s
return (-1)**(1 - (nd&1) + nd//2) * s
@defun
def _djacobi_theta3(ctx, z, q, nd):