{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Classify ALS diagnosis from white matter features\n",
    "\n",
    "This example demonstrates classification, a machine learning task in which a model is constructed and fit to  accurately discriminate between different classes of subjects. In this case, participants with amyotrophic lateral sclerosis (ALS) and healthy control. We use the dataset from ([Sarica et al, 2017](https://pmc.ncbi.nlm.nih.gov/articles/PMC6867092/)), which contains tractometry\n",
    "features from 24 patients with ALS and 24 demographically matched control subjects. We will use the sparse group lasso (SGL) algorithm implemented in AFQ-Insight, as also described in [Richie-Halford et al. 2021](https://journals.plos.org/ploscompbiol/article?id=10.1371/journal.pcbi.1009136)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from paths import afq_home\n",
    "\n",
    "import os.path as op\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from groupyr.decomposition import GroupPCA\n",
    "from sklearn.impute import SimpleImputer\n",
    "from sklearn.model_selection import cross_validate\n",
    "\n",
    "from afqinsight import load_afq_data, make_afq_classifier_pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the data from Sarica et al.\n",
    "\n",
    "The data is read into an AFQData object, which manages the division of subject information and outputs of pyAFQ (in the `nodes.csv` file) and merges between them.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "afqdata = load_afq_data(op.join(afq_home, \"afq-insight/sarica/nodes.csv\"), \n",
    "                        op.join(afq_home, \"afq-insight/sarica/subjects.csv\"),\n",
    "                        dwi_metrics=[\"fa\", \"md\"], \n",
    "                        target_cols=[\"class\"], \n",
    "                        label_encode_cols=[\"class\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Examine the data\n",
    "\n",
    "``afqdata`` is an ``AFQDataset`` object, with properties corresponding to\n",
    "the tractometry features and phenotypic targets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "X = afqdata.X\n",
    "y = afqdata.y.astype(float)  # SGL expects float targets\n",
    "groups = afqdata.groups\n",
    "feature_names = afqdata.feature_names\n",
    "group_names = afqdata.group_names\n",
    "subjects = afqdata.subjects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Reduce data dimensionality\n",
    "\n",
    "To save computational time, we take the first 10 principal components from each\n",
    "feature group (i.e. from each metric-bundle combination).\n",
    "For more details on this approach in a research setting, please see [2]_.\n",
    "\n",
    "Here we reduce computation time by taking the first 10 principal components of\n",
    "each feature group and performing SGL logistic regression on those components.\n",
    "If you want to train an SGL model without group PCA, set ``do_group_pca =\n",
    "False``. This will increase the number of features by an order of magnitude\n",
    "and slow down execution time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "do_group_pca = True\n",
    "\n",
    "if do_group_pca:\n",
    "    n_components = 10\n",
    "\n",
    "    # The next three lines retrieve the group structure of the group-wise PCA\n",
    "    # and store it in ``groups_pca``. We do not use the imputer or GroupPCA transformer\n",
    "    # for anything else\n",
    "    imputer = SimpleImputer(strategy=\"median\")\n",
    "    gpca = GroupPCA(n_components=n_components, groups=groups)\n",
    "    groups_pca = gpca.fit(imputer.fit_transform(X)).groups_out_\n",
    "\n",
    "    transformer = GroupPCA\n",
    "    transformer_kwargs = {\"groups\": groups, \"n_components\": n_components}\n",
    "else:\n",
    "    transformer = False\n",
    "    transformer_kwargs = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Create the classification pipeline\n",
    "\n",
    "The core computational machinery is a pipeline. These operate as scikit-learn\n",
    "compatible pipelines, so we can pass them to scikit-learn functions.\n",
    "There are many options that need to be set to configure the pipeline object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "pipe = make_afq_classifier_pipeline(\n",
    "    imputer_kwargs={\"strategy\": \"median\"},  # Use median imputation\n",
    "    use_cv_estimator=True,  # Automatically determine the best hyperparameters\n",
    "    feature_transformer=transformer,  # See note above about group PCA\n",
    "    feature_transformer_kwargs=transformer_kwargs,\n",
    "    scaler=\"standard\",  # Standard scale the features before regression\n",
    "    groups=(\n",
    "        groups_pca if do_group_pca else groups\n",
    "    ),  # SGL will use the original feature groups or the PCA feature groups depending on the choice above # noqa E501\n",
    "    verbose=0,  # Be quiet!\n",
    "    pipeline_verbosity=False,  # No really, be quiet!\n",
    "    tuning_strategy=\"bayes\",  # Use BayesSearchCV to determine optimal hyperparameters\n",
    "    n_bayes_iter=20,  # Consider only this many points in hyperparameter space\n",
    "    cv=3,  # Use three CV splits to evaluate each hyperparameter combination\n",
    "    l1_ratio=[0.0, 1.0],  # Explore the entire range of ``l1_ratio``\n",
    "    eps=5e-2,  # This is the ratio of the smallest to largest ``alpha`` value\n",
    "    tol=1e-2,  # Set a lenient convergence tolerance just for this example\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Fit and cross-validate\n",
    "The ``pipe`` object is a scikit-learn pipeline and can be used in other\n",
    "scikit-learn functions. Here, the generic `cross_validate` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "scores = cross_validate(\n",
    "    pipe, X, y, cv=5, return_train_score=True, return_estimator=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Display results \n",
    "\n",
    "Finally, we display the results, including both the training and test scores, as well as a visualization of the model coefficients: which tracts contributed the most to this classification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "print(f\"Mean train score: {np.mean(scores['train_score']):5.3f}\")\n",
    "print(f\"Mean test score:  {np.mean(scores['test_score']):5.3f}\")\n",
    "print(f\"Mean fit time:    {np.mean(scores['fit_time']):5.2f}s\")\n",
    "print(f\"Mean score time:  {np.mean(scores['score_time']):5.2f}s\")\n",
    "\n",
    "mean_coefs = np.mean(\n",
    "    np.abs([est.named_steps[\"estimate\"].coef_ for est in scores[\"estimator\"]]), axis=0\n",
    ")\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(8, 5))\n",
    "_ = ax.plot(mean_coefs[:180], color=\"black\", lw=2)\n",
    "_ = ax.set_xlim(0, 180)\n",
    "\n",
    "colors = plt.get_cmap(\"tab20\").colors\n",
    "for grp, grp_name, color in zip(groups_pca[:18], group_names, colors):\n",
    "    _ = ax.axvspan(grp.min(), grp.max() + 1, color=color, alpha=0.8, label=grp_name[1])\n",
    "\n",
    "box = ax.get_position()\n",
    "_ = ax.set_position(\n",
    "    [box.x0, box.y0 + box.height * 0.375, box.width, box.height * 0.625]\n",
    ")\n",
    "\n",
    "_ = ax.legend(loc=\"upper center\", bbox_to_anchor=(0.5, -0.2), ncol=3)\n",
    "_ = ax.set_ylabel(r\"$\\hat{\\beta}$\", fontsize=16)\n",
    "_ = ax.set_xlabel(\"Group principal component\", fontsize=16)\n",
    "_ = ax.set_title(\"Group Principal Regression Coefficients (FA only)\", fontsize=18)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}