torch-survival 0.1.0a2__tar.gz

torch_survival-0.1.0a2/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.3
+Name: torch-survival
+Version: 0.1.0a2
+Summary: Production-ready deep survival analysis with one line of code
+Requires-Dist: matplotlib>=3.10.8
+Requires-Dist: optuna>=4.5.0
+Requires-Dist: optuna-integration[sklearn]>=4.6.0
+Requires-Dist: pycox>=0.3.0
+Requires-Dist: rich>=14.2.0
+Requires-Dist: scikit-survival>=0.24.1
+Requires-Dist: torch>=2.8.0
+Requires-Dist: typing-extensions>=4.15.0
+Requires-Dist: furo>=2025.7.19 ; extra == 'docs'
+Requires-Dist: sphinx>=8.2.3 ; extra == 'docs'
+Requires-Dist: h5py>=3.14.0 ; extra == 'test'
+Requires-Dist: torchsurv>=0.1.5 ; extra == 'test'
+Requires-Python: >=3.12
+Provides-Extra: docs
+Provides-Extra: test
+Description-Content-Type: text/markdown
+

torch_survival-0.1.0a2/pyproject.toml

@@ -0,0 +1,40 @@
+[project]
+name = "torch-survival"
+version = "0.1.0-alpha.2"
+description = "Production-ready deep survival analysis with one line of code"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "matplotlib>=3.10.8",
+    "optuna>=4.5.0",
+    "optuna-integration[sklearn]>=4.6.0",
+    "pycox>=0.3.0",
+    "rich>=14.2.0",
+    "scikit-survival>=0.24.1",
+    "torch>=2.8.0",
+    "typing-extensions>=4.15.0",
+]
+
+[project.optional-dependencies]
+docs = [
+    "furo>=2025.7.19",
+    "sphinx>=8.2.3",
+]
+test = [
+    "h5py>=3.14.0",
+    "torchsurv>=0.1.5",
+]
+
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+
+[build-system]
+requires = ["uv_build>=0.8.11,<0.9.0"]
+build-backend = "uv_build"

torch_survival-0.1.0a2/src/torch_survival/config.py

@@ -0,0 +1,30 @@
+from typing_extensions import TypedDict
+
+
+class LayerConfig(TypedDict, closed=True):
+    # Maximum number of hidden layers
+    max_layers: int
+    # Maximum number of nodes per hidden layer
+    max_nodes_per_layer: int
+
+
+class NetworkConfig(TypedDict, closed=True):
+    #: Layers of the neural network
+    layers: LayerConfig | list[int]
+    #: Activation function used in hidden layers
+    activation: list[str] | str
+    #: Probability of dropout for dropout layers
+    dropout: tuple[float, float] | float
+
+
+class OptimizerConfig(TypedDict, closed=True):
+    #: Optimizer used to train neural network
+    name: list[str] | str
+    #: Learning rate used to train neural network
+    lr: tuple[float, float] | float
+    #: Learning rate momentum used to train neural network
+    momentum: tuple[float, float] | float
+    #: Scheduler applied to learning rate
+    scheduler: list[str] | str
+    #: Decay used by learning rate scheduler
+    decay: tuple[float, float] | float

torch_survival-0.1.0a2/src/torch_survival/losses.py

@@ -0,0 +1,219 @@
+import torch
+
+
+def cox_neg_log_likelihood(risk: torch.Tensor, event: torch.Tensor, time: torch.Tensor,
+                           sort: bool = True) -> torch.Tensor:
+    """ Cox negative log partial likelihood.
+    From DeepSurv: https://doi.org/10.1186/s12874-018-0482-1
+
+    Parameters
+    ----------
+    risk: torch.Tensor, of shape (..., n_samples)
+        Risks estimated by model
+
+    event: torch.Tensor, of shape (..., n_samples)
+        Event indicator denoting whether time is of observed event or dropout
+
+    time: torch.Tensor, of shape (..., n_samples)
+        Time of either observed event or dropout
+
+    sort: bool
+        Whether to sort by time, otherwise assumes pre-sorted ground truth times and events
+
+    """
+    if sort:
+        # Sort risk and events by time
+        sort_idx = torch.argsort(time, descending=True)
+        risk = torch.gather(risk, dim=-1, index=sort_idx)
+        event = torch.gather(event, dim=-1, index=sort_idx)
+    # Due to the sorting before, log_risk[i] = log(sum(e^risk[j=0:i]) with time[j] >= time[i]
+    log_risk = torch.logcumsumexp(risk, dim=-1)
+    likelihood = (risk - log_risk) * event.float()
+    return - likelihood.sum(dim=-1) / event.sum(dim=-1)
+
+
+def weibull_neg_log_likelihood(params: torch.Tensor, event: torch.Tensor, time: torch.Tensor,
+                               activations: bool = False) -> torch.Tensor:
+    """ Weibull negative log likelihood.
+    From DeepWeiSurv: https://doi.org/10.1007/978-3-030-47426-3_53
+
+    Parameters
+    ----------
+    params: torch.Tensor, of shape (..., n_samples, 2 | p * 3)
+        Parameters of the mixture of p Weibull distributions, namely weights, shapes, and scales
+
+    event: torch.Tensor, of shape (..., n_samples)
+        Event indicator denoting whether time is of observed event or dropout
+
+    time: torch.Tensor, of shape (..., n_samples)
+        Time of either observed event or dropout
+
+    activations: bool
+        Whether to apply ELU activations
+    """
+    # Extract mixture alpha, shape, and scale from params
+    if params.shape[-1] == 2 or params.shape[-1] == 3:
+        n_dists = 1
+        weight, shape, scale = None, params[..., -2], params[..., -1]
+    elif params.shape[-1] % 3 == 0:
+        n_dists = params.shape[-1] // 3
+        weight, shape, scale = params[..., 0:n_dists], params[..., n_dists:2 * n_dists], params[..., 2 * n_dists:]
+        time = time.unsqueeze(-1)  # for proper broadcasting with mixture params
+    else:
+        raise ValueError('Unexpected number of Weibull parameters: ' + str(params.shape[-1]))
+
+    # Apply ELU activations if requested
+    if activations:
+        shape = torch.nn.functional.elu(shape) + 2
+        scale = torch.nn.functional.elu(scale) + 1 + 1e-5
+
+    # Guard against time=0, where log(time) = -inf, since a * event_true = nan which is later disregarded in nansum
+    event_true = event & (time.squeeze(-1) != 0)
+    event_false = ~event
+
+    a = torch.log(shape) - torch.log(scale) + (shape - 1) * (torch.log(time) - torch.log(scale))
+    b = -torch.pow(time / scale, shape)
+    if n_dists == 1:
+        # Simplified equation for single distribution
+        return -(torch.nansum(a * event_true.float(), dim=-1) + torch.sum(b, dim=-1)) / event.shape[-1]
+    else:
+        # Extended equation for multiple distributions
+        b += weight
+        return -(torch.nansum(torch.logsumexp(a + b, dim=-1) * event_true.float(), dim=-1)
+                 + torch.sum(torch.logsumexp(b, dim=-1) * event_false.float(), dim=-1)
+                 - torch.sum(torch.logsumexp(weight, dim=-1), dim=-1)) / event.shape[-1]
+
+
+def weibull_neg_log_likelihood_original(params: torch.Tensor, event: torch.Tensor, time: torch.Tensor) -> torch.Tensor:
+    """ Weibull negative log likelihood.
+    From DeepWeiSurv: https://doi.org/10.1007/978-3-030-47426-3_53
+
+    Parameters
+    ----------
+    params: torch.Tensor, of shape (..., n_samples, 2 | p * 3)
+        Parameters of the mixture of p Weibull distributions, namely weights, shapes, and scales
+
+    event: torch.Tensor, of shape (..., n_samples)
+        Event indicator denoting whether time is of observed event or dropout
+
+    time: torch.Tensor, of shape (..., n_samples)
+        Time of either observed event or dropout
+    """
+    # Extract mixture alpha, shape, and scale from params
+    if params.shape[-1] % 3 == 0:
+        n_dists = params.shape[-1] // 3
+        alphas, shape, scale = params[..., 0:n_dists], params[..., n_dists:2 * n_dists], params[..., 2 * n_dists:]
+        alphas, shape, scale = alphas.t(), shape.t(), scale.t()
+    else:
+        raise ValueError('Unexpected number of Weibull parameters: ' + str(params.shape[-1]))
+
+    # https://github.com/AchrafB2015/pydpwte/blob/133aeb1004adf6bc0fd1e6985b42fc5986a77e02/pydpwte/utils/loss.py#L37-L44
+    t_over_eta = torch.div(time, scale)
+    h1 = torch.exp(-torch.pow(t_over_eta, shape))
+    h1_bis = torch.pow(t_over_eta, shape - 1)
+    params_aux = torch.div(torch.mul(alphas, shape), scale)
+    return -torch.mean(event * torch.log(torch.sum(torch.mul(torch.mul(params_aux, h1_bis), h1), 0))
+                       + (~event) * torch.log(torch.sum(alphas * h1, 0)))
+
+
+def weibull_survival_time(params: torch.Tensor, activations: bool = False):
+    """ Survival time from mean of mixture of Weibull distributions.
+    From DeepWeiSurv: https://doi.org/10.1007/978-3-030-47426-3_53
+
+    Parameters
+    ----------
+    params: torch.Tensor, of shape (..., n_samples, 2 | p * 3)
+        Parameters of the mixture of p Weibull distributions, namely weights, shapes, and scales
+
+    activations: bool
+        Whether to apply ELU activations
+    """
+    # Extract mixture alpha, shape, and scale from params
+    if params.shape[-1] == 2 or params.shape[-1] == 3:
+        n_dists = 1
+        weight, shape, scale = 1, params[..., [-2]], params[..., [-1]]
+    elif params.shape[-1] % 3 == 0:
+        n_dists = params.shape[-1] // 3
+        weight, shape, scale = params[..., 0:n_dists], params[..., n_dists:2 * n_dists], params[..., 2 * n_dists:]
+    else:
+        raise ValueError('Unexpected number of Weibull parameters: ' + str(params.shape[-1]))
+
+    # Apply softmax if needed
+    if activations and n_dists > 1:
+        weight = torch.softmax(weight, dim=-1)
+
+    # Apply ELU activations if requested
+    if activations:
+        shape = torch.nn.functional.elu(shape) + 2
+        scale = torch.nn.functional.elu(scale) + 1 + 1e-5
+
+    # Return weighted mean as survival time estimate
+    return torch.sum(weight * scale * torch.exp(torch.lgamma(1 + 1 / shape)), dim=-1)
+
+
+def mse_with_pairwise_rank(estimate: torch.Tensor, event: torch.Tensor, time: torch.Tensor) -> torch.Tensor:
+    """ Extended mean squared error and pairwise ranking loss.
+    From RankDeepSurv: https://doi.org/10.1016/j.artmed.2019.06.001
+
+    """
+    # First part of the loss function is a simple mean squared error
+    error = time - estimate
+    loss1 = torch.mean(torch.square(error) * (event | (estimate <= time)).float(), dim=-1)
+
+    # Here, we add extra dimensions to enable pairwise comparisons of (i,j)
+    event_i, event_j = event.unsqueeze(-2), event.unsqueeze(-1)
+    time_i, time_j = time.unsqueeze(-2), time.unsqueeze(-1)
+    # The compatibility matrix specifies which pairs (i,j) can be compared accounting for censoring
+    comp = event_i & (event_j | (time_i <= time_j))  # matrix C in report
+
+    # Second part of the loss function encourages correct ranking among compatible pairs based on relative distance
+    # To save computations, we can reuse the errors needed for loss1 as
+    # (time_j - time_i) - (estimate_j - estimate_i) = (time_j - estimate_j) - (time_i - estimate_i)
+    error_i, error_j = error.unsqueeze(-2), error.unsqueeze(-1)
+    diff = torch.clamp(error_j - error_i, min=0)  # matrix D in report, fused with condition
+    loss2 = torch.sum(diff * comp, dim=(-2, -1)) / event.shape[-1]
+
+    return loss1 + loss2
+
+
+def discrete_with_pairwise_rank(estimate: torch.Tensor, event: torch.Tensor, time: torch.Tensor,
+                                alpha: float, sigma: float) -> torch.Tensor:
+    """ Discrete negative log likelihood and pairwise ranking loss.
+    From DeepHit: https://doi.org/10.1609/aaai.v32i1.11842
+
+    Parameters
+    ----------
+    estimate: torch.Tensor, of shape (..., n_samples, n_times)
+        Discrete time probabilities estimated by model
+
+    event: torch.Tensor, of shape (..., n_samples)
+        Event indicator denoting whether time is of observed event or dropout
+
+    time: torch.Tensor, of shape (..., n_samples)
+        Time of either observed event or dropout
+
+    alpha: float
+        Weight for the ranking loss component
+
+    sigma: float
+        Bandwidth of the radial basis function in the ranking loss component
+    """
+    cum_incidence = torch.cumsum(estimate, dim=-1)
+    estimate_t = torch.gather(estimate, dim=-1, index=time.unsqueeze(-1)).squeeze(-1)
+    cum_incidence_t = torch.gather(cum_incidence, dim=-1, index=time.unsqueeze(-1)).squeeze(-1)
+    eps = torch.exp(
+        torch.tensor(-100, device=estimate.device))  # -100 is also used in PyTorch's binary cross entropy as a cut-off
+    loss1 = (- torch.sum(event * torch.log(torch.clamp(estimate_t, min=eps)), dim=-1)
+             - torch.sum(~event * torch.log(torch.clamp(1 - cum_incidence_t, min=eps)), dim=-1))
+
+    # Here, we add extra dimensions to enable pairwise comparisons of (i,j)
+    event_i, event_j = event.unsqueeze(-2), event.unsqueeze(-1)
+    time_i, time_j = time.unsqueeze(-2), time.unsqueeze(-1)
+    # The acceptability matrix specifies which (i, j) can be compared
+    acc = event_i & (time_i < time_j)
+
+    cum_incidence_ti = cum_incidence_t.unsqueeze(-2)
+    cum_incidence_tij = cum_incidence[:, time]
+    loss2 = torch.sum(torch.exp(-(cum_incidence_ti - cum_incidence_tij) / sigma) * acc, dim=(-2, -1))
+
+    return loss1 + alpha * loss2

torch_survival-0.1.0a2/src/torch_survival/metrics.py

@@ -0,0 +1,57 @@
+from typing import Literal
+
+import torch
+
+
+def concordance_index(estimate: torch.Tensor, event: torch.Tensor, time: torch.Tensor,
+                      mode: Literal['risk', 'time'] = 'risk') -> torch.Tensor:
+    r""" Compute Harrell's concordance index or c-index.
+
+    This essentially computes the ratio of correctly ordered pairs while accounting for censoring. Given estimates
+    :math:`\eta`, event indicators :math:`\delta`, and times :math:`t` it is described by
+
+    .. math::
+
+       C = \frac{\sum_{i,j} \mathbb{1}_{\eta_i < \eta_j} \mathbb{1}_{t_i > t_j} \delta_j}
+       {\sum_{i,j} \mathbb{1}_{t_i > t_j} \delta_j}.
+
+    If your model directly predicts survival time, rather than risk, you need to negate the estimates as this function
+    assumes an inverse relationship between estimates and times, i.e., if one increases the other decreases.
+
+    Note
+    ----
+    While Harrell's concordance index is easy to interpret, it is known to be biased in the presence of higher amounts
+    of censoring [1]_. An alternative is Uno's concordance index, as implemented in
+    `sksurv.metrics.concordance_index_ipcw`.
+
+    .. [1] H. Uno, T. Cai, M. J. Pencina, R. B. D’Agostino, and L. J. Wei, “On the C‐statistics for evaluating overall
+       adequacy of risk prediction procedures with censored survival data,” Statistics in Medicine, vol. 30, no. 10, pp.
+       1105–1117, Jan. 2011, doi: 10.1002/sim.4154. Available: http://dx.doi.org/10.1002/sim.4154
+
+
+    Parameters
+    ----------
+    estimate: torch.Tensor, of shape (..., n_samples)
+        Risks or times estimated by model
+
+    event: torch.Tensor, of shape (..., n_samples)
+        Event indicator denoting whether time is of observed event or dropout
+
+    time: torch.Tensor, of shape (..., n_samples)
+        Time of either observed event or dropout
+
+    mode: Literal['risk', 'time']
+        Whether the passed estimates are risks or survival times
+
+    Returns
+    -------
+    torch.Tensor, of shape (...,)
+        Concordance index score
+    """
+    if mode == 'risk':
+        estimate_comp = estimate.unsqueeze(-1) < estimate.unsqueeze(-2)
+    else:
+        estimate_comp = estimate.unsqueeze(-1) > estimate.unsqueeze(-2)
+    time_comp = time.unsqueeze(-1) > time.unsqueeze(-2)
+    event = event.unsqueeze(-2)
+    return torch.sum(estimate_comp & time_comp & event, (-2, -1)) / torch.sum(time_comp & event, (-2, -1))

torch_survival-0.1.0a2/src/torch_survival/models/__init__.py

@@ -0,0 +1,3 @@
+# alphabetical imports to make models more easily accessible
+from .deepsurv import DeepSurv, DeepSurvSearchSpace
+from .deephit import DeepHit

torch_survival-0.1.0a2/src/torch_survival/models/deephit.py

@@ -0,0 +1,192 @@
+import functools
+import warnings
+from copy import deepcopy
+from typing import TypedDict
+
+import numpy as np
+import optuna
+import torch
+import torchtuples as tt
+from optuna.samplers import TPESampler
+from pycox.evaluation import EvalSurv
+from pycox.models import DeepHitSingle
+from pycox.preprocessing.label_transforms import LabTransDiscreteTime
+from sklearn.base import BaseEstimator
+from sklearn.model_selection import StratifiedKFold, train_test_split
+from sklearn.utils.validation import check_is_fitted, validate_data
+from sksurv.base import SurvivalAnalysisMixin
+from sksurv.util import check_array_survival
+from torch import nn
+
+from torch_survival.progress import OptunaProgressCallback
+from torch_survival.utils import merge_configs
+
+
+class DeepHitNetwork(nn.Module):
+    def __init__(self, n_inputs, n_times):
+        super().__init__()
+        self.shared_network = nn.Sequential(
+            nn.Linear(n_inputs, 3 * n_inputs),
+            nn.ReLU(),
+            nn.Dropout(p=0.6),
+        )
+        self.cause_network = nn.Sequential(
+            # input consists of shared network output + features
+            nn.Linear(4 * n_inputs, 5 * n_inputs),
+            nn.ReLU(),
+            nn.Dropout(p=0.6),
+            nn.Linear(5 * n_inputs, 3 * n_inputs),
+            nn.ReLU(),
+            nn.Dropout(p=0.6),
+            nn.Linear(3 * n_inputs, n_times),
+        )
+
+    def forward(self, x):
+        x = torch.concat((self.shared_network(x), x), dim=-1)
+        x = self.cause_network(x)
+        return x
+
+
+class DeepHitSearchSpace(TypedDict):
+    #: Weight for the ranking loss component
+    alpha: tuple[float, float] | float
+    #: Bandwidth of the radial basis function in the ranking loss component
+    sigma: tuple[float, float] | float
+
+
+class DeepHit(SurvivalAnalysisMixin, BaseEstimator):
+    r""" Implements the DeepHit model presented by Lee et al. [1]_.
+
+    Uses a deep neural network trained with the negative log likelihood of the survival time probability distribution,
+    as estimated over discrete time intervals, combined with a ranking loss.
+
+    .. [1] C. Lee, W. Zame, J. Yoon, and M. Van der Schaar, “DeepHit: A Deep Learning Approach to Survival Analysis
+    With Competing Risks,” AAAI, vol. 32, no. 1, Apr. 2018, doi: 10.1609/aaai.v32i1.11842. Available:
+    http://dx.doi.org/10.1609/aaai.v32i1.11842
+    """
+
+    #: Default hyperparameter search space
+    default_search_space: DeepHitSearchSpace = {
+        'n_times': (10, 500),
+        'alpha': (0, 1),
+        'sigma': (1e-3, 1e1),
+    }
+
+    def __init__(self, search_space: DeepHitSearchSpace | None = None, n_epochs=100, random_state=None, device=None):
+        self.search_space = deepcopy(self.default_search_space)
+        if search_space:
+            self.search_space = merge_configs(self.search_space, search_space)
+        self.n_epochs = n_epochs
+        self.random_state = random_state
+        self.device = device
+        if self.device is None:
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+    def _optimize(self, trial: optuna.Trial, X, y_event, y_time):
+        # Do 5-fold cross validation
+        scores = []
+        for train_idx, test_idx in StratifiedKFold(n_splits=5).split(X, y_event):
+            model, _ = self._train(trial, X[train_idx], y_event[train_idx], y_time[train_idx])
+            surv = model.predict_surv_df(X[test_idx])
+            c_index = EvalSurv(surv, y_time[test_idx], y_event[test_idx], censor_surv='km').concordance_td('antolini')
+            scores.append(c_index)
+        return - sum(scores) / len(scores)
+
+    def _train(self, trial: optuna.Trial, X, y_event, y_time):
+        n_inputs, n_outputs = X.shape[-1], 1
+
+        # Discretize event times
+        n_times = self.search_space['n_times']
+        if not isinstance(n_times, int):
+            n_times = trial.suggest_int('n_times', *n_times)
+        y_trans = LabTransDiscreteTime(n_times)
+        y = y_trans.fit_transform(y_time, y_event)
+
+        # Split into training and validation for early stopping
+        X_train, X_val, y_time_train, y_time_val, y_event_train, y_event_val = \
+            train_test_split(X, *y, test_size=0.2, stratify=y_event, random_state=self.random_state)
+        y_train = (y_time_train, y_event_train)
+        y_val = (y_time_val, y_event_val)
+
+        # Sample loss parameters
+        alpha = self.search_space['alpha']
+        if not isinstance(alpha, float):
+            alpha = trial.suggest_float('alpha', *alpha)
+        sigma = self.search_space['sigma']
+        if not isinstance(sigma, float):
+            sigma = trial.suggest_float('sigma', *sigma, log=True)
+
+        # Train and return model
+        net = DeepHitNetwork(n_inputs, y_trans.out_features)
+        model = DeepHitSingle(net, tt.optim.Adam(lr=1e-4), duration_index=y_trans.cuts,
+                              alpha=alpha, sigma=sigma, device=self.device)
+        callbacks = [tt.callbacks.EarlyStopping(patience=10)]
+        model.fit(X_train, y_train, batch_size=50, epochs=100, callbacks=callbacks,
+                  val_data=(X_val, y_val), verbose=False)
+        return model, y_trans.cuts
+
+    def fit(self, X, y):
+        """ Fit the model to the given survival data.
+
+        Parameters
+        ----------
+        X: array-like, shape = (n_samples, n_features)
+            Data matrix.
+        y: structured array, shape = (n_samples,)
+            A structured array with two fields. The first field is a boolean where ``True`` indicates an event and
+            ``False`` indicates right-censoring. The second field is a float with the time of event or time of
+            censoring.
+
+        Returns
+        -------
+        self
+        """
+        # Validate and extract data
+        X, y = validate_data(self, X, y)
+        X = X.astype(np.float32)
+        y_event, y_time = check_array_survival(X, y)
+
+        # Seed PyTorch random number generator
+        if self.random_state:
+            torch.manual_seed(self.random_state)
+
+        # Optimize hyperparameters
+        optuna.logging.disable_default_handler()
+        warnings.filterwarnings('ignore', category=optuna.exceptions.ExperimentalWarning)
+        with OptunaProgressCallback(model_name='DeepHit', n_trials=50) as callback:
+            study = optuna.create_study(sampler=TPESampler(seed=self.random_state) if self.random_state else None)
+            objective = functools.partial(self._optimize, X=X, y_event=y_event, y_time=y_time)
+            study.optimize(objective, n_trials=50, callbacks=[callback])
+
+        # Train model
+        self.optuna_params_ = study.best_params
+        self.model_, self.disc_times_ = self._train(study.best_trial, X, y_event, y_time)
+
+        return self
+
+    @torch.no_grad()
+    def predict(self, X):
+        """ Predict survival times.
+
+        The survival time is estimated based on the mean of the mixture of Weibull distributions predicted by the
+        neural network.
+
+        Parameters
+        ----------
+        X: array-like, shape = (n_samples, n_features)
+            Data matrix.
+
+        Returns
+        -------
+        survival_time: array, shape = (n_samples,)
+            Predicted survival times.
+        """
+        check_is_fitted(self)
+        X = validate_data(self, X)
+        X = X.astype(np.float32)
+        pmf = self.model_.predict_pmf(X)
+        times = pmf @ self.disc_times_
+        return times
+
+    def get_optuna_params(self):
+        return self.optuna_params_

torch_survival-0.1.0a2/src/torch_survival/models/deepsurv.py

@@ -0,0 +1,180 @@
+import functools
+import warnings
+from copy import deepcopy
+from typing import TypedDict
+
+import optuna
+import torch
+from optuna.samplers import TPESampler
+from sklearn.base import BaseEstimator
+from sklearn.model_selection import StratifiedKFold
+from sklearn.utils.validation import check_is_fitted, validate_data
+from sksurv.base import SurvivalAnalysisMixin
+from sksurv.util import check_array_survival
+
+from torch_survival.config import NetworkConfig, OptimizerConfig
+from torch_survival.losses import cox_neg_log_likelihood
+from torch_survival.metrics import concordance_index
+from torch_survival.progress import OptunaProgressCallback
+from torch_survival.sample import sample_network, sample_optimizer
+from torch_survival.utils import merge_configs
+
+
+class DeepSurvSearchSpace(TypedDict):
+    #: Neural network configuration
+    network: NetworkConfig
+    #: Optimizer configuration
+    optimizer: OptimizerConfig
+
+
+class DeepSurv(SurvivalAnalysisMixin, BaseEstimator):
+    r""" Implements the DeepSurv model presented by Katzman et al. [1]_.
+
+    Uses a deep neural network trained with the Cox negative log partial likelihood to estimate the risk of each
+    individual. The network's configuration is tuned using the Sobol solver. This implementation tries to stay faithful
+    to the original paper, with the following deviations:
+
+    * Optuna's default TPE sampler is used in favor of the Sobol sampler with 5-fold internal cross-validation instead
+      of 3-fold internal cross-validation.
+    * The hyperparameter search space is not detailed in the original paper, and the reference implementation is
+      underspecified. We thus define our own shared search space in `DeepSurvSearchSpace`.
+    * Our implementation does not support or tune :math:`\ell_2` regularization. We found this to be detrimental to
+      performance and were unable to fully replicate the described weight regularization.
+
+    .. [1] J. L. Katzman, U. Shaham, A. Cloninger, J. Bates, T. Jiang, and Y. Kluger, “DeepSurv: personalized treatment
+       recommender system using a Cox proportional hazards deep neural network,” BMC Med Res Methodol, vol. 18, no. 1,
+       Feb. 2018, doi: 10.1186/s12874-018-0482-1. Available: http://dx.doi.org/10.1186/s12874-018-0482-1
+    """
+
+    #: Default hyperparameter search space
+    default_search_space: DeepSurvSearchSpace = {
+        'network': {
+            'layers': {
+                'max_layers': 4,
+                'max_nodes_per_layer': 50,
+            },
+            'activation': ['relu', 'selu'],
+            'dropout': (0.0, 0.5),
+        },
+        'optimizer': {
+            'name': ['sgd', 'adam'],
+            'lr': (1e-7, 1e-3),
+            'scheduler': 'inverse_time',
+            'decay': (0.0, 0.001),
+            'momentum': (0.8, 0.95),
+        },
+    }
+
+    def __init__(self, search_space: DeepSurvSearchSpace | None = None, n_epochs=500, random_state=None, device=None):
+        self.search_space = deepcopy(self.default_search_space)
+        if search_space:
+            self.search_space = merge_configs(self.search_space, search_space)
+        self.n_epochs = n_epochs
+        self.random_state = random_state
+        self.device = device
+        if self.device is None:
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+    def _optimize(self, trial: optuna.Trial, X, y_event, y_time):
+        # Do 5-fold cross validation
+        scores = []
+        for train_idx, test_idx in StratifiedKFold(n_splits=5).split(X, y_event.cpu()):
+            model = self._train(trial, X[train_idx], y_event[train_idx], y_time[train_idx])
+            model.eval()
+            risks = model(X[test_idx]).squeeze(dim=-1)
+            c_index = concordance_index(risks, y_event[test_idx], y_time[test_idx])
+            scores.append(c_index)
+        return - sum(scores) / len(scores)
+
+    def _train(self, trial: optuna.Trial, X, y_event, y_time):
+        # Set up model, optimizer, and scheduler
+        n_inputs, n_outputs = X.shape[-1], 1
+        model = sample_network(trial, self.search_space['network'], n_inputs, n_outputs)
+        optimizer, scheduler = sample_optimizer(trial, self.search_space['optimizer'], model)
+
+        # Pre-sort dataset based on time
+        sort_idx = torch.argsort(y_time, descending=True)
+        X = X[sort_idx]
+        y_event = y_event[sort_idx]
+        y_time = y_time[sort_idx]
+
+        # Train and return model
+        model.to(self.device)
+        for i in range(self.n_epochs):
+            optimizer.zero_grad()
+            risk = model(X).squeeze(-1)
+            loss = cox_neg_log_likelihood(risk, y_event, y_time, sort=False)
+            loss.backward()
+            optimizer.step()
+            if scheduler:
+                scheduler.step()
+        return model
+
+    def fit(self, X, y):
+        """ Fit the model to the given survival data.
+
+        Parameters
+        ----------
+        X: array-like, shape = (n_samples, n_features)
+            Data matrix.
+        y: structured array, shape = (n_samples,)
+            A structured array with two fields. The first field is a boolean where ``True`` indicates an event and
+            ``False`` indicates right-censoring. The second field is a float with the time of event or time of
+            censoring.
+
+        Returns
+        -------
+        self
+        """
+        # Validate and extract data
+        X, y = validate_data(self, X, y)
+        y_event, y_time = check_array_survival(X, y)
+
+        # Convert data to tensors
+        X = torch.as_tensor(X, dtype=torch.float32, device=self.device)
+        y_event = torch.as_tensor(y_event, dtype=torch.bool, device=self.device)
+        y_time = torch.as_tensor(y_time.copy(), dtype=torch.float32, device=self.device)
+
+        # Seed PyTorch random number generator
+        if self.random_state:
+            torch.manual_seed(self.random_state)
+
+        # Optimize hyperparameters
+        optuna.logging.disable_default_handler()
+        warnings.filterwarnings('ignore', category=optuna.exceptions.ExperimentalWarning)
+        with OptunaProgressCallback(model_name='DeepSurv', n_trials=50) as callback:
+            study = optuna.create_study(sampler=TPESampler(seed=self.random_state) if self.random_state else None)
+            objective = functools.partial(self._optimize, X=X, y_event=y_event, y_time=y_time)
+            study.optimize(objective, n_trials=50, callbacks=[callback])
+
+        # Train final model with best hyperparameters
+        self.optuna_params_ = study.best_params
+        self.model_ = self._train(study.best_trial, X, y_event, y_time)
+        self.model_.eval()
+
+        return self
+
+    @torch.no_grad()
+    def predict(self, X):
+        """ Predict risk scores.
+
+        The risk score is predicted directly by a neural network. A higher score indicates a higher risk of experiencing
+        the event.
+
+        Parameters
+        ----------
+        X: array-like, shape = (n_samples, n_features)
+            Data matrix.
+
+        Returns
+        -------
+        risk_score: array, shape = (n_samples,)
+            Predicted risk scores.
+        """
+        check_is_fitted(self)
+        X = validate_data(self, X)
+        X = torch.as_tensor(X, dtype=torch.float32, device=self.device)
+        return self.model_(X).detach().squeeze(dim=-1).cpu().numpy()
+
+    def get_optuna_params(self):
+        return self.optuna_params_

torch_survival-0.1.0a2/src/torch_survival/progress.py

@@ -0,0 +1,29 @@
+import optuna
+from rich.progress import Progress, TextColumn, BarColumn, MofNCompleteColumn, TimeRemainingColumn
+
+
+class OptunaProgressCallback:
+    def __init__(self, model_name, n_trials, verbose=1):
+        self.verbose = verbose
+        if self.verbose > 0:
+            # At verbosity = 1 we show a progress bar for trials
+            self.progress = Progress(
+                TextColumn('{task.description}'),
+                BarColumn(),
+                TextColumn('{task.fields[score]}'),
+                MofNCompleteColumn(),
+                TimeRemainingColumn(),
+            )
+            self.tune_task = self.progress.add_task('[green]Optimizing ' + model_name, total=n_trials, score='-.--')
+
+    def __call__(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> None:
+        if self.verbose > 0:
+            score = '{:.2f}'.format(abs(study.best_value))
+            self.progress.update(self.tune_task, advance=1, score=score)
+
+    def __enter__(self):
+        self.progress.start()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.progress.stop()

torch_survival-0.1.0a2/src/torch_survival/sample.py

@@ -0,0 +1,161 @@
+from typing import Mapping, Any
+
+import optuna
+import torch.optim as optim
+import torch.optim.lr_scheduler as sched
+from torch import nn
+
+from torch_survival.config import NetworkConfig, OptimizerConfig
+from torch_survival.utils import make_activation
+
+
+class SimpleNeuralNetwork(nn.Module):
+    def __init__(self, n_inputs, n_outputs, layers, activation, dropout):
+        super().__init__()
+        hidden = []
+        n_nodes = n_inputs
+        for nodes in layers:
+            hidden.append(nn.Linear(n_nodes, nodes))
+            hidden.append(make_activation(activation))
+            hidden.append(nn.Dropout(p=dropout))
+            n_nodes = nodes
+        self.hidden = nn.Sequential(*hidden)
+        self.output = nn.Linear(n_nodes, n_outputs)
+
+    def forward(self, x):
+        x = self.hidden(x)
+        return self.output(x)
+
+
+def sample_network(trial: optuna.Trial, config: NetworkConfig, n_inputs: int, n_outputs: int):
+    """
+    Sample a simple neural network using the provided configuration.
+
+    Parameters
+    ----------
+    trial: optuna.Trial
+        Active or fixed trial used to sample hyperparameters. Final configuration may be obtained by `study.best_trial`.
+    config: config.NetworkConfig
+        Network configuration specifying architecture and search constraints.
+    n_inputs: int
+        Number of inputs of the neural network. Typically depends on the data.
+    n_outputs: int
+        Number of outputs of the neural network. Typically depends on the modeling technique.
+
+    Returns
+    -------
+    nn.Module
+        PyTorch model consisting of linear layers.
+    """
+    layers = config['layers']
+    if not isinstance(layers, list):
+        n_layers = trial.suggest_int('layers', 0, layers['max_layers'])
+        layers = [trial.suggest_int('nodes_' + str(i + 1), 1, layers['max_nodes_per_layer']) for i in range(n_layers)]
+    activation = config['activation']
+    if not isinstance(activation, str):
+        activation = trial.suggest_categorical('activation', activation)
+    dropout = config['dropout']
+    if not isinstance(dropout, float):
+        dropout = trial.suggest_float('dropout', low=dropout[0], high=dropout[1])
+    return SimpleNeuralNetwork(n_inputs, n_outputs, layers, activation, dropout)
+
+
+def sample_optimizer(trial: optuna.Trial, config: OptimizerConfig, model: nn.Module):
+    """
+    Sample optimizer and scheduler for training using the provided configuration.
+
+    Parameters
+    ----------
+    trial: optuna.Trial
+        Active or fixed trial used to sample hyperparameters. Final configuration may be obtained by `study.best_trial`.
+    config: config.OptimizerConfig
+        Optimizer configuration specifying how the neural network should be trained.
+    model: nn.Module
+        Neural network whose parameters should be optimized.
+
+    Returns
+    -------
+    torch.optim.Optimizer
+        PyTorch optimizer that can be used for training.
+    torch.optim.lr_scheduler.LRScheduler or None
+        PyTorch scheduler that can be used to update learning rates.
+    """
+    # Sample optimizer parameters
+    lr = config['lr']
+    if not isinstance(lr, float):
+        lr = trial.suggest_float('lr', *lr, log=True)
+    momentum = config['momentum']
+    if not isinstance(momentum, float):
+        momentum = trial.suggest_float('momentum', *momentum)
+    # Initialize optimizer
+    optimizer = None
+    optimizer_name = config['name']
+    if not isinstance(optimizer_name, str):
+        optimizer_name = trial.suggest_categorical('optimizer', optimizer_name)
+    if optimizer_name == 'sgd':
+        optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)
+    if optimizer_name == 'adam':
+        optimizer = optim.Adam(model.parameters(), lr=lr, betas=(momentum, 0.999))
+    if optimizer is None:
+        raise ValueError('Optimizer with name `{}` is not supported'.format(config['name']))
+    # Sample optimizer parameters
+    decay = config['decay']
+    if not isinstance(decay, float):
+        decay = trial.suggest_float('decay', *decay)
+    # Initialize scheduler
+    scheduler = None
+    scheduler_name = config['scheduler']
+    if not isinstance(scheduler_name, str):
+        scheduler_name = trial.suggest_categorical('scheduler', scheduler_name)
+    if scheduler_name == 'inverse_time':
+        scheduler = sched.LambdaLR(optimizer, lr_lambda=lambda epoch: 1 / (1 + epoch * decay))
+    # Return both
+    return optimizer, scheduler
+
+
+def sample_int(trial: optuna.Trial, config: Mapping[str, Any], key: str) -> int:
+    """
+    Helper to sample an integer value from a search space configuration.
+
+    Parameters
+    ----------
+    trial: optuna.Trial
+        Active or fixed trial used to sample hyperparameters. Final configuration may be obtained by `study.best_trial`.
+    config: Mapping[str, Any]
+        Search space configuration dictionary.
+    key: str
+        Name of configuration parameter that should be sampled.
+
+    Returns
+    -------
+    int:
+        The sampled value.
+    """
+    value = config[key]
+    if not isinstance(value, int):
+        value = trial.suggest_int(key, *value)
+    return value
+
+
+def sample_float(trial: optuna.Trial, config: Mapping[str, Any], key: str) -> float:
+    """
+    Helper to sample a floating point value from a search space configuration.
+
+    Parameters
+    ----------
+    trial: optuna.Trial
+        Active or fixed trial used to sample hyperparameters. Final configuration may be obtained by `study.best_trial`.
+    config: Mapping[str, Any]
+        Search space configuration dictionary.
+    key: str
+        Name of configuration parameter that should be sampled.
+
+    Returns
+    -------
+    float:
+        The sampled value.
+    """
+    value = config[key]
+    if not isinstance(value, float):
+        value = trial.suggest_float(key, *value)
+    return value

torch_survival-0.1.0a2/src/torch_survival/utils.py

@@ -0,0 +1,87 @@
+import inspect
+from typing import TypeVar
+
+import torch.nn as nn
+
+_TypedDict = TypeVar('_TypedDict')
+
+
+def make_activation(query: str):
+    """
+    Resolves the name of an activation function to its PyTorch layer. This method uses inspection and should thus be
+    able to resolve any current PyTorch activation function, e.g., `relu` gets mapped to torch.nn.ReLU, and so forth.
+
+    Parameters
+    ----------
+    query: str
+        Name of activation function that should be created
+
+    Returns
+    -------
+    torch.nn.Module
+        PyTorch layer corresponding to query, initialized with default values
+    """
+    query = query.lower()
+    matches = [obj for name, obj in inspect.getmembers(nn) if query == name.lower()]
+    if len(matches) == 0:
+        raise ValueError('Found no candidate for `{}` activation'.format(query))
+    if len(matches) > 1:
+        raise ValueError('Found multiple candidates for `{}` activation'.format(query))
+    return matches[0]()
+
+
+def merge_configs(default_config: _TypedDict, user_config: _TypedDict) -> _TypedDict:
+    """
+    Recursively merges default configuration with user configuration.
+
+    Parameters
+    ----------
+    default_config: subclass of TypedDict
+        Default model configuration
+    user_config: subclass of TypedDict
+        User-provided model configuration overriding defaults
+
+    Returns
+    -------
+    subclass of TypedDict
+        Merged model configuration
+    """
+    for k, v in user_config.items():
+        if k in default_config:
+            if isinstance(default_config[k], dict) and isinstance(v, dict):
+                default_config[k] = merge_configs(default_config[k], v)
+            else:
+                default_config[k] = v
+    return default_config
+
+# def sample_params(search_space, trial: optuna.Trial):
+#     """
+#     Samples hyperparameter search space, allowing for fixed
+#
+#     Parameters
+#     ----------
+#     search_space: Mapping[str, ParamConfig]
+#         Dictionary of parameter configurations.
+#     trial: optuna.Trial
+#         Active trial for sampling parameters.
+#
+#     Returns
+#     -------
+#     params: dict[str, int | float | str]
+#         Sampled parameters.
+#     """
+#     params = {}
+#     for key, value in search_space.items():
+#         if type(value) in [int, float, str]:
+#             # The value is already fixed and will be passed on as-is
+#             params[key] = value
+#         elif type(value) is list:
+#             # The value is a list of possible categories
+#             params[key] = trial.suggest_categorical(key, value)
+#         elif type(value) is tuple and type(value[0]) is int:
+#             # The value are lower and upper bounds for an integer
+#             params[key] = trial.suggest_int(key, *value[:2], log=value[2] if len(value) > 2 else False)
+#         elif type(value) is tuple and type(value[0]) is float:
+#             # The value are lower and upper bounds for an integer
+#             params[key] = trial.suggest_float(key, *value[:2], log=value[2] if len(value) > 2 else False)
+#     return params