# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""Regression metric utilities built on top of NumPy."""
from __future__ import annotations
from typing import Any
import numpy as np
[docs]
def mean_squared_error(
y_true: Any,
y_pred: Any,
*,
sample_weight: Any | None = None,
squared: bool = True,
) -> float:
"""Compute the (weighted) mean squared error between true and predicted values.
Parameters
----------
y_true : array-like
True target values. Supports NumPy arrays, PyTorch tensors, JAX arrays,
and any array-like object convertible to NumPy.
y_pred : array-like
Predicted values. Must have the same shape as ``y_true``.
sample_weight : array-like, optional
Sample weights. If provided, the weighted mean is computed.
Must be broadcastable to the shape of ``y_true``.
squared : bool, default=True
If True, return the mean squared error (MSE).
If False, return the root mean squared error (RMSE).
Returns
-------
float
The (weighted) mean squared error or its square root.
See Also
--------
root_mean_squared_error : Convenience wrapper always returning RMSE.
mean_absolute_error : Mean absolute error (less sensitive to outliers).
r_squared : Coefficient of determination.
"""
y_true_np = _to_numpy(y_true)
y_pred_np = _to_numpy(y_pred)
_validate_shape(y_true_np, y_pred_np)
diff = y_true_np - y_pred_np
squared_error = diff * diff
mean = _weighted_mean(squared_error, sample_weight)
if squared:
return float(mean)
return float(np.sqrt(mean))
[docs]
def mean_absolute_error(
y_true: Any,
y_pred: Any,
*,
sample_weight: Any | None = None,
) -> float:
"""Compute the (weighted) mean absolute error between true and predicted values.
Parameters
----------
y_true : array-like
True target values. Supports NumPy arrays, PyTorch tensors, JAX arrays,
and any array-like object convertible to NumPy.
y_pred : array-like
Predicted values. Must have the same shape as ``y_true``.
sample_weight : array-like, optional
Sample weights. If provided, the weighted mean is computed.
Must be broadcastable to the shape of ``y_true``.
Returns
-------
float
The (weighted) mean absolute error.
See Also
--------
mean_squared_error : Mean squared error (more sensitive to outliers).
root_mean_squared_error : Root mean squared error.
r_squared : Coefficient of determination.
"""
y_true_np = _to_numpy(y_true)
y_pred_np = _to_numpy(y_pred)
_validate_shape(y_true_np, y_pred_np)
absolute_error = np.abs(y_true_np - y_pred_np)
mean = _weighted_mean(absolute_error, sample_weight)
return float(mean)
[docs]
def root_mean_squared_error(
y_true: Any,
y_pred: Any,
*,
sample_weight: Any | None = None,
) -> float:
"""Compute the (weighted) root mean squared error.
This is equivalent to ``mean_squared_error(..., squared=False)``.
Parameters
----------
y_true : array-like
True target values.
y_pred : array-like
Predicted values. Must have the same shape as ``y_true``.
sample_weight : array-like, optional
Sample weights. If provided, the weighted RMSE is computed.
Returns
-------
float
The (weighted) root mean squared error.
See Also
--------
mean_squared_error : Mean squared error (optionally squared or not).
mean_absolute_error : Mean absolute error.
r_squared : Coefficient of determination.
"""
return mean_squared_error(y_true, y_pred, sample_weight=sample_weight, squared=False)
[docs]
def r_squared(
y_true: Any,
y_pred: Any,
*,
sample_weight: Any | None = None,
) -> float:
"""Calculate R-squared (coefficient of determination).
Parameters
----------
y_true : array-like
True target values.
y_pred : array-like
Predicted values.
sample_weight : array-like, optional
Sample weights.
Returns
-------
float
R-squared value. Best possible score is 1.0, can be negative
(indicating the model performs worse than a horizontal line).
"""
y_true_np = _to_numpy(y_true)
y_pred_np = _to_numpy(y_pred)
_validate_shape(y_true_np, y_pred_np)
if sample_weight is None:
y_mean = np.mean(y_true_np)
ss_tot = np.sum((y_true_np - y_mean) ** 2)
ss_res = np.sum((y_true_np - y_pred_np) ** 2)
else:
weights = _to_numpy(sample_weight)
if weights.shape != y_true_np.shape:
weights = np.broadcast_to(weights, y_true_np.shape)
total_weight = np.sum(weights)
if total_weight <= 0:
raise ValueError("sample weights must have positive sum")
y_mean = np.sum(weights * y_true_np) / total_weight
ss_tot = np.sum(weights * (y_true_np - y_mean) ** 2)
ss_res = np.sum(weights * (y_true_np - y_pred_np) ** 2)
if ss_tot == 0:
# Constant model, R² is undefined, return 0
return 0.0
return float(1.0 - (ss_res / ss_tot))
def _weighted_mean(values: np.ndarray, sample_weight: Any | None) -> float:
if sample_weight is None:
return float(np.mean(values))
weights = _to_numpy(sample_weight)
if weights.shape != values.shape:
weights = np.broadcast_to(weights, values.shape)
total_weight = np.sum(weights)
if total_weight <= 0:
raise ValueError("sample weights must have positive sum")
return float(np.sum(weights * values) / total_weight)
def _validate_shape(y_true: np.ndarray, y_pred: np.ndarray) -> None:
if y_true.shape != y_pred.shape:
raise ValueError("y_true and y_pred must have identical shapes")
def _to_numpy(value: Any) -> np.ndarray:
if isinstance(value, np.ndarray):
return value.astype(np.float64, copy=False)
if hasattr(value, "detach"):
return np.asarray(value.detach().cpu().numpy().astype(np.float64, copy=False))
if hasattr(value, "cpu") and hasattr(value, "numpy"):
return np.asarray(value.cpu().numpy().astype(np.float64, copy=False))
return np.asarray(value, dtype=np.float64)
__all__ = [
"mean_squared_error",
"mean_absolute_error",
"root_mean_squared_error",
"r_squared",
]