{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predict age from white matter features\n",
    "\n",
    "This example uses data from the [HBN POD2 dataset](https://www.nature.com/articles/s41597-022-01695-7), \n",
    "which includes 1867 subjects ages 5-21. We will use the sparse group lasso implemented in AFQ-Insight to fit a predictive model that uses tractometry features to predict each subject's age. Because white matter develops dramatically during childhood and adolescence, this model can be fit to account for a substantial proportion of variance in a held-out dataset."
   ]
  },
  {
   "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 sklearn.model_selection import train_test_split\n",
    "from sklearn.impute import SimpleImputer\n",
    "from groupyr.decomposition import GroupPCA\n",
    "\n",
    "from afqinsight.neurocombat_sklearn import CombatModel\n",
    "from afqinsight import make_afq_regressor_pipeline\n",
    "from afqinsight import AFQDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read the data\n",
    "\n",
    "The `nodes.csv` file, which is the input here is the output of pyAFQ processing. The `subjects.tsv` file is a BIDS-compliant participants file, which includes subject identifiers that match those that are \n",
    "stored in the pyAFQ output. This allows AFQ-Insight to merge the data between the two files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "afqdata = AFQDataset.from_files(\n",
    "    fn_nodes=op.join(afq_home, \"afq-insight/hbn/nodes.csv\"),\n",
    "    fn_subjects=op.join(afq_home, \"afq-insight/hbn/subjects.tsv\"),\n",
    "    dwi_metrics=[\"dki_md\", \"dki_fa\"],\n",
    "    target_cols=[\"age\", \"sex\", \"scan_site_id\"],\n",
    "    label_encode_cols = [\"sex\", \"scan_site_id\"],\n",
    "    index_col=\"subject_id\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "afqdata.drop_target_na()\n",
    "print(afqdata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train / test split\n",
    "\n",
    "We can pass the `AFQDataset` class instance to scikit-learn's\n",
    "`train_test_split` function, just as we would with an array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train, dataset_test = train_test_split(afqdata, test_size=0.25)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Impute missing values\n",
    "\n",
    "Next we impute missing values using median imputation (some values are missing because of noisy MRI scans). We fit the imputer using the training set and then use it to transform both the training and test\n",
    "sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imputer = dataset_train.model_fit(SimpleImputer(strategy=\"median\"))\n",
    "dataset_train = dataset_train.model_transform(imputer)\n",
    "dataset_test = dataset_test.model_transform(imputer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Harmonize the sites and replot\n",
    "\n",
    "The HBN dataset contains measurements from four different sites. \n",
    "and there are substantial scan site differences in both the\n",
    "FA and MD profiles. We use neuroComBat ([Fortin et al., 2017](https://doi.org/10.1016/j.neuroimage.2017.08.047)) to harmonize \n",
    "the site differences and then replot the mean bundle profiles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fit the ComBat transformer to the training set\n",
    "\n",
    "combat = CombatModel()\n",
    "combat.fit(\n",
    "    dataset_train.X,\n",
    "    dataset_train.y[:, 2][:, np.newaxis],\n",
    "    dataset_train.y[:, 1][:, np.newaxis],\n",
    "    dataset_train.y[:, 0][:, np.newaxis],\n",
    ")\n",
    "\n",
    "\n",
    "# And then transform a copy of the test set and a copy of the train set:\n",
    "harmonized_test = dataset_test.copy()\n",
    "harmonized_test.X = combat.transform(\n",
    "    dataset_test.X,\n",
    "    dataset_test.y[:, 2][:, np.newaxis],\n",
    "    dataset_test.y[:, 1][:, np.newaxis],\n",
    "    dataset_test.y[:, 0][:, np.newaxis],\n",
    ")\n",
    "\n",
    "harmonized_train = dataset_train.copy()\n",
    "harmonized_train.X = combat.transform(\n",
    "    dataset_train.X,\n",
    "    dataset_train.y[:, 2][:, np.newaxis],\n",
    "    dataset_train.y[:, 1][:, np.newaxis],\n",
    "    dataset_train.y[:, 0][:, np.newaxis],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an analysis pipeline\n",
    "Finally, we can use the imputed and harmonized data. AFQ-Insight implements complex pipelines that include multiple analysis steps. Helper functions (such as `make_afq_regressor_pipeline`) create \n",
    "scikit-learn compatible pipelines that can then be used to fit, predict and score the model.\n"
   ]
  },
  {
   "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 GroupPCA transformer\n",
    "    # for anything else\n",
    "    imputer = SimpleImputer(strategy=\"median\")\n",
    "    gpca = GroupPCA(n_components=n_components, groups=afqdata.groups)\n",
    "    groups_pca = gpca.fit(harmonized_test.X).groups_out_\n",
    "\n",
    "    transformer = GroupPCA\n",
    "    transformer_kwargs = {\"groups\": afqdata.groups, \"n_components\": n_components}\n",
    "else:\n",
    "    transformer = False\n",
    "    transformer_kwargs = None\n",
    "\n",
    "pipe = make_afq_regressor_pipeline(\n",
    "    imputer_kwargs={\"strategy\": \"median\"},  # Use median imputation\n",
    "    use_cv_estimator=True,  # Automatically determine the best hyperparameters\n",
    "    scaler=\"standard\",  # Standard scale the features before regression\n",
    "    feature_transformer=transformer,  # See note above about group PCA\n",
    "    feature_transformer_kwargs=transformer_kwargs,\n",
    "    groups=(\n",
    "        groups_pca if do_group_pca else afqdata.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 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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe.fit(harmonized_train.X, harmonized_train.y[:, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_age = pipe.predict(harmonized_test.X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.scatter(harmonized_test.y[:, 0], pred_age)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "pipe.score(harmonized_test.X, harmonized_test.y[:, 0])"
   ]
  }
 ],
 "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
}