feat: add ml.metrics.mean_squared_error (#559)

* feat: add ml.metrics.mean_squared_error * fix docs * fix docs
googleapis · Apr 3, 2024 · 853c25e · 853c25e
1 parent 90bcec5
commit 853c25e
Show file tree

Hide file tree

Showing 4 changed files with 95 additions and 69 deletions.
diff --git a/bigframes/ml/metrics/__init__.py b/bigframes/ml/metrics/__init__.py
@@ -18,6 +18,7 @@
     auc,
     confusion_matrix,
     f1_score,
+    mean_squared_error,
     precision_score,
     r2_score,
     recall_score,
@@ -35,5 +36,6 @@
     "confusion_matrix",
     "precision_score",
     "f1_score",
+    "mean_squared_error",
     "pairwise",
 ]
diff --git a/bigframes/ml/metrics/_metrics.py b/bigframes/ml/metrics/_metrics.py
@@ -335,3 +335,17 @@ def f1_score(
 
 
 f1_score.__doc__ = inspect.getdoc(vendored_metrics_classification.f1_score)
+
+
+def mean_squared_error(
+    y_true: Union[bpd.DataFrame, bpd.Series],
+    y_pred: Union[bpd.DataFrame, bpd.Series],
+) -> float:
+    y_true_series, y_pred_series = utils.convert_to_series(y_true, y_pred)
+
+    return (y_pred_series - y_true_series).pow(2).sum() / len(y_true_series)
+
+
+mean_squared_error.__doc__ = inspect.getdoc(
+    vendored_metrics_regression.mean_squared_error
+)