torchsurv.stats.kaplan_meier

Classes

KaplanMeierEstimator()

Kaplan-Meier estimate of survival or censoring distribution for right-censored data [KM58].

class torchsurv.stats.kaplan_meier.KaplanMeierEstimator

Kaplan-Meier estimate of survival or censoring distribution for right-censored data [KM58].

__call__(event: tensor, time: tensor, censoring_dist: bool = False, check: bool = True)

Initialize Kaplan Meier estimator.

Parameters:

• event (torch.tensor, bool) – Event indicator of size n_samples (= True if event occured).

• time (torch.tensor, float) – Time-to-event or censoring of size n_samples.

• censoring_dist (bool, optional) – If False, returns the Kaplan-Meier estimate of the survival distribution. If True, returns the Kaplan-Meier estimate of the censoring distribution. Defaults to False.

• check (bool) – Whether to perform input format checks. Enabling checks can help catch potential issues in the input data. Defaults to True.

Examples

>>> _ = torch.manual_seed(42)
>>> n = 32
>>> time = torch.randint(low=0, high=8, size=(n,)).float()
>>> event = torch.randint(low=0, high=2, size=(n,)).bool()
>>> s = KaplanMeierEstimator() # estimate survival distribution
>>> s(event, time)
>>> s.km_est
tensor([1.0000, 1.0000, 0.8214, 0.7143, 0.6391, 0.6391, 0.5113, 0.2556])
>>> c = KaplanMeierEstimator() # estimate censoring distribution
>>> c(event, time, censoring_dist = True)
>>> c.km_est
tensor([0.9688, 0.8750, 0.8750, 0.8312, 0.6357, 0.4890, 0.3667, 0.0000])

References:

[KM58] (1,2,3)

E. L. Kaplan and Paul Meier. Nonparametric estimation from incomplete observations. Journal of the American Statistical Association, 53(282):457–481, June 1958.

plot_km(ax=None, **kwargs)

Plot the Kaplan-Meier estimate of the survival distribution.

Parameters:

• ax (matplotlib.axes.Axes, optional) – The axes to plot the Kaplan-Meier estimate. If None, a new figure and axes are created. Defaults to None.

• **kwargs – Additional keyword arguments to pass to the plot function.

Examples

>>> _ = torch.manual_seed(42)
>>> n = 32
>>> time = torch.randint(low=0, high=8, size=(n,)).float()
>>> event = torch.randint(low=0, high=2, size=(n,)).bool()
>>> km = KaplanMeierEstimator()
>>> km(event, time)
>>> km.plot_km()

predict(new_time: Tensor) → Tensor

Predicts the Kaplan-Meier estimate on new time points. If the new time points do not match any times used to fit, the left-limit is used.

Parameters:

new_time (torch.tensor) – New time points at which to predict the Kaplan-Meier estimate.

Returns:

Kaplan-Meier estimate evaluated at new_time.

Examples

>>> _ = torch.manual_seed(42)
>>> n = 8
>>> time = torch.randint(low=1, high=10, size=(n * 4,)).float()
>>> event = torch.randint(low=0, high=2, size=(n * 4,)).bool()
>>> km = KaplanMeierEstimator()
>>> km(event, time)
>>> km.predict(torch.randint(low=0, high=10, size=(n,))) # predict survival distribution
tensor([1.0000, 0.9062, 0.8700, 1.0000, 0.9062, 0.9062, 0.4386, 0.0000])

print_survival_table()

Prints the survival table with the unique times and Kaplan-Meier estimates.

Examples

>>> _ = torch.manual_seed(42)
>>> n = 32
>>> time = torch.randint(low=0, high=8, size=(n,)).float()
>>> event = torch.randint(low=0, high=2, size=(n,)).bool()
>>> s = KaplanMeierEstimator()