Skip to content

Commit

Permalink
GH-102670: Use sumprod() to simplify, speed up, and improve accuracy …
Browse files Browse the repository at this point in the history
…of statistics functions (GH-102649)
  • Loading branch information
rhettinger authored Mar 14, 2023
1 parent 61479d4 commit 457e4d1
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 13 deletions.
26 changes: 14 additions & 12 deletions Lib/statistics.py
Original file line number Diff line number Diff line change
Expand Up @@ -1036,7 +1036,7 @@ def covariance(x, y, /):
raise StatisticsError('covariance requires at least two data points')
xbar = fsum(x) / n
ybar = fsum(y) / n
sxy = fsum((xi - xbar) * (yi - ybar) for xi, yi in zip(x, y))
sxy = sumprod((xi - xbar for xi in x), (yi - ybar for yi in y))
return sxy / (n - 1)


Expand Down Expand Up @@ -1074,11 +1074,14 @@ def correlation(x, y, /, *, method='linear'):
start = (n - 1) / -2 # Center rankings around zero
x = _rank(x, start=start)
y = _rank(y, start=start)
xbar = fsum(x) / n
ybar = fsum(y) / n
sxy = fsum((xi - xbar) * (yi - ybar) for xi, yi in zip(x, y))
sxx = fsum((d := xi - xbar) * d for xi in x)
syy = fsum((d := yi - ybar) * d for yi in y)
else:
xbar = fsum(x) / n
ybar = fsum(y) / n
x = [xi - xbar for xi in x]
y = [yi - ybar for yi in y]
sxy = sumprod(x, y)
sxx = sumprod(x, x)
syy = sumprod(y, y)
try:
return sxy / sqrt(sxx * syy)
except ZeroDivisionError:
Expand Down Expand Up @@ -1131,14 +1134,13 @@ def linear_regression(x, y, /, *, proportional=False):
raise StatisticsError('linear regression requires that both inputs have same number of data points')
if n < 2:
raise StatisticsError('linear regression requires at least two data points')
if proportional:
sxy = fsum(xi * yi for xi, yi in zip(x, y))
sxx = fsum(xi * xi for xi in x)
else:
if not proportional:
xbar = fsum(x) / n
ybar = fsum(y) / n
sxy = fsum((xi - xbar) * (yi - ybar) for xi, yi in zip(x, y))
sxx = fsum((d := xi - xbar) * d for xi in x)
x = [xi - xbar for xi in x] # List because used three times below
y = (yi - ybar for yi in y) # Generator because only used once below
sxy = sumprod(x, y) + 0.0 # Add zero to coerce result to a float
sxx = sumprod(x, x)
try:
slope = sxy / sxx # equivalent to: covariance(x, y) / variance(x)
except ZeroDivisionError:
Expand Down
12 changes: 11 additions & 1 deletion Lib/test/test_statistics.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
"""Test suite for statistics module, including helper NumericTestCase and
x = """Test suite for statistics module, including helper NumericTestCase and
approx_equal function.
"""
Expand Down Expand Up @@ -2610,6 +2610,16 @@ def test_proportional(self):
self.assertAlmostEqual(slope, 20 + 1/150)
self.assertEqual(intercept, 0.0)

def test_float_output(self):
x = [Fraction(2, 3), Fraction(3, 4)]
y = [Fraction(4, 5), Fraction(5, 6)]
slope, intercept = statistics.linear_regression(x, y)
self.assertTrue(isinstance(slope, float))
self.assertTrue(isinstance(intercept, float))
slope, intercept = statistics.linear_regression(x, y, proportional=True)
self.assertTrue(isinstance(slope, float))
self.assertTrue(isinstance(intercept, float))

class TestNormalDist:

# General note on precision: The pdf(), cdf(), and overlap() methods
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Optimized fmean(), correlation(), covariance(), and linear_regression()
using the new math.sumprod() function.

0 comments on commit 457e4d1

Please sign in to comment.