API reference

class pytmle.PyTMLE(data: pandas.DataFrame, col_event_times: str = 'event_time', col_event_indicator: str = 'event_indicator', col_group: str = 'group', target_times: List[float] | None = None, g_comp: bool = True, evalues_benchmark: bool = False, key_1: int = 1, key_0: int = 0, initial_estimates: Dict[int, InitialEstimates] | None = None, verbose: int = 2, mlflow_logging: bool = False)
__init__(data: pandas.DataFrame, col_event_times: str = 'event_time', col_event_indicator: str = 'event_indicator', col_group: str = 'group', target_times: List[float] | None = None, g_comp: bool = True, evalues_benchmark: bool = False, key_1: int = 1, key_0: int = 0, initial_estimates: Dict[int, InitialEstimates] | None = None, verbose: int = 2, mlflow_logging: bool = False)

Initialize the PyTMLE class.

Parameters:

  • data (pd.DataFrame) – The input data containing event times, event indicators, treatment group information, and predictors.

  • col_event_times (str, optional) – The column name in the data that contains event times. Default is “event_time”.

  • col_event_indicator (str, optional) – The column name in the data that contains event indicators. Needs to contain consecutive integers starting at 0 (where 0 is for censored). Default is “event_indicator”.

  • col_group (str, optional) – The column name in the data that contains treatment group information. Needs to be binary. Default is “group”.

  • target_times (Optional[List[float]], optional) – Specific times at which to estimate the target parameter. If None, estimates for the last observed event time are used. Default is None.

  • g_comp (bool, optional) – Whether to store g-computation for initial estimates. Default is True.

  • evalues_benchmark (bool, optional) – Whether to compute E-values for measured confounders. Default is False.

  • key_1 (int, optional) – The key representing the treatment group. Default is 1.

  • key_0 (int, optional) – The key representing the control group. Default is 0.

  • initial_estimates (Optional[Dict[int, InitialEstimates]], optional) – Dict with pre-computed initial estimates for the two potential outcomes, which can be passed right to the second TMLE stage. If pre-computed hazards are given, make sure they have three dimensions and the last dimension corresponds to the number of non-zero elements in col_event_indicator. Default is None.

  • verbose (int, optional) – Verbosity level. 0: Absolutely so logging at all, 1: only warnings, 2: major execution steps, 3: execution steps, 4: everything for debugging. Default is 2.

  • mlflow_logging (bool, optional) – Whether to log the model fitting process to MLflow. Tracking URI and experiment should be set outside of the package, otherwise will log to default experiment in ./mlruns. Default is False.

fit(cv_folds: int = 10, max_updates: int = 500, min_nuisance: float | None = None, one_step_eps: float = 0.1, save_models: bool = False, alpha: float = 0.05, bootstrap: bool = False, n_bootstrap: int = 100, n_jobs: int = 4, stratified_bootstrap: bool = False, models=None, labtrans=None, propensity_score_models=None, propensity_score_calibration_method: Literal['isotonic', 'sigmoid'] | None = None, additional_inputs: Tuple | None = None, n_epochs: int = 100, batch_size: int = 128)

Fit the TMLE model.

Parameters:

  • cv_folds (int, optional) – Number of cross-validation folds for the initial estimate models. Default is 10.

  • max_updates (int) – Maximum number of updates to the estimates in the TMLE loop. Default is 500.

  • min_nuisance (Optional[float], optional) – Value between 0 and 1 for truncating the g-related denominator of the clever covariate. Default is None, which means no truncation at all.

  • one_step_eps (float) – Initial epsilon for the one-step update. Default is 0.1.

  • save_models (bool, optional) – Whether to store the models used for the initial estimates in a dictionary in self.models. Useful if you want to evaluate them later on. Default is False.

  • alpha (float, optional) – The alpha level for confidence intervals (relevant only for E-value benchmark). Default is 0.05.

  • bootstrap (bool, optional) – Whether to perform bootstrapping of the second TMLE stage for confidence intervals. Default is False.

  • n_bootstrap (int, optional) – Number of bootstrap samples. Has no effect if bootstrap is False. Default is 100.

  • n_jobs (int, optional) – Number of parallel jobs for bootstrapping. Has no effect if bootstrap is False. Default is 4.

  • stratified_bootstrap (bool, optional) – Whether to perform bootstrapping stratified by event indicator. Has no effect if bootstrap is False. Default is False.

  • models (Optional, optional) – A list of models to use for the state learner. If None, use the default library. Default is None.

  • labtrans (Optional, optional) – A list of labtrans objects to use for the risk model (if required; e.g., discretizer for DeepHit). If not None, needs to be one object for all models, or one object per model. Default is None.

  • propensity_score_models (Optional, optional) – A list of models to use for the propensity score stacking classifier. If None, use the default library. Default is None.

  • propensity_score_calibration_method (Optional[Literal["isotonic", "sigmoid"]], optional) – The calibration method to use for the propensity score model. If None, no calibration is performed. Default is None.

  • additional_inputs (Optional[Tuple], optional) – Additional inputs for the risk and censoring models. Can be tuple of torch.Tensors or np.ndarray, but has to be compatible with torchtuples. Default is None.

  • n_epochs (int, optional) – Number of epochs for training models in each cross fitting fold (if applicable). Default is 100.

  • batch_size (int, optional) – Batch size for training models in each cross fitting fold (if applicable). Default is 128.

plot(save_path: str | None = None, type: str = 'risks', alpha: float = 0.05, g_comp: bool = False, color_1: str | None = None, color_0: str | None = None, use_bootstrap: bool = False, only_converged: bool = False) → Tuple[Figure, ndarray] | None

Plot the counterfactual risks or average treatment effect (in terms of RR or RD).

Parameters:

  • save_path (Optional[str], optional) – Path to save the plot. If None, will return figure and axes. Default is None.

  • type (str, optional) – The type of prediction. “risks”, “rr” and “rd” are supported. Default is “risks”.

  • alpha (float, optional) – The alpha level for confidence intervals. Default is 0.05.

  • g_comp (bool, optional) – Whether to include the g-computation estimates in the plot. Default is False.

  • color_1 (Optional[str], optional) – Color for the potential outcome for “treated”. Pick None for standard matplotlib colors. Default is None.

  • color_0 (Optional[str], optional) – Color for the potential outcome for “untreated”. Pick None for standard matplotlib colors. Default is None.

  • use_bootstrap (bool, optional) – Whether to use the bootstrapped bounds instead of the theoretical bounds. Default is False.

  • only_converged (bool, optional) – Whether to plot only combinations of intervention/event/target time for which the TMLE update has converged. Default is False.

Returns:

The figure and axes of the plot. Only returned if save_path is None.

Return type:

Optional[Tuple[Figure, np.ndarray]]

plot_evalue_contours(save_dir_path: str | None = None, time: float | None = None, event: int | None = None, type: str = 'rr', use_bootstrap: bool = False, num_points_per_contour: int = 200, color_point_estimate: str = 'blue', color_ci: str = 'red', color_benchmarking: str = 'green', plot_size: Tuple[float, float] = (6.4, 4.8))

Plot the E-value contours for the estimated average treatment effect.

Parameters:

  • save_dir_path (Optional[str], optional) – Path to directory to save the plot. If None, will simply display the plot. Default is None.

  • time (Optional[float], optional) – Time at which to plot the E-value contours. If None, will plot for all target times. Default is None.

  • event (Optional[int], optional) – Event at which to plot the E-value contours. If None, will plot for all target events. Default is None.

  • type (str, optional) – The type of prediction. “rr” and “rd” are supported. Default is “rr”.

  • use_bootstrap (bool, optional) – Whether to use the bootstrapped bounds instead of the theoretical bounds. Default is False.

  • num_points_per_contour (int, optional) – Number of points per contour. Default is 200.

  • color_point_estimate (str, optional) – Color for the point estimate. Default is “blue”.

  • color_ci (str, optional) – Color for the confidence interval. Default is “red”.

  • color_benchmarking (str, optional) – Color for the benchmarking. Default is “green”.

  • plot_size (Tuple[float, float], optional) – Size of the plot. Default is (6.4, 4.8).

plot_norm_pn_eic(save_dir_path: str | None = None, plot_size: Tuple[float, float] = (6.4, 4.8))

Plot the norm of the empirical measure of the EIC.

Parameters:

  • save_dir_path (Optional[str], optional) – Path to directory to save the plot. If None, will simply display the plot. Default is None.

  • plot_size (Tuple[float, float], optional) – Size of the plot. Default is (6.4, 4.8).

plot_nuisance_weights(time: float | None = None, save_dir_path: str | None = None, color_1: str | None = None, color_0: str | None = None, plot_size: Tuple[float, float] = (6.4, 4.8)) → None

Plot the nuisance weights.

Parameters:

  • time (Optional[float], optional) – Time at which to plot the nuisance weights. If None, all target times are plotted. Default is None.

  • save_dir_path (Optional[str], optional) – Path to directory to save the plots. If None, will simply display the plots. Default is None.

  • color_1 (Optional[str], optional) – Color for the treatment group. Pick None for standard matplotlib colors. Default is None.

  • color_0 (Optional[str], optional) – Color for the control group. Pick None for standard matplotlib colors. Default is None.

plot_propensity_score_calibration(save_dir_path: str | None = None, plot_size: Tuple[float, float] = (6.4, 4.8), rolling_window_size: int = 50)

Plot the propensity score calibration (sorted propensity scores in comparison with empirically observed treatment probabilities).

Parameters:

  • save_dir_path (Optional[str], optional) – Path to directory to save the plot. If None, will simply display the plot. Default is None.

  • plot_size (Tuple[float, float], optional) – Size of the plot. Default is (6.4, 4.8).

  • rolling_window_size (int, optional) – Size of the rolling window for smoothing the observed treatment probabilities. Default is 50.

predict(type: str = 'risks', alpha: float = 0.05, g_comp: bool = False) → pandas.DataFrame

Predict the counterfactual risks or average treatment effect.

Parameters:

  • type (str, optional) – The type of prediction. “risks”, “rr” and “rd” are supported. Default is “risks”.

  • alpha (float, optional) – The alpha level for confidence intervals. Default is 0.05.

  • g_comp (bool, optional) – Whether to return the g-computation estimates instead of the updated estimates. Default is False.

Returns:

The predicted counterfactual risks or average treatment effect.

Return type:

pd.DataFrame

class pytmle.InitialEstimates(times: ndarray, g_star_obs: ndarray, propensity_scores: ndarray | None = None, hazards: ndarray | None = None, event_free_survival_function: ndarray | None = None, censoring_survival_function: ndarray | None = None)

Class to store the initial estimates for the TMLE algorithm, including the observed treatment, propensity scores, hazards, event-free survival function, and censoring survival function. The initial estimates can be set to pre-computed values to be passed right to the second stage of the TMLE algorithm. If left empty, they will be computed in the first stage of the TMLE algorithm. Make sure that the initial estimates are compatible with each other in terms of dimensions (same number of rows (observations) and columns (time points)). This will be checked prior to starting the TMLE update loop.

times

Array of time points (have to be available for all time-to-event functions).

Type:

np.ndarray

g_star_obs

Observed treatment values (binary) (n_observations,).

Type:

np.ndarray

propensity_scores

Propensity scores (n_observations,).

Type:

Optional[np.ndarray]

hazards

Hazards per competing event (n_observations, times, n_events). n_events must correspond to the number of non-zero events in the col_event_indicator of the given data frame.

Type:

Optional[np.ndarray]

event_free_survival_function

Event-free survival function (n_observations, times).

Type:

Optional[np.ndarray]

censoring_survival_function

Censoring survival function (n_observations, times).

Type:

Optional[np.ndarray]