{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Active Learning Protocols\n\nIn machine learning for quantum chemistry, molecular geometries\nare often abundant but labeling them with the properties (e.g., running\nexpensive quantum chemistry calculations for excitation energies) is costly.\nActive Learning (AL) solves this by iteratively selecting\nonly the most informative geometries from an unlabeled pool\nto add to the training set.\n\nThis tutorial demonstrates how to build a custom Active Learning loop using\nthree different Uncertainty Quantification (UQ) metrics:\n\n1. **GPR Variance**: Uses the exact analytical predictive variance of the kernel.\n2. **Ensemble Variance**: Uses the variance of predictions across an ensemble of models.\n3. **LFaB (Low Fidelity as Bias)**: Uses the prediction error on a cheap, low-fidelity baseline property.\n\nNote that GPR Variance and Ensemble Variance rely *only* on the structural features (Coulomb Matrices)\nof the pool to estimate uncertainty. **LFaB**, however, strictly requires that the cheap,\nlow-fidelity baseline property has already been computed for the entire pool.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Setting Up the Data\nWe will use the Benzene dataset. Our target (high-fidelity) property to learn\nis LC-DFTB, and our cheap baseline (low-fidelity) property for LFaB is ZINDO.\nBoth energies are in eV.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import train_test_split\nfrom mfml_qc.datasets import load_benzene_data\nfrom mfml_qc.krr import KRR\nfrom mfml_qc.active_learning import gpr_variance, ensemble_variance, lfab\n\ndataset = load_benzene_data()\n\nX_CM = dataset[\"X_CM\"]\ndata = dataset[\"energies\"]\n\n# Index 1: ZINDO (Low Fidelity)\n# Index 2: LC-DFTB (Target High Fidelity)\nzindo_energies = data[:, 1]\nlcdftb_energies = data[:, 2]\n\n# Filter out any rows where either calculation failed (NaNs)\nvalid_mask = ~(np.isnan(zindo_energies) | np.isnan(lcdftb_energies))\nX_valid = X_CM[valid_mask]\ny_low_valid = zindo_energies[valid_mask]\ny_target_valid = lcdftb_energies[valid_mask]\n\n# Split into a Training Pool (First 12,288) and a fixed Test Set\npool_mask = data[valid_mask, 0] < 12288\ntest_mask = data[valid_mask, 0] >= 12288\n\nX_pool_master = X_valid[pool_mask]\ny_low_pool_master = y_low_valid[pool_mask]\ny_target_pool_master = y_target_valid[pool_mask]\n\nX_test = X_valid[test_mask]\ny_test_true = y_target_valid[test_mask]\n\n# Mean-center the target fidelity\ny_target_mean = np.mean(y_target_pool_master)\ny_target_pool_master = y_target_pool_master - y_target_mean\ny_test_centered = y_test_true - y_target_mean\n\n# Mean-center the low-fidelity\ny_low_mean = np.mean(y_low_pool_master)\ny_low_pool_master = y_low_pool_master - y_low_mean\n\n# Initialize the starting training set with 10 random samples from the pool\n(\n    X_train_init,\n    X_pool_init,\n    y_train_init,\n    y_pool_init,\n    y_train_low_init,\n    y_pool_low_init,\n) = train_test_split(\n    X_pool_master,\n    y_target_pool_master,\n    y_low_pool_master,\n    train_size=10,\n    random_state=42,\n)\n\nprint(f\"Test Set Size:    {X_test.shape[0]} samples\")\nprint(f\"Total Pool Size:  {X_pool_init.shape[0]} samples\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Active Learning Loop\nWe write a flexible loop that runs for a fixed number of iterations.\nAt each step, it computes the chosen UQ metric for the entire pool, selects\nthe geometry with the highest uncertainty, moves it into the training set,\nand evaluates the new model's Mean Absolute Error (MAE) on the test set.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def run_active_learning(uq_mode: str, iters: int = 500) -> list:\n    \"\"\"Active learning loop\"\"\"\n\n    X_tr = X_train_init.copy()\n    y_tr = y_train_init.copy()\n    y_tr_low = y_train_low_init.copy()\n\n    X_p = X_pool_init.copy()\n    y_p = y_pool_init.copy()\n    y_p_low = y_pool_low_init.copy()\n\n    # We will use the inbuilt KRR model as the ML architecture\n    base_model = KRR(kernel_type=\"matern\", sigma=715.0, reg=1e-9)\n    maes = []\n\n    for i in range(iters):\n        # uncertainty function caller\n        if uq_mode == \"gpr\":\n            # GPR Variance only needs the structural representations\n            # we will use the same kernel parameters as the main KRR model\n            uq = gpr_variance(\n                X_train=X_tr, X_pool=X_p, kernel_type=\"matern\", sigma=715.0\n            )\n\n        elif uq_mode == \"ensemble\":\n            # Ensemble of models with 85% of training data each\n            uq = ensemble_variance(\n                base_model, X_p, X_tr, y_tr, n_ensemble=5, train_size=0.85\n            )\n\n        elif uq_mode == \"lfab\":\n            # LFaB strictly requires the low-fidelity baseline property for the pool\n            uq = lfab(base_model, X_p, y_p_low, X_tr, y_tr_low)\n\n        # geometry with the maximum uncertainty\n        # in principle you can add a batch of samples\n        # but we will stick to adding one sample at a time\n        best_idx = np.argmax(uq)\n\n        # update training data\n        X_tr = np.vstack([X_tr, X_p[best_idx]])\n\n        # in practice this is exactly where\n        # one would run the QC calculator to\n        # generate the labels\n        y_tr = np.append(y_tr, y_p[best_idx])\n        y_tr_low = np.append(y_tr_low, y_p_low[best_idx])\n\n        # delete from pool\n        X_p = np.delete(X_p, best_idx, axis=0)\n        y_p = np.delete(y_p, best_idx)\n        y_p_low = np.delete(y_p_low, best_idx)\n\n        # Train updated model\n        base_model.train(X_tr, y_tr)\n        preds = base_model.predict(X_test)\n\n        mae = np.mean(np.abs(preds - y_test_centered))\n        maes.append(mae)\n        if i % 50 == 0:\n            print(f\"{uq_mode} ---- Iteration {i}/{iters} ---- MAE: {mae}\")\n\n    return maes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Running the Active Learning Protocols\nWe will now run 200 iterations for each of the three protocols to observe\nwhich data acquisition strategy identifies the most important molecules quickest!\nIt is also wiorth noting that in practice, one rarely has the higher fidelity pre-computed\nand would generally run the LFaB to identify the best geometries where onw would compute the\nhigh fidelity properties as well. However, here since we use a pre-computed benzene dataset\nwe are going to pretend that each iteration of the AL loop we 'compute' the\nLC-DFTB energies for the selected sample.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "n_iters = 200\n\nmaes_gpr = run_active_learning(\"gpr\", iters=n_iters)\nmaes_ens = run_active_learning(\"ensemble\", iters=n_iters)\nmaes_lfab = run_active_learning(\"lfab\", iters=n_iters)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Learning Curves\nFinally, we plot the Test Set Prediction Error (MAE) against the number\nof samples iteratively added to the training set.\n\nA steeper, faster drop in MAE indicates a more efficient Active Learning strategy!\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "train_sizes = np.arange(11, 11 + n_iters)\n\nplt.figure(figsize=(5, 4))\n\nplt.loglog(train_sizes, maes_gpr, label=\"GPR Variance\", color=\"dodgerblue\")\nplt.loglog(train_sizes, maes_ens, label=\"Ensemble Variance\", color=\"darkorange\")\nplt.loglog(\n    train_sizes, maes_lfab, label=\"LFaB (Low Fidelity as Bias)\", color=\"forestgreen\"\n)\n\nplt.xlabel(\"Training Set Size\")\nplt.ylabel(\"Test Set MAE (eV)\")\nplt.title(\"Active Learning on MD-Benzene (LC-DFTB)\")\nplt.legend()\nplt.grid(True, linestyle=\"--\", alpha=0.6)\n\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
}