diff --git a/README.md b/README.md index 9392ffa..d583422 100644 --- a/README.md +++ b/README.md @@ -218,6 +218,24 @@ We obtain the concordance index for this batch with: ```python >>> from torchsurv.metrics.cindex import ConcordanceIndex >>> cindex = ConcordanceIndex() + +### Competing risks with cause-specific Cox + +For competing risks, the model outputs one log relative hazard per cause. Event labels are integer-coded with `0` for censoring and `1..K` for the observed cause. + +```python +>>> from torch import nn +>>> from torchsurv.loss import competing_risks +>>> n_causes = 2 +>>> event_cr = torch.randint(low=0, high=n_causes + 1, size=(n,), dtype=torch.long) +>>> model_cr = nn.Sequential(nn.Linear(16, n_causes)) +>>> log_hz_cr = model_cr(x) +>>> loss = competing_risks.neg_partial_log_likelihood(log_hz_cr, event_cr, time) +>>> baseline = competing_risks.baseline_cumulative_incidence_function(log_hz_cr.detach(), event_cr, time) +>>> cif = competing_risks.cumulative_incidence_function(baseline, log_hz_cr.detach(), torch.tensor([25.0, 50.0])) +>>> print(cif.shape) +torch.Size([64, 2, 2]) +``` >>> print(cindex(log_hz, event, time)) tensor(0.4062) ``` diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index bb8feef..2901439 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -1,6 +1,11 @@ Changelog ========= +Version 0.2.1 +------------- + +* Added a first competing-risks implementation based on cause-specific Cox models, including CIF and event-free survival helpers. + Version 0.1.6 ------------- diff --git a/docs/loss.rst b/docs/loss.rst index bbe5a3c..bb8596a 100644 --- a/docs/loss.rst +++ b/docs/loss.rst @@ -6,6 +6,7 @@ Loss :template: autosummary/module.rst torchsurv.loss.cox + torchsurv.loss.competing_risks torchsurv.loss.weibull torchsurv.loss.survival torchsurv.loss.momentum diff --git a/docs/notebooks/loss.md b/docs/notebooks/loss.md index 288434b..338f224 100644 --- a/docs/notebooks/loss.md +++ b/docs/notebooks/loss.md @@ -158,6 +158,38 @@ This model is particularly powerful when the true hazard does not follow a stand - Observations are conditionally independent given the covariates. - Numerical integration (trapezoidal rule) is sufficiently accurate given the time discretization. +### 5. Competing Risks with Cause-Specific Cox + +The first competing-risks loss is available through: + +```python +from torchsurv.loss.competing_risks import neg_partial_log_likelihood +``` + +In this setting, the event variable is integer-coded with `0` for censoring and `1, \ldots, K` for the observed cause. The model outputs one log relative hazard per cause, + +$$ +\log \lambda_{ik} = f_{\theta, k}(\mathbf{x}_i), \quad k \in \{1, \ldots, K\}. +$$ + +For each cause $k$, `TorchSurv` fits a cause-specific Cox model by treating cause $k$ as the event of interest and all other outcomes as censored at their observed times. The total loss is the sum of the per-cause Cox partial log-likelihoods: + +$$ +\text{npll}_{CR} = \sum_{k=1}^{K} \text{npll}_{k}. +$$ + +This parameterization supports neural networks with a final layer of width $K$, keeps the same Cox-style training loop, and enables prediction of: + +- cause-specific baseline hazard increments, +- event-free survival, +- cumulative incidence functions (CIFs) for each cause. + +**Assumptions.** +- Cause-specific proportional hazards within each cause. +- Independent right censoring. +- Correct specification of the cause-specific log-risk functions. +- Competing events are handled through separate cause-specific hazards rather than a direct subdistribution hazard model. + ### FAQ: Choosing the Right Survival Model @@ -214,3 +246,4 @@ Use the **Flexible Survival model** when you do not want to impose any parametri | **Weibull** | $h_i(t) = \frac{\exp(f_{\theta_1}(\mathbf{x}_i))}{\exp(f_{\theta_2}(\mathbf{x}_i))} \left(\frac{t}{\exp(f_{\theta_2}(\mathbf{x}_i))}\right)^{\exp(f_{\theta_1}(\mathbf{x}_i)) - 1}$ | ✗ | ✓ | You expect monotonic hazard shape | | **Exponential** | $h_i(t) = \frac{1}{\exp(f_\theta(\mathbf{x}_i))}$ | ✗ | ✓ | You expect constant risk over time | | **Flexible Survival** | $h_i(t) = \exp(f_{\theta}(\mathbf{x}_i, t))$ | ✓ | ✗ (numerical approximation) | You need full flexibility, no parametric form | +| **Competing Risks (Cause-Specific Cox)** | $h_{ik}(t) = \lambda_{0k}(t)\exp(f_{\theta,k}(\mathbf{x}_i))$ | ✗ | ✓ (per-cause partial likelihood) | You need CIFs with multiple mutually exclusive event types | diff --git a/docs/package_overview.md b/docs/package_overview.md index ecbd4ca..080d0c0 100644 --- a/docs/package_overview.md +++ b/docs/package_overview.md @@ -10,6 +10,7 @@ graph LR %% LOSS LOSS --> COX["Cox"]:::sub + LOSS --> CR["Competing\nRisks"]:::sub LOSS --> WEIBULL["Weibull"]:::sub LOSS --> SURVIVAL["Survival\n(discrete-time)"]:::sub LOSS --> MOMENTUM["Momentum"]:::sub @@ -18,6 +19,10 @@ graph LR COX --> C2["baseline_survival_function"]:::fn COX --> C3["survival_function"]:::fn + CR --> CR1["neg_partial_log_likelihood"]:::fn + CR --> CR2["baseline_cumulative_incidence_function"]:::fn + CR --> CR3["cumulative_incidence_function\n· survival_function"]:::fn + WEIBULL --> W1["neg_log_likelihood"]:::fn WEIBULL --> W2["log_hazard"]:::fn WEIBULL --> W3["survival_function"]:::fn diff --git a/src/torchsurv/loss/__init__.py b/src/torchsurv/loss/__init__.py index 73f2e41..e994aa1 100644 --- a/src/torchsurv/loss/__init__.py +++ b/src/torchsurv/loss/__init__.py @@ -2,6 +2,7 @@ from __future__ import annotations +from torchsurv.loss import competing_risks from torchsurv.loss.cox import baseline_survival_function, neg_partial_log_likelihood, survival_function_cox from torchsurv.loss.momentum import Momentum from torchsurv.loss.survival import neg_log_likelihood, survival_function @@ -9,6 +10,7 @@ __all__ = [ "baseline_survival_function", + "competing_risks", "log_hazard", "Momentum", "neg_log_likelihood", diff --git a/src/torchsurv/loss/competing_risks.py b/src/torchsurv/loss/competing_risks.py new file mode 100644 index 0000000..430204e --- /dev/null +++ b/src/torchsurv/loss/competing_risks.py @@ -0,0 +1,415 @@ +from __future__ import annotations + +import sys +import warnings + +import torch + +from torchsurv.loss.cox import neg_partial_log_likelihood as cox_neg_partial_log_likelihood +from torchsurv.tools.validators import CompetingRisksInputs, CompetingRisksModelInputs + +__all__ = [ + "baseline_cumulative_incidence_function", + "cumulative_incidence_function", + "neg_partial_log_likelihood", + "survival_function", +] + + +def _searchsorted(sorted_seq: torch.Tensor, values: torch.Tensor, right: bool = False) -> torch.Tensor: + """torch.searchsorted with CPU fallback for devices that don't support it.""" + return torch.searchsorted(sorted_seq.cpu(), values.cpu(), right=right).to(sorted_seq.device) + + +def _validate_new_prediction_inputs( + new_log_hz: torch.Tensor, new_time: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Coerce prediction inputs and validate simple shape constraints.""" + if not isinstance(new_log_hz, torch.Tensor): + raise ValueError("Input 'new_log_hz' must be a torch.Tensor.") + if not isinstance(new_time, torch.Tensor): + raise ValueError("Input 'new_time' must be a torch.Tensor.") + + new_log_hz = new_log_hz.float().to(new_log_hz.device).squeeze() + new_time = new_time.float().to(new_time.device).squeeze() + + if new_log_hz.dim() == 1: + new_log_hz = new_log_hz.unsqueeze(0) + if new_log_hz.dim() != 2: + raise ValueError("Input 'new_log_hz' must have shape (n_samples_new, n_causes).") + if new_log_hz.shape[1] == 0: + raise ValueError("Input 'new_log_hz' must contain at least one cause column.") + if new_time.dim() == 0: + new_time = new_time.unsqueeze(0) + if new_time.dim() != 1: + raise ValueError("Input 'new_time' must be one-dimensional.") + if torch.any(torch.isnan(new_time)) or torch.any(torch.isinf(new_time)): + raise ValueError("Input 'new_time' contains NaN or Inf values, which are not allowed.") + if torch.any(new_time < 0): + raise ValueError("Input 'new_time' must be non-negative.") + if torch.any(new_time[:-1] > new_time[1:]): + raise ValueError("Input 'new_time' must be sorted from smallest to largest.") + if len(new_time) != len(torch.unique(new_time)): + raise ValueError("Input 'new_time' must contain unique values.") + return new_log_hz, new_time + + +def _compute_baseline_curves(log_hz: torch.Tensor, event: torch.Tensor, time: torch.Tensor) -> dict[str, torch.Tensor]: + """Compute baseline cause-specific hazards, survival, and CIF on a stratum.""" + time_sorted, idx = torch.sort(time) + event_sorted = event[idx] + log_hz_sorted = log_hz[idx] + + time_unique = torch.unique(time_sorted) + n_times = len(time_unique) + n_causes = log_hz.shape[1] + + first_idx = _searchsorted(time_sorted, time_unique) + baseline_hazard = torch.zeros((n_times, n_causes), dtype=log_hz.dtype, device=log_hz.device) + + for cause_idx in range(n_causes): + cause = cause_idx + 1 + exp_hz = torch.exp(log_hz_sorted[:, cause_idx]) + reverse_cumsum = exp_hz.flip(0).cumsum(0).flip(0) + denominator = reverse_cumsum[first_idx] + + event_idx = _searchsorted(time_unique, time_sorted[event_sorted == cause]) + event_count = torch.zeros(n_times, dtype=log_hz.dtype, device=log_hz.device) + if len(event_idx) > 0: + event_count.scatter_add_(0, event_idx, torch.ones_like(event_idx, dtype=log_hz.dtype)) + + baseline_hazard[:, cause_idx] = torch.where( + denominator > 0, + event_count / denominator, + torch.zeros_like(event_count), + ) + + baseline_cumulative_hazard = torch.cumsum(baseline_hazard, dim=0) + baseline_survival = torch.ones(n_times, dtype=log_hz.dtype, device=log_hz.device) + baseline_cif = torch.zeros((n_times, n_causes), dtype=log_hz.dtype, device=log_hz.device) + + survival_prev = torch.tensor(1.0, dtype=log_hz.dtype, device=log_hz.device) + cif_prev = torch.zeros(n_causes, dtype=log_hz.dtype, device=log_hz.device) + for t_idx in range(n_times): + increments = baseline_hazard[t_idx] + total_increment = increments.sum() + if total_increment > 0: + delta = 1 - torch.exp(-total_increment) + cif_prev = cif_prev + survival_prev * delta * increments / total_increment + survival_prev = survival_prev * torch.exp(-total_increment) + baseline_survival[t_idx] = survival_prev + baseline_cif[t_idx] = cif_prev + + return { + "time": time_unique, + "baseline_hazard": baseline_hazard, + "baseline_cumulative_hazard": baseline_cumulative_hazard, + "baseline_survival": baseline_survival, + "baseline_cif": baseline_cif, + } + + +def neg_partial_log_likelihood( + log_hz: torch.Tensor, + event: torch.Tensor, + time: torch.Tensor, + ties_method: str = "efron", + reduction: str = "mean", + strata: torch.Tensor | None = None, +) -> torch.Tensor: + r"""Negative partial log-likelihood for a cause-specific Cox competing-risks model. + + Args: + log_hz (torch.Tensor, float): + Cause-specific log relative hazards of shape ``(n_samples, n_causes)``. + event (torch.Tensor, int): + Integer event indicator of shape ``(n_samples,)`` with ``0`` for censoring + and ``1..K`` for the observed cause. + time (torch.Tensor, float): + Event or censoring time of shape ``(n_samples,)``. + ties_method (str): + Method to handle ties in event time. Defaults to ``"efron"``. + reduction (str, optional): + Method to reduce losses. Defaults to ``"mean"``. + Must be one of ``"sum"`` or ``"mean"``. + strata (torch.Tensor, int, optional): + Integer tensor of shape ``(n_samples,)`` representing strata. + + Returns: + torch.Tensor: Negative cause-specific partial log-likelihood. + + Note: + The loss is defined as the sum of binary Cox partial log-likelihoods, one + per cause. For cause :math:`k`, subjects with ``event == k`` are treated as + events and all other subjects are treated as censored. + + Examples: + >>> log_hz = torch.tensor([[0.1, -0.2], [0.3, 0.1], [-0.4, 0.5], [0.0, -0.1]]) + >>> event = torch.tensor([1, 2, 0, 1]) + >>> time = torch.tensor([1.0, 2.0, 3.0, 4.0]) + >>> neg_partial_log_likelihood(log_hz, event, time) + tensor(0.8381) + """ + + if strata is None: + strata = torch.ones_like(event, dtype=torch.long) + + log_hz = log_hz.squeeze() + event = event.squeeze() + time = time.squeeze() + strata = strata.squeeze() + + if not (torch.jit.is_scripting() or torch.jit.is_tracing()): + _surv = CompetingRisksInputs(event=event, time=time, strata=strata) + event, time = _surv.event, _surv.time + strata = _surv.strata + _model = CompetingRisksModelInputs(log_hz=log_hz, event=event) + log_hz = _model.log_hz + + if torch.count_nonzero(event).item() == 0: + warnings.warn( + "No observed causes in the batch. Returning zero loss for the batch", + stacklevel=2, + ) + return torch.tensor(0.0, requires_grad=True, device=log_hz.device, dtype=log_hz.dtype) + + n_causes = log_hz.shape[1] + total_loss = torch.tensor(0.0, dtype=log_hz.dtype, device=log_hz.device) + total_events = 0 + + assert strata is not None # for mypy + for cause_idx in range(n_causes): + cause = cause_idx + 1 + cause_event = event == cause + n_events = int(cause_event.sum().item()) + if n_events == 0: + continue + total_loss = total_loss + cox_neg_partial_log_likelihood( + log_hz[:, cause_idx], + cause_event, + time, + ties_method=ties_method, + reduction="sum", + strata=strata, + ) + total_events += n_events + + if reduction.lower() == "sum": + return total_loss + if reduction.lower() == "mean": + return total_loss / total_events + raise ValueError(f"Reduction {reduction} is not implemented yet, should be one of ['mean', 'sum'].") + + +def baseline_cumulative_incidence_function( + log_hz: torch.Tensor, + event: torch.Tensor, + time: torch.Tensor, + strata: torch.Tensor | None = None, +) -> dict[str, torch.Tensor] | dict[int, dict[str, torch.Tensor]]: + r"""Estimate baseline cause-specific hazards, survival, and CIF curves. + + Args: + log_hz (torch.Tensor, float): + Cause-specific log relative hazards of shape ``(n_samples, n_causes)``. + event (torch.Tensor, int): + Integer event indicator of shape ``(n_samples,)`` with ``0`` for censoring + and ``1..K`` for the observed cause. + time (torch.Tensor, float): + Event or censoring time of shape ``(n_samples,)``. + strata (torch.Tensor, int, optional): + Integer tensor of shape ``(n_samples,)`` representing strata. + + Returns: + dict: + Baseline competing-risks curves. For a single stratum, the dictionary + contains ``time``, ``baseline_hazard``, ``baseline_cumulative_hazard``, + ``baseline_survival``, and ``baseline_cif``. With multiple strata, the + result is keyed by the integer stratum values. + + Examples: + >>> log_hz = torch.zeros((4, 2)) + >>> event = torch.tensor([1, 2, 0, 1]) + >>> time = torch.tensor([1.0, 2.0, 3.0, 4.0]) + >>> baseline = baseline_cumulative_incidence_function(log_hz, event, time) + >>> baseline["baseline_cif"].shape + torch.Size([4, 2]) + """ + + if strata is None: + strata = torch.ones_like(event, dtype=torch.long) + + log_hz = log_hz.squeeze() + event = event.squeeze() + time = time.squeeze() + strata = strata.squeeze() + + if not (torch.jit.is_scripting() or torch.jit.is_tracing()): + _surv = CompetingRisksInputs(event=event, time=time, strata=strata) + event, time = _surv.event, _surv.time + strata = _surv.strata + _model = CompetingRisksModelInputs(log_hz=log_hz, event=event) + log_hz = _model.log_hz + + time_sorted, idx = torch.sort(time) + event_sorted = event[idx] + log_hz_sorted = log_hz[idx] + assert strata is not None # for mypy + strata_sorted = strata[idx] + + strata_unique = torch.unique(strata_sorted) + if len(strata_unique) == 1: + mask = strata_sorted == strata_unique[0] + return _compute_baseline_curves( + log_hz_sorted[mask], + event_sorted[mask], + time_sorted[mask], + ) + + strata_results: dict[int, dict[str, torch.Tensor]] = {} + for stratum in strata_unique: + mask = strata_sorted == stratum + strata_results[int(stratum.item())] = _compute_baseline_curves( + log_hz_sorted[mask], + event_sorted[mask], + time_sorted[mask], + ) + + return strata_results + + +def cumulative_incidence_function( + baseline: dict[str, torch.Tensor] | dict[int, dict[str, torch.Tensor]], + new_log_hz: torch.Tensor, + new_time: torch.Tensor, + new_strata: torch.Tensor | None = None, +) -> torch.Tensor: + r"""Compute subject-specific cumulative incidence functions. + + Args: + baseline (dict): + Output of :func:`baseline_cumulative_incidence_function`. + new_log_hz (torch.Tensor, float): + Cause-specific log relative hazards for new subjects of shape + ``(n_samples_new, n_causes)``. + new_time (torch.Tensor, float): + Times at which to evaluate the CIF of shape ``(n_times,)``. + new_strata (torch.Tensor, int, optional): + Integer tensor of shape ``(n_samples_new,)`` representing strata. + + Returns: + torch.Tensor: + Subject-specific CIF values of shape ``(n_samples_new, n_times, n_causes)``. + + Examples: + >>> baseline = baseline_cumulative_incidence_function( + ... torch.zeros((4, 2)), + ... torch.tensor([1, 2, 0, 1]), + ... torch.tensor([1.0, 2.0, 3.0, 4.0]), + ... ) + >>> cumulative_incidence_function(baseline, torch.tensor([[0.0, 0.0]]), torch.tensor([1.0, 4.0])).shape + torch.Size([1, 2, 2]) + """ + + new_log_hz, new_time = _validate_new_prediction_inputs(new_log_hz, new_time) + + if new_strata is None: + new_strata = torch.ones(len(new_log_hz), device=new_log_hz.device, dtype=torch.long) + elif not isinstance(new_strata, torch.Tensor): + raise ValueError("Input 'new_strata' must be a torch.Tensor or None.") + else: + new_strata = new_strata.long().to(new_log_hz.device).squeeze() + + if new_strata.dim() == 0: + new_strata = new_strata.unsqueeze(0) + if len(new_strata) != len(new_log_hz): + raise ValueError( + f"Dimension mismatch: 'new_log_hz' has {len(new_log_hz)} samples but 'new_strata' has {len(new_strata)}." + ) + + n_samples = len(new_log_hz) + n_times = len(new_time) + n_causes = new_log_hz.shape[1] + cif = torch.empty((n_samples, n_times, n_causes), dtype=new_log_hz.dtype, device=new_log_hz.device) + + for stratum in torch.unique(new_strata): + mask = new_strata == stratum + new_log_hz_stratum = new_log_hz[mask] + + if isinstance(baseline, dict) and all(isinstance(v, dict) for v in baseline.values()): + baseline_stratum = baseline[int(stratum.item())] # type: ignore[index] + else: + baseline_stratum = baseline + + baseline_time = baseline_stratum["time"] + baseline_hazard = baseline_stratum["baseline_hazard"] + if baseline_hazard.shape[1] != n_causes: + raise ValueError( + f"Input 'new_log_hz' has {n_causes} causes but baseline was fitted with {baseline_hazard.shape[1]} causes." + ) + + scales = torch.exp(new_log_hz_stratum) + cif_path = torch.zeros( + (len(new_log_hz_stratum), len(baseline_time), n_causes), + dtype=new_log_hz.dtype, + device=new_log_hz.device, + ) + survival_prev = torch.ones(len(new_log_hz_stratum), dtype=new_log_hz.dtype, device=new_log_hz.device) + cif_prev = torch.zeros((len(new_log_hz_stratum), n_causes), dtype=new_log_hz.dtype, device=new_log_hz.device) + + for t_idx in range(len(baseline_time)): + increments = baseline_hazard[t_idx].unsqueeze(0) * scales + total_increment = increments.sum(dim=1) + has_increment = total_increment > 0 + delta = 1 - torch.exp(-total_increment) + ratio = torch.zeros_like(increments) + if has_increment.any(): + ratio[has_increment] = increments[has_increment] / total_increment[has_increment].unsqueeze(1) + cif_prev = cif_prev + survival_prev.unsqueeze(1) * delta.unsqueeze(1) * ratio + survival_prev = survival_prev * torch.exp(-total_increment) + cif_path[:, t_idx, :] = cif_prev + + time_index = _searchsorted(baseline_time, new_time, right=True) - torch.tensor(1, device=baseline_time.device) + time_index = time_index.clamp(min=0) + cif[mask] = cif_path[:, time_index, :] + + return cif + + +def survival_function( + baseline: dict[str, torch.Tensor] | dict[int, dict[str, torch.Tensor]], + new_log_hz: torch.Tensor, + new_time: torch.Tensor, + new_strata: torch.Tensor | None = None, +) -> torch.Tensor: + r"""Compute event-free survival under a cause-specific Cox competing-risks model. + + Args: + baseline (dict): + Output of :func:`baseline_cumulative_incidence_function`. + new_log_hz (torch.Tensor, float): + Cause-specific log relative hazards for new subjects of shape + ``(n_samples_new, n_causes)``. + new_time (torch.Tensor, float): + Times at which to evaluate the event-free survival of shape ``(n_times,)``. + new_strata (torch.Tensor, int, optional): + Integer tensor of shape ``(n_samples_new,)`` representing strata. + + Returns: + torch.Tensor: + Event-free survival probabilities of shape ``(n_samples_new, n_times)``. + """ + + cif = cumulative_incidence_function(baseline, new_log_hz, new_time, new_strata=new_strata) + return 1 - cif.sum(dim=2) + + +if __name__ == "__main__": + import doctest + + results = doctest.testmod() + if results.failed == 0: + print("All tests passed.") + else: + print("Some doctests failed.") + sys.exit(1) diff --git a/src/torchsurv/tools/validators.py b/src/torchsurv/tools/validators.py index 0693ca4..2e0a047 100644 --- a/src/torchsurv/tools/validators.py +++ b/src/torchsurv/tools/validators.py @@ -6,6 +6,8 @@ from pydantic import BaseModel, ConfigDict, field_validator, model_validator __all__ = [ + "CompetingRisksInputs", + "CompetingRisksModelInputs", "EvalTimeInputs", "ModelInputs", "NewTimeInputs", @@ -158,6 +160,60 @@ def check_dimensions(self) -> SurvivalInputs: return self +class CompetingRisksInputs(_TorchModel): + """Validates and coerces competing-risks survival inputs. + + ``event`` is encoded as integers with ``0`` for censoring and ``1..K`` for + observed causes. Unlike :class:`SurvivalInputs`, fully-censored data is + allowed so batch-level helpers can decide how to handle it. + """ + + event: torch.Tensor + time: torch.Tensor + strata: torch.Tensor | None = None + + @field_validator("event", mode="before") + @classmethod + def coerce_event(cls, v: object) -> torch.Tensor: + if not isinstance(v, torch.Tensor): + raise ValueError("Input 'event' must be a torch.Tensor.") + if torch.is_floating_point(v): + if not torch.allclose(v, v.round()): + raise ValueError("Input 'event' must contain integer-coded causes.") + coerced = v.long().to(v.device) + if torch.any(coerced < 0): + raise ValueError("Input 'event' must contain non-negative cause labels.") + return coerced + + @field_validator("time", mode="before") + @classmethod + def coerce_time(cls, v: object) -> torch.Tensor: + coerced = _to_float_tensor("time", v) + if torch.any(coerced < 0.0): + raise ValueError("Input 'time' must be non-negative.") + return coerced + + @field_validator("strata", mode="before") + @classmethod + def coerce_strata(cls, v: object) -> torch.Tensor | None: + if v is None: + return None + if not isinstance(v, torch.Tensor): + raise ValueError("Input 'strata' must be a torch.Tensor or None.") + return v.long().to(v.device) + + @model_validator(mode="after") + def check_dimensions(self) -> CompetingRisksInputs: + n = len(self.event) + if len(self.time) != n: + raise ValueError(f"Dimension mismatch: 'event' has {n} samples but 'time' has {len(self.time)}.") + if self.strata is not None and len(self.strata) != n: + raise ValueError(f"Dimension mismatch: 'event' has {n} samples but 'strata' has {len(self.strata)}.") + if self.strata is None: + self.strata = torch.ones_like(self.event, dtype=torch.long) + return self + + class ModelInputs(_TorchModel): """Validates and coerces model parameter inputs. @@ -250,6 +306,41 @@ def check_shape(self) -> ModelInputs: return self +class CompetingRisksModelInputs(_TorchModel): + """Validates multi-cause log-hazard tensors against competing-risks labels.""" + + log_hz: torch.Tensor + event: torch.Tensor + + @field_validator("log_hz", mode="before") + @classmethod + def coerce_log_hz(cls, v: object) -> torch.Tensor: + return _to_float_tensor("log_hz", v, allow_inf=True) + + @field_validator("event", mode="before") + @classmethod + def coerce_event(cls, v: object) -> torch.Tensor: + if not isinstance(v, torch.Tensor): + raise ValueError("Input 'event' must be a torch.Tensor.") + return v.long().to(v.device) + + @model_validator(mode="after") + def check_shape(self) -> CompetingRisksModelInputs: + n = len(self.event) + if self.log_hz.dim() != 2: + raise ValueError("For competing risks, 'log_hz' must have shape (n_samples, n_causes).") + if self.log_hz.shape[0] != n: + raise ValueError(f"Dimension mismatch: 'log_hz' has {self.log_hz.shape[0]} samples but 'event' has {n}.") + if self.log_hz.shape[1] == 0: + raise ValueError("For competing risks, 'log_hz' must contain at least one cause column.") + if len(self.event) > 0 and int(self.event.max().item()) > self.log_hz.shape[1]: + raise ValueError( + f"Input 'event' contains cause label {int(self.event.max().item())}, " + f"but 'log_hz' only has {self.log_hz.shape[1]} cause columns." + ) + return self + + class NewTimeInputs(_TorchModel): """Validates new evaluation times. diff --git a/tests/test_competing_risks.py b/tests/test_competing_risks.py new file mode 100644 index 0000000..6da25d0 --- /dev/null +++ b/tests/test_competing_risks.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import pytest +import torch + +from torchsurv.loss import competing_risks +from torchsurv.loss.cox import neg_partial_log_likelihood as cox + + +def test_competing_risks_loss_matches_sum_of_binary_cox_losses() -> None: + log_hz = torch.tensor( + [ + [0.1, -0.2], + [0.3, 0.4], + [-0.5, 0.2], + [0.2, -0.1], + [0.0, 0.5], + ], + dtype=torch.float32, + ) + event = torch.tensor([1, 2, 0, 1, 2], dtype=torch.long) + time = torch.tensor([1.0, 2.0, 3.0, 4.0, 4.0], dtype=torch.float32) + + observed = competing_risks.neg_partial_log_likelihood(log_hz, event, time, ties_method="efron", reduction="sum") + expected = cox(log_hz[:, 0], event == 1, time, ties_method="efron", reduction="sum") + cox( + log_hz[:, 1], event == 2, time, ties_method="efron", reduction="sum" + ) + + assert torch.allclose(observed, expected) + + +def test_competing_risks_loss_matches_stratified_binary_cox_losses() -> None: + log_hz = torch.tensor( + [ + [0.1, -0.2], + [0.3, 0.4], + [-0.5, 0.2], + [0.2, -0.1], + [0.0, 0.5], + [0.7, -0.4], + ], + dtype=torch.float32, + ) + event = torch.tensor([1, 2, 0, 1, 2, 1], dtype=torch.long) + time = torch.tensor([1.0, 2.0, 3.0, 4.0, 4.0, 2.0], dtype=torch.float32) + strata = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.long) + + observed = competing_risks.neg_partial_log_likelihood( + log_hz, + event, + time, + ties_method="breslow", + reduction="sum", + strata=strata, + ) + expected = cox(log_hz[:, 0], event == 1, time, ties_method="breslow", reduction="sum", strata=strata) + cox( + log_hz[:, 1], event == 2, time, ties_method="breslow", reduction="sum", strata=strata + ) + + assert torch.allclose(observed, expected) + + +def test_all_censored_batch_returns_zero_loss() -> None: + log_hz = torch.randn(5, 2) + event = torch.zeros(5, dtype=torch.long) + time = torch.arange(1, 6, dtype=torch.float32) + + loss = competing_risks.neg_partial_log_likelihood(log_hz, event, time) + + assert loss.item() == 0.0 + + +@pytest.mark.parametrize( + ("log_hz", "event", "match"), + [ + pytest.param(torch.randn(4), torch.tensor([1, 0, 1, 0]), "shape", id="log_hz_not_2d"), + pytest.param(torch.randn(4, 2), torch.tensor([1, 3, 0, 1]), "cause label", id="cause_exceeds_heads"), + pytest.param(torch.randn(4, 2), torch.tensor([1, -1, 0, 1]), "non-negative", id="negative_cause_label"), + ], +) +def test_invalid_competing_risks_inputs_raise(log_hz: torch.Tensor, event: torch.Tensor, match: str) -> None: + time = torch.arange(1, 5, dtype=torch.float32) + with pytest.raises(ValueError, match=match): + competing_risks.neg_partial_log_likelihood(log_hz, event, time) + + +def test_baseline_curves_match_hand_computed_toy_example() -> None: + log_hz = torch.zeros((4, 2), dtype=torch.float32) + event = torch.tensor([1, 2, 0, 1], dtype=torch.long) + time = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32) + + baseline = competing_risks.baseline_cumulative_incidence_function(log_hz, event, time) + + expected_hazard = torch.tensor( + [ + [0.25, 0.0], + [0.0, 1.0 / 3.0], + [0.0, 0.0], + [1.0, 0.0], + ], + dtype=torch.float32, + ) + expected_survival = torch.tensor( + [ + torch.exp(torch.tensor(-0.25)), + torch.exp(torch.tensor(-0.25 - 1.0 / 3.0)), + torch.exp(torch.tensor(-0.25 - 1.0 / 3.0)), + torch.exp(torch.tensor(-0.25 - 1.0 / 3.0 - 1.0)), + ], + dtype=torch.float32, + ) + expected_cif = torch.tensor( + [ + [1.0 - torch.exp(torch.tensor(-0.25)), 0.0], + [ + 1.0 - torch.exp(torch.tensor(-0.25)), + torch.exp(torch.tensor(-0.25)) * (1.0 - torch.exp(torch.tensor(-1.0 / 3.0))), + ], + [ + 1.0 - torch.exp(torch.tensor(-0.25)), + torch.exp(torch.tensor(-0.25)) * (1.0 - torch.exp(torch.tensor(-1.0 / 3.0))), + ], + [ + 1.0 + - torch.exp(torch.tensor(-0.25)) + + torch.exp(torch.tensor(-0.25 - 1.0 / 3.0)) * (1.0 - torch.exp(torch.tensor(-1.0))), + torch.exp(torch.tensor(-0.25)) * (1.0 - torch.exp(torch.tensor(-1.0 / 3.0))), + ], + ], + dtype=torch.float32, + ) + + assert torch.allclose(baseline["baseline_hazard"], expected_hazard, atol=1e-6) + assert torch.allclose(baseline["baseline_survival"], expected_survival, atol=1e-6) + assert torch.allclose(baseline["baseline_cif"], expected_cif, atol=1e-6) + assert torch.allclose(baseline["baseline_survival"], 1 - baseline["baseline_cif"].sum(dim=1), atol=1e-6) + + +def test_cif_predictions_are_monotone_and_match_survival_identity() -> None: + log_hz = torch.zeros((4, 2), dtype=torch.float32) + event = torch.tensor([1, 2, 0, 1], dtype=torch.long) + time = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32) + baseline = competing_risks.baseline_cumulative_incidence_function(log_hz, event, time) + + new_log_hz = torch.tensor([[0.0, 0.0], [1.0, 0.0]], dtype=torch.float32) + new_time = torch.tensor([1.0, 2.0, 4.0], dtype=torch.float32) + + cif = competing_risks.cumulative_incidence_function(baseline, new_log_hz, new_time) + survival = competing_risks.survival_function(baseline, new_log_hz, new_time) + + assert torch.all(cif[:, 1:, :] >= cif[:, :-1, :]) + assert torch.allclose(survival, 1 - cif.sum(dim=2), atol=1e-6) + assert torch.all(cif[1, :, 0] >= cif[0, :, 0]) + assert torch.all(survival[1] <= survival[0]) + + +def test_baseline_and_prediction_support_multiple_strata() -> None: + log_hz = torch.tensor( + [ + [0.0, 0.0], + [0.1, -0.2], + [0.0, 0.0], + [0.2, -0.1], + ], + dtype=torch.float32, + ) + event = torch.tensor([1, 0, 2, 1], dtype=torch.long) + time = torch.tensor([1.0, 2.0, 1.5, 3.0], dtype=torch.float32) + strata = torch.tensor([0, 0, 1, 1], dtype=torch.long) + + baseline = competing_risks.baseline_cumulative_incidence_function(log_hz, event, time, strata=strata) + cif = competing_risks.cumulative_incidence_function( + baseline, + torch.tensor([[0.0, 0.0], [0.0, 0.0]], dtype=torch.float32), + torch.tensor([1.5, 3.0], dtype=torch.float32), + new_strata=torch.tensor([0, 1], dtype=torch.long), + ) + + assert set(baseline.keys()) == {0, 1} + assert cif.shape == (2, 2, 2)