Estimator API

SurvFMRMSTRegressor dataclass

Bases: PseudoRMSTRegressor

SurvFM estimator using pseudo-RMST target construction.

Source code in survfm/model_wrappers.py

class SurvFMRMSTRegressor(PseudoRMSTRegressor):
    """SurvFM estimator using pseudo-RMST target construction."""

fit_survfm(X_train, time, event, tau=None, *, estimator=None, backbone='tabpfn', backbone_kwargs=None, preprocess='auto', tau_quantile=0.8, min_events_for_tau=5, clip_targets=True, clip_predictions=True)

Fit SurvFM in one call using pseudo-RMST targets.

Source code in survfm/model_wrappers.py

def fit_survfm(
    X_train,
    time,
    event,
    tau: float | None = None,
    *,
    estimator: object | None = None,
    backbone: str = "tabpfn",
    backbone_kwargs: dict | None = None,
    preprocess: str | object | None = "auto",
    tau_quantile: float = 0.8,
    min_events_for_tau: int = 5,
    clip_targets: bool = True,
    clip_predictions: bool = True,
) -> SurvFMRMSTRegressor:
    """Fit SurvFM in one call using pseudo-RMST targets."""

    model = SurvFMRMSTRegressor(
        estimator=estimator,
        backbone=backbone,
        backbone_kwargs={} if backbone_kwargs is None else dict(backbone_kwargs),
        tau=tau,
        tau_quantile=tau_quantile,
        min_events_for_tau=min_events_for_tau,
        preprocess=preprocess,
        clip_targets=clip_targets,
        clip_predictions=clip_predictions,
    )
    return model.fit(X_train, time, event)

Backbone registry

list_backbones()

Return the registered backbones in stable display order.

Source code in survfm/backbones.py

def list_backbones() -> list[BackboneInfo]:
    """Return the registered backbones in stable display order."""

    return list(_BACKBONES.values())

create_backbone(name='tabpfn', **kwargs)

Construct a registered regression backbone using lazy imports.

Source code in survfm/backbones.py

def create_backbone(name: str = "tabpfn", **kwargs):
    """Construct a registered regression backbone using lazy imports."""

    key = get_backbone_info(name).name
    if key == "tabpfn":
        try:
            from tabpfn import TabPFNRegressor
        except Exception as exc:  # pragma: no cover - optional dependency
            raise ImportError(
                "The default 'tabpfn' backbone requires TabPFN. Install the tabpfn extra."
            ) from exc
        os.environ.setdefault("TABPFN_DISABLE_TELEMETRY", "1")
        defaults = {
            "ignore_pretraining_limits": True,
            "n_estimators": 1,
            "show_progress_bar": False,
            "random_state": 20260709,
        }
        return _construct_with_supported_kwargs(TabPFNRegressor, {**defaults, **kwargs})
    if key == "tabicl":
        try:
            from tabicl import TabICLRegressor
        except Exception as exc:  # pragma: no cover - optional dependency
            raise ImportError(
                "The 'tabicl' backbone requires TabICL. Install the tabicl extra."
            ) from exc
        defaults = {
            "n_estimators": 8,
            "batch_size": 1,
            "allow_auto_download": True,
            "verbose": False,
            "random_state": 20260615,
        }
        return _construct_with_supported_kwargs(TabICLRegressor, {**defaults, **kwargs})
    if key == "tabdpt":
        return _TabDPTAdapter(**kwargs)
    if key == "tabh2o":
        return _TabH2OAdapter(**kwargs)
    if key == "mitra":
        return _MITRAAdapter(**kwargs)

    try:
        if key == "linear":
            from sklearn.linear_model import LinearRegression

            return LinearRegression(**kwargs)
        if key == "random_forest":
            from sklearn.ensemble import RandomForestRegressor

            defaults = {"n_estimators": 500, "min_samples_leaf": 3, "random_state": 42, "n_jobs": -1}
            return RandomForestRegressor(**{**defaults, **kwargs})
        if key == "gradient_boosting":
            from sklearn.ensemble import HistGradientBoostingRegressor

            defaults = {"random_state": 42}
            return HistGradientBoostingRegressor(**{**defaults, **kwargs})
    except Exception as exc:  # pragma: no cover - optional dependency
        raise ImportError(
            f"The {key!r} backbone requires scikit-learn. Install the sklearn extra."
        ) from exc
    raise AssertionError(f"Unhandled backbone registry entry: {key}")