Fix bootstrap scaling and selection_random standardization

.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"

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/

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

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

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,28 +294,39 @@ 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()
         warnings.filterwarnings("ignore", message=".*datetime.datetime.utcnow.*")
         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]
 
 

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

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")