Add modelling the baseline hazard in the weight models as a flexible function of time

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

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

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

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/

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.

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:

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

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)

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"

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",
 ]

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

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<var>[A-Za-z_.][A-Za-z0-9_.]*)\s*,\s*df\s*=\s*(?P<df>\d+)\s*
+        (?:,\s*constraints\s*=\s*(?P<quote>["'])(?P<constraints>center)(?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))