def zeta(s):
"""
zeta(s) -- calculate the Riemann zeta function of a real or complex
argument s.
"""
Float.store()
Float._prec += 8
si = s
s = ComplexFloat(s)
if s.real < 0:
# Reflection formula (XXX: gets bad around the zeros)
pi = pi_float()
y = power(2, s) * power(pi, s-1) * sin(pi*s/2) * gamma(1-s) * zeta(1-s)
else:
p = Float._prec
n = int((p + 2.28*abs(float(s.imag)))/2.54) + 3
d = _zeta_coefs(n)
if isinstance(si, (int, long)):
t = 0
for k in range(n):
t += (((-1)**k * (d[k] - d[n])) << p) // (k+1)**si
y = (Float((t, -p)) / -d[n]) / (Float(1) - Float(2)**(1-si))
else:
t = Float(0)
for k in range(n):
t += (-1)**k * Float(d[k]-d[n]) * exp(-_logk(k+1)*s)
y = (t / -d[n]) / (Float(1) - exp(log(2)*(1-s)))
Float.revert()
if isinstance(y, ComplexFloat) and s.imag == 0:
return +y.real
else:
return +y
def erf(x):
x = ComplexFloat(x)
if x == 0: return Float(0)
if x.real < 0: return -erf(-x)
Float.store()
Float._prec += 10
y = lower_gamma(0.5, x**2) / sqrt(pi_float())
if x.imag == 0:
y = y.real
Float.revert()
return +y
def lower_gamma(a, z):
Float.store()
prec = Float._prec
# XXX: may need more precision
Float._prec += 15
a = ComplexFloat(a)
z = ComplexFloat(z)
s = _lower_gamma_series(a.real, a.imag, z.real, z.imag, prec)
y = exp(log(z)*a) * exp(-z) * s / a
Float.revert()
return +y
def gamma(x):
"""
gamma(x) -- calculate the gamma function of a real or complex
number x.
x must not be a negative integer or 0
"""
Float.store()
Float._prec += 2
if isinstance(x, complex):
x = ComplexFloat(x)
elif not isinstance(x, (Float, ComplexFloat, Rational, int, long)):
x = Float(x)
if isinstance(x, (ComplexFloat, complex)):
re, im = x.real, x.imag
else:
re, im = x, 0
# For negative x (or positive x close to the pole at x = 0),
# we use the reflection formula
if re < 0.25:
if re == int(re) and im == 0:
raise ZeroDivisionError, "gamma function pole"
Float._prec += 3
p = pi_float()
g = p / (sin(p*x) * gamma(1-x))
else:
x -= 1
prec, a, c = _get_spouge_coefficients(Float.getprec()+7)
s = _spouge_sum(x, prec, a, c)
if not isinstance(x, (Float, ComplexFloat)):
x = Float(x)
# TODO: higher precision may be needed here when the precision
# and/or size of x are extremely large
Float._prec += 10
g = exp(log(x+a)*(x+Float(0.5))) * exp(-x-a) * s
Float.revert()
return +g
def adaptive(self, f, a, b, eps, steps=0, maxsteps=5000, verbose=False):
prec = Float.getprec()
for m in xrange(self.initial, 50):
if verbose:
print "using tanh-sinh rule with m =", m
ts = TanhSinh(eps, m, verbose=verbose)
s, err = ts(f, a, b, verbose=verbose)
steps = 2*len(ts.x)
if err <= eps or steps >= maxsteps:
return s, err, steps
if verbose:
print "error = ", err
def _calc_spouge_coefficients(a, prec):
"""
Calculate Spouge coefficients for approximation with parameter a.
Return a list of big integers representing the coefficients in
fixed-point form with a precision of prec bits.
"""
# We'll store the coefficients as fixed-point numbers but calculate
# them as Floats for convenience. The initial terms are huge, so we
# need to allocate extra bits to ensure full accuracy. The integer
# part of the largest term has size ~= exp(a) or 2**(1.4*a)
floatprec = prec + int(a*1.4)
Float.store()
Float.setprec(floatprec)
c = [0] * a
b = exp(a-1)
e = exp(1)
c[0] = make_fixed(sqrt(2*pi_float()), prec)
for k in range(1, a):
# print "%02f" % (100.0 * k / a), "% done"
c[k] = make_fixed(((-1)**(k-1) * (a-k)**k) * b / sqrt(a-k), prec)
# Divide off e and k instead of computing exp and k! from scratch
b = b / (e * k)
Float.revert()
return c
def __call__(self, f, a, b, extraprec=3, verbose=False):
Float.store()
Float.extraprec(extraprec)
f = self.transform(f, a, b)
try:
s, err = self._eval(f, verbose)
except Exception, e:
Float.revert()
raise e
def __init__(self, n=None, verbose=False):
self.n = n
prec = dps = Float.getdps()
self.prec = prec
# Reuse old nodes
if n in _gausscache:
cacheddps, xs, ws = _gausscache[n]
if cacheddps >= dps:
self.x = [x for x in xs]
self.w = [w for w in ws]
return
if verbose:
print ("calculating nodes for degree-%i Gauss-Legendre "
"quadrature..." % n)
Float.store()
Float.setdps(2*prec + 5)
self.x = [None] * n
self.w = [None] * n
pf = polyfunc(legendre(n, 'x'), True)
for k in xrange(n//2 + 1):
if verbose and k % 4 == 0:
print " node", k, "of", n//2
x, w = _gaussnode(pf, k, n)
self.x[k] = x
self.x[n-k-1] = -x
self.w[k] = self.w[n-k-1] = w
_gausscache[n] = (dps, self.x, self.w)
Float.revert()
def nintegrate(f, a, b, method=0, maxsteps=5000, verbose=False):
"""
Basic usage
===========
nintegrate(f, a, b) numerically evaluates the integral
- b
| f(x) dx.
- a
The integrand f should be a callable that accepts a Float as input
and outputs a Float or ComplexFloat; a and b should be Floats,
numbers that can be converted to Floats, or +/- infinity.
A simple example:
>>> Float.setdps(15)
>>> print nintegrate(lambda x: 2*x**2 - 4*x, 2, 3)
2.66666666666667
Calculating the area of a unit circle, with 30 digits:
>>> Float.setdps(30)
>>> print nintegrate(lambda x: 4*sqrt(1-x**2), 0, 1)
3.14159265358979323846264338328
The integration interval can be infinite or semi-infinite:
>>> Float.setdps(15)
>>> print nintegrate(lambda x: exp(-x)*sin(x), 0, oo)
0.5
Integration methods and accuracy
================================
Nintegrate attempts to obtain a value that is fully accurate within
the current working precision (i.e., correct to 15 decimal places
at the default precision level). If nintegrate fails to reach full
accuracy after a certain number of steps, it prints a warning
message.
This message signifies either that the integral is either divergent
or, if convergent, ill-behaved. It may still be possible to
evaluate an ill-behaved integral by increasing the 'maxsteps'
setting, changing the integration method, and/or manually
transforming the integrand.
Nintegrate currently supports the following integration methods:
method = 0 : Gaussian quadrature (default)
method = 1 : tanh-sinh quadrature
Gaussian quadrature is generally very efficient if the integration
interval is finite and the integrand is smooth on the entire range
(including the endpoints). It may fail if the integrand has many
discontinuities, is highly oscillatory, or possesses integrable
singularities.
The tanh-sinh algorithm is often better if the integration interval
is infinite or if singularities are present at the endpoints;
especially at very high precision levels. It does not perform well
if there are singularities between the endpoints or the integrand
is bumpy or oscillatory.
It may help to manually transform the integrand, e.g. changing
variables to remove singularities or breaking up the integration
interval so that singularities appear only at the endpoints.
The 'verbose' flag can be set to track the computation's progress.
"""
dps = Float.getdps()
Float.store()
Float.setdps(dps + 3)
prec = Float.getprec()
# Transform infinite or semi-infinite interval to a finite interval
if a == -oo or b == oo:
g = f
if a == -oo and b == oo:
def f(x):
# make adaptive quadrature work from the left
x = 1 - x
return g(x) + g(Float(1)/x)/x**2 + g(-x) + g(Float(-1)/x)/(-x)**2
elif b == oo:
aa = Float(a)
def f(x):
x = 1 - x
return g(x + aa) + g(Float(1)/x + aa)/x**2
elif a == -oo:
bb = Float(b)
def f(x):
x = 1 - x
return g(-x + bb) + g(Float(-1)/x + bb)/(-x)**2
a, b = Float(0), Float(1)
else:
a, b = Float(a), Float(b)
eps = Float((1, -prec+4))
if method == 0:
degree = int(5 + dps**0.8)
rule = CompositeGaussLegendre(degree, degree//2, verbose)
elif method == 1:
rule = AdaptiveTanhSinh(initial = int(3 + max(0, math.log(prec/30.0, 2))))
else:
Float.revert()
raise ValueError("unknown method")
s, err, steps = rule.adaptive(f, Float(a), Float(b), eps, steps=0, maxsteps=maxsteps, verbose=verbose)
Float.revert()
if not err.ae(0):
Float.store()
Float.setdps(1)
print "Warning: failed to reach full accuracy.", \
"Estimated magnitude of error: ", str(err)
Float.revert()
return +s
def __init__(self, eps, m, verbose=False):
# Compute abscissas and weights
prec = Float.getdps()
self.prec = prec
self.eps = eps
self.m = m
self.h = h = Float(1) / 2**m
self.x = []
self.w = []
if (eps, m) in _tscache:
self.x, self.w = _tscache[(eps, m)]
return
if verbose:
print ("calculating nodes for tanh-sinh quadrature with "
"epsilon %s and degree %i" % (eps, m))
Float.store()
Float.setdps(prec + 10)
for k in xrange(20 * 2**m + 1):
x, w = _tsnode(k, h)
self.x.append(x)
self.w.append(w)
diff = abs(self.x[-1] - Float(1))
if diff <= eps:
break
if verbose and m > 6 and k % 300 == 299:
Float.store(); Float.setdps(5)
print " progress", -log(diff, 10) / prec
Float.revert()
_tscache[(eps, m)] = self.x, self.w
Float.revert()
def evalf(expr):
"""
evalf(expr) attempts to evaluate a SymPy expression to a Float or
ComplexFloat with an error smaller than 10**(-Float.getdps())
"""
if isinstance(expr, (Float, ComplexFloat)):
return expr
elif isinstance(expr, (int, float)):
return Float(expr)
elif isinstance(expr, complex):
return ComplexFloat(expr)
expr = Basic.sympify(expr)
if isinstance(expr, (Rational)):
y = Float(expr)
elif isinstance(expr, Real):
y = Float(str(expr))
elif expr is I:
y = ComplexFloat(0,1)
elif expr is pi:
y = constants.pi_float()
elif expr is E:
y = functions.exp(1)
elif isinstance(expr, Mul):
factors = expr[:]
workprec = Float.getprec() + 1 + len(factors)
Float.store()
Float.setprec(workprec)
y = Float(1)
for f in factors:
y *= evalf(f)
Float.revert()
elif isinstance(expr, Pow):
base, expt = expr[:]
workprec = Float.getprec() + 8 # may need more
Float.store()
Float.setprec(workprec)
base = evalf(base)
expt = evalf(expt)
if expt == 0.5:
y = functions.sqrt(base)
else:
y = functions.exp(functions.log(base) * expt)
Float.revert()
elif isinstance(expr, Basic.exp):
Float.store()
Float.setprec(Float.getprec() + 3)
#XXX: how is it possible, that this works:
x = evalf(expr[0])
#and this too:
#x = evalf(expr[1])
#?? (Try to uncomment it and you'll see)
y = functions.exp(x)
Float.revert()
elif isinstance(expr, Add):
# TODO: this doesn't yet work as it should.
# We need some way to handle sums whose results are
# very close to 0, and when necessary, repeat the
# summation with higher precision
reqprec = Float.getprec()
Float.store()
Float.setprec(10)
terms = expr[:]
approxterms = [abs(evalf(x)) for x in terms]
min_mag = min(x.exp for x in approxterms)
max_mag = max(x.exp+bitcount(x.man) for x in approxterms)
Float.setprec(reqprec - 10 + max_mag - min_mag + 1 + len(terms))
workprec = Float.getdps()
y = 0
for t in terms:
y += evalf(t)
Float.revert()
else:
# print expr, expr.__class__
raise NotImplementedError
# print expr, y
return +y
def secant(f, x0, x1=None, eps=None, maxsteps=50, bracket=None,
on_error=1, verbose=False):
"""
secant(f, x0) uses the secant method to find a numerical root of f,
starting from the initial approximation/guess x0.
This algorithm is very efficient for finding regular roots of
smooth functions; nearly as efficient as Newton's method. Each
iteration roughly multiplies the number of correct digits in the
answer by 1.6 (in practice, a 15-digit value can usually be
found in less than 10 steps).
The secant method can be inefficient or may fail to converge in
the presence of high-order roots or if the initial approximation
is far from the true root. Generally, it works best if f is
monotonic and steeply sloped everywhere in the surrounding of
the root (and x0 is in that surrounding).
Advanced settings
=================
The secant method actually requires two initial values x0 and x1.
By default, an x1 value is generated by perturbing x0 slightly;
in some cases this perturbation may be too small and it may then
help to specify a custom value for x1.
Several other advanced options are supported as guards against
poor convergence:
eps
Only attempt to locate the root to within this epsilon.
By default, eps is set to give a full-precision value.
maxsteps
Break when this many iterations have been performed.
bracket
If bracket = [a, b], break if the iteration ends up somewhere
outside the interval [a, b]. This is useful to protect from
the iterate x_n "shooting off to space" in the presence of
a nearly horizontal function value. If the bracket is a
single number, break if abs(x_0-x_n) exceeds this value.
The 'on_error' parameter can be set to control what happens in the
case of a convergence failure:
on_error = 0 print warning if in verbose mode; return value
on_error = 1 print warning; return value (default)
on_error = 2 raise an exception
Examples
========
A simple example, calculating pi:
>>> Float.setdps(15)
>>> print secant(sin, 3)
3.14159265358979
If we try the same for cosine, starting from x=0, we get the
"wrong" root because the cosine is nearly horizontal around x=0,
sending the initial iterate far away (verbose=True shows what
happens in more detail):
>>> Float.setdps(15)
>>> print secant(cos, 0)
391.128285371929
This can be helped by providing a second point that better matches
the geometry of the function's graph or by providing a closer
initial estimate:
>>> print secant(cos, x0=0, x1=1)
1.57079632679490
>>> print secant(cos, 1.5)
1.57079632679490
As another example, a high-precision calculation of log(3):
>>> Float.setdps(50)
>>> print secant(lambda x: exp(x)-3, 1)
1.0986122886681096913952452369225257046474905578227
>>> Float.revert()
"""
prec = Float.getprec()
if eps is None:
eps = Float((1, -prec+2))
else:
eps = Float(eps)
x = x0 = x0 * Float(1)
if x1 is None:
# Perturb the initial value to generate a second point as
# needed to start the secant method iteration
xprev = perturb(x, 0.01)
else:
xprev = x1 * Float(1)
if isinstance(bracket, (list, tuple)):
brackets = Float(bracket[0]), Float(bracket[1])
else:
brackets = None
exception = None
bracketmsg = "Failed to converge (outside bracket)"
stepsmsg = "Failed to converge (maxsteps exceeded)"
derivmsg = "Failed to converge (inft./zero derivative or " \
"loss of precision in derivative)"
for i in xrange(1, maxsteps+1):
if verbose:
try:
print " secant: x_%i=%8f delta=%g" % (i,x,x-xprev)
except TypeError:
print " secant: x_%i=%s delta=%s" % (i,x,x-xprev)
# Calculate function values
fx = f(x)
fdiff = fx - f(xprev)
xdiff = x - xprev
if xdiff.ae(0, eps):
break
# Try to calculate the finite difference approximation for the
# derivative and update accordingly. In floating-point
# arithmetic, this may cause a division by zero when f(x) is
# extremely close to 0, which we have to watch out for
try:
deriv = xdiff/fdiff
x, xprev = x - fx*deriv, x
except ZeroDivisionError:
exception = derivmsg
break
# Check if within reasonable bounds
if brackets:
if not brackets[0] <= x <= brackets[1]:
exception = bracketmsg
break
elif bracket is not None:
print "heh", x, x0, x-x0, bracket
if abs(x-x0) >= bracket:
exception = bracketmsg
break
if i == maxsteps:
exception = stepsmsg
if exception:
if on_error == 2: raise ConvergenceError(exception)
if on_error == 1 or verbose: print exception
return x
def bisect(f, a, b, eps=None, maxsteps=None, verbose=False):
"""
Numerical root-finding using the bisection method.
Given a real-valued function f and an interval [a, b] (not
necessarily real) such that f(a) and f(b) have opposite signs,
narrow the interval where f crosses the x axis to a relative
width at most equal to 'eps' through repeated bisections. If
not specified, eps is set to give a full precision estimate
of the root.
A tuple for the narrowed interval is returned.
If f is continuous, bisect is guaranteed to find a root. If f
jumps discontinuously from positive negative values, bisect will
locate the point of the discontinuity. bisect quits early and
returns an interval of width zero if it encounters a point x
where f(x) = 0 exactly.
Optionally, perform no more than 'maxsteps' bisections (by
default perform as many as needed); if the 'verbose' flag is set,
print status at each step.
Examples
========
Find a bounding interval for pi/2 (which is a root of cos):
>>> Float.setdps(15)
>>> a, b = bisect(cos, 1, 2, 1e-2, verbose=True)
bisect step 1: a=1.000000 b=2.000000 delta=1
bisect step 2: a=1.500000 b=2.000000 delta=0.5
bisect step 3: a=1.500000 b=1.750000 delta=0.25
bisect step 4: a=1.500000 b=1.625000 delta=0.125
bisect step 5: a=1.562500 b=1.625000 delta=0.0625
bisect step 6: a=1.562500 b=1.593750 delta=0.03125
>>> print a, b
1.5625 1.578125
Calculate the same value to full precision:
>>> a, b = bisect(cos, 1, 2)
>>> a, b
(Float('1.5707963267948966'), Float('1.5707963267948968'))
>>> print cos(a)
6.12320117570134E-17
>>> print cos(b)
-1.60812593168018E-16
Although generally reliable, the bisection method has a slow rate
of convergence. The previous computation required about 50 steps,
which can be compared to the secant method which only requires
5-6 steps to obtain a full precision value from an initial value
near 1.5.
"""
a, b = a*Float(1), b*Float(1) # Convert to Float or ComplexFloat
fa0, fb0 = f(a), f(b)
# Sanity check
if not fa0 * fb0 <= 0:
raise ValueError("bisect: f(a) and f(b) have the same sign at a=%r "
" and b=%r" % (a, b))
fa0sign = sign(fa0)
# Default eps / set eps to something sane if zero
mineps = Float((1, -Float.getprec()+1))
if not eps or eps < mineps:
eps = mineps
if maxsteps is None:
# Use maxsteps as a safeguard should the loop condition fail
abdelta = abs(a-b)
maxsteps = bitcount(abdelta.man) - abdelta.exp + Float.getprec()
i = 1
while not a.ae(b, eps) and i <= maxsteps:
if verbose:
print " bisect step %2i: a=%8f b=%8f delta=%g" % (i,a,b,b-a)
mid = (a+b)*0.5
fmid = f(mid)
if fmid == 0:
return (mid, mid)
if sign(fmid) == fa0sign:
a = mid
else:
b = mid
i += 1
return a, b
