Extend _sqrtprod() to cover the full range of inputs. Add tests. (GH-107855)

1 files changed, 18 insertions, 6 deletions

diff --git a/Lib/statistics.py b/Lib/statistics.py
index 93c44f1..a8036e9 100644
--- a/Lib/statistics.py
+++ b/Lib/statistics.py
@@ -137,6 +137,7 @@ from decimal import Decimal
 from itertools import count, groupby, repeat
 from bisect import bisect_left, bisect_right
 from math import hypot, sqrt, fabs, exp, erf, tau, log, fsum, sumprod
+from math import isfinite, isinf
 from functools import reduce
 from operator import itemgetter
 from collections import Counter, namedtuple, defaultdict
@@ -1005,14 +1006,25 @@ def _mean_stdev(data):
         return float(xbar), float(xbar) / float(ss)
 
 def _sqrtprod(x: float, y: float) -> float:
-    "Return sqrt(x * y) computed with high accuracy."
-    # Square root differential correction:
-    # https://www.wolframalpha.com/input/?i=Maclaurin+series+sqrt%28h**2+%2B+x%29+at+x%3D0
+    "Return sqrt(x * y) computed with improved accuracy and without overflow/underflow."
     h = sqrt(x * y)
+    if not isfinite(h):
+        if isinf(h) and not isinf(x) and not isinf(y):
+            # Finite inputs overflowed, so scale down, and recompute.
+            scale = 2.0 ** -512  # sqrt(1 / sys.float_info.max)
+            return _sqrtprod(scale * x, scale * y) / scale
+        return h
     if not h:
-        return 0.0
-    x = sumprod((x, h), (y, -h))
-    return h + x / (2.0 * h)
+        if x and y:
+            # Non-zero inputs underflowed, so scale up, and recompute.
+            # Scale:  1 / sqrt(sys.float_info.min * sys.float_info.epsilon)
+            scale = 2.0 ** 537
+            return _sqrtprod(scale * x, scale * y) / scale
+        return h
+    # Improve accuracy with a differential correction.
+    # https://www.wolframalpha.com/input/?i=Maclaurin+series+sqrt%28h**2+%2B+x%29+at+x%3D0
+    d = sumprod((x, h), (y, -h))
+    return h + d / (2.0 * h)
 
 
 # === Statistics for relations between two inputs ===