pythongh-121149: improve accuracy of builtin sum() for complex inputs

Author: skirpichev

Doc/library/functions.rst

@@ -1934,6 +1934,9 @@ are always available.  They are listed here in alphabetical order.
    .. versionchanged:: 3.12 Summation of floats switched to an algorithm
       that gives higher accuracy and better commutativity on most builds.
 
+   .. versionchanged:: 3.14 Added specialization for summation of complexes,
+      using same algorithm as for summation of floats.
+
 
 .. class:: super()
            super(type, object_or_type=None)

Lib/test/test_builtin.py

@@ -1775,6 +1775,10 @@ def __getitem__(self, index):
     def test_sum_accuracy(self):
         self.assertEqual(sum([0.1] * 10), 1.0)
         self.assertEqual(sum([1.0, 10E100, 1.0, -10E100]), 2.0)
+        self.assertEqual(sum([1.0, 10E100, 1.0, -10E100, 2j]), 2+2j)
+        self.assertEqual(sum([2+1j, 10E100j, 1j, -10E100j]), 2+2j)
+        self.assertEqual(sum([1j, 1, 10E100j, 1j, 1.0, -10E100j]), 2+2j)
+        self.assertEqual(sum([0.1j]*10 + [fractions.Fraction(1, 10)]), 0.1+1j)
 
     def test_type(self):
         self.assertEqual(type(''),  type('123'))

Misc/NEWS.d/next/Core and Builtins/2024-06-30-03-48-10.gh-issue-121149.lLBMKe.rst

@@ -0,0 +1,2 @@
+Added specialization for summation of complexes, this also improves accuracy
+of builtin :func:`sum` for such inputs.  Patch by Sergey B Kirpichev.

Python/bltinmodule.c

@@ -2516,6 +2516,41 @@ Without arguments, equivalent to locals().\n\
 With an argument, equivalent to object.__dict__.");
 
 
+/* Improved Kahan–Babuška algorithm by Arnold Neumaier
+   Neumaier, A. (1974), Rundungsfehleranalyse einiger Verfahren
+   zur Summation endlicher Summen.  Z. angew. Math. Mech.,
+   54: 39-51. https://doi.org/10.1002/zamm.19740540106
+   https://en.wikipedia.org/wiki/Kahan_summation_algorithm#Further_enhancements
+ */
+
+typedef struct {
+    double sum;  /* accumulator */
+    double c;  /* a running compensation for lost low-order bits */
+} _csum;
+
+static inline void
+_csum_neumaier_step(_csum *v, double x)
+{
+    double t = v->sum + x;
+    if (fabs(v->sum) >= fabs(x)) {
+        v->c += (v->sum - t) + x;
+    } else {
+        v->c += (x - t) + v->sum;
+    }
+    v->sum = t;
+}
+
+static inline void
+_csum_neumaier_finalize(_csum *v)
+{
+    /* Avoid losing the sign on a negative result,
+       and don't let adding the compensation convert
+       an infinite or overflowed sum to a NaN. */
+    if (v->c && isfinite(v->c)) {
+        v->sum += v->c;
+    }
+}
+
 /*[clinic input]
 sum as builtin_sum
 
@@ -2628,37 +2663,69 @@ builtin_sum_impl(PyObject *module, PyObject *iterable, PyObject *start)
     }
 
     if (PyFloat_CheckExact(result)) {
-        double f_result = PyFloat_AS_DOUBLE(result);
-        double c = 0.0;
+        _csum f_result = {PyFloat_AS_DOUBLE(result), 0.0};
         Py_SETREF(result, NULL);
         while(result == NULL) {
             item = PyIter_Next(iter);
             if (item == NULL) {
                 Py_DECREF(iter);
                 if (PyErr_Occurred())
                     return NULL;
-                /* Avoid losing the sign on a negative result,
-                   and don't let adding the compensation convert
-                   an infinite or overflowed sum to a NaN. */
-                if (c && isfinite(c)) {
-                    f_result += c;
-                }
-                return PyFloat_FromDouble(f_result);
+                _csum_neumaier_finalize(&f_result);
+                return PyFloat_FromDouble(f_result.sum);
             }
             if (PyFloat_CheckExact(item)) {
-                // Improved Kahan–Babuška algorithm by Arnold Neumaier
-                // Neumaier, A. (1974), Rundungsfehleranalyse einiger Verfahren
-                // zur Summation endlicher Summen.  Z. angew. Math. Mech.,
-                // 54: 39-51. https://doi.org/10.1002/zamm.19740540106
-                // https://en.wikipedia.org/wiki/Kahan_summation_algorithm#Further_enhancements
-                double x = PyFloat_AS_DOUBLE(item);
-                double t = f_result + x;
-                if (fabs(f_result) >= fabs(x)) {
-                    c += (f_result - t) + x;
-                } else {
-                    c += (x - t) + f_result;
+                _csum_neumaier_step(&f_result, PyFloat_AS_DOUBLE(item));
+                _Py_DECREF_SPECIALIZED(item, _PyFloat_ExactDealloc);
+                continue;
+            }
+            if (PyLong_Check(item)) {
+                long value;
+                int overflow;
+                value = PyLong_AsLongAndOverflow(item, &overflow);
+                if (!overflow) {
+                    f_result.sum += (double)value;
+                    Py_DECREF(item);
+                    continue;
                 }
-                f_result = t;
+            }
+            _csum_neumaier_finalize(&f_result);
+            result = PyFloat_FromDouble(f_result.sum);
+            if (result == NULL) {
+                Py_DECREF(item);
+                Py_DECREF(iter);
+                return NULL;
+            }
+            temp = PyNumber_Add(result, item);
+            Py_DECREF(result);
+            Py_DECREF(item);
+            result = temp;
+            if (result == NULL) {
+                Py_DECREF(iter);
+                return NULL;
+            }
+        }
+    }
+
+    if (PyComplex_CheckExact(result)) {
+        Py_complex z = PyComplex_AsCComplex(result);
+        _csum cr_result = {z.real, 0.0};
+        _csum ci_result = {z.imag, 0.0};
+        Py_SETREF(result, NULL);
+        while(result == NULL) {
+            item = PyIter_Next(iter);
+            if (item == NULL) {
+                Py_DECREF(iter);
+                if (PyErr_Occurred())
+                    return NULL;
+                _csum_neumaier_finalize(&cr_result);
+                _csum_neumaier_finalize(&ci_result);
+                return PyComplex_FromDoubles(cr_result.sum, ci_result.sum);
+            }
+            if (PyComplex_CheckExact(item)) {
+                z = PyComplex_AsCComplex(item);
+                _csum_neumaier_step(&cr_result, z.real);
+                _csum_neumaier_step(&ci_result, z.imag);
                 _Py_DECREF_SPECIALIZED(item, _PyFloat_ExactDealloc);
                 continue;
             }
@@ -2667,15 +2734,22 @@ builtin_sum_impl(PyObject *module, PyObject *iterable, PyObject *start)
                 int overflow;
                 value = PyLong_AsLongAndOverflow(item, &overflow);
                 if (!overflow) {
-                    f_result += (double)value;
+                    cr_result.sum += (double)value;
+                    ci_result.sum += 0.0;
                     Py_DECREF(item);
                     continue;
                 }
             }
-            if (c && isfinite(c)) {
-                f_result += c;
+            if (PyFloat_Check(item)) {
+                double value = PyFloat_AS_DOUBLE(item);
+                cr_result.sum += value;
+                ci_result.sum += 0.0;
+                Py_DECREF(item);
+                continue;
             }
-            result = PyFloat_FromDouble(f_result);
+            _csum_neumaier_finalize(&cr_result);
+            _csum_neumaier_finalize(&ci_result);
+            result = PyComplex_FromDoubles(cr_result.sum, ci_result.sum);
             if (result == NULL) {
                 Py_DECREF(item);
                 Py_DECREF(iter);