{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# A Multi-Fidelity Dataset: MD Benzene\n\nThis tutorial covers the basics of data handling, visualization, and structural\nrepresentation generation in MFML-QC. We will use an inbuilt dataset for the benzene molecule.\n\n## Dataset Details\nThe dataset for benzene is taken from the work presented in\nJ. Chem. Theory Comput. 2023, 19, 21, 7658\u20137670 (10.1021/acs.jctc.3c00882).\nThe dataset consists of 15,000 fs MD-simulation of benzene with details of simulation presented\nin the article above. For the geometries of this simulation, a sparse multifidelity dataset of excitation energies is provided.\nBelow (and in a few more tutorials) we will interact with this toy dataset to better understand the MFML-4-QC package.\n\n## Downloading the Dataset\nThe Benzene dataset is built into the package to allow for quick benchmarking.\nIf you encounter a ``FileNotFoundError`` during this tutorial, ensure you have\ninstalled the package using ``pip install .`` from the root directory so the\ndataset is correctly packaged, or manually place the data files inside the\n``data/benzene`` directory.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Loading the Built-in Dataset\nWe start by loading the built-in dataset using our provided utility function.\nThis automatically handles path resolution and caching for you. Note that a\ncached unsorted Coulomb Matrix for all 15,000 geometries will be generated.\n\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\n\ndataset = load_benzene_data()\n\n# Let's inspect what the dataset dictionary contains.\nprint(f\"Keys in dataset: {list(dataset.keys())}\")\nprint(f\"Columns in CSV: {dataset['columns']}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Exploring the Data\nThe data consists of vertical excitation energies computed at multiple fidelities\nand their corresponding computational time costs. Let's extract them and look\nat the distributions.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "energies = dataset[\"energies\"]\ntimecosts = dataset[\"timecosts\"]\n\n# From our columns list, LC-DFTB is at index 2, and def2-TZVP is at index 7.\nlc_dftb_energies = energies[:, 2]\ndef2_tzvp_energies = energies[:, 7]\n\n# Drop NaN values for the plot\n# (high-fidelity calculations are sparsely computed)\nlc_dftb_clean = lc_dftb_energies[~np.isnan(lc_dftb_energies)]\ndef2_tzvp_clean = def2_tzvp_energies[~np.isnan(def2_tzvp_energies)]\n\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 4))\n\nax1.hist(lc_dftb_clean, bins=50, color=\"skyblue\", edgecolor=\"black\")\nax1.set_title(\"TD-LC-DFTB Energies\")\nax1.set_xlabel(\"Energy (eV)\")\nax1.set_ylabel(\"Frequency\")\n\nax2.hist(def2_tzvp_clean, bins=50, color=\"salmon\", edgecolor=\"black\")\nax2.set_title(\"TD-DFT def2-TZVP Energies\")\nax2.set_xlabel(\"Energy (eV)\")\n\nplt.tight_layout()\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Why Multi-Fidelity Machine Learning?\nMulti-fidelity machine learning relies on the fact that lower levels of theory\nare significantly cheaper to compute. Let's visualize the computational time cost\ndistributions across all fidelities to see if we can make a case for MFML.\n\nTo ensure a fair comparison, we will extract a strict subset of geometries from\nthe training set (the first 12,288 samples) where the most expensive\ncalculation (def2-QZVP) is present.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "train_energies = energies[:12288]\ntrain_timecosts = timecosts[:12288]\n\n# def2-QZVP is at column index 8.\n# We create a mask for rows where it is present.\nvalid_qzvp_mask = ~np.isnan(train_energies[:, 8])\nfiltered_timecosts = train_timecosts[valid_qzvp_mask]\n\nprint(\n    f\"Found {np.sum(valid_qzvp_mask)} training geometries with complete hierarchy data.\"\n)\n\nmethods = [\n    \"ZINDO\",\n    \"LC-DFTB\",\n    \"STO-3G\",\n    \"3-21G\",\n    \"6-31G\",\n    \"def2-SVP\",\n    \"def2-TZVP\",\n    \"def2-QZVP\",\n]\nmethod_indices = [1, 2, 3, 4, 5, 6, 7, 8]\n\n# Calculate means and standard deviations\nmeans = [np.mean(filtered_timecosts[:, i]) for i in method_indices]\nstds = [np.std(filtered_timecosts[:, i]) for i in method_indices]\n\n# ---------------------------------------------------------\n# Plotting Mean and Std Dev\n# ---------------------------------------------------------\nplt.figure(figsize=(8, 4))\n\n# Plotting with error bars\nplt.errorbar(\n    range(len(methods)),\n    means,\n    yerr=stds,\n    fmt=\"o-\",\n    capsize=5,\n    elinewidth=2,\n    markersize=8,\n    color=\"darkblue\",\n    ecolor=\"gray\",\n)\n\nplt.xticks(ticks=range(len(methods)), labels=methods)\nplt.yscale(\"log\")\nplt.title(\"Mean Computational Cost per Fidelity\")\nplt.ylabel(\"Time (seconds, log scale)\")\nplt.grid(axis=\"y\", linestyle=\"--\", alpha=0.7)\nplt.tight_layout()\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generating Molecular Representations\nMachine learning models generally cannot learn from raw 3D Cartesian (XYZ) coordinates directly,\nas they are not translationally or rotationally invariant. Instead, we map the\n3D structures into a molecular descriptor called the **Coulomb Matrix (CM)**.\nThis is one of the simplest molecular descriptors and we will use this throughout\nthe tutorials.\n\nThe ``load_benzene_data()`` function automatically invoked\nroutines to compute and load these representations for us.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X_CM = dataset[\"X_CM\"]\nprint(f\"Total Geometries: {X_CM.shape[0]}\")\nprint(f\"Coulomb Matrix Features per geometry: {X_CM.shape[1]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Custom Datasets\nIf you are working with your own custom molecular dynamics trajectory (e.g., an\n``.xyz`` file) instead of the built-in dataset, you can generate these representations\nmanually.\n\nThe representations module handles the file parsing and memory allocation automatically:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from mfml_qc.representations import generate_coulomb_matrices\n\n# This is the exact command you would use on your own data:\n# X_custom_CM = generate_coulomb_matrices(\n#     xyz_filepath=\"path/to/your/custom_trajectory.xyz\",\n#     save_path=\"optional/path/to/cache_coulomb_matrices.npy\"\n# )"
      ]
    }
  ],
  "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
}