{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Optimized MFML\nOptimized MFML is an extension of MFML which uses data-adaptive combination of sub-models for superior model precision. Further details on the mechanism can be found in Mach. Learn.: Sci. Technol. 5 015054 (10.1088/2632-2153/ad2cef).\n\nThis example demonstrates the use of optimized multifidelity machine learning (o-MFML) with a comparison of learning curves for:\n\n1. **Single-Fidelity Kernel Ridge Regression (SF-KRR)**: Trained exclusively on high-fidelity data.\n2. **Multi-Fidelity Machine Learning (MFML)**: Uses the default summation of sub-model predictions.\n3. **Optimized MFML (o-MFML)**: Uses an Ordinary Least Squares (OLS) optimizer to learn the optimal combination weights for the sub-models based on a validation set.\n\n## Data Setup\n- **Training Pool**: 12,288 samples (used to draw nested training subsets).\n- **Validation Set**: 712 high-fidelity samples (used strictly validation set for the o-MFML OLS optimizer).\n- **Test Set**: 2,000 samples (used for final MAE evaluation).\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\nfrom mfml_qc.utils import build_hierarchy_arrays, top_down_subsetting\nfrom mfml_qc.krr import KRR\nfrom mfml_qc.mfml import ModelMFML"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Loading and Splitting the Dataset\nWe load the Benzene trajectory and explicitly slice it into the three\nrequired cohorts: Training Pool, Validation Set, and Test Set.\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 the strict boundaries\ntrain_mask = data[:, 0] < 12288\nval_mask = (data[:, 0] >= 12288) & (data[:, 0] < 13000)\ntest_mask = data[:, 0] >= 13000\n\n# training data\nX_train_parent = X_CM[train_mask]\ndata_train = data[train_mask]\n\n# validation data\nX_val = X_CM[val_mask]\ny_val_true = data[val_mask, hierarchy_cols[-1]]\n\n# Test Set\nX_test = X_CM[test_mask]\ny_test_true = data[test_mask, hierarchy_cols[-1]]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Preparing the Data Arrays\nWe extract the fully populated hierarchy arrays for the training pool.\nWe also mean-center the validation targets using the exact same\nhigh-fidelity training mean.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "y_trains, indexes, means = build_hierarchy_arrays(data_train, hierarchy_cols)\ny_val_centered = y_val_true - means[-1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generating the Learning Curves\nWe will loop over increasing high-fidelity training capacities. At each step,\nwe evaluate all three models and average the results over multiple random shuffles.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "hf_train_sizes = 2 ** np.arange(1, 9)\nnavg = 5\n\n# Arrays to store the Mean Absolute Errors\nsf_maes = np.zeros((len(hf_train_sizes), navg))\nmfml_maes = np.zeros((len(hf_train_sizes), navg))\nomfml_maes = np.zeros((len(hf_train_sizes), navg))\n\nfor n in range(navg):\n\n    for s_idx, hf_size in enumerate(hf_train_sizes):\n\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\n        # Perform top-down cascade to guarantee nestedness\n        subset_y, subset_idx = top_down_subsetting(\n            y_trains, indexes, n_trains_target, seed=42 + n\n        )\n\n        # single fidelity KRR\n        X_sf = X_train_parent[subset_idx[-1][:, 0]]\n        y_sf = subset_y[-1]\n\n        krr_sf = KRR(kernel_type=\"matern\", sigma=715.0, reg=1e-9)\n        krr_sf.train(X_sf, y_sf)\n\n        preds_sf = krr_sf.predict(X_test) + means[-1]\n        sf_maes[s_idx, n] = np.mean(np.abs(preds_sf - y_test_true)) * 23\n\n        # basic MFML\n        mfml_model = ModelMFML(kernel=\"matern\", sigma=715.0, reg=1e-9, p_bar=False)\n        mfml_model.train(\n            X_train_parent=X_train_parent, y_trains=subset_y, indexes=subset_idx\n        )\n\n        preds_mfml = mfml_model.predict(X_test=X_test, optimiser=\"default\") + means[-1]\n        mfml_maes[s_idx, n] = np.mean(np.abs(preds_mfml - y_test_true)) * 23\n\n        # o-MFML\n        preds_omfml = (\n            mfml_model.predict(\n                X_test=X_test, X_val=X_val, y_val=y_val_centered, optimiser=\"OLS\"\n            )\n            + means[-1]\n        )\n\n        omfml_maes[s_idx, n] = np.mean(np.abs(preds_omfml - y_test_true)) * 23"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualizing the Comparison\nFinally, we plot the averaged learning curves to see how standard MFML\ndominates at extremely low data regimes, and how o-MFML optimizes the\nweights as the training set size increases.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Calculate means over the random shuffles\nmean_sf = np.mean(sf_maes, axis=1)\nmean_mfml = np.mean(mfml_maes, axis=1)\nmean_omfml = np.mean(omfml_maes, axis=1)\n\n# Calculate standard deviations for error bars\nstd_sf = np.std(sf_maes, axis=1)\nstd_mfml = np.std(mfml_maes, axis=1)\nstd_omfml = np.std(omfml_maes, axis=1)\n\nplt.figure(figsize=(5, 4))\n\n# Plotting the three curves\nplt.errorbar(\n    hf_train_sizes,\n    mean_sf,\n    yerr=std_sf,\n    fmt=\"o-\",\n    color=\"gray\",\n    capsize=5,\n    linewidth=2,\n    label=\"Single-Fidelity KRR\",\n)\n\nplt.errorbar(\n    hf_train_sizes,\n    mean_mfml,\n    yerr=std_mfml,\n    fmt=\"s-\",\n    color=\"dodgerblue\",\n    capsize=5,\n    linewidth=2,\n    label=\"MFML\",\n)\n\nplt.errorbar(\n    hf_train_sizes,\n    mean_omfml,\n    yerr=std_omfml,\n    fmt=\"d-\",\n    color=\"darkorange\",\n    capsize=5,\n    linewidth=2,\n    label=\"o-MFML (OLS)\",\n)\n\n# Formatting the plot\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_{\\mathrm{train}}^{\\mathrm{TZVP}}$)\")\nplt.ylabel(\"Test Set MAE (kcal/mol)\")\nplt.title(\"Learning Curves\")\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
}