{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# The Basics of MFML\n\nThis example demonstrates how to use the MFML-QC package to load the\nbuilt-in Benzene trajectory dataset, manually extract a multi-fidelity\nsubset using a top-down approach (that is start with the highest fidelity\nthen move down the fidelity hierarchy), and train an MFML model\nto predict high-fidelity excitation energies.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Imports and Helper Functions\nBefore we begin the actual MFML workflow, let's define two helper functions.\nThese are present in the utils script of the MFML package.\n\nThe first function parses our raw CSV arrays, mean-centers\nthe target properties,and establishes the foundational\n``indexes`` mapping arrays that link high-fidelity calculations back to their\nlow-fidelity entries and corresponding geometries.\n\nThe second helper function handles **Top-Down Nested Extraction**.\nBecause the provided benzene datasets is sparsely populated at high fidelities,\na top-down extraction of the data guarantees that any geometry selected for a high-fidelity\ntraining set is strictly included in all lower-fidelity training sets beneath it.\nNote that in practice, however, one would start from the lowest fidelity and try to\nbuild their way up from there (Such as Adaptive-MFML). This will form a separate tutorial.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nimport numpy as np\nfrom mfml_qc.datasets import load_benzene_data\nfrom mfml_qc.mfml import ModelMFML\nfrom mfml_qc.utils import build_hierarchy_arrays, top_down_subsetting"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Loading and Splitting the Dataset\nLet's load the built-in Benzene dataset and split it chronologically into\na training pool (the first 12,288 geometries) and a test set.\nWe will use a 4-fidelity example which in increasing order of accuracy are\nLC-DFTB, TD-DFT (STO-3G), TD-DFT (def2-SVP), and TD-DFT (def2-TZVP).\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# LC-DFTB (Col 2) -> STO-3G (Col 3) -> def2-SVP (Col 6) -> def2-TZVP (Col 7)\nhierarchy_cols = [2, 3, 6, 7]\nnum_fids = len(hierarchy_cols)\n\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\nprint(data_train.shape, data_test.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Preparing the Multi-Fidelity Data\nWe use our helper functions to extract the full valid arrays, and then apply\nthe top-down nested extraction to enforce our specific target sizes.\nHere we will use 1024 samples at LC-DFTB, 512 at STO-3G, 256 at def2-SVP, and 128 samples at def2-TZVP\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "y_trains, indexes, means = build_hierarchy_arrays(data_train, hierarchy_cols)\nn_trains_target = [1024, 512, 256, 128]\nsubset_y_trains, subset_indexes = top_down_subsetting(\n    y_trains, indexes, n_trains_target, seed=42\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Instantiating the MFML Model\nWith our data strictly nested and formatted, we can instantiate the\n``ModelMFML`` class. It will automatically build and train the\n2N - 1 underlying sub-models required.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "mfml_model = ModelMFML(kernel=\"matern\", sigma=715.0, reg=1e-9, p_bar=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We pass the subset arrays here to train on our specific N_trains targets.\nIn a different example where we carry out a bottom-to-top data nestedness,\nwe will come across an internal shuffler built in for the MFML.\nIt must be noted that shuffling of training data here is what is\noften called random sub-sampling. The order of the data itself in training\nkernel based models does not affect the model.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "mfml_model.train(\n    X_train_parent=X_train_parent, y_trains=subset_y_trains, indexes=subset_indexes\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Predicting and Evaluating\nWe predict the test geometries using the standard +1 and -1\nweights and un-center the predictions using the highest-fidelity mean.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Pick def2-TZVP (highest fidelity) as our \"ground truth\" test targets\ny_test_true = data_test[:, hierarchy_cols[-1]]\npreds = mfml_model.predict(X_test=X_test, optimiser=\"default\")\npreds += means[-1]\n\n# Calculate Mean Absolute Error\nmae = np.mean(np.abs(preds - y_test_true))\nmae_kcal = mae * 23  # Conversion factor: eV to kcal/mol\n\nprint(f\"MFML Test Set MAE: {mae:.6f} eV ({mae_kcal:.4f} kcal/mol)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Often, it is worth visualising a parity plot\nwhich is simply a scatter plot between the predictions\nand the true reference of the test set.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n\nplt.figure(figsize=(5, 5))\nplt.scatter(y_test_true, preds, alpha=0.6, color=\"dodgerblue\", edgecolor=\"k\", s=25)\n\n# the 45 degree diagonal indicates the ideal of everything matches perfectly\nmin_val = min(np.min(y_test_true), np.min(preds))\nmax_val = max(np.max(y_test_true), np.max(preds))\nplt.plot([min_val, max_val], [min_val, max_val], \"k--\", lw=2)\n\nplt.xlabel(\"True def2-TZVP Energy (eV)\")\nplt.ylabel(\"MFML Predicted Energy (eV)\")\nplt.title(\"Parity Plot (MFML-KRR)\")\nplt.grid(True, linestyle=\"--\", alpha=0.6)\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
}