shap-recommender

Exclusion / non-linearity / interaction recommendations from saved SHAP attribution files, and application of them to a design matrix.

One shared idea runs through two of the three rules. A feature's own contribution is represented flexibly (indicator columns when it takes few values, a restricted cubic spline when it is continuous), and a model built on that flexible basis is compared against a straight line in the feature. That comparison answers two different questions:

• non-linearity -- does the attribution deviate from a linear function of the feature? This directly tests the linear-trend assumption, rather than relying on a raw correlation coefficient, which conflates "no effect" with "non-linear effect".
• interaction -- does the attribution vary among subjects who share the same feature value? Under additivity the attribution is a deterministic function of the feature, so residual dispersion implies effect modification.

The interaction screen routes each candidate partner to the statistic that is valid for its type:

• Binary partners go to a within-stratum contrast. The stratifying feature is split at a pre-specified cut point (never on the attribution itself, which would condition on the candidate modifier), and E[phi_y | y=1] - E[phi_y | y=0] is compared across strata. For a binary partner that two-point contrast equals g(1) - g(0) exactly, so it is free of the within-stratum distribution of y.
• Continuous and ordinal partners go to a product-term regression, with no stratification at all: the stratifying feature's own attribution is regressed on a flexible basis in x, the (standardised) partners, and the products partner_z * x_z, and the product coefficient is the modification signal. A within-stratum OLS slope is Cov(y, h(y)) / Var(y) under additivity, which moves with the distribution of y inside the stratum whenever h is non-linear, so comparing that slope across strata can declare an interaction where none exists. Any other binary covariate in the data is folded in as an adjustment term (with its own product term) so a reported partner correlated with it isn't credited with an interaction that is really the covariate's.

Both branches report a two-sided Wald p-value and 95% CI. se_method chooses how the standard error behind that Wald test is obtained:

• 'sandwich' (the default) -- an HC3 heteroskedasticity-robust standard error, computed analytically from a single fit. A patient-level pairs bootstrap of an OLS coefficient converges to the (HC0) sandwich SE, so this estimates the same quantity without Monte Carlo error: on data of this design the bootstrap and sandwich SEs agreed to within 1%, and the screen ran about 170x faster. HC3 rather than HC0: HC0 is anticonservative when the design is wide relative to the sample (type-I error 0.21 with 200 rows and 31 columns in our tests), while HC3 held the nominal level at every size checked.
• 'bootstrap' -- the standard error is instead the bootstrap standard deviation of the estimate (n_boot_screen resamples, minimum 30, recommended >= 300). Degenerate or non-finite resamples are excluded rather than silently biasing the estimate, and a column with too few usable resamples is reported as non-significant instead of assigned a spurious p-value of zero.

Both branches are also flagged against min_abs_effect applied to the effect scaled by the typical attribution magnitude (phi_scale) rather than by the estimate's own magnitude, since a ratio to the estimate itself blows up whenever it sits near zero. The test table records which statistic was used for each pair in a method column and its p-value in p_screen.

Note on df_spline for the product-term screen. The product column partner_z * x_z is close to a quadratic in x whenever the partner is itself close to a linear function of x, so it competes with the basis for whatever curvature in h(x) the basis cannot represent. With a strongly collinear partner and the default df_spline=4, that leftover curvature can inflate the product coefficient's significance and flag an additive design as an interaction; df_spline=6 removes it in our tests. Standardising the inputs keeps the leaked coefficient itself small, so the practical risk is a miscalibrated p-value rather than a large spurious effect size, but the p-value is exactly what the FDR step consumes. Consider df_spline >= 6 when screening partners strongly correlated with the stratifying feature.

Two-reference subcohort design

Recommender.generate screens a "low" and a "high" reference subcohort separately and combines the results, because each subcohort's attribution is a contrast against its own reference point -- pooling the two before screening would average those contrasts away, which is what a two-reference design exists to avoid:

• exclusion -- intersection of the two subcohorts. A feature is dropped only if it is negligible in both; negligible in just one is not grounds for removing it from the model.
• non-linearity -- union of the two subcohorts. Curvature visible in one subcohort is still curvature.
• interaction -- union, with multiplicity controlled once over the pooled tests from both subcohorts (Recommender.merge_interaction_candidates), rather than twice over each half. A pair can be tested up to four times (either feature as the stratifier, in either subcohort); those replicate p-values are combined per pair by Bonferroni (valid under any dependence among the replicates), and BH is then applied once across unique pairs.

Recommender.top_k_spec(tests, k) picks the k selected pairs with the largest effect as a fixed-size interaction spec, for when the number of interaction pairs entering a model must be fixed in advance, without reference to a held-out test set.

Recommender.spec_within_budget(tests, base_frame, budget) does the same thing but fixes the number of parameter columns instead: it adds selected pairs, in order of effect size, until the columns they actually expand into in base_frame (one-hot/indicator-expanded features can cost more than one column each) would exceed budget. The two only agree when every pair costs the same number of terms; spec_within_budget is the one to use when you need to cap the design's total size rather than its pair count.

Choosing the budget: riley_parameter_budget

Recommender.riley_parameter_budget(X, y, shrinkage=0.9) sizes that budget (or top_k_spec's k) from the data itself, rather than picking a number by hand. It implements criterion (i) of Riley et al. (Stat Med 2019;38:1276-96): the largest number of predictor parameters for which the expected global shrinkage factor stays at or above shrinkage, given a Cox model's sample size and its (heuristically shrunk) Cox-Snell R^2. X should be the design before any interaction terms are added, so the budget errs toward fewer terms; y is a mapping with 'time' and 'event' entries. It returns the apparent and shrinkage-adjusted R^2, the headline p_max, and p_max at a small sensitivity grid of S values (the one constant the bound depends on), since that grid makes the sensitivity of the budget to S visible rather than hidden behind a single number.

Riley's criterion is defined on candidate parameters, so it is meant to size a budget for spec_within_budget/top_k_spec up front -- once terms have actually been selected from a larger candidate set (as merge_interaction_candidates does), the resulting model is more optimistic than this bound implies, and only a held-out evaluation, not this criterion, guards against that.

This method needs lifelines, an optional dependency (pip install "shap-recommender[survival]"); it is imported lazily, so installing it is not required to use the rest of the package.

Install

pip install shap-recommender

# to also use Recommender.riley_parameter_budget:
pip install "shap-recommender[survival]"

Expected input files

For each dataset "tag" you want to load, Recommender.load(tag) (and the CLI's --low-tag / --high-tag) expects two tab-separated files in res_dir:

• shap_values_<tag>.tsv -- SHAP values, one row per subject, one column per feature, first column = row index.
• sel_data_<tag>.tsv -- the corresponding feature values (design matrix), same row index.

Command-line use

shap-recommender \
    --res-dir ./shap_results \
    --low-tag cohort_low --high-tag cohort_high \
    --nonlinear-candidates age bmi creatinine \
    --out ./recommendations

This writes, per subcohort, exclusion_tests_{low,high}.tsv, exclusion_sensitivity_{low,high}.tsv, nonlinear_tests_{low,high}.tsv, and attribution_patterns_{low,high}.tsv, plus the pooled interaction_tests.tsv, a combined attribution_patterns.tsv, and recommendations.json to --out. Run shap-recommender --help for all options (thresholds, --se-method sandwich|bootstrap, bootstrap count, spline degrees of freedom, a --cutpoints JSON file for pre-specified stratification cut points, etc).

Library use

from shap_recommender import Recommender

rec = Recommender(res_dir="./shap_results")
recommendations = rec.generate(
    candidates_nonlinear=["age", "bmi", "creatinine"],
    low_tag="cohort_low", high_tag="cohort_high",
    out="./recommendations",
)

# apply the recommendations to a design matrix
X_train_adj, X_test_adj = Recommender.apply(
    X_train, X_test, recommendations, variant="all",
)

Recommender.apply(..., variant=...) accepts "baseline", "exclusion", "nonlinear", "interaction", or "all", so each rule's effect on downstream model performance can be evaluated separately.

A feature flagged non-linear is expanded one of two ways, chosen by how many distinct values it takes in the training data (n_bins, default 10): a low-cardinality (ordinal) feature is expanded into per-level indicator columns (<feature>_lvl<value>, one per level after a reference level), since a quadratic in the level code cannot represent an arbitrary threshold effect; a feature with more distinct values than n_bins is centred on its training mean and given a <feature>_quad column instead. An interaction partner that is itself such an expanded ordinal feature attaches to every one of its indicator columns, rather than being dropped for a missing main effect.

If you want to run a single subcohort's screen directly (rather than through generate), Recommender.interaction(shap_df, X_df, ...) returns (tests, diag) for that subcohort alone; pass its output (optionally concatenated with a second subcohort's) to merge_interaction_candidates to get the final {stratifying_variable: [partner, ...]} spec.

Validating the interaction rule

Recommender.null_sim() runs two small simulations under an additive null (no true interaction) and reports type-I error:

• a binary partner at 5% prevalence in one stratum and 25% in the other (the paper's own stratum sizes), comparing the original stratum-reference
  • Wilcoxon test against the within-stratum contrast with the Wald p-value (sandwich or bootstrap, per se_method) used by stratified_screen;
• a continuous partner correlated with a non-linear stratifying feature, with binary adjustment covariates present and no interaction of any kind, exercising the product-term regression directly.

from shap_recommender import Recommender
Recommender(res_dir=".").null_sim()

License

MIT