survival

survival

Survival analysis support for ngboost-lightning.

Provides right-censored data handling, the CensoredLogScore scoring rule, and the Y_from_censored helper for creating structured target arrays.

Censoring convention: event=1 means the time is observed (uncensored), event=0 means right-censored (the true event time is unknown, but greater than the observed time).

CensoredLogScore

Censored negative log-likelihood scoring rule.

For uncensored observations (event=1): loss = -logpdf(T)

For right-censored observations (event=0): loss = -logsf(T) (= -log(1 - CDF(T)))

The metric (Fisher Information) is the same as uncensored LogScore — this matches NGBoost's approach.

score

score(
    dist: Distribution, y: NDArray[void]
) -> NDArray[floating]

Per-sample censored NLL.

PARAMETER	DESCRIPTION
dist	Predicted distribution instance. TYPE: Distribution
y	Structured target array with 'Event' and 'Time' fields. TYPE: NDArray[void]

RETURNS	DESCRIPTION
NDArray[floating]	Censored NLL values, shape [n_samples].

Source code in ngboost_lightning/survival.py

def score(
    self,
    dist: Distribution,
    y: NDArray[np.void],
) -> NDArray[np.floating]:
    """Per-sample censored NLL.

    Args:
        dist: Predicted distribution instance.
        y: Structured target array with 'Event' and 'Time' fields.

    Returns:
        Censored NLL values, shape ``[n_samples]``.
    """
    E = y["Event"].astype(np.float64)
    T = y["Time"]
    uncens = dist.logpdf(T)
    cens = dist.logsf(T)
    return -(E * uncens + (1.0 - E) * cens)

d_score

d_score(
    dist: Distribution, y: NDArray[void]
) -> NDArray[floating]

Gradient of censored NLL w.r.t. distribution parameters.

Uses the uncensored analytical gradient for observed samples and numerical finite differences for censored samples. This avoids distribution-specific censored gradient derivations while remaining accurate.

PARAMETER	DESCRIPTION
dist	Predicted distribution instance. TYPE: Distribution
y	Structured target array with 'Event' and 'Time' fields. TYPE: NDArray[void]

RETURNS	DESCRIPTION
NDArray[floating]	Gradient array, shape [n_samples, n_params].

Source code in ngboost_lightning/survival.py

def d_score(
    self,
    dist: Distribution,
    y: NDArray[np.void],
) -> NDArray[np.floating]:
    """Gradient of censored NLL w.r.t. distribution parameters.

    Uses the uncensored analytical gradient for observed samples and
    numerical finite differences for censored samples. This avoids
    distribution-specific censored gradient derivations while remaining
    accurate.

    Args:
        dist: Predicted distribution instance.
        y: Structured target array with 'Event' and 'Time' fields.

    Returns:
        Gradient array, shape ``[n_samples, n_params]``.
    """
    E = y["Event"]
    T = y["Time"]
    n = len(T)
    n_params = dist._params.shape[1]

    # Uncensored gradient: standard d_score
    uncens_grad = dist.d_score(T)

    # Censored gradient via finite differences on -logsf
    eps = 1e-5
    cens_grad = np.zeros((n, n_params))
    for k in range(n_params):
        params_plus = dist._params.copy()
        params_plus[:, k] += eps
        params_minus = dist._params.copy()
        params_minus[:, k] -= eps

        logsf_plus = type(dist)(params_plus).logsf(T)
        logsf_minus = type(dist)(params_minus).logsf(T)
        cens_grad[:, k] = -(logsf_plus - logsf_minus) / (2.0 * eps)

    # Blend: E * uncens_grad + (1-E) * cens_grad
    E_f = E.astype(np.float64)[:, np.newaxis]
    return E_f * uncens_grad + (1.0 - E_f) * cens_grad

metric

metric(dist: Distribution) -> NDArray[floating]

Fisher Information matrix (same as uncensored).

PARAMETER	DESCRIPTION
dist	Predicted distribution instance. TYPE: Distribution

RETURNS	DESCRIPTION
NDArray[floating]	Metric tensor, shape [n_samples, n_params, n_params].

Source code in ngboost_lightning/survival.py

def metric(
    self,
    dist: Distribution,
) -> NDArray[np.floating]:
    """Fisher Information matrix (same as uncensored).

    Args:
        dist: Predicted distribution instance.

    Returns:
        Metric tensor, shape ``[n_samples, n_params, n_params]``.
    """
    return dist.metric()

natural_gradient

natural_gradient(
    dist: Distribution, y: NDArray[void]
) -> NDArray[floating]

Natural gradient for censored NLL.

PARAMETER	DESCRIPTION
dist	Predicted distribution instance. TYPE: Distribution
y	Structured target array with 'Event' and 'Time' fields. TYPE: NDArray[void]

RETURNS	DESCRIPTION
NDArray[floating]	Natural gradient, shape [n_samples, n_params].

Source code in ngboost_lightning/survival.py

def natural_gradient(
    self,
    dist: Distribution,
    y: NDArray[np.void],
) -> NDArray[np.floating]:
    """Natural gradient for censored NLL.

    Args:
        dist: Predicted distribution instance.
        y: Structured target array with 'Event' and 'Time' fields.

    Returns:
        Natural gradient, shape ``[n_samples, n_params]``.
    """
    grad = self.d_score(dist, y)
    fi = self.metric(dist)
    result: NDArray[np.floating] = np.linalg.solve(fi, grad[..., np.newaxis])[
        ..., 0
    ]
    return result

total_score

total_score(
    dist: Distribution,
    y: NDArray[void],
    sample_weight: NDArray[floating] | None = None,
) -> float

Weighted mean censored NLL.

PARAMETER	DESCRIPTION
dist	Predicted distribution instance. TYPE: Distribution
y	Structured target array with 'Event' and 'Time' fields. TYPE: NDArray[void]
sample_weight	Per-sample weights. TYPE: NDArray[floating] | None DEFAULT: None

RETURNS	DESCRIPTION
float	Scalar (weighted) mean censored NLL.

Source code in ngboost_lightning/survival.py

def total_score(
    self,
    dist: Distribution,
    y: NDArray[np.void],
    sample_weight: NDArray[np.floating] | None = None,
) -> float:
    """Weighted mean censored NLL.

    Args:
        dist: Predicted distribution instance.
        y: Structured target array with 'Event' and 'Time' fields.
        sample_weight: Per-sample weights.

    Returns:
        Scalar (weighted) mean censored NLL.
    """
    return float(np.average(self.score(dist, y), weights=sample_weight))

Y_from_censored

Y_from_censored(
    T: NDArray[floating],
    E: NDArray[integer] | NDArray[bool_],
) -> NDArray[void]

Create a structured target array from times and event indicators.

PARAMETER	DESCRIPTION
T	Times to event or censoring, shape [n_samples]. TYPE: NDArray[floating]
E	Event indicators, shape [n_samples]. E[i] = 1 (or True) means T[i] is an observed event time. E[i] = 0 (or False) means T[i] is a right-censored time. TYPE: NDArray[integer] | NDArray[bool_]

RETURNS	DESCRIPTION
NDArray[void]	Structured array with fields 'Event' (bool) and 'Time'
NDArray[void]	(float64), shape [n_samples].

Source code in ngboost_lightning/survival.py

def Y_from_censored(
    T: NDArray[np.floating],
    E: NDArray[np.integer] | NDArray[np.bool_],
) -> NDArray[np.void]:
    """Create a structured target array from times and event indicators.

    Args:
        T: Times to event or censoring, shape ``[n_samples]``.
        E: Event indicators, shape ``[n_samples]``.
            ``E[i] = 1`` (or True) means ``T[i]`` is an observed event time.
            ``E[i] = 0`` (or False) means ``T[i]`` is a right-censored time.

    Returns:
        Structured array with fields ``'Event'`` (bool) and ``'Time'``
        (float64), shape ``[n_samples]``.
    """
    T = np.asarray(T, dtype=np.float64)
    E = np.asarray(E, dtype=bool)
    n = len(T)
    Y = np.empty(n, dtype=SURVIVAL_DTYPE)
    Y["Event"] = E
    Y["Time"] = T
    return Y