Correctly rounded stdev results for Decimal inputs

python · rhettinger · Dec 1, 2021 · Mar 16, 2021 · Mar 22, 2021 · Mar 22, 2021
commit 037b5fea25bf59df4e513992a7cc34a79d18425b
diff --git a/Lib/statistics.py b/Lib/statistics.py
@@ -305,15 +305,18 @@ def _fail_neg(values, errmsg='negative value'):
             raise StatisticsError(errmsg)
         yield x
 
+
 def _isqrt_frac_rto(n: int, m: int) -> float:
     """Square root of n/m, rounded to the nearest integer using round-to-odd."""
     # Reference: https://www.lri.fr/~melquion/doc/05-imacs17_1-expose.pdf
     a = math.isqrt(n // m)
     return a | (a*a*m != n)
 
+
 # For 53 bit precision floats, the _sqrt_frac() shift is 109.
 _sqrt_shift: int = 2 * sys.float_info.mant_dig + 3
 
+
 def _sqrt_frac(n: int, m: int) -> float:
     """Square root of n/m as a float, correctly rounded."""
     # See principle and proof sketch at: https://bugs.python.org/msg407078
@@ -327,6 +330,31 @@ def _sqrt_frac(n: int, m: int) -> float:
     return numerator / denominator   # Convert to float
 
 
+def _deci_sqrt(n: int, m: int) -> Decimal:
+    """Square root of n/m as a float, correctly rounded."""
+    # Premise:  For decimal, computing sqrt(m / n) can be off by 1 ulp.
+    # Method:   Check the result, moving up or down a step if needed.
+    if not n:
+        return 0.0
+
+    f_square = Fraction(n, m)
+
+    d_mid = (Decimal(n) / Decimal(m)).sqrt()
+    f_mid = Fraction(*d_mid.as_integer_ratio())
+
+    d_plus = d_mid.next_plus()
+    f_plus = Fraction(*d_plus.as_integer_ratio())
+    if f_square > ((f_mid + f_plus) / 2) ** 2:
+        return d_plus
+
+    d_minus = d_mid.next_minus()
+    f_minus = Fraction(*d_minus.as_integer_ratio())
+    if f_square < ((f_mid + f_minus) / 2) ** 2:
+        return d_minus
+
+    return d_mid
+
+
 # === Measures of central tendency (averages) ===
 
 def mean(data):
@@ -888,9 +916,8 @@ def pstdev(data, mu=None):
         raise StatisticsError('pstdev requires at least one data point')
     T, ss = _ss(data, mu)
     mss = ss / n
-    if hasattr(T, 'sqrt'):
-        var = _convert(mss, T)
-        return var.sqrt()
+    if issubclass(T, Decimal):
+        return _deci_sqrt(mss.numerator, mss.denominator)
     return _sqrt_frac(mss.numerator, mss.denominator)
 
 

diff --git a/Lib/test/test_statistics.py b/Lib/test/test_statistics.py
@@ -2219,7 +2219,42 @@ def is_root_correctly_rounded(x: Fraction, root: float) -> bool:
             statistics._sqrt_frac(1, 0)
 
         # The result is well defined if both inputs are negative
-        self.assertAlmostEqual(statistics._sqrt_frac(-2, -1), math.sqrt(2.0))
+        self.assertEqual(statistics._sqrt_frac(-2, -1), statistics._sqrt_frac(2, 1))
+
+    def test_deci_sqrt(self):
+        root: Decimal
+        numerator: int
+        denominator: int
+
+        for root, numerator, denominator in [
+            (Decimal('0.4481904599041192673635338663'), 200874688349065940678243576378, 1000000000000000000000000000000),  # No adj
+            (Decimal('0.7924949131383786609961759598'), 628048187350206338833590574929, 1000000000000000000000000000000),  # Adj up
+            (Decimal('0.8500554152289934068192208727'), 722594208960136395984391238251, 1000000000000000000000000000000),  # Adj down
+        ]:
+            with decimal.localcontext(decimal.DefaultContext):
+                self.assertEqual(statistics._deci_sqrt(numerator, denominator), root)
+
+            # Confirm expected root with a quad precision decimal computation
+            with decimal.localcontext(decimal.DefaultContext) as ctx:
+                ctx.prec *= 4
+                high_prec_root = (Decimal(numerator) / Decimal(denominator)).sqrt()
+            with decimal.localcontext(decimal.DefaultContext):
+                target_root = +high_prec_root
+            self.assertEqual(root, target_root)
+
+        # Verify that corner cases and error handling match Decimal.sqrt()
+        self.assertEqual(statistics._deci_sqrt(0, 1), 0.0)
+        with self.assertRaises(decimal.InvalidOperation):
+            statistics._deci_sqrt(-1, 1)
+        with self.assertRaises(decimal.InvalidOperation):
+            statistics._deci_sqrt(1, -1)
+
+        # Error handling for zero denominator matches that for Fraction(1, 0)
+        with self.assertRaises(ZeroDivisionError):
+            statistics._deci_sqrt(1, 0)
+
+        # The result is well defined if both inputs are negative
+        self.assertEqual(statistics._deci_sqrt(-2, -1), statistics._deci_sqrt(2, 1))
 
 
 class TestStdev(VarianceStdevMixin, NumericTestCase):