Add weight_stabilized, fix bootstrap memory growth, point PyPI at the docs

.github/workflows/python-app.yml

@@ -23,7 +23,7 @@ jobs:
     - uses: actions/checkout@v7
 
     - name: Install uv
-      uses: astral-sh/setup-uv@v10.2.0
+      uses: astral-sh/setup-uv@v10.3.3
 
     - name: Set up Python ${{ matrix.python-version }}
       run: uv python install ${{ matrix.python-version }}

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.3
+version: 0.14.4
 url: https://pyseqtarget.readthedocs.io/
 repository-code: https://github.com/CausalInference/pySEQTarget
 repository-artifact: https://pypi.org/project/pySEQTarget/

pySEQTarget/SEQopts.py

@@ -81,8 +81,12 @@ class SEQopts:
         followup is always included. Defaults to ``None`` (report at the maximum followup only).
     :param seed: RNG seed
     :param selection_first_trial: Boolean to only use first trial for analysis (similar to non-expanded)
-    :param selection_sample: Subsampling proportion of ID-trials which did not initiate a treatment
-    :param selection_random: Boolean to randomly downsample ID-trials which did not initiate a treatment
+    :param selection_sample: Proportion of the ID-trials starting on the first level of
+        ``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
     :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
@@ -95,6 +99,10 @@ class SEQopts:
     :param weight_lag_condition: Boolean to fit weights based on their treatment lag
     :param weight_p99: Boolean to force weight min and max to be 1st and 99th percentile respectively
     :param weight_preexpansion: Boolean to fit weights on preexpanded data
+    :param weight_stabilized: Boolean to use stabilized IP weights, default ``True``.
+        ``False`` fits no numerator models, so the treatment weights are cumulative
+        products of ``1 / denominator`` and any ``numerator`` is ignored; LTFU and
+        visit weights keep their numerators
     :param weight_spline: Boolean to model time in the default weight models with a
         centred natural cubic spline (patsy's ``cr()``) instead of a quadratic,
         making the baseline hazard a flexible function of time, default ``False``.
@@ -164,6 +172,7 @@ class SEQopts:
     weight_lag_condition: bool = True
     weight_p99: bool = False
     weight_preexpansion: bool = True
+    weight_stabilized: bool = True
     weight_spline: bool = False
     weight_spline_df: int = 4
     verbose: bool = False
@@ -187,6 +196,7 @@ def _validate_bools(self):
             "weight_lag_condition",
             "weight_p99",
             "weight_preexpansion",
+            "weight_stabilized",
             "weight_spline",
             "weighted",
         ]

pySEQTarget/SEQuential.py

@@ -1,5 +1,6 @@
 import datetime
 import time
+import warnings
 from collections import Counter
 from dataclasses import asdict
 from typing import List, Literal, Optional
@@ -11,7 +12,7 @@
 from .analysis import (_calculate_hazard, _calculate_survival, _clamp,
                        _create_endoffup, _eof_counts, _eof_estimate,
                        _eof_frame, _eof_summary, _outcome_fit, _pred_risk,
-                       _risk_estimates, _subgroup_fit)
+                       _risk_estimates, _slim_model, _subgroup_fit)
 from .error import _data_checker, _param_checker
 from .expansion import _binder, _diagnostics, _dynamic, _random_selection
 from .helpers import (Offloader, _bake_model_formulas, _col_string,
@@ -111,8 +112,25 @@ def __init__(
         if self.covariates is None:
             self.covariates = _outcome(self)
 
+        if not self.weight_stabilized and (not self.weighted or self.method == "ITT"):
+            warnings.warn(
+                "'weight_stabilized=False' only affects the treatment weights of a "
+                "weighted, non-ITT analysis and has no effect here",
+                UserWarning,
+                stacklevel=2,
+            )
+
         if self.weighted:
-            if self.numerator is None:
+            if not self.weight_stabilized and self.method != "ITT":
+                # No numerator model is fit; "1" is for display only
+                if self.numerator is not None:
+                    warnings.warn(
+                        "'numerator' is ignored when weight_stabilized=False",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+                self.numerator = "1"
+            elif self.numerator is None:
                 self.numerator = _numerator(self)
 
             if self.denominator is None:
@@ -365,7 +383,11 @@ def fit(self) -> None:
 
         if self.subgroup_colname is not None:
             models_list = _subgroup_fit(self, start_params=start)
-            if not is_boot:
+            if is_boot:
+                for sg in models_list.values():
+                    for m in sg.values():
+                        _slim_model(m)
+            else:
                 self._outcome_start_params = {
                     val: {
                         key: (m.params.values, list(m.model.exog_names))
@@ -400,7 +422,12 @@ def fit(self) -> None:
                 start_params=start_compevent,
             )
 
-        if not is_boot:
+        if is_boot:
+            # Replicate models are kept until the end of the run; drop their
+            # per-row data so memory does not grow with bootstrap_nboot
+            for m in models.values():
+                _slim_model(m)
+        else:
             self._outcome_start_params = {
                 k: (m.params.values, list(m.model.exog_names))
                 for k, m in models.items()

pySEQTarget/analysis/__init__.py

@@ -3,6 +3,7 @@
 from ._hazard import _calculate_hazard
 from ._outcome_fit import _outcome_fit
 from ._risk_estimates import _risk_estimates
+from ._slim_model import _slim_model
 from ._subgroup_fit import _subgroup_fit
 from ._survival_pred import (_calculate_survival, _clamp,
                              _get_outcome_predictions, _pred_risk)
@@ -16,6 +17,7 @@
     "_eof_summary",
     "_outcome_fit",
     "_risk_estimates",
+    "_slim_model",
     "_subgroup_fit",
     "_calculate_survival",
     "_clamp",

pySEQTarget/analysis/_slim_model.py

@@ -0,0 +1,74 @@
+import pickle
+import re
+
+import pandas as pd
+
+from ..helpers._fix_categories import _model_design_info
+from ._survival_pred import _same_state, _transform_state
+
+# Columns with at most this many distinct values contribute one row per value
+# to the reference frame, so the levels patsy infers from data are all present
+_MAX_LEVELS = 1000
+
+
+def _reference_frame(frame, formula):
+    """
+    A few rows of the formula's columns of ``frame``, holding every value of
+    each low-cardinality column.
+    """
+    names = set(re.findall(r"\w+", formula))
+    frame = frame[[c for c in frame.columns if c in names]]
+    keep = pd.Index(frame.index[:2])
+    for col in frame.columns:
+        values = frame[col]
+        if values.nunique(dropna=False) <= _MAX_LEVELS:
+            keep = keep.union(values.drop_duplicates().index)
+    return frame.loc[keep]
+
+
+def _same_design(a, b):
+    return list(a.column_names) == list(b.column_names) and _same_state(
+        _transform_state(a), _transform_state(b)
+    )
+
+
+def _slim_model(model):
+    """
+    Drop a fitted outcome model's per-row data in place, keeping what
+    prediction, ``bse``, ``summary()`` and pickling need.
+
+    One outcome model is kept per bootstrap replicate until the end of the run,
+    and a statsmodels GLM holds its full data frame, design matrix, weights and
+    fitted values — about 260 bytes per expanded row — so memory grew with the
+    expanded data times ``bootstrap_nboot``. The glum and jax wrappers keep the
+    design matrix and weights in the same way.
+    """
+    if hasattr(model, "_X_design"):
+        # glum/jax wrappers: cache the covariance before dropping its inputs
+        model.remove_data()
+        return model
+
+    # Fill summary()'s caches first; several of its statistics (e.g. the
+    # null-model log-likelihood) are computed from the data being dropped.
+    model.summary()
+    model.remove_data()
+
+    # remove_data() keeps the formula frame, which statsmodels re-evaluates the
+    # formula on when unpickling (process pool, offload, users saving results).
+    # Swap it for a few reference rows, but only if the design rebuilt from them
+    # matches the fitted one; stateful transforms such as a centred spline
+    # depend on all the data, so those models keep the full frame.
+    data = model.model.data
+    data.orig_endog = None
+    data.orig_exog = None
+    frame = data.frame
+    design_info = _model_design_info(model)
+    data.frame = _reference_frame(frame, str(data.formula))
+    try:
+        restored = _model_design_info(pickle.loads(pickle.dumps(model)))
+        ok = restored is not None and _same_design(restored, design_info)
+    except Exception:
+        ok = False
+    if not ok:
+        data.frame = frame
+    return model

pySEQTarget/error/_param_checker.py

@@ -134,6 +134,7 @@ def _param_checker(self):
     if (
         self.weighted
         and self.method != "ITT"
+        and self.weight_stabilized
         and self.numerator is not None
         and self.denominator is not None
     ):

pySEQTarget/helpers/_glum_fit.py

@@ -128,6 +128,12 @@ def predict(self, data, transform=True):
             X_arr = np.asarray(data)[:, 1:]
         return self._glum.predict(X_arr)
 
+    def remove_data(self):
+        """Drop the per-row design matrix and weights, as statsmodels' remove_data."""
+        self.cov_params()
+        self._X_design = None
+        self._sample_weight = None
+
     def cov_params(self):
         if self._cov_cached is not None:
             return self._cov_cached

pySEQTarget/helpers/_jax_fit.py

@@ -159,6 +159,14 @@ def predict(self, data, transform=True):
         probs = np.asarray(self._jax.predict(X_arr), dtype=np.float64)
         return probs[:, 1] if self._n_classes == 2 else probs
 
+    def remove_data(self):
+        """Drop the per-row design matrix and weights, as statsmodels' remove_data."""
+        # Covariance is only implemented for binary fits
+        if self._n_classes == 2:
+            self.cov_params()
+        self._X_design = None
+        self._sample_weight = None
+
     def cov_params(self):
         if self._n_classes != 2:
             raise NotImplementedError(

pySEQTarget/weighting/_weight_fit.py

@@ -102,6 +102,10 @@ def _fit_numerator(self, WDT):
         return
     if self.method == "ITT":
         return
+    if not self.weight_stabilized:
+        # Unstabilized weights: no numerator models, numerator fixed at 1
+        self.numerator_model = [None] * len(self.treatment_level)
+        return
     predictor = "switch" if self.excused else self.treatment_col
     tx_lag_col = (
         f"{self.treatment_col}{self.indicator_baseline}" if self.excused else "tx_lag"

pySEQTarget/weighting/_weight_pred.py

@@ -115,7 +115,7 @@ def _weight_predict(self, WDT):
                     .alias("denominator")
                 )
 
-            if self.weight_preexpansion:
+            if self.weight_preexpansion or not self.weight_stabilized:
                 WDT = WDT.with_columns(pl.lit(1.0).alias("numerator"))
             else:
                 for i, level in enumerate(self.treatment_level):

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pySEQTarget"
-version = "0.14.3"
+version = "0.14.4"
 description = "Sequential Target Trial Emulation"
 readme = "README.md"
 license = {text = "MIT"}
@@ -63,7 +63,8 @@ dev = [
 ]
 
 [project.urls]
-Homepage = "https://github.com/CausalInference/pySEQTarget"
+Homepage = "https://pyseqtarget.readthedocs.io"
+Documentation = "https://pyseqtarget.readthedocs.io"
 Repository = "https://github.com/CausalInference/pySEQTarget"
 "Bug Tracker" = "https://github.com/CausalInference/pySEQTarget/issues"
 "Ryan O'Dea (ORCID)" = "https://orcid.org/0009-0000-0103-9546"

tests/test_bootstrap_memory.py

@@ -0,0 +1,85 @@
+"""Bootstrap replicate outcome models keep no per-row data, so memory does not
+grow with the expanded data times bootstrap_nboot."""
+
+import pickle
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from pySEQTarget import SEQopts, SEQuential
+from pySEQTarget.data import load_data
+
+
+def _fit(**kwargs):
+    s = SEQuential(
+        load_data("SEQdata"),
+        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="censoring",
+        parameters=SEQopts(weighted=True, bootstrap_nboot=2, seed=7, **kwargs),
+    )
+    s.expand()
+    s.bootstrap()
+    s.fit()
+    return s
+
+
+@pytest.mark.filterwarnings("ignore")
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        dict(glm_package="statsmodels"),
+        dict(glm_package="glum"),
+        dict(glm_package="statsmodels", parallel=True, ncores=2),
+        dict(glm_package="statsmodels", offload=True),
+    ],
+)
+def test_replicate_models_keep_no_per_row_data(kwargs, tmp_path):
+    if kwargs.get("offload"):
+        kwargs["offload_dir"] = str(tmp_path)
+    s = _fit(**kwargs)
+    load = s._offloader.load_model
+    main = load(s.outcome_model[0]["outcome"])
+    reps = [load(m["outcome"]) for m in s.outcome_model[1:]]
+    assert len(reps) == 2
+
+    # ~100k expanded rows: a model holding per-row data pickles to megabytes
+    for m in reps:
+        assert len(pickle.dumps(m)) < 100_000
+
+    # The replicates still predict, give standard errors and summarise
+    df = s.DT.head(20).to_pandas()
+    df["tx_init_bas"] = pd.Categorical(df["tx_init_bas"], categories=[0, 1])
+    for m in reps:
+        assert len(m.predict(df)) == 20
+        assert m.bse.notna().all()
+        m.summary()
+
+    # Unpickling rebuilds the formula design from the reference rows kept in
+    # place of the data; predictions must not change
+    for m in reps:
+        restored = pickle.loads(pickle.dumps(m))
+        np.testing.assert_allclose(restored.predict(df), m.predict(df))
+
+    # The main fit is untouched
+    if kwargs["glm_package"] == "statsmodels" and not kwargs.get("offload"):
+        assert main.model.data.frame is not None
+
+
+@pytest.mark.filterwarnings("ignore")
+def test_stateful_transform_keeps_the_full_frame():
+    # center() memorises the mean of all the data, which a few reference rows
+    # cannot reproduce, so the frame is kept for unpickling
+    s = _fit(covariates="tx_init_bas+followup+trial+sex+center(N)")
+    m = s.outcome_model[1]["outcome"]
+    assert len(m.model.data.frame) > 10_000
+    df = s.DT.head(20).to_pandas()
+    df["tx_init_bas"] = pd.Categorical(df["tx_init_bas"], categories=[0, 1])
+    restored = pickle.loads(pickle.dumps(m))
+    np.testing.assert_allclose(restored.predict(df), m.predict(df))

tests/test_no_variation.py

@@ -10,6 +10,7 @@
 def _mock_self(**overrides):
     attrs = {
         "weight_preexpansion": False,
+        "weight_stabilized": True,
         "excused": False,
         "method": "censoring",
         "treatment_col": "tx",