Targets and horizons API

survfm.pseudo_rmst.jackknife_pseudo_rmst_fast_exact(time, event, tau, *, clip=True, batch_size=256)

Compute exact, memory-bounded leave-one-out pseudo-RMST targets.

This evaluates the same full jackknife definition used by the submitted benchmark, but computes the full event-time grid once and evaluates delete-one curves in bounded patient batches. It is not an influence- function, subsampled, grouped, split, or cross-fitted approximation.

The submitted Kaplan-Meier convention keeps observations with time >= event_time in the risk set, including censoring tied with an event. Events at tau do not alter the integral over [0, tau]. Final clipping to [0, tau] defaults to the submitted behavior.

Source code in survfm/pseudo_rmst.py

def jackknife_pseudo_rmst_fast_exact(
    time,
    event,
    tau: float,
    *,
    clip: bool = True,
    batch_size: int = 256,
) -> np.ndarray:
    """Compute exact, memory-bounded leave-one-out pseudo-RMST targets.

    This evaluates the same full jackknife definition used by the submitted
    benchmark, but computes the full event-time grid once and evaluates
    delete-one curves in bounded patient batches. It is not an influence-
    function, subsampled, grouped, split, or cross-fitted approximation.

    The submitted Kaplan-Meier convention keeps observations with
    ``time >= event_time`` in the risk set, including censoring tied with an
    event. Events at ``tau`` do not alter the integral over ``[0, tau]``.
    Final clipping to ``[0, tau]`` defaults to the submitted behavior.
    """

    time_arr, event_arr, tau = _validated_fast_exact_inputs(time, event, tau)
    if not isinstance(batch_size, (int, np.integer)) or batch_size <= 0:
        raise ValueError("batch_size must be a positive integer")

    n = len(time_arr)
    all_event_times, all_event_counts = np.unique(
        time_arr[event_arr == 1], return_counts=True
    )
    before_tau = all_event_times < tau
    event_times = all_event_times[before_tau].astype(np.float64, copy=False)
    event_counts = all_event_counts[before_tau].astype(np.float64, copy=False)

    if len(event_times) == 0:
        raw = np.full(n, tau, dtype=np.float64)
        return np.clip(raw, 0.0, tau) if clip else raw

    risk_counts = np.array(
        [np.count_nonzero(time_arr >= event_time) for event_time in event_times],
        dtype=np.float64,
    )
    survival_full = np.cumprod(1.0 - event_counts / risk_counts)
    interval_widths = np.diff(np.concatenate((event_times, [tau])))
    rmst_full = event_times[0] + float(interval_widths @ survival_full)

    raw = np.empty(n, dtype=np.float64)
    event_times_row = event_times[None, :]
    risk_counts_row = risk_counts[None, :]
    event_counts_row = event_counts[None, :]

    for start in range(0, n, int(batch_size)):
        stop = min(start + int(batch_size), n)
        batch_time = time_arr[start:stop, None]
        batch_event = event_arr[start:stop, None]

        deleted_risk = risk_counts_row - (batch_time >= event_times_row)
        deleted_events = event_counts_row - (
            (batch_event == 1) & (batch_time == event_times_row)
        )

        # A full-grid time can disappear after deletion. Its neutral factor is 1.
        hazard = np.zeros_like(deleted_risk, dtype=np.float64)
        np.divide(deleted_events, deleted_risk, out=hazard, where=deleted_risk > 0)
        survival_deleted = np.cumprod(1.0 - hazard, axis=1)
        rmst_deleted = event_times[0] + survival_deleted @ interval_widths
        raw[start:stop] = n * rmst_full - (n - 1) * rmst_deleted

    return np.clip(raw, 0.0, tau) if clip else raw

survfm.pseudo_rmst.jackknife_pseudo_rmst(time_or_frame, event=None, tau=None, *, clip=True)

Construct jackknife pseudo-observation RMST targets.

Parameters:

Name	Type	Description	Default
time_or_frame		Either a sequence of observed times or a DataFrame with time and event columns.	required
event		Binary event indicators when time_or_frame is an array.	None
tau	float | None	Restriction horizon. If the first argument is a DataFrame, tau must be passed by keyword or as the second positional argument using jackknife_pseudo_rmst(df, tau=...).	None
clip	bool	If true, clip pseudo-targets to [0, tau]. This matches the manuscript benchmark implementation.	True

Source code in survfm/pseudo_rmst.py

def jackknife_pseudo_rmst(time_or_frame, event=None, tau: float | None = None, *, clip: bool = True) -> np.ndarray:
    """Construct jackknife pseudo-observation RMST targets.

    Parameters
    ----------
    time_or_frame:
        Either a sequence of observed times or a DataFrame with ``time`` and
        ``event`` columns.
    event:
        Binary event indicators when ``time_or_frame`` is an array.
    tau:
        Restriction horizon. If the first argument is a DataFrame, ``tau`` must
        be passed by keyword or as the second positional argument using
        ``jackknife_pseudo_rmst(df, tau=...)``.
    clip:
        If true, clip pseudo-targets to ``[0, tau]``. This matches the manuscript
        benchmark implementation.
    """

    if tau is None:
        if event is None:
            raise ValueError("tau must be provided.")
        if np.isscalar(event):
            tau = float(event)
            event = None
        else:
            raise ValueError("tau must be provided when event indicators are passed.")

    time_arr, event_arr = _extract_time_event(time_or_frame, event)
    tau = float(tau)
    n = len(time_arr)
    if n < 2:
        raise ValueError("At least two observations are required for jackknife pseudo-RMST.")

    rmst_all = km_rmst(time_arr, event_arr, tau)
    pseudo = np.empty(n, dtype=float)
    for i in range(n):
        mask = np.ones(n, dtype=bool)
        mask[i] = False
        rmst_minus_i = km_rmst(time_arr[mask], event_arr[mask], tau)
        pseudo[i] = n * rmst_all - (n - 1) * rmst_minus_i

    if clip:
        pseudo = np.clip(pseudo, 0.0, tau)
    return pseudo

survfm.pseudo_rmst.km_rmst(time, event, tau)

Compute Kaplan-Meier restricted mean survival time up to tau.

Source code in survfm/pseudo_rmst.py

def km_rmst(time, event, tau: float) -> float:
    """Compute Kaplan-Meier restricted mean survival time up to ``tau``."""

    if not np.isfinite(tau) or tau <= 0:
        raise ValueError("tau must be a positive finite value.")
    time_arr, event_arr = _validate_time_event(time, event)
    times, survival = kaplan_meier_curve(time_arr, event_arr)
    grid = np.unique(np.concatenate([times[times < tau], [float(tau)]]))
    if grid[0] != 0.0:
        grid = np.insert(grid, 0, 0.0)

    rmst = 0.0
    for left, right in zip(grid[:-1], grid[1:]):
        if right <= left:
            continue
        idx = np.searchsorted(times, left, side="right") - 1
        idx = max(idx, 0)
        rmst += (right - left) * survival[idx]
    return float(np.clip(rmst, 0.0, tau))

survfm.horizons.select_rmst_horizon(time, event, *, quantile=0.8, min_events=5)

Select tau from a quantile of observed training-fold event times.

Only the training outcomes supplied to this function are used. This helper implements the benchmark convention and is intended for generic benchmark use. In a disease-specific application, tau should preferably be prespecified from the clinical decision horizon and follow-up support.

Source code in survfm/horizons.py

def select_rmst_horizon(
    time,
    event,
    *,
    quantile: float = 0.8,
    min_events: int = 5,
) -> HorizonSelection:
    """Select ``tau`` from a quantile of observed training-fold event times.

    Only the training outcomes supplied to this function are used. This helper
    implements the benchmark convention and is intended for generic benchmark
    use. In a disease-specific application, ``tau`` should preferably be
    prespecified from the clinical decision horizon and follow-up support.
    """

    time_arr = np.asarray(time, dtype=float).reshape(-1)
    event_arr = np.asarray(event).reshape(-1)
    if len(time_arr) != len(event_arr) or len(time_arr) == 0:
        raise ValueError("time and event must be non-empty and have equal length.")
    if not np.all(np.isfinite(time_arr)) or np.any(time_arr < 0):
        raise ValueError("time must contain finite, non-negative values.")
    if not np.isin(event_arr, [0, 1, False, True]).all():
        raise ValueError("event must be binary with 1=event and 0=censored.")
    if not 0 < float(quantile) <= 1:
        raise ValueError("quantile must lie in (0, 1].")
    if int(min_events) < 1:
        raise ValueError("min_events must be at least 1.")

    event_times = time_arr[np.asarray(event_arr, dtype=bool)]
    if len(event_times) < int(min_events):
        raise ValueError(
            f"At least {int(min_events)} observed training events are required; "
            f"received {len(event_times)}."
        )

    tau = float(np.quantile(event_times, float(quantile)))
    if not np.isfinite(tau) or tau <= 0:
        raise ValueError("The selected horizon is not positive and finite.")
    return HorizonSelection(
        tau=tau,
        quantile=float(quantile),
        n_observations=int(len(time_arr)),
        n_events=int(len(event_times)),
        minimum_event_time=float(np.min(event_times)),
        maximum_event_time=float(np.max(event_times)),
    )

survfm.horizons.HorizonSelection dataclass

Audit record for a training-derived RMST horizon.

Source code in survfm/horizons.py

@dataclass(frozen=True)
class HorizonSelection:
    """Audit record for a training-derived RMST horizon."""

    tau: float
    quantile: float
    n_observations: int
    n_events: int
    minimum_event_time: float
    maximum_event_time: float

    def to_dict(self) -> dict[str, float | int]:
        return asdict(self)