fix(r_squared): implement the documented zero_division parameter

metrics/r_squared/r_squared.py

@@ -15,6 +15,8 @@
 """R squared metric."""
 
 
+import warnings
+
 import datasets
 import numpy as np
 
@@ -61,6 +63,14 @@
     >>> r_squared = r2_metric.compute(predictions=[1, 2, 3, 4], references=[0.9, 2.1, 3.2, 3.8])
     >>> print(r_squared)
     0.98
+
+    R^2 is undefined when the references have zero variance, since there is no
+    variance to explain. `zero_division` selects what to return instead:
+
+    >>> r2_metric.compute(predictions=[1, 2, 3, 4], references=[5, 5, 5, 5], zero_division=0)
+    0.0
+    >>> r2_metric.compute(predictions=[5, 5, 5, 5], references=[5, 5, 5, 5], zero_division=1)
+    1.0
 """
 
 
@@ -83,17 +93,25 @@ def _info(self):
             ],
         )
 
-    def _compute(self, predictions=None, references=None):
+    def _compute(self, predictions=None, references=None, zero_division="warn"):
         """
         Computes the coefficient of determination (R-squared) of predictions with respect to references.
 
         Parameters:
             predictions (List or np.ndarray): The predicted values.
             references (List or np.ndarray): The true/reference values.
+            zero_division (0, 1 or "warn"): Value to return when `references` has zero
+                variance, which makes the sum of squared total zero and R^2 undefined.
+                "warn" returns 0.0 and raises a warning.
 
         Returns:
             float: The R-squared value, rounded to 3 decimal places.
         """
+        if zero_division not in (0, 1, "warn"):
+            raise ValueError(
+                f'zero_division must be one of 0, 1 or "warn", got {zero_division!r}'
+            )
+
         predictions = np.array(predictions)
         references = np.array(references)
 
@@ -106,6 +124,22 @@ def _compute(self, predictions=None, references=None):
         # Calculate sum of squared total
         sst = np.sum((references - mean_references) ** 2)
 
+        # R^2 is undefined when the references never vary: there is no variance to
+        # explain, so the ratio divides by zero. Left to numpy this returns -inf
+        # (or nan when the predictions match exactly), which then propagates
+        # silently through any downstream aggregation.
+        if sst == 0:
+            if zero_division == "warn":
+                warnings.warn(
+                    "R^2 is undefined when all `references` are identical (zero "
+                    "variance); returning 0.0. Pass zero_division=0 or 1 to choose "
+                    "the value and silence this warning.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                return 0.0
+            return float(zero_division)
+
         # Calculate R Squared
         r_squared = 1 - (ssr / sst)