Scikit-learn Interface Tutorial

This tutorial will walk you through using the sklearn style DAGRegressor and DAGClassifier models on two familiar datasets: Diabetes and Breast Cancer.

How it fits into the bigger causalnex picture

The sklearn wrappers are an alternative to the structure learning step. The fitted model objects include a model.graph_ attribute that can be used for visualisation, post-processing, and most importantly the probability fitting step of the BayesianNetwork, the second stage of the causalnex workflow.


This section demonstrates the performance of the DAGRegressor on a real-world dataset. The main things to note in this section are:

  • The scale sensitivity of the algorithm

  • Interpretability of nonlinear .coef_

The Diabetes dataset

The diabetes dataset presents a standard benchmark regression exercise. The objective is to predict disease progression, given a set of features.

Note: a previous version of this tutorial used the Boston housing data for its demonstration. For more information about the racial discrimination present in the Boston housing data, see the github issue that triggered the removal. To learn more about this dataset, we suggest checking out a sklearn issue that has resulted in its deprecation.

The meaning of available features for the diabetes dataset is shown below.

%load_ext autoreload
%autoreload 2
import os
import sys
module_path = os.path.abspath(os.path.join("../../../"))
if module_path not in sys.path:

import warnings
from sklearn.exceptions import ConvergenceWarning
warnings.simplefilter(action='ignore', category=FutureWarning)
warnings.simplefilter(action='ignore', category=RuntimeWarning)
warnings.simplefilter(action='ignore', category=ConvergenceWarning)
from IPython.display import display, Markdown
import pandas as pd

def variable_score_card(protected_variable: pd.Series, target: pd.Series, n_bins: int = 3):

    def percentile(n):
        def percentile_(x):
            return np.percentile(x, n)
        percentile_.__name__ = 'percentile_%s' % n
        return percentile_

    def pct_size():
        def pct_size_(x):
            return np.round(x.size / len(target) * 100, 1)
        pct_size_.__name__ = 'pct_size'
        return pct_size_

    display(Markdown("**Score card for ``%s``**" %

    agg_funcs = ["size", pct_size(), "mean", "std", "min", percentile(10), "median", percentile(90), "max"]
    if protected_variable.dtype != 'O' and protected_variable.nunique() < 10:
        return y.groupby(protected_variable).agg(agg_funcs).reset_index()
        bin_array = pd.cut(protected_variable, bins=n_bins)
        return y.groupby(bin_array).agg(agg_funcs).reset_index()
import numpy as np
import pandas as pd
from sklearn.datasets import load_diabetes
data = load_diabetes()
.. _diabetes_dataset:

Diabetes dataset

Ten baseline variables, age, sex, body mass index, average blood
pressure, and six blood serum measurements were obtained for each of n =
442 diabetes patients, as well as the response of interest, a
quantitative measure of disease progression one year after baseline.

**Data Set Characteristics:**

  :Number of Instances: 442

  :Number of Attributes: First 10 columns are numeric predictive values

  :Target: Column 11 is a quantitative measure of disease progression one year after baseline

  :Attribute Information:
      - age     age in years
      - sex
      - bmi     body mass index
      - bp      average blood pressure
      - s1      tc, total serum cholesterol
      - s2      ldl, low-density lipoproteins
      - s3      hdl, high-density lipoproteins
      - s4      tch, total cholesterol / HDL
      - s5      ltg, possibly log of serum triglycerides level
      - s6      glu, blood sugar level

Note: Each of these 10 feature variables have been mean centered and scaled by the standard deviation times `n_samples` (i.e. the sum of squares of each column totals 1).

Source URL:

For more information see:
Bradley Efron, Trevor Hastie, Iain Johnstone and Robert Tibshirani (2004) "Least Angle Regression," Annals of Statistics (with discussion), 407-499.

Dataset bias evaluation

As we’re dealing with the data of the individuals and the predictions of the model may have a profound impact on their lives, we should evaluate the dataset on the potential presence of bias.

  1. Sample bias/Data collection:

  • The papers do not explain the protocol of how the samples were generated, therefore unintended biases that could have been introduced can not be detected.

  • Data biases could be a result of inequalities in access to healthcare, e.g. due to insurance coverage, limited access to diabetes screenings in underdeveloped regions or neighbourhoods. Undocumented sensitive variables, e.g. ethnicity, would need to be statistically independent of the choice to be included in the dataset.

  1. Data bias estimation with respect to available sensitive attributes: The dataset includes direct information on age and sex, note that both are standardized. A careful evaluation of the possible bias in the sensitive attributes, includes the comparison of ratios in groups in the data with their population rates, or benchmarks from literature. In our case, without information about geography nor ethnicity and the masking of the actual values the variables take, there is hardly any conclusion to be made from the bias estimation. The breakdown below shows roughly uniform distribution across the two variables. Follow up questions to be assessed include: Which one of the variables is each sex? can we expect the disease to progress similarly in both sexes? and similarly for age distribution.

  2. Risks for model deployment: We cannot determine the distribution across some omitted sensitive variables, nor have the information regarding the data gathering. This poses fairness risks when deploying a model trained on the dataset on populations different to the study cohort.

When deploying the model in the context of healthcare, make sure it is equally performant in the subgroups with respect to sensitive attributes and their intersection.

We recommend always assessing the bias and fairness risks at each step of the process (from problem understanding, data collection, processing, modelling and deployment), when working on models to be deployed, to minimize undesired outcomes

from sklearn.datasets import load_diabetes
data = load_diabetes()
X, y =,
names = data["feature_names"]

X = pd.DataFrame(X, columns=names)
y = pd.Series(y, name="DPROG")

variable_score_card(protected_variable=X["sex"], target=y)

Score card for ``sex``

sex size pct_size mean std min percentile_10 median percentile_90 max
0 -0.044642 235 53.2 149.021277 75.905781 25.0 59.4 140.0 266.8 346.0
1 0.050680 207 46.8 155.666667 78.453313 39.0 62.2 141.0 264.4 341.0
variable_score_card(protected_variable=X["age"], target=y)

Score card for ``age``

age size pct_size mean std min percentile_10 median percentile_90 max
0 (-0.107, -0.0346] 117 26.5 134.273504 69.369777 39.0 57.6 118.0 225.8 346.0
1 (-0.0346, 0.0381] 212 48.0 153.278302 82.337721 25.0 53.0 140.5 272.0 341.0
2 (0.0381, 0.111] 113 25.6 168.477876 70.996617 39.0 78.4 163.0 272.8 332.0

Lets initially benchmark the performance of an ElasticNetCV fitted across the entire dataset.

from sklearn.datasets import load_diabetes
data = load_diabetes()
X, y =,

from sklearn.preprocessing import StandardScaler
ss = StandardScaler()
X = ss.fit_transform(X)
y = (y - y.mean()) / y.std()

from sklearn.linear_model import ElasticNetCV
reg = ElasticNetCV(l1_ratio=[.1, .5, .7, .9, .95, .99, 1], fit_intercept=True)

from sklearn.model_selection import cross_val_score
from sklearn.model_selection import KFold
scores = cross_val_score(reg, X, y, cv=KFold(shuffle=True, random_state=42))
print(f'MEAN R2: {np.mean(scores).mean():.3f}')
MEAN R2: 0.475

Linear DAGRegressor

The DAGRegressor has several parameters which can be used to better fit a more complicated noisy DAG:

  • alpha: The l1 (lasso) regularisation parameter. Increasing this creates a sparser DAG.

  • beta: The l2 (ridge) regularisation parameter.

It was decided to use alpha and beta rather than alpha and l1_ratio like in sklearn elasticnet to uncouple the parameters during optimisation.

There are several parameters which are also of interest which have good defaults, but we highlight here:

  • dependent_target: This forces the target variable y to be only a child node. This is important for performance because in some cases X -> y is indistinguishable from y -> X. Enabling this (default enabled) ensures that the regressor performance at least matches linear regression. The trade-off is that the learned structure might be less accurate if y does cause other features.

  • enforce_dag: This thresholds the learned structure model until the system is a DAG. This is useful for removing the small straggler connections which enables the DAG to be visualised easier. It does not impact performance, because the regressor still uses those connections under the hood.

from sklearn.datasets import load_diabetes
import torch
data = load_diabetes()
X, y =,
names = data["feature_names"]

from sklearn.preprocessing import StandardScaler
ss = StandardScaler()
X = ss.fit_transform(X)
y = (y - y.mean()) / y.std()

from causalnex.structure import DAGRegressor
reg = DAGRegressor(

from sklearn.model_selection import cross_val_score
from sklearn.model_selection import KFold
scores = cross_val_score(reg, X, y, cv=KFold(shuffle=True, random_state=42))
print(f'MEAN R2: {np.mean(scores).mean():.3f}')

X = pd.DataFrame(X, columns=names)
y = pd.Series(y, name="DPROG"), y)
print(pd.Series(reg.coef_, index=names))
reg.plot_dag(output_filename="supporting_files/reg_plot1.html", enforce_dag=True)
/Users/gabriel_azevedo_ferreira/opt/miniconda3/envs/private-causalnex-39/lib/python3.9/site-packages/tqdm/ TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See
  from .autonotebook import tqdm as notebook_tqdm
MEAN R2: 0.479
age    0.000000
sex    0.000000
bmi    0.304309
bp     0.180043
s1     0.000000
s2     0.000000
s3     0.000000
s4     0.000000
s5     0.277159
s6     0.000000
dtype: float64