{ "cells": [ { "cell_type": "markdown", "id": "intro", "metadata": {}, "source": [ "# Prediction with CSlearn: Mushroom Example\n", "\n", "This notebook shows how to apply CSlearn to a real dataset and use the fitted model for prediction. We use the [UCI Mushroom dataset](https://archive.ics.uci.edu/dataset/73/mushroom) (8124 samples, 22 categorical features) and predict whether a mushroom is poisonous or not.\n", "\n", "The workflow has three steps:\n", "1. **Fit** — learn a CStree from training data via `CStree.fit()`.\n", "2. **Save/load** — pickle the fitted model so the slow fit step can be skipped on subsequent runs.\n", "3. **Predict** — call `tree.predict(partial_observations)` on held-out test data." ] }, { "cell_type": "code", "execution_count": 1, "id": "imports", "metadata": {}, "outputs": [], "source": [ "import pickle\n", "import warnings\n", "\n", "warnings.simplefilter(action=\"ignore\", category=FutureWarning)\n", "\n", "import numpy as np\n", "import pandas as pd\n", "from sklearn.model_selection import train_test_split\n", "from sklearn.preprocessing import OrdinalEncoder\n", "from sklearn.metrics import confusion_matrix, accuracy_score\n", "from ucimlrepo import fetch_ucirepo\n", "\n", "from cslearn.cstree import CStree" ] }, { "cell_type": "markdown", "id": "data-section", "metadata": {}, "source": [ "## Load and preprocess the data\n", "\n", "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." ] }, { "cell_type": "code", "execution_count": 2, "id": "load-data", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Train: 3950 rows | Test: 1694 rows | Variables: 22\n", "Cardinalities: [6, 4, 8, 2, 7, 2, 2, 2, 9, 2, 4, 4, 4, 7, 7, 2, 3, 4, 6, 6, 6, 2]\n" ] } ], "source": [ "mushroom = fetch_ucirepo(id=73)\n", "data = pd.concat([mushroom.data.features, mushroom.data.targets], axis=1)\n", "data.dropna(inplace=True)\n", "data.drop(\"veil-type\", axis=1, inplace=True) # constant column\n", "\n", "ordinalized = OrdinalEncoder().fit_transform(data.values).astype(int)\n", "data = pd.DataFrame(ordinalized, columns=data.columns)\n", "\n", "cards = data.nunique().values\n", "train, test = train_test_split(data, test_size=0.3, random_state=0)\n", "\n", "print(f\"Train: {len(train)} rows | Test: {len(test)} rows | Variables: {len(cards)}\")\n", "print(f\"Cardinalities: {cards.tolist()}\")" ] }, { "cell_type": "markdown", "id": "fit-section", "metadata": {}, "source": [ "## Fit the model\n", "\n", "`CStree.fit()` runs the three-phase CSlearn pipeline:\n", "1. GRaSP to identify candidate parent sets.\n", "2. Order MCMC (Gibbs sampler) to find the MAP variable ordering.\n", "3. Exact staging search under the β ≤ 2 constraint.\n", "\n", "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." ] }, { "cell_type": "code", "execution_count": 3, "id": "fit-model", "metadata": {}, "outputs": [], "source": [ "# The cardinalities row must be the first row of the DataFrame passed to fit().\n", "train_with_cards = train.copy()\n", "train_with_cards.loc[-1] = cards\n", "train_with_cards = train_with_cards.reindex(np.roll(train_with_cards.index, 1)).reset_index(drop=True)\n", "\n", "# Fit and save (takes several minutes on first run).\n", "# opt_tree = CStree().fit(train_with_cards, save_as=\"mushroom\")\n", "\n", "# Load from a previously saved pickle instead:\n", "with open(\"mushroom-opt_tree.pkl\", \"rb\") as f:\n", " opt_tree = pickle.load(f)" ] }, { "cell_type": "markdown", "id": "predict-section", "metadata": {}, "source": [ "## Predict on test data\n", "\n", "`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." ] }, { "cell_type": "code", "execution_count": 4, "id": "predict", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
poisonous
01
10
20
30
40
\n", "
" ], "text/plain": [ " poisonous\n", "0 1\n", "1 0\n", "2 0\n", "3 0\n", "4 0" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "predict_from = test.drop(\"poisonous\", axis=1)\n", "true_labels = test[\"poisonous\"].values\n", "\n", "predictions = opt_tree.predict(predict_from)\n", "predictions.head()" ] }, { "cell_type": "markdown", "id": "eval-section", "metadata": {}, "source": [ "## Evaluate\n", "\n", "We report overall accuracy and the confusion matrix." ] }, { "cell_type": "code", "execution_count": 5, "id": "evaluate", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Accuracy: 0.9581\n", "Confusion matrix (rows=true, cols=predicted):\n", " pred edible pred poisonous\n", "edible 1023 30\n", "poisonous 41 600\n" ] } ], "source": [ "pred_labels = predictions[\"poisonous\"].values\n", "acc = accuracy_score(true_labels, pred_labels)\n", "print(f\"Accuracy: {acc:.4f}\")\n", "\n", "cm = confusion_matrix(true_labels, pred_labels)\n", "print(\"Confusion matrix (rows=true, cols=predicted):\")\n", "print(pd.DataFrame(cm, index=[\"edible\", \"poisonous\"], columns=[\"pred edible\", \"pred poisonous\"]))" ] }, { "cell_type": "markdown", "id": "probs-section", "metadata": {}, "source": [ "## Prediction with confidence scores\n", "\n", "Pass `return_probs=True` to get the conditional probability of the MAP prediction given the observed features." ] }, { "cell_type": "code", "execution_count": 6, "id": "probs", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
poisonousPROB
010.999255
100.999967
200.999939
300.998355
400.999957
\n", "
" ], "text/plain": [ " poisonous PROB\n", "0 1 0.999255\n", "1 0 0.999967\n", "2 0 0.999939\n", "3 0 0.998355\n", "4 0 0.999957" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "predictions_with_probs = opt_tree.predict(predict_from.head(5), return_probs=True)\n", "predictions_with_probs" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.11.0" }, "nbsphinx": { "execute": "never" } }, "nbformat": 4, "nbformat_minor": 5 }