{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Kernel Ridge Regression\n\nIn this tutorial, we will build our first kernel based machine learning model with\nthe in-built Kernel Ridge Regressor module. We will focus on a single fidelity KRR\nmodel that predicts the excitation energies of the benzene molecule from the\ndataset provided in this package.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Loading the Built-in Dataset\nLet's load the dataset and isolate only those\ngeometries which have the def2-TZVP fidelity\nwe will build a single fidelity KRR model to\npredict these energies.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport matplotlib.pyplot as plt\nfrom mfml_qc.datasets import load_benzene_data\n\ndataset = load_benzene_data()\n\nX = dataset[\"X_CM\"]\ny = dataset[\"energies\"][\n    :, 7\n]  # Selecting TZVP (index 7; recall that index 0 is the time step)\n\n# Filter out NaNs\nmask = ~np.isnan(y)\nX, y = X[mask], y[mask]\n\n# we will follow the protocol in\n# J. Chem. Theory Comput. 2023, 19, 21, 7658\u20137670 (10.1021/acs.jctc.3c00882)\n# and set the last 2712 geometries/entries as the test set\n\nX_train = X[:-2712]\ny_train = y[:-2712]\nX_test = X[-2712:]\ny_test = y[-2712:]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Learning Curve Generator Function\nWhile training a model, it is worthwhile to also\nconsider how well the prediction error is on an\nunseen test dataset. This assessment is made\nwith the help of a learning curve which plots\nmodel complexity against the model prediction error.\nFor KRR, model complexity is directly related to the number of training samples.\nTherefore we will look at the model prediction error (given as mean absolute error; MAE)\nas a function of the nunmber of training samples used.\nTo get a statistical idea of the error, we will generate learning curves\nfor 10 random sub-sampling of the training dataset.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# import the inbuilt KRR module\nfrom mfml_qc.krr import KRR\nfrom sklearn.utils import shuffle\n\n\n# let's create a learning curve helper function\ndef sf_LC(\n    X_train: np.ndarray,\n    y_train: np.ndarray,\n    X_test: np.ndarray,\n    y_test: np.ndarray,\n    kernel_type: str = \"matern\",\n    sigma: float = 100.0,\n    reg: float = 1e-10,\n    nu=1.5,\n    p=1.5,\n    q=1.0,\n    navg: int = 10,\n    nmax: int = 10,\n    seed: int = 42,\n) -> np.ndarray:\n    \"\"\"\n    A simple learning curve generation function.\n\n    Args:\n        X_train (np.ndarray): training feature data.\n        y_train (np.ndarray): training target data.\n        X_test (np.ndarray): testing feature data.\n        y_test (np.ndarray): testing target data.\n        kernel_type (str, optional): kernel type. Defaults to 'matern'.\n        sigma (float, optional): kernel width/sigma. Defaults to 100.0.\n        reg (float, optional): regularization parameter. Defaults to 1e-10.\n        nu (1.0, 1.5, or 2, optional): nu paraemter for matern kernels. Defaults to 1.5.\n        p (float, optional): Wasserstein kernel parameter. Defaults to 1.5.\n        1 (float, optional): Wasserstein kernel parameter, Defaults to 1.0.\n        navg (int, optional): number of averaging loops/shuffles. Defaults to 10.\n        nmax (int, optional): maximum training size (2^nmax). Defaults to 10.\n        seed (int, optional): random state seed for reproducibility. Defaults to 42.\n\n    Returns:\n        full_maes (np.ndarray): array of shape (nmax, navg) containing Mean Absolute Errors.\n    \"\"\"\n\n    full_maes = np.zeros((nmax, navg), dtype=float)\n    # init the KRR model\n    model = KRR(sigma=sigma, reg=reg, kernel_type=kernel_type)\n    for n in range(navg):\n        maes = []\n        X_train_shuf, y_train_shuf = shuffle(X_train, y_train, random_state=seed + n)\n        # training size loop\n        for i in range(1, 1 + nmax):\n            X_tr_subset = X_train_shuf[: 2**i]\n            y_tr_subset = y_train_shuf[: 2**i]\n            model.train(X_tr_subset, y_tr_subset)\n            preds_test = model.predict(X_test)\n            maes.append(np.mean(np.abs(preds_test - y_test)))\n        full_maes[:, n] = np.asarray(maes)\n    return full_maes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## The learning curve\nget the maes\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "maes = sf_LC(\n    X_train,\n    y_train,\n    X_test,\n    y_test,\n    reg=1e-9,\n    sigma=715.0,\n    kernel_type=\"matern\",\n    nu=1.5,\n    navg=5,\n    nmax=8,\n)\n\n# Plot the MAE vs training set size\nmean_mae = np.mean(maes, axis=1)\nstd_mae = np.std(maes, axis=1)\ntrain_sizes = [2**i for i in range(1, 9)]\n\nplt.figure(figsize=(5, 4))\nplt.errorbar(\n    train_sizes, mean_mae, yerr=std_mae, fmt=\"o-\", capsize=5, label=\"KRR (TZVP)\"\n)\nplt.xscale(\"log\")\nplt.yscale(\"log\")\nplt.xlabel(\"Training Size\")\nplt.ylabel(\"MAE (eV)\")\nplt.title(\"Single Fidelity Learning Curve\")\nplt.legend()\nplt.grid(True, which=\"both\", ls=\"-\", alpha=0.2)\nplt.show()"
      ]
    }
  ],
  "metadata": {
    "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.10.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}