diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml index ec4e173..9ff701b 100644 --- a/.github/workflows/python-app.yml +++ b/.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 }} diff --git a/CITATION.cff b/CITATION.cff index 9315375..f3cfe26 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.3 +version: 0.14.4 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 fdd2427..aa29019 100644 --- a/pySEQTarget/SEQopts.py +++ b/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", ] diff --git a/pySEQTarget/SEQuential.py b/pySEQTarget/SEQuential.py index f2f578b..368a8a9 100644 --- a/pySEQTarget/SEQuential.py +++ b/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() diff --git a/pySEQTarget/analysis/__init__.py b/pySEQTarget/analysis/__init__.py index 7e7097c..aa11373 100644 --- a/pySEQTarget/analysis/__init__.py +++ b/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", diff --git a/pySEQTarget/analysis/_slim_model.py b/pySEQTarget/analysis/_slim_model.py new file mode 100644 index 0000000..888d358 --- /dev/null +++ b/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 diff --git a/pySEQTarget/error/_param_checker.py b/pySEQTarget/error/_param_checker.py index a4b575b..a98519a 100644 --- a/pySEQTarget/error/_param_checker.py +++ b/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 ): diff --git a/pySEQTarget/helpers/_glum_fit.py b/pySEQTarget/helpers/_glum_fit.py index 4263467..fc490f5 100644 --- a/pySEQTarget/helpers/_glum_fit.py +++ b/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 diff --git a/pySEQTarget/helpers/_jax_fit.py b/pySEQTarget/helpers/_jax_fit.py index e7bbf63..a995f12 100644 --- a/pySEQTarget/helpers/_jax_fit.py +++ b/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( diff --git a/pySEQTarget/weighting/_weight_fit.py b/pySEQTarget/weighting/_weight_fit.py index a198f99..513aabf 100644 --- a/pySEQTarget/weighting/_weight_fit.py +++ b/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" diff --git a/pySEQTarget/weighting/_weight_pred.py b/pySEQTarget/weighting/_weight_pred.py index c6310f4..d568614 100644 --- a/pySEQTarget/weighting/_weight_pred.py +++ b/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): diff --git a/pyproject.toml b/pyproject.toml index 2cea093..99ed58e 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/test_bootstrap_memory.py b/tests/test_bootstrap_memory.py new file mode 100644 index 0000000..96e70c9 --- /dev/null +++ b/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)) diff --git a/tests/test_no_variation.py b/tests/test_no_variation.py index 86a607b..5f03091 100644 --- a/tests/test_no_variation.py +++ b/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", diff --git a/tests/test_unstabilized.py b/tests/test_unstabilized.py new file mode 100644 index 0000000..3e6a455 --- /dev/null +++ b/tests/test_unstabilized.py @@ -0,0 +1,108 @@ +"""Unstabilized IP weights: weight_stabilized=False fixes the treatment-weight +numerator at 1, so the weights are cumulative products of 1/denominator.""" + +import numpy as np +import polars as pl +import pytest + +from pySEQTarget import SEQopts, SEQuential +from pySEQTarget.data import load_data + +KEY = ["ID", "trial", "followup"] + + +def _run(data="SEQdata", method="censoring", fit=True, **kwargs): + s = SEQuential( + load_data(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=method, + parameters=SEQopts(**kwargs), + ) + if fit: + s.expand() + s.fit() + return s + + +def _merge(stab, unstab): + return stab.DT.select(KEY + ["numerator", "denominator"]).join( + unstab.DT.select(KEY + ["numerator", "denominator"]), + on=KEY, + suffix="_unstab", + ) + + +@pytest.mark.filterwarnings("ignore") +def test_unstabilized_weights_use_the_same_denominators(): + opts = dict(weighted=True, denominator="N+L+P+sex", seed=1636) + stab = _run(numerator="sex", **opts) + unstab = _run(weight_stabilized=False, **opts) + + merged = _merge(stab, unstab) + assert (merged["numerator_unstab"] == 1.0).all() + # Same denominator model, so identical denominators; only the numerator differs + np.testing.assert_allclose(merged["denominator"], merged["denominator_unstab"]) + + # Weights are the running product of 1/denominator within each trial + DT = unstab.DT.sort(KEY) + wt = DT["wt"].to_numpy() + inv_denom = 1.0 / DT["denominator"].to_numpy() + assert np.all(np.isclose(wt, 1.0) | np.isclose(wt, inv_denom)) + cum = DT.select(pl.col("wt").cum_prod().over(["ID", "trial"]))["wt"] + np.testing.assert_allclose(DT["weight"], cum) + + # No numerator model is fit, and the reported numerator formula is "1" + assert unstab.numerator == "1" + assert all(m is None for m in unstab.numerator_model) + assert unstab.collect().summary("numerator") == [] + + +def test_unstabilized_warns_where_it_cannot_apply(): + with pytest.warns(UserWarning, match="has no effect"): + _run(method="ITT", fit=False, weight_stabilized=False) + with pytest.warns(UserWarning, match="has no effect"): + _run(fit=False, weight_stabilized=False) + + +def test_unstabilized_warns_when_numerator_supplied(): + with pytest.warns(UserWarning, match="'numerator' is ignored"): + s = _run( + fit=False, + weighted=True, + weight_stabilized=False, + numerator="sex", + denominator="N+L+P+sex", + ) + assert s.numerator == "1" + + +def test_weight_stabilized_must_be_bool(): + with pytest.raises(TypeError, match="weight_stabilized"): + SEQopts(weight_stabilized="no") + + +@pytest.mark.filterwarnings("ignore") +@pytest.mark.parametrize( + "data, extra", + [ + ( + "SEQdata", + dict(excused=True, excused_colnames=["excusedZero", "excusedOne"]), + ), + ("SEQdata_multitreatment", dict(treatment_level=[0, 1, 2])), + ], +) +def test_unstabilized_keeps_denominators_excused_and_multinomial(data, extra): + opts = dict(weighted=True, weight_preexpansion=False, seed=1636, **extra) + stab = _run(data, **opts) + unstab = _run(data, weight_stabilized=False, **opts) + + merged = _merge(stab, unstab) + assert (merged["numerator_unstab"] == 1.0).all() + np.testing.assert_allclose(merged["denominator"], merged["denominator_unstab"])