diff --git a/metrics/r_squared/r_squared.py b/metrics/r_squared/r_squared.py index 75e0f7e00..62dc28e2e 100644 --- a/metrics/r_squared/r_squared.py +++ b/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)