Skip to content
stratx is a library for A Stratification Approach to Partial Dependence for Codependent Variables
Jupyter Notebook Other
Branch: master
Clone or download
Fetching latest commit…
Cannot retrieve the latest commit at this time.
Type Name Latest commit message Commit time
Failed to load latest commit information.
images update readme Jul 15, 2019
LICENSE Initial commit May 6, 2019

A Stratification Approach to Partial Dependence for Codependent Variables


This repo provides pip package stratx for the StratPD and CatStratPD algorithms for computing model-independent partial dependence. Source is Python 3 and MIT license.

Model interpretability is important to machine learning practitioners and a key component of interpretation is the characterization of partial dependence of the response variable on any subset of features used in the model. The two most common strategies suffer from a number of critical weaknesses. In the first strategy, linear regression model coefficients describe how a unit change in an explanatory variable changes the response, while holding other variables constant. But, linear regression is inapplicable for high dimensional (p>n) data sets or when a linear model is insufficient to capture the relationship between explanatory variables and the response. In the second strategy, Partial Dependence (PD) plots and Individual Conditional Expectation (ICE) plots give biased results for the common situation of codependent variables and they rely on fitted models provided by the user. When the supplied model is a poor choice due to systematic bias or overfitting, PD/ICE plots provide little (if any) useful information.

To address these issues, we introduce a new strategy, called StratPD, that does not depend on a user's fitted model, provides accurate results in the presence codependent variables, and is applicable to high dimensional settings. The strategy works by stratifying a data set into groups of observations that are similar, except in the variable of interest, through the use of a decision tree. Any fluctuations of the response variable within a group is likely due to the variable of interest. We apply StratPD to a collection of simulations and case studies to show that StratPD is a fast, reliable, and robust method for assessing partial dependence with clear advantages over state-of-the-art methods.


pip install stratx


Some simple usages for numerical and categorical target variables.


from sklearn.datasets import load_boston, load_diabetes
from stratx.partdep import *

boston = load_boston()
df = pd.DataFrame(, columns=boston.feature_names)
df['MEDV'] =

X = df.drop('MEDV', axis=1)
y = df['MEDV']

plot_stratpd(X, y, 'AGE', 'MEDV', yrange=(-10,10))


diabetes = load_diabetes()
df = pd.DataFrame(, columns=diabetes.feature_names)
df['y'] =
X = df.drop('y', axis=1)
y = df['y']

foo = plot_catstratpd(X, y, 'sex', 'y', catnames=['female','male']) # not sure which is male/female actually!


(See notebooks/examples.ipynb for lots more stuff.)

Kaggle Blue book for bulldozer data set.

Kaggle Two sigma connect: Rental listing inquiries. data set:

Using a synthetic weather data set where temperature varies in sinusoidal fashion over the year and with different baseline temperatures per state:

Here's an example of a categorical variable from the weather data set:

Comparing to PDP / ICE plots

Plots of bedrooms versus rent price using New York City apartment rent data. (a) marginal plot, (b) PD/ICE plot derived from random forest, (c) PD/ICE plot derived from SVM, and (d) StratPD plot; sample size is 10,000 observations of ~400k. The PD/ICE plots are radically different for the same data set, depending on the chosen user model.

Plots of height versus weight using synthetic data body weight data (see academic paper). The PD/ICE on the left is biased by codependent features since pregnant women, who are typically shorter than men, have a jump in weight.

You can’t perform that action at this time.