def mpf_pow(s, t, prec, rnd=round_fast): """ Compute s**t. Raises ComplexResult if s is negative and t is fractional. """ ssign, sman, sexp, sbc = s tsign, tman, texp, tbc = t if ssign and texp < 0: raise ComplexResult("negative number raised to a fractional power") if texp >= 0: return mpf_pow_int(s, (-1)**tsign * (tman<<texp), prec, rnd) # s**(n/2) = sqrt(s)**n if texp == -1: if tman == 1: if tsign: return mpf_div(fone, mpf_sqrt(s, prec+10, reciprocal_rnd[rnd]), prec, rnd) return mpf_sqrt(s, prec, rnd) else: if tsign: return mpf_pow_int(mpf_sqrt(s, prec+10, reciprocal_rnd[rnd]), -tman, prec, rnd) return mpf_pow_int(mpf_sqrt(s, prec+10, rnd), tman, prec, rnd) # General formula: s**t = exp(t*log(s)) # TODO: handle rnd direction of the logarithm carefully c = mpf_log(s, prec+10, rnd) return mpf_exp(mpf_mul(t, c), prec, rnd)
def mpf_asin(x, prec, rnd=round_fast): sign, man, exp, bc = x if bc+exp > 0 and x not in (fone, fnone): raise ComplexResult("asin(x) is real only for -1 <= x <= 1") flag_nr = True if prec < 1000 or exp+bc < -13: flag_nr = False else: ebc = exp + bc if ebc < -13: flag_nr = False elif ebc < -3: if prec < 3000: flag_nr = False if not flag_nr: # asin(x) = 2*atan(x/(1+sqrt(1-x**2))) wp = prec + 15 a = mpf_mul(x, x) b = mpf_add(fone, mpf_sqrt(mpf_sub(fone, a, wp), wp), wp) c = mpf_div(x, b, wp) return mpf_shift(mpf_atan(c, prec, rnd), 1) # use Newton's method extra = 10 extra_p = 10 prec2 = prec + extra r = math.asin(to_float(x)) r = from_float(r, 50, rnd) for p in giant_steps(50, prec2): wp = p + extra_p c, s = cos_sin(r, wp, rnd) tmp = mpf_sub(x, s, wp, rnd) tmp = mpf_div(tmp, c, wp, rnd) r = mpf_add(r, tmp, wp, rnd) sign, man, exp, bc = r return normalize(sign, man, exp, bc, prec, rnd)
def mpf_acosh(x, prec, rnd=round_fast): # acosh(x) = log(x+sqrt(x**2-1)) wp = prec + 15 if mpf_cmp(x, fone) == -1: raise ComplexResult("acosh(x) is real only for x >= 1") q = mpf_sqrt(mpf_add(mpf_mul(x,x), fnone, wp), wp) return mpf_log(mpf_add(x, q, wp), prec, rnd)
def mpf_agm(a, b, prec, rnd=round_fast): """ Computes the arithmetic-geometric mean agm(a,b) for nonnegative mpf values a, b. """ asign, aman, aexp, abc = a bsign, bman, bexp, bbc = b if asign or bsign: raise ComplexResult("agm of a negative number") # Handle inf, nan or zero in either operand if not (aman and bman): if a == fnan or b == fnan: return fnan if a == finf: if b == fzero: return fnan return finf if b == finf: if a == fzero: return fnan return finf # agm(0,x) = agm(x,0) = 0 return fzero wp = prec + 20 amag = aexp+abc bmag = bexp+bbc mag_delta = amag - bmag # Reduce to roughly the same magnitude using floating-point AGM abs_mag_delta = abs(mag_delta) if abs_mag_delta > 10: while abs_mag_delta > 10: a, b = mpf_shift(mpf_add(a,b,wp),-1), \ mpf_sqrt(mpf_mul(a,b,wp),wp) abs_mag_delta //= 2 asign, aman, aexp, abc = a bsign, bman, bexp, bbc = b amag = aexp+abc bmag = bexp+bbc mag_delta = amag - bmag #print to_float(a), to_float(b) # Use agm(a,b) = agm(x*a,x*b)/x to obtain a, b ~= 1 min_mag = min(amag,bmag) max_mag = max(amag,bmag) n = 0 # If too small, we lose precision when going to fixed-point if min_mag < -8: n = -min_mag # If too large, we waste time using fixed-point with large numbers elif max_mag > 20: n = -max_mag if n: a = mpf_shift(a, n) b = mpf_shift(b, n) #print to_float(a), to_float(b) af = to_fixed(a, wp) bf = to_fixed(b, wp) g = agm_fixed(af, bf, wp) return from_man_exp(g, -wp-n, prec, rnd)
def mpf_agm(a, b, prec, rnd=round_fast): """ Computes the arithmetic-geometric mean agm(a,b) for nonnegative mpf values a, b. """ asign, aman, aexp, abc = a bsign, bman, bexp, bbc = b if asign or bsign: raise ComplexResult("agm of a negative number") # Handle inf, nan or zero in either operand if not (aman and bman): if a == fnan or b == fnan: return fnan if a == finf: if b == fzero: return fnan return finf if b == finf: if a == fzero: return fnan return finf # agm(0,x) = agm(x,0) = 0 return fzero wp = prec + 20 amag = aexp + abc bmag = bexp + bbc mag_delta = amag - bmag # Reduce to roughly the same magnitude using floating-point AGM abs_mag_delta = abs(mag_delta) if abs_mag_delta > 10: while abs_mag_delta > 10: a, b = mpf_shift(mpf_add(a,b,wp),-1), \ mpf_sqrt(mpf_mul(a,b,wp),wp) abs_mag_delta //= 2 asign, aman, aexp, abc = a bsign, bman, bexp, bbc = b amag = aexp + abc bmag = bexp + bbc mag_delta = amag - bmag #print to_float(a), to_float(b) # Use agm(a,b) = agm(x*a,x*b)/x to obtain a, b ~= 1 min_mag = min(amag, bmag) max_mag = max(amag, bmag) n = 0 # If too small, we lose precision when going to fixed-point if min_mag < -8: n = -min_mag # If too large, we waste time using fixed-point with large numbers elif max_mag > 20: n = -max_mag if n: a = mpf_shift(a, n) b = mpf_shift(b, n) #print to_float(a), to_float(b) af = to_fixed(a, wp) bf = to_fixed(b, wp) g = agm_fixed(af, bf, wp) return from_man_exp(g, -wp - n, prec, rnd)
def calc_spouge_coefficients(a, prec): wp = prec + int(a * 1.4) c = [0] * a # b = exp(a-1) b = mpf_exp(from_int(a - 1), wp) # e = exp(1) e = mpf_exp(fone, wp) # sqrt(2*pi) sq2pi = mpf_sqrt(mpf_shift(mpf_pi(wp), 1), wp) c[0] = to_fixed(sq2pi, prec) for k in xrange(1, a): # c[k] = ((-1)**(k-1) * (a-k)**k) * b / sqrt(a-k) term = mpf_mul_int(b, ((-1)**(k - 1) * (a - k)**k), wp) term = mpf_div(term, mpf_sqrt(from_int(a - k), wp), wp) c[k] = to_fixed(term, prec) # b = b / (e * k) b = mpf_div(b, mpf_mul(e, from_int(k), wp), wp) return c
def calc_spouge_coefficients(a, prec): wp = prec + int(a*1.4) c = [0] * a # b = exp(a-1) b = mpf_exp(from_int(a-1), wp) # e = exp(1) e = mpf_exp(fone, wp) # sqrt(2*pi) sq2pi = mpf_sqrt(mpf_shift(mpf_pi(wp), 1), wp) c[0] = to_fixed(sq2pi, prec) for k in xrange(1, a): # c[k] = ((-1)**(k-1) * (a-k)**k) * b / sqrt(a-k) term = mpf_mul_int(b, ((-1)**(k-1) * (a-k)**k), wp) term = mpf_div(term, mpf_sqrt(from_int(a-k), wp), wp) c[k] = to_fixed(term, prec) # b = b / (e * k) b = mpf_div(b, mpf_mul(e, from_int(k), wp), wp) return c
def mpf_asin(x, prec, rnd=round_fast): sign, man, exp, bc = x if bc + exp > 0 and x not in (fone, fnone): raise ComplexResult("asin(x) is real only for -1 <= x <= 1") # asin(x) = 2*atan(x/(1+sqrt(1-x**2))) wp = prec + 15 a = mpf_mul(x, x) b = mpf_add(fone, mpf_sqrt(mpf_sub(fone, a, wp), wp), wp) c = mpf_div(x, b, wp) return mpf_shift(mpf_atan(c, prec, rnd), 1)
def mpc_erf(z, prec, rnd=round_fast): re, im = z if im == fzero: return (mpf_erf(re, prec, rnd), fzero) wp = prec + 20 z2 = mpc_mul(z, z, prec+20) v = mpc_hyp1f1_rat((1,2), (3,2), mpc_neg(z2), wp, rnd) sqrtpi = mpf_sqrt(mpf_pi(wp), wp) c = mpf_rdiv_int(2, sqrtpi, wp) c = mpc_mul_mpf(z, c, wp) return mpc_mul(c, v, prec, rnd)
def mpc_erf(z, prec, rnd=round_fast): re, im = z if im == fzero: return (mpf_erf(re, prec, rnd), fzero) wp = prec + 20 z2 = mpc_square(z, prec + 20) v = mpc_hyp1f1_rat((1, 2), (3, 2), mpc_neg(z2), wp, rnd) sqrtpi = mpf_sqrt(mpf_pi(wp), wp) c = mpf_rdiv_int(2, sqrtpi, wp) c = mpc_mul_mpf(z, c, wp) return mpc_mul(c, v, prec, rnd)
def mpc_sqrt(z, prec, rnd=round_fast): """Complex square root (principal branch). We have sqrt(a+bi) = sqrt((r+a)/2) + b/sqrt(2*(r+a))*i where r = abs(a+bi), when a+bi is not a negative real number.""" a, b = z if b == fzero: if a == fzero: return (a, b) # When a+bi is a negative real number, we get a real sqrt times i if a[0]: im = mpf_sqrt(mpf_neg(a), prec, rnd) return (fzero, im) else: re = mpf_sqrt(a, prec, rnd) return (re, fzero) wp = prec + 20 if not a[0]: # case a positive t = mpf_add(mpc_abs((a, b), wp), a, wp) # t = abs(a+bi) + a u = mpf_shift(t, -1) # u = t/2 re = mpf_sqrt(u, prec, rnd) # re = sqrt(u) v = mpf_shift(t, 1) # v = 2*t w = mpf_sqrt(v, wp) # w = sqrt(v) im = mpf_div(b, w, prec, rnd) # im = b / w else: # case a negative t = mpf_sub(mpc_abs((a, b), wp), a, wp) # t = abs(a+bi) - a u = mpf_shift(t, -1) # u = t/2 im = mpf_sqrt(u, prec, rnd) # im = sqrt(u) v = mpf_shift(t, 1) # v = 2*t w = mpf_sqrt(v, wp) # w = sqrt(v) re = mpf_div(b, w, prec, rnd) # re = b/w if b[0]: re = mpf_neg(re) im = mpf_neg(im) return re, im
def mpc_sqrt(z, prec, rnd=round_fast): """Complex square root (principal branch). We have sqrt(a+bi) = sqrt((r+a)/2) + b/sqrt(2*(r+a))*i where r = abs(a+bi), when a+bi is not a negative real number.""" a, b = z if b == fzero: if a == fzero: return (a, b) # When a+bi is a negative real number, we get a real sqrt times i if a[0]: im = mpf_sqrt(mpf_neg(a), prec, rnd) return (fzero, im) else: re = mpf_sqrt(a, prec, rnd) return (re, fzero) wp = prec+20 if not a[0]: # case a positive t = mpf_add(mpc_abs((a, b), wp), a, wp) # t = abs(a+bi) + a u = mpf_shift(t, -1) # u = t/2 re = mpf_sqrt(u, prec, rnd) # re = sqrt(u) v = mpf_shift(t, 1) # v = 2*t w = mpf_sqrt(v, wp) # w = sqrt(v) im = mpf_div(b, w, prec, rnd) # im = b / w else: # case a negative t = mpf_sub(mpc_abs((a, b), wp), a, wp) # t = abs(a+bi) - a u = mpf_shift(t, -1) # u = t/2 im = mpf_sqrt(u, prec, rnd) # im = sqrt(u) v = mpf_shift(t, 1) # v = 2*t w = mpf_sqrt(v, wp) # w = sqrt(v) re = mpf_div(b, w, prec, rnd) # re = b/w if b[0]: re = mpf_neg(re) im = mpf_neg(im) return re, im
def mpf_acos(x, prec, rnd=round_fast): # acos(x) = 2*atan(sqrt(1-x**2)/(1+x)) sign, man, exp, bc = x if bc + exp > 0: if x not in (fone, fnone): raise ComplexResult("acos(x) is real only for -1 <= x <= 1") if x == fnone: return mpf_pi(prec, rnd) wp = prec + 15 a = mpf_mul(x, x) b = mpf_sqrt(mpf_sub(fone, a, wp), wp) c = mpf_div(b, mpf_add(fone, x, wp), wp) return mpf_shift(mpf_atan(c, prec, rnd), 1)
def mpf_asinh(x, prec, rnd=round_fast): wp = prec + 20 sign, man, exp, bc = x mag = exp+bc if mag < -8: if mag < -wp: return mpf_perturb(x, 1-sign, prec, rnd) wp += (-mag) # asinh(x) = log(x+sqrt(x**2+1)) # use reflection symmetry to avoid cancellation q = mpf_sqrt(mpf_add(mpf_mul(x, x), fone, wp), wp) q = mpf_add(mpf_abs(x), q, wp) if sign: return mpf_neg(mpf_log(q, prec, negative_rnd[rnd])) else: return mpf_log(q, prec, rnd)
def mpf_ellipk(x, prec, rnd=round_fast): if not x[1]: if x == fzero: return mpf_shift(mpf_pi(prec, rnd), -1) if x == fninf: return fzero if x == fnan: return x if x == fone: return finf # TODO: for |x| << 1/2, one could use fall back to # pi/2 * hyp2f1_rat((1,2),(1,2),(1,1), x) wp = prec + 15 # Use K(x) = pi/2/agm(1,a) where a = sqrt(1-x) # The sqrt raises ComplexResult if x > 0 a = mpf_sqrt(mpf_sub(fone, x, wp), wp) v = mpf_agm1(a, wp) r = mpf_div(mpf_pi(wp), v, prec, rnd) return mpf_shift(r, -1)
def mpf_erf(x, prec, rnd=round_fast): sign, man, exp, bc = x if not man: if x == fzero: return fzero if x == finf: return fone if x == fninf: return fnone return fnan size = exp + bc lg = math.log # The approximation erf(x) = 1 is accurate to > x^2 * log(e,2) bits if size > 3 and 2 * (size - 1) + 0.528766 > lg(prec, 2): if sign: return mpf_perturb(fnone, 0, prec, rnd) else: return mpf_perturb(fone, 1, prec, rnd) # erf(x) ~ 2*x/sqrt(pi) close to 0 if size < -prec: # 2*x x = mpf_shift(x, 1) c = mpf_sqrt(mpf_pi(prec + 20), prec + 20) # TODO: interval rounding return mpf_div(x, c, prec, rnd) wp = prec + abs(size) + 20 # Taylor series for erf, fixed-point summation t = abs(to_fixed(x, wp)) t2 = (t * t) >> wp s, term, k = t, 12345, 1 while term: t = ((t * t2) >> wp) // k term = t // (2 * k + 1) if k & 1: s -= term else: s += term k += 1 s = (s << (wp + 1)) // sqrt_fixed(pi_fixed(wp), wp) if sign: s = -s return from_man_exp(s, -wp, wp, rnd)
def mpf_erf(x, prec, rnd=round_fast): sign, man, exp, bc = x if not man: if x == fzero: return fzero if x == finf: return fone if x== fninf: return fnone return fnan size = exp + bc lg = math.log # The approximation erf(x) = 1 is accurate to > x^2 * log(e,2) bits if size > 3 and 2*(size-1) + 0.528766 > lg(prec,2): if sign: return mpf_perturb(fnone, 0, prec, rnd) else: return mpf_perturb(fone, 1, prec, rnd) # erf(x) ~ 2*x/sqrt(pi) close to 0 if size < -prec: # 2*x x = mpf_shift(x,1) c = mpf_sqrt(mpf_pi(prec+20), prec+20) # TODO: interval rounding return mpf_div(x, c, prec, rnd) wp = prec + abs(size) + 25 # Taylor series for erf, fixed-point summation t = abs(to_fixed(x, wp)) t2 = (t*t) >> wp s, term, k = t, 12345, 1 while term: t = ((t * t2) >> wp) // k term = t // (2*k+1) if k & 1: s -= term else: s += term k += 1 s = (s << (wp+1)) // sqrt_fixed(pi_fixed(wp), wp) if sign: s = -s return from_man_exp(s, -wp, prec, rnd)
def mpf_asinh(x, prec, rnd=round_fast): # asinh(x) = log(x+sqrt(x**2+1)) wp = prec + 15 q = mpf_sqrt(mpf_add(mpf_mul(x, x), fone, wp), wp) return mpf_log(mpf_add(x, q, wp), prec, rnd)
def mpf_asinh(x, prec, rnd=round_fast): # asinh(x) = log(x+sqrt(x**2+1)) wp = prec + 15 q = mpf_sqrt(mpf_add(mpf_mul(x,x), fone, wp), wp) return mpf_log(mpf_add(x, q, wp), prec, rnd)
def acos_asin(z, prec, rnd, n): """ complex acos for n = 0, asin for n = 1 The algorithm is described in T.E. Hull, T.F. Fairgrieve and P.T.P. Tang 'Implementing the Complex Arcsine and Arcosine Functions using Exception Handling', ACM Trans. on Math. Software Vol. 23 (1997), p299 The complex acos and asin can be defined as acos(z) = acos(beta) - I*sign(a)* log(alpha + sqrt(alpha**2 -1)) asin(z) = asin(beta) + I*sign(a)* log(alpha + sqrt(alpha**2 -1)) where z = a + I*b alpha = (1/2)*(r + s); beta = (1/2)*(r - s) = a/alpha r = sqrt((a+1)**2 + y**2); s = sqrt((a-1)**2 + y**2) These expressions are rewritten in different ways in different regions, delimited by two crossovers alpha_crossover and beta_crossover, and by abs(a) <= 1, in order to improve the numerical accuracy. """ a, b = z wp = prec + 10 # special cases with real argument if b == fzero: am = mpf_sub(fone, mpf_abs(a), wp) # case abs(a) <= 1 if not am[0]: if n == 0: return mpf_acos(a, prec, rnd), fzero else: return mpf_asin(a, prec, rnd), fzero # cases abs(a) > 1 else: # case a < -1 if a[0]: pi = mpf_pi(prec, rnd) c = mpf_acosh(mpf_neg(a), prec, rnd) if n == 0: return pi, mpf_neg(c) else: return mpf_neg(mpf_shift(pi, -1)), c # case a > 1 else: c = mpf_acosh(a, prec, rnd) if n == 0: return fzero, c else: pi = mpf_pi(prec, rnd) return mpf_shift(pi, -1), mpf_neg(c) asign = bsign = 0 if a[0]: a = mpf_neg(a) asign = 1 if b[0]: b = mpf_neg(b) bsign = 1 am = mpf_sub(fone, a, wp) ap = mpf_add(fone, a, wp) r = mpf_hypot(ap, b, wp) s = mpf_hypot(am, b, wp) alpha = mpf_shift(mpf_add(r, s, wp), -1) beta = mpf_div(a, alpha, wp) b2 = mpf_mul(b, b, wp) # case beta <= beta_crossover if not mpf_sub(beta_crossover, beta, wp)[0]: if n == 0: re = mpf_acos(beta, wp) else: re = mpf_asin(beta, wp) else: # to compute the real part in this region use the identity # asin(beta) = atan(beta/sqrt(1-beta**2)) # beta/sqrt(1-beta**2) = (alpha + a) * (alpha - a) # alpha + a is numerically accurate; alpha - a can have # cancellations leading to numerical inaccuracies, so rewrite # it in differente ways according to the region Ax = mpf_add(alpha, a, wp) # case a <= 1 if not am[0]: # c = b*b/(r + (a+1)); d = (s + (1-a)) # alpha - a = (1/2)*(c + d) # case n=0: re = atan(sqrt((1/2) * Ax * (c + d))/a) # case n=1: re = atan(a/sqrt((1/2) * Ax * (c + d))) c = mpf_div(b2, mpf_add(r, ap, wp), wp) d = mpf_add(s, am, wp) re = mpf_shift(mpf_mul(Ax, mpf_add(c, d, wp), wp), -1) if n == 0: re = mpf_atan(mpf_div(mpf_sqrt(re, wp), a, wp), wp) else: re = mpf_atan(mpf_div(a, mpf_sqrt(re, wp), wp), wp) else: # c = Ax/(r + (a+1)); d = Ax/(s - (1-a)) # alpha - a = (1/2)*(c + d) # case n = 0: re = atan(b*sqrt(c + d)/2/a) # case n = 1: re = atan(a/(b*sqrt(c + d)/2) c = mpf_div(Ax, mpf_add(r, ap, wp), wp) d = mpf_div(Ax, mpf_sub(s, am, wp), wp) re = mpf_shift(mpf_add(c, d, wp), -1) re = mpf_mul(b, mpf_sqrt(re, wp), wp) if n == 0: re = mpf_atan(mpf_div(re, a, wp), wp) else: re = mpf_atan(mpf_div(a, re, wp), wp) # to compute alpha + sqrt(alpha**2 - 1), if alpha <= alpha_crossover # replace it with 1 + Am1 + sqrt(Am1*(alpha+1))) # where Am1 = alpha -1 # if alpha <= alpha_crossover: if not mpf_sub(alpha_crossover, alpha, wp)[0]: c1 = mpf_div(b2, mpf_add(r, ap, wp), wp) # case a < 1 if mpf_neg(am)[0]: # Am1 = (1/2) * (b*b/(r + (a+1)) + b*b/(s + (1-a)) c2 = mpf_add(s, am, wp) c2 = mpf_div(b2, c2, wp) Am1 = mpf_shift(mpf_add(c1, c2, wp), -1) else: # Am1 = (1/2) * (b*b/(r + (a+1)) + (s - (1-a))) c2 = mpf_sub(s, am, wp) Am1 = mpf_shift(mpf_add(c1, c2, wp), -1) # im = log(1 + Am1 + sqrt(Am1*(alpha+1))) im = mpf_mul(Am1, mpf_add(alpha, fone, wp), wp) im = mpf_log(mpf_add(fone, mpf_add(Am1, mpf_sqrt(im, wp), wp), wp), wp) else: # im = log(alpha + sqrt(alpha*alpha - 1)) im = mpf_sqrt(mpf_sub(mpf_mul(alpha, alpha, wp), fone, wp), wp) im = mpf_log(mpf_add(alpha, im, wp), wp) if asign: if n == 0: re = mpf_sub(mpf_pi(wp), re, wp) else: re = mpf_neg(re) if not bsign and n == 0: im = mpf_neg(im) if bsign and n == 1: im = mpf_neg(im) re = normalize(re[0], re[1], re[2], re[3], prec, rnd) im = normalize(im[0], im[1], im[2], im[3], prec, rnd) return re, im
re = from_man_exp(re, int(n*aexp), prec, rnd) im = from_man_exp(im, int(n*bexp), prec, rnd) return re, im return mpc_exp(mpc_mul_int(mpc_log(z, prec+10), n, prec+10), prec, rnd) def mpc_sqrt((a, b), prec, rnd=round_fast): """Complex square root (principal branch). We have sqrt(a+bi) = sqrt((r+a)/2) + b/sqrt(2*(r+a))*i where r = abs(a+bi), when a+bi is not a negative real number.""" if a == b == fzero: return (a, b) # When a+bi is a negative real number, we get a real sqrt times i if b == fzero: if a[0]: im = mpf_sqrt(mpf_neg(a), prec, rnd) return (fzero, im) else: re = mpf_sqrt(a, prec, rnd) return (re, fzero) wp = prec+20 if not a[0]: # case a positive t = mpf_add(mpc_abs((a, b), wp), a, wp) # t = abs(a+bi) + a u = mpf_shift(t, -1) # u = t/2 re = mpf_sqrt(u, prec, rnd) # re = sqrt(u) v = mpf_shift(t, 1) # v = 2*t w = mpf_sqrt(v, wp) # w = sqrt(v) im = mpf_div(b, w, prec, rnd) # im = b / w else: # case a negative t = mpf_sub(mpc_abs((a, b), wp), a, wp) # t = abs(a+bi) - a u = mpf_shift(t, -1) # u = t/2
im = from_man_exp(im, int(n * bexp), prec, rnd) return re, im return mpc_exp(mpc_mul_int(mpc_log(z, prec + 10), n, prec + 10), prec, rnd) def mpc_sqrt((a, b), prec, rnd=round_fast): """Complex square root (principal branch). We have sqrt(a+bi) = sqrt((r+a)/2) + b/sqrt(2*(r+a))*i where r = abs(a+bi), when a+bi is not a negative real number.""" if a == b == fzero: return (a, b) # When a+bi is a negative real number, we get a real sqrt times i if b == fzero: if a[0]: im = mpf_sqrt(mpf_neg(a), prec, rnd) return (fzero, im) else: re = mpf_sqrt(a, prec, rnd) return (re, fzero) wp = prec + 20 if not a[0]: # case a positive t = mpf_add(mpc_abs((a, b), wp), a, wp) # t = abs(a+bi) + a u = mpf_shift(t, -1) # u = t/2 re = mpf_sqrt(u, prec, rnd) # re = sqrt(u) v = mpf_shift(t, 1) # v = 2*t w = mpf_sqrt(v, wp) # w = sqrt(v) im = mpf_div(b, w, prec, rnd) # im = b / w else: # case a negative t = mpf_sub(mpc_abs((a, b), wp), a, wp) # t = abs(a+bi) - a u = mpf_shift(t, -1) # u = t/2