mpmath 0.2__zip → 0.4__zip

mpmath-0.4/mpmath/lib/squareroot.py

@@ -0,0 +1,199 @@
+"""
+This module provides functions for computing approximate square roots
+(in fixed-point or floating-point form) of positive real numbers.
+"""
+
+from util import *
+from floatop import *
+
+
+def sqrt_initial(y, prec):
+    """Given y = floor(x * 2**prec), compute floor(sqrt(x) * 2**50),
+    i.e. calculate a 50-bit approximation of the square root. This is
+    done quickly using regular floating-point arithmetic. It is
+    assumed that x ~= 1."""
+
+    # Two cases; the second avoids overflow
+    if prec < 200: return int(y**0.5 * 2.0**(50 - prec*0.5))
+    else:          return int((y >> (prec-100))**0.5)
+
+
+# XXX: doesn't work
+def invsqrt_initial(y, prec):
+    """Like sqrt_initial, but computes 1/sqrt(y) instead of sqrt(y)."""
+    if prec < 200: return int(y**-0.5 * 2.0**(50 + prec*0.5))
+    else:          return int((y >> (prec-100)) ** -0.5)
+
+
+def sqrt_fixed(y, prec):
+    """
+    Square root of a fixed-point number.
+
+    Given the big integer y = floor(x * 2**prec), this function returns
+    floor(r * 2**prec) where r = sqrt(x).
+
+    We start from a 50-bit estimate for r generated with ordinary
+    floating-point arithmetic, and then refines the value to full
+    accuracy using the iteration
+
+                 1  /        y  \
+        r     = --- | r  +  --- |
+         n+1     2  \ n     r_n /
+
+    which is simply Newton's method applied to the equation r**2 = y.
+
+    Newton's method doubles the accuracy with each step. We make use of
+    this fact by only using a precision corresponding to the current
+    accuracy during intermediate iterations. For example, with a 50-bit
+    accurate r_1, r_2 can be computed using 100-bit precision, r_3
+    using 200-bit precision, and so on. (In practice, the precision
+    levels must be chosen slightly more conservatively to account for
+    rounding errors in the last one or two bits.)
+
+    It is assumed that x ~= 1 (the main fsqrt() function fiddles with
+    the exponent of the input to reduce it to unit magnitude before
+    passing it here.)
+
+    TODO: it would be possible to specify separate precision levels
+    for the input and output, which could be useful when calculating
+    pure-integer square roots.
+    """
+
+    r = sqrt_initial(y, prec)
+    extra = 10
+    prevp = 50
+
+    for p in giant_steps(50, prec+extra):
+
+        # Explanation: in the first term, we shift by the appropriate number
+        # of bits to convert r from the previous precision to the current one.
+        # The "-1" divides by two in the same step.
+
+        # In the second term, we do a fixed-point division using the usual
+        # formula (y<<precision)//r. The precision term is a bit messy and
+        # takes into account the fact that y, r_n and r_{n+1} all have
+        # different precision levels. As before, the "-1" divides by two.
+        r = lshift_quick(r, p-prevp-1) + (lshift_quick(y, p+prevp-prec-1)//r)
+
+        prevp = p
+
+    return r >> extra
+
+
+def sqrt_fixed2(y, prec):
+    """
+    This function is essentially equivalent to sqrt_fixed (see its
+    documentation), but uses an asymptotically faster algorithm.
+
+    Instead of using Newton's method to calculate sqrt(y) directly,
+    we calculate 1/sqrt(y) with Newton's method and multiply by y to
+    obtain sqrt(y). The Newton iteration for 1/sqrt(y) is
+
+                 r
+                  n      /            2 \
+        r    =  ----  *  | 3  - y * r   |.
+         n+1      2      \           n  /
+
+    This is slightly slower at low precision levels since it requires
+    three multiplications in each step, as opposed to the single
+    division in the Newton iteration for sqrt(y).
+
+    However, since Python uses Karatsuba algorithm for multiplication,
+    three multiplications can be performed much more quickly than a
+    single division at high precision levels. In practice, the cutoff
+    where sqrt_fixed2 becomes faster than sqrt_fixed seems to be around
+    60,000 bits.
+    """
+
+    # XXX
+    from mpmath.lib import to_float, normalize, ROUND_FLOOR
+    r = to_float(normalize(y, -prec, 64, ROUND_FLOOR)) ** -0.5
+    r = int(r * 2**50)
+
+    # r = invsqrt_initial(y, prec)
+
+    extra = 10
+    prevp = 50
+
+    for p in giant_steps(50, prec+extra):
+
+        # This is even messier than in sqrt_fixed. As many shifts as possible
+        # have been combined together for optimal speed, at a slight expense
+        # of legibility.
+
+        # Compute r**2 at precision p.
+        r2 = rshift_quick(r*r, 2*prevp - p)
+
+        # A = r, converted from precision prevp to p
+        A = lshift_quick(r, p-prevp)
+
+        # S = y * r2, computed at precision p. We shift y by '-prec' to
+        # account for its initial precision, and by 'p' for the fixed-point
+        # multiplication
+        S = (lshift_quick(y, p-prec) * r2) >> p
+
+        # B = (3-S) and finally the outer product, both done at precision p
+        B = (3<<p) - S
+        r = (A*B) >> (p+1)
+
+        prevp = p
+
+    # Finally, sqrt(y) = y * (1/sqrt(y))
+    r = (r * y) >> prec
+
+    return r >> extra
+
+
+
+
+def fsqrt(s, prec, rounding):
+    """
+    Floating-point square root.
+
+    Returns a tuple representing the square root of s, rounded to the
+    nearest floating-point number in the specified rounding direction.
+    The input must be a tuple representing a nonnegative floating-point
+    number.
+    """
+    if s == fone:
+        return fone
+
+    man, exp, bc = s
+
+    if not man:
+        return fzero
+
+    # Convert to a fixed-point number with prec2 bits. Adjust
+    # exponents to be even so that they can be divided in half
+    prec2 = prec + 10 + (prec & 1)
+
+    if exp & 1:
+        exp -= 1
+        man <<= 1
+        bc += 1
+
+    # Mantissa may have more bits than we need. Trim it down.
+    shift = bc - prec2
+    shift -= shift & 1
+    man = rshift_quick(man, shift)
+
+    if prec < 65000:
+        man = sqrt_fixed(man, prec2)
+    else:
+        man = sqrt_fixed2(man, prec2)
+
+    return normalize(man, (exp+shift-prec2)>>1, prec, rounding)
+
+
+
+def fhypot(x, y, prec, rounding):
+    if y == fzero:
+        return fabs(x, prec, rounding)
+    if x == fzero:
+        return fabs(y, prec, rounding)
+
+    RF = ROUND_FLOOR
+    hypot2 = fadd(fmul(x,x,prec+4,RF), fmul(y,y,prec+4,RF), prec+4, RF)
+
+    return fsqrt(hypot2, prec, rounding)
+

mpmath-0.4/mpmath/lib/util.py

@@ -0,0 +1,268 @@
+"""
+Contents of this module:
+
+* Integer and bit-level operations
+* Miscellaneous utilities
+
+"""
+
+import math
+import decimal
+
+# Same as standard Python float
+STANDARD_PREC = 53
+
+LOG2_10 = math.log(10,2)  # 3.3219...
+
+
+# All supported rounding modes. We define them as integer constants for easy
+# management, but change __repr__ to give more information on inspection
+
+class Rounding(int):
+
+    def __new__(cls, level, name):
+        a = int.__new__(cls, level)
+        a.name = name
+        return a
+
+    def __repr__(self):
+        return self.name
+
+    __str__ = __repr__
+
+ROUND_DOWN    = Rounding(1, 'ROUND_DOWN')
+ROUND_UP      = Rounding(2, 'ROUND_UP')
+ROUND_FLOOR   = Rounding(3, 'ROUND_FLOOR')
+ROUND_CEILING = Rounding(4, 'ROUND_CEILING')
+ROUND_HALF_UP = Rounding(5, 'ROUND_HALF_UP')
+ROUND_HALF_DOWN = Rounding(6, 'ROUND_HALF_DOWN')
+ROUND_HALF_EVEN = Rounding(7, 'ROUND_HALF_EVEN')
+
+
+def giant_steps(start, target):
+    """Generate a list of precisions ranging from 'start' to 'target'
+    that doubles with each step. This is used by quadratically
+    convergent iterations (that is, Newton iterations), where we want
+    to keep the precision at the same level as the accuracy in each
+    step to minimize work.
+
+    For example, to find a sequence of precisions to reach 1000 bits
+    starting from a 53-bit estimate, giant_steps(53, 1000) gives
+
+        [64, 126, 251, 501, 1000]
+
+    So, if we start Newton's method with a 53-bit accurate initial
+    guess, the first iteration should be carried out at 64-bit
+    precision, the second at 126-bit precision, and so on.
+
+    Note the conservative rounding (1000 to 501, etc); this is used
+    guard against unit errors in the last place."""
+    L = [target]
+    while L[-1] > start*2:
+        L = L + [L[-1]//2 + 1]
+    return L[::-1]
+
+
+def rshift_quick(x, n):
+    """For an integer x, calculate x >> n with the fastest (floor)
+    rounding. Unlike the plain Python expression (x >> n), n is
+    allowed to be negative, in which case a left shift is performed."""
+    if n >= 0: return x >> n
+    else:      return x << (-n)
+
+
+def lshift_quick(x, n):
+    """For an integer x, calculate x << n. Unlike the plain Python
+    expression (x << n), n is allowed to be negative, in which case a
+    right shift with default (floor) rounding is performed."""
+    if n >= 0: return x << n
+    else:      return x >> (-n)
+
+
+def make_fixed(s, prec):
+    """Convert a floating-point number to a fixed-point big integer"""
+    man, exp, bc = s
+    offset = exp + prec
+    if offset >= 0:
+        return man << offset
+    else:
+        return man >> (-offset)
+
+
+def bitcount(n, log=math.log, table=(0,1,2,2,3,3,3,3,4,4,4,4,4,4,4,4)):
+    """Give size of n in bits; i.e. the position of the highest set bit
+    in n. If n is negative, the absolute value is used. The bitcount of
+    zero is taken to be 0."""
+
+    if not n: return 0
+    if n < 0: n = -n
+
+    # math.log gives a good estimate, and never overflows, but
+    # is not always exact. Subtract 2 to underestimate, then
+    # count remaining bits by table lookup
+    bc = int(log(n, 2)) - 2
+    if bc < 0:
+        bc = 0
+    return bc + table[n >> bc]
+
+
+# from decimal.py -- faster for small precs
+def bitcount2(n, correction = {
+        '0': 4, '1': 3, '2': 2, '3': 2,
+        '4': 1, '5': 1, '6': 1, '7': 1,
+        '8': 0, '9': 0, 'a': 0, 'b': 0,
+        'c': 0, 'd': 0, 'e': 0, 'f': 0}):
+    if n < 0:
+        n = -n
+    hex_n = "%x" % n
+    return 4*len(hex_n) - correction[hex_n[0]]
+
+
+def trailing_zeros(n):
+    """Count trailing zero bits in an integer. If n is negative, it is
+    replaced by its absolute value."""
+    if n & 1: return 0
+    if not n: return 0
+    if n < 0: n = -n
+    t = 0
+    while not n & 0xffffffffffffffff: n >>= 64; t += 64
+    while not n & 0xff: n >>= 8; t += 8
+    while not n & 1: n >>= 1; t += 1
+    return t
+
+#----------------------------------------------------------------------------
+# Integer shifting with directed rounding
+#
+
+def round_floor(x, n):
+    if not n or not x: return x
+    if n < 0: return x << -n
+    return x >> n
+
+def round_ceiling(x, n):
+    if not n or not x: return x
+    if n < 0: return x << -n
+    return -((-x) >> n)
+
+def round_down(x, n):
+    if not n or not x: return x
+    if n < 0: return x << -n
+    if x > 0:
+        return x >> n
+    else:
+        return -((-x) >> n)
+
+def round_up(x, n):
+    if not n or not x: return x
+    if n < 0: return x << -n
+    if x > 0:
+        return -((-x) >> n)
+    else:
+        return x >> n
+
+def round_half_up(x, n):
+    if not n or not x: return x
+    if n < 0: return x << -n
+    positive = x > 0
+    if positive: t = x >> (n-1)
+    else:        t = (-x) >> (n-1)
+    if t & 1:
+        if positive: return (t>>1)+1
+        else:        return -((t>>1)+1)
+    if positive: return t>>1
+    else:        return -(t>>1)
+
+def round_half_down(x, n):
+    if not n or not x: return x
+    if n < 0: return x << -n
+    positive = x > 0
+    if positive: t = x >> (n-1)
+    else:        t = (-x) >> (n-1)
+    if t & 1 and x & ((1<<(n-1))-1):
+        if positive: return (t>>1)+1
+        else:        return -((t>>1)+1)
+    if positive: return t>>1
+    else:        return -(t>>1)
+
+def round_half_even(x, n):
+    if not n or not x: return x
+    if n < 0: return x << -n
+    positive = x > 0
+    if positive: t = x >> (n-1)
+    else:        t = (-x) >> (n-1)
+    if t & 1 and (t&2 or x & ((1<<(n-1))-1)):
+        if positive: return (t>>1)+1
+        else:        return -((t>>1)+1)
+    if positive:  return t>>1
+    else:         return -(t>>1)
+
+
+rounding_functions = {
+  ROUND_FLOOR : round_floor,
+  ROUND_CEILING : round_ceiling,
+  ROUND_DOWN : round_down,
+  ROUND_UP : round_up,
+  ROUND_HALF_UP : round_half_up,
+  ROUND_HALF_DOWN : round_half_down,
+  ROUND_HALF_EVEN : round_half_even
+}
+
+
+def rshift(x, n, rounding):
+    """Shift x (a plain Python integer) n bits to the right (i.e.,
+    calculate x/(2**n)), and round to the nearest integer in accordance
+    with the specified rounding mode. The exponent n may be negative,
+    in which case x is shifted to the left (and no rounding is
+    necessary)."""
+    return rounding_functions[rounding](x, n)
+
+
+
+# TODO: speed up for bases 2, 4, 8, 16, ...
+
+def bin_to_radix(x, xbits, base, bdigits):
+    """
+    Radix conversion for fixed-point numbers. That is, convert
+    x * 2**xbits to floor(x * 10**bdigits).
+    """
+    return x * (base**bdigits) >> xbits
+
+stddigits = '0123456789abcdefghijklmnopqrstuvwxyz'
+
+def small_numeral(n, base=10, digits=stddigits):
+    """
+    Return the string numeral of a positive integer in an arbitrary
+    base. Most efficient for small input.
+    """
+    if base == 10:
+        return str(n)
+    digs = []
+    while n:
+        n, digit = divmod(n, base)
+        digs.append(digits[digit])
+    return "".join(digs[::-1])
+
+def numeral(n, base=10, size=0, digits=stddigits):
+    """
+    Represent the integer n as a string of digits in the given base.
+    Recursive division is used to make this function about 3x faster
+    than Python's str() for converting integers to decimal strings.
+
+    The 'size' parameters specifies the number of digits in n; this
+    number is only used to determine splitting points and need not
+    be exact.
+    """
+
+    if n < 0:
+        return "-" + numeral(-n, base, size, digits)
+
+    # Fast enough to do directly
+    if size < 250:
+        return small_numeral(n, base, digits)
+
+    # Divide in half
+    half = (size // 2) + (size & 1)
+    A, B = divmod(n, base**half)
+    ad = numeral(A, base, half, digits)
+    bd = numeral(B, base, half, digits).rjust(half, "0")
+    return ad + bd