Avoid cancellation in linear regression with large coordinate offsets

algorithms/math/linear_regression.py

@@ -14,6 +14,9 @@
 def linear_regression(x: list[float], y: list[float]) -> tuple[float, float]:
     """Return (slope, intercept) for the best-fit line.
 
+    Shift observations by the first point before centering. This avoids
+    cancellation and preserves midpoints near large common offsets.
+
     >>> m, b = linear_regression([1, 2, 3, 4, 5], [2, 4, 5, 4, 5])
     >>> round(m, 4)
     0.6
@@ -24,16 +27,20 @@ def linear_regression(x: list[float], y: list[float]) -> tuple[float, float]:
     if n != len(y) or n < 2:
         msg = "x and y must have at least 2 equal-length elements"
         raise ValueError(msg)
-    sum_x = sum(x)
-    sum_y = sum(y)
-    sum_xy = sum(xi * yi for xi, yi in zip(x, y, strict=False))
-    sum_x2 = sum(xi * xi for xi in x)
-    denom = n * sum_x2 - sum_x * sum_x
+    x_shifted = [xi - x[0] for xi in x]
+    y_shifted = [yi - y[0] for yi in y]
+    x_mean = math.fsum(x_shifted) / n
+    y_mean = math.fsum(y_shifted) / n
+    denom = math.fsum((xi - x_mean) ** 2 for xi in x_shifted)
     if denom == 0:
         msg = "Vertical line — slope is undefined"
         raise ValueError(msg)
-    slope = (n * sum_xy - sum_x * sum_y) / denom
-    intercept = (sum_y - slope * sum_x) / n
+    numerator = math.fsum(
+        (xi - x_mean) * (yi - y_mean)
+        for xi, yi in zip(x_shifted, y_shifted, strict=False)
+    )
+    slope = numerator / denom
+    intercept = (y[0] - slope * x[0]) + (y_mean - slope * x_mean)
     return slope, intercept
 
 

tests/test_linear_regression_offsets.py

@@ -0,0 +1,43 @@
+"""Large-offset regressions must not be mistaken for vertical lines."""
+
+import statistics
+
+import pytest
+
+from algorithms.math.linear_regression import linear_regression, r_squared, rmse
+
+
+@pytest.mark.parametrize("x_offset", [0, 10**8, 10**12, -(10**12)])
+@pytest.mark.parametrize("y_offset", [0, 10**12, -(10**12)])
+def test_exact_line_after_translation(x_offset, y_offset):
+    x = [float(x_offset + i) for i in range(5)]
+    y = [float(y_offset + 2 * i) for i in range(5)]
+    slope, intercept = linear_regression(x, y)
+    assert slope == pytest.approx(2)
+    assert intercept == pytest.approx(y_offset - 2 * x_offset)
+    assert r_squared(x, y) == pytest.approx(1)
+    assert rmse(x, y) == pytest.approx(0)
+
+
+def test_noisy_offset_matches_standard_library():
+    x = [10**12 + i for i in range(7)]
+    y = [1, 4, 4, 8, 7, 12, 13]
+    assert linear_regression(x, y) == pytest.approx(statistics.linear_regression(x, y))
+
+
+def test_constant_x_still_raises():
+    with pytest.raises(ValueError, match="undefined"):
+        linear_regression([10**12, 10**12], [1, 2])
+
+
+@pytest.mark.parametrize("offset", [1e16, -1e16])
+def test_two_point_midpoint_is_not_rounded(offset):
+    slope, intercept = linear_regression([offset, offset + 2], [0, 4])
+    assert slope == 2
+    assert intercept == -2 * offset
+
+
+def test_large_offsets_on_both_axes():
+    slope, intercept = linear_regression([1e16, 1e16 + 2], [1e16, 1e16 + 4])
+    assert slope == 2
+    assert intercept == -1e16