survivalpredict 0.0.1__py3-none-any.whl

survivalpredict/_allen_additive.py

@@ -0,0 +1,68 @@
+import numba as nb
+import numpy as np
+
+_estimate_allen_additive_hazard_time_weights_sig = nb.types.Tuple(
+    (
+        nb.types.Array(nb.types.float64, 2, "C", False, aligned=True),
+        nb.types.Array(nb.types.int64, 1, "C", False, aligned=True),
+    )
+)(
+    nb.types.Array(nb.types.float64, 2, "C", False, aligned=True),
+    nb.types.Array(nb.types.int64, 1, "C", False, aligned=True),
+    nb.types.Array(nb.types.bool_, 1, "C", False, aligned=True),
+    nb.types.Array(nb.types.int64, 1, "C", False, aligned=True),
+    nb.types.float64,
+)
+
+
+@nb.njit(_estimate_allen_additive_hazard_time_weights_sig, cache=True)
+def _estimate_allen_additive_hazard_time_weights(X, times, events, times_start, alpha):
+    hazard_weights_times = np.unique(times[events])
+    hazard_weights = np.empty((hazard_weights_times.shape[0], X.shape[1]))
+
+    for i, t in enumerate(hazard_weights_times):
+        survived_at_time = times >= t
+        exits = times == t
+        deaths_mask = np.logical_and(exits, events)
+        not_right_censored = np.logical_and(survived_at_time, deaths_mask)
+        not_left_censored = times_start < t
+        not_censored = np.logical_or(not_right_censored, not_left_censored)
+
+        if not_censored.any() and not_censored.sum() > 1:
+            death_as_target = deaths_mask[not_censored].astype(np.float64)
+
+            X_mask = X[not_censored, :]
+
+            a = np.dot(X_mask.T, X_mask) + alpha
+            b = np.dot(X_mask.T, death_as_target)
+
+            try:
+                w = np.linalg.solve(a, b)
+            except:
+                w = np.zeros(X.shape[1])
+   
+            hazard_weights[i, :] = w
+
+    return hazard_weights, hazard_weights_times
+
+
+_generate_hazards_at_times_from_allen_additive_hazard_weights_sig = nb.types.Array(
+    nb.types.float64, 2, "C", False, aligned=True
+)(
+    nb.types.Array(nb.types.float64, 2, "C", False, aligned=True),
+    nb.types.Array(nb.types.float64, 2, "C", False, aligned=True),
+    nb.types.Array(nb.types.int64, 1, "C", False, aligned=True),
+    nb.types.int64,
+)
+
+
+@nb.njit(_generate_hazards_at_times_from_allen_additive_hazard_weights_sig, cache=True)
+def _generate_hazards_at_times_from_allen_additive_hazard_weights(
+    X, hazard_weights, hazard_weights_times, max_time
+):
+    hazards = np.zeros((X.shape[0], max_time))
+
+    for t, w in zip(hazard_weights_times, hazard_weights):
+        hazards[:, t] = np.dot(X, w)
+
+    return hazards

survivalpredict/_base_hazard.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+
+def _get_breslow_base_hazard(
+    risk,
+    times,
+    events,
+    max_time,
+):
+    unique_times = np.arange(1, max_time + 1)
+    rows_at_risk_at_time = times[:, np.newaxis] > unique_times
+
+    failure_per_unique_time = np.bincount(
+        times.astype(np.int64), events, minlength=max_time + 1
+    )[1:]
+
+    risk_per_time = np.dot(risk, rows_at_risk_at_time)
+    base_hazard = np.divide(
+        failure_per_unique_time,
+        risk_per_time,
+        out=np.zeros(max_time),
+        where=risk_per_time != 0,
+    )
+    return base_hazard

survivalpredict/_cox_net_ph.py

@@ -0,0 +1,254 @@
+from functools import reduce
+from itertools import pairwise
+from typing import Literal, Optional
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+from optax import GradientTransformationExtraArgs
+
+
+def get_gradient_updater(
+    gradient_updater: Literal[
+        "adadelta",
+        "adagrad",
+        "adam",
+        "adamax",
+        "rmsprop",
+    ] = "adam",
+    learning_rate: float = 0.01,
+    beta1: float = 0.9,
+    beta2: float = 0.999,
+    epsilon: float = 0.0000001,
+    rho: float = 0.95,
+    decay: float = 0.9,
+) -> GradientTransformationExtraArgs:
+    if gradient_updater == "adadelta":
+        return optax.adadelta(learning_rate=learning_rate, rho=rho, esp=epsilon)
+
+    elif gradient_updater == "adagrad":
+        return optax.adagrad(learning_rate=learning_rate, eps=epsilon)
+
+    elif gradient_updater == "adam":
+        return optax.adam(learning_rate=learning_rate, b1=beta1, b2=beta2, eps=epsilon)
+
+    elif gradient_updater == "adamax":
+        return optax.adamax(
+            learning_rate=learning_rate, b1=beta1, b2=beta1, eps=epsilon
+        )
+
+    else:  # gradient_updater == 'rmsprop':
+        return optax.rmsprop(learning_rate=learning_rate, decay=decay, eps=epsilon)
+
+
+def relu_jax(x):
+    return jnp.where(x < 0, 0, x)
+
+
+def _reverse_cumsum_jax(a):
+    return jnp.flip(jnp.cumsum(jnp.flip(a)))
+
+
+def get_init_weights(
+    input_n_cols: int,
+    hidden_layers: list[int],
+    init_dis: Literal["uniform", "normal"] = "uniform",
+):
+    weight_matrix_shapes = list(pairwise([input_n_cols] + hidden_layers))
+    if init_dis == "uniform":
+        initializer = jax.nn.initializers.he_uniform()
+    else:
+        initializer = jax.nn.initializers.he_normal()
+
+    jax_key = jax.random.key(np.random.randint(-10000, 10000))
+
+    weight_matrix_shapes.append((hidden_layers[-1], 1))
+    weights = [initializer(jax_key, shape=ws) for ws in weight_matrix_shapes]
+    weights[-1] = weights[-1].flatten()
+    return weights
+
+
+def _get_cox_net_ph_loss(
+    weights,
+    X_strata: list[np.ndarray[tuple[int, int], np.dtype[np.floating]]],
+    n_strata: int,
+    events_strata: list[np.ndarray[tuple[int], np.dtype[np.bool_]]],
+    time_end_return_inverse_strata: list[np.ndarray[tuple[int], np.dtype[np.integer]]],
+    n_unique_times_strata: list[int],
+    alpha=0.0,
+    l1_ratio=0.5,
+    time_start_return_inverse_strata: (
+        list[np.ndarray[tuple[int], np.dtype[np.integer]]] | None
+    ) = None,
+):
+
+    uses_left_censorship = time_start_return_inverse_strata is not None
+
+    partial_log_likelihood_per_strata = []
+
+    abs_weights_sum = reduce(lambda a, b: a + b, [jnp.sum(jnp.abs(w)) for w in weights])
+    square_weights_sum = reduce(
+        lambda a, b: a + b, [jnp.sum(jnp.square(w)) for w in weights]
+    )
+
+    for s_i in range(n_strata):
+        X = X_strata[s_i]
+        time_end_return_inverse = time_end_return_inverse_strata[s_i]
+        n_unique_times = n_unique_times_strata[s_i]
+        events = events_strata[s_i]
+        if uses_left_censorship:
+            time_start_return_inverse = time_start_return_inverse_strata[s_i]
+
+        matrixs_for_reduce_dot = [X] + weights[:-1]
+
+        second_to_last_layer = reduce(
+            lambda a, b: relu_jax(
+                jnp.dot(
+                    a,
+                    b,
+                )
+            ),
+            matrixs_for_reduce_dot,
+        )
+
+        o = jnp.dot(second_to_last_layer, weights[-1])
+        o_exp = jnp.exp(o)
+
+        if uses_left_censorship:
+            risk_removed_at_time = jnp.bincount(
+                time_end_return_inverse, weights=o_exp, minlength=n_unique_times
+            )
+            risk_added_at_time = jnp.bincount(
+                time_start_return_inverse, weights=o_exp, minlength=n_unique_times
+            )
+            risk_at_time = jnp.cumsum(risk_added_at_time - risk_removed_at_time)
+            risk_set = risk_at_time[time_end_return_inverse - 1]
+        else:
+            risk_set = _reverse_cumsum_jax(
+                jnp.bincount(
+                    time_end_return_inverse, weights=o_exp, minlength=n_unique_times
+                )
+            )[time_end_return_inverse]
+
+        loss = -jnp.sum(events * (o - jnp.log(risk_set)))
+
+        partial_log_likelihood_per_strata.append(loss)
+
+    l1 = alpha * l1_ratio * abs_weights_sum
+    l2 = 0.5 * alpha * (1.0 - l1_ratio) * square_weights_sum
+
+    return reduce(lambda a, b: a + b, partial_log_likelihood_per_strata + [l1, l2])
+
+
+def train_cox_net_ph(
+    X_strata: list[np.ndarray[tuple[int, int], np.dtype[np.floating]]],
+    n_strata: int,
+    events_strata: list[np.ndarray[tuple[int], np.dtype[np.bool_]]],
+    time_end_return_inverse_strata: list[np.ndarray[tuple[int], np.dtype[np.integer]]],
+    n_unique_times_strata: list[int],
+    hidden_layers: list[int],
+    weights: Optional[list[np.ndarray]] = None,
+    alpha: float = 0.0,
+    l1_ratio: float = 0.5,
+    init_dis: Literal["uniform", "normal"] = "uniform",
+    track_loss=True,
+    max_iter=100,
+    gradient_updater: Literal[
+        "adadelta",
+        "adagrad",
+        "adam",
+        "adamax",
+        "rmsprop",
+    ] = "adam",
+    learning_rate: float = 0.01,
+    beta1: float = 0.9,
+    beta2: float = 0.999,
+    epsilon: float = 0.0000001,
+    rho: float = 0.95,
+    decay: float = 0.9,
+    time_start_return_inverse_strata: (
+        list[np.ndarray[tuple[int], np.dtype[np.integer]]] | None
+    ) = None,
+) -> tuple[list[np.ndarray], float, list[float]]:
+    if weights is None:
+        weights = get_init_weights(X_strata[0].shape[1], hidden_layers, init_dis)
+
+    grad_updater = get_gradient_updater(
+        gradient_updater, learning_rate, beta1, beta2, epsilon, rho, decay
+    )
+
+    opt_state = grad_updater.init(weights)
+    get_cox_net_ph_grad = jax.grad(_get_cox_net_ph_loss)
+
+    losses_per_steps = []
+    loss = None
+
+    for i in range(max_iter):
+        jacobian = get_cox_net_ph_grad(
+            weights,
+            X_strata,
+            n_strata,
+            events_strata,
+            time_end_return_inverse_strata,
+            n_unique_times_strata,
+            alpha,
+            l1_ratio,
+            time_start_return_inverse_strata,
+        )
+        updates, opt_state = grad_updater.update(jacobian, opt_state, weights)
+        weights = optax.apply_updates(weights, updates)
+
+        if track_loss:
+            loss = _get_cox_net_ph_loss(
+                weights,
+                X_strata,
+                n_strata,
+                events_strata,
+                time_end_return_inverse_strata,
+                n_unique_times_strata,
+                alpha,
+                l1_ratio,
+                time_start_return_inverse_strata,
+            ).item()
+
+            losses_per_steps.append(loss)
+
+    if loss is None:
+        loss = _get_cox_net_ph_loss(
+            weights,
+            X_strata,
+            n_strata,
+            events_strata,
+            time_end_return_inverse_strata,
+            n_unique_times_strata,
+            alpha,
+            l1_ratio,
+            time_start_return_inverse_strata,
+        ).item()
+
+    weights_np = [np.array(w) for w in weights]
+
+    return weights_np, loss, losses_per_steps
+
+
+def relu_np(x):
+    return np.where(x < 0, 0, x)
+
+
+def get_relative_risk_from_cox_net_ph_weights(
+    X: np.ndarray[tuple[int, int], np.dtype[np.floating]], weights: list[np.ndarray]
+) -> np.ndarray[tuple[int], np.dtype[np.floating]]:
+    matrixs_for_reduce_dot_np = [X] + weights[:-1]
+
+    second_to_last_layer = reduce(
+        lambda a, b: relu_np(
+            np.dot(
+                a,
+                b,
+            )
+        ),
+        matrixs_for_reduce_dot_np,
+    )
+
+    return np.exp(np.dot(second_to_last_layer, weights[-1]))

survivalpredict/_cox_ph_elastic_net.py

@@ -0,0 +1,338 @@
+import numba as nb
+import numpy as np
+
+from ._stratification import _unique_with_return_inverse
+
+get_breslow_neg_log_likelihood_with_elasticnet_penalty_signature = nb.types.float64(
+    nb.types.Array(nb.types.float64, 1, "C", False, aligned=True),
+    nb.types.Array(nb.types.float64, 2, "C", False, aligned=True),
+    nb.types.Array(nb.types.bool_, 1, "C", False, aligned=True),
+    nb.types.Array(nb.types.int64, 1, "C", False, aligned=True),
+    nb.types.int64,
+    nb.types.float64,
+    nb.types.float64,
+    nb.types.bool_,
+)
+
+
+@nb.njit(get_breslow_neg_log_likelihood_with_elasticnet_penalty_signature, cache=True)
+def get_breslow_neg_log_likelihood_with_elasticnet_penalty(
+    weights,
+    X,
+    event,
+    time_return_inverse,
+    n_unique_times,
+    alpha,
+    l1_ratio,
+    scaled,
+):
+    n = float(X.shape[0])
+
+    p = np.dot(X, weights)
+    p_exp = np.exp(p)
+
+    risk_set_at_time = np.flip(
+        np.cumsum(
+            np.flip(
+                np.bincount(
+                    time_return_inverse, weights=p_exp, minlength=n_unique_times
+                )
+            )
+        )
+    )
+
+    risk_set = risk_set_at_time[time_return_inverse]
+
+    breslow_neg_log_likelihood = -np.sum(event * (p - np.log(risk_set)))
+
+    l1 = alpha * l1_ratio * np.abs(weights).sum()
+    l2 = 0.5 * alpha * (1.0 - l1_ratio) * np.square(weights).sum()
+    elasticnet_loss = l1 + l2
+
+    if scaled:
+        return (2 / n * breslow_neg_log_likelihood) + elasticnet_loss
+    else:
+        return breslow_neg_log_likelihood + elasticnet_loss
+
+
+@nb.njit
+def soft_threasholding_operator(z, t):
+    return np.fmax((np.abs(z) - t), 0) * np.sign(z)
+
+
+train_cox_elastic_net_signature = nb.types.Tuple(
+    (
+        nb.types.Array(nb.types.float64, 1, "C", False, aligned=True),
+        nb.types.float64,
+    )
+)(
+    nb.types.Array(nb.types.float64, 2, "C", False, aligned=True),
+    nb.types.Array(nb.types.int64, 1, "C", False, aligned=True),
+    nb.types.Array(nb.types.bool_, 1, "C", False, aligned=True),
+    nb.types.float64,
+    nb.types.float64,
+    nb.types.float64,
+    nb.types.int64,
+)
+
+
+@nb.njit(train_cox_elastic_net_signature, cache=True)
+def train_cox_elastic_net_regularization_paths(
+    X: np.ndarray,
+    times: np.ndarray,
+    events: np.ndarray,
+    alpha: float,
+    l1_ratio: float,
+    tol: float,
+    n_iter: int,
+) -> tuple[np.ndarray, float]:
+    """A direct implementation of,
+    Simon N, Friedman J, Hastie T, Tibshirani R. Regularization Paths for Cox's Proportional Hazards Model via Coordinate Descent. J Stat Softw. 2011 Mar;39(5):1-13. doi: 10.18637/jss.v039.i05. PMID: 27065756; PMCID: PMC4824408.
+    """
+    unique_times, time_return_inverse = _unique_with_return_inverse(times)
+    n_unique_times = len(unique_times)
+
+    n = X.shape[0]
+
+    col_indexes = np.arange(X.shape[1])
+
+    weights = np.zeros(X.shape[1])
+
+    last_loss = np.inf
+
+    for _ in range(n_iter):
+
+        new_weights = weights.copy()
+        n_hat = np.dot(X, new_weights)
+        n_hat_exp = np.exp(n_hat)
+        n_hat_exp_risk_set = np.flip(
+            np.cumsum(
+                np.flip(
+                    np.bincount(
+                        time_return_inverse, n_hat_exp, minlength=n_unique_times
+                    )
+                )
+            )
+        )[time_return_inverse]
+        wn = np.sum(
+            events
+            * (
+                (n_hat_exp * n_hat_exp_risk_set - n_hat_exp**2)
+                / (n_hat_exp_risk_set**2)
+            )
+        )
+        zn = n_hat + 1 / wn * (
+            events - np.sum(events * (n_hat_exp / n_hat_exp_risk_set))
+        )
+        for j in range(X.shape[1]):
+            not_j = col_indexes[col_indexes != j]
+            not_j_n = np.dot(X[:, not_j], new_weights[not_j])
+            left = 1 / n * np.sum(wn * X[:, j] * (zn - not_j_n))
+            right = alpha * l1_ratio
+            top = soft_threasholding_operator(left, right)
+            bottom = 1 / n * np.sum(wn * (X[:, j]) ** 2 + alpha * (1 - l1_ratio))
+            new_weights[j] = top / bottom
+
+        loss = get_breslow_neg_log_likelihood_with_elasticnet_penalty(
+            weights,
+            X,
+            events,
+            time_return_inverse,
+            n_unique_times,
+            alpha,
+            l1_ratio,
+            True,
+        )
+
+        if (loss - last_loss) > tol:
+            break
+        last_loss = loss
+        weights = new_weights
+
+    final_loss = get_breslow_neg_log_likelihood_with_elasticnet_penalty(
+        weights, X, events, time_return_inverse, n_unique_times, alpha, l1_ratio, False
+    )
+
+    return weights, final_loss
+
+
+get_breslow_neg_log_likelihood_with_elasticnet_penalty_with_left_censorship_signature = nb.types.float64(
+    nb.types.Array(nb.types.float64, 1, "C", False, aligned=True),
+    nb.types.Array(nb.types.float64, 2, "C", False, aligned=True),
+    nb.types.Array(nb.types.bool_, 1, "C", False, aligned=True),
+    nb.types.Array(nb.types.int64, 1, "C", False, aligned=True),
+    nb.types.Array(nb.types.int64, 1, "C", False, aligned=True),
+    nb.types.int64,
+    nb.types.float64,
+    nb.types.float64,
+    nb.types.bool_,
+)
+
+
+@nb.njit(
+    get_breslow_neg_log_likelihood_with_elasticnet_penalty_with_left_censorship_signature,
+    cache=True,
+)
+def get_breslow_neg_log_likelihood_with_elasticnet_penalty_with_left_censorship(
+    weights,
+    X,
+    event,
+    time_end_return_inverse,
+    time_start_return_inverse,
+    n_unique_times,
+    alpha,
+    l1_ratio,
+    scaled,
+):
+    n = float(X.shape[0])
+
+    p = np.dot(X, weights)
+    p_exp = np.exp(p)
+
+    risk_removed_at_time = np.bincount(
+        time_end_return_inverse, weights=p_exp, minlength=n_unique_times
+    )
+    risk_added_at_time = np.bincount(
+        time_start_return_inverse, weights=p_exp, minlength=n_unique_times
+    )
+    risk_at_time = np.cumsum(risk_added_at_time - risk_removed_at_time)
+    risk_set = risk_at_time[time_end_return_inverse - 1]
+
+    breslow_neg_log_likelihood = -np.sum(event * (p - np.log(risk_set)))
+
+    l1 = alpha * l1_ratio * np.abs(weights).sum()
+    l2 = 0.5 * alpha * (1.0 - l1_ratio) * np.square(weights).sum()
+    elasticnet_loss = l1 + l2
+
+    if scaled:
+        return (2 / n * breslow_neg_log_likelihood) + elasticnet_loss
+    else:
+        return breslow_neg_log_likelihood + elasticnet_loss
+
+
+train_cox_elastic_net_with_left_censorship_signature = nb.types.Tuple(
+    (
+        nb.types.Array(nb.types.float64, 1, "C", False, aligned=True),
+        nb.types.float64,
+    )
+)(
+    nb.types.Array(nb.types.float64, 2, "C", False, aligned=True),
+    nb.types.Array(nb.types.int64, 1, "C", False, aligned=True),
+    nb.types.Array(nb.types.int64, 1, "C", False, aligned=True),
+    nb.types.Array(nb.types.bool_, 1, "C", False, aligned=True),
+    nb.types.float64,
+    nb.types.float64,
+    nb.types.float64,
+    nb.types.int64,
+)
+
+
+@nb.njit(train_cox_elastic_net_with_left_censorship_signature, cache=True)
+def train_cox_elastic_net_with_left_censorship(
+    X: np.ndarray,
+    times: np.ndarray,
+    times_start: np.ndarray,
+    events: np.ndarray,
+    alpha: float,
+    l1_ratio: float,
+    tol: float,
+    n_iter: int,
+) -> tuple[np.ndarray, float]:
+    """A direct implementation of,
+    Simon N, Friedman J, Hastie T, Tibshirani R. Regularization Paths for Cox's Proportional Hazards Model via Coordinate Descent. J Stat Softw. 2011 Mar;39(5):1-13. doi: 10.18637/jss.v039.i05. PMID: 27065756; PMCID: PMC4824408.
+    """
+    all_times = np.concatenate((times_start, times))
+    unique_times, unique_times_return_inverse = _unique_with_return_inverse(all_times)
+
+    time_end_return_inverse = unique_times_return_inverse[len(times) :]
+    time_start_return_inverse = unique_times_return_inverse[: len(times_start)]
+
+    n_unique_times = len(unique_times)
+
+    n = X.shape[0]
+
+    col_indexes = np.arange(X.shape[1])
+
+    weights = np.zeros(X.shape[1])
+
+    last_loss = np.inf
+
+    for _ in range(n_iter):
+
+        new_weights = weights.copy()
+        n_hat = np.dot(X, new_weights)
+
+        n_hat_exp = np.exp(n_hat)
+
+        # n_hat_exp_risk_set = np.flip(
+        #     np.cumsum(
+        #         np.flip(
+        #             np.bincount(
+        #                 time_return_inverse, n_hat_exp, minlength=n_unique_times
+        #             )
+        #         )
+        #     )
+        # )[time_return_inverse]
+
+        risk_removed_at_time = np.bincount(
+            time_end_return_inverse, weights=n_hat_exp, minlength=n_unique_times
+        )
+        risk_added_at_time = np.bincount(
+            time_start_return_inverse, weights=n_hat_exp, minlength=n_unique_times
+        )
+        risk_at_time = np.cumsum(risk_added_at_time - risk_removed_at_time)
+        n_hat_exp_risk_set = risk_at_time[time_end_return_inverse - 1]
+
+        wn = np.sum(
+            events
+            * (
+                (n_hat_exp * n_hat_exp_risk_set - n_hat_exp**2)
+                / (n_hat_exp_risk_set**2)
+            )
+        )
+        zn = n_hat + 1 / wn * (
+            events - np.sum(events * (n_hat_exp / n_hat_exp_risk_set))
+        )
+        for j in range(X.shape[1]):
+            not_j = col_indexes[col_indexes != j]
+            not_j_n = np.dot(X[:, not_j], new_weights[not_j])
+            left = 1 / n * np.sum(wn * X[:, j] * (zn - not_j_n))
+            right = alpha * l1_ratio
+            top = soft_threasholding_operator(left, right)
+            bottom = 1 / n * np.sum(wn * (X[:, j]) ** 2 + alpha * (1 - l1_ratio))
+            new_weights[j] = top / bottom
+
+        loss = (
+            get_breslow_neg_log_likelihood_with_elasticnet_penalty_with_left_censorship(
+                weights,
+                X,
+                events,
+                time_end_return_inverse,
+                time_start_return_inverse,
+                n_unique_times,
+                alpha,
+                l1_ratio,
+                True,
+            )
+        )
+
+        if (loss - last_loss) > tol:
+            break
+        last_loss = loss
+        weights = new_weights
+
+    final_loss = (
+        get_breslow_neg_log_likelihood_with_elasticnet_penalty_with_left_censorship(
+            weights,
+            X,
+            events,
+            time_end_return_inverse,
+            time_start_return_inverse,
+            n_unique_times,
+            alpha,
+            l1_ratio,
+            False,
+        )
+    )
+
+    return weights, final_loss