{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "np.random.seed(13)\n",
    "from sklearn.datasets import fetch_openml\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.tree import plot_tree, DecisionTreeClassifier\n",
    "from sklearn import metrics\n",
    "from scipy.io.arff import loadarff\n",
    "\n",
    "# installable with: `pip install imodels`\n",
    "from imodels import SLIMRegressor, BayesianRuleListClassifier, RuleFitRegressor, GreedyRuleListClassifier\n",
    "from imodels import SLIMClassifier, OneRClassifier, BoostedRulesClassifier\n",
    "\n",
    "# change working directory to project root\n",
    "if os.getcwd().split('/')[-1] != 'imodels':\n",
    "    os.chdir('..')\n",
    "\n",
    "\n",
    "def get_ames_data():\n",
    "    housing = fetch_openml(name=\"house_prices\", as_frame=True)\n",
    "    housing_target = housing['target'].values\n",
    "    housing_data_numeric = housing['data'].select_dtypes('number').drop(columns=['Id']).dropna(axis=1)\n",
    "    feature_names = housing_data_numeric.columns.values\n",
    "    X_train_reg, X_test_reg, y_train_reg, y_test_reg = train_test_split(\n",
    "        housing_data_numeric.values, housing_target, test_size=0.75)\n",
    "    return X_train_reg, X_test_reg, y_train_reg, y_test_reg, feature_names\n",
    "    \n",
    "\n",
    "def get_diabetes_data():\n",
    "    '''load (classification) data on diabetes\n",
    "    '''\n",
    "    data = loadarff(\"imodels/tests/test_data/diabetes.arff\")\n",
    "    data_np = np.array(list(map(lambda x: np.array(list(x)), data[0])))\n",
    "    X = data_np[:, :-1].astype('float32')\n",
    "    y_text = data_np[:, -1].astype('str')\n",
    "    y = (y_text == 'tested_positive').astype(int)  # labels 0-1\n",
    "    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.75)  # split\n",
    "    feature_names = [\"#Pregnant\", \"Glucose concentration test\", \"Blood pressure(mmHg)\",\n",
    "                     \"Triceps skin fold thickness(mm)\",\n",
    "                     \"2-Hour serum insulin (mu U/ml)\", \"Body mass index\", \"Diabetes pedigree function\", \"Age (years)\"]\n",
    "    return X_train, X_test, y_train, y_test, feature_names\n",
    "\n",
    "\n",
    "X_train_reg, X_test_reg, y_train_reg, y_test_reg, feat_names_reg = get_ames_data()\n",
    "X_train, X_test, y_train, y_test, feat_names = get_diabetes_data()\n",
    "\n",
    "\n",
    "def viz_classification_preds(probs, y_test):\n",
    "    '''look at prediction breakdown\n",
    "    '''\n",
    "    plt.subplot(121)\n",
    "    plt.hist(probs[:, 1][y_test == 0], label='Class 0')\n",
    "    plt.hist(probs[:, 1][y_test == 1], label='Class 1', alpha=0.8)\n",
    "    plt.ylabel('Count')\n",
    "    plt.xlabel('Predicted probability of class 1')\n",
    "    plt.legend()\n",
    "\n",
    "    plt.subplot(122)\n",
    "    preds = np.argmax(probs, axis=1)\n",
    "    plt.title('ROC curve')\n",
    "    fpr, tpr, thresholds = metrics.roc_curve(y_test, preds)\n",
    "    plt.xlabel('False positive rate')\n",
    "    plt.ylabel('True positive rate')\n",
    "    plt.plot(fpr, tpr)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "# load some data\n",
    "print('regression data', X_train_reg.shape, 'classification data', X_train.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# rule sets\n",
    "Rule sets are models that create a set of (potentially overlapping) rules."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### rulefit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# fit a rulefit model\n",
    "rulefit = RuleFitRegressor(max_rules=10)\n",
    "rulefit.fit(X_train_reg, y_train_reg, feature_names=feat_names_reg)\n",
    "\n",
    "# get test performance\n",
    "preds = rulefit.predict(X_test_reg)\n",
    "print(f'test r2: {metrics.r2_score(y_test_reg, preds):0.2f}')\n",
    "\n",
    "# inspect and print the rules\n",
    "rules = rulefit.get_rules()\n",
    "rules = rules[rules.coef != 0].sort_values(\"support\", ascending=False)\n",
    "\n",
    "# 'rule' is how the feature is constructed\n",
    "# 'coef' is its weight in the final linear model\n",
    "# 'support' is the fraction of points it applies to\n",
    "rules[['rule', 'coef', 'support']].style.background_gradient(cmap='viridis')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## boosted stumps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fit boosted stumps\n",
    "brc = BoostedRulesClassifier(n_estimators=10)\n",
    "brc.fit(X_train, y_train, feature_names=feat_names)\n",
    "\n",
    "print(brc)\n",
    "\n",
    "# look at performance\n",
    "probs = brc.predict_proba(X_test)\n",
    "viz_classification_preds(probs, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# rule lists"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### greedy rule lists\n",
    "**like a decision tree that only ever splits going left**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# fit a greedy rule list\n",
    "m = GreedyRuleListClassifier()\n",
    "m.fit(X_train, y=y_train, feature_names=feat_names)  # stores into m.rules_\n",
    "probs = m.predict_proba(X_test)\n",
    "\n",
    "# print the list\n",
    "print(m)\n",
    "\n",
    "# look at prediction breakdown\n",
    "viz_classification_preds(probs, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### oneR\n",
    "**fits a rule list restricted to use only one feature**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fit a oneR model\n",
    "m = OneRClassifier()\n",
    "m.fit(X_train, y=y_train, feature_names=feat_names)  # stores into m.rules_\n",
    "probs = m.predict_proba(X_test)\n",
    "\n",
    "# print the rule list\n",
    "print(m)\n",
    "\n",
    "# look at prediction breakdown\n",
    "viz_classification_preds(probs, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### scalable bayesian rule lists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train classifier (allow more iterations for better accuracy; use BigDataRuleListClassifier for large datasets)\n",
    "print('training...')\n",
    "m = BayesianRuleListClassifier(max_iter=3000, class1label=\"diabetes\", verbose=False)\n",
    "m.fit(X_train, y_train)\n",
    "probs = m.predict_proba(X_test)\n",
    "print(\"learned model:\\n\", m)\n",
    "viz_classification_preds(probs, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# rule trees"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### short decision tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# specify a decision tree with a maximum depth\n",
    "dt = DecisionTreeClassifier(max_depth=3)\n",
    "dt.fit(X_train, y_train)\n",
    "\n",
    "# calculate mse on the training data\n",
    "probs = dt.predict_proba(X_test)\n",
    "# print(f'test mse: {np.mean(np.square(preds-y)):0.2f}')\n",
    "\n",
    "plot_tree(dt)\n",
    "# plt.savefig('tree.pdf')\n",
    "plt.show()\n",
    "\n",
    "viz_classification_preds(probs, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### optimal classification tree\n",
    "- docs [here](https://github.com/csinva/interpretability-workshop/tree/master/imodels/optimal_classification_tree)\n",
    "- note: this implementation is still somewhat unstable, and can be made faster by installing either `cplex` or `gurobi`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sys.path.append('../imodels/optimal_classification_tree/pyoptree')\n",
    "# sys.path.append('../imodels/optimal_classification_tree/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from optree import OptimalTreeModel\n",
    "# feature_names = np.array([\"x1\", \"x2\"])\n",
    "\n",
    "# X = np.array([[1, 2, 2, 2, 3], [1, 2, 1, 0, 1]]).T\n",
    "# y = np.array([1, 1, 0, 0, 0]).reshape(-1, 1)\n",
    "# X_test = np.array([[1, 1, 2, 2, 2, 3, 3], [1, 2, 2, 1, 0, 1, 0]]).T\n",
    "# y_test = np.array([1, 1, 1, 0, 0, 0, 0])\n",
    "\n",
    "# np.random.seed(13)\n",
    "# model = OptimalTreeModel(tree_depth=3, N_min=1, alpha=0.1) #, solver_name='baron'\n",
    "# model.fit(X_test, y_test) # this method is currently using the fast, but not optimal solver\n",
    "# preds = model.predict(X_test)\n",
    "\n",
    "# # fit on the bigger diabetes dset from above\n",
    "# # model.fit(Xtrain, ytrain) # this method is currently using the fast, but not optimal solver\n",
    "# # preds = model.predict(Xtest)\n",
    "\n",
    "# print('acc', np.mean(preds == y_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.print_tree(feature_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# algebraic models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### integer linear models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "\n",
    "# generate X and y\n",
    "n, p = 500, 10\n",
    "X_sim = np.random.randn(n, p)\n",
    "y_sim = 1 * X_sim[:, 0] + 2 * X_sim[:, 1] - 1 * X_sim[:, 2] + np.random.randn(n)\n",
    "\n",
    "# fit linear models with different regularization parameters\n",
    "print('groundtruth weights should be 1, 2, -1...')\n",
    "model = SLIMRegressor()\n",
    "for lambda_reg in [1e-3, 1e-2, 5e-2, 1e-1, 1, 2, 5, 10]:\n",
    "    model.fit(X_sim, y_sim, lambda_reg)\n",
    "    mse = np.mean(np.square(y_sim - model.predict(X_sim)))\n",
    "    print(f'lambda: {lambda_reg}\\tmse: {mse: 0.2f}\\tweights: {model.model_.coef_}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_sim = 1 / (1 + np.exp(-y_sim))\n",
    "y_sim = np.round(y_sim)\n",
    "\n",
    "# fit linear models with different regularization parameters\n",
    "print('groundtruth weights should be 1, 2, -1...')\n",
    "model = SLIMClassifier()\n",
    "for lambda_reg in [1e-3, 1e-2, 5e-2, 1e-1, 1, 2, 5, 10]:\n",
    "    model.fit(X_sim, y_sim, lambda_reg)\n",
    "    mll = np.mean(metrics.log_loss(y_sim, model.predict(X_sim)))\n",
    "    print(f'lambda: {lambda_reg}\\tmlogloss: {mll: 0.2f}\\tweights: {model.model_.coef_}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,../imodels/tests/notebooks//py:percent"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
