surv_bart

bart_survival.surv_bart

High-level API for Bayesian Additive Regression Trees (BART) survival analysis.

Provides: - SurvBART (alias BartSurvModel): Estimator class with fit, predict, and serialization. - format_survival_data (alias get_surv_pre_train): Preprocess time-to-event data into pooled cross-sections. - format_prediction_data (alias get_posterior_test): Prepare test covariates across prediction time intervals. - compute_survival_curves (alias get_sv_prob): Cumulative product survival probability transformation. - compute_pdp (alias get_pdp): Partial dependence data generation. - compute_causal_pdp (alias get_causal_pdp): Causal G-Computation with support violation audit. - compute_hazard_ratio (alias pdp_rr_metric): Posterior Risk Ratio / Hazard Ratio extraction. - compute_pdp_difference (alias pdp_diff_metric): Posterior marginal survival difference extraction.

class bart_survival.surv_bart.BartSurvModel(trees: int = 50, draws: int = 2500, tune: int = 1000, chains: int = 4, cores: int | None = None, thin: int = 10, thin_idata: bool = False, use_sparsity: bool = False, split_method: str = 'empirical', max_cutpoints: int = 512, update_fraction: float = 1.0, min_node_size: int = 5, k: float = 2.0, alpha: float = 0.95, beta: float = 2.0, sparsity_warmup: int | None = None, unique_midpoints: bool | None = None, split_rules: List[Any] | None = None, split_prior: Any | None = None, model_config: Dict[str, Any] | None = None, sampler_config: Dict[str, Any] | None = None)

Bayesian Additive Regression Trees (BART) Survival Model (SurvBART).

Implements discrete-time probit hazard modeling using an ensemble of sum-of-trees. Supports warm-started Dirichlet sparsity selection (DART), empirical unique midpoint cutpoints (512 resolution), MCMC thinning, and support-audited Causal G-Computation.

bart_predict(X_pred: ndarray, coords: ndarray, size: int = None, rng: RandomState = None, **kwargs) → DataArray

Derives posterior predictions on updated dataset. Alternative method for re-loaded models.

Parameters:

X_pred (np.ndarray) – Covariate matrix in long-time format.
coords (np.ndarray) – Coordinates associated with long-time format.
size (int, optional) – Sets sample of posterior draws. Defaults to None.
rng (RandomState, optional) – Random number generator for reproducable results. Defaults to None.

Returns:

DataArray containing the predicted outputs from the model and covariate matrix.

Return type:

xr.DataArray

build_model(X: ndarray, y: ndarray, weights: ndarray | None = None, coords: ndarray | None = None, predictor_names: List[str] | None = None, **kwargs: Any) → Model

Builds the PyMC probabilistic model for BART discrete-time survival.

Parameters:

X (np.ndarray) – Covariate matrix in long-time form [time, x1, x2, …].
y (np.ndarray) – Event status in long-time form (0=censored/alive, 1=event).
weights (Optional[np.ndarray], optional) – Sample observation weights. Defaults to None.
coords (Optional[np.ndarray], optional) – Subject coordinate index alignment. Defaults to None.
predictor_names (Optional[List[str]], optional) – Variable names. Defaults to None.
**kwargs – Additional model parameters.

Returns:

The compiled PyMC model.

Return type:

pm.Model

close() → None: Explicitly releases background SyncManager server processes and frees IPC memory.

Estimates counterfactual treatment effects using Causal G-Computation and Modified Treatment Policies.

Supports both: 1. Binary / Categorical exposures: Audits conditional support for positivity violations (e.g. deterministic

zeros like LOS=0 for outpatients) and re-routes dependent features to conditional donor resampling.

Continuous exposures: Evaluates a counterfactual shift intervention (Modified Treatment Policy) comparing natural course (X) vs shifted exposure (X + delta). Alternatively evaluates fixed 2-point contrasts if explicit values=[v0, v1] are supplied.

Metrics computed: - Event-weighted summary Hazard Ratio (default, Cox-aligned via complementary log-log link). - Unweighted summary Hazard Ratio across follow-up intervals. - Time-varying Hazard Ratio curve HR(t) via complementary log-log link. - Time-varying discrete risk ratio curve RR_discrete(t) = h1(t) / h0(t) (not averaged). - Time-varying Risk Difference curve RD(t) = Risk_1(t) - Risk_0(t) = S_0(t) - S_1(t) (1 over 0). - Survival Difference curve S_1(t) - S_0(t) = -RD(t). - Event-weighted and unweighted summary Risk Differences. - Final horizon Risk Difference RD(t_end). - Restricted Mean Survival Time (RMST) difference.

Parameters:

X (Union[np.ndarray, pd.DataFrame]) – Baseline cohort covariates.
primary_col (Union[str, int], optional) – Primary exposure variable name or column index. Defaults to “inpatient”.
values (Optional[List[Union[int, float]]], optional) – Counterfactual intervention values. If None, auto-detects binary exposures ([0, 1]) or triggers continuous shift intervention.
delta (Optional[float], optional) – Shift increment for continuous exposures (X + delta). Defaults to 1.0.
shift (Optional[bool], optional) – Force continuous shift intervention mode. Defaults to None.
times (Optional[Union[np.ndarray, List[float]]], optional) – Evaluation time points. Defaults to None.
audit_support (bool, optional) – Whether to audit and flag support violations. Defaults to True.
donor_seed (Optional[int], optional) – Random seed for reproducible donor resampling. Defaults to 42.
event_weights (Optional[np.ndarray], optional) – Normalized interval weights for summary HR/RD. Defaults to model’s empirical event_weights_ if available.
clip_to_range (bool, optional) – Whether to clip continuous shifted values to observed maximum. Defaults to False.
**kwargs – Additional prediction keyword arguments.

Returns:

Dictionary containing weighted and unweighted HRs, RDs, RMST, curves, credible intervals, and draws.

Return type:

Dict[str, Any]

feature_importance(sort: bool = True) → Series

Extracts normalized variable split inclusion proportions from the posterior tree ensemble.

Parameters:: sort (bool, optional) – Whether to sort features in descending order of importance. Defaults to True.
Returns:: Feature importance series indexed by predictor names.
Return type:: pd.Series

Fit the BART survival model using MCMC sampling.

Parameters:

y (np.ndarray) – Binary event indicator in long-time pooled format (1=event, 0=censored/alive).
X (np.ndarray) – Covariate matrix in long-time format where column 0 is discrete time, followed by baseline covariates [time, x1, x2, …].
weights (Optional[np.ndarray], optional) – Sample observation weights (e.g. for case-cohort designs). Defaults to None.
coords (Optional[np.ndarray], optional) – Subject coordinate index alignment for observations. Defaults to None.
progressbar (bool, optional) – Whether to display MCMC progress bar. Defaults to True.
predictor_names (Optional[List[str]], optional) – Names of features in X matrix. Defaults to None.
random_seed (Optional[Union[int, RandomState]], optional) – Random seed for reproducibility. Defaults to None.
**kwargs – Additional sampling configuration overrides (e.g. draws, tune, cores, chains).

Returns:

Posterior inference data object containing sampled trees and chain statistics.

Return type:

az.InferenceData

Fits SurvBART directly from baseline tabular covariates and survival outcomes.

Eliminates manual boilerplate by automatically discretizing follow-up times and expanding tabular data into long-format person-time pooled cross-sections.

Parameters:

X (Union[np.ndarray, pd.DataFrame]) – Baseline covariates.
time (Union[np.ndarray, pd.Series, List[float]]) – Observed event or censoring times.
status (Union[np.ndarray, pd.Series, List[int]]) – Binary event indicators (1=event, 0=censored).
time_scale (Optional[float], optional) – Interval discretization scale factor (e.g. 5 or 30). Defaults to None.
weights (Optional[np.ndarray], optional) – Sample observation weights. Defaults to None.
progressbar (bool, optional) – Whether to display MCMC progress bar. Defaults to True.
random_seed (Optional[Union[int, RandomState]], optional) – Random seed. Defaults to None.
**kwargs – Additional sampling keyword overrides.

Returns:

Fitted posterior inference data object.

Return type:

az.InferenceData

static get_default_model_config() → Dict[str, Any]: Default model prior and tree ensemble configuration.

static get_default_sampler_config() → Dict[str, Any]: Default sampler and MCMC inference configuration.

property id: str

Generate a unique hash value for the model.

The hash value is created using the last 16 characters of the SHA256 hash encoding, based on the model configuration, version, and model type.

Returns:: A string of length 16 characters containing a unique hash of the model.
Return type:: str

classmethod load(fname: str, treename: str)

Loads a saved model.

Parameters:

fname (str) – Path to saved model object.
treename (str) – Path to saved tree object.

Returns:

Returns object of SurvBartModel class.

Return type:

_type_

predict(X_pred: ndarray, coords: ndarray | None = None, **kwargs: Any) → ndarray

Alias for sample_posterior_predictive following scikit-learn convention.

Parameters:

X_pred (np.ndarray) – Covariate matrix in long-time format [time, x1, x2, …].
coords (Optional[np.ndarray], optional) – Coordinates associated with subject indices. Defaults to None.
**kwargs – Additional prediction keyword options.

Returns:

Array of shape (n_draws, n_obs) containing predicted hazard probabilities.

Return type:

np.ndarray

predict_survival(X_test: ndarray | DataFrame, times: ndarray | List[float] | None = None, return_quantiles: bool = False, qntile: List[float] = [0.025, 0.975], **kwargs: Any) → ndarray | Dict[str, ndarray]

Calculates out-of-sample cumulative survival curves S(t) for new subjects.

Parameters:

X_test (Union[np.ndarray, pd.DataFrame]) – Baseline covariates for test cohort.
times (Optional[Union[np.ndarray, List[float]]], optional) – Discrete time evaluation grid. If None, uses unique event times observed during training. Defaults to None.
return_quantiles (bool, optional) – If True, returns mean and credible interval quantiles. If False, returns full posterior draws of shape (n_draws, n_subjects, n_times). Defaults to False.
qntile (List[float], optional) – Credible interval quantiles. Defaults to [0.025, 0.975].
**kwargs – Additional prediction keyword arguments.

Returns:

Full survival curve array or summary dict.

Return type:

Union[np.ndarray, Dict[str, np.ndarray]]

classmethod preset_default(**overrides: Any) → BartSurvModel: Default recommended preset (alias for preset_standard).

classmethod preset_fast(**overrides: Any) → BartSurvModel: Fast exploration preset: 25 trees, 500 draws, 250 tune, 4 chains, thin 5, update_fraction 1.0.

classmethod preset_high_precision(**overrides: Any) → BartSurvModel: High-precision publication preset: 100 trees, 5000 draws, 2000 tune, 4 chains, thin 10, update_fraction 1.0.

classmethod preset_production(**overrides: Any) → BartSurvModel: Legacy alias for preset_standard.

classmethod preset_standard(**overrides: Any) → BartSurvModel: Standard scientific preset: 50 trees, 1,000 draws, 250 tune, 4 chains, thin 4 (1,000 final draws), update_fraction 1.0, empirical cutpoints.

sample_model(**kwargs: Any) → DataTree

Initiates training/sampling of the model.

Returns:: Posterior data collected from training the model.
Return type:: az.InferenceData

sample_posterior_predictive(X_pred: ndarray, coords: ndarray | None = None, extend_idata: bool = False, **kwargs: Any) → ndarray

Derives posterior predictive hazard probabilities on new covariates.

Parameters:

X_pred (np.ndarray) – Covariate matrix in long-time format [time, x1, x2, …].
coords (Optional[np.ndarray], optional) – Coordinates associated with subject indices. Defaults to None.
extend_idata (bool, optional) – Legacy compatibility flag. Defaults to False.
**kwargs – Additional keyword arguments.

Returns:

Array of shape (n_draws, n_obs) containing predicted hazard probabilities.

Return type:

np.ndarray

save(idata_name: str, all_tree_name: str) → None

Saves a trained model and tree object.

Parameters:

idata_name (str) – Path to saving model object.
all_tree_name (str) – Path to saving tree object.

set_idata_attrs(idata: DataTree | None = None) → DataTree

Sets the additional information in the idata object.

Parameters:: idata (Optional[az.InferenceData], optional) – Idata. Defaults to None.
Returns:: Idata
Return type:: az.InferenceData

bart_survival.surv_bart.SurvBART: alias of BartSurvModel

bart_survival.surv_bart.compute_causal_pdp(x: DataFrame, primary_col: str | int = 'inpatient', values: list = [0, 1], audit_support: bool = True, donor_seed: int | None = 42) → Tuple[ndarray, Dict, List[str]]

Generates counterfactual data for Causal G-Computation with support violation audit.

Identifies dependent features or post-exposure mediators that have structural support or positivity violations with the primary exposure (e.g., deterministic zeros like LOS=0 when Inpatient=0). Performs conditional donor resampling for dependent features while leaving baseline confounders intact, strictly avoiding off-support extrapolation.

Parameters:

x (pd.DataFrame) – Baseline covariate dataframe.
primary_col (Union[str, int]) – Primary exposure column name or index.
values (list, optional) – Counterfactual intervention values. Defaults to [0, 1].
audit_support (bool, optional) – Whether to audit and flag support violations. Defaults to True.
donor_seed (int, optional) – Random seed for reproducible donor resampling. Defaults to 42.

Returns:

out_sk: Stacked counterfactual covariate matrix.
c_idx: Index coordinate dictionary compatible with get_posterior_test.
dependent_features: List of audited features exhibiting support violations.

Return type:

Tuple[np.ndarray, Dict, List[str]]

bart_survival.surv_bart.compute_hazard_and_risk_difference(pdp_val: Dict[str, ndarray], idx: int, event_weights: ndarray | None = None, qntile: List[float] = [0.025, 0.975], eps: float = 1e-07) → Dict[str, Any]

Computes event-weighted and unweighted Hazard Ratios and Risk Differences from counterfactual predictions.

Contrasts Intervention Arm 1 (obs >= idx) over Reference Arm 0 (obs < idx).

Calculates: - Event-weighted summary Hazard Ratio (default, Cox-aligned via complementary log-log link). - Unweighted summary Hazard Ratio across follow-up intervals. - Time-varying Hazard Ratio curve HR(t) via complementary log-log link. - Time-varying discrete risk ratio curve RR_discrete(t) = h1(t) / h0(t) (not averaged). - Time-varying Risk Difference curve RD(t) = Risk_1(t) - Risk_0(t) = S_0(t) - S_1(t) (1 over 0). - Survival Difference curve S_1(t) - S_0(t) = -RD(t). - Event-weighted and unweighted summary Risk Differences. - Final horizon Risk Difference RD(t_end). - Restricted Mean Survival Time (RMST) difference.

Parameters:

pdp_val (Dict[str, np.ndarray]) – Dictionary containing ‘prob’ (hazards) and ‘sv’ (survival curves). Arrays have shape (n_draws, 2 * N, n_times).
idx (int) – Split index demarcating Reference Arm 0 (:idx) and Intervention Arm 1 (idx:).
event_weights (Optional[np.ndarray], optional) – Normalized event weights per interval. If None, defaults to uniform interval weighting.
qntile (List[float], optional) – Credible interval quantiles. Defaults to [0.025, 0.975].
eps (float, optional) – Epsilon threshold for numerical safety. Defaults to 1e-7.

Returns:

Comprehensive dictionary containing summary statistics, credible intervals, and draw arrays.

Return type:

Dict[str, Any]

bart_survival.surv_bart.compute_hazard_ratio(pdp_val: Dict[str, ndarray], idx: int, qntile: List[float] = [0.025, 0.975]) → Dict[str, ndarray]

Generates posterior Hazard Ratio / Risk Ratio between counterfactual intervention groups.

Parameters:

pdp_val (Dict[str, np.ndarray]) – Posterior output dictionary containing ‘prob’.
idx (int) – Split index demarcating counterfactual groups in stacked PDP matrix.
qntile (List[float], optional) – Quantiles for credible interval. Defaults to [0.025, 0.975].

Returns:

‘rr_m’: Posterior mean hazard ratio curve across time points.
’rr_q’: Posterior credible interval quantiles of the hazard ratio curve.

Return type:

Dict[str, np.ndarray]

bart_survival.surv_bart.compute_pdp(x: ndarray | DataFrame, var_col: List[int | str] = [], values: List[int | float] = [], qt: List[float] = [0.25, 0.5, 0.75], sample_n=None, method: str = 'marginal', audit_support: bool = False) → Tuple[ndarray, Dict] | Tuple[ndarray, Dict, List[str]]

Generates data for Partial Dependency Plots or Causal G-Computation.

Parameters:

x (Union[np.ndarray, pd.DataFrame]) – Covariate matrix or DataFrame.
var_col (List[Union[int, str]], optional) – Covariate column to generate pdp values. Defaults to [].
values (List[Union[int, float]], optional) – Values to test on for each covariate. Defaults to [].
qt (List[float], optional) – Quantiles to generate values on if none are given. Defaults to [0.25, 0.5, 0.75].
sample_n (int, optional) – Sample size of large dataset. Useful for working with large datasets. Defaults to None.
method (str, optional) – ‘marginal’ for standard PDP, or ‘causal’ for support-audited G-Computation. Defaults to ‘marginal’.
audit_support (bool, optional) – If True, audits for positivity/support violations. Defaults to False.

Returns:

PDP covariate matrix and tracking dictionary.

Return type:

Union[Tuple[np.ndarray, Dict], Tuple[np.ndarray, Dict, List[str]]]

bart_survival.surv_bart.compute_pdp_difference(pdp_val: Dict[str, ndarray], idx: int, qntile: List[float] = [0.025, 0.975]) → Dict[str, ndarray]

Generates posterior estimate of marginal survival difference between counterfactual groups.

Parameters:

pdp_val (Dict[str, np.ndarray]) – Posterior output dictionary containing ‘sv’.
idx (int) – Split index demarcating counterfactual groups in stacked PDP matrix.
qntile (List[float], optional) – Quantiles for credible interval. Defaults to [0.025, 0.975].

Returns:

‘diff_m’: Posterior mean survival difference curve across time points.
’diff_q’: Posterior credible interval quantiles of the difference curve.

Return type:

Dict[str, np.ndarray]

bart_survival.surv_bart.compute_shift_pdp(x: DataFrame | ndarray, primary_col: str | int, delta: float = 1.0, clip_to_range: bool = False) → Tuple[ndarray, Dict[str, Any]]

Generates counterfactual datasets for a continuous shift intervention (Modified Treatment Policy).

Contrasts the natural course cohort (Arm 0: X) against a cohort where every individual’s continuous exposure is shifted by delta units (Arm 1: X + delta).

Parameters:

x (Union[pd.DataFrame, np.ndarray]) – Baseline covariate matrix or DataFrame.
primary_col (Union[str, int]) – Continuous variable name or column index to shift.
delta (float, optional) – Increment added to primary_col. Defaults to 1.0.
clip_to_range (bool, optional) – Whether to clip shifted values to observed maximum. Defaults to False.

Returns:

out_sk: Stacked counterfactual matrix of shape (2 * N, n_features).
c_idx: Index coordinate dictionary demarcating Arm 0 and Arm 1.

Return type:

Tuple[np.ndarray, Dict[str, Any]]

bart_survival.surv_bart.compute_survival_curves(post: ndarray | DataArray | Dataset, coords: ndarray) → Dict[str, ndarray]

Generates discrete-time hazard and cumulative survival probability estimates.

Computes S(t) = prod_{s=1}^t (1 - h(s)) across all posterior draws and subjects.

Parameters:

post (Union[np.ndarray, xr.DataArray, xr.Dataset]) – Posterior predictive hazard draws.
coords (np.ndarray) – Subject coordinate indices to reshape predictions per subject and time.

Returns:

‘prob’: Hazard probability array of shape (n_draws, n_subjects, n_time_points).
’sv’: Cumulative survival probability array of shape (n_draws, n_subjects, n_time_points).

Return type:

Dict[str, np.ndarray]

bart_survival.surv_bart.compute_survival_quantiles(sv: ndarray, msk: ndarray, draws: bool = True, qntile: List[float] = [0.025, 0.975]) → Dict[str, ndarray]

Generates mean and quantile estimates of Survival Probabilities.

Parameters:

sv (np.ndarray) – Survival probabilities of shape (n_draws, n_obs, n_times).
msk (np.ndarray) – Boolean mask selecting subpopulation of interest.
draws (bool, optional) – If True, aggregates across posterior draws. Defaults to True.
qntile (List[float], optional) – Credible interval quantiles. Defaults to [0.025, 0.975].

Returns:

Mean and quantile curves for masked and unmasked groups.

Return type:

Dict[str, np.ndarray]

bart_survival.surv_bart.format_prediction_data(y_time: ndarray, y_status: ndarray, x: ndarray, time_scale: float | None = None) → Dict[str, ndarray]

Generates evaluation grid in long-time format for out-of-sample posterior predictions.

Constructs counterfactual cross-product of discrete evaluation times and subject covariates.

Parameters:

y_time (np.ndarray) – Follow-up times defining the discrete time evaluation grid.
y_status (np.ndarray) – Event indicators (used for historical reference).
x (np.ndarray) – Baseline covariate matrix to evaluate.
time_scale (Optional[float], optional) – Time discretization scale. Defaults to None.

Returns:

‘post_x’: Evaluated covariate matrix with time in column 0.
’coords’: Subject coordinate identifier array.

Return type:

Dict[str, np.ndarray]

bart_survival.surv_bart.format_survival_data(y_time: ndarray, y_status: ndarray, x: ndarray, weight: ndarray | None = None, time_scale: float | None = None) → Dict[str, ndarray]

Generates pooled cross-sectional long-format dataset for discrete-time survival training.

Transforms right-censored time-to-event outcomes into person-time binary event observations. Each individual contributes records for discrete intervals $t = 1, 2, dots, T_i$.

Parameters:

y_time (np.ndarray) – Array of observed follow-up times.
y_status (np.ndarray) – Binary event indicator (1 if event occurred, 0 if right-censored).
x (np.ndarray) – Baseline covariate matrix of shape (n_individuals, n_features).
weight (Optional[np.ndarray], optional) – Observation weights (e.g. for case-cohort designs). Defaults to 1.0.
time_scale (Optional[float], optional) – Time discretization scale (e.g. 5 for 5-day intervals). Defaults to None.

Returns:

‘y’: Binary event indicator array of shape (n_person_time, 1).
’x’: Long-format covariate matrix where column 0 is discrete time, followed by covariates.
’w’: Observation weights aligned with person-time rows.
’coord’: Individual identifier integer array linking person-time rows to subjects.

Return type:

Dict[str, np.ndarray]

bart_survival.surv_bart.get_causal_pdp(x: DataFrame, primary_col: str | int = 'inpatient', values: list = [0, 1], audit_support: bool = True, donor_seed: int | None = 42) → Tuple[ndarray, Dict, List[str]]

Generates counterfactual data for Causal G-Computation with support violation audit.

Identifies dependent features or post-exposure mediators that have structural support or positivity violations with the primary exposure (e.g., deterministic zeros like LOS=0 when Inpatient=0). Performs conditional donor resampling for dependent features while leaving baseline confounders intact, strictly avoiding off-support extrapolation.

Parameters:

x (pd.DataFrame) – Baseline covariate dataframe.
primary_col (Union[str, int]) – Primary exposure column name or index.
values (list, optional) – Counterfactual intervention values. Defaults to [0, 1].
audit_support (bool, optional) – Whether to audit and flag support violations. Defaults to True.
donor_seed (int, optional) – Random seed for reproducible donor resampling. Defaults to 42.

Returns:

out_sk: Stacked counterfactual covariate matrix.
c_idx: Index coordinate dictionary compatible with get_posterior_test.
dependent_features: List of audited features exhibiting support violations.

Return type:

Tuple[np.ndarray, Dict, List[str]]

bart_survival.surv_bart.get_pdp(x: ndarray | DataFrame, var_col: List[int | str] = [], values: List[int | float] = [], qt: List[float] = [0.25, 0.5, 0.75], sample_n=None, method: str = 'marginal', audit_support: bool = False) → Tuple[ndarray, Dict] | Tuple[ndarray, Dict, List[str]]

Generates data for Partial Dependency Plots or Causal G-Computation.

Parameters:

x (Union[np.ndarray, pd.DataFrame]) – Covariate matrix or DataFrame.
var_col (List[Union[int, str]], optional) – Covariate column to generate pdp values. Defaults to [].
values (List[Union[int, float]], optional) – Values to test on for each covariate. Defaults to [].
qt (List[float], optional) – Quantiles to generate values on if none are given. Defaults to [0.25, 0.5, 0.75].
sample_n (int, optional) – Sample size of large dataset. Useful for working with large datasets. Defaults to None.
method (str, optional) – ‘marginal’ for standard PDP, or ‘causal’ for support-audited G-Computation. Defaults to ‘marginal’.
audit_support (bool, optional) – If True, audits for positivity/support violations. Defaults to False.

Returns:

PDP covariate matrix and tracking dictionary.

Return type:

Union[Tuple[np.ndarray, Dict], Tuple[np.ndarray, Dict, List[str]]]

bart_survival.surv_bart.get_posterior_test(y_time: ndarray, y_status: ndarray, x: ndarray, time_scale: float | None = None) → Dict[str, ndarray]

Generates evaluation grid in long-time format for out-of-sample posterior predictions.

Constructs counterfactual cross-product of discrete evaluation times and subject covariates.

Parameters:

y_time (np.ndarray) – Follow-up times defining the discrete time evaluation grid.
y_status (np.ndarray) – Event indicators (used for historical reference).
x (np.ndarray) – Baseline covariate matrix to evaluate.
time_scale (Optional[float], optional) – Time discretization scale. Defaults to None.

Returns:

‘post_x’: Evaluated covariate matrix with time in column 0.
’coords’: Subject coordinate identifier array.

Return type:

Dict[str, np.ndarray]

bart_survival.surv_bart.get_shift_pdp(x: DataFrame | ndarray, primary_col: str | int, delta: float = 1.0, clip_to_range: bool = False) → Tuple[ndarray, Dict[str, Any]]

Generates counterfactual datasets for a continuous shift intervention (Modified Treatment Policy).

Contrasts the natural course cohort (Arm 0: X) against a cohort where every individual’s continuous exposure is shifted by delta units (Arm 1: X + delta).

Parameters:

x (Union[pd.DataFrame, np.ndarray]) – Baseline covariate matrix or DataFrame.
primary_col (Union[str, int]) – Continuous variable name or column index to shift.
delta (float, optional) – Increment added to primary_col. Defaults to 1.0.
clip_to_range (bool, optional) – Whether to clip shifted values to observed maximum. Defaults to False.

Returns:

out_sk: Stacked counterfactual matrix of shape (2 * N, n_features).
c_idx: Index coordinate dictionary demarcating Arm 0 and Arm 1.

Return type:

Tuple[np.ndarray, Dict[str, Any]]

bart_survival.surv_bart.get_surv_pre_train(y_time: ndarray, y_status: ndarray, x: ndarray, weight: ndarray | None = None, time_scale: float | None = None) → Dict[str, ndarray]

Generates pooled cross-sectional long-format dataset for discrete-time survival training.

Transforms right-censored time-to-event outcomes into person-time binary event observations. Each individual contributes records for discrete intervals $t = 1, 2, dots, T_i$.

Parameters:

y_time (np.ndarray) – Array of observed follow-up times.
y_status (np.ndarray) – Binary event indicator (1 if event occurred, 0 if right-censored).
x (np.ndarray) – Baseline covariate matrix of shape (n_individuals, n_features).
weight (Optional[np.ndarray], optional) – Observation weights (e.g. for case-cohort designs). Defaults to 1.0.
time_scale (Optional[float], optional) – Time discretization scale (e.g. 5 for 5-day intervals). Defaults to None.

Returns:

‘y’: Binary event indicator array of shape (n_person_time, 1).
’x’: Long-format covariate matrix where column 0 is discrete time, followed by covariates.
’w’: Observation weights aligned with person-time rows.
’coord’: Individual identifier integer array linking person-time rows to subjects.

Return type:

Dict[str, np.ndarray]

bart_survival.surv_bart.get_sv_mean_quant(sv: ndarray, msk: ndarray, draws: bool = True, qntile: List[float] = [0.025, 0.975]) → Dict[str, ndarray]

Generates mean and quantile estimates of Survival Probabilities.

Parameters:

sv (np.ndarray) – Survival probabilities of shape (n_draws, n_obs, n_times).
msk (np.ndarray) – Boolean mask selecting subpopulation of interest.
draws (bool, optional) – If True, aggregates across posterior draws. Defaults to True.
qntile (List[float], optional) – Credible interval quantiles. Defaults to [0.025, 0.975].

Returns:

Mean and quantile curves for masked and unmasked groups.

Return type:

Dict[str, np.ndarray]

bart_survival.surv_bart.get_sv_prob(post: ndarray | DataArray | Dataset, coords: ndarray) → Dict[str, ndarray]

Generates discrete-time hazard and cumulative survival probability estimates.

Computes S(t) = prod_{s=1}^t (1 - h(s)) across all posterior draws and subjects.

Parameters:

post (Union[np.ndarray, xr.DataArray, xr.Dataset]) – Posterior predictive hazard draws.
coords (np.ndarray) – Subject coordinate indices to reshape predictions per subject and time.

Returns:

‘prob’: Hazard probability array of shape (n_draws, n_subjects, n_time_points).
’sv’: Cumulative survival probability array of shape (n_draws, n_subjects, n_time_points).

Return type:

Dict[str, np.ndarray]

bart_survival.surv_bart.pdp_diff_metric(pdp_val: Dict[str, ndarray], idx: int, qntile: List[float] = [0.025, 0.975]) → Dict[str, ndarray]

Generates posterior estimate of marginal survival difference between counterfactual groups.

Parameters:

pdp_val (Dict[str, np.ndarray]) – Posterior output dictionary containing ‘sv’.
idx (int) – Split index demarcating counterfactual groups in stacked PDP matrix.
qntile (List[float], optional) – Quantiles for credible interval. Defaults to [0.025, 0.975].

Returns:

‘diff_m’: Posterior mean survival difference curve across time points.
’diff_q’: Posterior credible interval quantiles of the difference curve.

Return type:

Dict[str, np.ndarray]

bart_survival.surv_bart.pdp_hr_rd_metric(pdp_val: Dict[str, ndarray], idx: int, event_weights: ndarray | None = None, qntile: List[float] = [0.025, 0.975], eps: float = 1e-07) → Dict[str, Any]

Computes event-weighted and unweighted Hazard Ratios and Risk Differences from counterfactual predictions.

Contrasts Intervention Arm 1 (obs >= idx) over Reference Arm 0 (obs < idx).

Calculates: - Event-weighted summary Hazard Ratio (default, Cox-aligned via complementary log-log link). - Unweighted summary Hazard Ratio across follow-up intervals. - Time-varying Hazard Ratio curve HR(t) via complementary log-log link. - Time-varying discrete risk ratio curve RR_discrete(t) = h1(t) / h0(t) (not averaged). - Time-varying Risk Difference curve RD(t) = Risk_1(t) - Risk_0(t) = S_0(t) - S_1(t) (1 over 0). - Survival Difference curve S_1(t) - S_0(t) = -RD(t). - Event-weighted and unweighted summary Risk Differences. - Final horizon Risk Difference RD(t_end). - Restricted Mean Survival Time (RMST) difference.

Parameters:

pdp_val (Dict[str, np.ndarray]) – Dictionary containing ‘prob’ (hazards) and ‘sv’ (survival curves). Arrays have shape (n_draws, 2 * N, n_times).
idx (int) – Split index demarcating Reference Arm 0 (:idx) and Intervention Arm 1 (idx:).
event_weights (Optional[np.ndarray], optional) – Normalized event weights per interval. If None, defaults to uniform interval weighting.
qntile (List[float], optional) – Credible interval quantiles. Defaults to [0.025, 0.975].
eps (float, optional) – Epsilon threshold for numerical safety. Defaults to 1e-7.

Returns:

Comprehensive dictionary containing summary statistics, credible intervals, and draw arrays.

Return type:

Dict[str, Any]

bart_survival.surv_bart.pdp_rr_metric(pdp_val: Dict[str, ndarray], idx: int, qntile: List[float] = [0.025, 0.975]) → Dict[str, ndarray]

Generates posterior Hazard Ratio / Risk Ratio between counterfactual intervention groups.

Parameters:

pdp_val (Dict[str, np.ndarray]) – Posterior output dictionary containing ‘prob’.
idx (int) – Split index demarcating counterfactual groups in stacked PDP matrix.
qntile (List[float], optional) – Quantiles for credible interval. Defaults to [0.025, 0.975].

Returns:

‘rr_m’: Posterior mean hazard ratio curve across time points.
’rr_q’: Posterior credible interval quantiles of the hazard ratio curve.

Return type:

Dict[str, np.ndarray]