{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Duck Typing for MFML\n\nThis example demonstrates the \"duck typing\" flexibility of the ``ModelMFML``\norchestrator. DUck typing is summarized well with 'If it walks like a duck, quacks like a duck, and swims like a duck, it must be a duck.'\nHere, we use it to mean that we can use any ML architecture that has certain attributes.\nWhile MFML-QC provides an ultra-fast, built-in Kernel Ridge\nRegressor (KRR), you are not forced into using it.\n\nBecause the ModelMFML class is flexible, you can pass any custom machine\nlearning model (such as a Neural Network or Random Forest from ``scikit-learn``)\nusing the ``base_estimator`` argument, provided it has standard ``.fit(X, y)`` or ``.train(X,y)``,\nand ``.predict(X)`` methods!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Imports and Helper Functions\nWe reuse the exact same helper functions from the previous tutorial to extract\nand strictly nest our multi-fidelity training data.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nimport numpy as np\nimport matplotlib.pyplot as plt\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\n\n# We will use tha random forest regressor from scikit learn\nfrom sklearn.ensemble import RandomForestRegressor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Loading and Splitting Data\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\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\ny_trains, indexes, means = build_hierarchy_arrays(data_train, hierarchy_cols)\n\n# We use the same target sample sizes as the previous tutorial\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": [
        "## Passing a Custom Base Estimator\nInstead of initializing the default Kernel Ridge Regressor, we will initialize\na Random Forest Regressor from ``scikit-learn``.\nBefore we get into the code, let's briefly review the mechanics of a\nRandom Forest Regressor. Unlike Kernel Ridge Regression, which relies on\nmeasuring continuous structural distances (kernels), a Random Forest is an\nensemble method constructed from multiple decision trees.\n\nFor a given input geometry representation $x$, each individual decision\ntree $k$ in the forest makes an independent prediction $h_k(x)$.\nThe final ensemble prediction $\\hat{y}$ is simply the average of all\n$K$ trees:\n\n\\begin{align}\\hat{y}(x) = \\frac{1}{K} \\sum_{k=1}^{K} h_k(x)\\end{align}\n\nBy training each tree on a random subset of the data (bootstrapping) and\nconsidering a random subset of features at each split, the Random Forest\ndrastically reduces the variance and overfitting that is common in single\ndecision trees.\n\nNotice how we pass the initialized ``sk_model`` directly to the ``base_estimator``\nargument of the ``ModelMFML`` object. The script will automatically\nduplicate this tree-based model for all 2N-1 required sub-models!\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "sk_model = RandomForestRegressor(\n    n_estimators=50, max_depth=10, random_state=42, n_jobs=-1\n)\n\n# the mfml model class takes the random forest regressor as the base-estimator\nmfml_model = ModelMFML(base_estimator=sk_model, p_bar=False)\n\n# The class internally handles calling `.fit()` on the Random Forest models natively\nmfml_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\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "y_test_true = data_test[:, hierarchy_cols[-1]]\n\n# the `default` optimiser carries out basic MFML\npreds = mfml_model.predict(X_test=X_test, optimiser=\"default\")\n\n# Un-center the predictions\npreds += means[-1]\n\nmae = np.mean(np.abs(preds - y_test_true))\nmae_kcal = mae * 23  # eV to kcal/mol\n\nprint(f\"MFML (Random Forest) Test Set MAE: {mae:.6f} eV ({mae_kcal:.4f} kcal/mol)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Parity Plot\nWe can visualize the performance of our Random Forest-based MFML model.\nWhile kernels generally perform better for molecular representations like\nthe Coulomb Matrix, this example shows that the framework is completely\nmodel-agnostic and can be used in a modular fashion.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(5, 5))\n\n# Plot the scatter points\nplt.scatter(y_test_true, preds, alpha=0.6, color=\"lightcoral\", edgecolor=\"k\", s=25)\n\n# Calculate the limits to draw a perfect 45-degree diagonal line\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-Random Forest)\")\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
}