diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml index 9ff701b..a79130f 100644 --- a/.github/workflows/python-app.yml +++ b/.github/workflows/python-app.yml @@ -13,6 +13,7 @@ permissions: jobs: test: runs-on: ${{ matrix.os }} + timeout-minutes: 30 strategy: fail-fast: false matrix: @@ -28,7 +29,9 @@ jobs: - name: Set up Python ${{ matrix.python-version }} run: uv python install ${{ matrix.python-version }} - # tabmat and glum have no Python 3.15 wheels yet, so they build from source. + # glum has no Python 3.15 wheel yet, so it builds from source against + # Homebrew libomp. One OpenMP thread stops it deadlocking with the libomp + # bundled in the scikit-learn and tabmat wheels. # lifelines caps pandas below 3, and pandas 2.x is broken on 3.15, so override # the cap and stop uv run from re-syncing it back. - name: Prepare Python 3.15 build @@ -37,6 +40,7 @@ jobs: brew install libomp xsimd jemalloc echo "CPPFLAGS=-I$(brew --prefix libomp)/include -I$(brew --prefix xsimd)/include -I$(brew --prefix jemalloc)/include" >> "$GITHUB_ENV" echo "LDFLAGS=-L$(brew --prefix libomp)/lib -L$(brew --prefix jemalloc)/lib" >> "$GITHUB_ENV" + echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV" echo "pandas>=3" > "$RUNNER_TEMP/uv-overrides.txt" echo "UV_OVERRIDE=$RUNNER_TEMP/uv-overrides.txt" >> "$GITHUB_ENV" echo "UV_NO_SYNC=1" >> "$GITHUB_ENV" diff --git a/CITATION.cff b/CITATION.cff index f3cfe26..9305e48 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -18,7 +18,7 @@ authors: affiliation: "CAUSALab, Department of Epidemiology, Harvard T.H. Chan School of Public Health; Department of Biostatistics, Harvard T.H. Chan School of Public Health" orcid: https://orcid.org/0000-0003-1619-8456 title: "pySEQTarget: Sequential Target Trial Emulation" -version: 0.14.4 +version: 0.14.5 url: https://pyseqtarget.readthedocs.io/ repository-code: https://github.com/CausalInference/pySEQTarget repository-artifact: https://pypi.org/project/pySEQTarget/ diff --git a/pySEQTarget/SEQopts.py b/pySEQTarget/SEQopts.py index aa29019..f4e7fbb 100644 --- a/pySEQTarget/SEQopts.py +++ b/pySEQTarget/SEQopts.py @@ -10,7 +10,11 @@ class SEQopts: Parameter builder for ``pySEQTarget.SEQuential`` analysis :param bootstrap_nboot: Number of bootstraps to perform - :param bootstrap_sample: Subsampling proportion of ID-Trials gathered for each bootstrapping iteration + :param bootstrap_sample: Proportion of subjects drawn (with replacement) in each bootstrap + resample, in (0, 1], default ``0.8``. Values below ``1`` give an m-out-of-n bootstrap, which + is faster on large data; the bootstrap estimates are rescaled by ``sqrt(bootstrap_sample)`` + around the full-data estimate (ratios on the log scale) so that standard errors and + confidence intervals refer to the full sample size :param bootstrap_CI: If bootstrapped, confidence interval level :param bootstrap_CI_method: If bootstrapped, confidence interval method ['SE' or 'percentile'] :param cense_colname: Column name for censoring effect (LTFU, etc.) @@ -85,8 +89,14 @@ class SEQopts: ``treatment_level`` to keep when ``selection_random=True``, in (0, 1], default ``0.8`` :param selection_random: Boolean to randomly subsample the ID-trials starting on the first level of ``treatment_level`` (e.g. the untreated), keeping a proportion ``selection_sample`` - of them, sampled without replacement, and all other ID-trials. The sampled ID-trials are - not reweighted + of them, sampled without replacement, and all other ID-trials. This reduces computation. + The models are fit on the kept ID-trials unweighted - selection depends only on baseline + treatment, which the outcome model conditions on - while the risks and hazard ratios are + standardized with the kept ID-trials of the first level weighted by + ``1 / selection_sample``, so they refer to all ID-trials (the weighted hazard-ratio Cox + fit always uses lifelines, as scikit-survival's Cox model takes no weights). Selection is + made once, from ``seed``; bootstrap confidence intervals, which resample subjects from the + kept ID-trials, include the variability it adds :param subgroup_colname: Column name for subgroups to share the same weighting but different outcome model fits :param treatment_level: List of eligible treatment levels within treatment_col, default ``[0, 1]`` :param trial_include: Boolean to force trial values into model covariates diff --git a/pySEQTarget/analysis/_endoffup.py b/pySEQTarget/analysis/_endoffup.py index bb2e9de..4d7f4c7 100644 --- a/pySEQTarget/analysis/_endoffup.py +++ b/pySEQTarget/analysis/_endoffup.py @@ -12,6 +12,7 @@ import polars as pl from scipy import stats +from ..helpers._bootstrap import _rescale_boot from ..weighting._weight_stats import _weight_bounds from ._risk_estimates import _ci_label @@ -279,13 +280,16 @@ def boot_estimates(key): if has_ci: se_rows, lci_rows, uci_rows = [], [], [] for key in keys: - draws = boot_estimates(key) - se = float(np.nanstd(draws, ddof=1)) point = float( full.filter( pl.all_horizontal([pl.col(c) == v for c, v in zip(key_cols, key)]) )["estimate"][0] ) + # Rescale to the full sample size (m-out-of-n bootstrap) for the + # per-arm SE and CI; the between-arm contrasts below are rescaled on + # their own scales + draws = _rescale_boot(boot_estimates(key), point, self.bootstrap_sample) + se = float(np.nanstd(draws, ddof=1)) if use_se: lci, uci = point - z * se, point + z * se else: @@ -348,15 +352,23 @@ def boot_estimates(key): if has_ci: kx = ((g,) if g is not None else ()) + (a_x,) ky = ((g,) if g is not None else ()) + (a_y,) - d = boot_estimates(ky) - boot_estimates(kx) + d = _rescale_boot( + boot_estimates(ky) - boot_estimates(kx), + diff, + self.bootstrap_sample, + ) d_se = float(np.nanstd(d, ddof=1)) + # Ratio summarised (and rescaled) on the log scale, as in + # _risk_estimates; non-positive draws dropped with np.errstate(divide="ignore", invalid="ignore"): r = boot_estimates(ky) / boot_estimates(kx) - r_valid = r[np.isfinite(r) & (r > 0)] + log_r = np.log(r[np.isfinite(r) & (r > 0)]) + if 0 < ratio < float("inf"): + log_r = _rescale_boot( + log_r, np.log(ratio), self.bootstrap_sample + ) r_logse = ( - float(np.std(np.log(r_valid), ddof=1)) - if r_valid.size > 1 - else float("nan") + float(np.std(log_r, ddof=1)) if log_r.size > 1 else float("nan") ) if use_se: d_lci, d_uci = diff - z * d_se, diff + z * d_se @@ -368,9 +380,9 @@ def boot_estimates(key): else: d_lci = float(np.nanquantile(d, alpha)) d_uci = float(np.nanquantile(d, 1 - alpha)) - if r_valid.size > 1: - r_lci = float(np.quantile(r_valid, alpha)) - r_uci = float(np.quantile(r_valid, 1 - alpha)) + if log_r.size > 1: + r_lci = float(np.exp(np.quantile(log_r, alpha))) + r_uci = float(np.exp(np.quantile(log_r, 1 - alpha))) else: r_lci = r_uci = float("nan") row[f"Difference {ci_label} LCI"] = d_lci diff --git a/pySEQTarget/analysis/_hazard.py b/pySEQTarget/analysis/_hazard.py index a3cbc52..cd44eeb 100644 --- a/pySEQTarget/analysis/_hazard.py +++ b/pySEQTarget/analysis/_hazard.py @@ -6,6 +6,8 @@ import polars as pl from lifelines import CoxPHFitter +from ..expansion._selection import _selection_weight +from ..helpers._bootstrap import _rescale_boot from ..helpers._predict_model import _safe_predict from ._outcome_fit import _cast_categories from ._subgroup_fit import _warn_missing_subgroup @@ -59,6 +61,9 @@ def _calculate_hazard_single(self, data, val=None): if len(boot_log_hrs) == 0: return _create_hazard_output(np.exp(full_log_hr), None, None, val, self) + boot_log_hrs = _rescale_boot( + np.asarray(boot_log_hrs), full_log_hr, self.bootstrap_sample + ) if self.bootstrap_CI_method == "se": from scipy.stats import norm @@ -203,6 +208,12 @@ def _hazard_handler(self, data, val, boot_idx, rng): ] if self.compevent_colname: exclude_cols.append(self.compevent_colname) + # Weight ID-trials kept by selection_random by the inverse of their + # selection probability, taken from the observed baseline treatment before + # each ID-trial is copied into every arm below + sel_w = _selection_weight(self) + if sel_w is not None: + data = data.with_columns(sel_w.alias("_sel_w")) keep_cols = [col for col in data.columns if col not in exclude_cols] trials = ( @@ -283,6 +294,7 @@ def _hazard_handler(self, data, val, boot_idx, rng): sim_data_pd = sim_data.to_pandas() tx_bas = f"{self.treatment_col}{self.indicator_baseline}" + weights_col = "_sel_w" if sel_w is not None else None try: # COXPHFITTER CURRENTLY HAS DEPRECATED datetime.datetime.utcnow() @@ -290,21 +302,31 @@ def _hazard_handler(self, data, val, boot_idx, rng): if ce_model is not None: cox_data = sim_data_pd[sim_data_pd["event"].isin([0, 1])].copy() cox_data["event_binary"] = (cox_data["event"] == 1).astype(int) - return _cox_log_hr(self, cox_data, "followup", "event_binary", tx_bas) - return _cox_log_hr(self, sim_data_pd, "followup", "event", tx_bas) + return _cox_log_hr( + self, cox_data, "followup", "event_binary", tx_bas, weights_col + ) + return _cox_log_hr(self, sim_data_pd, "followup", "event", tx_bas, weights_col) except Exception as e: print(f"Cox model fitting failed: {e}") return None -def _cox_log_hr(self, data_pd, duration_col, event_col, covariate_col): +def _cox_log_hr( + self, data_pd, duration_col, event_col, covariate_col, weights_col=None +): """ Fit a univariate Cox model (single covariate = baseline treatment) and return the log hazard ratio, dispatching on self.cox_package. scikit-survival uses Efron tie handling to match lifelines, which matters here because integer follow-up produces many tied event times. + + scikit-survival's Cox model takes no case weights, so a weighted fit + (``weights_col``, from selection_random) always uses lifelines. """ - if getattr(self, "cox_package", "lifelines") == "scikit-survival": + if ( + getattr(self, "cox_package", "lifelines") == "scikit-survival" + and weights_col is None + ): from sksurv.linear_model import CoxPHSurvivalAnalysis y = np.empty(len(data_pd), dtype=[("event", bool), ("time", "float64")]) @@ -316,12 +338,17 @@ def _cox_log_hr(self, data_pd, duration_col, event_col, covariate_col): return float(cox.coef_[0]) cph = CoxPHFitter() - cph.fit( - data_pd, - duration_col=duration_col, - event_col=event_col, - formula=f"`{covariate_col}`", - ) + with warnings.catch_warnings(): + # lifelines suggests robust=True for non-integer weights; only the point + # estimate is used here, the CI comes from the bootstrap + warnings.filterwarnings("ignore", message=".*not integers.*") + cph.fit( + data_pd, + duration_col=duration_col, + event_col=event_col, + weights_col=weights_col, + formula=f"`{covariate_col}`", + ) return cph.params_.values[0] diff --git a/pySEQTarget/analysis/_outcome_fit.py b/pySEQTarget/analysis/_outcome_fit.py index ec64d4b..c5a0728 100644 --- a/pySEQTarget/analysis/_outcome_fit.py +++ b/pySEQTarget/analysis/_outcome_fit.py @@ -9,14 +9,11 @@ from ..helpers._check_design import _check_design from ..helpers._fix_categories import _cast_fixed_categories -from ..helpers._spline import _compute_spline_knots, _cr_term +from ..helpers._spline import _bake_spline_knots from ..weighting._weight_stats import _weight_bounds -def _apply_spline_formula(formula, indicator_squared, spline_knots): - inner_knots, lower, upper = spline_knots - spline = _cr_term("followup", inner_knots, lower, upper) - +def _apply_spline_formula(formula, indicator_squared, spline): # Mark interactions first, so dropping the bare followup terms below # cannot reach into the spline, which contains "followup" itself mark = "_SPLINE_" @@ -105,8 +102,10 @@ def _outcome_fit( if self.followup_spline: if getattr(self, "_current_boot_idx", None) is None: - self._followup_spline_knots = _compute_spline_knots( - self.DT["followup"].to_numpy(), df=self.followup_spline_df + # Centred, so the spline is not collinear with the intercept + self._followup_spline_knots = _bake_spline_knots( + f'cr(followup, df={self.followup_spline_df}, constraints="center")', + self.DT, ) formula = _apply_spline_formula( formula, self.indicator_squared, self._followup_spline_knots diff --git a/pySEQTarget/analysis/_risk_estimates.py b/pySEQTarget/analysis/_risk_estimates.py index 41b5204..027780c 100644 --- a/pySEQTarget/analysis/_risk_estimates.py +++ b/pySEQTarget/analysis/_risk_estimates.py @@ -3,6 +3,8 @@ import polars as pl from scipy import stats +from ..helpers._bootstrap import _rescale_boot + def _ci_label(bootstrap_CI): """Format the CI level as a column-name fragment, e.g. 0.95 -> '95%'.""" @@ -188,31 +190,35 @@ def _risk_estimates(self): risk_x_val / risk_y_val if risk_y_val != 0 else float("inf") ) - # Filter degenerate RR bootstrap values (risk_y == 0 or negative) - valid_rr = paired.filter( - (pl.col("risk_y") > 0) & (pl.col("risk_x") >= 0) - ).with_columns((pl.col("risk_x") / pl.col("risk_y")).alias("RR")) + # Rescale to the full sample size (m-out-of-n bootstrap): + # the RD on its own scale, the RR on the log scale. Drop + # degenerate RR bootstrap values (non-positive or infinite) + rd_i = _rescale_boot(paired["RD"], rd_point, self.bootstrap_sample) + rr_i = paired["risk_x"] / paired["risk_y"] + log_rr_i = rr_i.filter(rr_i.is_finite() & (rr_i > 0)).log() + if 0 < rr_point < float("inf"): + log_rr_i = _rescale_boot( + log_rr_i, math.log(rr_point), self.bootstrap_sample + ) - n_valid_rr = len(valid_rr) + n_valid_rr = len(log_rr_i) # Bootstrap SEs, retained regardless of CI method: the risk # difference SE is on the natural scale, the risk ratio SE # on the log scale (the scale ratio measures are pooled on # for inverse-variance meta-analysis: combine the log Risk # Ratio with log(RR) SE, then exponentiate). - rd_se = float(paired["RD"].std()) + rd_se = float(rd_i.std()) log_rr_se = ( - float(valid_rr["RR"].log().std()) - if n_valid_rr >= 2 - else float("nan") + float(log_rr_i.std()) if n_valid_rr >= 2 else float("nan") ) if self.bootstrap_CI_method == "percentile": - rd_lci = float(paired["RD"].quantile(alpha / 2)) - rd_uci = float(paired["RD"].quantile(1 - alpha / 2)) + rd_lci = float(rd_i.quantile(alpha / 2)) + rd_uci = float(rd_i.quantile(1 - alpha / 2)) if n_valid_rr >= 2: - rr_lci = float(valid_rr["RR"].quantile(alpha / 2)) - rr_uci = float(valid_rr["RR"].quantile(1 - alpha / 2)) + rr_lci = math.exp(log_rr_i.quantile(alpha / 2)) + rr_uci = math.exp(log_rr_i.quantile(1 - alpha / 2)) else: rr_lci = float("nan") rr_uci = float("nan") diff --git a/pySEQTarget/analysis/_survival_pred.py b/pySEQTarget/analysis/_survival_pred.py index 2d41db1..d5ea1e7 100644 --- a/pySEQTarget/analysis/_survival_pred.py +++ b/pySEQTarget/analysis/_survival_pred.py @@ -2,6 +2,8 @@ import polars as pl from patsy import PatsyError, dmatrix +from ..expansion._selection import _selection_weight +from ..helpers._bootstrap import _rescale_boot from ..helpers._fix_categories import (_fix_categories_for_predict, _model_design_info) from ..helpers._predict_model import _safe_predict @@ -167,9 +169,7 @@ def _calculate_risk(self, data, val=None): followup_range = list(range(0, self.followup_max + 1)) SDT = ( - data.with_columns( - [pl.concat_str([pl.col(self.id_col), pl.col("trial")]).alias("TID")] - ) + data.with_columns([pl.struct([self.id_col, "trial"]).alias("TID")]) .group_by("TID") .first() .drop(["followup", f"followup{self.indicator_squared}"]) @@ -180,6 +180,42 @@ def _calculate_risk(self, data, val=None): ) ).sort([self.id_col, "trial", "followup"]) + # Standardize over the ID-trials before selection_random: weight the kept + # ID-trials by the inverse of their selection probability, taken from the + # observed baseline treatment before it is set to each arm below + sel_w = _selection_weight(self) + if sel_w is not None: + SDT = SDT.with_columns(sel_w.alias("_sel_w")) + + def _mean(col): + return (pl.col(col) * pl.col("_sel_w")).sum() / pl.col("_sel_w").sum() + + else: + + def _mean(col): + return pl.col(col).mean() + + def _boot_summary(frame, cols, main, name): + # Rescale to the full sample size (m-out-of-n bootstrap) for the SE and + # percentile CI of each curve. _store_boot_risks() takes the draws + # before this: _risk_estimates() rescales the RD and log RR itself. + return ( + frame.select(["followup", main] + cols) + .unpivot( + index=["followup", main], + on=cols, + variable_name="bootID", + value_name=name, + ) + .with_columns( + _rescale_boot(pl.col(name), pl.col(main), self.bootstrap_sample).alias( + name + ) + ) + .drop(main) + .group_by("followup") + ) + risks = [] for treatment_val in self.treatment_level: TxDT = SDT.with_columns( @@ -206,7 +242,9 @@ def _calculate_risk(self, data, val=None): preds = _get_outcome_predictions(self, TxDT, val=val) # Drop original data columns — only followup and TID needed from here - TxDT = TxDT.select(["followup", "TID"]) + TxDT = TxDT.select( + ["followup", "TID"] + (["_sel_w"] if sel_w is not None else []) + ) pred_series = [pl.Series("pred_outcome", preds["outcome"][0])] @@ -252,7 +290,7 @@ def _calculate_risk(self, data, val=None): inc_names = [n.replace("pred_outcome", "inc") for n in outcome_names] TxDT = ( TxDT.group_by("followup") - .agg([pl.col(col).mean() for col in surv_names + inc_names]) + .agg([_mean(col).alias(col) for col in surv_names + inc_names]) .sort("followup") .with_columns(pl.col("followup") + 1) ) @@ -271,7 +309,7 @@ def _calculate_risk(self, data, val=None): ] ) .group_by("followup") - .agg([pl.col(col).mean() for col in outcome_names]) + .agg([_mean(col).alias(col) for col in outcome_names]) .sort("followup") .with_columns([(1 - pl.col(col)).alias(col) for col in outcome_names]) .with_columns(pl.col("followup") + 1) @@ -283,14 +321,7 @@ def _calculate_risk(self, data, val=None): if boot_cols: risk = ( - TxDT.select(["followup"] + boot_cols) - .unpivot( - index="followup", - on=boot_cols, - variable_name="bootID", - value_name="risk", - ) - .group_by("followup") + _boot_summary(TxDT, boot_cols, main_col, "risk") .agg( [ pl.col("risk").std().cast(pl.Float64).alias("SE"), @@ -318,14 +349,7 @@ def _calculate_risk(self, data, val=None): inc_boot_cols = [col for col in inc_names if col != "inc"] if inc_boot_cols: inc_risk = ( - TxDT.select(["followup"] + inc_boot_cols) - .unpivot( - index="followup", - on=inc_boot_cols, - variable_name="bootID", - value_name="inc_val", - ) - .group_by("followup") + _boot_summary(TxDT, inc_boot_cols, "inc", "inc_val") .agg( [ pl.col("inc_val") @@ -344,6 +368,13 @@ def _calculate_risk(self, data, val=None): ) .join(TxDT.select(["followup", "inc"]), on="followup") ) + if self.bootstrap_CI_method == "se": + inc_risk = inc_risk.with_columns( + [ + (pl.col("inc") - z * pl.col("inc_SE")).alias("inc_LCI"), + (pl.col("inc") + z * pl.col("inc_SE")).alias("inc_UCI"), + ] + ) risk = risk.join(inc_risk, on="followup") final_cols = [ "followup", @@ -491,9 +522,15 @@ def _calculate_risk(self, data, val=None): def _calculate_survival(self, risk_data): + if (risk_data["estimate"] == "survival").any(): + return risk_data.clear() if self.bootstrap_nboot > 0: surv = risk_data.with_columns( - [(1 - pl.col(col)).alias(col) for col in ["pred", "LCI", "UCI"]] + [ + (1 - pl.col("pred")).alias("pred"), + (1 - pl.col("UCI")).alias("LCI"), + (1 - pl.col("LCI")).alias("UCI"), + ] ).with_columns(pl.lit("survival").alias("estimate")) else: surv = risk_data.with_columns( diff --git a/pySEQTarget/expansion/__init__.py b/pySEQTarget/expansion/__init__.py index b11771b..ceaba1b 100644 --- a/pySEQTarget/expansion/__init__.py +++ b/pySEQTarget/expansion/__init__.py @@ -2,7 +2,7 @@ from ._diagnostics import _diagnostics from ._dynamic import _dynamic from ._mapper import _mapper -from ._selection import _random_selection +from ._selection import _random_selection, _selection_weight __all__ = [ "_binder", @@ -10,4 +10,5 @@ "_dynamic", "_mapper", "_random_selection", + "_selection_weight", ] diff --git a/pySEQTarget/expansion/_selection.py b/pySEQTarget/expansion/_selection.py index adaedf2..4cc5be2 100644 --- a/pySEQTarget/expansion/_selection.py +++ b/pySEQTarget/expansion/_selection.py @@ -1,6 +1,30 @@ import polars as pl +def _selection_weight(self): + """Inverse-probability weights for the ID-trials kept by ``selection_random``. + + ``selection_random`` keeps a proportion ``selection_sample`` of the + ID-trials starting on the first level of ``treatment_level`` and all other + ID-trials. Selection depends only on baseline treatment, which the outcome + model conditions on, so the model needs no correction, but standardizing + over the kept ID-trials under-represents the first arm's covariate + distribution. Each kept ID-trial of that arm stands for + ``1 / selection_sample`` ID-trials. + + Returns a polars expression on the observed baseline treatment, or ``None`` + when no selection was made. + """ + if not self.selection_random: + return None + tx_bas = f"{self.treatment_col}{self.indicator_baseline}" + return ( + pl.when(pl.col(tx_bas) == self.treatment_level[0]) + .then(pl.lit(1 / self.selection_sample)) + .otherwise(pl.lit(1.0)) + ) + + def _random_selection(self): """ Handles the case where random selection is applied for data from diff --git a/pySEQTarget/helpers/__init__.py b/pySEQTarget/helpers/__init__.py index 2244d71..adee8f3 100644 --- a/pySEQTarget/helpers/__init__.py +++ b/pySEQTarget/helpers/__init__.py @@ -1,4 +1,4 @@ -from ._bootstrap import bootstrap_loop +from ._bootstrap import _rescale_boot, bootstrap_loop from ._col_string import _col_string from ._format_time import _format_time from ._offloader import Offloader @@ -11,6 +11,7 @@ __all__ = [ "bootstrap_loop", + "_rescale_boot", "_col_string", "_format_time", "Offloader", diff --git a/pySEQTarget/helpers/_bootstrap.py b/pySEQTarget/helpers/_bootstrap.py index ff1fec8..c365551 100644 --- a/pySEQTarget/helpers/_bootstrap.py +++ b/pySEQTarget/helpers/_bootstrap.py @@ -1,4 +1,5 @@ import copy +import math import time import warnings from concurrent.futures import ProcessPoolExecutor, as_completed @@ -27,6 +28,23 @@ ) +def _rescale_boot(draws, estimate, bootstrap_sample): + """Rescale bootstrap draws to the full sample size. + + Each resample draws ``bootstrap_sample * n`` of the ``n`` subjects (an + m-out-of-n bootstrap), so the spread of the draws reflects a sample of size + ``m``, not ``n``: for a root-n consistent estimator their standard deviation + is about ``1 / sqrt(bootstrap_sample)`` times too large. Shrinking each draw + towards the full-data estimate by ``sqrt(bootstrap_sample)`` corrects this + for both the standard error and the percentile interval, and keeps draws of + a bounded quantity (a risk) within its bounds. Works on numpy arrays and + polars expressions alike; no-op when ``bootstrap_sample = 1``. + """ + if bootstrap_sample >= 1: + return draws + return estimate + math.sqrt(bootstrap_sample) * (draws - estimate) + + def _prepare_boot_data(self, data, boot_id): id_counts = self._boot_samples[boot_id] diff --git a/pyproject.toml b/pyproject.toml index 99ed58e..deca546 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pySEQTarget" -version = "0.14.4" +version = "0.14.5" description = "Sequential Target Trial Emulation" readme = "README.md" license = {text = "MIT"} @@ -98,6 +98,8 @@ SEQdata = ["data/*.csv"] [tool.pytest.ini_options] pythonpath = ["."] testpaths = ["tests"] +# Dump stack traces if a test hangs, e.g. an OpenMP deadlock +faulthandler_timeout = 300 filterwarnings = [ "ignore:FigureCanvasAgg is non-interactive:UserWarning", # joblib internals still assign array.shape; deprecated in NumPy 2.5, diff --git a/tests/test_bootstrap_rescale.py b/tests/test_bootstrap_rescale.py new file mode 100644 index 0000000..57a4f62 --- /dev/null +++ b/tests/test_bootstrap_rescale.py @@ -0,0 +1,130 @@ +from types import SimpleNamespace + +import numpy as np +import polars as pl +import pytest +from pytest import approx + +from pySEQTarget.analysis import _create_endoffup, _risk_estimates +from pySEQTarget.helpers import _rescale_boot + +PROBS = [0.025, 0.975] + + +def test_rescale_boot_corrects_m_out_of_n_spread(): + # Bootstrap the mean drawing half the sample each time: the raw draws are + # sqrt(2) times too spread out; the rescaled ones match the full-sample SE + rng = np.random.default_rng(1) + x = rng.normal(size=4000) + draws = np.array([rng.choice(x, 2000, replace=True).mean() for _ in range(2000)]) + rescaled = _rescale_boot(draws, x.mean(), 0.5) + assert np.std(rescaled, ddof=1) == approx(np.std(x) / np.sqrt(4000), rel=0.05) + assert np.std(rescaled, ddof=1) == approx(np.std(draws, ddof=1) * np.sqrt(0.5)) + # A bounded quantity stays within its bounds + bounded = _rescale_boot(np.array([0.0, 1.0]), 0.3, 0.5) + assert np.all((bounded >= 0) & (bounded <= 1)) + # No-op for the standard bootstrap + assert _rescale_boot(draws, x.mean(), 1.0) is draws + + +def test_rescale_boot_polars_expression_keeps_draw_column(): + df = pl.DataFrame({"est": [0.5, 0.5], "draw": [0.3, 0.9]}) + out = df.with_columns( + _rescale_boot(pl.col("draw"), pl.col("est"), 0.25).alias("draw") + ) + assert out["draw"].to_list() == approx([0.4, 0.7]) + assert out["est"].to_list() == [0.5, 0.5] + + +def _shrink(q, est): + # bootstrap_sample = 0.25 shrinks each draw's distance from the estimate by + # sqrt(0.25) = 0.5; with an affine map the percentiles shrink the same way + return est + 0.5 * (q - est) + + +@pytest.fixture +def draws(): + rng = np.random.default_rng(1) + return rng.uniform(0.2, 0.4, 200), rng.uniform(0.3, 0.6, 200) + + +def test_risk_estimates_percentile_ci_rescaled(draws): + r0, r1 = draws + boot = lambda r: pl.DataFrame( # noqa: E731 + {"followup": [10] * 200, "risk": r, "boot_idx": list(range(1, 201))} + ) + obj = SimpleNamespace( + km_data=pl.DataFrame( + { + "followup": [10, 10], + "tx_init": [0, 1], + "pred": [0.3, 0.45], + "estimate": ["risk", "risk"], + } + ), + risk_times=None, + subgroup_colname=None, + bootstrap_nboot=200, + bootstrap_sample=0.25, + bootstrap_CI=0.95, + bootstrap_CI_method="percentile", + treatment_level=[0, 1], + treatment_col="tx_init", + _boot_risks={0: boot(r0), 1: boot(r1)}, + ) + out = _risk_estimates(obj) + rd = out["risk_difference"].filter(pl.col("A_x") == 1) + rr = out["risk_ratio"].filter(pl.col("A_x") == 1) + rd_draws = pl.Series(r1 - r0) + log_rr_draws = pl.Series(np.log(r1 / r0)) + assert [rd["RD 95% LCI"][0], rd["RD 95% UCI"][0]] == approx( + [_shrink(rd_draws.quantile(p), 0.15) for p in PROBS] + ) + assert [rr["RR 95% LCI"][0], rr["RR 95% UCI"][0]] == approx( + [np.exp(_shrink(log_rr_draws.quantile(p), np.log(1.5))) for p in PROBS] + ) + assert rd["RD SE"][0] == approx(0.5 * rd_draws.std()) + assert rr["log(RR) SE"][0] == approx(0.5 * log_rr_draws.std()) + + +def test_endoffup_percentile_ci_rescaled(draws): + r0, r1 = draws + + def eof(e0, e1): + return { + "eof": pl.DataFrame( + { + "tx_init_bas": [0, 1], + "estimate": [e0, e1], + "n": [10, 10], + "n_subjects": [10, 10], + "n_eligible": [10, 10], + "n_censored": [0, 0], + "n_nomeasure": [0, 0], + } + ) + } + + obj = SimpleNamespace( + treatment_col="tx_init", + indicator_baseline="_bas", + subgroup_colname=None, + end_of_fup_type="binary", + end_of_fup_time=5, + bootstrap_CI=0.95, + bootstrap_CI_method="percentile", + bootstrap_sample=0.25, + outcome_model=[eof(0.3, 0.45)] + [eof(a, b) for a, b in zip(r0, r1)], + ) + out = _create_endoffup(obj) + est = out["eof_data"].filter(pl.col("A") == 0) + assert [est["95% LCI"][0], est["95% UCI"][0]] == approx( + [_shrink(np.quantile(r0, p), 0.3) for p in PROBS] + ) + cmp = out["eof_comparison"].filter(pl.col("A_x") == 0) + assert [cmp["Difference 95% LCI"][0], cmp["Difference 95% UCI"][0]] == approx( + [_shrink(np.quantile(r1 - r0, p), 0.15) for p in PROBS] + ) + assert [cmp["Ratio 95% LCI"][0], cmp["Ratio 95% UCI"][0]] == approx( + [np.exp(_shrink(np.quantile(np.log(r1 / r0), p), np.log(1.5))) for p in PROBS] + ) diff --git a/tests/test_followup_options.py b/tests/test_followup_options.py index c33ecc5..ba45223 100644 --- a/tests/test_followup_options.py +++ b/tests/test_followup_options.py @@ -1,3 +1,4 @@ +import numpy as np import pytest from pySEQTarget import SEQopts, SEQuential @@ -40,15 +41,6 @@ def test_followup_class(): assert [round(x, 3) for x in matrix] == [round(x, 3) for x in expected] -# patsy's natural-cubic cr() basis spans the constant function, so cr() plus -# the model Intercept is always collinear by exactly one dimension. Fitted -# values (all the spline path is used for) are unique regardless; only the -# individual spline coefficients are non-identified, resolved deterministically -# by statsmodels' pinv. statsmodels >= 0.15 reports the redundancy as a -# SingularMatrixWarning. -@pytest.mark.filterwarnings( - "ignore::statsmodels.tools.sm_exceptions.SingularMatrixWarning" -) def test_followup_spline(): data = load_data("SEQdata") @@ -68,20 +60,22 @@ def test_followup_spline(): s.fit() matrix = s.outcome_model[0]["outcome"].summary2().tables[1]["Coef."].to_list() expected = [ - -4.804282252748607, - 0.19115933860001255, - 0.044310717515918724, - 0.0005814999431447507, - 0.12717121164606823, - 0.0032948355025455216, - -0.013371824500839971, - 0.19972467861548412, - -2.027245615586753, - -1.395729861856384, - -0.9397731941281695, - -0.4415335811772879, + -6.275803730636, + 0.190835581131, + 0.044181701547, + 0.000583701185, + 0.127096501708, + 0.003280255156, + -0.013270653552, + 0.199471860366, + 0.281342694425, + 0.453405473706, + 0.758640425731, + 0.966164355851, ] assert [round(x, 3) for x in matrix] == [round(x, 3) for x in expected] + X = s.outcome_model[0]["outcome"].model.exog + assert np.linalg.matrix_rank(X) == X.shape[1] def test_no_followup(): @@ -115,9 +109,6 @@ def test_no_followup(): assert [round(x, 3) for x in matrix] == [round(x, 3) for x in expected] -@pytest.mark.filterwarnings( - "ignore::statsmodels.tools.sm_exceptions.SingularMatrixWarning" -) @pytest.mark.parametrize("method", ["ITT", "dose-response"]) def test_followup_spline_with_km_curves_interactions(method): # km_curves adds treatment- or dose-by-followup interactions, which must diff --git a/tests/test_selection_random.py b/tests/test_selection_random.py index 3607dcc..3d68ca1 100644 --- a/tests/test_selection_random.py +++ b/tests/test_selection_random.py @@ -1,4 +1,6 @@ +import numpy as np import polars as pl +import pytest from pySEQTarget import SEQopts, SEQuential from pySEQTarget.data import load_data @@ -87,3 +89,77 @@ def build(**opts): # subsampled to the requested fraction — not dropped entirely. assert sel[2] == base[2] assert sel[1] == int(prob * base[1]) + + +def _final_risks(**opts): + s = _build(km_curves=True, **opts) + s.fit() + s.survival() + risk = s.km_data.filter(pl.col("estimate") == "risk") + risk = risk.filter(pl.col("followup") == risk["followup"].max()) + return dict(zip(risk["tx_init"].to_list(), risk["pred"].to_list())) + + +def test_selection_random_standardizes_risks_over_all_trials(): + # Unweighted standardization over the kept trials under-represented the + # subsampled arm's covariate distribution and shifted both risks by ~0.009 + # (selection_sample = 0.3, 5 seeds); weighting the kept control trials by + # 1 / selection_sample removes the shift. + full = _final_risks() + selected = [ + _final_risks(selection_random=True, selection_sample=0.3, seed=s) + for s in range(1, 6) + ] + for tx in (0, 1): + mean = sum(r[tx] for r in selected) / len(selected) + assert abs(mean - full[tx]) < 0.005 + + +def test_selection_weight_none_without_selection(): + from pySEQTarget.expansion import _selection_weight + + assert _selection_weight(_build()) is None + w = _selection_weight(_build(selection_random=True, selection_sample=0.25)) + out = pl.DataFrame({"tx_init_bas": [0, 1]}).select(w.alias("w")) + assert out["w"].to_list() == [4.0, 1.0] + + +def test_cox_log_hr_uses_weights_on_both_backends(): + # scikit-survival's Cox model takes no weights, so a weighted fit uses + # lifelines whichever backend is selected + from types import SimpleNamespace + + import pandas as pd + + from pySEQTarget.analysis._hazard import _cox_log_hr + + df = pd.DataFrame( + { + "followup": [1, 2, 3, 4, 5, 6, 2, 3, 5, 6, 7, 8], + "event": [1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1], + "tx": [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1], + "w": [4.0, 1.0, 4.0, 1.0, 4.0, 1.0] + [1.0] * 6, + } + ) + fit = { + pkg: lambda weights_col, pkg=pkg: _cox_log_hr( + SimpleNamespace(cox_package=pkg), df, "followup", "event", "tx", weights_col + ) + for pkg in ("lifelines", "scikit-survival") + } + weighted = fit["lifelines"]("w") + assert weighted != pytest.approx(fit["lifelines"](None), abs=1e-3) + assert fit["scikit-survival"]("w") == pytest.approx(weighted) + + +def test_selection_random_hazard_runs_on_both_cox_backends(): + for pkg in ("lifelines", "scikit-survival"): + s = _build( + hazard_estimate=True, + cox_package=pkg, + selection_random=True, + selection_sample=0.5, + ) + s.fit() + s.hazard() + assert np.isfinite(s.hazard_ratio["Hazard ratio"][0]) diff --git a/tests/test_survival.py b/tests/test_survival.py index af0a55a..0d19b2e 100644 --- a/tests/test_survival.py +++ b/tests/test_survival.py @@ -357,3 +357,61 @@ def test_risk_se_columns_present_subgroup_delta_path(): assert "log(RR) SE" in rr.columns assert "RD 90% LCI" in rd.columns assert "RR 90% LCI" in rr.columns + + +def _curves(data, **opts): + s = SEQuential( + data, + id_col="ID", + time_col="time", + eligible_col="eligible", + treatment_col="tx_init", + outcome_col="outcome", + time_varying_cols=["N", "L", "P"], + fixed_cols=["sex"], + method="ITT", + parameters=SEQopts(km_curves=True, **opts), + ) + s.expand() + if opts.get("bootstrap_nboot"): + s.bootstrap() + s.fit() + s.survival() + return s + + +def test_every_id_trial_enters_the_curves(): + # Pasted together, ID 1 trial 12 and ID 11 trial 2 are the same key; with + # IDs of equal width no two ID-trials can collide, so relabelling the + # people must not change the curves + data = load_data("SEQdata") + keys = ["estimate", "tx_init", "followup"] + as_given = _curves(data).km_data.sort(keys) + relabelled = _curves(data.with_columns(pl.col("ID") + 1000)).km_data.sort(keys) + assert (as_given["pred"] - relabelled["pred"]).abs().max() < 1e-12 + + +def test_compevent_curves_have_one_row_per_estimate(): + data = load_data("SEQdata_LTFU").filter(pl.col("ID") <= 300) + km = _curves(data, compevent_colname="LTFU").km_data + keys = ["estimate", "tx_init", "followup"] + assert set(km["estimate"].unique()) == {"risk", "survival", "incidence"} + assert km.select(keys).is_duplicated().sum() == 0 + wide = km.pivot(on="estimate", index=["tx_init", "followup"], values="pred") + assert (wide["survival"] + wide["risk"] - 1).abs().max() < 1e-12 + + +def test_bootstrapped_bounds_are_ordered_and_follow_the_method(): + from scipy.stats import norm + + data = load_data("SEQdata_LTFU").filter(pl.col("ID") <= 300) + km = _curves(data, compevent_colname="LTFU", bootstrap_nboot=3, seed=42).km_data + assert (km["LCI"] <= km["UCI"]).all() + # bootstrap_CI_method defaults to "se": normal bounds for every estimate + inc = km.filter((pl.col("estimate") == "incidence") & (pl.col("followup") > 0)) + z = norm.ppf(0.975) + lower = (inc["pred"] - z * inc["SE"]).clip(0.0, 1.0) + assert (inc["LCI"] - lower).abs().max() < 1e-12 + + km = _curves(load_data("SEQdata"), bootstrap_nboot=3, seed=42).km_data + assert (km["LCI"] <= km["UCI"]).all()