{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# MFML Cross-Validation and Learning Curves\n\nAs we saw in on of the previous examples, learning curves are a\nvital metric for evaluating kernel-based machine learning\nmethods. They depict the relationship between the model's prediction error and\nthe amount of training data provided.\n\nHowever, validating a Multi-Fidelity Machine Learning (MFML) model requires a\nhighly specialized approach due to the nested structural constraints of the data.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## The Nested Cross-Validation Strategy\nDue to the strictly nested structure of the training data used in MFML,\nconventional cross-validation methods (like standard K-Fold random shuffling)\n**cannot be used**. If we were to randomly shuffle the training sets for each\nfidelity independently, we would break the hierarchy guarantee\u2014meaning a high-fidelity\ngeometry might no longer exist in the baseline training set!\n\nTo ensure our results are robust to the choice of training data while preserving\nthis nested structure, we must use a **Nested Validation** algorithm.\n\nFor each random shuffle iteration, the procedure is as follows:\n\n1. **Highest Fidelity Sampling**: Randomly select a target number of training samples,\n   $N_{\\rm train}^{F}$, from the available highest fidelity dataset. This forms\n   the apex subset, $\\mathcal{G}^{F}$.\n2. **Upper and Lower Sub-models**: Train the highest fidelity sub-model on this set.\n   Crucially, extract the exact same geometries at the *lower* fidelity, $F-1$,\n   to train the corresponding lower sub-model.\n3. **Going Downwards**: At the next lower fidelity ($f = F-1$), build the\n   new training set by strictly inheriting the geometries selected in step 1, and\n   then randomly sampling the remaining required geometries from the available pool:\n\n\\begin{align}\\mathcal{G}^{F-1} := \\mathcal{G}^{F}_{\\mathrm{inherited}} \\cup \\mathcal{G}^{F-1}_{\\mathrm{random}}\\end{align}\n\n4. **Recursion**: Repeat this nested sampling process recursively downwards until\n   the baseline fidelity ($f = f_b$) is reached.\n\nMore details can be found in the publication:\nMach. Learn.: Sci. Technol. 5 015054 (10.1088/2632-2153/ad2cef)\n\nFinally, all models are evaluated on an unseen test set, and the prediction errors\nare reported as Mean Absolute Errors (MAE) using a discrete $L_1$ norm:\n\n\\begin{align}MAE = \\frac{1}{N_{\\mathrm{test}}}\\sum_{q=1}^{N_{\\mathrm{test}}}\\left\\lvert P_{\\mathrm{ML}}\\left(\\boldsymbol{X}_q^{\\mathrm{ref}}\\right) - {y}^{\\mathrm{ref}}_q\\right\\rvert\\end{align}\n\n\nThe MFML class natively handles the shuffling of training\ndata for cross-validation while maintaining the nested structure\nrequirements for the multifidelity data.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nfrom mfml_qc.datasets import load_benzene_data\nfrom mfml_qc.utils import build_hierarchy_arrays, top_down_subsetting\nfrom mfml_qc.mfml import ModelMFML"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Loading and Formatting Data\nWe load the dataset and establish our fixed test set of the last 2,712 geometries.\nThe remaining 12,288 geometries act as our sampling pool.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "dataset = load_benzene_data()\n\nX_CM = dataset[\"X_CM\"]\ndata = dataset[\"energies\"]\n\n# Fidelities: LC-DFTB, STO-3G, def2-SVP, def2-TZVP\nhierarchy_cols = [2, 3, 6, 7]\n\n# Establish fixed train/test masks\ntrain_mask = data[:, 0] < 12288\ntest_mask = data[:, 0] >= 12288\n\nX_train_parent = X_CM[train_mask]\nX_test = X_CM[test_mask]\ndata_train = data[train_mask]\ndata_test = data[test_mask]\n\ny_trains, indexes, means = build_hierarchy_arrays(data_train, hierarchy_cols)\n\n# Extract test (reference) energies\ny_test_true = data_test[:, hierarchy_cols[-1]]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generating the Learning Curve\nWe will generate a learning curve by training the MFML model at different\ntraining set sizes. We simply\ndouble the training size at each lower fidelity (e.g., if the highest fidelity\ngets 64 samples, the hierarchy will be 512, 256, 128, 64).\n\nWe perform this for ``navg = 5`` random shuffles to calculate variance.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# High fidelity target sizes\nhf_train_sizes = 2 ** np.arange(1, 9)\nnavg = 5\n\n# Storage array for MAEs: shape (num_sizes, navg)\nmfml_maes = np.zeros((len(hf_train_sizes), navg))\n\n# Initialize a base model once\nmfml_model = ModelMFML(kernel=\"matern\", sigma=715.0, reg=1e-9, p_bar=False)\n\nfor n in range(navg):\n    for s_idx, hf_size in enumerate(hf_train_sizes):\n        # Define the target hierarchy sizes: [8x, 4x, 2x, 1x]\n        n_trains_target = [hf_size * 8, hf_size * 4, hf_size * 2, hf_size]\n        # Nested shuffling of data\n        subset_y_trains, subset_indexes = top_down_subsetting(\n            y_trains, indexes, n_trains_target, seed=42 + n\n        )\n        # Train the MFML model\n        mfml_model.train(\n            X_train_parent=X_train_parent,\n            y_trains=subset_y_trains,\n            indexes=subset_indexes,\n        )\n        # Predict on the test set\n        preds = mfml_model.predict(X_test=X_test, optimiser=\"default\")\n        preds += means[-1]  # Un-center predictions\n\n        # MAE in kcal/mol\n        mae = np.mean(np.abs(preds - y_test_true)) * 23\n\n        mfml_maes[s_idx, n] = mae"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualizing the Learning Curve\nFinally, we can plot the average MAE with standard deviation error bars\nagainst the number of high-fidelity training samples.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "mean_mae = np.mean(mfml_maes, axis=1)\nstd_mae = np.std(mfml_maes, axis=1)\n\nplt.figure(figsize=(5, 4))\n\n# individual learning curves\nfor n in range(navg):\n    plt.plot(\n        hf_train_sizes,\n        mfml_maes[:, n],\n        color=\"lightgray\",\n        linestyle=\"-\",\n        alpha=0.7,\n        zorder=1,\n        label=\"Individual Runs\" if n == 0 else None,\n    )\n\n# the average learning curve\nplt.errorbar(\n    hf_train_sizes,\n    mean_mae,\n    yerr=std_mae,\n    fmt=\"s-\",\n    color=\"darkorange\",\n    ecolor=\"gray\",\n    capsize=5,\n    linewidth=2,\n    label=\"Averaged Run\",\n)\n\n# plot cosmetics\nplt.xscale(\"log\")\nplt.yscale(\"log\")\nplt.xticks(hf_train_sizes, labels=[str(s) for s in hf_train_sizes])\n\nplt.xlabel(r\"Target Fidelity Training Size ($N_{\\rm train}^{TZVP}$)\")\nplt.ylabel(\"Mean Absolute Error (kcal/mol)\")\nplt.title(\"MFML Nested Cross-Validation Learning Curve\")\nplt.legend()\nplt.grid(True, which=\"both\", ls=\"--\", alpha=0.4)\nplt.tight_layout()\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
}