Loss R2#

r2 #

CLASS DESCRIPTION
R2Loss

Compute \(R^2\) loss with optional weighting.

FUNCTION DESCRIPTION
r2_loss

Compute \(R^2\) loss with optional weighting.

Classes#

R2Loss(reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0, context_prefix: str | None = None) #

Bases: Loss

Compute \(R^2\) loss with optional weighting.

Computes \(R^2 = 1 - \frac{\sum w_i (y_i - \hat{y_i})^2}{\sum w_i (y_i - \bar{y})^2}\), where \(y_i\) are the target values, \(\hat{y_i}\) are the predicted values, \(\bar{y}\) is the mean of the predicted values, and \(w_i\) are optional weights.

For multidimensional targets, computes per-feature means and aggregates the total sum of squares across all features.

PARAMETER DESCRIPTION
reduce

Reduction method.

TYPE: Literal["sum", "mean"], optional, by default "mean" DEFAULT: 'mean'

sign

Loss sign multiplier.

TYPE: float, optional, by default 1.0 DEFAULT: 1.0

METHOD DESCRIPTION
forward

Compute \(R^2\) loss.

Source code in spectre/loss/r2.py
def __init__(
    self,
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
    context_prefix: str | None = None,
):
    super().__init__(reduce=reduce, sign=sign, context_prefix=context_prefix)
Functions#
forward(context: dict, **kwargs) -> torch.Tensor #

Compute \(R^2\) loss.

PARAMETER DESCRIPTION
context

Context dictionary containing target, target_pred, and weights (optional).

TYPE: dict

**kwargs

Additional keyword arguments.

TYPE: dict DEFAULT: {}

RETURNS DESCRIPTION
Tensor

\(R^2\) loss.

Source code in spectre/loss/r2.py
def forward(self, context: dict, **kwargs) -> torch.Tensor:
    """
    Compute $R^2$ loss.

    Parameters
    ----------
    context : dict
        Context dictionary containing `target`, `target_pred`, and `weights`
        (optional).

    **kwargs : dict
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        $R^2$ loss.
    """
    target = self._get_context(context, "target", None)
    target_pred = self._get_context(context, "target_pred", None)
    weights = self._get_context(context, "weights", None)

    if target is None or target_pred is None:
        raise ValueError(
            "`R2Loss` requires `target` and `target_pred` must be provided "
            "in the context."
        )
    return r2_loss(
        target, target_pred, weights=weights, reduce=self.reduce, sign=self.sign
    )

Functions#

r2_loss(target: torch.Tensor, target_pred: torch.Tensor, weights: torch.Tensor | None = None, reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0) -> torch.Tensor #

Compute \(R^2\) loss with optional weighting.

Computes \(R^2 = 1 - \frac{\sum w_i (y_i - \hat{y_i})^2}{\sum w_i (y_i - \bar{y})^2}\), where \(y_i\) are the target values, \(\hat{y_i}\) are the predicted values, \(\bar{y}\) is the mean of the predicted values, and \(w_i\) are optional weights.

For multidimensional targets, computes per-feature means and aggregates the total sum of squares across all features.

PARAMETER DESCRIPTION
target

Target values.

TYPE: Tensor

target_pred

Predicted values.

TYPE: Tensor

weights

Sample weights.

TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

reduce

Reduction method.

TYPE: Literal["sum", "mean"], optional, by default "mean" DEFAULT: 'mean'

sign

Loss sign multiplier.

TYPE: float, optional, by default 1.0 DEFAULT: 1.0

RETURNS DESCRIPTION
Tensor

Loss value.

Source code in spectre/loss/r2.py
def r2_loss(
    target: torch.Tensor,
    target_pred: torch.Tensor,
    weights: torch.Tensor | None = None,
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
) -> torch.Tensor:
    """
    Compute $R^2$ loss with optional weighting.

    Computes
    $R^2 = 1 - \\frac{\\sum w_i (y_i - \\hat{y_i})^2}{\\sum w_i (y_i - \\bar{y})^2}$,
    where $y_i$ are the target values, $\\hat{y_i}$ are the predicted values,
    $\\bar{y}$ is the mean of the predicted values, and $w_i$ are optional weights.

    For multidimensional targets, computes per-feature means and aggregates the total
    sum of squares across all features.

    Parameters
    ----------
    target : torch.Tensor
        Target values.

    target_pred : torch.Tensor
        Predicted values.

    weights : torch.Tensor | None, optional, by default None
        Sample weights.

    reduce : Literal["sum", "mean"], optional, by default "mean"
        Reduction method.

    - "sum": Sum over all elements
    - "mean": Mean over all elements

    sign : float, optional, by default 1.0
        Loss sign multiplier.

    Returns
    -------
    torch.Tensor
        Loss value.
    """
    check_same_shape(target, target_pred)

    if weights is not None:
        weight_sum = torch.sum(weights)
        weights = weights.view(-1, 1)
        target_pred_mean = torch.sum(weights * target_pred, dim=0) / weight_sum
        ss_res = torch.sum(weights * (target - target_pred) ** 2)
        ss_tot = torch.sum(weights * (target - target_pred_mean) ** 2)
    else:
        target_pred_mean = torch.mean(target_pred, dim=0)
        ss_res = torch.sum((target - target_pred) ** 2)
        ss_tot = torch.sum((target - target_pred_mean) ** 2)

    if ss_tot == 0:
        result = 1.0 if ss_res == 0 else 0.0
        return torch.tensor(result, device=target.device, dtype=target.dtype)

    loss = ss_res / ss_tot

    if reduce == "sum":
        loss = loss.sum()
    elif reduce == "mean":
        loss = loss.mean()
    else:
        raise ValueError(
            f"Unknown reduce operation '{reduce}'. Options: 'sum', 'mean'."
        )

    return sign * (1 - loss)