Add fsum to the math module

carbotaniuman · carbotaniuman · commit 0ff37540ab30 · 2020-12-19T19:32:08.000-06:00
diff --git a/Lib/test/test_math.py b/Lib/test/test_math.py
@@ -599,102 +599,101 @@ def testfrexp(name, result, expected):
         self.assertTrue(math.isnan(math.frexp(NAN)[0]))
 
 
-    # TODO: RUSTPYTHON
-    # @requires_IEEE_754
-    # @unittest.skipIf(HAVE_DOUBLE_ROUNDING,
-    #                      "fsum is not exact on machines with double rounding")
-    # def testFsum(self):
-    #     # math.fsum relies on exact rounding for correct operation.
-    #     # There's a known problem with IA32 floating-point that causes
-    #     # inexact rounding in some situations, and will cause the
-    #     # math.fsum tests below to fail; see issue #2937.  On non IEEE
-    #     # 754 platforms, and on IEEE 754 platforms that exhibit the
-    #     # problem described in issue #2937, we simply skip the whole
-    #     # test.
-
-    #     # Python version of math.fsum, for comparison.  Uses a
-    #     # different algorithm based on frexp, ldexp and integer
-    #     # arithmetic.
-    #     from sys import float_info
-    #     mant_dig = float_info.mant_dig
-    #     etiny = float_info.min_exp - mant_dig
-
-    #     def msum(iterable):
-    #         """Full precision summation.  Compute sum(iterable) without any
-    #         intermediate accumulation of error.  Based on the 'lsum' function
-    #         at http://code.activestate.com/recipes/393090/
-    #         """
-    #         tmant, texp = 0, 0
-    #         for x in iterable:
-    #             mant, exp = math.frexp(x)
-    #             mant, exp = int(math.ldexp(mant, mant_dig)), exp - mant_dig
-    #             if texp > exp:
-    #                 tmant <<= texp-exp
-    #                 texp = exp
-    #             else:
-    #                 mant <<= exp-texp
-    #             tmant += mant
-    #         # Round tmant * 2**texp to a float.  The original recipe
-    #         # used float(str(tmant)) * 2.0**texp for this, but that's
-    #         # a little unsafe because str -> float conversion can't be
-    #         # relied upon to do correct rounding on all platforms.
-    #         tail = max(len(bin(abs(tmant)))-2 - mant_dig, etiny - texp)
-    #         if tail > 0:
-    #             h = 1 << (tail-1)
-    #             tmant = tmant // (2*h) + bool(tmant & h and tmant & 3*h-1)
-    #             texp += tail
-    #         return math.ldexp(tmant, texp)
-
-    #     test_values = [
-    #         ([], 0.0),
-    #         ([0.0], 0.0),
-    #         ([1e100, 1.0, -1e100, 1e-100, 1e50, -1.0, -1e50], 1e-100),
-    #         ([2.0**53, -0.5, -2.0**-54], 2.0**53-1.0),
-    #         ([2.0**53, 1.0, 2.0**-100], 2.0**53+2.0),
-    #         ([2.0**53+10.0, 1.0, 2.0**-100], 2.0**53+12.0),
-    #         ([2.0**53-4.0, 0.5, 2.0**-54], 2.0**53-3.0),
-    #         ([1./n for n in range(1, 1001)],
-    #          float.fromhex('0x1.df11f45f4e61ap+2')),
-    #         ([(-1.)**n/n for n in range(1, 1001)],
-    #          float.fromhex('-0x1.62a2af1bd3624p-1')),
-    #         ([1e16, 1., 1e-16], 10000000000000002.0),
-    #         ([1e16-2., 1.-2.**-53, -(1e16-2.), -(1.-2.**-53)], 0.0),
-    #         # exercise code for resizing partials array
-    #         ([2.**n - 2.**(n+50) + 2.**(n+52) for n in range(-1074, 972, 2)] +
-    #          [-2.**1022],
-    #          float.fromhex('0x1.5555555555555p+970')),
-    #         ]
-
-    #     # Telescoping sum, with exact differences (due to Sterbenz)
-    #     terms = [1.7**i for i in range(1001)]
-    #     test_values.append((
-    #         [terms[i+1] - terms[i] for i in range(1000)] + [-terms[1000]],
-    #         -terms[0]
-    #     ))
-
-    #     for i, (vals, expected) in enumerate(test_values):
-    #         try:
-    #             actual = math.fsum(vals)
-    #         except OverflowError:
-    #             self.fail("test %d failed: got OverflowError, expected %r "
-    #                       "for math.fsum(%.100r)" % (i, expected, vals))
-    #         except ValueError:
-    #             self.fail("test %d failed: got ValueError, expected %r "
-    #                       "for math.fsum(%.100r)" % (i, expected, vals))
-    #         self.assertEqual(actual, expected)
-
-    #     from random import random, gauss, shuffle
-    #     for j in range(1000):
-    #         vals = [7, 1e100, -7, -1e100, -9e-20, 8e-20] * 10
-    #         s = 0
-    #         for i in range(200):
-    #             v = gauss(0, random()) ** 7 - s
-    #             s += v
-    #             vals.append(v)
-    #         shuffle(vals)
-
-    #         s = msum(vals)
-    #         self.assertEqual(msum(vals), math.fsum(vals))
+    @requires_IEEE_754
+    @unittest.skipIf(HAVE_DOUBLE_ROUNDING,
+                         "fsum is not exact on machines with double rounding")
+    def testFsum(self):
+        # math.fsum relies on exact rounding for correct operation.
+        # There's a known problem with IA32 floating-point that causes
+        # inexact rounding in some situations, and will cause the
+        # math.fsum tests below to fail; see issue #2937.  On non IEEE
+        # 754 platforms, and on IEEE 754 platforms that exhibit the
+        # problem described in issue #2937, we simply skip the whole
+        # test.
+
+        # Python version of math.fsum, for comparison.  Uses a
+        # different algorithm based on frexp, ldexp and integer
+        # arithmetic.
+        from sys import float_info
+        mant_dig = float_info.mant_dig
+        etiny = float_info.min_exp - mant_dig
+
+        def msum(iterable):
+            """Full precision summation.  Compute sum(iterable) without any
+            intermediate accumulation of error.  Based on the 'lsum' function
+            at http://code.activestate.com/recipes/393090/
+            """
+            tmant, texp = 0, 0
+            for x in iterable:
+                mant, exp = math.frexp(x)
+                mant, exp = int(math.ldexp(mant, mant_dig)), exp - mant_dig
+                if texp > exp:
+                    tmant <<= texp-exp
+                    texp = exp
+                else:
+                    mant <<= exp-texp
+                tmant += mant
+            # Round tmant * 2**texp to a float.  The original recipe
+            # used float(str(tmant)) * 2.0**texp for this, but that's
+            # a little unsafe because str -> float conversion can't be
+            # relied upon to do correct rounding on all platforms.
+            tail = max(len(bin(abs(tmant)))-2 - mant_dig, etiny - texp)
+            if tail > 0:
+                h = 1 << (tail-1)
+                tmant = tmant // (2*h) + bool(tmant & h and tmant & 3*h-1)
+                texp += tail
+            return math.ldexp(tmant, texp)
+
+        test_values = [
+            ([], 0.0),
+            ([0.0], 0.0),
+            ([1e100, 1.0, -1e100, 1e-100, 1e50, -1.0, -1e50], 1e-100),
+            ([2.0**53, -0.5, -2.0**-54], 2.0**53-1.0),
+            ([2.0**53, 1.0, 2.0**-100], 2.0**53+2.0),
+            ([2.0**53+10.0, 1.0, 2.0**-100], 2.0**53+12.0),
+            ([2.0**53-4.0, 0.5, 2.0**-54], 2.0**53-3.0),
+            ([1./n for n in range(1, 1001)],
+             float.fromhex('0x1.df11f45f4e61ap+2')),
+            ([(-1.)**n/n for n in range(1, 1001)],
+             float.fromhex('-0x1.62a2af1bd3624p-1')),
+            ([1e16, 1., 1e-16], 10000000000000002.0),
+            ([1e16-2., 1.-2.**-53, -(1e16-2.), -(1.-2.**-53)], 0.0),
+            # exercise code for resizing partials array
+            ([2.**n - 2.**(n+50) + 2.**(n+52) for n in range(-1074, 972, 2)] +
+             [-2.**1022],
+             float.fromhex('0x1.5555555555555p+970')),
+            ]
+
+        # Telescoping sum, with exact differences (due to Sterbenz)
+        terms = [1.7**i for i in range(1001)]
+        test_values.append((
+            [terms[i+1] - terms[i] for i in range(1000)] + [-terms[1000]],
+            -terms[0]
+        ))
+
+        for i, (vals, expected) in enumerate(test_values):
+            try:
+                actual = math.fsum(vals)
+            except OverflowError:
+                self.fail("test %d failed: got OverflowError, expected %r "
+                          "for math.fsum(%.100r)" % (i, expected, vals))
+            except ValueError:
+                self.fail("test %d failed: got ValueError, expected %r "
+                          "for math.fsum(%.100r)" % (i, expected, vals))
+            self.assertEqual(actual, expected)
+
+        from random import random, gauss, shuffle
+        for j in range(1000):
+            vals = [7, 1e100, -7, -1e100, -9e-20, 8e-20] * 10
+            s = 0
+            for i in range(200):
+                v = gauss(0, random()) ** 7 - s
+                s += v
+                vals.append(v)
+            shuffle(vals)
+
+            s = msum(vals)
+            self.assertEqual(msum(vals), math.fsum(vals))
 
 
     # Python 3.9
@@ -1559,12 +1558,12 @@ def testIsinf(self):
     #     self.assertTrue(math.isnan(math.nan))
 
     # TODO: RUSTPYTHON
-    # @requires_IEEE_754
-    # def test_inf_constant(self):
-    #     self.assertTrue(math.isinf(math.inf))
-    #     self.assertGreater(math.inf, 0.0)
-    #     self.assertEqual(math.inf, float("inf"))
-    #     self.assertEqual(-math.inf, float("-inf"))
+    @requires_IEEE_754
+    def test_inf_constant(self):
+        self.assertTrue(math.isinf(math.inf))
+        self.assertGreater(math.inf, 0.0)
+        self.assertEqual(math.inf, float("inf"))
+        self.assertEqual(-math.inf, float("-inf"))
 
     # RED_FLAG 16-Oct-2000 Tim
     # While 2.0 is more consistent about exceptions than previous releases, it
diff --git a/vm/src/stdlib/math.rs b/vm/src/stdlib/math.rs
@@ -11,7 +11,7 @@ use statrs::function::gamma::{gamma, ln_gamma};
 use crate::builtins::float::{self, IntoPyFloat, PyFloatRef};
 use crate::builtins::int::{self, PyInt, PyIntRef};
 use crate::function::{Args, OptionalArg};
-use crate::pyobject::{BorrowValue, Either, PyObjectRef, PyResult, TypeProtocol};
+use crate::pyobject::{BorrowValue, Either, PyIterable, PyObjectRef, PyResult, TypeProtocol};
 use crate::vm::VirtualMachine;
 use rustpython_common::float_ops;
 
@@ -372,6 +372,100 @@ fn math_lcm(args: Args<PyIntRef>) -> BigInt {
     math_perf_arb_len_int_op(args, |x, y| x.lcm(y.borrow_value()), BigInt::one())
 }
 
+fn math_fsum(iter: PyIterable<IntoPyFloat>, vm: &VirtualMachine) -> PyResult<f64> {
+    let mut iter = iter.iter(vm)?;
+    let mut partials = vec![];
+    let mut special_sum = 0.0;
+    let mut inf_sum = 0.0;
+
+    while let Some(obj) = iter.next() {
+        let mut x = obj?.to_f64();
+
+        let xsave = x;
+        let mut j = 0;
+        // This inner loop applies `hi`/`lo` summation to each
+        // partial so that the list of partial sums remains exact.
+        for i in 0..partials.len() {
+            let mut y: f64 = partials[i];
+            if x.abs() < y.abs() {
+                let t = x;
+                x = y;
+                y = t;
+            }
+            // Rounded `x+y` is stored in `hi` with round-off stored in
+            // `lo`. Together `hi+lo` are exactly equal to `x+y`.
+            let hi = x + y;
+            let lo = y - (hi - x);
+            if lo != 0.0 {
+                partials[j] = lo;
+                j += 1;
+            }
+            x = hi;
+        }
+
+        if !x.is_finite() {
+            // a nonfinite x could arise either as
+            // a result of intermediate overflow, or
+            // as a result of a nan or inf in the
+            // summands
+            if xsave.is_finite() {
+                return Err(vm.new_overflow_error("intermediate overflow in fsum".to_owned()));
+            }
+            if xsave.is_infinite() {
+                inf_sum += xsave;
+            }
+            special_sum += xsave;
+            // reset partials
+            partials.clear();
+        }
+
+        if j >= partials.len() {
+            partials.push(x);
+        } else {
+            partials[j] = x;
+            partials.truncate(j + 1);
+        }
+    }
+    if special_sum != 0.0 {
+        return if inf_sum.is_nan() {
+            Err(vm.new_overflow_error("-inf + inf in fsum".to_owned()))
+        } else {
+            Ok(special_sum)
+        };
+    }
+
+    let mut n = partials.len();
+    if n > 0 {
+        n -= 1;
+        let mut hi = partials[n];
+
+        let mut lo = 0.0;
+        while n > 0 {
+            let x = hi;
+
+            n -= 1;
+            let y = partials[n];
+
+            hi = x + y;
+            lo = y - (hi - x);
+            if lo != 0.0 {
+                break;
+            }
+        }
+        if n > 0 && ((lo < 0.0 && partials[n - 1] < 0.0) || (lo > 0.0 && partials[n - 1] > 0.0)) {
+            let y = lo + lo;
+            let x = hi + y;
+            if y == x - hi {
+                hi = x;
+            }
+        }
+
+        Ok(hi)
+    } else {
+        Ok(0.0)
+    }
+}
+
 fn math_factorial(value: PyIntRef, vm: &VirtualMachine) -> PyResult<BigInt> {
     let value = value.borrow_value();
     if value.is_negative() {
@@ -516,6 +610,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "ldexp" => named_function!(ctx, math, ldexp),
         "modf" => named_function!(ctx, math, modf),
         "fmod" => named_function!(ctx, math, fmod),
+        "fsum" => named_function!(ctx, math, fsum),
         "remainder" => named_function!(ctx, math, remainder),
 
         // Rounding functions:
