def arg(x):
"""Returns the complex argument (phase) of x. The returned value is
an mpf instance. The argument is here defined to satisfy
-pi < arg(x) <= pi. On the negative real half-axis, it is taken to
be +pi."""
x = mpc(x)
return atan2(x.imag, x.real)
def cplot(f,
re=[-5, 5],
im=[-5, 5],
points=2000,
color=default_color_function,
verbose=False,
file=None,
dpi=None):
"""
Plots the given complex-valued function over a rectangular part
of the complex plane given by the pairs of intervals re and im.
For example:
cplot(lambda z: z, [-2, 2], [-10, 10])
cplot(exp)
cplot(zeta, [0, 1], [0, 50])
By default, the complex argument (phase) is shown as color and
the magnitude is show as brightness. You can also supply a
custom color function ('color'). This function should take a
complex number as input and return an RGB 3-tuple containing
floats in the range 0.0-1.0.
To obtain a sharp image, the number of points may need to be
increased to 100,000 or thereabout. Since evaluating the
function that many times is likely to be slow, the 'verbose'
option is useful to display progress.
"""
import pylab
pylab.clf()
rea, reb = re
ima, imb = im
dre = reb - rea
dim = imb - ima
M = int(sqrt(points * dre / dim) + 1)
N = int(sqrt(points * dim / dre) + 1)
x = pylab.linspace(rea, reb, M)
y = pylab.linspace(ima, imb, N)
# Note: we have to be careful to get the right rotation.
# Test with these plots:
# cplot(lambda z: z if z.real < 0 else 0)
# cplot(lambda z: z if z.imag < 0 else 0)
w = pylab.zeros((N, M, 3))
for n in xrange(N):
for m in xrange(M):
z = mpc(x[m], y[n])
try:
v = color(f(z))
except plot_ignore:
v = (0.5, 0.5, 0.5)
w[n, m] = v
if verbose:
print n, "of", N
pylab.imshow(w, extent=(rea, reb, ima, imb), origin='lower')
pylab.xlabel('Re(z)')
pylab.ylabel('Im(z)')
if file:
pylab.savefig(file, dpi=dpi)
else:
pylab.show()
def arg(x):
"""Returns the complex argument (phase) of x. The returned value is
an mpf instance. The argument is here defined to satisfy
-pi < arg(x) <= pi. On the negative real half-axis, it is taken to
be +pi."""
x = mpc(x)
return atan2(x.imag, x.real)
def cplot(f, re=[-5,5], im=[-5,5], points=2000, color=default_color_function,
verbose=False, file=None, dpi=None):
"""
Plots the given complex-valued function *f* over a rectangular part
of the complex plane specified by the pairs of intervals *re* and *im*.
For example::
cplot(lambda z: z, [-2, 2], [-10, 10])
cplot(exp)
cplot(zeta, [0, 1], [0, 50])
By default, the complex argument (phase) is shown as color (hue) and
the magnitude is show as brightness. You can also supply a
custom color function (*color*). This function should take a
complex number as input and return an RGB 3-tuple containing
floats in the range 0.0-1.0.
To obtain a sharp image, the number of points may need to be
increased to 100,000 or thereabout. Since evaluating the
function that many times is likely to be slow, the 'verbose'
option is useful to display progress.
NOTE: This function requires matplotlib (pylab).
"""
import pylab
pylab.clf()
rea, reb = re
ima, imb = im
dre = reb - rea
dim = imb - ima
M = int(sqrt(points*dre/dim)+1)
N = int(sqrt(points*dim/dre)+1)
x = pylab.linspace(rea, reb, M)
y = pylab.linspace(ima, imb, N)
# Note: we have to be careful to get the right rotation.
# Test with these plots:
# cplot(lambda z: z if z.real < 0 else 0)
# cplot(lambda z: z if z.imag < 0 else 0)
w = pylab.zeros((N, M, 3))
for n in xrange(N):
for m in xrange(M):
z = mpc(x[m], y[n])
try:
v = color(f(z))
except plot_ignore:
v = (0.5, 0.5, 0.5)
w[n,m] = v
if verbose:
print n, "of", N
pylab.imshow(w, extent=(rea, reb, ima, imb), origin='lower')
pylab.xlabel('Re(z)')
pylab.ylabel('Im(z)')
if file:
pylab.savefig(file, dpi=dpi)
else:
pylab.show()
def _jacobi_theta2(z, q):
extra1 = 10
extra2 = 20
# 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 == zero:
if isinstance(q, mpf):
wp = mp.prec + extra1
x = to_fixed(q._mpf_, 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 = mpf(from_man_exp(s, -wp, mp.prec, 'n'))
else:
wp = mp.prec + extra1
xre, xim = q._mpc_
xre = to_fixed(xre, wp)
xim = to_fixed(xim, 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 = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
else:
if isinstance(q, mpf) and isinstance(z, mpf):
wp = mp.prec + extra1
x = to_fixed(q._mpf_, wp)
x2 = (x*x) >> wp
a = b = x2
c1, s1 = cos_sin(z._mpf_, wp)
cn = c1 = to_fixed(c1, wp)
sn = s1 = 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 = mpf(from_man_exp(s, -wp, mp.prec, 'n'))
s *= nthroot(q, 4)
return s
# case z real, q complex
elif isinstance(z, mpf):
wp = mp.prec + extra2
xre, xim = q._mpc_
xre = to_fixed(xre, wp)
xim = to_fixed(xim, wp)
x2re = (xre*xre - xim*xim) >> wp
x2im = (xre*xim) >> (wp - 1)
are = bre = x2re
aim = bim = x2im
c1, s1 = cos_sin(z._mpf_, wp)
cn = c1 = to_fixed(c1, wp)
sn = s1 = 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 = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
#case z complex, q real
elif isinstance(q, mpf):
wp = mp.prec + extra2
x = to_fixed(q._mpf_, wp)
x2 = (x*x) >> wp
a = b = x2
prec0 = mp.prec
mp.prec = wp
c1 = cos(z)
s1 = sin(z)
mp.prec = prec0
cnre = c1re = to_fixed(c1.real._mpf_, wp)
cnim = c1im = to_fixed(c1.imag._mpf_, wp)
snre = s1re = to_fixed(s1.real._mpf_, wp)
snim = s1im = to_fixed(s1.imag._mpf_, 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 = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
# case z and q complex
else:
wp = mp.prec + extra2
xre, xim = q._mpc_
xre = to_fixed(xre, wp)
xim = to_fixed(xim, wp)
x2re = (xre*xre - xim*xim) >> wp
x2im = (xre*xim) >> (wp - 1)
are = bre = x2re
aim = bim = x2im
prec0 = mp.prec
mp.prec = wp
# cos(z), siz(z) with z complex
c1 = cos(z)
s1 = sin(z)
mp.prec = prec0
cnre = c1re = to_fixed(c1.real._mpf_, wp)
cnim = c1im = to_fixed(c1.imag._mpf_, wp)
snre = s1re = to_fixed(s1.real._mpf_, wp)
snim = s1im = to_fixed(s1.imag._mpf_, 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 = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
s *= nthroot(q, 4)
return s
def _djacobi_theta3(z, q, nd):
"""nd=1,2,3 order of the derivative with respect to z"""
MIN = 2
extra1 = 10
extra2 = 20
if isinstance(q, mpf) and isinstance(z, mpf):
s = MP_ZERO
wp = mp.prec + extra1
x = to_fixed(q._mpf_, wp)
a = b = x
x2 = (x*x) >> wp
c1, s1 = cos_sin(mpf_shift(z._mpf_, 1), wp)
c1 = to_fixed(c1, wp)
s1 = to_fixed(s1, wp)
cn = c1
sn = s1
if (nd&1):
s += (a * sn) >> wp
else:
s += (a * cn) >> wp
n = 2
while abs(a) > MIN:
b = (b*x2) >> wp
a = (a*b) >> wp
cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
if nd&1:
s += (a * sn * n**nd) >> wp
else:
s += (a * cn * n**nd) >> wp
n += 1
s = -(s << (nd+1))
s = mpf(from_man_exp(s, -wp, mp.prec, 'n'))
# case z real, q complex
elif isinstance(z, mpf):
wp = mp.prec + extra2
xre, xim = q._mpc_
xre = to_fixed(xre, wp)
xim = to_fixed(xim, wp)
x2re = (xre*xre - xim*xim) >> wp
x2im = (xre*xim) >> (wp - 1)
are = bre = xre
aim = bim = xim
c1, s1 = cos_sin(mpf_shift(z._mpf_, 1), wp)
c1 = to_fixed(c1, wp)
s1 = to_fixed(s1, wp)
cn = c1
sn = s1
if (nd&1):
sre = (are * sn) >> wp
sim = (aim * sn) >> wp
else:
sre = (are * cn) >> wp
sim = (aim * cn) >> wp
n = 2
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
if nd&1:
sre += (are * sn * n**nd) >> wp
sim += (aim * sn * n**nd) >> wp
else:
sre += (are * cn * n**nd) >> wp
sim += (aim * cn * n**nd) >> wp
n += 1
sre = -(sre << (nd+1))
sim = -(sim << (nd+1))
sre = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
#case z complex, q real
elif isinstance(q, mpf):
wp = mp.prec + extra2
x = to_fixed(q._mpf_, wp)
a = b = x
x2 = (x*x) >> wp
prec0 = mp.prec
mp.prec = wp
c1 = cos(2*z)
s1 = sin(2*z)
mp.prec = prec0
cnre = c1re = to_fixed(c1.real._mpf_, wp)
cnim = c1im = to_fixed(c1.imag._mpf_, wp)
snre = s1re = to_fixed(s1.real._mpf_, wp)
snim = s1im = to_fixed(s1.imag._mpf_, wp)
if (nd&1):
sre = (a * snre) >> wp
sim = (a * snim) >> wp
else:
sre = (a * cnre) >> wp
sim = (a * cnim) >> wp
n = 2
while abs(a) > MIN:
b = (b*x2) >> wp
a = (a*b) >> wp
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
if (nd&1):
sre += (a * snre * n**nd) >> wp
sim += (a * snim * n**nd) >> wp
else:
sre += (a * cnre * n**nd) >> wp
sim += (a * cnim * n**nd) >> wp
n += 1
sre = -(sre << (nd+1))
sim = -(sim << (nd+1))
sre = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
# case z and q complex
else:
wp = mp.prec + extra2
xre, xim = q._mpc_
xre = to_fixed(xre, wp)
xim = to_fixed(xim, wp)
x2re = (xre*xre - xim*xim) >> wp
x2im = (xre*xim) >> (wp - 1)
are = bre = xre
aim = bim = xim
prec0 = mp.prec
mp.prec = wp
# cos(2*z), sin(2*z) with z complex
c1 = cos(2*z)
s1 = sin(2*z)
mp.prec = prec0
cnre = c1re = to_fixed(c1.real._mpf_, wp)
cnim = c1im = to_fixed(c1.imag._mpf_, wp)
snre = s1re = to_fixed(s1.real._mpf_, wp)
snim = s1im = to_fixed(s1.imag._mpf_, wp)
if (nd&1):
sre = (are * snre - aim * snim) >> wp
sim = (aim * snre + are * snim) >> wp
else:
sre = (are * cnre - aim * cnim) >> wp
sim = (aim * cnre + are * cnim) >> wp
n = 2
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
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
if(nd&1):
sre += ((are * snre - aim * snim) * n**nd) >> wp
sim += ((aim * snre + are * snim) * n**nd) >> wp
else:
sre += ((are * cnre - aim * cnim) * n**nd) >> wp
sim += ((aim * cnre + are * cnim) * n**nd) >> wp
n += 1
sre = -(sre << (nd+1))
sim = -(sim << (nd+1))
sre = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
if (nd&1):
return (-1)**(nd//2) * s
else:
return (-1)**(1 + nd//2) * s
def _jacobi_theta3(z, q):
extra1 = 10
extra2 = 20
MIN = 2
if z == zero:
if isinstance(q, mpf):
wp = mp.prec + extra1
x = to_fixed(q._mpf_, 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 = mpf(from_man_exp(s, -wp, mp.prec, 'n'))
return s
else:
wp = mp.prec + extra1
xre, xim = q._mpc_
xre = to_fixed(xre, wp)
xim = to_fixed(xim, 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 = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
return s
else:
if isinstance(q, mpf) and isinstance(z, mpf):
s = MP_ZERO
wp = mp.prec + extra1
x = to_fixed(q._mpf_, wp)
a = b = x
x2 = (x*x) >> wp
c1, s1 = cos_sin(mpf_shift(z._mpf_, 1), wp)
c1 = to_fixed(c1, wp)
s1 = to_fixed(s1, wp)
cn = c1
sn = s1
s += (a * cn) >> wp
while abs(a) > MIN:
b = (b*x2) >> wp
a = (a*b) >> wp
cn, sn = (cn*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
s += (a * cn) >> wp
s = (1 << wp) + (s << 1)
s = mpf(from_man_exp(s, -wp, mp.prec, 'n'))
return s
# case z real, q complex
elif isinstance(z, mpf):
wp = mp.prec + extra2
xre, xim = q._mpc_
xre = to_fixed(xre, wp)
xim = to_fixed(xim, wp)
x2re = (xre*xre - xim*xim) >> wp
x2im = (xre*xim) >> (wp - 1)
are = bre = xre
aim = bim = xim
c1, s1 = cos_sin(mpf_shift(z._mpf_, 1), wp)
c1 = to_fixed(c1, wp)
s1 = to_fixed(s1, wp)
cn = c1
sn = s1
sre = (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*c1 - sn*s1) >> wp, (sn*c1 + cn*s1) >> wp
sre += (are * cn) >> wp
sim += (aim * cn) >> wp
sre = (1 << wp) + (sre << 1)
sim = (sim << 1)
sre = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
return s
#case z complex, q real
elif isinstance(q, mpf):
wp = mp.prec + extra2
x = to_fixed(q._mpf_, wp)
a = b = x
x2 = (x*x) >> wp
prec0 = mp.prec
mp.prec = wp
c1 = cos(2*z)
s1 = sin(2*z)
mp.prec = prec0
cnre = c1re = to_fixed(c1.real._mpf_, wp)
cnim = c1im = to_fixed(c1.imag._mpf_, wp)
snre = s1re = to_fixed(s1.real._mpf_, wp)
snim = s1im = to_fixed(s1.imag._mpf_, wp)
sre = (a * cnre) >> wp
sim = (a * cnim) >> wp
while abs(a) > MIN:
b = (b*x2) >> wp
a = (a*b) >> wp
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 = (1 << wp) + (sre << 1)
sim = (sim << 1)
sre = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
return s
# case z and q complex
else:
wp = mp.prec + extra2
xre, xim = q._mpc_
xre = to_fixed(xre, wp)
xim = to_fixed(xim, wp)
x2re = (xre*xre - xim*xim) >> wp
x2im = (xre*xim) >> (wp - 1)
are = bre = xre
aim = bim = xim
prec0 = mp.prec
mp.prec = wp
# cos(2*z), sin(2*z) with z complex
c1 = cos(2*z)
s1 = sin(2*z)
mp.prec = prec0
cnre = c1re = to_fixed(c1.real._mpf_, wp)
cnim = c1im = to_fixed(c1.imag._mpf_, wp)
snre = s1re = to_fixed(s1.real._mpf_, wp)
snim = s1im = to_fixed(s1.imag._mpf_, wp)
sre = (are * cnre - aim * cnim) >> wp
sim = (aim * cnre + are * cnim) >> 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
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 = (1 << wp) + (sre << 1)
sim = (sim << 1)
sre = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
return s
def _djacobi_theta2(z, q, nd):
MIN = 2
extra1 = 10
extra2 = 20
if isinstance(q, mpf) and isinstance(z, mpf):
wp = mp.prec + extra1
x = to_fixed(q._mpf_, wp)
x2 = (x*x) >> wp
a = b = x2
c1, s1 = cos_sin(z._mpf_, wp)
cn = c1 = to_fixed(c1, wp)
sn = s1 = 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
if (nd&1):
s = s1 + ((a * sn * 3**nd) >> wp)
else:
s = c1 + ((a * cn * 3**nd) >> wp)
n = 2
while abs(a) > MIN:
b = (b*x2) >> wp
a = (a*b) >> wp
cn, sn = (cn*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
if nd&1:
s += (a * sn * (2*n+1)**nd) >> wp
else:
s += (a * cn * (2*n+1)**nd) >> wp
n += 1
s = -(s << 1)
s = mpf(from_man_exp(s, -wp, mp.prec, 'n'))
# case z real, q complex
elif isinstance(z, mpf):
wp = mp.prec + extra2
xre, xim = q._mpc_
xre = to_fixed(xre, wp)
xim = to_fixed(xim, wp)
x2re = (xre*xre - xim*xim) >> wp
x2im = (xre*xim) >> (wp - 1)
are = bre = x2re
aim = bim = x2im
c1, s1 = cos_sin(z._mpf_, wp)
cn = c1 = to_fixed(c1, wp)
sn = s1 = 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
if (nd&1):
sre = s1 + ((are * sn * 3**nd) >> wp)
sim = ((aim * sn * 3**nd) >> wp)
else:
sre = c1 + ((are * cn * 3**nd) >> wp)
sim = ((aim * cn * 3**nd) >> 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*c2 - sn*s2) >> wp, (sn*c2 + cn*s2) >> wp
if (nd&1):
sre += ((are * sn * n**nd) >> wp)
sim += ((aim * sn * n**nd) >> wp)
else:
sre += ((are * cn * n**nd) >> wp)
sim += ((aim * cn * n**nd) >> wp)
n += 2
sre = -(sre << 1)
sim = -(sim << 1)
sre = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
#case z complex, q real
elif isinstance(q, mpf):
wp = mp.prec + extra2
x = to_fixed(q._mpf_, wp)
x2 = (x*x) >> wp
a = b = x2
prec0 = mp.prec
mp.prec = wp
c1 = cos(z)
s1 = sin(z)
mp.prec = prec0
cnre = c1re = to_fixed(c1.real._mpf_, wp)
cnim = c1im = to_fixed(c1.imag._mpf_, wp)
snre = s1re = to_fixed(s1.real._mpf_, wp)
snim = s1im = to_fixed(s1.imag._mpf_, 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
if (nd&1):
sre = s1re + ((a * snre * 3**nd) >> wp)
sim = s1im + ((a * snim * 3**nd) >> wp)
else:
sre = c1re + ((a * cnre * 3**nd) >> wp)
sim = c1im + ((a * cnim * 3**nd) >> wp)
n = 5
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
if (nd&1):
sre += ((a * snre * n**nd) >> wp)
sim += ((a * snim * n**nd) >> wp)
else:
sre += ((a * cnre * n**nd) >> wp)
sim += ((a * cnim * n**nd) >> wp)
n += 2
sre = -(sre << 1)
sim = -(sim << 1)
sre = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
# case z and q complex
else:
wp = mp.prec + extra2
xre, xim = q._mpc_
xre = to_fixed(xre, wp)
xim = to_fixed(xim, wp)
x2re = (xre*xre - xim*xim) >> wp
x2im = (xre*xim) >> (wp - 1)
are = bre = x2re
aim = bim = x2im
prec0 = mp.prec
mp.prec = wp
# cos(2*z), siz(2*z) with z complex
c1 = cos(z)
s1 = sin(z)
mp.prec = prec0
cnre = c1re = to_fixed(c1.real._mpf_, wp)
cnim = c1im = to_fixed(c1.imag._mpf_, wp)
snre = s1re = to_fixed(s1.real._mpf_, wp)
snim = s1im = to_fixed(s1.imag._mpf_, 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
if (nd&1):
sre = s1re + (((are * snre - aim * snim) * 3**nd) >> wp)
sim = s1im + (((are * snim + aim * snre)* 3**nd) >> wp)
else:
sre = c1re + (((are * cnre - aim * cnim) * 3**nd) >> wp)
sim = c1im + (((are * cnim + aim * cnre)* 3**nd) >> 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
if (nd&1):
sre += (((are * snre - aim * snim) * n**nd) >> wp)
sim += (((aim * snre + are * snim) * n**nd) >> wp)
else:
sre += (((are * cnre - aim * cnim) * n**nd) >> wp)
sim += (((aim * cnre + are * cnim) * n**nd) >> wp)
n += 2
sre = -(sre << 1)
sim = -(sim << 1)
sre = from_man_exp(sre, -wp, mp.prec, 'n')
sim = from_man_exp(sim, -wp, mp.prec, 'n')
s = mpc(sre, sim)
s *= nthroot(q, 4)
if (nd&1):
return (-1)**(nd//2) * s
else:
return (-1)**(1 + nd//2) * s
def polyroots(coeffs, maxsteps=50, cleanup=True, extraprec=10, error=False):
"""
Numerically locate all (complex) roots of a polynomial using the
Durand-Kerner method.
With error=True, this function returns a tuple (roots, err) where roots
is a list of complex numbers sorted by absolute value, and err is an
estimate of the maximum error. The polynomial should be given as a list
of coefficients, in the same format as accepted by polyval(). The
leading coefficient must be nonzero.
These are the roots of x^3 - x^2 - 14*x + 24 and 4x^2 + 3x + 2:
>>> nprint(polyroots([1,-1,-14,24]), 4)
[-4.0, 2.0, 3.0]
>>> nprint(polyroots([4,3,2], error=True))
([(-0.375 - 0.599479j), (-0.375 + 0.599479j)], 2.22045e-16)
"""
if len(coeffs) <= 1:
if not coeffs or not coeffs[0]:
raise ValueError("Input to polyroots must not be the zero polynomial")
# Constant polynomial with no roots
return []
orig = mp.prec
weps = +eps
try:
mp.prec += 10
deg = len(coeffs) - 1
# Must be monic
lead = convert_lossless(coeffs[0])
if lead == 1:
coeffs = map(convert_lossless, coeffs)
else:
coeffs = [c/lead for c in coeffs]
f = lambda x: polyval(coeffs, x)
roots = [mpc((0.4+0.9j)**n) for n in xrange(deg)]
err = [mpf(1) for n in xrange(deg)]
for step in xrange(maxsteps):
if max(err).ae(0):
break
for i in xrange(deg):
if not err[i].ae(0):
p = roots[i]
x = f(p)
for j in range(deg):
if i != j:
try:
x /= (p-roots[j])
except ZeroDivisionError:
continue
roots[i] = p - x
err[i] = abs(x)
if cleanup:
for i in xrange(deg):
if abs(roots[i].imag) < weps:
roots[i] = roots[i].real
elif abs(roots[i].real) < weps:
roots[i] = roots[i].imag * 1j
roots.sort(key=lambda x: (abs(x.imag), x.real))
finally:
mp.prec = orig
if error:
err = max(err)
err = max(err, ldexp(1, -orig+1))
return [+r for r in roots], +err
else:
return [+r for r in roots]
