DeepHealth/losses.py

from __future__ import annotations

from typing import Iterable

import torch
import torch.nn as nn
import torch.nn.functional as F


PAD_IDX = 0
RESERVED_IDX = 1
NO_EVENT_IDX = 2


def _make_ignore_mask(
    vocab_size: int,
    ignored_idx: Iterable[int],
    device: torch.device,
) -> torch.Tensor:
    ignore_mask = torch.zeros(vocab_size, dtype=torch.bool, device=device)
    for idx in ignored_idx:
        idx = int(idx)
        if 0 <= idx < vocab_size:
            ignore_mask[idx] = True
    return ignore_mask


def _valid_vocab_mask(
    vocab_size: int,
    ignored_idx: Iterable[int],
    device: torch.device,
) -> torch.Tensor:
    return ~_make_ignore_mask(vocab_size, ignored_idx, device)


def _zero_loss_like(logits: torch.Tensor) -> torch.Tensor:
    return logits.sum() * 0.0


def _all_future_at_risk_mask(
    vocab_size: int,
    ignored_idx: Iterable[int],
    logits: torch.Tensor,
    history: torch.Tensor | None,
) -> torch.Tensor:
    """Return outcomes that have not occurred by the query time."""
    batch_size = logits.shape[0]
    at_risk = _valid_vocab_mask(
        vocab_size,
        ignored_idx,
        logits.device,
    ).unsqueeze(0).expand(batch_size, -1).clone()
    if history is None or history.numel() == 0:
        return at_risk
    if history.dim() != 2 or history.shape[0] != batch_size:
        raise ValueError(
            "history must be (B, L). "
            f"Got logits={tuple(logits.shape)}, history={tuple(history.shape)}"
        )

    history = history.to(device=logits.device, dtype=torch.long)
    history_valid = (history >= 0) & (history < vocab_size)
    for idx in ignored_idx:
        history_valid &= history != int(idx)
    safe_history = history.clamp(min=0, max=vocab_size - 1)
    prevalent = torch.zeros(
        (batch_size, vocab_size),
        dtype=torch.long,
        device=logits.device,
    )
    prevalent.scatter_add_(1, safe_history, history_valid.long())
    return at_risk & ~prevalent.bool()


def _all_future_target_mask(
    targets: torch.Tensor,
    at_risk: torch.Tensor,
    ignored_idx: Iterable[int],
) -> tuple[torch.Tensor, torch.Tensor]:
    vocab_size = at_risk.shape[1]
    in_vocab = (targets >= 0) & (targets < vocab_size)
    for idx in ignored_idx:
        in_vocab &= targets != int(idx)
    safe_targets = targets.clamp(min=0, max=vocab_size - 1)
    return in_vocab & at_risk.gather(1, safe_targets), safe_targets


class Delphi2MLoss(nn.Module):
    """Next-token plus exponential time-to-next-token supervision."""

    def __init__(
        self,
        t_min: float = 1.0 / 365.25,
        ignored_tokens: Iterable[int] | None = None,
        exclude_ignored_from_intensity: bool = True,
        max_exp_input: float = 60.0,
        ce_weight: float = 1.0,
        time_weight: float = 1.0,
    ):
        super().__init__()
        self.t_min = float(t_min)
        self.ignored_tokens = (
            [PAD_IDX, RESERVED_IDX]
            if ignored_tokens is None
            else [int(x) for x in ignored_tokens]
        )
        self.exclude_ignored_from_intensity = bool(exclude_ignored_from_intensity)
        self.max_exp_input = float(max_exp_input)
        self.ce_weight = float(ce_weight)
        self.time_weight = float(time_weight)

    def forward(
        self,
        logits: torch.Tensor,
        target_events: torch.Tensor,
        target_times: torch.Tensor,
        current_times: torch.Tensor,
        padding_mask: torch.Tensor,
        return_components: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
        if logits.dim() != 3:
            raise ValueError(f"logits must be (B, L, K), got {tuple(logits.shape)}")
        bsz, seq_len, vocab_size = logits.shape

        expected = (bsz, seq_len)
        if target_events.shape != expected:
            raise ValueError(f"target_events must be {expected}, got {tuple(target_events.shape)}")
        if target_times.shape != expected:
            raise ValueError(f"target_times must be {expected}, got {tuple(target_times.shape)}")
        if current_times.shape != expected:
            raise ValueError(f"current_times must be {expected}, got {tuple(current_times.shape)}")
        if padding_mask.shape != expected:
            raise ValueError(f"padding_mask must be {expected}, got {tuple(padding_mask.shape)}")

        valid_mask = padding_mask.bool()
        for idx in self.ignored_tokens:
            valid_mask = valid_mask & (target_events != int(idx))
        valid_mask = valid_mask & (target_events > PAD_IDX)

        if not valid_mask.any():
            total_loss = _zero_loss_like(logits)
            if return_components:
                return total_loss, {
                    "ce": total_loss.detach(),
                    "time": total_loss.detach(),
                    "total": total_loss.detach(),
                }
            return total_loss

        logits_valid = logits[valid_mask]
        target_events_valid = target_events[valid_mask]
        target_times_valid = target_times[valid_mask]
        current_times_valid = current_times[valid_mask]

        logits_safe = torch.nan_to_num(
            logits_valid,
            nan=0.0,
            posinf=self.max_exp_input,
            neginf=-self.max_exp_input,
        )

        loss_ce = F.cross_entropy(
            logits_safe,
            target_events_valid,
            reduction="mean",
        )

        logits_for_lse = logits_safe
        if self.exclude_ignored_from_intensity:
            ignore_mask = _make_ignore_mask(vocab_size, self.ignored_tokens, logits.device)
            logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf"))

        log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1)
        log_lambda_total = -torch.log(torch.exp(-log_lambda_total) + self.t_min)

        dt = torch.clamp(target_times_valid - current_times_valid, min=self.t_min)
        log_dt_inv = -torch.log(dt + self.t_min)
        loss_dt = -(
            log_lambda_total
            - torch.exp(
                torch.clamp(log_lambda_total - log_dt_inv, max=self.max_exp_input)
            )
        )
        loss_dt = loss_dt.mean()

        total_loss = self.ce_weight * loss_ce + self.time_weight * loss_dt
        if return_components:
            return total_loss, {
                "ce": loss_ce.detach(),
                "time": loss_dt.detach(),
                "total": total_loss.detach(),
            }
        return total_loss


class ExponentialLoss(nn.Module):
    """First-onset all-future exponential survival likelihood."""

    def __init__(
        self,
        ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX),
        eps: float = 1e-8,
    ):
        super().__init__()
        self.ignored_idx = tuple(int(i) for i in ignored_idx)
        self.eps = eps

    def forward(
        self,
        logits: torch.Tensor,
        targets: torch.Tensor,
        exposure: torch.Tensor,
        dt: torch.Tensor,
        history: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch_size, vocab_size = logits.shape
        if targets.dim() != 2 or targets.shape[0] != batch_size:
            raise ValueError(
                "targets must be (B, M). "
                f"Got logits={tuple(logits.shape)}, targets={tuple(targets.shape)}"
            )
        if exposure.shape != (batch_size,):
            raise ValueError(
                f"exposure must be ({batch_size},), got {tuple(exposure.shape)}"
            )
        if dt.shape != targets.shape:
            raise ValueError(
                f"dt must match targets, got dt={tuple(dt.shape)}, "
                f"targets={tuple(targets.shape)}"
            )

        rate = F.softplus(logits) + self.eps
        at_risk = _all_future_at_risk_mask(
            vocab_size,
            self.ignored_idx,
            logits,
            history,
        )
        target_valid, safe_targets = _all_future_target_mask(
            targets,
            at_risk,
            self.ignored_idx,
        )

        censor_time = exposure.to(rate.dtype).clamp_min(self.eps)
        cumulative_at_censor = rate * censor_time.unsqueeze(1)
        penalty = (cumulative_at_censor * at_risk.to(rate.dtype)).sum(dim=-1)

        # A first-onset outcome leaves the risk set at its event time, not at
        # the common end of follow-up.
        if targets.numel() > 0:
            event_time = dt.to(rate.dtype).clamp_min(self.eps)
            event_time = torch.minimum(event_time, censor_time.unsqueeze(1))
            event_cumulative = rate.gather(1, safe_targets) * event_time
            censor_cumulative = cumulative_at_censor.gather(1, safe_targets)
            penalty = penalty + (
                (event_cumulative - censor_cumulative)
                * target_valid.to(rate.dtype)
            ).sum(dim=-1)

        observed = rate.log().gather(1, safe_targets)
        observed = (observed * target_valid.to(rate.dtype)).sum(dim=-1)
        return (-observed + penalty).mean()


class WeibullLoss(nn.Module):
    """First-onset all-future Weibull survival likelihood."""

    def __init__(
        self,
        ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX),
        eps: float = 1e-8,
    ):
        super().__init__()
        self.ignored_idx = tuple(int(i) for i in ignored_idx)
        self.eps = eps

    def forward(
        self,
        logits: torch.Tensor,
        weibull_rho: torch.Tensor,
        targets: torch.Tensor,
        dt: torch.Tensor,
        exposure: torch.Tensor,
        history: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch_size, vocab_size = logits.shape
        if weibull_rho is None:
            raise ValueError("weibull_rho is required for WeibullLoss")
        if weibull_rho.shape != logits.shape:
            raise ValueError(
                "weibull_rho must have the same shape as logits. "
                f"Got logits={tuple(logits.shape)}, weibull_rho={tuple(weibull_rho.shape)}"
            )
        if targets.dim() != 2 or targets.shape[0] != batch_size:
            raise ValueError(
                "targets must be (B, M). "
                f"Got logits={tuple(logits.shape)}, targets={tuple(targets.shape)}"
            )
        if dt.shape != targets.shape:
            raise ValueError(
                f"dt must match targets, got dt={tuple(dt.shape)}, "
                f"targets={tuple(targets.shape)}"
            )
        if exposure.shape != (batch_size,):
            raise ValueError(
                f"exposure must be ({batch_size},), got {tuple(exposure.shape)}"
            )

        dtype = logits.dtype
        rate = F.softplus(logits) + self.eps
        rho = weibull_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps)
        at_risk = _all_future_at_risk_mask(
            vocab_size,
            self.ignored_idx,
            logits,
            history,
        )
        target_valid, safe_targets = _all_future_target_mask(
            targets,
            at_risk,
            self.ignored_idx,
        )

        censor_time = exposure.to(dtype).clamp_min(self.eps)
        cumulative_at_censor = rate * torch.pow(censor_time.unsqueeze(1), rho)
        penalty = (cumulative_at_censor * at_risk.to(dtype)).sum(dim=-1)

        target_rate = rate.gather(1, safe_targets)
        target_rho = rho.gather(1, safe_targets)
        target_dt = dt.to(dtype).clamp_min(self.eps)
        target_dt = torch.minimum(target_dt, censor_time.unsqueeze(1))
        event_cumulative = target_rate * torch.pow(target_dt, target_rho)
        censor_cumulative = cumulative_at_censor.gather(1, safe_targets)
        penalty = penalty + (
            (event_cumulative - censor_cumulative) * target_valid.to(dtype)
        ).sum(dim=-1)

        log_intensity = (
            target_rate.log()
            + target_rho.log()
            + (target_rho - 1.0) * target_dt.log()
        )
        observed = (log_intensity * target_valid.to(dtype)).sum(dim=-1)
        return (-observed + penalty).mean()


def build_loss(name: str, **kwargs) -> nn.Module:
    name = name.lower()
    if name == "delphi2m":
        return Delphi2MLoss(**kwargs)
    if name == "exponential":
        return ExponentialLoss(**kwargs)
    if name == "weibull":
        return WeibullLoss(**kwargs)
    raise ValueError(
        f"Unknown loss {name!r}. Available: delphi2m, exponential, weibull."
    )