BART Binary Regression Issues

Hi,

I’ve been working on a classification problem recently. Its highly skew data (~8% of observations in the positive class) so used a cloglog link to fit a bambi model. Results were reasonable but not amazing so I thought I’d give a BART model a go. Since most features in the model are binary I thought it might be able to find some complex interactions that I’d missed.

When trying to predict out of samples from the BART model I get 0 in the positive class. To isolate data issues I went back to the classic titanic dataset to try benchmark BART vs a Bambi model;

import pymc as pm
import pymc_bart as pmb
import bambi as bmb
import arviz as az
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix

df=pd.read_csv(...'Titanic-Dataset.csv')
##lazy feature selection
df['Male']=np.where(df['Sex']=='male', 1,0)
train_df,test_df=train_test_split(df,test_size=0.2)
SS=StandardScaler()
SS.fit(train_df[['Fare']])
train_df['FareScaled']=SS.transform(train_df[['Fare']])
test_df['FareScaled']=SS.transform(test_df[['Fare']])
AgeMean=train_df['Age'].mean()
train_df['AgeFilled']=train_df['Age'].fillna(AgeMean)
test_df['AgeFilled']=test_df['Age'].fillna(AgeMean)
AS=StandardScaler()
AS.fit(train_df[['AgeFilled']])
train_df['AgeScaled']=AS.transform(train_df[['AgeFilled']])
test_df['AgeScaled']=AS.transform(test_df[['AgeFilled']])
##Bambi
train_df.head()
bmbMdl=bmb.Model('Survived~Pclass+Male+FareScaled+AgeFilled+SibSp+Parch', data=train_df, link='cloglog', family='bernoulli')
bmbMdl_fitted=bmbMdl.fit()
bmbMdl.predict(bmbMdl_fitted,data=test_df, kind='response')
test_df['Pred Surv']=bmbMdl_fitted.posterior_predictive.Survived.mean(dim=['chain','draw'])
test_df['Surv']=np.where(test_df['Pred Surv']>=0.5,1,0)
g=ConfusionMatrixDisplay(confusion_matrix(test_df['Survived'], test_df['Surv']))
g.plot()

##BART
features=['Pclass','Male','FareScaled','AgeScaled','SibSp','Parch']
with pm.Model() as BartMdl:
    x_data=pm.Data("x_data", test_df[features])
    y_data=pm.Data("y_data", test_df['Survived'])
    mu=pmb.BART("mu", x_data, y_data,m=50)
    y_pred=pm.Bernoulli("y_pred", p=1.0-pm.math.exp(-pm.math.exp(mu)), observed=y_data)
    idata=pm.sample(draws=1000, tune=1000, chains=4)
az.plot_convergence_dist(idata)
with BartMdl:
    pm.set_data({"x_data":test_df[features], "y_data":test_df['Survived']})
    pm.sample_posterior_predictive(idata, extend_inferencedata=True, predictions=True)
test_df['Bart Pred']=idata.predictions.y_pred.mean(dim=['chain','draw'])
test_df['Bart Surv']=np.where(test_df['Bart Pred']>=0.5,1,0)
g=ConfusionMatrixDisplay(confusion_matrix(test_df['Survived'], test_df['Bart Surv']))
g.plot()

We get fairly poor BART convergence, but good predictive power and I understand that convergence is slightly less of an issue for BART models than standard Bayesian GLMs.

BART Confusion Matrix:

BAMBI Confusion Matrix:

So far so good. Running my BART model on real data gives:

import pandas as pd
import arviz as az
import numpy as np
import pymc as pm
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix
import pymc_bart as pmb

df=pd.read_csv(...)
df_test=pd.read_csv(...)
features=['f1', 'f2','f3',...]
with pm.Model() as logitbart:
    x_data=pm.Data("x_data", df[features])
    y_data=pm.Data("y_data", df['Target'])
    mu_bart=pmb.BART("mu_bart", X=x_data,Y=y_data, m=50)
    y_pred=pm.Bernoulli("y_pred",p=1.0-pm.math.exp(-pm.math.exp(mu_bart)), observed=y_data)
    idata=pm.sample(draws=1000, tune=1000, chains=4,random_seed=42)
az.plot_convergence_dist(idata)
with logitbart:
    pm.set_data({"x_data":df_test[features], "y_data":df_test['Target']})
    pm.sample_posterior_predictive(idata, extend_inferencedata=True,predictions=True, sample_vars=['mu_bart', 'y_pred'])

test=idata.predictions.y_pred.mean(dim=['chain','draw'])
df_test['Pred P']=test
df_test['Pred Val']=np.where(df_test['Pred P']>=0.5,1,0)
g=ConfusionMatrixDisplay(confusion_matrix(df_test['Target'], df_test['Pred Val']))
g.plot()
    

Again, relatively iffy convergence but its really the predictions which are the issue. With a Bambi cloglog model I get poor, but some results;

So I was expecting similar at least.

In fitting the BART model this line is interesting to me;

pm.sample_posterior_predictive(idata, extend_inferencedata=True,predictions=True, sample_vars=['mu_bart', 'y_pred'])

If I leave out the sample_vars argument, like I did in the titanic example, I get an error message here:

ImplicitFreezeWarning: The following trace variables were implicitly frozen, but something in their upstream suggests you may have wanted them resampled: ‘mu_bart’ (upstream Data/coords changed (x_data, y_data)). To resample them, pass them in `sample_vars`; to silence this warning while keeping the trace values, pass them in `freeze_vars`.

I didnt get this error for the titanic data, which makes me think I’m doing something wrong with this model, but I cant spot any differences in its formatting.

So have I done something wrong in this second BART model, or is it an issue with my data?

cc @aloctavodia

I’d be interested in seeing what the probabilities for the positive cases are. Can you show the probability distribution, PR-AUC and calibration?

Also make sure you are not using pandas or polars as inputs. Sometimes

import plotly.express as px
from sklearn.metrics import PrecisionRecallDisplay
from sklearn.calibration import CalibrationDisplay

fig1=px.histogram(data_frame=df_test[df_test['Target']==1], x='Pred P', nbins=40)
fig2=PrecisionRecallDisplay.from_predictions(df_test['Target'], df_test['Pred P'])
fig3=CalibrationDisplay.from_predictions(df_test['Target'],df_test['Pred P'])
fig1.show()
fig2.show()
fig3.show()

Distribution

PR-AUC

Calibration

Also changed the data in my model (both in the initial model creation and set_data) to be passed via to_numpy() but no changes.

So yeah, the model is doing something, looking at the PR curve. I think the issue is that with prevalence near 6% in your data, a calibrated model with weak/modest discrimination is going to have a tough time getting to 0.5. What do the probabilities from Bambi look like relative to this?

BTW, averaging draws of the simulated y_pred adds Monte Carlo noise to the probability estimate. You should transform mu_bart and averaging that instead – should give a smoother estimate.

Thanks for the tip. Adding the probability in as a deterministic does produce far smoother outputs. It makes a lot of sense, so I’ll keep it in mind for the future with binary regression problems.

with pm.Model() as logitbart:
    x_data=pm.Data("x_data", df[features].to_numpy())
    y_data=pm.Data("y_data", df['Target'].to_numpy())
    mu_bart=pmb.BART("mu_bart", X=x_data,Y=y_data, m=50)
    mu_p=pm.Deterministic("mu_p",1.0-pm.math.exp(-pm.math.exp(mu_bart)))
    y_pred=pm.Bernoulli("y_pred",p=mu_p, observed=y_data)
    idata=pm.sample(draws=1000, tune=1000, chains=4,random_seed=42)
with logitbart:
    pm.set_data({"x_data":df_test[features].to_numpy(), "y_data":df_test['Target'].to_numpy()})
    pm.sample_posterior_predictive(idata, extend_inferencedata=True,predictions=True, sample_vars=['y_pred','mu_bart', 'mu_p'])
test=idata.predictions.mu_p.mean(dim=['chain','draw'])
df_test['Pred P']=test
df_test['Pred Val']=np.where(df_test['Pred P']>=0.5,1,0)
PrecisionRecallDisplay.from_predictions(df_test['Target'], df_test['Pred P'])

The Bambi model provides predictions over a much wider range:

Formulae='Targer ~ 1 + Stuff'
BMBModel=bmb.Model(Formulae,data=df,family='bernoulli', link='cloglog')
BMBModel_fit=BMBModel.fit(idata_kwargs=dict(log_likelihood = True))
BMBModel.predict(FinalIMDModel_fit,data=df_test, inplace=True, kind="response_params")
y_pred=FinalIMDModel_fit.posterior['p'].mean(dim=['chain','draw']).values
df_test['Pred P']=y_pred
px.histogram(data_frame=df_test[df_test['Target']==1], x='Pred P', nbins=50)
PrecisionRecallDisplay.from_predictions(df_test['Target'], df_test['Pred P'])
CalibrationDisplay.from_predictions(df_test['Target'],df_test['Pred P'])

This is mostly for curiosity and trying to understand where BART models fit into the types of problem I encounter. Most of the factors (11 total) are binary so I was hoping that the trees might find some interesting interactions I hadn’t spotted and I was surprised when the performance on that positive class was far worse. Realistically this problem is missing truly descriptive factors and is going to be one that goes back into a development stack, I just wanted to tinker a bit before that happened.