diff --git a/.github/workflows/autoformat.yml b/.github/workflows/autoformat.yml index c194a09..4a87201 100644 --- a/.github/workflows/autoformat.yml +++ b/.github/workflows/autoformat.yml @@ -14,7 +14,7 @@ jobs: token: ${{ secrets.GITHUB_TOKEN }} - name: Set up Python - uses: actions/setup-python@v6 + uses: actions/setup-python@v7 with: python-version: '3.11' diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index ccb6bbd..2d24a6e 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -23,7 +23,7 @@ jobs: steps: - uses: actions/checkout@v7 - - uses: actions/setup-python@v6 + - uses: actions/setup-python@v7 with: python-version: "3.x" diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml index f104b35..2421ad6 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.0.1 + uses: astral-sh/setup-uv@v10.2.0 - name: Set up Python ${{ matrix.python-version }} run: uv python install ${{ matrix.python-version }} diff --git a/CITATION.cff b/CITATION.cff index 59d3f48..28d5beb 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.1 +version: 0.14.2 url: https://pyseqtarget.readthedocs.io/ repository-code: https://github.com/CausalInference/pySEQTarget repository-artifact: https://pypi.org/project/pySEQTarget/ diff --git a/docs/vignettes/more_advanced_models.md b/docs/vignettes/more_advanced_models.md index d31bd06..d281a45 100644 --- a/docs/vignettes/more_advanced_models.md +++ b/docs/vignettes/more_advanced_models.md @@ -108,6 +108,40 @@ my_options = SEQopts( The rest of the analytical pipeline is unchanged — `expand()`, `fit()`, `survival()`, and `collect()` all work exactly as before. +## Modelling the Baseline Hazard in the Weight Models with a Natural Cubic Spline + +The same choice arises in the weight models. By default they are quadratic in time: `followup`, `followup_sq`, `trial` and `trial_sq` when `weight_preexpansion=False`, and the time column and its square when `weight_preexpansion=True`. A quadratic only lets the hazard of treatment rise or flatten off, which can fit poorly when the cumulative incidence of treatment moves through several phases over followup. + +Setting `weight_spline=True` replaces each of those quadratics with a natural cubic spline basis (`cr()`) of `weight_spline_df` degrees of freedom, letting the baseline hazard take a flexible shape over time. This applies to the treatment weight models and, when `cense_colname` or `visit_colname` is set, to those models too. + +```python +my_options = SEQopts( + km_curves=True, + weighted=True, + weight_preexpansion=False, + weight_spline=True, # model time in the weight models as a natural cubic spline + weight_spline_df=4, # 4 degrees of freedom, i.e. 3 interior knots at percentiles +) +``` + +The default is unchanged (`weight_spline=False`), so existing analyses are unaffected. + +For finer control — a spline in `followup` but not `trial`, a different number of knots per term, or a spline in a time-varying confounder — write the `cr()` terms into `numerator` and `denominator` yourself. Any `cr(x, df=N)` term in a model formula, wherever it came from, has its knots fixed from the full data the model is fit on before fitting, so every bootstrap resample builds the same basis as the main fit. + +Pass `constraints="center"`, as the generated terms do. An unconstrained `cr()` basis spans the constant function, so it is collinear with the model intercept by exactly one dimension; statsmodels absorbs that with a pseudo-inverse (after reporting a convergence warning), but the `glum` backend cannot fit it at all. + +```python +my_options = SEQopts( + km_curves=True, + weighted=True, + weight_preexpansion=False, + numerator='sex+N_bas+L_bas+P_bas+trial+trial_sq+cr(followup, df=5, constraints="center")', + denominator='sex+N+L+P+N_bas+L_bas+P_bas+trial+trial_sq+cr(followup, df=5, constraints="center")', +) +``` + +After `fit()`, the `numerator` and `denominator` attributes show the formulas with their knots written out. + ## That's it? Yes! There are very few differences between the code for more straightforward and more difficult analyses using this package. Our hope is that through utilizing almost only the SEQopts to work with your analysis, that this is a streamlined process that is also easy to manipulate. diff --git a/pySEQTarget/SEQopts.py b/pySEQTarget/SEQopts.py index e6f542f..fdd2427 100644 --- a/pySEQTarget/SEQopts.py +++ b/pySEQTarget/SEQopts.py @@ -95,6 +95,14 @@ 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_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``. + Ignored for weight models given an explicit ``numerator``/``denominator``/ + ``cense_*`` formula — write ``cr(..., constraints="center")`` terms into those + directly + :param weight_spline_df: Degrees of freedom for ``cr()`` when ``weight_spline=True``, + at least ``2``. Default ``4`` :param verbose: Boolean to print dataset size summaries and bootstrap information :param weighted: Boolean to weight analysis """ @@ -156,6 +164,8 @@ class SEQopts: weight_lag_condition: bool = True weight_p99: bool = False weight_preexpansion: bool = True + weight_spline: bool = False + weight_spline_df: int = 4 verbose: bool = False weighted: bool = False @@ -177,6 +187,7 @@ def _validate_bools(self): "weight_lag_condition", "weight_p99", "weight_preexpansion", + "weight_spline", "weighted", ] for i in bools: @@ -242,6 +253,32 @@ def _validate_choices(self): if self.end_of_fup_type not in ["binary", "continuous"]: raise ValueError("end_of_fup_type must be 'binary' or 'continuous'") + def _validate_formulas(self): + # A list where a single formula is expected only fails much further + # downstream, so name the offending argument here. Only numerator and + # denominator may be per-arm lists (length checked in _param_checker). + for i in ("covariates", "cense_numerator", "cense_denominator"): + attr = getattr(self, i) + if attr is not None and not isinstance(attr, str): + raise TypeError( + f"{i} must be a single patsy formula string or None, but a " + f"{type(attr).__name__} was supplied." + ) + for i in ("numerator", "denominator"): + attr = getattr(self, i) + if attr is None or isinstance(attr, str): + continue + if not isinstance(attr, (list, tuple)): + raise TypeError( + f"{i} must be a patsy formula string, None, or one formula " + f"per treatment_level, but a {type(attr).__name__} was supplied." + ) + if len(attr) == 0: + raise ValueError( + f"{i} must be a patsy formula string, None, or one formula " + "per treatment_level, but an empty sequence was supplied." + ) + def _normalize_formulas(self): for i in ( "covariates", @@ -270,6 +307,7 @@ def __post_init__(self): self._validate_bools() self._validate_ranges() self._validate_choices() + self._validate_formulas() self._normalize_formulas() if self.offload: diff --git a/pySEQTarget/SEQuential.py b/pySEQTarget/SEQuential.py index eafb5d3..784d37f 100644 --- a/pySEQTarget/SEQuential.py +++ b/pySEQTarget/SEQuential.py @@ -14,7 +14,8 @@ _risk_estimates, _subgroup_fit) from .error import _data_checker, _param_checker from .expansion import _binder, _diagnostics, _dynamic, _random_selection -from .helpers import Offloader, _col_string, _format_time, bootstrap_loop +from .helpers import (Offloader, _bake_model_formulas, _col_string, + _format_time, bootstrap_loop) from .initialization import (_cense_denominator, _cense_numerator, _denominator, _numerator, _outcome) from .plot import _survival_plot @@ -274,6 +275,11 @@ def fit(self) -> None: if hasattr(self, "_current_boot_idx"): boot_idx = self._current_boot_idx + # Fix every cr(x, df = N) term's knots from the full data the model is + # fit on, so each replicate builds the same basis as the main fit. + if boot_idx is None: + _bake_model_formulas(self) + if self.weighted: # With weight_preexpansion the weight models are fit on the # un-resampled pre-expansion data, so every bootstrap replicate diff --git a/pySEQTarget/analysis/_outcome_fit.py b/pySEQTarget/analysis/_outcome_fit.py index 9c13656..99cf945 100644 --- a/pySEQTarget/analysis/_outcome_fit.py +++ b/pySEQTarget/analysis/_outcome_fit.py @@ -1,32 +1,17 @@ import re -import numpy as np import pandas as pd import polars as pl import statsmodels.api as sm import statsmodels.formula.api as smf from pandas.api.types import is_numeric_dtype - -def _compute_spline_knots(followup_arr, df=3): - lower = float(np.min(followup_arr)) - upper = float(np.max(followup_arr)) - n_inner = df - 2 - if n_inner == 0: - inner_knots = [] - else: - # Replicate patsy's knot placement: percentiles of unique values in [lower, upper] - x = np.unique(followup_arr[(lower <= followup_arr) & (followup_arr <= upper)]) - q = np.linspace(0, 100, n_inner + 2)[1:-1] - inner_knots = np.percentile(x, q.tolist()).tolist() - return inner_knots, lower, upper +from ..helpers._spline import _compute_spline_knots, _cr_term def _apply_spline_formula(formula, indicator_squared, spline_knots): inner_knots, lower, upper = spline_knots - spline = ( - f"cr(followup, knots={inner_knots}, lower_bound={lower}, upper_bound={upper})" - ) + spline = _cr_term("followup", inner_knots, lower, upper) formula = re.sub(r"(\w+)\s*\*\s*followup\b", rf"\1*{spline}", formula) formula = re.sub(r"\bfollowup\s*\*\s*(\w+)", rf"{spline}*\1", formula) diff --git a/pySEQTarget/error/_param_checker.py b/pySEQTarget/error/_param_checker.py index 4ef1637..a4b575b 100644 --- a/pySEQTarget/error/_param_checker.py +++ b/pySEQTarget/error/_param_checker.py @@ -90,6 +90,9 @@ def _param_checker(self): if self.followup_spline_df < 2: raise ValueError("followup_spline_df must be at least 2.") + if not isinstance(self.weight_spline_df, int) or self.weight_spline_df < 2: + raise ValueError("weight_spline_df must be an integer of at least 2.") + if ( self.weighted and self.method == "ITT" diff --git a/pySEQTarget/helpers/__init__.py b/pySEQTarget/helpers/__init__.py index c0c867d..2244d71 100644 --- a/pySEQTarget/helpers/__init__.py +++ b/pySEQTarget/helpers/__init__.py @@ -6,6 +6,8 @@ from ._pad import _pad from ._predict_model import _predict_model from ._prepare_data import _prepare_data +from ._spline import (_bake_model_formulas, _bake_spline_knots, + _compute_spline_knots) __all__ = [ "bootstrap_loop", @@ -17,4 +19,7 @@ "_pad", "_predict_model", "_prepare_data", + "_bake_model_formulas", + "_bake_spline_knots", + "_compute_spline_knots", ] diff --git a/pySEQTarget/helpers/_col_string.py b/pySEQTarget/helpers/_col_string.py index 8c67709..1f6c07c 100644 --- a/pySEQTarget/helpers/_col_string.py +++ b/pySEQTarget/helpers/_col_string.py @@ -1,3 +1,26 @@ +import ast + + +def _term_columns(expression): + """Column names a patsy right-hand-side formula refers to (function-call + terms like ``cr(followup, df=4)`` contribute the variable they wrap).""" + try: + tree = ast.parse(expression, mode="eval") + except SyntaxError: + return set(expression.replace("+", " ").replace("*", " ").split()) + + called = { + node.func + for node in ast.walk(tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + return { + node.id + for node in ast.walk(tree) + if isinstance(node, ast.Name) and node not in called + } + + def _col_string(expressions): cols = set() for expression in expressions: @@ -8,5 +31,5 @@ def _col_string(expressions): parts = expression if isinstance(expression, (list, tuple)) else [expression] for part in parts: if part is not None: - cols.update(part.replace("+", " ").replace("*", " ").split()) + cols.update(_term_columns(part)) return cols diff --git a/pySEQTarget/helpers/_spline.py b/pySEQTarget/helpers/_spline.py new file mode 100644 index 0000000..2333e2e --- /dev/null +++ b/pySEQTarget/helpers/_spline.py @@ -0,0 +1,108 @@ +import re + +import numpy as np + +# The only cr() form whose basis depends on the rows, so the only one to bake. +_CR_DF = re.compile( + r"""\bcr\(\s*(?P[A-Za-z_.][A-Za-z0-9_.]*)\s*,\s*df\s*=\s*(?P\d+)\s* + (?:,\s*constraints\s*=\s*(?P["'])(?Pcenter)(?P=quote)\s*)?\)""", + re.VERBOSE, +) + + +def _compute_spline_knots(arr, df=3): + """Knots patsy's ``cr()`` places over ``arr``: ``df - 2`` interior knots at + percentiles of the unique values and two boundary knots.""" + lower = float(np.min(arr)) + upper = float(np.max(arr)) + n_inner = df - 2 + if n_inner <= 0: + inner_knots = [] + else: + # Replicate patsy's knot placement: percentiles of unique values in [lower, upper] + x = np.unique(arr[(lower <= arr) & (arr <= upper)]) + q = np.linspace(0, 100, n_inner + 2)[1:-1] + inner_knots = np.percentile(x, q.tolist()).tolist() + return inner_knots, lower, upper + + +def _cr_term(var, inner_knots, lower, upper, constraints=None): + """A ``cr()`` term with its knots written out.""" + constraint = f', constraints="{constraints}"' if constraints else "" + return ( + f"cr({var}, knots={list(inner_knots)}, " + f"lower_bound={lower}, upper_bound={upper}{constraint})" + ) + + +def _bake_spline_knots(formula, data): + """ + Rewrite every ``cr(x, df=N)`` in ``formula`` to explicit knots taken from + ``data``. ``formula``""" + if formula is None: + return formula + if isinstance(formula, (list, tuple)): + return [_bake_spline_knots(part, data) for part in formula] + + tokens = { + match.group(0): ( + match.group("var"), + int(match.group("df")), + match.group("constraints"), + ) + for match in _CR_DF.finditer(formula) + } + if not tokens: + return formula + + for token, (var, df, constraints) in tokens.items(): + if var not in data.columns: + continue + column = data[var] + if not column.dtype.is_numeric(): + continue + arr = column.drop_nulls().to_numpy() + if arr.size == 0: + continue + + # A centering constraint absorbs one df, so patsy adds an interior knot. + inner_knots, lower, upper = _compute_spline_knots( + arr, df=df + 1 if constraints else df + ) + # Too few distinct values ties knots to each other or to a boundary; + # patsy rejects that basis, without naming the term that caused it. + if len(set(inner_knots)) < len(inner_knots) or not all( + lower < knot < upper for knot in inner_knots + ): + raise ValueError( + f"'{var}' has too few distinct values to place " + f"{len(inner_knots)} interior knot(s) for cr(df={df}); use a " + "smaller df, or model this term without a spline." + ) + formula = formula.replace( + token, _cr_term(var, inner_knots, lower, upper, constraints) + ) + + return formula + + +def _bake_model_formulas(self): + """ + Fix every ``cr(x, df=N)`` term's knots from the full data each model is fit + on, so bootstrap resamples and the unpickled glum offload/parallel paths all + build the same basis as the main fit. + """ + if not self.followup_spline: + # followup_spline rewrites the outcome's followup terms (with fixed + # knots) at fit time, and would mangle a cr() term baked in here. + self.covariates = _bake_spline_knots(self.covariates, self.DT) + + if self.weighted: + weight_data = self.data if self.weight_preexpansion else self.DT + for name in ( + "numerator", + "denominator", + "cense_numerator", + "cense_denominator", + ): + setattr(self, name, _bake_spline_knots(getattr(self, name), weight_data)) diff --git a/pySEQTarget/initialization/_censoring.py b/pySEQTarget/initialization/_censoring.py index 828d584..38eeac2 100644 --- a/pySEQTarget/initialization/_censoring.py +++ b/pySEQTarget/initialization/_censoring.py @@ -1,15 +1,10 @@ +from ._time_terms import _time_terms + + def _cense_numerator(self) -> str: - trial = ( - "+".join(["trial", f"trial{self.indicator_squared}"]) - if self.trial_include - else None - ) - followup = ( - "+".join(["followup", f"followup{self.indicator_squared}"]) - if self.followup_include - else None - ) - time = "+".join([self.time_col, f"{self.time_col}{self.indicator_squared}"]) + trial = _time_terms(self, "trial") if self.trial_include else None + followup = _time_terms(self, "followup") if self.followup_include else None + time = _time_terms(self, self.time_col) tv_bas = ( "+".join([f"{v}{self.indicator_baseline}" for v in self.time_varying_cols]) if self.time_varying_cols @@ -26,17 +21,9 @@ def _cense_numerator(self) -> str: def _cense_denominator(self) -> str: - trial = ( - "+".join(["trial", f"trial{self.indicator_squared}"]) - if self.trial_include - else None - ) - followup = ( - "+".join(["followup", f"followup{self.indicator_squared}"]) - if self.followup_include - else None - ) - time = "+".join([self.time_col, f"{self.time_col}{self.indicator_squared}"]) + trial = _time_terms(self, "trial") if self.trial_include else None + followup = _time_terms(self, "followup") if self.followup_include else None + time = _time_terms(self, self.time_col) tv = "+".join(self.time_varying_cols) if self.time_varying_cols else None tv_bas = ( "+".join([f"{v}{self.indicator_baseline}" for v in self.time_varying_cols]) diff --git a/pySEQTarget/initialization/_denominator.py b/pySEQTarget/initialization/_denominator.py index 0a081b3..5ce67e1 100644 --- a/pySEQTarget/initialization/_denominator.py +++ b/pySEQTarget/initialization/_denominator.py @@ -1,17 +1,12 @@ +from ._time_terms import _time_terms + + def _denominator(self) -> str: if self.method == "ITT": return - trial = ( - "+".join(["trial", f"trial{self.indicator_squared}"]) - if self.trial_include - else None - ) - followup = ( - "+".join(["followup", f"followup{self.indicator_squared}"]) - if self.followup_include - else None - ) - time = "+".join([self.time_col, f"{self.time_col}{self.indicator_squared}"]) + trial = _time_terms(self, "trial") if self.trial_include else None + followup = _time_terms(self, "followup") if self.followup_include else None + time = _time_terms(self, self.time_col) tv = "+".join(self.time_varying_cols) if self.time_varying_cols else None tv_bas = ( diff --git a/pySEQTarget/initialization/_numerator.py b/pySEQTarget/initialization/_numerator.py index b47ba23..07ab1fa 100644 --- a/pySEQTarget/initialization/_numerator.py +++ b/pySEQTarget/initialization/_numerator.py @@ -1,17 +1,12 @@ +from ._time_terms import _time_terms + + def _numerator(self) -> str: if self.method == "ITT": return - trial = ( - "+".join(["trial", f"trial{self.indicator_squared}"]) - if self.trial_include - else None - ) - followup = ( - "+".join(["followup", f"followup{self.indicator_squared}"]) - if self.followup_include - else None - ) - time = "+".join([self.time_col, f"{self.time_col}{self.indicator_squared}"]) + trial = _time_terms(self, "trial") if self.trial_include else None + followup = _time_terms(self, "followup") if self.followup_include else None + time = _time_terms(self, self.time_col) tv_bas = ( "+".join([f"{v}{self.indicator_baseline}" for v in self.time_varying_cols]) diff --git a/pySEQTarget/initialization/_time_terms.py b/pySEQTarget/initialization/_time_terms.py new file mode 100644 index 0000000..ebdd0df --- /dev/null +++ b/pySEQTarget/initialization/_time_terms.py @@ -0,0 +1,9 @@ +def _time_terms(self, var: str) -> str: + """ + Default weight-model terms for a time axis ``var``: a quadratic + (``var``, ``var_sq``), or a centred ``cr()`` spline of ``weight_spline_df`` + df when ``weight_spline`` is set. Knots are fixed later from the fit data. + """ + if self.weight_spline: + return f'cr({var}, df={self.weight_spline_df}, constraints="center")' + return "+".join([var, f"{var}{self.indicator_squared}"]) diff --git a/pyproject.toml b/pyproject.toml index 44e87be..592513b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pySEQTarget" -version = "0.14.1" +version = "0.14.2" description = "Sequential Target Trial Emulation" readme = "README.md" license = {text = "MIT"} diff --git a/tests/test_formula_arguments.py b/tests/test_formula_arguments.py new file mode 100644 index 0000000..74accdb --- /dev/null +++ b/tests/test_formula_arguments.py @@ -0,0 +1,59 @@ +"""SEQopts formula arguments: the ones that are a single formula by +construction are checked where they are supplied, rather than failing further +downstream in the column extraction or the patsy formula build. +""" + +import pytest + +from pySEQTarget import SEQopts +from pySEQTarget.helpers._col_string import _col_string + + +@pytest.mark.parametrize("name", ["covariates", "cense_numerator", "cense_denominator"]) +def test_single_formula_arguments_reject_a_sequence(name): + with pytest.raises(TypeError, match=f"{name} must be a single patsy formula"): + SEQopts(**{name: ["sex", "N"]}) + + +@pytest.mark.parametrize("name", ["numerator", "denominator"]) +def test_weight_formulas_reject_an_empty_sequence(name): + with pytest.raises(ValueError, match="empty sequence"): + SEQopts(**{name: []}) + + +@pytest.mark.parametrize("name", ["numerator", "denominator"]) +def test_weight_formulas_reject_a_non_formula(name): + with pytest.raises(TypeError, match=f"{name} must be a patsy formula"): + SEQopts(**{name: 3}) + + +def test_valid_formula_arguments_are_accepted(): + # numerator/denominator may still be one formula per treatment_level, + # validated against treatment_level in _param_checker + opts = SEQopts( + covariates="sex+N", + numerator=["sex", "sex+N"], + denominator=["sex+L", "sex+N+L"], + ) + assert opts.covariates == "sex+N" + assert opts.numerator == ["sex", "sex+N"] + + +def test_col_string_reads_through_function_call_terms(): + # The expansion works out which columns to carry from the model formulas, + # so a term that wraps a column has to contribute the column itself + assert _col_string(["sex+N_bas+followup+followup_sq"]) == { + "sex", + "N_bas", + "followup", + "followup_sq", + } + assert _col_string(["cr(followup, df=4)+sex"]) == {"followup", "sex"} + assert _col_string( + ["cr(followup, knots=[1.0, 2.0], lower_bound=0.0, upper_bound=60.0)"] + ) == {"followup"} + assert _col_string(["tx_init_bas*followup"]) == {"tx_init_bas", "followup"} + # An intercept-only weight model refers to no columns at all + assert _col_string(["1", "", None]) == set() + # Per-treatment-level formulas contribute the union of their columns + assert _col_string([["sex", "sex+N"]]) == {"sex", "N"} diff --git a/tests/test_weight_spline.py b/tests/test_weight_spline.py new file mode 100644 index 0000000..61c1996 --- /dev/null +++ b/tests/test_weight_spline.py @@ -0,0 +1,298 @@ +"""weight_spline: model the baseline hazard in the weight models as a flexible +function of time, and fix the knots of every cr(x, df=N) term — in the weight +models or the outcome model — from the data that model is fit on, so the basis +is constant across bootstrap resamples. +""" + +import re + +import numpy as np +import patsy +import polars as pl +import pytest + +from pySEQTarget import SEQopts, SEQuential +from pySEQTarget.data import load_data +from pySEQTarget.helpers._spline import _bake_spline_knots +from pySEQTarget.initialization import _cense_denominator, _cense_numerator +from pySEQTarget.initialization._denominator import _denominator +from pySEQTarget.initialization._numerator import _numerator + +# Two expected warnings, neither marking a bad fit: _check_separation trips on the +# large (but precise, z ~ 1.8) coefficient of a small-scale cr() basis column, and +# an unconstrained hand-written cr() is collinear with the intercept, which pinv +# resolves after newton reports it. +pytestmark = [ + pytest.mark.filterwarnings("ignore:Possible perfect or quasi-complete"), + pytest.mark.filterwarnings("ignore:Maximum Likelihood optimization failed"), +] + +BAKED_PLAIN = r"cr\({var}, knots=\[[^]]*\], lower_bound=[-\d.]+, upper_bound=[-\d.]+" +BAKED = BAKED_PLAIN + r', constraints="center"\)' + + +def _model(**opts): + 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, seed=42, **opts), + ) + s.expand() + s.fit() + return s + + +class _Opts: + """Minimal stand-in for the formula-building attributes of SEQuential.""" + + def __init__(self, **kwargs): + defaults = dict( + method="censoring", + excused=False, + time_col="time", + time_varying_cols=["N", "L", "P"], + fixed_cols=["sex"], + indicator_baseline="_bas", + indicator_squared="_sq", + trial_include=True, + followup_include=True, + weight_preexpansion=True, + weight_spline=False, + weight_spline_df=4, + ) + defaults.update(kwargs) + for key, value in defaults.items(): + setattr(self, key, value) + + +def _default_formulas(**kwargs): + opts = _Opts(**kwargs) + return [ + _numerator(opts), + _denominator(opts), + _cense_numerator(opts), + _cense_denominator(opts), + ] + + +def test_weight_spline_replaces_the_time_quadratics(): + post = _default_formulas(weight_spline=True, weight_preexpansion=False) + # Post-expansion weight models run on two time axes: both become splines + assert all('cr(followup, df=4, constraints="center")' in f for f in post) + assert all('cr(trial, df=4, constraints="center")' in f for f in post) + assert not any(re.search(r"followup_sq|trial_sq", f) for f in post) + + pre = _default_formulas(weight_spline=True, weight_preexpansion=True) + # Pre-expansion weight models only see the subject's own time column + assert all('cr(time, df=4, constraints="center")' in f for f in pre) + assert not any("time_sq" in f for f in pre) + + +def test_weight_spline_off_leaves_the_weight_models_quadratic(): + covs = _default_formulas(weight_preexpansion=False)[1] + assert "cr(" not in covs + assert {"followup", "followup_sq", "trial", "trial_sq"} <= set(covs.split("+")) + + +def test_weight_spline_df_is_passed_through(): + covs = _default_formulas(weight_spline=True, weight_spline_df=6)[1] + assert 'cr(time, df=6, constraints="center")' in covs + + +def test_weight_spline_df_validation(): + for df in (1, 0, -1): + with pytest.raises(ValueError, match="weight_spline_df"): + _model(weight_spline=True, weight_spline_df=df) + + +def test_weight_model_knots_are_baked_from_the_data_fit_on(): + s = _model(weight_preexpansion=False, weight_spline=True) + + for covs in (s.numerator, s.denominator): + assert re.search(BAKED.format(var="followup"), covs) + assert re.search(BAKED.format(var="trial"), covs) + assert "df=4" not in covs + + # The percentiles patsy would place over the post-expansion data the weight + # models are fit on, df - 1 = 3 of them once the basis is centred + followup = s.DT["followup"].to_numpy() + expected = np.percentile(np.unique(followup), [25, 50, 75]).tolist() + knots = re.search(r"cr\(followup, knots=\[([^]]*)\]", s.denominator).group(1) + assert [float(k) for k in knots.split(",")] == pytest.approx(expected) + assert f"lower_bound={float(followup.min())}" in s.denominator + assert f"upper_bound={float(followup.max())}" in s.denominator + + +def test_preexpansion_knots_come_from_the_preexpansion_data(): + s = _model(weight_preexpansion=True, weight_spline=True) + time = s.data["time"].to_numpy() + assert f"upper_bound={float(time.max())}" in s.denominator + # followup tops out one below the pre-expansion time column, so the two + # would be indistinguishable if the wrong frame were used + assert float(time.max()) != float(s.DT["followup"].max()) + + +def test_user_supplied_df_form_terms_are_baked(): + s = _model( + weight_preexpansion=False, + numerator="sex+cr(followup, df=3)", + denominator="sex+N+L+P+cr(followup, df=3)", + ) + # A hand-written term is baked as written — the centering constraint is + # only added to the ones weight_spline generates + assert re.search(BAKED_PLAIN.format(var="followup") + r"\)", s.denominator) + assert "df=3" not in s.denominator + + +def test_arm_specific_weight_models_each_get_their_knots_baked(): + s = _model( + weight_preexpansion=False, + numerator=["sex+cr(followup, df=3)", "sex+cr(followup, df=5)"], + denominator=["sex+N+L+P+cr(followup, df=3)", "sex+N+L+P+cr(followup, df=5)"], + ) + assert len(s.denominator) == 2 + assert not any("df=" in covs for covs in s.denominator) + # unconstrained: df = 3 gives 1 interior knot, df = 5 gives 3 + n_knots = [ + len(re.search(r"knots=\[([^]]*)\]", covs).group(1).split(",")) + for covs in s.denominator + ] + assert n_knots == [1, 3] + + +def test_baked_basis_matches_the_df_form_on_the_full_data(): + # Baking has to reproduce patsy's own knot placement exactly, or the + # weights would shift purely because the term was rewritten. + s = _model(weight_preexpansion=False, weight_spline=True) + DT = s.DT.to_pandas() + baked = patsy.dmatrix( + re.search(BAKED.format(var="followup"), s.denominator).group(0), DT + ) + df_form = patsy.dmatrix('cr(followup, df=4, constraints="center")', DT) + assert np.allclose(np.asarray(baked), np.asarray(df_form)) + + +def test_weights_are_unchanged_by_the_default_options(): + # weight_spline is opt-in: the default weight models must be untouched + baseline = _model(weight_preexpansion=False) + assert "cr(" not in baseline.denominator + spline = _model(weight_preexpansion=False, weight_spline=True) + assert not np.allclose( + baseline.DT["weight"].to_numpy(), spline.DT["weight"].to_numpy() + ) + + +def test_bootstrap_replicates_share_the_main_fit_basis(): + 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, + weight_preexpansion=False, + weight_spline=True, + bootstrap_nboot=3, + bootstrap_sample=0.8, + seed=42, + ), + ) + s.expand() + s.bootstrap() + s.fit() + baked = s.denominator + assert "df=" not in baked + # Refitting (the replicates run through the same body) leaves the basis + # fixed at the one the main fit chose + s.fit() + assert s.denominator == baked + + +def test_expansion_keeps_the_columns_a_spline_term_wraps(): + # The expansion works out which columns to carry from the model formulas; + # a cr() term has to contribute its variable, not fragments of the call. + s = _model( + weight_preexpansion=False, + numerator="sex+cr(N_bas, df=3)", + denominator="sex+N+L+P+cr(N_bas, df=3)", + ) + assert "N_bas" in s.DT.columns + + +def test_bake_leaves_terms_it_cannot_bake_alone(): + DT = pl.DataFrame( + {"followup": range(101), "grp": ["a", "b"] * 50 + ["a"]}, + ) + # Variable absent from the data + assert _bake_spline_knots("cr(missing, df=4)", DT) == "cr(missing, df=4)" + # Non-numeric variable + assert _bake_spline_knots("cr(grp, df=4)", DT) == "cr(grp, df=4)" + # Already-explicit knots + explicit = "cr(followup, knots=[1.0, 2.0], lower_bound=0.0, upper_bound=100.0)" + assert _bake_spline_knots(explicit, DT) == explicit + # None and per-arm lists pass through in the same shape + assert _bake_spline_knots(None, DT) is None + assert _bake_spline_knots(["sex", "cr(followup, df=3)"], DT)[0] == "sex" + # Several terms in one formula, each baked against its own variable + baked = _bake_spline_knots("cr(followup, df=3)+cr(followup, df=5)", DT) + assert baked.count("lower_bound") == 2 + assert "df=" not in baked + + +def test_bake_rejects_a_term_with_too_few_distinct_values(): + # Every knot lands on the single value, which patsy would reject with an + # error naming neither the term nor the reason + DT = pl.DataFrame({"x": [5.0] * 50}) + with pytest.raises(ValueError, match="too few distinct values"): + _bake_spline_knots("cr(x, df=4)", DT) + + +@pytest.mark.parametrize("glm_package", ["statsmodels", "glum"]) +def test_weight_spline_fits_on_either_glm_backend(glm_package): + # The basis is centred, which keeps it full rank alongside the model + # intercept — glum's solver, unlike statsmodels' pinv, needs that + s = _model(weight_preexpansion=False, weight_spline=True, glm_package=glm_package) + assert 'constraints="center"' in s.denominator + assert s.DT["weight"].to_numpy().min() > 0 + + +def test_user_supplied_outcome_spline_is_baked(): + 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="ITT", + parameters=SEQopts( + covariates=( + "tx_init_bas+trial+trial_sq+sex+N_bas+L_bas+P_bas" + '+cr(followup, df=4, constraints="center")' + ), + km_curves=True, + seed=42, + ), + ) + s.expand() + s.fit() + s.survival() + assert re.search(BAKED.format(var="followup"), s.covariates) + assert "df=4" not in s.covariates + # df = 4 centred contributes 4 columns to the outcome model + coef_names = list(s.outcome_model[0]["outcome"].params.index) + assert len([n for n in coef_names if n.startswith("cr(followup,")]) == 4