Skip to content

tausurv.discretization

time_grid(event_time: ArrayLike, event_indicator: ArrayLike, n_bins: int = 20) -> NDArray[np.float64]

Choose right bin edges for discrete-time models from training data.

Interior edges sit at quantiles of the uncensored event times, so each bin captures roughly the same number of events — the spacing that keeps per-bin hazard estimates stable. The last edge is the largest observed time, event or censored, so the grid spans the training data and no observation falls beyond the final bin.

Edge tkt_k closes the interval (tk−1,tk](t_{k-1}, t_k]; bin_index places observed times under this convention. Pass the returned grid both as time_bins to the discrete-time losses and to model.set_time_grid so training and prediction agree.

  • event_time — (n,) array — Observed time Y=min⁡(T,C)Y = \min(T, C).
  • event_indicator — (n,) array — δ=0\delta = 0 if censored; any positive value counts as an event, so competing-risks cause labels can be passed as-is.
  • n_bins — int, default 20 — Number of edges requested. Tied quantiles are collapsed, so the grid can come back shorter — size the model to the grid (n_bins=len(grid)), not the other way around.
  • (K,) float array — Strictly increasing right bin edges, K <= n_bins, ending at max(event_time).

Kvamme & Borgan (2021) place the edges at quantiles of the Kaplan-Meier estimate instead, which spaces them evenly in survival probability.

Kvamme, H., Borgan, Ø. (2021). Continuous and discrete-time survival prediction with neural networks. Lifetime Data Analysis, 27(4).

bin_index(time_grid: ArrayLike, times: ArrayLike) -> NDArray[np.intp]

Map times onto a grid of right bin edges.

Bin kk covers (tk−1,tk](t_{k-1}, t_k]: a time equal to an edge belongs to the bin that edge closes, and anything at or below the first edge lands in bin 0. Times beyond the last edge return KK — one past the final bin — leaving the caller to fold them into the last bin (.clip(max=K - 1), what the discrete-time losses do) or treat them as beyond the model’s horizon.

  • time_grid — (K,) array — Right bin edges, sorted ascending, e.g. from time_grid.
  • times — array — Times to place, any shape.
  • integer array — Bin indices in 0..K, same shape as times.