Prediction with CSlearn: Mushroom Example

This notebook shows how to apply CSlearn to a real dataset and use the fitted model for prediction. We use the UCI Mushroom dataset (8124 samples, 22 categorical features) and predict whether a mushroom is poisonous or not.

The workflow has three steps:

1. Fit — learn a CStree from training data via CStree.fit().

2. Save/load — pickle the fitted model so the slow fit step can be skipped on subsequent runs.

3. Predict — call tree.predict(partial_observations) on held-out test data.

import pickle
import warnings

warnings.simplefilter(action="ignore", category=FutureWarning)

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OrdinalEncoder
from sklearn.metrics import confusion_matrix, accuracy_score
from ucimlrepo import fetch_ucirepo

from cslearn.cstree import CStree

Load and preprocess the data

We fetch the dataset from the UCI repository, drop the constant veil-type column, ordinally encode all features, and split into 70% train / 30% test.

mushroom = fetch_ucirepo(id=73)
data = pd.concat([mushroom.data.features, mushroom.data.targets], axis=1)
data.dropna(inplace=True)
data.drop("veil-type", axis=1, inplace=True)  # constant column

ordinalized = OrdinalEncoder().fit_transform(data.values).astype(int)
data = pd.DataFrame(ordinalized, columns=data.columns)

cards = data.nunique().values
train, test = train_test_split(data, test_size=0.3, random_state=0)

print(f"Train: {len(train)} rows | Test: {len(test)} rows | Variables: {len(cards)}")
print(f"Cardinalities: {cards.tolist()}")

Train: 3950 rows | Test: 1694 rows | Variables: 22
Cardinalities: [6, 4, 8, 2, 7, 2, 2, 2, 9, 2, 4, 4, 4, 7, 7, 2, 3, 4, 6, 6, 6, 2]

Fit the model

CStree.fit() runs the three-phase CSlearn pipeline:

1. GRaSP to identify candidate parent sets.

2. Order MCMC (Gibbs sampler) to find the MAP variable ordering.

3. Exact staging search under the β ≤ 2 constraint.

With 22 variables and 3950 training rows this takes several minutes. save_as="mushroom" pickles all intermediate results so the fit can be skipped on re-runs.

# The cardinalities row must be the first row of the DataFrame passed to fit().
train_with_cards = train.copy()
train_with_cards.loc[-1] = cards
train_with_cards = train_with_cards.reindex(np.roll(train_with_cards.index, 1)).reset_index(drop=True)

# Fit and save (takes several minutes on first run).
# opt_tree = CStree().fit(train_with_cards, save_as="mushroom")

# Load from a previously saved pickle instead:
with open("mushroom-opt_tree.pkl", "rb") as f:
    opt_tree = pickle.load(f)

Predict on test data

predict(partial_observations) takes a DataFrame of partial observations and returns a DataFrame of predicted values for the unobserved variables. Here we predict the poisonous label from all other features.

predict_from = test.drop("poisonous", axis=1)
true_labels = test["poisonous"].values

predictions = opt_tree.predict(predict_from)
predictions.head()

	poisonous
0	1
1	0
2	0
3	0
4	0

Evaluate

We report overall accuracy and the confusion matrix.

pred_labels = predictions["poisonous"].values
acc = accuracy_score(true_labels, pred_labels)
print(f"Accuracy: {acc:.4f}")

cm = confusion_matrix(true_labels, pred_labels)
print("Confusion matrix (rows=true, cols=predicted):")
print(pd.DataFrame(cm, index=["edible", "poisonous"], columns=["pred edible", "pred poisonous"]))

Accuracy: 0.9581
Confusion matrix (rows=true, cols=predicted):
           pred edible  pred poisonous
edible            1023              30
poisonous           41             600

Prediction with confidence scores

Pass return_probs=True to get the conditional probability of the MAP prediction given the observed features.

predictions_with_probs = opt_tree.predict(predict_from.head(5), return_probs=True)
predictions_with_probs

	poisonous	PROB
0	1	0.999255
1	0	0.999967
2	0	0.999939
3	0	0.998355
4	0	0.999957